set srun --cpus-per-task in custom eessi-srun launcher

Author: Samuel Moors

eessi/testsuite/common_config.py

@@ -1,5 +1,7 @@
 import os
 
+from reframe.core.backends import register_launcher
+from reframe.core.launchers.mpi import SrunLauncher
 import reframe.core.logging as rflog
 
 perflog_format = '|'.join([
@@ -143,3 +145,24 @@ 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
+
+
+@register_launcher('eessi-srun')
+class EESSISrunLauncher(SrunLauncher):
+    """
+    Custom srun launcher for the EESSI test suite
+    Sets srun cpus per task to the actually used cpus per task (used_cpus_per_task)
+    """
+
+    def command(self, job):
+        ret = ['srun']
+        if hasattr(job, 'used_cpus_per_task') and job.used_cpus_per_task:
+            ret.append(f'--cpus-per-task={job.used_cpus_per_task}')
+        elif self.use_cpus_per_task and job.num_cpus_per_task:
+            ret.append(f'--cpus-per-task={job.num_cpus_per_task}')
+
+        if self.env_vars:
+            env_vars = ','.join(f'{k}={v}' for k, v in self.env_vars.items())
+            ret.append(f'--export={env_vars}')
+
+        return ret

eessi/testsuite/eessi_mixin.py

@@ -58,7 +58,7 @@ class EESSI_Mixin(RegressionTestPlugin):
     bench_name = None
     is_ci_test = False
     num_tasks_per_compute_unit = 1
-    used_cpus_per_task = None  # effectively used cpus per task
+    used_cpus_per_task = None  # actually used cpus per task
     always_request_gpus = None
     require_buildenv_module = False
     require_internet = False
@@ -229,6 +229,14 @@ def EESSI_mixin_validate_setup(self):
         # i.e. exists in their respective dict from eessi.testsuite.constants
         self.EESSI_mixin_validate_item_in_list('compute_unit', COMPUTE_UNITS[:])
 
+        # Check that default ReFrame srun launcher is not used
+        launcher = self.current_partition.launcher_type().registered_name
+        if launcher == 'srun':
+            msg = ('The default ReFrame srun launcher is not fully supported by the EESSI test suite.'
+                   ' Please use `eessi.testsuite.common_config.eessi-srun` instead.')
+            log_once(self, msg, msg_id='3', level='warning')
+
+
     @run_after('setup')
     def EESSI_mixin_assign_tasks_per_compute_unit(self):
         """Call hooks to assign tasks per compute unit, set OMP_NUM_THREADS, and set compact process binding"""

eessi/testsuite/hooks.py

@@ -30,6 +30,7 @@ def _set_job_resources(test: rfm.RegressionTest):
     test.job.num_tasks_per_socket = test.num_tasks_per_socket
     test.job.num_cpus_per_task = test.num_cpus_per_task
     test.job.use_smt = test.use_multithreading
+    test.job.used_cpus_per_task = test.used_cpus_per_task
 
 
 def _assign_default_num_cpus_per_node(test: rfm.RegressionTest):
@@ -161,12 +162,7 @@ def assign_tasks_per_compute_unit(test: rfm.RegressionTest):
     if not test.used_cpus_per_task:
         test.used_cpus_per_task = test.num_cpus_per_task
 
-    # Tell srun to use the number of actually used cpus per tasks, not the number of allocated cpus per task.
-    # With the default srun launcher in ReFrame this doesn’t work because it sets `srun --cpus-per-task`.
-    # To make it work, a custom srun launcher should be used that doesn’t set `srun --cpus-per-task`.
-    test.env_vars['SLURM_CPUS_PER_TASK'] = test.used_cpus_per_task
-
-    if test.current_partition.launcher_type().registered_name == 'srun':
+    if test.current_partition.launcher_type().registered_name in ['eessi-srun', 'srun']:
         # 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.
         # https://bugs.schedmd.com/show_bug.cgi?id=13351
@@ -705,8 +701,9 @@ def set_compact_process_binding(test: rfm.RegressionTest):
     check_proc_attribute_defined(test, 'num_cpus_per_core')
     num_cpus_per_core = test.current_partition.processor.num_cpus_per_core
     physical_cpus_per_task = int(test.used_cpus_per_task / num_cpus_per_core)
+    launcher = test.current_partition.launcher_type().registered_name
 
-    if test.current_partition.launcher_type().registered_name == 'mpirun':
+    if launcher == 'mpirun':
         # Do binding for intel and OpenMPI's mpirun, and srun
         env_vars = {
             'I_MPI_PIN_CELL': 'core',  # Don't bind to hyperthreads, only to physcial cores
@@ -725,17 +722,16 @@ def set_compact_process_binding(test: rfm.RegressionTest):
         if any(re.search(pattern, x) for x in test.modules):
             test.job.launcher.options.append(f'--map-by slot:PE={physical_cpus_per_task} --report-bindings')
             log(f'Set launcher command to {test.job.launcher.run_command(test.job)}')
-    elif test.current_partition.launcher_type().registered_name == 'srun':
+    elif launcher in ['eessi-srun', 'srun']:
         # Set compact binding for SLURM. Only effective if the task/affinity plugin is enabled
         # and when number of tasks times cpus per task equals either socket, core or thread count
         env_vars = {
             'SLURM_DISTRIBUTION': 'block:block',
-            'SLURM_CPU_BIND': 'verbose',
+            'SLURM_CPU_BIND': 'verbose,cores',
         }
     else:
         env_vars = {}
-        msg = "hooks.set_compact_process_binding does not support the current launcher"
-        msg += f" ({test.current_partition.launcher_type().registered_name})."
+        msg = f"hooks.set_compact_process_binding does not support the current launcher ({launcher})."
         msg += " The test will run, but using the default binding strategy of your parallel launcher."
         msg += " This may lead to suboptimal performance."
         msg += " Please expand the functionality of hooks.set_compact_process_binding for your parallel launcher."