From c8c4df423dbca0a65b5d7a6ccb963b627379c15f 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] Probe DCGM hostengine GPU visibility in the liveness check

A hostengine that starts while the nvidia kernel module is not loaded gets
plain runc from the toolkit's host wrapper, comes up without NVML, and never
re-initialises. The tcpSocket liveness probe on 5555 still passes because the
process keeps listening, so the pod stays 1/1 Running and exports no DCGM_FI_*
series for the node.

Make the liveness probe run `dcgmi discovery -l` and require it to report at
least one active GPU, so kubelet restarts a deaf hostengine once the driver is
back. The pattern matches the current "N GPUs found (Active)." summary line and
the plain "N GPUs found." form, and rejects "0 GPUs found". The check is skipped
while no driver is present (no `nvidia` module and no WSL2 /dev/dxg), so a
planned driver restart does not loop the hostengine; once the driver returns the
probe fails and kubelet restarts the deaf hostengine. Readiness is unchanged.

The init-container half of the original change is now redundant: #2881 already
requires a live nvidia module, or the WSL2 /dev/dxg path, for every operand that
uses the NVIDIA runtime, so this rebase keeps only the probe change.

Refs #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/0400_dcgm.yml | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/assets/state-dcgm/0400_dcgm.yml b/assets/state-dcgm/0400_dcgm.yml
index 6f5b348f6..14d10ad7d 100644
--- a/assets/state-dcgm/0400_dcgm.yml
+++ b/assets/state-dcgm/0400_dcgm.yml
@@ -43,10 +43,23 @@ 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. Check GPU visibility, but only while a driver
+        # is present, so a planned driver restart does not loop the hostengine.
         livenessProbe:
-          tcpSocket:
-            port: 5555
-          initialDelaySeconds: 15
+          exec:
+            command:
+            - sh
+            - -c
+            - >-
+              if [ -e /dev/dxg ] || grep -q '^nvidia ' /proc/modules;
+              then dcgmi discovery -l | grep -qE '^[1-9][0-9]* GPUs? found';
+              else exit 0; fi
+          initialDelaySeconds: 30
+          periodSeconds: 60
+          timeoutSeconds: 15
+          failureThreshold: 3
         readinessProbe:
           tcpSocket:
             port: 5555