add intelmpi slurm container support

Author: stdweird

lib/vsc/mympirun/mpi/intelmpi.py

@@ -174,6 +174,25 @@ def mpirun_prepare_execution(self):
 
         return super().mpirun_prepare_execution()
 
+    def prepare(self):
+        super().prepare()
+
+        slurm_container = os.environ.get('SLURM_CONTAINER', None)
+        if slurm_container is not None:
+            logging.debug(f"Found SLURM_CONTAINER {slurm_container}")
+
+            # intel mpi has some support for running with singularity, but not for slurm container mode
+            for key in [k for k in os.environ.keys() if k.startswith('APPTAINER') or k.startswith('SINGULARITY')]:
+                logging(f"Removing environment variable {key} (value {os.environ[key]}) in slurm container mode")
+                del os.environ[key]
+
+            if self._has_hydra() and self.HYDRA_RMK and self.HYDRA_RMK[0] == 'slurm':
+                # need to pass full environment to tasks with slurm container wrapper
+                if slurm_container.startswith('HPCWN'):
+                    slurm_container += ':fullenv'
+                # pass the container option to srun in slurm bootstrap
+                os.environ['I_MPI_HYDRA_BOOTSTRAP_EXEC_EXTRA_ARGS'] = f"--container {slurm_container}"
+
     def pinning_override(self):
         """ pinning """