From 10fe10ee79bb8b30a4d3904d9a4e42c7f534999b Mon Sep 17 00:00:00 2001
From: Nathaniel Dean <nathaniel.dean@nebius.com>
Date: Thu, 2 Jul 2026 09:26:42 -0400
Subject: [PATCH] adding nvidia library active checks

---
 .../slurm_scripts/check_nvidia_libraries.sh   | 40 +++++++++++++++++++
 .../check_nvidia_libraries.sh.json            | 17 ++++++++
 helm/slurm-cluster/values.yaml                |  4 ++
 3 files changed, 61 insertions(+)
 create mode 100644 helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh
 create mode 100644 helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh.json

diff --git a/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh b/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh
new file mode 100644
index 000000000..8cc874983
--- /dev/null
+++ b/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh
@@ -0,0 +1,40 @@
+#!/bin/bash
+
+set -euxo pipefail
+
+fail_if_lib_dne () {
+	checklib=$1
+	if python -c "import ctypes; ctypes.CDLL('$checklib')"; then
+		echo "...found $checklib, continuing"
+	else
+		echo "ERROR: missing $checklib"
+		exit 1
+	fi
+}
+
+# first check if I can load libnvidia-ml.so.1, because nothing works without it
+# (also not a typo, looks specifically for .1)
+
+fail_if_lib_dne "libnvidia-ml.so.1"
+
+# get list of base Nvidia libraries to check from nvidia-container-cli
+full_library_list=$(nvidia-container-cli list -l | grep \\.so)
+for x in $full_library_list; do
+	fail_if_lib_dne "$x"
+done
+
+# check against a list of commonly needed CUDA libraries that are common in a
+# standard Pytorch training session
+
+cudalibs="libcuda libcudart libcublas libcublasLt libcufft libcurand libcusolver libcusparse libcudnn libnccl libucp libucs libuct libgdrapi libmpi libnuma"
+libs_w_version="libibverbs.so.1 librdmacm.so.1 libgomp.so.1"
+
+for x in $cudalibs; do
+	fail_if_lib_dne "$x.so"
+done
+for x in $libs_w_version; do
+	fail_if_lib_dne "$x"
+done
+
+echo "OK: all libraries exist"
+exit 0 
diff --git a/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh.json b/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh.json
new file mode 100644
index 000000000..8507595fb
--- /dev/null
+++ b/helm/slurm-cluster/slurm_scripts/check_nvidia_libraries.sh.json
@@ -0,0 +1,17 @@
+{
+  "name": "check_nvidia_libraries",
+  "command": "./check_nvidia_libraries.sh",
+  "platforms": ["8xGPU", "1xGPU", "4xGPU"],
+  "skip_for_cpu_jobs": false,
+  "skip_for_partial_gpu_jobs": false,
+  "skip_for_reservation_prefixes": [],
+  "contexts": ["prolog", "hc_program"],
+  "node_states": ["any"],
+  "on_fail": "drain",
+  "on_ok": "undrain",
+  "reason_base": "[node_problem] $name",
+  "reason_append_details": true,
+  "run_in_jail": true,
+  "log": "slurm_scripts/$worker.$name.$context.out",
+  "need_env": []
+}
diff --git a/helm/slurm-cluster/values.yaml b/helm/slurm-cluster/values.yaml
index 1843b5789..3dfadb705 100644
--- a/helm/slurm-cluster/values.yaml
+++ b/helm/slurm-cluster/values.yaml
@@ -744,6 +744,10 @@ slurmScripts:
       enabled: true
       customContent: null
       customConfig: null
+    check_nvidia_libraries.sh:
+      enabled: true
+      customContent: null
+      customConfig: null
     chmod_enroot_layers.sh:
       enabled: true
       customContent: null