preliminary device health monitor

Signed-off-by: Swati Gupta <[email protected]>

Author: guptaNswati

cmd/gpu-kubelet-plugin/allocatable.go

@@ -22,11 +22,19 @@ import (
 	resourceapi "k8s.io/api/resource/v1"
 )
 
+const (
+	// Healthy means that the device is healthy
+	Healthy = "Healthy"
+	// Unhealthy means that the device is unhealthy
+	Unhealthy = "Unhealthy"
+)
+
 type AllocatableDevices map[string]*AllocatableDevice
 
 type AllocatableDevice struct {
-	Gpu *GpuInfo
-	Mig *MigDeviceInfo
+	Gpu    *GpuInfo
+	Mig    *MigDeviceInfo
+	Health string // from device-plugin
 }
 
 func (d AllocatableDevice) Type() string {
@@ -69,6 +77,10 @@ func (d *AllocatableDevice) GetDevice() resourceapi.Device {
 	panic("unexpected type for AllocatableDevice")
 }
 
+func (d *AllocatableDevice) IsHealthy() bool {
+	return d.Health == Healthy
+}
+
 func (d AllocatableDevices) GpuUUIDs() []string {
 	var uuids []string
 	for _, device := range d {

cmd/gpu-kubelet-plugin/device_health.go

@@ -0,0 +1,196 @@
+/*
+ * Copyright 2025 The Kubernetes Authors.
+ * Copyright (c) 2025, 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 main
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"k8s.io/klog/v2"
+)
+
+type deviceHealthMonitor struct {
+	nvdevlib    *deviceLib
+	allocatable AllocatableDevices
+	eventSet    nvml.EventSet
+	unhealthy   chan *AllocatableDevice
+	stop        chan struct{}
+	wg          sync.WaitGroup
+}
+
+func newDeviceHealthMonitor(ctx context.Context, config *Config) (*deviceHealthMonitor, error) {
+	containerDriverRoot := root(config.flags.containerDriverRoot)
+	nvdevlib, err := newDeviceLib(containerDriverRoot)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create device library: %w", err)
+	}
+
+	if err := nvdevlib.Init(); err != nil {
+		return nil, fmt.Errorf("failed to initialize NVML: %w", err)
+	}
+	//defer nvdevlib.alwaysShutdown()
+
+	allocatable, err := nvdevlib.enumerateAllPossibleDevices(config)
+	if err != nil {
+		return nil, fmt.Errorf("error enumerating all possible devices: %w", err)
+	}
+
+	eventSet, err := nvdevlib.nvmllib.EventSetCreate()
+	if err != nvml.SUCCESS {
+		return nil, fmt.Errorf("failed to create event set: %w", err)
+	}
+
+	monitor := &deviceHealthMonitor{
+		nvdevlib:    nvdevlib,
+		allocatable: allocatable,
+		eventSet:    eventSet,
+		unhealthy:   make(chan *AllocatableDevice, len(allocatable)),
+		stop:        make(chan struct{}),
+	}
+
+	if err := monitor.registerDevicesForEvents(); err != nil {
+		monitor.eventSet.Free()
+		return nil, fmt.Errorf("failed to register devices for health monitoring: %w", err)
+	}
+
+	monitor.start()
+	return monitor, nil
+}
+
+func (m *deviceHealthMonitor) registerDevicesForEvents() error {
+	nvmllib := m.nvdevlib.nvmllib
+	eventMask := uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError)
+
+	for _, uuid := range m.allocatable.UUIDs() {
+		gpu, err := nvmllib.DeviceGetHandleByUUID(uuid)
+		if err != nvml.SUCCESS {
+			klog.Infof("Unable to get NVML handle for UUID %s: %v; skipping health check for this device", uuid, err)
+			continue
+		}
+
+		if err := gpu.RegisterEvents(eventMask, m.eventSet); err != nvml.SUCCESS {
+			klog.Infof("Failed to register events for device %s: %v; skipping health check for this device", uuid, err)
+		}
+	}
+	return nil
+}
+
+func (m *deviceHealthMonitor) start() {
+	m.wg.Add(1)
+	go m.run()
+}
+
+func (m *deviceHealthMonitor) Stop() {
+	if m == nil {
+		return
+	}
+	close(m.stop)
+	m.wg.Wait()
+	close(m.unhealthy)
+	m.eventSet.Free()
+
+	if m.nvdevlib != nil {
+		m.nvdevlib.alwaysShutdown()
+	}
+}
+
+func (m *deviceHealthMonitor) run() {
+	defer m.wg.Done()
+
+	uuidToDeviceMap := make(map[string]*AllocatableDevice)
+	for _, device := range m.allocatable {
+		uuid := device.getUUID()
+		if uuid != "" {
+			uuidToDeviceMap[uuid] = device
+		}
+	}
+
+	klog.Info("Starting event-driven GPU health monitor...")
+
+	for {
+		select {
+		case <-m.stop:
+			klog.Info("Stopping event-driven GPU health monitor...")
+			return
+		default:
+			event, err := m.eventSet.Wait(5000)
+			if err == nvml.ERROR_TIMEOUT {
+				continue
+			}
+			if err != nvml.SUCCESS {
+				klog.Infof("Error waiting for event: %v; Marking all devices as unhealthy", err)
+				for _, dev := range m.allocatable {
+					m.unhealthy <- dev
+				}
+				continue
+			}
+
+			// Process health events
+			switch event.EventType {
+			case nvml.EventTypeXidCriticalError:
+				klog.Warningf("Critical XID error detected on device")
+			case nvml.EventTypeDoubleBitEccError:
+				klog.Warningf("Double-bit ECC error detected on device")
+			case nvml.EventTypeSingleBitEccError:
+				klog.Infof("Single-bit ECC error detected on device")
+			default:
+				continue
+			}
+
+			eventUUID, err := event.Device.GetUUID()
+			if err != nvml.SUCCESS {
+				klog.Infof("Failed to determine uuid for event %v: %v; Marking all devices as unhealthy.", event, err)
+				for _, dev := range m.allocatable {
+					m.unhealthy <- dev
+				}
+				continue
+			}
+
+			device, exists := uuidToDeviceMap[eventUUID]
+			if !exists {
+				continue
+			}
+
+			// Send notification to driver
+			klog.Infof("Sending unhealthy notification for device %s due to event type %v", eventUUID, event.EventType)
+			select {
+			case m.unhealthy <- device:
+				// Successfully sent notification
+			default:
+				// Channel full, log and continue
+				klog.Warningf("Health notification channel full, dropping event for device %s", eventUUID)
+			}
+		}
+	}
+}
+
+func (m *deviceHealthMonitor) Unhealthy() <-chan *AllocatableDevice {
+	return m.unhealthy
+}
+
+func (d *AllocatableDevice) getUUID() string {
+	if d.Gpu != nil {
+		return d.Gpu.UUID
+	}
+	if d.Mig != nil {
+		return d.Mig.UUID
+	}
+	return ""
+}

cmd/gpu-kubelet-plugin/driver.go

@@ -39,11 +39,12 @@ import (
 const DriverPrepUprepFlockFileName = "pu.lock"
 
 type driver struct {
-	client       coreclientset.Interface
-	pluginhelper *kubeletplugin.Helper
-	state        *DeviceState
-	pulock       *flock.Flock
-	healthcheck  *healthcheck
+	client              coreclientset.Interface
+	pluginhelper        *kubeletplugin.Helper
+	state               *DeviceState
+	pulock              *flock.Flock
+	healthcheck         *healthcheck
+	deviceHealthMonitor *deviceHealthMonitor
 }
 
 func NewDriver(ctx context.Context, config *Config) (*driver, error) {
@@ -93,10 +94,18 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	}
 	driver.healthcheck = healthcheck
 
+	deviceHealthMonitor, err := newDeviceHealthMonitor(ctx, config)
+	if err != nil {
+		return nil, fmt.Errorf("start deviceHealthMonitor: %w", err)
+	}
+	driver.deviceHealthMonitor = deviceHealthMonitor
+
 	if err := driver.pluginhelper.PublishResources(ctx, resources); err != nil {
 		return nil, err
 	}
 
+	go driver.handleHealthNotifications(ctx, config.flags.nodeName)
+
 	return driver, nil
 }
 
@@ -109,6 +118,10 @@ func (d *driver) Shutdown() error {
 		d.healthcheck.Stop()
 	}
 
+	if d.deviceHealthMonitor != nil {
+		d.deviceHealthMonitor.Stop()
+	}
+
 	d.pluginhelper.Stop()
 	return nil
 }
@@ -177,6 +190,50 @@ func (d *driver) nodeUnprepareResource(ctx context.Context, claimNs kubeletplugi
 	return nil
 }
 
+func (d *driver) handleHealthNotifications(ctx context.Context, nodeName string) {
+	for {
+		select {
+		case <-ctx.Done():
+			klog.Info("Stopping health notification handler")
+			return
+		case device, ok := <-d.deviceHealthMonitor.Unhealthy():
+			if !ok {
+				klog.Info("Health monitor channel closed")
+				return
+			}
+
+			uuid := device.getUUID()
+			klog.Warningf("Received unhealthy notification for device: %s", uuid)
+
+			// Mark device as unhealthy in state
+			//d.state.MarkUnhealthy(device)
+
+			// Rebuild resource slice with only healthy devices
+			var resourceSlice resourceslice.Slice
+			for _, dev := range d.state.allocatable {
+				if dev.IsHealthy() {
+					resourceSlice.Devices = append(resourceSlice.Devices, dev.GetDevice())
+				} else {
+					klog.Errorf("device:%v with uuid:%s is unhealthy", uuid, dev)
+				}
+			}
+
+			// Republish updated resources
+			resources := resourceslice.DriverResources{
+				Pools: map[string]resourceslice.Pool{
+					nodeName: {Slices: []resourceslice.Slice{resourceSlice}},
+				},
+			}
+
+			if err := d.pluginhelper.PublishResources(ctx, resources); err != nil {
+				klog.Errorf("Failed to publish resources after health change: %v", err)
+			} else {
+				klog.Infof("Successfully republished resources after marking device %s unhealthy", uuid)
+			}
+		}
+	}
+}
+
 // TODO: implement loop to remove CDI files from the CDI path for claimUIDs
 //       that have been removed from the AllocatedClaims map.
 // func (d *driver) cleanupCDIFiles(wg *sync.WaitGroup) chan error {

cmd/gpu-kubelet-plugin/nvlib.go

@@ -129,7 +129,8 @@ func (l deviceLib) enumerateGpusAndMigDevices(config *Config) (AllocatableDevice
 		}
 
 		deviceInfo := &AllocatableDevice{
-			Gpu: gpuInfo,
+			Gpu:    gpuInfo,
+			Health: Healthy,
 		}
 		devices[gpuInfo.CanonicalName()] = deviceInfo
 
@@ -140,7 +141,8 @@ func (l deviceLib) enumerateGpusAndMigDevices(config *Config) (AllocatableDevice
 
 		for _, migDeviceInfo := range migs {
 			deviceInfo := &AllocatableDevice{
-				Mig: migDeviceInfo,
+				Mig:    migDeviceInfo,
+				Health: Healthy,
 			}
 			devices[migDeviceInfo.CanonicalName()] = deviceInfo
 		}