Merge branch 'main' into update_enroot_import_dockerhub

Author: ekouts

checks/apps/pytorch/pytorch_megatronlm.py

@@ -501,37 +501,15 @@ def throughput_per_gpu(self):
         ))
 
 
-class pytorch_image_import(rfm.RunOnlyRegressionTest):
-    image = variable(
-        str,
-        value=('docker://jfrog.svc.cscs.ch#reframe-oci/'
-               'pytorch:25.01-py3_nvrtc-12.9')
-    )
-    archive_name = 'pytorch.sqsh'
-    executable = 'enroot'
-    valid_systems = ['+ce']
-    valid_prog_environs = ['builtin']
-
-    @run_before('run')
-    def set_executable_opts(self):
-        self.executable_opts = ['import', '-o', self.archive_name, self.image]
-
-    @sanity_function
-    def assert_image_imported(self):
-        return sn.path_exists(os.path.join(self.stagedir, self.archive_name))
-
-
 @rfm.simple_test
 class PyTorchMegatronLM_CE(PyTorchMegatronLM, ContainerEngineMixin):
     valid_systems = ['+nvgpu +ce']
     valid_prog_environs = ['builtin']
     maintainers = ['ml-team']
-    pytorch_image = fixture(pytorch_image_import, scope='session')
+    container_image = 'docker://jfrog.svc.cscs.ch#reframe-oci/pytorch:25.01-py3_nvrtc-12.9'
 
     @run_after('setup')
     def set_container_config(self):
-        self.container_image = os.path.join(self.pytorch_image.stagedir,
-                                            self.pytorch_image.archive_name)
         self.container_env_table = {
             'annotations.com.hooks': {
                 'aws_ofi_nccl.enabled': 'true',

checks/apps/pytorch/pytorch_megatronlm_amd.py

@@ -18,7 +18,7 @@
 class PyTorchMegatronLM_AMD(rfm.RunOnlyRegressionTest):
     num_tasks_per_node = 1
     default_num_nodes = variable(int, type(None), value=None)
-    time_limit = '30m'
+    time_limit = '50m'
     megatron_repo = variable(
         str, value='https://github.com/ROCm/Megatron-LM'
     )
@@ -379,37 +379,15 @@ def throughput_per_gpu(self):
         ))
 
 
-class pytorch_image_import(rfm.RunOnlyRegressionTest):
-    sourcesdir = None
-    image = variable(
-        str,
-        value=('docker://rocm/megatron-lm:v25.6_py312')
-    )
-    archive_name = 'pytorch.sqsh'
-    executable = 'enroot'
-    valid_systems = ['+ce']
-    valid_prog_environs = ['builtin']
-
-    @run_before('run')
-    def set_executable_opts(self):
-        self.executable_opts = ['import', '-o', self.archive_name, self.image]
-
-    @sanity_function
-    def assert_image_imported(self):
-        return sn.path_exists(os.path.join(self.stagedir, self.archive_name))
-
-
 @rfm.simple_test
 class PyTorchMegatronLM_AMD_CE(PyTorchMegatronLM_AMD, ContainerEngineMixin):
     valid_systems = ['+amdgpu +ce']
     valid_prog_environs = ['builtin']
     maintainers = ['ml-team']
-    pytorch_image = fixture(pytorch_image_import, scope='session')
+    container_image = 'rocm/megatron-lm:v25.6_py312'
 
     @run_after('setup')
     def set_container_config(self):
-        self.container_image = os.path.join(self.pytorch_image.stagedir,
-                                            self.pytorch_image.archive_name)
         self.container_env_table = {
             'annotations.com.hooks': {
                 'aws_ofi_nccl.enabled': 'true',

checks/prgenv/affinity_check.py

@@ -27,7 +27,7 @@ class CompileAffinityTool(rfm.CompileOnlyRegressionTest,
     env_vars = {'MPICH_GPU_SUPPORT_ENABLED': 0}
 
     sourcesdir = 'https://github.com/vkarak/affinity'
-    tags = {'production', 'scs', 'maintenance', 'craype'}
+    tags = {'scs', 'craype'}
 
     @run_before('compile')
     def set_build_opts(self):
@@ -76,7 +76,7 @@ class AffinityTestBase(rfm.RunOnlyRegressionTest,
         '+openmp +prgenv'
     ]
 
-    tags = {'production', 'scs', 'maintenance', 'craype'}
+    tags = {'scs', 'craype'}
 
     @run_after('setup')
     def skip_cpe_2312(self):

checks/prgenv/mpi.py

@@ -11,10 +11,11 @@
 
 sys.path.append(str(pathlib.Path(__file__).parent.parent / 'mixins'))
 from container_engine import ContainerEngineCPEMixin  # noqa: E402
+from uenv_slurm_mpi_options import UenvSlurmMpiOptionsMixin  # noqa: E402
 
 
 @rfm.simple_test
-class MpiInitTest(rfm.RegressionTest, ContainerEngineCPEMixin):
+class MpiInitTest(rfm.RegressionTest, ContainerEngineCPEMixin, UenvSlurmMpiOptionsMixin):
     '''
     This test checks the value returned by calling MPI_Init_thread.
     '''
@@ -55,9 +56,11 @@ def set_sanity(self):
         # - 7.7.15 (ANL base 3.2)
         # - 8.0.16.17 (ANL base 3.3)
         # - 8.1.4.31,8.1.5.32,8.1.18.4,8.1.21.11,8.1.25.17 (ANL base 3.4a2)
-        regex = r'= MPI VERSION\s+: CRAY MPICH version \S+ \(ANL base (\S+)\)'
+        # OpenMPI version:
+        # - MPI-3.1 = Open MPI v5.0.9
+        regex = r'= (MPI VERSION\s+: CRAY MPICH version \S+ \(ANL base |Open MPI v)([\S^\)]+)'
         stdout = os.path.join(self.stagedir, sn.evaluate(self.stdout))
-        mpich_version = sn.extractsingle(regex, stdout, 1)
+        mpich_version = sn.extractsingle(regex, stdout, 2)
         self.mpithread_version = {
             '3.2': {
                 'MPI_THREAD_SINGLE': 0,
@@ -125,7 +128,7 @@ def set_sanity(self):
 
 
 @rfm.simple_test
-class MpiGpuDirectOOM(rfm.RegressionTest, ContainerEngineCPEMixin):
+class MpiGpuDirectOOM(rfm.RegressionTest, ContainerEngineCPEMixin, UenvSlurmMpiOptionsMixin):
     '''
     This test checks the issue reported in:
     https://github.com/eth-cscs/alps-gh200-reproducers/tree/main/gpudirect-oom

checks/system/slurm/slurm.py

@@ -4,13 +4,18 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import os
+import pathlib
 import re
+import sys
 
 import reframe as rfm
 import reframe.core.runtime as rt
 import reframe.utility.osext as osext
 import reframe.utility.sanity as sn
 
+sys.path.append(str(pathlib.Path(__file__).parent / 'mixins'))
+from uenv_slurm_mpi_options import UenvSlurmMpiOptionsMixin  # noqa: E402
+
 
 class SlurmSimpleBaseCheck(rfm.RunOnlyRegressionTest):
     '''Base class for Slurm simple binary tests'''
@@ -192,7 +197,7 @@ def assert_found_exceeded_memory(self):
 
 
 @rfm.simple_test
-class MemoryOverconsumptionCheckMPI(SlurmCompiledBaseCheck):
+class MemoryOverconsumptionCheckMPI(SlurmCompiledBaseCheck, UenvSlurmMpiOptionsMixin):
     # TODO: maintainers = ['@jgphpc', '@ekouts']
     descr = 'Tests for max allocatable memory'
     valid_systems = ['+remote']
@@ -207,8 +212,11 @@ class MemoryOverconsumptionCheckMPI(SlurmCompiledBaseCheck):
     def set_num_tasks(self):
         self.skip_if_no_procinfo()
         cpu = self.current_partition.processor
-        self.num_tasks_per_node = int(
-            cpu.info['num_cpus'] / cpu.info['num_cpus_per_core'])
+        # Limit number of tasks because PMIx/OpenMPI can take very long to
+        # initialize with e.g. 288 ranks on one GH200 node. The test still
+        # fails in a reasonable time with a limited number of ranks.
+        self.num_tasks_per_node = min(16, int(
+            cpu.info['num_cpus'] / cpu.info['num_cpus_per_core']))
         self.num_tasks = self.num_tasks_per_node
         self.job.launcher.options += ['-u']