diff --git a/docs/config_reference.rst b/docs/config_reference.rst index 4f48e20a3..3ab38af0c 100644 --- a/docs/config_reference.rst +++ b/docs/config_reference.rst @@ -518,8 +518,10 @@ System Partition Configuration The program will be launched locally. - ``lrun``: Parallel programs will be launched using `LC Launcher `__'s ``lrun`` command. - ``lrun-gpu``: Parallel programs will be launched using `LC Launcher `__'s ``lrun -M "-gpu"`` command that enables the CUDA-aware Spectrum MPI. - - ``mpirun``: Parallel programs will be launched using the ``mpirun`` command. - ``mpiexec``: Parallel programs will be launched using the ``mpiexec`` command. + - ``mpirun``: Parallel programs will be launched using the (generic) ``mpirun`` command. + - ``mpirun-intelmpi``: Parallel programs will be launched using the Intel MPI's ``mpirun`` command. + - ``mpirun-openmpi``: Parallel programs will be launched using the OpenMPI's ``mpirun`` command. - ``pdsh``: Parallel programs will be launched using the ``pdsh`` command. This launcher uses the partition's :attr:`~config.systems.partitions.access` property in order to determine the options to be passed to ``pdsh``. - ``srun``: Parallel programs will be launched using `Slurm `__'s ``srun`` command. - ``srunalloc``: Parallel programs will be launched using `Slurm `__'s ``srun`` command, but job allocation options will also be emitted. @@ -542,6 +544,10 @@ System Partition Configuration - ``upcrun``: Parallel programs will be launched using the `UPC `__ ``upcrun`` command. - ``upcxx-run``: Parallel programs will be launched using the `UPC++ `__ ``upcxx-run`` command. + .. versionadded:: 4.9 + + The ``mpirun-intelmpi`` and ``mpirun-openmpi`` parallel launchers were added. + .. tip:: .. versionadded:: 4.0.0 diff --git a/reframe/core/launchers/__init__.py b/reframe/core/launchers/__init__.py index dd1cf768b..ec10c0d27 100644 --- a/reframe/core/launchers/__init__.py +++ b/reframe/core/launchers/__init__.py @@ -36,6 +36,22 @@ class JobLauncher(metaclass=_JobLauncherMeta): #: :default: ``[]`` options = variable(typ.List[str], value=[]) + #: Dictionary of environment variables to be passed via the job launcher + #: invocation. The keys are the variable names and the values are their + #: corresponding values. + #: + #: This is supported by the following launchers only: + #: + #: - ``srun`` + #: - ``mpirun-openmpi`` + #: - ``mpirun-intelmpi`` + #: + #: :type: :class:`Dict[str, str]` + #: :default: ``{}`` + #: + #: .. versionadded:: 4.9 + env_vars = variable(typ.Dict[str, str], value={}) + #: Optional modifier of the launcher command. #: #: This will be combined with the :attr:`modifier_options` and prepended to @@ -59,6 +75,7 @@ class JobLauncher(metaclass=_JobLauncherMeta): def __init__(self): self.options = [] + self.env_vars = {} @abc.abstractmethod def command(self, job): diff --git a/reframe/core/launchers/mpi.py b/reframe/core/launchers/mpi.py index 6e6acc56b..9616b8ac6 100644 --- a/reframe/core/launchers/mpi.py +++ b/reframe/core/launchers/mpi.py @@ -48,6 +48,10 @@ def command(self, job): if 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 @@ -109,6 +113,24 @@ def command(self, job): return ['mpirun', '-np', str(job.num_tasks)] +@register_launcher('mpirun-openmpi') +class MpirunOpenMPILauncher(JobLauncher): + def command(self, job): + cmd = ['mpirun', '-np', str(job.num_tasks)] + for name, value in self.env_vars.items(): + cmd += ['-x', f'{name}={value}'] + return cmd + + +@register_launcher('mpirun-intelmpi') +class MpirunIntelMPILauncher(JobLauncher): + def command(self, job): + cmd = ['mpirun', '-np', str(job.num_tasks)] + for name, value in self.env_vars.items(): + cmd += ['-genv', name, value] + return cmd + + @register_launcher('mpiexec') class MpiexecLauncher(JobLauncher): def command(self, job): diff --git a/unittests/test_launchers.py b/unittests/test_launchers.py index 561dffd86..a3675ba77 100644 --- a/unittests/test_launchers.py +++ b/unittests/test_launchers.py @@ -92,6 +92,7 @@ def job(make_job, launcher): job.exclusive_access = True job.options = ['--gres=gpu:4', '#DW jobdw anything'] job.launcher.options = ['--foo'] + job.launcher.env_vars = {'FOO': 'bar', 'BAR': 'baz'} return job @@ -127,7 +128,20 @@ def test_run_command(job): elif launcher_name == 'mpirun': assert command == 'mpirun -np 4 --foo' elif launcher_name == 'srun': - assert command == 'srun --cpus-per-task=2 --foo' + assert command == ('srun ' + '--cpus-per-task=2 ' + '--export=FOO=bar,BAR=baz ' + '--foo') + elif launcher_name == 'mpirun-openmpi': + assert command == ('mpirun -np 4 ' + '-x FOO=bar ' + '-x BAR=baz ' + '--foo') + elif launcher_name == 'mpirun-intelmpi': + assert command == ('mpirun -np 4 ' + '-genv FOO bar ' + '-genv BAR baz ' + '--foo') elif launcher_name == 'srunalloc': assert command == ('srun ' '--job-name=fake_job '