diff --git a/reframe/core/schedulers/slurm.py b/reframe/core/schedulers/slurm.py index 5892a936f..b96a6137a 100644 --- a/reframe/core/schedulers/slurm.py +++ b/reframe/core/schedulers/slurm.py @@ -68,10 +68,40 @@ def slurm_state_pending(state): _run_strict = functools.partial(osext.run_command, check=True) +def _count_array_tasks(array_spec): + # The optional throttle does not affect the number of array tasks. + array_spec = array_spec.split('%', maxsplit=1)[0] + num_tasks = 0 + for task_range in array_spec.split(','): + range_spec, *step_spec = task_range.split(':', maxsplit=1) + step = int(step_spec[0]) if step_spec else 1 + bounds = [int(x) for x in range_spec.split('-', maxsplit=1)] + if len(bounds) == 1: + num_tasks += 1 + else: + start, stop = bounds + num_tasks += len(range(start, stop + 1, step)) + + return num_tasks + + +def _count_array_tasks_from_jobid(jobid): + try: + array_spec = jobid.split('_', maxsplit=1)[1] + except IndexError: + return 0 + + if array_spec.startswith('[') and array_spec.endswith(']'): + array_spec = array_spec[1:-1] + + return _count_array_tasks(array_spec) + + class _SlurmJob(sched.Job): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._is_array = False + self._array_size = None self._is_cancelling = False # The compacted nodelist as reported by Slurm. This must be updated in @@ -93,6 +123,10 @@ def nodelist(self): def is_array(self): return self._is_array + @property + def array_size(self): + return self._array_size + @property def is_cancelling(self): return self._is_cancelling @@ -107,9 +141,11 @@ class SlurmJobScheduler(sched.JobScheduler): # (https://slurm.schedmd.com/heterogeneous_jobs.html) # For job arrays the job_id has one of the following formats: # * _ - # * _[-] + # * _[] + # An array expression may contain lists, ranges, steps and a task + # throttle, e.g., ``1,3-7:2%2``. # (https://slurm.schedmd.com/job_array.html) - _jobid_patt = r'\d+(?:\+\d+|_\d+|_\[\d+-\d+\])?' + _jobid_patt = r'\d+(?:\+\d+|_(?:\d+|\[[\d,:%-]+\]))?' def __init__(self): self._prefix = '#SBATCH' @@ -190,6 +226,14 @@ def emit_preamble(self, job): ) if parsed_args.array: job._is_array = True + try: + job._array_size = _count_array_tasks(parsed_args.array) + except (TypeError, ValueError): + self.log( + f'could not determine the number of tasks in Slurm ' + f'array expression {parsed_args.array!r}' + ) + self.log('Slurm job is a job array') # Slurm replaces '%a' by the corresponding SLURM_ARRAY_TASK_ID @@ -480,6 +524,24 @@ def _update_completion_time(self, job, timestamps): if ct: job._completion_time = max(ct) + def _get_job_states(self, job, jobarr_info): + states = [m.group('state') for m in jobarr_info] + if not job.is_array or job.array_size is None: + return states + + num_tasks = sum( + _count_array_tasks_from_jobid(m.group('jobid')) + for m in jobarr_info + ) + if num_tasks < job.array_size: + self.log( + f'Slurm reports {num_tasks} of {job.array_size} tasks for ' + f'job array {job.jobid}; keeping the job pending' + ) + states.append('PENDING') + + return states + def poll(self, *jobs): '''Update the status of the jobs.''' @@ -542,7 +604,7 @@ def poll(self, *jobs): continue # Join the states with ',' in case of job arrays|heterogeneous jobs - job._state = ','.join(m.group('state') for m in jobarr_info) + job._state = ','.join(self._get_job_states(job, jobarr_info)) if slurm_state_completed(job.state): # Since Slurm exitcodes are positive take the maximum one @@ -718,7 +780,7 @@ def poll(self, *jobs): continue # Join the states with ',' in case of job arrays - job._state = ','.join(s.group('state') for s in job_match) + job._state = ','.join(self._get_job_states(job, job_match)) # Use ',' to join nodes to be consistent with Slurm syntax job._nodespec = ','.join(m.group('nodespec') for m in job_match) diff --git a/unittests/test_schedulers.py b/unittests/test_schedulers.py index cf5737b51..b2aeaa3d6 100644 --- a/unittests/test_schedulers.py +++ b/unittests/test_schedulers.py @@ -9,8 +9,10 @@ import signal import socket import time +from types import SimpleNamespace import reframe.core.runtime as rt +import reframe.core.schedulers.slurm as slurm import reframe.utility.osext as osext import unittests.utility as test_util from reframe.core.backends import (getlauncher, getscheduler) @@ -579,6 +581,69 @@ def test_submit_job_array(make_job, slurm_only, exec_ctx): re.search('Task id: 1', output)]) +@pytest.mark.parametrize('array_spec,num_tasks', [ + ('0-1', 2), + ('0-1%1', 2), + ('1,3,5,7', 4), + ('1-7:2', 4), + ('1,3-7:2%2', 4), +]) +def test_slurm_count_array_tasks(array_spec, num_tasks): + assert slurm._count_array_tasks(array_spec) == num_tasks + + +@pytest.mark.parametrize('jobid', [ + '123_1', + '123_[1%1]', + '123_[0-1%1]', + '123_[1,3-7:2%2]', +]) +def test_slurm_array_jobid_pattern(jobid): + assert re.fullmatch(slurm.SlurmJobScheduler._jobid_patt, jobid) + + +@pytest.mark.parametrize('sacct_output,expected_state,expected_finished', [ + ( + '123_0|COMPLETED|0:0|1|nid001\n', + 'COMPLETED,PENDING', + False, + ), + ( + '123_0|COMPLETED|0:0|1|nid001\n' + '123_[1%1]|PENDING|0:0|Unknown|\n', + 'COMPLETED,PENDING', + False, + ), + ( + '123_0|COMPLETED|0:0|1|nid001\n' + '123_1|COMPLETED|0:0|2|nid002\n', + 'COMPLETED,COMPLETED', + True, + ), +]) +def test_slurm_poll_job_array(make_job, slurm_only, testsys_exec_ctx, + monkeypatch, + sacct_output, expected_state, + expected_finished): + job = make_job() + if job.scheduler.registered_name != 'slurm': + pytest.skip('test exercises the sacct Slurm backend') + + job.options = ['--array=0-1%1'] + job.scheduler.emit_preamble(job) + job._jobid = '123' + job._submit_time = time.time() + monkeypatch.setattr( + slurm, '_run_strict', + lambda *args, **kwargs: SimpleNamespace(stdout=sacct_output) + ) + + job.scheduler.poll(job) + assert job.array_size == 2 + assert job.state == expected_state + assert job.finished() is expected_finished + + def test_cancel(make_job, exec_ctx): minimal_job = make_job(sched_access=exec_ctx.access) prepare_job(minimal_job, 'sleep 5')