From 119d901c507abcec7da9d458084df4257dec84c3 Mon Sep 17 00:00:00 2001
From: Weigang Geng <3356786+gengwg@users.noreply.github.com>
Date: Thu, 3 Sep 2026 23:02:51 -0700
Subject: [PATCH] Gate dcgm and dcgm-exporter on the driver module and probe
 hostengine GPU visibility

The toolkit-validation init container in state-dcgm and state-dcgm-exporter
only waits for /run/nvidia/validations/toolkit-ready, which survives an
in-place driver container restart. A hostengine that starts while the nvidia
module is unloaded gets plain runc from the toolkit's host wrapper, has no
NVML, and never recovers; the tcpSocket liveness probe passes because it
still listens on 5555. Both pods then sit 1/1 Running and export no GPU
metrics for the node.

Make the init gate also require the nvidia module in /proc/modules, and make
the hostengine liveness probe check that dcgmi discovery sees at least one
GPU so kubelet restarts a deaf hostengine once the driver is back.

Fixes #2854

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Weigang Geng <3356786+gengwg@users.noreply.github.com>
---
 assets/state-dcgm-exporter/0800_daemonset.yaml |  2 +-
 assets/state-dcgm/0400_dcgm.yml                | 17 +++++++++++++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/assets/state-dcgm-exporter/0800_daemonset.yaml b/assets/state-dcgm-exporter/0800_daemonset.yaml
index eb0dab1635..ba802a9f29 100644
--- a/assets/state-dcgm-exporter/0800_daemonset.yaml
+++ b/assets/state-dcgm-exporter/0800_daemonset.yaml
@@ -29,7 +29,7 @@ spec:
       - name: toolkit-validation
         image: "FILLED BY THE OPERATOR"
         command: ['sh', '-c']
-        args: ["until [ -f /run/nvidia/validations/toolkit-ready ]; do echo waiting for nvidia container stack to be setup; sleep 5; done"]
+        args: ["until [ -f /run/nvidia/validations/toolkit-ready ] && grep -q '^nvidia ' /proc/modules; do echo waiting for nvidia container stack and driver modules to be setup; sleep 5; done"]
         securityContext:
           privileged: true
         volumeMounts:
diff --git a/assets/state-dcgm/0400_dcgm.yml b/assets/state-dcgm/0400_dcgm.yml
index 0c27fd4e60..f043e5833f 100644
--- a/assets/state-dcgm/0400_dcgm.yml
+++ b/assets/state-dcgm/0400_dcgm.yml
@@ -28,7 +28,7 @@ spec:
       - name: toolkit-validation
         image: "FILLED BY THE OPERATOR"
         command: ['sh', '-c']
-        args: ["until [ -f /run/nvidia/validations/toolkit-ready ]; do echo waiting for nvidia container stack to be setup; sleep 5; done"]
+        args: ["until [ -f /run/nvidia/validations/toolkit-ready ] && grep -q '^nvidia ' /proc/modules; do echo waiting for nvidia container stack and driver modules to be setup; sleep 5; done"]
         securityContext:
           privileged: true
         volumeMounts:
@@ -43,10 +43,19 @@ spec:
         ports:
         - name: "dcgm"
           containerPort: 5555
+        # A hostengine that started without NVML (e.g. before the driver
+        # modules were loaded) keeps listening on 5555 but sees no GPUs and
+        # never re-initialises, so the probe checks GPU visibility.
         livenessProbe:
-          tcpSocket:
-            port: 5555
-          initialDelaySeconds: 15
+          exec:
+            command:
+            - sh
+            - -c
+            - dcgmi discovery -l | grep -qE '^[1-9][0-9]* GPUs? found'
+          initialDelaySeconds: 30
+          periodSeconds: 60
+          timeoutSeconds: 15
+          failureThreshold: 3
         readinessProbe:
           tcpSocket:
             port: 5555