Repository: jcklie/slurmee

Issue:
Job array support
It would greatly help if there was support for job arrays, see https://slurm.schedmd.com/job_array.html.

Discussion:
jcklie: Can you add unit tests and format your code with [black](https://github.com/psf/black) if not done already?

Patch:
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..409afb5
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,107 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+.static_storage/
+.media/
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# pycharm
+.idea/
\ No newline at end of file
diff --git a/slurmee/slurmee.py b/slurmee/slurmee.py
index d716b53..b4e268b 100644
--- a/slurmee/slurmee.py
+++ b/slurmee/slurmee.py
@@ -4,7 +4,7 @@ https://slurm.schedmd.com/sbatch.html#SECTION_OUTPUT-ENVIRONMENT-VARIABLES
 """
 
 import os
-from typing import List, Optional
+from typing import Optional, Dict
 
 
 def get_job_id() -> Optional[int]:
@@ -55,7 +55,7 @@ def get_submit_host() -> Optional[str]:
     """ Gets the name of the submission host.
 
     Returns:
-        job_id: The submission host name if running in slurm else `None`
+        submit_host: The submission host name if running in slurm else `None`
     """
     return os.getenv("SLURM_SUBMIT_HOST")
 
@@ -82,7 +82,7 @@ def get_ntasks() -> Optional[int]:
     """ Get the number of tasks per node.
 
     Returns:
-        ntaskts: The number of tasks per node if running in slurm else `None`.
+        ntasks: The number of tasks per node if running in slurm else `None`.
     """
     return _to_int(os.getenv("SLURM_NTASKS"))
 
@@ -93,11 +93,34 @@ def get_nodeid() -> Optional[int]:
     The node id is an index to node running on relative to nodes assigned to job.
 
     Returns:
-        ntaskts: The node id if running in slurm else `None`.
+        nodeid: The node id if running in slurm else `None`.
     """
 
     return _to_int(os.getenv("SLURM_NODEID"))
 
 
+def get_job_array_info() -> Optional[Dict[str, int]]:
+    """ Get job array information.
+
+    Refer to https://slurm.schedmd.com/job_array.html for interpreting the returned values.
+
+    Returns:
+        job_array_info: A dict with job array information if running in slurm and inside a job array else `None`.
+    :return:
+    """
+    d = {
+        "array_job_id": "SLURM_ARRAY_JOB_ID",
+        "task_id": "SLURM_ARRAY_TASK_ID",
+        "task_count": "SLURM_ARRAY_TASK_COUNT",
+        "task_max": "SLURM_ARRAY_TASK_MAX",
+        "task_min": "SLURM_ARRAY_TASK_MIN",
+    }
+
+    if os.getenv(d["array_job_id"]) is None:
+        return None
+
+    return {k: _to_int(os.getenv(env_var)) for k, env_var in d.items()}
+
+
 def _to_int(s: Optional[str]) -> Optional[int]:
     return int(s) if s else None


Test patch:
diff --git a/tests/test_slurmee.py b/tests/test_slurmee.py
index 0fb9b92..e0192ae 100644
--- a/tests/test_slurmee.py
+++ b/tests/test_slurmee.py
@@ -83,3 +83,32 @@ class TestSlurmee(unittest.TestCase):
     def test_get_nodeid_when_absent(self):
         os.environ.clear()
         self.assertEqual(slurmee.get_nodeid(), None)
+
+
+class TestSlurmeeJobArrays(unittest.TestCase):
+    def tearDown(self):
+        os.environ.clear()
+
+    def test_get_job_array_info(self):
+        os.environ["SLURM_JOB_ID"] = "37"
+        os.environ["SLURM_ARRAY_JOB_ID"] = "36"
+        os.environ["SLURM_ARRAY_TASK_ID"] = "2"
+        os.environ["SLURM_ARRAY_TASK_COUNT"] = "3"
+        os.environ["SLURM_ARRAY_TASK_MAX"] = "3"
+        os.environ["SLURM_ARRAY_TASK_MIN"] = "1"
+
+        self.assertEqual(slurmee.get_job_id(), 37)
+
+        job_array_info = slurmee.get_job_array_info()
+        self.assertEqual(job_array_info["array_job_id"], 36)
+        self.assertEqual(job_array_info["task_id"], 2)
+        self.assertEqual(job_array_info["task_count"], 3)
+        self.assertEqual(job_array_info["task_max"], 3)
+        self.assertEqual(job_array_info["task_min"], 1)
+
+    def test_get_job_array_info_when_not_in_slurm(self):
+        self.assertEqual(slurmee.get_job_array_info(), None)
+
+    def test_get_job_array_info_when_not_in_job_array(self):
+        os.environ["SLURM_JOB_ID"] = "42"
+        self.assertEqual(slurmee.get_job_array_info(), None)