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 <[email protected]>

Author: jack-morrison

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.
         '''

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

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 '