Repository: reframe-hpc/reframe

Issue:
Feature request: support for exporting environment variables with parallel launchers
Some software requires environment variables to run - e.g. PyTorch's distributed framework requires `MASTER_PORT` (among others) to be set. As discussed on Slack, this is currently challenging if the test developer doesn't know the configured launcher in advance.

I.e. if we _know_ that OpenMPI's `mpirun` will be the launcher, we can do

```
self.job.launcher.options = ['-x MASTER_PORT']
```

But if we are writing a test with the purpose of it being reused (e.g. a test for the `hpctestlib`), it would be nice to have a way of specifying this in a launcher-agnostic way. E.g.

```
test.env_vars['MASTER_PORT'] = '1234'
self.job.launcher.export_var = ['MASTER_PORT']
```
or
```
self.job.launcher.export_var['MASTER_PORT] = ['1234']
```
(the 2nd is probably more convenient, but not sure which API is easiest to support from the ReFrame side).

ReFrame would then abstract how each particular launcher exports environment variables. E.g. for OpenMPI, the ReFrame backend would add `-x MASTER_PORT=1234` as extra launcher argument, whereas for `srun` it would add `--export=MASTERPORT=1234`.

Note that right now, I worked around this issue by making a wrapper shell script that sets the environment variables, similar to what is used [here](https://github.com/reframe-hpc/cscs-reframe-tests/blob/33ef7eab2988950551fc19f0e5816c55dc23500f/checks/apps/pytorch/pytorch_distr_check.py#L67) by CSCS in their PyTorch test.

Discussion:
@vkarak: I think that

```python
self.job.launcher.env_vars = {'MASTER_PORT': '1234'}
```

or

```python
self.job.launcher.env_vars['MASTER_PORT] = '1234'
```

is the best and matches the test's `env_vars` in the syntax.@vkarak: @casparvl I guess that your suggestion is for cases that the scheduler does not forward the environment to the compute nodes, right? If that's the case, maybe we would just need a configuration parameter to tell ReFrame to export the test's `env_vars` when launching the job. Do you see a case that you would need to pass a different set of environment variables to the launcher?@casparvl: > @casparvl I guess that your suggestion is for cases that the scheduler does not forward the environment to the compute nodes, right? 

Hmm, that was actually not my problem. What you're talking about (I think) is exporting variables from the submission environment to the batch job. What I'm talking about is exporting the environment from the job's head-node of the allocation, to the processes launched by a parallel launcher like `srun` or `mpirun`. The trouble is that there is no launcher-agnostic way to do that, and it would be nice if ReFrame would abstract that away, so that we can define e.g. in the test:

```
self.job.launcher.env_vars['MASTER_PORT] = '1234'
self.executable = python
self.executable_opts = ['pytorch_benchmark.py']
```
And that, when configured with e.g. `mpirun` as parallel launcher, ReFrame would create the following job script:
```
#SBATCH ...
module load ...
mpirun -x MASTER_PORT=1234 python pytorch_benchmark.py
```
But, when configured with `srun` as parallel launcher, ReFrame would created:
```
#SBATCH ...
module load ...
srun --export=MASTER_PORT=1234 python pytorch_benchmark.py
```
I.e. in that way, we can write tests that can tell the parallel launcher to export an environment variable to the parallel processes being launched _without_ the test developer having to know which parallel launcher will be used.@vkarak: Thanks @casparvl for the clarification; it makes sense. I was confused at the beginning.@jack-morrison: @vkarak, we discussed modifying the `JobLauncher` base class and implementing MPI-specific launchers, like an IntelMPI mpirun (`impi_mpirun`) and an OpenMPI mpirun (`ompi_mpirun`).

After thinking about this some more, doesn't this just shift the reuse problem from a test definition into the config files?

If I'm not mistaken, in order to use these MPI-specific launchers, a user would have to have multiple `system.partition.launcher`s, meaning new `system:partition`s. This seems inconvenient. 

Am I thinking about this correctly?

Edit: Maybe this approach is fine... `system.partition.launcher` already implies a single launcher is configured. I'm just thinking about situations where a user may want to use either OpenMPI or IntelMPI (or others) on a single partition.@vkarak: > Edit: Maybe this approach is fine... system.partition.launcher already implies a single launcher is configured. I'm just thinking about situations where a user may want to use either OpenMPI or IntelMPI (or others) on a single partition.

Yes, this is a general assumption of the system partitions, so we are not doing anything different. The key problem we are trying to solve here is that exporting variables is specific to the MPI launcher flavor. We will not replace the generic `mpirun` launcher, though it won't support/interpret the discussed `env_vars`. If users want to export variables at the launcher level, would then have to pick a more specific launcher for their partition.

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


Test patch:
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 '