diff --git a/internal/plugin/health_init_test.go b/internal/plugin/health_init_test.go new file mode 100644 index 000000000..cc3a7e71c --- /dev/null +++ b/internal/plugin/health_init_test.go @@ -0,0 +1,210 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plugin + +import ( + "testing" + + "github.com/NVIDIA/go-nvlib/pkg/nvlib/device" + "github.com/NVIDIA/go-nvlib/pkg/nvlib/info" + "github.com/NVIDIA/go-nvml/pkg/nvml" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + v1 "github.com/NVIDIA/k8s-device-plugin/api/config/v1" + "github.com/NVIDIA/k8s-device-plugin/internal/rm" +) + +func TestListAndWatchHealthInitializationFailure(t *testing.T) { + t.Setenv("DP_DISABLE_HEALTHCHECKS", "") + t.Setenv("DP_ENABLE_HEALTHCHECKS", "") + + for _, tc := range []struct { + name string + initResult nvml.Return + failOnInitError bool + expectError bool + }{ + { + name: "NVML initialization failure", + initResult: nvml.ERROR_UNKNOWN, + failOnInitError: true, + expectError: true, + }, + { + name: "NVML initialization failure with fail-on-init-error disabled", + initResult: nvml.ERROR_UNKNOWN, + }, + { + name: "event set creation failure", + initResult: nvml.SUCCESS, + failOnInitError: true, + expectError: true, + }, + } { + t.Run(tc.name, func(t *testing.T) { + config := &v1.Config{ + Flags: v1.Flags{ + CommandLineFlags: v1.CommandLineFlags{ + MigStrategy: new(v1.MigStrategyNone), + FailOnInitError: &tc.failOnInitError, + }, + }, + Resources: v1.Resources{ + GPUs: []v1.Resource{{Pattern: "*", Name: "nvidia.com/gpu"}}, + }, + } + nvmllib := &healthInitializationNVML{initResult: nvml.SUCCESS} + managers, err := rm.NewNVMLResourceManagers(healthInitializationInfo{}, nvmllib, healthInitializationDevices{}, config) + require.NoError(t, err) + require.Len(t, managers, 1) + + // Discovery succeeds, but initializing health monitoring subsequently fails. + nvmllib.initResult = tc.initResult + plugin := &nvidiaDevicePlugin{ + rm: managers[0], + health: make(chan *rm.Device), + stop: make(chan any), + } + stream := &healthInitializationStream{started: make(chan struct{})} + watchDone := make(chan error, 1) + go func() { + watchDone <- plugin.ListAndWatch(&pluginapi.Empty{}, stream) + }() + <-stream.started + + err = plugin.rm.CheckHealth(plugin.stop, plugin.health) + close(plugin.stop) + require.NoError(t, <-watchDone) + if tc.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + } + + // CheckHealth returns only after ListAndWatch receives both notifications. + // ListAndWatch must publish the last update before observing the stop signal. + require.Len(t, stream.snapshots, 3) + require.Equal(t, map[string]string{ + "GPU-0": pluginapi.Healthy, + "GPU-1": pluginapi.Healthy, + }, stream.snapshots[0]) + unhealthy := map[string]string{ + "GPU-0": pluginapi.Unhealthy, + "GPU-1": pluginapi.Unhealthy, + } + require.Equal(t, unhealthy, stream.snapshots[2]) + + // A new stream must retain the updated state even after CheckHealth exits. + // The closed stop channel lets it return immediately after its first report. + reconnected := &healthInitializationStream{started: make(chan struct{})} + require.NoError(t, plugin.ListAndWatch(&pluginapi.Empty{}, reconnected)) + require.Equal(t, []map[string]string{unhealthy}, reconnected.snapshots) + }) + } +} + +type healthInitializationStream struct { + grpc.ServerStream + started chan struct{} + snapshots []map[string]string +} + +func (s *healthInitializationStream) Send(response *pluginapi.ListAndWatchResponse) error { + // Copy health values because subsequent notifications mutate the same devices. + health := make(map[string]string) + for _, d := range response.Devices { + health[d.ID] = d.Health + } + s.snapshots = append(s.snapshots, health) + if len(s.snapshots) == 1 { + close(s.started) + } + return nil +} + +type healthInitializationNVML struct { + nvml.Interface + initResult nvml.Return +} + +func (n *healthInitializationNVML) Init() nvml.Return { + return n.initResult +} + +func (n *healthInitializationNVML) Shutdown() nvml.Return { + return nvml.SUCCESS +} + +func (n *healthInitializationNVML) EventSetCreate() (nvml.EventSet, nvml.Return) { + return nil, nvml.ERROR_UNKNOWN +} + +type healthInitializationInfo struct { + info.Interface +} + +func (healthInitializationInfo) ResolvePlatform() info.Platform { + return info.PlatformNVML +} + +type healthInitializationDevices struct { + device.Interface +} + +func (healthInitializationDevices) VisitDevices(visit func(int, device.Device) error) error { + for i, uuid := range []string{"GPU-0", "GPU-1"} { + if err := visit(i, healthInitializationDevice{uuid: uuid}); err != nil { + return err + } + } + return nil +} + +type healthInitializationDevice struct { + device.Device + uuid string +} + +func (healthInitializationDevice) GetName() (string, nvml.Return) { + return "test GPU", nvml.SUCCESS +} + +func (healthInitializationDevice) IsMigEnabled() (bool, error) { + return false, nil +} + +func (d healthInitializationDevice) GetUUID() (string, nvml.Return) { + return d.uuid, nvml.SUCCESS +} + +func (healthInitializationDevice) GetMinorNumber() (int, nvml.Return) { + return 0, nvml.SUCCESS +} + +func (healthInitializationDevice) GetPciInfo() (nvml.PciInfo, nvml.Return) { + return nvml.PciInfo{}, nvml.SUCCESS +} + +func (healthInitializationDevice) GetMemoryInfo() (nvml.Memory, nvml.Return) { + return nvml.Memory{}, nvml.SUCCESS +} + +func (healthInitializationDevice) GetCudaComputeCapability() (int, int, nvml.Return) { + return 9, 0, nvml.SUCCESS +} diff --git a/internal/plugin/server.go b/internal/plugin/server.go index 9281b48ef..188ca5eae 100644 --- a/internal/plugin/server.go +++ b/internal/plugin/server.go @@ -148,13 +148,13 @@ func (plugin *nvidiaDevicePlugin) Start(kubeletSocket string) error { } klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource()) - go func() { + go func(stop <-chan any, health chan<- *rm.Device) { // TODO: add MPS health check - err := plugin.rm.CheckHealth(plugin.stop, plugin.health) + err := plugin.rm.CheckHealth(stop, health) if err != nil { - klog.Errorf("Failed to start health check: %v; continuing with health checks disabled", err) + klog.Errorf("Failed to start health check: %v", err) } - }() + }(plugin.stop, plugin.health) return nil } diff --git a/internal/rm/health.go b/internal/rm/health.go index dda16015a..006cf6ea1 100644 --- a/internal/rm/health.go +++ b/internal/rm/health.go @@ -49,6 +49,8 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan any, devices Devices, unhe ret := r.nvml.Init() if ret != nvml.SUCCESS { + klog.Errorf("Failed to initialize NVML for health checks: %v; marking all devices as unhealthy", ret) + reportUnhealthyDevices(stop, devices, unhealthy) if *r.config.Flags.FailOnInitError { return fmt.Errorf("failed to initialize NVML: %v", ret) } @@ -65,6 +67,8 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan any, devices Devices, unhe eventSet, ret := r.nvml.EventSetCreate() if ret != nvml.SUCCESS { + klog.Errorf("Failed to create event set for health checks: %v; marking all devices as unhealthy", ret) + reportUnhealthyDevices(stop, devices, unhealthy) return fmt.Errorf("failed to create event set: %v", ret) } defer func() { @@ -171,6 +175,18 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan any, devices Devices, unhe } } +// reportUnhealthyDevices reports devices whose health cannot be monitored. +// Stop must release the sender if ListAndWatch is no longer receiving updates. +func reportUnhealthyDevices(stop <-chan any, devices Devices, unhealthy chan<- *Device) { + for _, d := range devices { + select { + case <-stop: + return + case unhealthy <- d: + } + } +} + const allXIDs = 0 // disabledXIDs stores a map of explicitly disabled XIDs. diff --git a/internal/rm/health_init_test.go b/internal/rm/health_init_test.go new file mode 100644 index 000000000..280dceb59 --- /dev/null +++ b/internal/rm/health_init_test.go @@ -0,0 +1,282 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package rm + +import ( + "testing" + "time" + + "github.com/NVIDIA/go-nvml/pkg/nvml" + "github.com/stretchr/testify/require" + pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" +) + +// These tests verify notifications sent to the plugin. ListAndWatch is +// responsible for applying those notifications to the advertised device health. +func TestCheckHealthInitializationFailure(t *testing.T) { + t.Setenv(envDisableHealthChecks, "") + t.Setenv(envEnableHealthChecks, "") + + testCases := []struct { + name string + initResult nvml.Return + eventSetResult nvml.Return + failOnInitError bool + deviceIDs []string + expectedError string + }{ + { + name: "NVML initialization failure", + initResult: nvml.ERROR_UNKNOWN, + failOnInitError: true, + deviceIDs: []string{"GPU-0", "GPU-1", "GPU-2"}, + expectedError: "failed to initialize NVML", + }, + { + name: "NVML initialization failure with fail on init disabled", + initResult: nvml.ERROR_UNKNOWN, + deviceIDs: []string{"GPU-0", "GPU-1", "GPU-2"}, + }, + { + name: "event set creation failure", + initResult: nvml.SUCCESS, + eventSetResult: nvml.ERROR_UNKNOWN, + failOnInitError: true, + deviceIDs: []string{"GPU-0", "GPU-1", "GPU-2"}, + expectedError: "failed to create event set", + }, + { + name: "event set creation failure with fail on init disabled", + initResult: nvml.SUCCESS, + eventSetResult: nvml.ERROR_UNKNOWN, + deviceIDs: []string{"GPU-0", "GPU-1", "GPU-2"}, + expectedError: "failed to create event set", + }, + { + name: "NVML initialization failure without devices", + initResult: nvml.ERROR_UNKNOWN, + failOnInitError: true, + expectedError: "failed to initialize NVML", + }, + { + name: "event set creation failure without devices", + initResult: nvml.SUCCESS, + eventSetResult: nvml.ERROR_UNKNOWN, + expectedError: "failed to create event set", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + eventSet := &healthInitEventSet{} + lib := &healthInitNvmlLib{ + initResult: tc.initResult, + eventSetResult: tc.eventSetResult, + eventSet: eventSet, + } + r := newHealthInitResourceManager(lib, tc.failOnInitError, tc.deviceIDs...) + unhealthy := make(chan *Device, len(tc.deviceIDs)) + + err := r.CheckHealth(make(chan any), unhealthy) + if tc.expectedError == "" { + require.NoError(t, err) + } else { + require.ErrorContains(t, err, tc.expectedError) + } + + // Drain after CheckHealth returns so a missing notification fails + // an assertion instead of blocking the test on the original code. + var notifiedIDs []string + for len(unhealthy) > 0 { + d := <-unhealthy + require.Same(t, r.devices[d.ID], d) + notifiedIDs = append(notifiedIDs, d.ID) + } + require.ElementsMatch(t, tc.deviceIDs, notifiedIDs) + require.Equal(t, 1, lib.initCalls) + if tc.initResult == nvml.SUCCESS { + require.Equal(t, 1, lib.eventSetCreateCalls) + require.Equal(t, 1, lib.shutdownCalls) + } else { + require.Zero(t, lib.eventSetCreateCalls) + require.Zero(t, lib.shutdownCalls) + } + require.Zero(t, eventSet.freeCalls) + }) + } +} + +func TestCheckHealthInitializationDisabled(t *testing.T) { + for _, disabled := range []string{"all", "xids"} { + t.Run(disabled, func(t *testing.T) { + t.Setenv(envDisableHealthChecks, disabled) + t.Setenv(envEnableHealthChecks, "") + lib := &healthInitNvmlLib{initResult: nvml.ERROR_UNKNOWN} + r := newHealthInitResourceManager(lib, true, "GPU-0", "GPU-1") + unhealthy := make(chan *Device, len(r.devices)) + + require.NoError(t, r.CheckHealth(make(chan any), unhealthy)) + require.Empty(t, unhealthy) + require.Zero(t, lib.initCalls) + require.Zero(t, lib.eventSetCreateCalls) + require.Zero(t, lib.shutdownCalls) + }) + } +} + +func TestCheckHealthInitializationSuccessCleanup(t *testing.T) { + t.Setenv(envDisableHealthChecks, "") + t.Setenv(envEnableHealthChecks, "") + + stop := make(chan any) + waitCalls := 0 + eventSet := &healthInitEventSet{ + waitFunc: func(timeout uint32) (nvml.EventData, nvml.Return) { + waitCalls++ + require.EqualValues(t, 5000, timeout) + close(stop) + return nvml.EventData{}, nvml.ERROR_TIMEOUT + }, + } + lib := &healthInitNvmlLib{ + initResult: nvml.SUCCESS, + eventSetResult: nvml.SUCCESS, + eventSet: eventSet, + } + // An empty inventory isolates successful monitor initialization, an + // ordinary event timeout, and cleanup from per-device event handling. + r := newHealthInitResourceManager(lib, true) + require.NoError(t, r.CheckHealth(stop, make(chan *Device))) + require.Equal(t, 1, lib.initCalls) + require.Equal(t, 1, lib.eventSetCreateCalls) + require.Equal(t, 1, waitCalls) + require.Equal(t, 1, eventSet.freeCalls) + require.Equal(t, 1, lib.shutdownCalls) +} + +func TestCheckHealthInitializationFailureStopsDuringNotification(t *testing.T) { + t.Setenv(envDisableHealthChecks, "") + t.Setenv(envEnableHealthChecks, "") + + for _, initResult := range []nvml.Return{nvml.ERROR_UNKNOWN, nvml.SUCCESS} { + t.Run(initResult.String(), func(t *testing.T) { + lib := &healthInitNvmlLib{ + initResult: initResult, + eventSetResult: nvml.ERROR_UNKNOWN, + } + r := newHealthInitResourceManager(lib, true, "GPU-0", "GPU-1") + stop := make(chan any) + t.Cleanup(func() { + select { + case <-stop: + default: + close(stop) + } + }) + unhealthy := make(chan *Device) + done := make(chan error, 1) + go func() { + done <- r.CheckHealth(stop, unhealthy) + }() + + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + select { + case d := <-unhealthy: + require.Same(t, r.devices[d.ID], d) + case err := <-done: + t.Fatalf("health check returned without notifying an unhealthy device: %v", err) + case <-timer.C: + t.Fatal("health check did not send its first failure notification") + } + + // No receiver remains for the second device. Closing stop must + // release the sender even when ListAndWatch is disconnected. + close(stop) + select { + case err := <-done: + require.Error(t, err) + case <-timer.C: + t.Fatal("health check did not stop while sending failure notifications") + } + if initResult == nvml.SUCCESS { + require.Equal(t, 1, lib.shutdownCalls) + } else { + require.Zero(t, lib.shutdownCalls) + } + }) + } +} + +func newHealthInitResourceManager(lib nvml.Interface, failOnInitError bool, ids ...string) *nvmlResourceManager { + config := &spec.Config{} + config.Flags.FailOnInitError = &failOnInitError + devices := make(Devices) + for _, id := range ids { + devices[id] = &Device{Device: pluginapi.Device{ID: id, Health: pluginapi.Healthy}} + } + return &nvmlResourceManager{ + resourceManager: resourceManager{ + config: config, + resource: "nvidia.com/gpu", + devices: devices, + }, + nvml: lib, + } +} + +type healthInitNvmlLib struct { + nvml.Interface + initResult nvml.Return + eventSetResult nvml.Return + eventSet nvml.EventSet + initCalls int + eventSetCreateCalls int + shutdownCalls int +} + +func (l *healthInitNvmlLib) Init() nvml.Return { + l.initCalls++ + return l.initResult +} + +func (l *healthInitNvmlLib) EventSetCreate() (nvml.EventSet, nvml.Return) { + l.eventSetCreateCalls++ + return l.eventSet, l.eventSetResult +} + +func (l *healthInitNvmlLib) Shutdown() nvml.Return { + l.shutdownCalls++ + return nvml.SUCCESS +} + +type healthInitEventSet struct { + nvml.EventSet + freeCalls int + waitFunc func(uint32) (nvml.EventData, nvml.Return) +} + +func (e *healthInitEventSet) Wait(timeout uint32) (nvml.EventData, nvml.Return) { + return e.waitFunc(timeout) +} + +func (e *healthInitEventSet) Free() nvml.Return { + e.freeCalls++ + return nvml.SUCCESS +}