Tweaks and fixes to the SLURM header. (#263)

* Tweaks and fixes to the SLURM header. Adds the ability to specify GPUs, fixes reservation pairing with accounts, and now uses a ChainMap to handle the internal key conflicts between the batch and step settings. Also introduces the exclusive key. Changed to full parameter names for clarity. * Addition of chainmap package for python2.7 * A check to see if procs is included in the header. Fixes #234 and includes ntasks in the header if the batch section includes the key.
LLNL · May 17, 2020 · c5d1f69 · c5d1f69
1 parent 95049f1
commit c5d1f69
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 23 deletions.
diff --git a/maestrowf/datastructures/schemas.json b/maestrowf/datastructures/schemas.json
@@ -42,7 +42,8 @@
           "gpus": {"type": "integer"},
           "cores per task": {"type": "integer"},
           "walltime": {"type": "string", "minLength": 1},
-          "reservation": {"type": "string", "minLength": 1}
+          "reservation": {"type": "string", "minLength": 1},
+          "exclusive": {"type": "boolean"}
         },
         "additionalProperties": false,
         "required": [

diff --git a/maestrowf/interfaces/script/slurmscriptadapter.py b/maestrowf/interfaces/script/slurmscriptadapter.py
@@ -32,6 +32,11 @@
 import logging
 import os
 import re
+# In order to support Python2.7, we need to catch the import error.
+try:
+    from collections import ChainMap
+except ImportError:
+    from chainmap import ChainMap
 
 from maestrowf.abstracts.interfaces import SchedulerScriptAdapter
 from maestrowf.abstracts.enums import JobStatusCode, State, SubmissionCode, \
@@ -74,24 +79,31 @@ def __init__(self, **kwargs):
         self.add_batch_parameter("nodes", kwargs.pop("nodes", "1"))
         self.add_batch_parameter("reservation", kwargs.pop("reservation", ""))
 
+        # Check for procs separately, as we don't want it in the header if it's
+        # not present.
+        procs = kwargs.get("procs", None)
+        if procs:
+            self.add_batch_parameter("procs", procs)
+
         self._header = {
-            "nodes": "#SBATCH -N {nodes}",
-            "queue": "#SBATCH -p {queue}",
-            "bank": "#SBATCH -A {bank}",
-            "walltime": "#SBATCH -t {walltime}",
+            "nodes": "#SBATCH --nodes={nodes}",
+            "queue": "#SBATCH --partition={queue}",
+            "bank": "#SBATCH --account={bank}",
+            "walltime": "#SBATCH --time={walltime}",
             "job-name":
-                "#SBATCH -J \"{job-name}\"\n"
-                "#SBATCH -o \"{job-name}.out\"\n"
-                "#SBATCH -e \"{job-name}.err\"",
-            "comment": "#SBATCH --comment \"{comment}\""
+                "#SBATCH --job-name=\"{job-name}\"\n"
+                "#SBATCH --output=\"{job-name}.out\"\n"
+                "#SBATCH --error=\"{job-name}.err\"",
+            "comment": "#SBATCH --comment \"{comment}\"",
+            "reservation": "#SBATCH --reservation=\"{reservation}\"",
+            "gpus": "#SBATCH --gres=gpu:{gpus}"
         }
 
         self._cmd_flags = {
             "cmd": "srun",
             "depends": "--dependency",
             "ntasks": "-n",
             "nodes": "-N",
-            "reservation": "--reservation",
             "cores per task": "-c",
         }
         self._unsupported = set(["cmd", "depends", "ntasks", "nodes"])
@@ -104,22 +116,26 @@ def get_header(self, step):
         :returns: A string of the header based on internal batch parameters and
             the parameter step.
         """
-        run = dict(step.run)
-        batch_header = dict(self._batch)
-        batch_header["walltime"] = run.pop("walltime")
-        if run["nodes"]:
-            batch_header["nodes"] = run.pop("nodes")
-        batch_header["job-name"] = step.name.replace(" ", "_")
-        batch_header["comment"] = step.description.replace("\n", " ")
+        resources = ChainMap(step.run, self._batch)
+        resources["job-name"] = step.name.replace(" ", "_")
+        resources["comment"] = step.description.replace("\n", " ")
 
         modified_header = ["#!{}".format(self._exec)]
         for key, value in self._header.items():
-            # If we're looking at the bank and the reservation header exists,
-            # skip the bank to prefer the reservation.
-            if key == "bank" and "reservation" in self._batch:
-                if self._batch["reservation"]:
-                    continue
-            modified_header.append(value.format(**batch_header))
+            if key not in resources:
+                continue
+
+            if resources[key]:
+                modified_header.append(value.format(**resources))
+
+        if "procs" in self._batch:
+            modified_header.append(
+                "#SBATCH --ntasks={}".format(resources["procs"])
+            )
+
+        exclusive = resources.get("exclusive", False)
+        if exclusive:
+            modified_header.append("#SBATCH --exclusive")
 
         return "\n".join(modified_header)
 

diff --git a/requirements.txt b/requirements.txt
@@ -6,6 +6,7 @@ enum34; python_version<'3.4'
 dill
 jsonschema>=3.2.0
 coloredlogs
+chainmap; python_version<'3'
 -e .
 
 fabric

diff --git a/samples/lulesh/lulesh_sample1_unix_slurm.yaml b/samples/lulesh/lulesh_sample1_unix_slurm.yaml
@@ -21,6 +21,7 @@ batch:
     bank        : baasic
     queue       : pbatch
     gres        : ignore
+    reservation : test_reservation
 
 study:
     - name: make-lulesh
@@ -43,6 +44,7 @@ study:
           depends: [make-lulesh]
           nodes: 2
           procs: 27
+          exclusive   : True
           walltime: "00:10:00"
 
 global.parameters:

diff --git a/setup.py b/setup.py
@@ -27,6 +27,7 @@
     "dill",
     "jsonschema>=3.2.0",
     "coloredlogs",
+    "chainmap ; python_version<'3'",
   ],
   extras_require={},
   long_description_content_type='text/markdown',