[bugfix] Fix Slurm job array completion detection

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:
     #   * <job_id>_<array_task_id>
-    #   * <job_id>_[<array_task_id_start>-<array_task_id_end>]
+    #   * <job_id>_[<array_expression>]
+    # 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)

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')