Merge pull request #448 from msimberg/mpi-tests-slurm-flag

Set `--mpi` slurm flag appropriately flag on some MPI tests

Author: ekouts

checks/microbenchmarks/mpi/osu/osu_run.py

@@ -3,17 +3,22 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+import pathlib
+import sys
+
 import reframe as rfm
 import reframe.utility.sanity as sn
 
+sys.path.append(str(pathlib.Path(__file__).parent.parent.parent.parent / 'mixins'))
+from uenv_slurm_mpi_options import UenvSlurmMpiOptionsMixin
+
 
-class BaseCheck(rfm.RunOnlyRegressionTest):
+class BaseCheck(rfm.RunOnlyRegressionTest, UenvSlurmMpiOptionsMixin):
     valid_systems = ['+remote']
     valid_prog_environs = ['+osu-micro-benchmarks +uenv']
     sourcesdir = None
     num_tasks = 2
     num_tasks_per_node = 1
-    pmi = variable(str, value='')
     env_vars = {
         # Disable GPU support for mpich
         'MPIR_CVAR_ENABLE_GPU': 0,
@@ -38,11 +43,6 @@ def prepare_test(self):
         processor = self.current_partition.processor
         self.num_cpus_per_task = processor.num_cpus_per_socket
 
-    @run_after('setup')
-    def set_launcher_options(self):
-        if self.pmi != '':
-            self.job.launcher.options = [f'--mpi={self.pmi}']
-
     @sanity_function
     def assert_sanity(self):
         # Only check for the last entry in the latency test,

checks/microbenchmarks/xccl/xccl_tests.py

@@ -12,6 +12,8 @@
 sys.path.append(str(pathlib.Path(__file__).parent.parent.parent / 'mixins'))
 
 from container_engine import ContainerEngineMixin  # noqa: E402
+from slurm_mpi_pmix import SlurmMpiPmixMixin
+from uenv_slurm_mpi_options import UenvSlurmMpiOptionsMixin
 
 
 class XCCLTestsBase(rfm.RunOnlyRegressionTest):
@@ -88,36 +90,17 @@ def set_perf(self):
         }
 
 
-class XCCLTestsBaseCE(XCCLTestsBase, ContainerEngineMixin):
+class XCCLTestsBaseCE(XCCLTestsBase, ContainerEngineMixin, SlurmMpiPmixMixin):
     container_env_table = {
         'annotations.com.hooks': {
             'aws_ofi_nccl.enabled': 'true'
         }
     }
 
-    @run_before('run')
-    def set_pmix(self):
-        self.job.launcher.options += ['--mpi=pmix']
 
-        # Disable MCA components to avoid warnings
-        self.env_vars.update(
-            {
-                'PMIX_MCA_psec': '^munge',
-                'PMIX_MCA_gds': '^shmem2'
-            }
-        )
-
-
-class XCCLTestsBaseUENV(XCCLTestsBase):
+class XCCLTestsBaseUENV(XCCLTestsBase, UenvSlurmMpiOptionsMixin):
     tags.add('bencher')
 
-    @run_before('run')
-    def set_pmix(self):
-        # Some clusters, like clariden, don't use cray_shasta as default.
-        # cray_shasta is required for cray-mpich, which most uenvs use.  This
-        # will need to be updated when uenvs can have OpenMPI in them.
-        self.job.launcher.options += ['--mpi=cray_shasta']
-
 
 def _set_xccl_uenv_env_vars(env_vars):
     env_vars.update(

checks/mixins/slurm_mpi_pmix.py

@@ -0,0 +1,28 @@
+# Copyright 2016-2023 Swiss National Supercomputing Centre (CSCS/ETH Zurich)
+# ReFrame Project Developers. See the top-level LICENSE file for details.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import reframe as rfm
+
+
+class SlurmMpiPmixMixin(rfm.RegressionMixin):
+    """
+    Set slurm --mpi flags for containers that require PMIx.
+
+    Containers that use OpenMPI require the --mpi=pmix flag.
+
+    Additionally silence some warnings triggered by OpenMPI.
+    """
+
+    @run_after('setup')
+    def set_slurm_mpi_pmix(self):
+        self.job.launcher.options += ['--mpi=pmix']
+
+        # Disable MCA components to avoid warnings
+        self.env_vars.update(
+            {
+                'PMIX_MCA_psec': 'native',
+                'PMIX_MCA_gds': '^shmem2'
+            }
+        )

checks/mixins/uenv_slurm_mpi_options.py

@@ -0,0 +1,37 @@
+# Copyright 2016-2023 Swiss National Supercomputing Centre (CSCS/ETH Zurich)
+# ReFrame Project Developers. See the top-level LICENSE file for details.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import reframe as rfm
+
+
+class UenvSlurmMpiOptionsMixin(rfm.RegressionMixin):
+    """
+    Set slurm --mpi flags based on uenv features.
+
+    For uenvs that announce that they use OpenMPI, --mpi=pmix needs to be used.
+    For uenvs that advertise Cray MPICH, or uenvs that don't specify, we assume
+    that --mpi=cray_shasta is required.
+
+    Additionally for OpenMPI we silence some warnings.
+    """
+
+    @run_after('setup')
+    def set_uenv_slurm_mpi_options(self):
+        features = self.current_environ.features
+        if "openmpi" in features:
+            self.job.launcher.options += ['--mpi=pmix']
+
+            # Disable MCA components to avoid warnings
+            self.env_vars.update(
+                {
+                    'PMIX_MCA_psec': 'native',
+                    'PMIX_MCA_gds': '^shmem2'
+                }
+            )
+        elif "cray-mpich" in features:
+            self.job.launcher.options += ['--mpi=cray_shasta']
+        else:
+            # Assume cray-mpich is used if nothing is specified.
+            self.job.launcher.options += ['--mpi=cray_shasta']

checks/prgenv/helloworld.py

@@ -11,10 +11,11 @@
 sys.path.append(str(pathlib.Path(__file__).parent.parent / 'mixins'))
 from extra_launcher_options import ExtraLauncherOptionsMixin
 from container_engine import ContainerEngineCPEMixin
+from uenv_slurm_mpi_options import UenvSlurmMpiOptionsMixin
 
 
 class HelloWorldBaseTest(rfm.RegressionTest, ExtraLauncherOptionsMixin,
-                         ContainerEngineCPEMixin):
+                         ContainerEngineCPEMixin, UenvSlurmMpiOptionsMixin):
     linking = parameter(['dynamic'])
     lang = parameter(['c', 'cpp', 'F90'])
     sourcesdir = 'src/hello'