From fd56ed4e8916007384a0ada9c9ba2d96417291ff Mon Sep 17 00:00:00 2001 From: Jack Morrison Date: Tue, 8 Apr 2025 10:55:27 -0400 Subject: [PATCH 1/3] Enable launchers to forward environment variables Support for environment variable forwarding is added to the srun launcher, as well as two new launchers (mpirun-openmpi, mpirun-intelmpi). Signed-off-by: Jack Morrison --- reframe/core/launchers/__init__.py | 11 +++++++++-- reframe/core/launchers/mpi.py | 20 +++++++++++++++++++- unittests/test_launchers.py | 7 ++++++- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/reframe/core/launchers/__init__.py b/reframe/core/launchers/__init__.py index dd1cf768b0..26081b1de6 100644 --- a/reframe/core/launchers/__init__.py +++ b/reframe/core/launchers/__init__.py @@ -36,6 +36,14 @@ 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. + #: + #: :type: :class:`Dict[str, str]` + #: :default: ``{}`` + environment_variables = 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 +67,7 @@ class JobLauncher(metaclass=_JobLauncherMeta): def __init__(self): self.options = [] + self.environment_variables = {} @abc.abstractmethod def command(self, job): @@ -73,8 +82,6 @@ def command(self, job): def run_command(self, job): '''The full launcher command to be emitted for a specific job. - This includes any user options. - :param job: a job descriptor. :returns: the launcher command as a string. ''' diff --git a/reframe/core/launchers/mpi.py b/reframe/core/launchers/mpi.py index 6e6acc56b8..04e2365700 100644 --- a/reframe/core/launchers/mpi.py +++ b/reframe/core/launchers/mpi.py @@ -48,8 +48,11 @@ 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}') - return ret + if self.environment_variables: + env_vars = ','.join(f'{k}={v}' for k, v in self.environment_variables.items()) + ret.append(f'--export={env_vars}') + return ret @register_launcher('ibrun') class IbrunLauncher(JobLauncher): @@ -108,6 +111,21 @@ class MpirunLauncher(JobLauncher): 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.environment_variables.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.environment_variables.items(): + cmd += ['-genv', name, value] + return cmd @register_launcher('mpiexec') class MpiexecLauncher(JobLauncher): diff --git a/unittests/test_launchers.py b/unittests/test_launchers.py index 561dffd86e..85a4741d71 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.environment_variables = {'FOO': 'bar', 'BAR': 'baz'} return job @@ -127,7 +128,11 @@ 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 ' From ad7af822c3f0a43bcd97b57643f3bfcd6dbe1cb9 Mon Sep 17 00:00:00 2001 From: Jack Morrison Date: Tue, 29 Apr 2025 22:02:57 -0400 Subject: [PATCH 2/3] Address PR comments for formatting, naming, line length, and docstrings Signed-off-by: Jack Morrison --- reframe/core/launchers/__init__.py | 12 ++++++++++-- reframe/core/launchers/mpi.py | 11 +++++++---- unittests/test_launchers.py | 17 +++++++++++++---- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/reframe/core/launchers/__init__.py b/reframe/core/launchers/__init__.py index 26081b1de6..ff89805a5f 100644 --- a/reframe/core/launchers/__init__.py +++ b/reframe/core/launchers/__init__.py @@ -40,9 +40,15 @@ class JobLauncher(metaclass=_JobLauncherMeta): #: invocation. The keys are the variable names and the values are their #: corresponding values. #: + #: This is supported by the following launchers: + #: + #: - ``srun`` + #: - ``mpirun-openmpi`` + #: - ``mpirun-intelmpi`` + #: #: :type: :class:`Dict[str, str]` #: :default: ``{}`` - environment_variables = variable(typ.Dict[str, str], value={}) + env_vars = variable(typ.Dict[str, str], value={}) #: Optional modifier of the launcher command. #: @@ -67,7 +73,7 @@ class JobLauncher(metaclass=_JobLauncherMeta): def __init__(self): self.options = [] - self.environment_variables = {} + self.env_vars = {} @abc.abstractmethod def command(self, job): @@ -82,6 +88,8 @@ def command(self, job): def run_command(self, job): '''The full launcher command to be emitted for a specific job. + This includes any user options. + :param job: a job descriptor. :returns: the launcher command as a string. ''' diff --git a/reframe/core/launchers/mpi.py b/reframe/core/launchers/mpi.py index 04e2365700..c0e94fd7f2 100644 --- a/reframe/core/launchers/mpi.py +++ b/reframe/core/launchers/mpi.py @@ -48,8 +48,8 @@ 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.environment_variables: - env_vars = ','.join(f'{k}={v}' for k, v in self.environment_variables.items()) + 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 @@ -111,22 +111,25 @@ class MpirunLauncher(JobLauncher): 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.environment_variables.items(): + 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.environment_variables.items(): + 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 85a4741d71..a3675ba771 100644 --- a/unittests/test_launchers.py +++ b/unittests/test_launchers.py @@ -92,7 +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.environment_variables = {'FOO': 'bar', 'BAR': 'baz'} + job.launcher.env_vars = {'FOO': 'bar', 'BAR': 'baz'} return job @@ -128,11 +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 --export=FOO=bar,BAR=baz --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' + 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' + assert command == ('mpirun -np 4 ' + '-genv FOO bar ' + '-genv BAR baz ' + '--foo') elif launcher_name == 'srunalloc': assert command == ('srun ' '--job-name=fake_job ' From 995a9e73b6094bb7b6dc23a944fc6a7c303d5279 Mon Sep 17 00:00:00 2001 From: Vasileios Karakasis Date: Tue, 6 May 2025 23:58:52 +0200 Subject: [PATCH 3/3] Update parallel launcher docs --- docs/config_reference.rst | 8 +++++++- reframe/core/launchers/__init__.py | 4 +++- reframe/core/launchers/mpi.py | 1 + 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/docs/config_reference.rst b/docs/config_reference.rst index 4f48e20a3a..3ab38af0c9 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 ff89805a5f..ec10c0d275 100644 --- a/reframe/core/launchers/__init__.py +++ b/reframe/core/launchers/__init__.py @@ -40,7 +40,7 @@ class JobLauncher(metaclass=_JobLauncherMeta): #: invocation. The keys are the variable names and the values are their #: corresponding values. #: - #: This is supported by the following launchers: + #: This is supported by the following launchers only: #: #: - ``srun`` #: - ``mpirun-openmpi`` @@ -48,6 +48,8 @@ class JobLauncher(metaclass=_JobLauncherMeta): #: #: :type: :class:`Dict[str, str]` #: :default: ``{}`` + #: + #: .. versionadded:: 4.9 env_vars = variable(typ.Dict[str, str], value={}) #: Optional modifier of the launcher command. diff --git a/reframe/core/launchers/mpi.py b/reframe/core/launchers/mpi.py index c0e94fd7f2..9616b8ac6f 100644 --- a/reframe/core/launchers/mpi.py +++ b/reframe/core/launchers/mpi.py @@ -54,6 +54,7 @@ def command(self, job): return ret + @register_launcher('ibrun') class IbrunLauncher(JobLauncher): '''TACC's custom parallel job launcher.'''