Merge pull request #3156 from ekouts/feat/sched_option_sbatch

[feat] Allow ReFrame to pass the access options in command line instead of the script for Slurm

Author: vkarak

docs/config_reference.rst

@@ -367,6 +367,20 @@ System Partition Configuration
    .. warning::
       This option is broken in 4.0.
 
+
+.. py:attribute:: systems.partitions.sched_options.sched_access_in_submit
+
+   :required: No
+   :default: ``false``
+
+   Normally, ReFrame will pass the :attr:`~config.systems.partitions.access` options to the job script only.
+   When this attribute is ``true`` the options are passed in the submission command instead.
+
+   This option is relevant for the LSF, OAR, PBS and Slurm backends.
+
+   .. versionadded:: 4.7
+
+
 .. py:attribute:: systems.partitions.sched_options.ssh_hosts
 
    :required: No

docs/manpage.rst

@@ -1226,6 +1226,21 @@ Here is an alphabetical list of the environment variables recognized by ReFrame.
 Whenever an environment variable is associated with a configuration option, its default value is omitted as it is the same.
 
 
+.. envvar:: RFM_SCHED_ACCESS_IN_SUBMIT
+
+   Pass access options in the submission command (relevant for LSF, OAR, PBS and Slurm).
+
+   .. table::
+      :align: left
+
+      ================================== ==================
+      Associated command line option     N/A
+      Associated configuration parameter :attr::attr:`~config.systems.partitions.sched_options.sched_access_in_submit`
+      ================================== ==================
+
+.. versionadded:: 4.7
+
+
 .. envvar:: RFM_AUTODETECT_FQDN
 
    Use the fully qualified domain name as the hostname.

reframe/core/schedulers/lsf.py

@@ -27,6 +27,9 @@ class LsfJobScheduler(PbsJobScheduler):
     def __init__(self):
         self._prefix = '#BSUB'
         self._submit_timeout = self.get_option('job_submit_timeout')
+        self._sched_access_in_submit = self.get_option(
+            'sched_access_in_submit'
+        )
 
     def _format_option(self, var, option):
         if var is not None:
@@ -57,8 +60,9 @@ def emit_preamble(self, job):
                 f'{self._prefix} -W {int(job.time_limit // 60)}'
             )
 
-        for opt in job.sched_access:
-            preamble.append(f'{self._prefix} {opt}')
+        if not self._sched_access_in_submit:
+            for opt in job.sched_access:
+                preamble.append(f'{self._prefix} {opt}')
 
         # emit the rest of the options
         options = job.options + job.cli_options
@@ -76,7 +80,13 @@ def emit_preamble(self, job):
 
     def submit(self, job):
         with open(job.script_filename, 'r') as fp:
-            completed = _run_strict('bsub', stdin=fp)
+            cmd_parts = ['bsub']
+            if self._sched_access_in_submit:
+                cmd_parts += job.sched_access
+
+            cmd = ' '.join(cmd_parts)
+            completed = _run_strict(cmd, stdin=fp)
+
         jobid_match = re.search(r'^Job <(?P<jobid>\S+)> is submitted',
                                 completed.stdout)
         if not jobid_match:

reframe/core/schedulers/oar.py

@@ -60,6 +60,9 @@ class OarJobScheduler(PbsJobScheduler):
     def __init__(self):
         self._prefix = '#OAR'
         self._submit_timeout = self.get_option('job_submit_timeout')
+        self._sched_access_in_submit = self.get_option(
+            'sched_access_in_submit'
+        )
 
     def emit_preamble(self, job):
         # host is de-facto nodes and core is number of cores requested per node
@@ -88,8 +91,11 @@ def emit_preamble(self, job):
             num_nodes=num_nodes, num_tasks_per_node=num_tasks_per_node,
         )]
 
+        if not self._sched_access_in_submit:
+            options += job.sched_access
+
         # Emit the rest of the options
-        options += job.sched_access + job.options + job.cli_options
+        options += job.options + job.cli_options
         for opt in options:
             if opt.startswith('#'):
                 preamble.append(opt)
@@ -101,9 +107,13 @@ def emit_preamble(self, job):
     def submit(self, job):
         # OAR batch submission mode needs full path to the job script
         job_script_fullpath = os.path.join(job.workdir, job.script_filename)
+        cmd_parts = ['oarsub']
+        if self._sched_access_in_submit:
+            cmd_parts += job.sched_access
 
         # OAR needs -S to submit job in batch mode
-        cmd = f'oarsub -S {job_script_fullpath}'
+        cmd_parts += ['-S', job_script_fullpath]
+        cmd = ' '.join(cmd_parts)
         completed = _run_strict(cmd, timeout=self._submit_timeout)
         jobid_match = re.search(r'.*OAR_JOB_ID=(?P<jobid>\S+)',
                                 completed.stdout)

reframe/core/schedulers/pbs.py

@@ -76,6 +76,9 @@ class PbsJobScheduler(sched.JobScheduler):
     def __init__(self):
         self._prefix = '#PBS'
         self._submit_timeout = self.get_option('job_submit_timeout')
+        self._sched_access_in_submit = self.get_option(
+            'sched_access_in_submit'
+        )
 
     def _emit_lselect_option(self, job):
         num_tasks = job.num_tasks or 1
@@ -92,7 +95,12 @@ def _emit_lselect_option(self, job):
         # Options starting with `-` are emitted in separate lines
         rem_opts = []
         verb_opts = []
-        for opt in (*job.sched_access, *job.options, *job.cli_options):
+        if self._sched_access_in_submit:
+            all_opts = (*job.options, *job.cli_options)
+        else:
+            all_opts = (*job.sched_access, *job.options, *job.cli_options)
+
+        for opt in all_opts:
             if opt.startswith('-'):
                 rem_opts.append(opt)
             elif opt.startswith('#'):
@@ -139,9 +147,14 @@ def filternodes(self, job, nodes):
                                   'node filtering')
 
     def submit(self, job):
+        cmd_parts = ['qsub']
+        if self._sched_access_in_submit:
+            cmd_parts += job.sched_access
+
         # `-o` and `-e` options are only recognized in command line by the PBS
         # Slurm wrappers.
-        cmd = f'qsub -o {job.stdout} -e {job.stderr} {job.script_filename}'
+        cmd_parts += ['-o', job.stdout, '-e', job.stderr, job.script_filename]
+        cmd = ' '.join(cmd_parts)
         completed = _run_strict(cmd, timeout=self._submit_timeout)
         jobid_match = re.search(r'^(?P<jobid>\S+)', completed.stdout)
         if not jobid_match:

reframe/core/schedulers/slurm.py

@@ -140,6 +140,9 @@ def __init__(self):
         self._submit_timeout = self.get_option('job_submit_timeout')
         self._use_nodes_opt = self.get_option('use_nodes_option')
         self._resubmit_on_errors = self.get_option('resubmit_on_errors')
+        self._sched_access_in_submit = self.get_option(
+            'sched_access_in_submit'
+        )
 
     def make_job(self, *args, **kwargs):
         return _SlurmJob(*args, **kwargs)
@@ -209,39 +212,33 @@ def emit_preamble(self, job):
                 )
             )
 
-        for opt in job.sched_access:
-            if not opt.strip().startswith(('-C', '--constraint')):
-                preamble.append('%s %s' % (self._prefix, opt))
-
-        # To avoid overriding a constraint that's passed into `sched_access`,
-        # we AND it with the `--constraint` option passed either in `options`
-        # or in `cli_options`
-        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 + job.cli_options
-        )
-        if parsed_options.constraint:
-            constraints.append(parsed_options.constraint.strip())
-
-        if constraints:
-            if len(constraints) == 1:
-                constr = constraints[0]
+        # Combine constraints in `sched_access`
+        #
+        # We AND the constraints defined in `sched_access` with those in
+        # either the `job.options` or `job.cli_options`. We essentially "move"
+        # the option from the source option list to `sched_access` as if the
+        # user has specified all the constraint in `sched_access`. We can then
+        # move with the preamble generation or the submission normally.
+        c_parser = ArgumentParser()
+        c_parser.add_argument('-C', '--constraint')
+        access, access_other = c_parser.parse_known_args(job.sched_access)
+        job_opts, other_job_opts = c_parser.parse_known_args(job.options)
+        cli_opts, other_cli_opts = c_parser.parse_known_args(job.cli_options)
+        if access.constraint and (job_opts.constraint or cli_opts.constraint):
+            constraints = [access.constraint]
+            if job_opts.constraint:
+                constraints.append(job_opts.constraint)
+                job.options = other_job_opts
             else:
-                # Parenthesize the constraints prior to joining them with `&`
-                # to make sure that precedence is respected.
-                constr = '&'.join(f'({c})' for c in constraints)
+                constraints.append(cli_opts.constraint)
+                job._cli_options = other_cli_opts
 
-            preamble.append(self._format_option(constr, '--constraint={0}'))
+            arg = '&'.join(f'({c.strip()})' for c in constraints)
+            job._sched_access = [f'--constraint={arg}']
+
+        if not self._sched_access_in_submit:
+            for opt in job.sched_access:
+                preamble.append(f'{self._prefix} {opt}')
 
         preamble.append(self._format_option(hint, '--hint={0}'))
         prefix_patt = re.compile(r'(#\w+)')
@@ -259,7 +256,12 @@ def emit_preamble(self, job):
         return list(filter(None, preamble))
 
     def submit(self, job):
-        cmd = f'sbatch {job.script_filename}'
+        cmd_parts = ['sbatch']
+        if self._sched_access_in_submit:
+            cmd_parts += job.sched_access
+
+        cmd_parts += [job.script_filename]
+        cmd = ' '.join(cmd_parts)
         intervals = itertools.cycle([1, 2, 3])
         while True:
             try:

reframe/frontend/cli.py

@@ -582,6 +582,13 @@ def main():
     )
 
     # Options not associated with command-line arguments
+    argparser.add_argument(
+        dest='sched_access_in_submit',
+        envvar='RFM_SCHED_ACCESS_IN_SUBMIT',
+        configvar='systems*/sched_options/sched_access_in_submit',
+        action='store_true',
+        help='Pass access options in the submission command (only for Slurm)'
+    )
     argparser.add_argument(
         dest='autodetect_fqdn',
         envvar='RFM_AUTODETECT_FQDN',

reframe/schemas/config.json

@@ -109,6 +109,7 @@
         "sched_options": {
             "type": "object",
             "properties": {
+                "sched_access_in_submit": {"type": "boolean"},
                 "hosts": {
                     "type": "array",
                     "items": {"type": "string"}
@@ -635,6 +636,7 @@
         "systems/partitions/time_limit": null,
         "systems/partitions/devices": [],
         "systems/partitions/extras": {},
+        "systems*/sched_options/sched_access_in_submit": false,
         "systems*/sched_options/ssh_hosts": [],
         "systems*/sched_options/ignore_reqnodenotavail": false,
         "systems*/sched_options/job_submit_timeout": 60,

unittests/test_schedulers.py

@@ -618,6 +618,7 @@ def test_combined_access_constraint(make_job, slurm_only):
     with open(job.script_filename) as fp:
         script_content = fp.read()
 
+    print(script_content)
     assert re.search(r'(?m)--constraint=\(c1\)&\(c2&c3\)$', script_content)
     assert re.search(r'(?m)--constraint=(c1|c2&c3)$', script_content) is None
 
@@ -645,6 +646,28 @@ def test_combined_access_verbatim_constraint(make_job, slurm_only):
     assert re.search(r'(?m)^#SBATCH -C c3$', script_content)
 
 
+def test_sched_access_in_submit(make_job):
+    job = make_job(sched_access=['--constraint=c1', '--foo=bar'])
+    job.options = ['--constraint=c2', '--xyz']
+    job.scheduler._sched_access_in_submit = True
+
+    if job.scheduler.registered_name in ('flux', 'local', 'ssh'):
+        pytest.skip(f'not relevant for this scheduler backend')
+
+    prepare_job(job)
+    with open(job.script_filename) as fp:
+        script_content = fp.read()
+
+    print(script_content)
+    assert '--xyz' in script_content
+    assert '--foo=bar' not in script_content
+    if job.scheduler.registered_name in ('slurm', 'squeue'):
+        # Constraints are combined in `sched_access` for Slurm backends
+        assert '--constraint' not in script_content
+    else:
+        assert '--constraint=c1' not in script_content
+
+
 def test_guess_num_tasks(minimal_job, scheduler):
     minimal_job.num_tasks = 0
     if scheduler.registered_name == 'local':