Add opt-in /health check for GPU collector presence

deployment/README.md

@@ -243,6 +243,26 @@ exporter-toolkit web config and set
 file, because `/debug/pprof/` can expose runtime profiling details and must be
 protected by exporter-toolkit authentication or TLS.
 
+### Restarting a blind exporter
+
+A DCGM hostengine that starts before the NVIDIA driver is ready comes up
+without NVML and reports no GPUs. An exporter that connects to it in that state
+registers no GPU collector and serves an empty metrics page for the life of the
+process, while `/health` still returns 200.
+
+Adding `--health-require-gpus` to `arguments` makes `/health` return 503 in that
+state, so the liveness probe restarts the pod and it re-enumerates. It is off by
+default because the exporter cannot distinguish a node with no GPUs from a
+hostengine reporting none.
+
+Two caveats:
+
+- When `basicAuth.users` is set, the chart degrades both probes to `tcpSocket`,
+  which cannot observe the 503. The flag then has no effect on restarts.
+- The check counts collectors registered under `FE_GPU`. A counters file that
+  contributes no `FE_GPU` fields also reads as zero, so do not enable this with
+  a switch- or CPU-only counter set.
+
 ## Troubleshooting
 
 ### Debug Dump Files

deployment/values.yaml

@@ -46,6 +46,12 @@ arguments: []
 # If adding "--enable-pprof", also enable tlsServerConfig and/or basicAuth.
 # Pprof requires a web config file so /debug/pprof/ is protected by
 # exporter-toolkit authentication or TLS.
+# Use "--health-require-gpus" to have /health report 503 while no GPU collector
+# is registered, so the liveness probe restarts an exporter that started against
+# a hostengine with no GPUs visible. Off by default.
+# Example arguments: ["--health-require-gpus"]
+# NOTE: this has no effect when basicAuth.users is set - the chart then degrades
+# both probes to tcpSocket, which cannot see the 503.
 
 # Optional dcgm-exporter YAML configuration. YAML is read at exporter startup.
 # Changes to the YAML file require restarting the exporter pod.

internal/pkg/appconfig/types.go

@@ -104,6 +104,7 @@ type Config struct {
 	CPUDeviceOptions                 DeviceOptions
 	NoHostname                       bool
 	UseFakeGPUs                      bool
+	HealthRequireGPUs                bool
 	ConfigMapData                    string
 	MetricSource                     MetricSource
 	WatchGroups                      []WatchGroup

internal/pkg/registry/registry.go

@@ -67,6 +67,21 @@ func (r *Registry) Register(entityCollectorTuples collector.EntityCollectorTuple
 	r.collectorGroupsSeen[entityCollectorTuples] = struct{}{}
 }
 
+// CollectorCount returns the number of collectors registered for an entity
+// group. Zero for dcgm.FE_GPU means this registry was built without any GPU
+// collector, which happens when the hostengine it connected to had no GPUs
+// visible. Absent a reload or a GPU bind event the registry is not rebuilt, so
+// that state persists.
+//
+// Safe to call once the registry has been published via Store/Swap: Register
+// runs before publication and does not lock.
+func (r *Registry) CollectorCount(entityGroup dcgm.Field_Entity_Group) int {
+	r.mtx.RLock()
+	defer r.mtx.RUnlock()
+
+	return len(r.collectorGroups[entityGroup])
+}
+
 // Gather gathers metrics from all registered collectors.
 func (r *Registry) Gather() (MetricsByCounterGroup, error) {
 	// Check if registry is shutting down

internal/pkg/registry/registry_test.go

@@ -193,3 +193,25 @@ func TestRegistry_Register_Accepts_Duplicates_(t *testing.T) {
 	assert.Len(t, reg.collectorGroups, 1)
 	assert.Len(t, reg.collectorGroupsSeen, 1)
 }
+
+func TestRegistryCollectorCountIsKeyedByEntityGroup(t *testing.T) {
+	reg := NewRegistry()
+
+	cpuTuple := collectorpkg.EntityCollectorTuple{}
+	cpuTuple.SetEntity(dcgm.FE_CPU)
+	cpuTuple.SetCollector(&mockCollector{})
+	reg.Register(cpuTuple)
+
+	// One collector is registered, but none of them is a GPU collector. A
+	// count that ignored the entity group would report 1 for FE_GPU here.
+	assert.Equal(t, 1, reg.CollectorCount(dcgm.FE_CPU))
+	assert.Equal(t, 0, reg.CollectorCount(dcgm.FE_GPU))
+
+	gpuTuple := collectorpkg.EntityCollectorTuple{}
+	gpuTuple.SetEntity(dcgm.FE_GPU)
+	gpuTuple.SetCollector(&mockCollector{})
+	reg.Register(gpuTuple)
+
+	assert.Equal(t, 1, reg.CollectorCount(dcgm.FE_GPU))
+	assert.Equal(t, 0, reg.CollectorCount(dcgm.FE_SWITCH))
+}

internal/pkg/server/server.go

@@ -28,6 +28,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/NVIDIA/go-dcgm/pkg/dcgm"
 	"github.com/gorilla/mux"
 	"github.com/prometheus/exporter-toolkit/web"
 
@@ -423,14 +424,37 @@ func (s *MetricsServer) Health(w http.ResponseWriter, _ *http.Request) {
 	}
 
 	// Check the raw atomic value to see if registry is nil
-	if s.registry.Load() == nil {
+	reg := s.registry.Load()
+	if reg == nil {
 		w.Header().Set("X-Registry-Available", "false")
 		w.Header().Set("X-Reload-In-Progress", "true")
 		_, _ = w.Write([]byte("OK - reload in progress"))
 		return
 	}
 
 	w.Header().Set("X-Registry-Available", "true")
+
+	// A registry built while the hostengine had no GPUs visible carries no GPU
+	// collector, and is only rebuilt on reload or a bind event, so on its own the
+	// exporter serves an empty metrics page while still reporting healthy.
+	// Reporting that as unhealthy lets a liveness probe already pointed at /health
+	// restart the pod so it re-enumerates once the hostengine itself is sighted.
+	// While the hostengine stays blind this is a restart loop rather than a silent
+	// blind pod - the loud failure is the point.
+	//
+	// Note this counts collectors registered under FE_GPU, which is a proxy for
+	// GPU visibility: a counters file contributing no FE_GPU fields also yields
+	// zero. See the flag's usage text.
+	if s.config != nil && s.config.HealthRequireGPUs && !s.IsReloadInProgress() {
+		// reg is non-nil: the nil case returned above. Load once so an unbind
+		// swapping it to nil between loads cannot skip the check.
+		if reg.CollectorCount(dcgm.FE_GPU) == 0 {
+			w.Header().Set("X-GPU-Collectors", "0")
+			http.Error(w, "KO - no GPU collector registered", http.StatusServiceUnavailable)
+			return
+		}
+	}
+
 	_, err := w.Write([]byte("OK"))
 	if err != nil {
 		slog.Error(failedWriteResponseError, slog.String(logging.ErrorKey, err.Error()))

internal/pkg/server/server_test.go

@@ -404,6 +404,74 @@ func TestHealthReturnsOKWithRegistryAvailable(t *testing.T) {
 	assert.NotEqual(t, "true", recorder.Header().Get("X-Reload-In-Progress"))
 }
 
+func TestHealthReturnsOKWithNoGPUCollectorsWhenNotRequired(t *testing.T) {
+	// Default behaviour must not change: an exporter with no GPU collector still
+	// reports healthy unless HealthRequireGPUs is set.
+	metricServer := &MetricsServer{config: &appconfig.Config{}}
+	metricServer.registry.Store(registry.NewRegistry())
+	recorder := httptest.NewRecorder()
+	metricServer.Health(recorder, nil)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+	assert.Empty(t, recorder.Header().Get("X-GPU-Collectors"))
+}
+
+func TestHealthReturnsUnavailableWhenGPUCollectorsRequiredButAbsent(t *testing.T) {
+	metricServer := &MetricsServer{config: &appconfig.Config{HealthRequireGPUs: true}}
+	metricServer.registry.Store(registry.NewRegistry())
+	recorder := httptest.NewRecorder()
+	metricServer.Health(recorder, nil)
+	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
+	assert.Equal(t, "0", recorder.Header().Get("X-GPU-Collectors"))
+	assert.Contains(t, recorder.Body.String(), "no GPU collector registered")
+}
+
+func TestHealthReturnsUnavailableWhenOnlyNonGPUCollectorsRegistered(t *testing.T) {
+	// The incident shape: the registry is not empty, it holds a single non-GPU
+	// collector (collector_count=1). A predicate that counted collectors rather
+	// than keying on FE_GPU would report healthy here and be a no-op in exactly
+	// the scenario this flag exists for.
+	ctrl := gomock.NewController(t)
+	reg := registry.NewRegistry()
+	tuple := collector.EntityCollectorTuple{}
+	tuple.SetEntity(dcgm.FE_CPU)
+	tuple.SetCollector(mockcollectorpkg.NewMockCollector(ctrl))
+	reg.Register(tuple)
+
+	metricServer := &MetricsServer{config: &appconfig.Config{HealthRequireGPUs: true}}
+	metricServer.registry.Store(reg)
+	recorder := httptest.NewRecorder()
+	metricServer.Health(recorder, nil)
+	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
+	assert.Equal(t, "0", recorder.Header().Get("X-GPU-Collectors"))
+}
+
+func TestHealthReturnsOKWhenGPUCollectorsRequiredAndPresent(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	reg := registry.NewRegistry()
+	tuple := collector.EntityCollectorTuple{}
+	tuple.SetEntity(dcgm.FE_GPU)
+	tuple.SetCollector(mockcollectorpkg.NewMockCollector(ctrl))
+	reg.Register(tuple)
+
+	metricServer := &MetricsServer{config: &appconfig.Config{HealthRequireGPUs: true}}
+	metricServer.registry.Store(reg)
+	recorder := httptest.NewRecorder()
+	metricServer.Health(recorder, nil)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+	assert.Equal(t, "true", recorder.Header().Get("X-Registry-Available"))
+}
+
+func TestHealthReturnsOKDuringReloadEvenWhenGPUCollectorsRequired(t *testing.T) {
+	// Reload takes precedence: the registry is legitimately empty mid-reload and
+	// must not be reported unhealthy, or every reload would restart the pod.
+	metricServer := &MetricsServer{config: &appconfig.Config{HealthRequireGPUs: true}}
+	metricServer.registry.Store(registry.NewRegistry())
+	metricServer.SetReloadInProgress(true)
+	recorder := httptest.NewRecorder()
+	metricServer.Health(recorder, nil)
+	assert.Equal(t, http.StatusOK, recorder.Code)
+}
+
 func TestPprofEndpointsDisabledByDefault(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	mockManager := mockdevicewatchlistmanager.NewMockManager(ctrl)

llms.txt

@@ -91,6 +91,13 @@ CLI and environment behavior comes from `pkg/cmd/app.go`.
 Important contracts:
 
 - The default listen address is `:9400`.
+- `--health-require-gpus` / `DCGM_EXPORTER_HEALTH_REQUIRE_GPUS` should stay
+  opt-in (default off). When on, `/health` returns 503 while no collector is
+  registered under `FE_GPU`, so a liveness probe on `/health` restarts an
+  exporter that came up against a hostengine with no GPUs visible. A reload in
+  progress takes precedence and still returns 200. Note the predicate counts
+  FE_GPU collectors, so a counters file contributing no FE_GPU fields also
+  reads as zero.
 - The default collectors file is `/etc/dcgm-exporter/default-counters.csv`.
 - Remote hostengine values support `<HOST>:<PORT>`,
   `tcp://<HOST>:<PORT>`, `unix:///<SOCKET_PATH>`,

pkg/cmd/app.go

@@ -82,6 +82,7 @@ const (
 	CLIGPUDevices                       = "devices"
 	CLISwitchDevices                    = "switch-devices"
 	CLICPUDevices                       = "cpu-devices"
+	CLIHealthRequireGPUs                = "health-require-gpus"
 	CLINoHostname                       = "no-hostname"
 	CLIUseFakeGPUs                      = "fake-gpus"
 	CLIConfigMapData                    = "configmap-data"
@@ -256,6 +257,16 @@ func NewApp(buildVersion ...string) *cli.App {
 			Usage:   "Omit the hostname information from the output, matching older versions.",
 			EnvVars: []string{"DCGM_EXPORTER_NO_HOSTNAME"},
 		},
+		&cli.BoolFlag{
+			Name:  CLIHealthRequireGPUs,
+			Value: false,
+			Usage: "Report /health as unhealthy when no GPU collector is registered. " +
+				"Off by default so nodes that legitimately expose no GPUs keep passing. " +
+				"Enable where every instance is expected to see at least one GPU, so that " +
+				"a liveness probe on /health restarts an exporter that came up against a " +
+				"hostengine with no GPUs visible.",
+			EnvVars: []string{"DCGM_EXPORTER_HEALTH_REQUIRE_GPUS"},
+		},
 		&cli.StringFlag{
 			Name:    CLISwitchDevices,
 			Aliases: []string{"s"},
@@ -1448,6 +1459,7 @@ func defaultConfig() (*appconfig.Config, error) {
 		GPUDeviceOptions:                 gOpt,
 		SwitchDeviceOptions:              sOpt,
 		CPUDeviceOptions:                 cOpt,
+		HealthRequireGPUs:                false,
 		NoHostname:                       false,
 		UseFakeGPUs:                      false,
 		ConfigMapData:                    undefinedConfigMapData,
@@ -1539,6 +1551,9 @@ func applyExplicitConfigOverrides(c *cli.Context, config *appconfig.Config) erro
 		}
 		config.CPUDeviceOptions = opt
 	}
+	if c.IsSet(CLIHealthRequireGPUs) {
+		config.HealthRequireGPUs = c.Bool(CLIHealthRequireGPUs)
+	}
 	if c.IsSet(CLINoHostname) {
 		config.NoHostname = c.Bool(CLINoHostname)
 	}

pkg/cmd/app_test.go

@@ -476,6 +476,7 @@ func TestNewAppDefaultsMatchDefaultConfig(t *testing.T) {
 		cpuDeviceOptions, err := parseDeviceOptions(c.String(CLICPUDevices))
 		require.NoError(t, err)
 		assert.Equal(t, defaults.CPUDeviceOptions, cpuDeviceOptions)
+		assert.Equal(t, defaults.HealthRequireGPUs, c.Bool(CLIHealthRequireGPUs))
 		assert.Equal(t, defaults.NoHostname, c.Bool(CLINoHostname))
 		assert.Equal(t, defaults.UseFakeGPUs, c.Bool(CLIUseFakeGPUs))
 		assert.Equal(t, defaults.ConfigMapData, c.String(CLIConfigMapData))
@@ -1337,6 +1338,43 @@ func TestContextToConfigHonorsConfigurationFlags(t *testing.T) {
 	assert.Equal(t, appconfig.DeviceOptions{MajorRange: []int{2, 3}}, cfg.GPUDeviceOptions)
 }
 
+func TestContextToConfigHonorsHealthRequireGPUs(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		args []string
+		env  string
+		want bool
+	}{
+		{name: "default off", args: []string{"dcgm-exporter"}, want: false},
+		{name: "flag on", args: []string{"dcgm-exporter", "--health-require-gpus"}, want: true},
+		{name: "flag explicitly off", args: []string{"dcgm-exporter", "--health-require-gpus=false"}, want: false},
+		{name: "env on", args: []string{"dcgm-exporter"}, env: "true", want: true},
+		{name: "env off", args: []string{"dcgm-exporter"}, env: "false", want: false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var cfg *appconfig.Config
+			app := NewApp("test-version")
+			// Neutralise ambient flag env vars first, so an exported
+			// DCGM_EXPORTER_HEALTH_REQUIRE_GPUS cannot make the "off" cases
+			// pass or fail for the wrong reason. t.Setenv below still wins.
+			unsetFlagEnvVars(t, app.Flags)
+			if tc.env != "" {
+				t.Setenv("DCGM_EXPORTER_HEALTH_REQUIRE_GPUS", tc.env)
+			}
+
+			app.Action = func(c *cli.Context) error {
+				var err error
+				cfg, err = contextToConfig(c)
+				return err
+			}
+
+			require.NoError(t, app.Run(tc.args))
+			require.NotNil(t, cfg)
+			assert.Equal(t, tc.want, cfg.HealthRequireGPUs)
+		})
+	}
+}
+
 func TestContextToConfigHonorsConfigurationEnvironment(t *testing.T) {
 	t.Setenv("DCGM_EXPORTER_COLLECTORS", "/tmp/env-counters.csv")
 	t.Setenv("DCGM_EXPORTER_LISTEN", ":19501")