add used_cpus_per_task parameter

Author: Samuel Moors

eessi/testsuite/eessi_mixin.py

@@ -58,6 +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
     always_request_gpus = None
     require_buildenv_module = False
     require_internet = False
@@ -231,8 +232,7 @@ def EESSI_mixin_validate_setup(self):
     @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"""
-        hooks.assign_tasks_per_compute_unit(test=self, compute_unit=self.compute_unit,
-                                            num_per=self.num_tasks_per_compute_unit)
+        hooks.assign_tasks_per_compute_unit(self)
 
         # Set OMP_NUM_THREADS environment variable
         hooks.set_omp_num_threads(self)
@@ -281,7 +281,7 @@ def EESSI_check_proc_binding(self):
         get_binding = os.path.join(os.path.dirname(check_binding_script), 'get_process_binding.sh')
         check_binding = ' '.join([
             f'{check_binding_script}',
-            f'--cpus-per-proc {self.num_cpus_per_task}',
+            f'--cpus-per-proc {self.used_cpus_per_task}',
             f'--procs {self.num_tasks}',
             f'--nodes {self.num_tasks // self.num_tasks_per_node}',
         ])

eessi/testsuite/hooks.py

@@ -77,7 +77,7 @@ def _assign_default_num_gpus_per_node(test: rfm.RegressionTest):
     log(f'default_num_gpus_per_node set to {test.default_num_gpus_per_node}')
 
 
-def assign_tasks_per_compute_unit(test: rfm.RegressionTest, compute_unit: str, num_per: int = 1):
+def assign_tasks_per_compute_unit(test: rfm.RegressionTest):
     """
     Assign one task per compute unit. More than 1 task per compute unit can be assigned with
     num_per for compute units that support it.
@@ -102,7 +102,9 @@ def assign_tasks_per_compute_unit(test: rfm.RegressionTest, compute_unit: str, n
     - assign_tasks_per_compute_unit(test, COMPUTE_UNITS.CPU_SOCKET) will launch 2 tasks with 64 threads per task
 
     """
-    log(f'assign_tasks_per_compute_unit called with compute_unit: {compute_unit} and num_per: {num_per}')
+    compute_unit = test.compute_unit
+    num_per = test.num_tasks_per_compute_unit
+    log(f'assign_tasks_per_compute_unit with compute_unit: {compute_unit} and num_per: {num_per}')
 
     if num_per != 1 and compute_unit not in [COMPUTE_UNITS.NODE]:
         raise NotImplementedError(
@@ -150,7 +152,7 @@ def assign_tasks_per_compute_unit(test: rfm.RegressionTest, compute_unit: str, n
     elif compute_unit == COMPUTE_UNITS.NUMA_NODE:
         _assign_one_task_per_numa_node(test)
     elif compute_unit == COMPUTE_UNITS.NODE:
-        _assign_num_tasks_per_node(test, num_per)
+        _assign_num_tasks_per_node(test)
     else:
         raise ValueError(f'compute unit {compute_unit} is currently not supported')
 
@@ -165,10 +167,13 @@ def assign_tasks_per_compute_unit(test: rfm.RegressionTest, compute_unit: str, n
         test.env_vars['SRUN_CPUS_PER_TASK'] = test.num_cpus_per_task
         log(f'Set environment variable SRUN_CPUS_PER_TASK to {test.env_vars["SRUN_CPUS_PER_TASK"]}')
 
+    if not test.used_cpus_per_task:
+        test.used_cpus_per_task = test.num_cpus_per_task
+
     _set_job_resources(test)
 
 
-def _assign_num_tasks_per_node(test: rfm.RegressionTest, num_per: int = 1):
+def _assign_num_tasks_per_node(test: rfm.RegressionTest):
     """
     Sets num_tasks_per_node and num_cpus_per_task such that it will run
     'num_per' tasks per node, unless specified with:
@@ -183,6 +188,7 @@ def _assign_num_tasks_per_node(test: rfm.RegressionTest, num_per: int = 1):
     - num_tasks_per_node = num_per
     - num_cpus_per_task = test.default_num_cpus_per_node / num_tasks_per_node
     """
+    num_per = test.num_tasks_per_compute_unit
 
     # neither num_tasks_per_node nor num_cpus_per_task are set
     if not test.num_tasks_per_node and not test.num_cpus_per_task:
@@ -693,7 +699,7 @@ def set_compact_process_binding(test: rfm.RegressionTest):
     # TODO: check if this also leads to sensible binding when using COMPUTE_UNITS.HWTHREAD
     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.num_cpus_per_task / num_cpus_per_core)
+    physical_cpus_per_task = int(test.used_cpus_per_task / num_cpus_per_core)
 
     if test.current_partition.launcher_type().registered_name == 'mpirun':
         # Do binding for intel and OpenMPI's mpirun, and srun
@@ -760,7 +766,7 @@ def set_omp_num_threads(test: rfm.RegressionTest):
     """
     Set number of OpenMP threads equal to number of CPUs per task
     """
-    test.env_vars['OMP_NUM_THREADS'] = test.num_cpus_per_task
+    test.env_vars['OMP_NUM_THREADS'] = test.used_cpus_per_task
     log(f'Set environment variable OMP_NUM_THREADS to {test.env_vars["OMP_NUM_THREADS"]}')
 
 

eessi/testsuite/tests/apps/osu.py

@@ -58,6 +58,7 @@ class EESSI_OSU_Base(osu_benchmark):
     """ base class for OSU tests """
     time_limit = '30m'
     module_name = parameter(find_modules('OSU-Micro-Benchmarks'))
+    used_cpus_per_task = 1
 
     # reset num_tasks_per_node from the hpctestlib: we handle it ourselves
     num_tasks_per_node = None