monkeypatch srun launcher

Author: Samuel Moors

eessi/testsuite/common_config.py

@@ -1,5 +1,6 @@
 import os
 
+from reframe.core.launchers.mpi import SrunLauncher
 import reframe.core.logging as rflog
 
 perflog_format = '|'.join([
@@ -143,3 +144,16 @@ def get_sbatch_account():
         err_msg += " It is required to set `SBATCH_ACCOUNT` to run on this system."
         raise ValueError(err_msg)
     return sbatch_account
+
+
+# Monkeypatch the srun launcher: remove '--cpus-per-task' option:
+# we add it ourselves, see the `assign_tasks_per_compute_unit` hook
+_srunlauncher_command = SrunLauncher.command
+
+
+def _srunlauncher_patched_command(self, job):
+    cmd = _srunlauncher_command(self, job)
+    return [x for x in cmd if not x.startswith('--cpus-per-task=')]
+
+
+SrunLauncher.command = _srunlauncher_patched_command

eessi/testsuite/hooks.py

@@ -162,7 +162,7 @@ def assign_tasks_per_compute_unit(test: rfm.RegressionTest):
         test.used_cpus_per_task = test.num_cpus_per_task
 
     if test.current_partition.launcher_type().registered_name == 'srun':
-        # Add again --cpus-per-task to srun, which wins because it is set last:
+        # Add (again) --cpus-per-task to srun, which wins because it is set last:
         test.job.launcher.options += ['--cpus-per-task', str(test.used_cpus_per_task)]
         # Make sure srun inherits --cpus-per-task from the job environment for Slurm versions >= 22.05 < 23.11,
         # ensuring the same task binding across all Slurm versions.