Merge pull request #3464 from jack-morrison/jackm/launcher-fwd-env-vars

[feat] Enable launchers to forward environment variables

Author: vkarak

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  <https://hpc.llnl.gov/training/tutorials/using-lcs-sierra-system#lrun>`__'s ``lrun`` command.
    - ``lrun-gpu``: Parallel programs will be launched using `LC Launcher <https://hpc.llnl.gov/training/tutorials/using-lcs-sierra-system#lrun>`__'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 <https://slurm.schedmd.com/srun.html>`__'s ``srun`` command.
    - ``srunalloc``: Parallel programs will be launched using `Slurm <https://slurm.schedmd.com/srun.html>`__'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 <https://upc.lbl.gov/>`__ ``upcrun`` command.
    - ``upcxx-run``: Parallel programs will be launched using the `UPC++ <https://bitbucket.org/berkeleylab/upcxx/wiki/Home>`__ ``upcxx-run`` command.
 
+   .. versionadded:: 4.9
+
+      The ``mpirun-intelmpi`` and ``mpirun-openmpi`` parallel launchers were added.
+
    .. tip::
 
       .. versionadded:: 4.0.0

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):

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):

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 '