Combine Slurm constraints also when sched_access_in_submit is set

Author: vkarak

reframe/core/schedulers/slurm.py

@@ -212,40 +212,33 @@ def emit_preamble(self, job):
  )
  )
 
- constraints = []
- constraint_parser = ArgumentParser()
- constraint_parser.add_argument('-C', '--constraint')
- if not self._sched_access_in_submit:
- 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`
- 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
+
+ arg = '&'.join(f'({c.strip()})' for c in constraints)
+ job._sched_access = [f'--constraint={arg}']
 
- preamble.append(self._format_option(constr, '--constraint={0}'))
+ 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+)')

unittests/test_schedulers.py

@@ -617,6 +617,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
 
@@ -644,6 +645,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 ('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':