Merge pull request #1356 from rsarm/feat/slurm-N-option

[feat] Add configuration option for always emitting `--nodes` in Slurm job scripts

Author: Vasileios Karakasis

docs/config_reference.rst

@@ -960,6 +960,14 @@ Common scheduler options
    A list of systems or system/partitions combinations that this scheduler configuration is valid for.
    For a detailed description of this property, you may refer `here <#.environments[].target_systems>`__.
 
+.. js:attribute:: .schedulers[].use_nodes_option
+
+   :required: No
+   :default: ``false``
+
+   Always emit the ``--nodes`` Slurm option in the preamble of the job script.
+   This option is relevant to Slurm backends only.
+
 
 .. js:attribute:: .schedulers[].ignore_reqnodenotavail
 

reframe/core/schedulers/slurm.py

@@ -112,6 +112,9 @@ def __init__(self):
         self._job_submit_timeout = rt.runtime().get_option(
             f'schedulers/@{self.registered_name}/job_submit_timeout'
         )
+        self._use_nodes_opt = rt.runtime().get_option(
+            f'schedulers/@{self.registered_name}/use_nodes_option'
+        )
 
     def completion_time(self, job):
         if (self._completion_time or
@@ -182,6 +185,10 @@ def emit_preamble(self, job):
                 self._format_option(job.sched_exclusive_access, '--exclusive')
             )
 
+        if self._use_nodes_opt:
+            num_nodes = job.num_tasks // job.num_tasks_per_node
+            preamble.append(self._format_option(num_nodes, '--nodes={0}'))
+
         if job.use_smt is None:
             hint = None
         else:

reframe/schemas/config.json

@@ -279,7 +279,8 @@
                     },
                     "ignore_reqnodenotavail": {"type": "boolean"},
                     "job_submit_timeout": {"type": "number"},
-                    "target_systems": {"$ref": "#/defs/system_ref"}
+                    "target_systems": {"$ref": "#/defs/system_ref"},
+                    "use_nodes_option": {"type": "boolean"}
                 },
                 "required": ["name"],
                 "additionalProperties": false
@@ -424,6 +425,7 @@
         "schedulers/ignore_reqnodenotavail": false,
         "schedulers/job_submit_timeout": 60,
         "schedulers/target_systems": ["*"],
+        "schedulers/use_nodes_option": false,
         "systems/descr": "",
         "systems/modules_system": "nomod",
         "systems/modules": [],

unittests/test_schedulers.py

@@ -261,6 +261,18 @@ def test_prepare_without_smt(fake_job, slurm_only):
         assert re.search(r'--hint=nomultithread', fp.read()) is not None
 
 
+def test_prepare_nodes_option(temp_runtime, make_job, slurm_only):
+    rt = temp_runtime(fixtures.TEST_CONFIG_FILE, 'generic',
+                      {'schedulers/use_nodes_option': True})
+    next(rt)
+    job = make_job()
+    job.num_tasks = 16
+    job.num_tasks_per_node = 2
+    prepare_job(job)
+    with open(job.script_filename) as fp:
+        assert re.search(r'--nodes=8', fp.read()) is not None
+
+
 def test_submit(make_job, exec_ctx):
     minimal_job = make_job(sched_access=exec_ctx.access)
     prepare_job(minimal_job)