Added NVML and cgroup support to the slurm package (#40638)

* Added NVML support to the slurm package
* dbus package is required for cgroup support
* Fixing formatting
* Style fix
* Added PAM support
* Added ROCm SMI support

1 files changed, 18 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/slurm/package.py b/var/spack/repos/builtin/packages/slurm/package.py
index 61214702b0..aa4f126018 100644
--- a/var/spack/repos/builtin/packages/slurm/package.py
+++ b/var/spack/repos/builtin/packages/slurm/package.py
@@ -129,6 +129,10 @@ class Slurm(AutotoolsPackage):
         description="Set system configuration path (possibly /etc/slurm)",
     )
     variant("restd", default=False, description="Enable the slurmrestd server")
+    variant("nvml", default=False, description="Enable NVML autodetection")
+    variant("cgroup", default=False, description="Enable cgroup plugin")
+    variant("pam", default=False, description="Enable PAM support")
+    variant("rsmi", default=False, description="Enable ROCm SMI support")
 
     # TODO: add variant for BG/Q and Cray support
 
@@ -156,6 +160,11 @@ class Slurm(AutotoolsPackage):
     depends_on("libyaml", when="+restd")
     depends_on("libjwt", when="+restd")
 
+    depends_on("cuda", when="+nvml")
+    depends_on("dbus", when="+cgroup")
+    depends_on("linux-pam", when="+pam")
+    depends_on("rocm-smi-lib", when="+rsmi")
+
     executables = ["^srun$", "^salloc$"]
 
     @classmethod
@@ -213,6 +222,15 @@ class Slurm(AutotoolsPackage):
         else:
             args.append("--without-pmix")
 
+        if spec.satisfies("+nvml"):
+            args.append(f"--with-nvml={spec['cuda'].prefix}")
+
+        if spec.satisfies("+pam"):
+            args.append(f"--with-pam_dir={spec['linux-pam'].prefix}")
+
+        if spec.satisfies("+rsmi"):
+            args.append(f"--with-rsmi={spec['rocm-smi-lib'].prefix}")
+
         sysconfdir = spec.variants["sysconfdir"].value
         if sysconfdir != "PREFIX/etc":
             args.append("--sysconfdir={0}".format(sysconfdir))