add PRTE_MCA_rmaps_base_mapping_policy

eessi/testsuite/hooks.py

@@ -667,7 +667,7 @@ def set_compact_process_binding(test: rfm.RegressionTest):
 
     It is hard to do this in a portable way. Currently supported for process binding are:
     - Intel MPI (through I_MPI_PIN_DOMAIN)
-    - OpenMPI (through OMPI_MCA_rmaps_base_mapping_policy)
+    - OpenMPI (through {OMPI,PRTE}_MCA_rmaps_base_mapping_policy)
     - srun (LIMITED SUPPORT: through SLURM_CPU_BIND, but only effective if task/affinity plugin is enabled)
     """
 
@@ -680,21 +680,21 @@ def set_compact_process_binding(test: rfm.RegressionTest):
 
     if test.current_partition.launcher_type().registered_name == 'mpirun':
         # Do binding for intel and OpenMPI's mpirun, and srun
-        test.env_vars['I_MPI_PIN_CELL'] = 'core'  # Don't bind to hyperthreads, only to physcial cores
-        test.env_vars['I_MPI_PIN_DOMAIN'] = '%s:compact' % physical_cpus_per_task
-        test.env_vars['OMPI_MCA_rmaps_base_mapping_policy'] = 'slot:PE=%s' % physical_cpus_per_task
-        log(f'Set environment variable I_MPI_PIN_CELL to {test.env_vars["I_MPI_PIN_CELL"]}')
-        log(f'Set environment variable I_MPI_PIN_DOMAIN to {test.env_vars["I_MPI_PIN_DOMAIN"]}')
-        log('Set environment variable OMPI_MCA_rmaps_base_mapping_policy to '
-            f'{test.env_vars["OMPI_MCA_rmaps_base_mapping_policy"]}')
+        env_vars = {
+            'I_MPI_PIN_CELL': 'core',  # Don't bind to hyperthreads, only to physcial cores
+            'I_MPI_PIN_DOMAIN': f'{physical_cpus_per_task}:compact',
+            'OMPI_MCA_rmaps_base_mapping_policy': f'slot:PE={physical_cpus_per_task}',
+            'PRTE_MCA_rmaps_base_mapping_policy': f'slot:PE={physical_cpus_per_task}',
+        }
     elif test.current_partition.launcher_type().registered_name == '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
-        test.env_vars['SLURM_DISTRIBUTION'] = 'block:block'
-        test.env_vars['SLURM_CPU_BIND'] = 'verbose'
-        log(f'Set environment variable SLURM_DISTRIBUTION to {test.env_vars["SLURM_DISTRIBUTION"]}')
-        log(f'Set environment variable SLURM_CPU_BIND to {test.env_vars["SLURM_CPU_BIND"]}')
+        env_vars = {
+            'SLURM_DISTRIBUTION': 'block:block',
+            'SLURM_CPU_BIND': 'verbose',
+        }
     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 += " The test will run, but using the default binding strategy of your parallel launcher."
@@ -703,6 +703,10 @@ def set_compact_process_binding(test: rfm.RegressionTest):
         # Warnings will, at default loglevel, be printed on stdout when executing the ReFrame command
         rflog.getlogger().warning(msg)
 
+    for key, value in env_vars.items():
+        test.env_vars[key] = value
+        log(f'Set environment variable {key} to {test.env_vars[key]}')
+
 
 def set_compact_thread_binding(test: rfm.RegressionTest):
     """