SLURM scheduler: using scheduler's env, but without mpi4py pollution

pytest_parallel/plugin.py

@@ -36,6 +36,19 @@ def pytest_addoption(parser):
     parser.addoption('--_scheduler_port', dest='_scheduler_port', type=int, help='Internal pytest_parallel option')
     parser.addoption('--_test_idx'    , dest='_test_idx'    , type=int, help='Internal pytest_parallel option')
 
+@pytest.hookimpl(tryfirst=True)
+def pytest_load_initial_conftests():
+    import sys
+    assert 'mpi4py.MPI' not in sys.modules, 'Internal pytest_parallel error: mpi4py.MPI should not be imported' \
+                                            ' when we are about to register and environment for SLURM' \
+                                            ' (because importing mpi4py.MPI makes the current process look like and MPI process,' \
+                                            ' and SLURM does not like that)'
+
+    import subprocess
+    r = subprocess.run(['env','--null'], stdout=subprocess.PIPE) # `--null`: end each output line with NUL, required by `sbatch --export-file`
+    assert r.returncode==0, 'SLURM scheduler: error when writing `env` to `pytest_slurm/env_vars.sh`'
+    pytest._pytest_parallel_env_vars = r.stdout
+
 # --------------------------------------------------------------------------
 @pytest.hookimpl(trylast=True)
 def pytest_configure(config):

pytest_parallel/process_scheduler.py

@@ -91,8 +91,11 @@ def submit_items(items_to_run, socket, main_invoke_params, slurm_ntasks, slurm_c
       f.write(pytest_slurm)
 
     # submit SLURM job
+    with open('pytest_slurm/env_vars.sh','wb') as f:
+      f.write(pytest._pytest_parallel_env_vars)
+
     if slurm_conf['sub_command'] is None:
-        sbatch_cmd = 'sbatch --parsable pytest_slurm/job.sh'
+        sbatch_cmd = 'sbatch --parsable --export-file=pytest_slurm/env_vars.sh pytest_slurm/job.sh'
     else:
         sbatch_cmd = slurm_conf['sub_command'] + ' pytest_slurm/job.sh'