address review comment of eventMap and health string

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

Author: guptaNswati

cmd/gpu-kubelet-plugin/allocatable.go

@@ -22,11 +22,13 @@ import (
 	resourceapi "k8s.io/api/resource/v1"
 )
 
+type HealthStatus string
+
 const (
 	// Healthy means that the device is healthy.
-	Healthy = "Healthy"
+	Healthy HealthStatus = "Healthy"
 	// Unhealthy means that the device is unhealthy.
-	Unhealthy = "Unhealthy"
+	Unhealthy HealthStatus = "Unhealthy"
 )
 
 type AllocatableDevices map[string]*AllocatableDevice
@@ -35,7 +37,7 @@ type AllocatableDevice struct {
 	Gpu *GpuInfo
 	Mig *MigDeviceInfo
 	// Defined similarly as https://pkg.go.dev/k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1#Healthy
-	Health string
+	Health HealthStatus
 }
 
 func (d AllocatableDevice) Type() string {

cmd/gpu-kubelet-plugin/device_health.go

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ * 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.
@@ -31,23 +31,24 @@ const (
 	FullGPUInstanceID uint32 = 0xFFFFFFFF
 )
 
-type deviceHealthMonitor struct {
-	nvmllib         nvml.Interface
-	eventSet        nvml.EventSet
-	unhealthy       chan *AllocatableDevice
-	cancelContext   context.CancelFunc
-	uuidToDeviceMap map[string]*AllocatableDevice
-	wg              sync.WaitGroup
+type nvmlDeviceHealthMonitor struct {
+	nvmllib                  nvml.Interface
+	eventSet                 nvml.EventSet
+	unhealthy                chan *AllocatableDevice
+	cancelContext            context.CancelFunc
+	uuidToDeviceMap          map[string]*AllocatableDevice
+	getDeviceByParentGiCiMap map[string]map[uint32]map[uint32]*AllocatableDevice
+	wg                       sync.WaitGroup
 }
 
-func newDeviceHealthMonitor(ctx context.Context, config *Config, allocatable AllocatableDevices, nvdevlib *deviceLib) (*deviceHealthMonitor, error) {
+func newNvmlDeviceHealthMonitor(ctx context.Context, config *Config, allocatable AllocatableDevices, nvdevlib *deviceLib) (*nvmlDeviceHealthMonitor, error) {
 	if nvdevlib.nvmllib == nil {
 		return nil, fmt.Errorf("nvml library is nil")
 	}
 
 	ctx, cancel := context.WithCancel(ctx)
 
-	m := &deviceHealthMonitor{
+	m := &nvmlDeviceHealthMonitor{
 		nvmllib:       nvdevlib.nvmllib,
 		unhealthy:     make(chan *AllocatableDevice, len(allocatable)),
 		cancelContext: cancel,
@@ -69,8 +70,10 @@ func newDeviceHealthMonitor(ctx context.Context, config *Config, allocatable All
 
 	m.uuidToDeviceMap = getUUIDToDeviceMap(allocatable)
 
+	m.getDeviceByParentGiCiMap = getDeviceByParentGiCiMap(allocatable)
+
 	klog.V(6).Info("registering NVML events for device health monitor")
-	m.registerDevicesForEvents()
+	m.registerEventsForDevices()
 
 	skippedXids := m.xidsToSkip(config.flags.additionalXidsToIgnore)
 	klog.V(6).Info("started device health monitoring")
@@ -80,7 +83,7 @@ func newDeviceHealthMonitor(ctx context.Context, config *Config, allocatable All
 	return m, nil
 }
 
-func (m *deviceHealthMonitor) registerDevicesForEvents() {
+func (m *nvmlDeviceHealthMonitor) registerEventsForDevices() {
 	eventMask := uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError)
 
 	processedUUIDs := make(map[string]bool)
@@ -122,7 +125,7 @@ func (m *deviceHealthMonitor) registerDevicesForEvents() {
 	}
 }
 
-func (m *deviceHealthMonitor) Stop() {
+func (m *nvmlDeviceHealthMonitor) Stop() {
 	if m == nil {
 		return
 	}
@@ -153,7 +156,7 @@ func getUUIDToDeviceMap(allocatable AllocatableDevices) map[string]*AllocatableD
 	return uuidToDeviceMap
 }
 
-func (m *deviceHealthMonitor) run(ctx context.Context, skippedXids map[uint64]bool) {
+func (m *nvmlDeviceHealthMonitor) run(ctx context.Context, skippedXids map[uint64]bool) {
 	defer m.wg.Done()
 	for {
 		select {
@@ -194,10 +197,12 @@ func (m *deviceHealthMonitor) run(ctx context.Context, skippedXids map[uint64]bo
 			}
 
 			var affectedDevice *AllocatableDevice
-			if event.GpuInstanceId != FullGPUInstanceID && event.ComputeInstanceId != FullGPUInstanceID {
-				affectedDevice = m.findMigDevice(eventUUID, event.GpuInstanceId, event.ComputeInstanceId)
-			} else {
-				affectedDevice = m.findGpuDevice(eventUUID)
+			pMap, ok1 := m.getDeviceByParentGiCiMap[eventUUID]
+			if ok1 {
+				giMap, ok2 := pMap[event.GpuInstanceId]
+				if ok2 {
+					affectedDevice, _ = giMap[event.ComputeInstanceId]
+				}
 			}
 
 			if affectedDevice == nil {
@@ -211,31 +216,46 @@ func (m *deviceHealthMonitor) run(ctx context.Context, skippedXids map[uint64]bo
 	}
 }
 
-func (m *deviceHealthMonitor) Unhealthy() <-chan *AllocatableDevice {
+func (m *nvmlDeviceHealthMonitor) Unhealthy() <-chan *AllocatableDevice {
 	return m.unhealthy
 }
 
-func (m *deviceHealthMonitor) findMigDevice(parentUUID string, giID uint32, ciID uint32) *AllocatableDevice {
-	for _, device := range m.uuidToDeviceMap {
-		if device.Type() != MigDeviceType {
+func getDeviceByParentGiCiMap(allocatable AllocatableDevices) map[string]map[uint32]map[uint32]*AllocatableDevice {
+	deviceByParentGiCiMap := make(map[string]map[uint32]map[uint32]*AllocatableDevice)
+
+	for _, d := range allocatable {
+		var parentUUID string
+		var giID, ciID uint32
+
+		switch d.Type() {
+		case GpuDeviceType:
+			parentUUID = d.UUID()
+			if parentUUID == "" {
+				continue
+			}
+			giID = FullGPUInstanceID
+			ciID = FullGPUInstanceID
+		case MigDeviceType:
+			parentUUID = d.Mig.parent.UUID
+			if parentUUID == "" {
+				continue
+			}
+			giID = d.Mig.giInfo.Id
+			ciID = d.Mig.ciInfo.Id
+		default:
+			klog.Errorf("Skipping device with unknown type: %s", d.UUID())
 			continue
 		}
 
-		if device.Mig.parent.UUID == parentUUID &&
-			device.Mig.giInfo.Id == giID &&
-			device.Mig.ciInfo.Id == ciID {
-			return device
+		if _, ok := deviceByParentGiCiMap[parentUUID]; !ok {
+			deviceByParentGiCiMap[parentUUID] = make(map[uint32]map[uint32]*AllocatableDevice)
 		}
+		if _, ok := deviceByParentGiCiMap[parentUUID][giID]; !ok {
+			deviceByParentGiCiMap[parentUUID][giID] = make(map[uint32]*AllocatableDevice)
+		}
+		deviceByParentGiCiMap[parentUUID][giID][ciID] = d
 	}
-	return nil
-}
-
-func (m *deviceHealthMonitor) findGpuDevice(uuid string) *AllocatableDevice {
-	device, exists := m.uuidToDeviceMap[uuid]
-	if exists && device.Type() == GpuDeviceType {
-		return device
-	}
-	return nil
+	return deviceByParentGiCiMap
 }
 
 // getAdditionalXids returns a list of additional Xids to skip from the specified string.
@@ -264,7 +284,8 @@ func getAdditionalXids(input string) []uint64 {
 	return additionalXids
 }
 
-func (m *deviceHealthMonitor) xidsToSkip(additionalXids string) map[uint64]bool {
+// Refer https://docs.nvidia.com/deploy/xid-errors/analyzing-xid-catalog.html for information on xids.
+func (m *nvmlDeviceHealthMonitor) xidsToSkip(additionalXids string) map[uint64]bool {
 	ignoredXids := []uint64{
 		13,  // Graphics Engine Exception
 		31,  // GPU memory page fault

cmd/gpu-kubelet-plugin/device_state.go

@@ -550,12 +550,12 @@ func GetOpaqueDeviceConfigs(
 	return resultConfigs, nil
 }
 
-func (s *DeviceState) UpdateDeviceHealthStatus(device *AllocatableDevice, healthstatus string) {
+func (s *DeviceState) UpdateDeviceHealthStatus(device *AllocatableDevice, hs HealthStatus) {
 	s.Lock()
 	defer s.Unlock()
 
-	device.Health = healthstatus
-	klog.Infof("Update device sattus:%s healthstatus", device.UUID())
+	device.Health = hs
+	klog.Infof("Updated device: %s health status to %s", device.UUID(), hs)
 }
 
 // TODO: Dynamic MIG is not yet supported with structured parameters.

cmd/gpu-kubelet-plugin/driver.go

@@ -41,13 +41,13 @@ import (
 const DriverPrepUprepFlockFileName = "pu.lock"
 
 type driver struct {
-	client              coreclientset.Interface
-	pluginhelper        *kubeletplugin.Helper
-	state               *DeviceState
-	pulock              *flock.Flock
-	healthcheck         *healthcheck
-	deviceHealthMonitor *deviceHealthMonitor
-	wg                  sync.WaitGroup
+	client                  coreclientset.Interface
+	pluginhelper            *kubeletplugin.Helper
+	state                   *DeviceState
+	pulock                  *flock.Flock
+	healthcheck             *healthcheck
+	nvmlDeviceHealthMonitor *nvmlDeviceHealthMonitor
+	wg                      sync.WaitGroup
 }
 
 func NewDriver(ctx context.Context, config *Config) (*driver, error) {
@@ -98,12 +98,12 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	driver.healthcheck = healthcheck
 
 	if featuregates.Enabled(featuregates.DeviceHealthCheck) {
-		deviceHealthMonitor, err := newDeviceHealthMonitor(ctx, config, state.allocatable, state.nvdevlib)
+		nvmlDeviceHealthMonitor, err := newNvmlDeviceHealthMonitor(ctx, config, state.allocatable, state.nvdevlib)
 		if err != nil {
-			return nil, fmt.Errorf("start deviceHealthMonitor: %w", err)
+			return nil, fmt.Errorf("start nvmlDeviceHealthMonitor: %w", err)
 		}
 
-		driver.deviceHealthMonitor = deviceHealthMonitor
+		driver.nvmlDeviceHealthMonitor = nvmlDeviceHealthMonitor
 
 		driver.wg.Add(1)
 		go func() {
@@ -128,8 +128,8 @@ func (d *driver) Shutdown() error {
 		d.healthcheck.Stop()
 	}
 
-	if d.deviceHealthMonitor != nil {
-		d.deviceHealthMonitor.Stop()
+	if d.nvmlDeviceHealthMonitor != nil {
+		d.nvmlDeviceHealthMonitor.Stop()
 	}
 
 	d.wg.Wait()
@@ -203,13 +203,13 @@ func (d *driver) nodeUnprepareResource(ctx context.Context, claimNs kubeletplugi
 }
 
 func (d *driver) deviceHealthEvents(ctx context.Context, nodeName string) {
-	klog.Info("Processing device health notifications")
+	klog.Info("Starting to watch for device health notifications")
 	for {
 		select {
 		case <-ctx.Done():
 			klog.V(6).Info("Stop processing device health notifications")
 			return
-		case device, ok := <-d.deviceHealthMonitor.Unhealthy():
+		case device, ok := <-d.nvmlDeviceHealthMonitor.Unhealthy():
 			if !ok {
 				klog.V(6).Info("Health monitor channel closed")
 				return
@@ -219,7 +219,7 @@ func (d *driver) deviceHealthEvents(ctx context.Context, nodeName string) {
 			klog.Warningf("Received unhealthy notification for device: %s", uuid)
 
 			if !device.IsHealthy() {
-				klog.V(6).Infof("Device: %s is aleady marked unhealthy. Skip republishing resourceslice", uuid)
+				klog.V(6).Infof("Device: %s is aleady marked unhealthy. Skip republishing ResourceSlice", uuid)
 				continue
 			}
 
@@ -234,15 +234,15 @@ func (d *driver) deviceHealthEvents(ctx context.Context, nodeName string) {
 
 			// Republish resource slice with only healthy devices
 			// There is no remediation loop right now meaning if the unhealthy device is fixed,
-			// driver needs to be restarted to publish the resourceslice with all devices
+			// driver needs to be restarted to publish the ResourceSlice with all devices
 			var resourceSlice resourceslice.Slice
 			for _, dev := range d.state.allocatable {
 				uuid := dev.UUID()
 				if dev.IsHealthy() {
-					klog.V(6).Infof("Device: %s is healthy, added to resoureslice", uuid)
+					klog.V(6).Infof("Device: %s is healthy, added to ResoureSlice", uuid)
 					resourceSlice.Devices = append(resourceSlice.Devices, dev.GetDevice())
 				} else {
-					klog.Warningf("Device: %s is unhealthy, will be removed from resoureslice", uuid)
+					klog.Warningf("Device: %s is unhealthy, will be removed from ResoureSlice", uuid)
 				}
 			}