Add unit test for health events

Signed-off-by: Swati Gupta <swatig@nvidia.com>

Signed-off-by: Swati Gupta <swatig@nvidia.com>

Signed-off-by: Swati Gupta <swatig@nvidia.com>

Signed-off-by: Swati Gupta <swatig@nvidia.com>

Author: guptaNswati

cmd/gpu-kubelet-plugin/allocatable.go

@@ -325,31 +325,24 @@ func (d *AllocatableDevice) Taints() []resourceapi.DeviceTaint {
 // Meaning, if a device receives multiple events for the same taint dimension
 // (e.g., XID 48 followed by XID 63), the value is overwritten and only the most recent event data is retained.
 // Returns true if the taint set was modified.
-func (d *AllocatableDevice) AddOrUpdateTaint(taint resourceapi.DeviceTaint) bool {
+func (d *AllocatableDevice) AddOrUpdateTaint(taint *resourceapi.DeviceTaint) bool {
 	for i, existing := range d.taints {
 		if existing.Key == taint.Key {
-			changed := false
 
-			// Overwrite Value if it changed
-			if existing.Value != taint.Value {
-				d.taints[i].Value = taint.Value
-				changed = true
+			// 1. If nothing actually changed, exit early to avoid API calls
+			if existing.Value == taint.Value && existing.Effect == taint.Effect {
+				return false
 			}
 
-			if existing.Effect != taint.Effect {
-				d.taints[i].Effect = taint.Effect
-				changed = true
-			}
-
-			if changed {
-				d.taints[i].TimeAdded = nil // reset timestamp for the API server
-			}
-
-			return changed
+			// 2. Otherwise, update the fields and the timestamp
+			d.taints[i].Value = taint.Value
+			d.taints[i].Effect = taint.Effect
+			d.taints[i].TimeAdded = nil // reset timestamp for the API server
+			return true
 		}
 	}
 
-	// Key doesn't exist yet, append the new taint
-	d.taints = append(d.taints, taint)
+	// 3. Key doesn't exist yet, append the dereferenced struct
+	d.taints = append(d.taints, *taint)
 	return true
 }

cmd/gpu-kubelet-plugin/device_health.go

@@ -72,31 +72,31 @@ type DeviceHealthEvent struct {
 // healthEventToTaint maps a DeviceHealthEvent to the corresponding DRA
 // DeviceTaint using the Option A taint key schema: one key per health
 // dimension under the gpu.nvidia.com domain.
-func healthEventToTaint(event *DeviceHealthEvent, monitor deviceHealthMonitor) resourceapi.DeviceTaint {
+func healthEventToTaint(monitor deviceHealthMonitor, event *DeviceHealthEvent) *resourceapi.DeviceTaint {
 	switch event.EventType {
 	case HealthEventXID:
 		effect := resourceapi.DeviceTaintEffectNoSchedule
 		if monitor != nil && monitor.IsEventNonFatal(event) {
 			effect = DeviceTaintEffectNone
 		}
-		return resourceapi.DeviceTaint{
+		return &resourceapi.DeviceTaint{
 			Key:    TaintKeyXID,
 			Value:  strconv.FormatUint(event.EventData, 10),
 			Effect: effect,
 		}
 	case HealthEventGPULost:
-		return resourceapi.DeviceTaint{
+		return &resourceapi.DeviceTaint{
 			Key:    TaintKeyGPULost,
 			Effect: resourceapi.DeviceTaintEffectNoSchedule,
 		}
 	case HealthEventUnmonitored:
-		return resourceapi.DeviceTaint{
+		return &resourceapi.DeviceTaint{
 			Key:    TaintKeyUnmonitored,
 			Effect: DeviceTaintEffectNone,
 		}
 	default:
 		klog.Errorf("Unknown health event type %q, defaulting to unmonitored taint", event.EventType)
-		return resourceapi.DeviceTaint{
+		return &resourceapi.DeviceTaint{
 			Key:    TaintKeyUnmonitored,
 			Effect: DeviceTaintEffectNone,
 		}
@@ -193,8 +193,8 @@ func (m *nvmlDeviceHealthMonitor) registerEventsForDevices() {
 		ret = gpu.RegisterEvents(eventMask&supportedEvents, m.eventSet)
 		if ret == nvml.ERROR_NOT_SUPPORTED {
 			klog.Warningf("Device %v is too old to support healthchecking.", parentUUID)
-		}
-		if ret != nvml.SUCCESS {
+			m.sendHealthEventForDevices(giMap, HealthEventUnmonitored)
+		} else if ret != nvml.SUCCESS {
 			klog.Warningf("unable to register events for %s: %v; marking devices as unmonitored", parentUUID, ret)
 			m.sendHealthEventForDevices(giMap, HealthEventUnmonitored)
 		}
@@ -288,19 +288,19 @@ func (m *nvmlDeviceHealthMonitor) Unhealthy() <-chan *DeviceHealthEvent {
 func (m *nvmlDeviceHealthMonitor) sendHealthEventForAllDevices(eventType DeviceHealthEventType) {
 	var devices []*AllocatableDevice
 	for _, giMap := range m.deviceByPlacement {
-		devices = append(devices, collectDevices(giMap)...)
+		devices = append(devices, flattenMIGDeviceMap(giMap)...)
 	}
 	m.sendBatchedHealthEvent(devices, eventType)
 }
 
 // sendHealthEventForDevices aggregates all devices under a single parent GPU
 // into one batched DeviceHealthEvent.
 func (m *nvmlDeviceHealthMonitor) sendHealthEventForDevices(giMap map[uint32]map[uint32]*AllocatableDevice, eventType DeviceHealthEventType) {
-	m.sendBatchedHealthEvent(collectDevices(giMap), eventType)
+	m.sendBatchedHealthEvent(flattenMIGDeviceMap(giMap), eventType)
 }
 
-// collectDevices flattens a GI→CI device map into a slice.
-func collectDevices(giMap map[uint32]map[uint32]*AllocatableDevice) []*AllocatableDevice {
+// flattenMIGDeviceMap flattens a GI→CI device map into a slice.
+func flattenMIGDeviceMap(giMap map[uint32]map[uint32]*AllocatableDevice) []*AllocatableDevice {
 	var devices []*AllocatableDevice
 	for _, ciMap := range giMap {
 		for _, dev := range ciMap {