Merge pull request #1363 from teojgo/bugfix/slurm_combined_constraints

[bugfix] Combine access and job options for slurm '--constraint'

Author: Vasileios Karakasis

docs/manpage.rst

@@ -337,11 +337,19 @@ Options controlling job submission
 .. option:: -J, --job-option=OPTION
 
    Pass ``OPTION`` directly to the job scheduler backend.
-   The syntax for this option is ``-J key=value``.
-   If ``key`` starts with ``-`` or ``#``, the option will be passed verbatim to the job script.
-   Otherwise, ReFrame will add ``-`` or ``--`` as well as the directive corresponding to the current scheduler.
-   This option will be emitted after any options specified in the :js:attr:`access` system partition configuration parameter.
-
+   The syntax of ``OPTION`` is ``-J key=value``.
+   If ``OPTION`` starts with ``-`` it will be passed verbatim to the backend job scheduler.
+   If ``OPTION`` starts with ``#`` it will be emitted verbatim in the job script.
+   Otherwise, ReFrame will pass ``--key=value`` or ``-k value`` (if ``key`` is a single character) to the backend scheduler.
+   Any job options specified with this command-line option will be emitted after any job options specified in the :js:attr:`access` system partition configuration parameter.
+
+   Especially for the Slurm backends, constraint options, such as ``-J constraint=value``, ``-J C=value``, ``-J --constraint=value`` or ``-J -C=value``, are going to be combined with any constraint options specified in the :js:attr:`access` system partition configuration parameter.
+   For example, if ``-C x`` is specified in the :js:attr:`access` and ``-J C=y`` is passed to the command-line, ReFrame will pass ``-C x,y`` as a constraint to the scheduler.
+   Notice, however, that if constraint options are specified through multiple :option:`-J` options, only the last one will be considered.
+   If you wish to completely overwrite any constraint options passed in :js:attr:`access`, you should consider passing explicitly the Slurm directive with ``-J '#SBATCH --constraint=new'``.
+
+   .. versionchanged:: 3.0
+      This option has become more flexible.
 
 ------------------------
 Flexible node allocation

reframe/core/schedulers/slurm.py

@@ -195,11 +195,35 @@ def emit_preamble(self, job):
             hint = 'multithread' if job.use_smt else 'nomultithread'
 
         for opt in job.sched_access:
-            preamble.append('%s %s' % (self._prefix, opt))
+            if not opt.strip().startswith(('-C', '--constraint')):
+                preamble.append('%s %s' % (self._prefix, opt))
+
+        constraints = []
+        constraint_parser = ArgumentParser()
+        constraint_parser.add_argument('-C', '--constraint')
+        parsed_options, _ = constraint_parser.parse_known_args(
+            job.sched_access)
+        if parsed_options.constraint:
+            constraints.append(parsed_options.constraint.strip())
+
+        # NOTE: Here last of the passed --constraint job options is taken
+        # into account in order to respect the behavior of slurm.
+        parsed_options, _ = constraint_parser.parse_known_args(job.options)
+        if parsed_options.constraint:
+            constraints.append(parsed_options.constraint.strip())
+
+        if constraints:
+            preamble.append(
+                self._format_option(','.join(constraints), '--constraint={0}')
+            )
 
         preamble.append(self._format_option(hint, '--hint={0}'))
         prefix_patt = re.compile(r'(#\w+)')
         for opt in job.options:
+            if opt.strip().startswith(('-C', '--constraint')):
+                # Constraints are already processed
+                continue
+
             if not prefix_patt.match(opt):
                 preamble.append('%s %s' % (self._prefix, opt))
             else:

unittests/test_schedulers.py

@@ -67,7 +67,7 @@ def exec_ctx(temp_runtime, scheduler):
     next(rt)
     if scheduler.registered_name == 'squeue':
         # slurm backend fulfills the functionality of the squeue backend, so
-        # if squeue is not configured, use slurrm instead
+        # if squeue is not configured, use slurm instead
         partition = (fixtures.partition_by_scheduler('squeue') or
                      fixtures.partition_by_scheduler('slurm'))
     else:
@@ -370,6 +370,40 @@ def test_no_empty_lines_in_preamble(minimal_job):
         assert line != ''
 
 
+def test_combined_access_constraint(make_job, slurm_only):
+    job = make_job(sched_access=['--constraint=c1'])
+    job.options = ['-C c2,c3']
+    prepare_job(job)
+    with open(job.script_filename) as fp:
+        script_content = fp.read()
+
+    assert re.search(r'(?m)--constraint=c1,c2,c3$', script_content)
+    assert re.search(r'(?m)--constraint=(c1|c2,c3)$', script_content) is None
+
+
+def test_combined_access_multiple_constraints(make_job, slurm_only):
+    job = make_job(sched_access=['--constraint=c1'])
+    job.options = ['--constraint=c2', '-C c3']
+    prepare_job(job)
+    with open(job.script_filename) as fp:
+        script_content = fp.read()
+
+    assert re.search(r'(?m)--constraint=c1,c3$', script_content)
+    assert re.search(r'(?m)--constraint=(c1|c2|c3)$', script_content) is None
+
+
+def test_combined_access_verbatim_constraint(make_job, slurm_only):
+    job = make_job(sched_access=['--constraint=c1'])
+    job.options = ['#SBATCH --constraint=c2', '#SBATCH -C c3']
+    prepare_job(job)
+    with open(job.script_filename) as fp:
+        script_content = fp.read()
+
+    assert re.search(r'(?m)--constraint=c1$', script_content)
+    assert re.search(r'(?m)^#SBATCH --constraint=c2$', script_content)
+    assert re.search(r'(?m)^#SBATCH -C c3$', script_content)
+
+
 def test_guess_num_tasks(minimal_job, scheduler):
     minimal_job.num_tasks = 0
     if scheduler.registered_name == 'local':