From 81997ae6d606d458bf88d7e755c4226ec49a5a3f Mon Sep 17 00:00:00 2001
From: G-Ragghianti <33492707+G-Ragghianti@users.noreply.github.com>
Date: Mon, 30 Oct 2023 22:12:09 -0400
Subject: 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
---
 var/spack/repos/builtin/packages/dbus/package.py  |  4 ++++
 var/spack/repos/builtin/packages/slurm/package.py | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/var/spack/repos/builtin/packages/dbus/package.py b/var/spack/repos/builtin/packages/dbus/package.py
index 69cda7b477..37a1b8a694 100644
--- a/var/spack/repos/builtin/packages/dbus/package.py
+++ b/var/spack/repos/builtin/packages/dbus/package.py
@@ -29,6 +29,7 @@ class Dbus(AutotoolsPackage):
     version("1.8.2", sha256="5689f7411165adc953f37974e276a3028db94447c76e8dd92efe910c6d3bae08")
 
     variant("xml_docs", default=False, description="Build XML documentation")
+    variant("system-socket", default="default", description="Location for the DBus system socket")
 
     depends_on("pkgconfig", type="build")
     depends_on("docbook-xml", type="build")
@@ -41,6 +42,9 @@ class Dbus(AutotoolsPackage):
     def configure_args(self):
         args = ["--disable-systemd", "--disable-launchd"]
         args += self.enable_or_disable("xml-docs", variant="xml_docs")
+        socket = self.spec.variants["system-socket"].value
+        if socket != "default":
+            args += ["--with-system-socket={0}".format(socket)]
         return args
 
     @run_after("install")
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))
-- 
cgit v1.2.3-70-g09d2