Merge pull request #2656 from ekouts/slurm_cpus_per_task_bugfix

[bugfix] Pass explicitly the `--cpus-per-task` option to `srun` for Slurm >= 22.05

Author: vkarak

reframe/core/launchers/mpi.py

@@ -3,15 +3,52 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+import semver
+import re
+
+import reframe.utility.osext as osext
 from reframe.core.backends import register_launcher
 from reframe.core.launchers import JobLauncher
+from reframe.core.logging import getlogger
 from reframe.utility import seconds_to_hms
 
 
 @register_launcher('srun')
 class SrunLauncher(JobLauncher):
+    def __init__(self):
+        self.options = []
+        self.use_cpus_per_task = True
+        try:
+            out = osext.run_command('srun --version')
+            match = re.search('slurm (\d+)\.(\d+)\.(\d+)', out.stdout)
+            if match:
+                # We cannot pass to semver strings like 22.05.1 directly
+                # because it is not a valid version string for semver. We
+                # need to remove all the leading zeros.
+                slurm_version = (
+                    semver.VersionInfo(
+                        match.group(1), match.group(2), match.group(3)
+                    )
+                )
+                if slurm_version < semver.VersionInfo(22, 5, 0):
+                    self.use_cpus_per_task = False
+            else:
+                getlogger().warning(
+                    'could not get version of Slurm, --cpus-per-task will be '
+                    'set according to the num_cpus_per_task attribute'
+                )
+        except Exception:
+            getlogger().warning(
+                'could not get version of Slurm, --cpus-per-task will be set '
+                'according to the num_cpus_per_task attribute'
+            )
+
     def command(self, job):
-        return ['srun']
+        ret = ['srun']
+        if self.use_cpus_per_task and job.num_cpus_per_task:
+            ret.append(f'--cpus-per-task={job.num_cpus_per_task}')
+
+        return ret
 
 
 @register_launcher('ibrun')

unittests/test_launchers.py

@@ -104,6 +104,9 @@ def minimal_job(make_job, launcher):
 
 def test_run_command(job):
     launcher_name = type(job.launcher).registered_name
+    # This is relevant only for the srun launcher, because it may
+    # run in different platforms with older versions of Slurm
+    job.launcher.use_cpus_per_task = True
     command = job.launcher.run_command(job)
     if launcher_name == 'alps':
         assert command == 'aprun -n 4 -N 2 -d 2 -j 0 --foo'
@@ -116,7 +119,7 @@ def test_run_command(job):
     elif launcher_name == 'mpirun':
         assert command == 'mpirun -np 4 --foo'
     elif launcher_name == 'srun':
-        assert command == 'srun --foo'
+        assert command == 'srun --cpus-per-task=2 --foo'
     elif launcher_name == 'srunalloc':
         assert command == ('srun '
                            '--job-name=fake_job '
@@ -147,6 +150,9 @@ def test_run_command(job):
 
 def test_run_command_minimal(minimal_job):
     launcher_name = type(minimal_job.launcher).registered_name
+    # This is relevant only for the srun launcher, because it may
+    # run in different platforms with older versions of Slurm
+    minimal_job.launcher.use_cpus_per_task = True
     command = minimal_job.launcher.run_command(minimal_job)
     if launcher_name == 'alps':
         assert command == 'aprun -n 1 --foo'