Feature request: support for exporting environment variables with parallel launchers

[casparvl]

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 by CSCS in their PyTorch test.

[vkarak]

I think that

self.job.launcher.env_vars = {'MASTER_PORT': '1234'}

or

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.