mvapich2: add pmi_version variant for pmix, pmi1 support (#40665)

1 files changed, 24 insertions, 6 deletions

diff --git a/var/spack/repos/builtin/packages/mvapich2/package.py b/var/spack/repos/builtin/packages/mvapich2/package.py
index 246d013e3b..7e9204ffbd 100644
--- a/var/spack/repos/builtin/packages/mvapich2/package.py
+++ b/var/spack/repos/builtin/packages/mvapich2/package.py
@@ -84,6 +84,15 @@ class Mvapich2(AutotoolsPackage):
     )
 
     variant(
+        "pmi_version",
+        description=("The pmi version to be used with slurm"),
+        when="process_managers=slurm",
+        default="pmi2",
+        values=("pmi1", "pmi2", "pmix"),
+        multi=False,
+    )
+
+    variant(
         "fabrics",
         description="Select the fabric to be enabled for this build."
         "If you have verbs (either from OFED or MOFED), PSM or PSM2 "
@@ -133,6 +142,7 @@ class Mvapich2(AutotoolsPackage):
     depends_on("rdma-core", when="fabrics=nemesisibtcp")
     depends_on("libfabric", when="fabrics=nemesisofi")
     depends_on("slurm", when="process_managers=slurm")
+    depends_on("pmix", when="pmi_version=pmix")
 
     # Fix segmentation fault in `MPIR_Attr_delete_list`:
     # <https://lists.osu.edu/pipermail/mvapich-discuss/2023-January/010695.html>.
@@ -277,11 +287,14 @@ class Mvapich2(AutotoolsPackage):
 
         # See: http://slurm.schedmd.com/mpi_guide.html#mvapich2
         if "process_managers=slurm" in spec:
-            opts = [
-                "--with-pmi=pmi2",
-                "--with-pm=slurm",
-                "--with-slurm={0}".format(spec["slurm"].prefix),
-            ]
+            opts = ["--with-pm=slurm", "--with-slurm={0}".format(spec["slurm"].prefix)]
+            if "pmi_version=pmi1" in spec:
+                opts.append("--with-pmi=pmi1")
+            elif "pmi_version=pmi2" in spec:
+                opts.append("--with-pmi=pmi2")
+            elif "pmi_version=pmix" in spec:
+                opts.append("--with-pmi=pmix")
+                opts.append("--with-pmix={0}".format(spec["pmix"].prefix))
 
         return opts
 
@@ -347,7 +360,12 @@ class Mvapich2(AutotoolsPackage):
 
     def setup_run_environment(self, env):
         if "process_managers=slurm" in self.spec:
-            env.set("SLURM_MPI_TYPE", "pmi2")
+            if "pmi_version=pmi1" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmi1")
+            elif "pmi_version=pmi2" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmi2")
+            elif "pmi_version=pmix" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmix")
 
         env.set("MPI_ROOT", self.prefix)