add taints on health events

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

Author: guptaNswati

cmd/gpu-kubelet-plugin/allocatable.go

@@ -41,6 +41,11 @@ type AllocatableDevice struct {
 	MigDynamic *MigSpec
 	MigStatic  *MigDeviceInfo
 	Vfio       *VfioDeviceInfo
+
+	// taints holds DRA device taints set by the health monitor. Published
+	// as part of the device's ResourceSlice entry so the scheduler and
+	// kubelet honour them per KEP-5055.
+	taints []resourceapi.DeviceTaint
 }
 
 type AllocatableDeviceList []*AllocatableDevice
@@ -272,18 +277,44 @@ func (d AllocatableDevices) RemoveSiblingDevices(device *AllocatableDevice) {
 	}
 }
 
-func (d *AllocatableDevice) IsHealthy() bool {
-	switch d.Type() {
-	case GpuDeviceType:
-		return d.Gpu.health == Healthy
-	case MigStaticDeviceType:
-		// TODO: review -- what about the parent?
-		return d.MigStatic.health == Healthy
-	case MigDynamicDeviceType:
-		// TODOMIG: For now, pretend health -- this device maybe hasn't
-		// manifested yet. Or has it? We could adopt the health status of the
-		// parent, but that's also not meaningful I think.
-		return true
+func (d *AllocatableDevice) Taints() []resourceapi.DeviceTaint {
+	return d.taints
+}
+
+// AddOrUpdateTaint adds a new taint or updates an existing one with the same
+// key. Within the same key, the effect can only escalate
+// (None < NoSchedule < NoExecute), and the value is always updated to the
+// latest. Returns true if the taint set was modified.
+func (d *AllocatableDevice) AddOrUpdateTaint(taint resourceapi.DeviceTaint) bool {
+	for i, existing := range d.taints {
+		if existing.Key != taint.Key {
+			continue
+		}
+		changed := false
+		if taintEffectSeverity(taint.Effect) > taintEffectSeverity(existing.Effect) {
+			d.taints[i].Effect = taint.Effect
+			d.taints[i].TimeAdded = nil // reset so the API server sets a fresh timestamp
+			changed = true
+		}
+		if d.taints[i].Value != taint.Value {
+			d.taints[i].Value = taint.Value
+			changed = true
+		}
+		return changed
+	}
+	d.taints = append(d.taints, taint)
+	return true
+}
+
+func taintEffectSeverity(effect resourceapi.DeviceTaintEffect) int {
+	switch effect {
+	case DeviceTaintEffectNone:
+		return 0
+	case resourceapi.DeviceTaintEffectNoSchedule:
+		return 1
+	case resourceapi.DeviceTaintEffectNoExecute:
+		return 2
+	default:
+		return -1
 	}
-	panic("unexpected type for AllocatableDevice")
 }

cmd/gpu-kubelet-plugin/device_health.go

@@ -24,21 +24,85 @@ import (
 	"sync"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	resourceapi "k8s.io/api/resource/v1"
 	"k8s.io/klog/v2"
 )
 
 const (
 	FullGPUInstanceID uint32 = 0xFFFFFFFF
 )
 
+const (
+	TaintKeyXID         = DriverName + "/xid"
+	TaintKeyGPULost     = DriverName + "/gpu-lost"
+	TaintKeyUnmonitored = DriverName + "/unmonitored"
+)
+
+// TODO: remove later, as its not in vendored api
+// DeviceTaintEffectNone is an informational effect that does not affect
+// scheduling or eviction. Available in k8s 1.36+ (KEP-5055 beta).
+const DeviceTaintEffectNone resourceapi.DeviceTaintEffect = "None"
+
+// DeviceHealthEventType classifies the category of health event detected by
+// the NVML health monitor.
+type DeviceHealthEventType string
+
+const (
+	HealthEventXID         DeviceHealthEventType = "xid"
+	HealthEventGPULost     DeviceHealthEventType = "gpu-lost"
+	HealthEventUnmonitored DeviceHealthEventType = "unmonitored"
+)
+
+// DeviceHealthEvent carries a typed health notification from the NVML health
+// monitor to the driver's event handler, enabling the driver to set the
+// appropriate DRA device taint per the Option A schema (KEP-5055).
+type DeviceHealthEvent struct {
+	Device    *AllocatableDevice
+	EventType DeviceHealthEventType
+	// inspired by NVML Event type and only meaningful for xid errors.
+	// may have to create a custom type based on future device-api
+	EventData uint64
+}
+
+// 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.
+// TODO: revisit this to double check on the desired effects
+func healthEventToTaint(event *DeviceHealthEvent) resourceapi.DeviceTaint {
+	switch event.EventType {
+	case HealthEventXID:
+		return resourceapi.DeviceTaint{
+			Key:    TaintKeyXID,
+			Value:  strconv.FormatUint(event.EventData, 10),
+			Effect: DeviceTaintEffectNone,
+		}
+	case HealthEventGPULost:
+		return resourceapi.DeviceTaint{
+			Key:    TaintKeyGPULost,
+			Effect: resourceapi.DeviceTaintEffectNoExecute,
+		}
+	case HealthEventUnmonitored:
+		return resourceapi.DeviceTaint{
+			Key:    TaintKeyUnmonitored,
+			Effect: DeviceTaintEffectNone,
+		}
+	default:
+		klog.Errorf("Unknown health event type %q, defaulting to unmonitored taint", event.EventType)
+		return resourceapi.DeviceTaint{
+			Key:    TaintKeyUnmonitored,
+			Effect: DeviceTaintEffectNone,
+		}
+	}
+}
+
 // For a MIG device the placement is defined by the 3-tuple <parent UUID, GI, CI>.
 // For a full device the returned 3-tuple is the device's uuid and (FullGPUInstanceID) 0xFFFFFFFF for the other two elements.
 type devicePlacementMap map[string]map[uint32]map[uint32]*AllocatableDevice
 
 type nvmlDeviceHealthMonitor struct {
 	nvmllib           nvml.Interface
 	eventSet          nvml.EventSet
-	unhealthy         chan *AllocatableDevice
+	unhealthy         chan *DeviceHealthEvent
 	deviceByPlacement devicePlacementMap
 	skippedXids       map[uint64]bool
 	wg                sync.WaitGroup
@@ -57,7 +121,7 @@ func newNvmlDeviceHealthMonitor(config *Config, allocatable AllocatableDevices,
 
 	m := &nvmlDeviceHealthMonitor{
 		nvmllib:           nvdevlib.nvmllib,
-		unhealthy:         make(chan *AllocatableDevice, len(allocatable)),
+		unhealthy:         make(chan *DeviceHealthEvent, len(allocatable)),
 		deviceByPlacement: getDevicePlacementMap(allocatable),
 		skippedXids:       xidsToSkip(config.flags.additionalXidsToIgnore),
 	}
@@ -102,15 +166,15 @@ func (m *nvmlDeviceHealthMonitor) registerEventsForDevices() {
 	for parentUUID, giMap := range m.deviceByPlacement {
 		gpu, ret := m.nvmllib.DeviceGetHandleByUUID(parentUUID)
 		if ret != nvml.SUCCESS {
-			klog.Warningf("Unable to get device handle from UUID[%s]: %v; marking it as unhealthy", parentUUID, ret)
-			m.markAllMigDevicesUnhealthy(giMap)
+			klog.Warningf("Unable to get device handle from UUID[%s]: %v; marking devices as unmonitored", parentUUID, ret)
+			m.sendHealthEventsForDevices(giMap, HealthEventUnmonitored)
 			continue
 		}
 
 		supportedEvents, ret := gpu.GetSupportedEventTypes()
 		if ret != nvml.SUCCESS {
-			klog.Warningf("unable to determine the supported events for %s: %v; marking it as unhealthy", parentUUID, ret)
-			m.markAllMigDevicesUnhealthy(giMap)
+			klog.Warningf("unable to determine the supported events for %s: %v; marking devices as unmonitored", parentUUID, ret)
+			m.sendHealthEventsForDevices(giMap, HealthEventUnmonitored)
 			continue
 		}
 
@@ -119,8 +183,8 @@ func (m *nvmlDeviceHealthMonitor) registerEventsForDevices() {
 			klog.Warningf("Device %v is too old to support healthchecking.", parentUUID)
 		}
 		if ret != nvml.SUCCESS {
-			klog.Warningf("unable to register events for %s: %v; marking it as unhealthy", parentUUID, ret)
-			m.markAllMigDevicesUnhealthy(giMap)
+			klog.Warningf("unable to register events for %s: %v; marking devices as unmonitored", parentUUID, ret)
+			m.sendHealthEventsForDevices(giMap, HealthEventUnmonitored)
 		}
 	}
 }
@@ -158,8 +222,8 @@ func (m *nvmlDeviceHealthMonitor) run(ctx context.Context) {
 			// Ref doc: [https://docs.nvidia.com/deploy/nvml-api/group__nvmlEvents.html#group__nvmlEvents_1g9714b0ca9a34c7a7780f87fee16b205c].
 			if ret != nvml.SUCCESS {
 				if ret == nvml.ERROR_GPU_IS_LOST {
-					klog.Warningf("GPU is lost error: %v; Marking all devices as unhealthy", ret)
-					m.markAllDevicesUnhealthy()
+					klog.Warningf("GPU is lost error: %v; Tainting all devices with %s", ret, TaintKeyGPULost)
+					m.sendHealthEventsForAllDevices(HealthEventGPULost)
 					continue
 				}
 				klog.V(6).Infof("Error waiting for NVML event: %v. Retrying...", ret)
@@ -187,8 +251,8 @@ func (m *nvmlDeviceHealthMonitor) run(ctx context.Context) {
 			// TODO: look into how to properly handle this error.
 			eventUUID, ret := event.Device.GetUUID()
 			if ret != nvml.SUCCESS {
-				klog.Warningf("Failed to determine uuid for event %v: %v; Marking all devices as unhealthy.", event, ret)
-				m.markAllDevicesUnhealthy()
+				klog.Warningf("Failed to determine uuid for event %v: %v; Tainting all devices with %s", event, ret, TaintKeyGPULost)
+				m.sendHealthEventsForAllDevices(HealthEventGPULost)
 				continue
 			}
 			affectedDevice := m.deviceByPlacement.get(eventUUID, gi, ci)
@@ -197,38 +261,41 @@ func (m *nvmlDeviceHealthMonitor) run(ctx context.Context) {
 				continue
 			}
 
-			klog.V(4).Infof("Sending unhealthy notification for device %s due to event type:%v and event data:%d", affectedDevice.UUID(), eType, xid)
-			m.unhealthy <- affectedDevice
+			klog.V(4).Infof("Sending XID=%d health event for device %s", xid, affectedDevice.UUID())
+			m.unhealthy <- &DeviceHealthEvent{
+				Device:    affectedDevice,
+				EventType: HealthEventXID,
+				EventData: xid,
+			}
 		}
 	}
 }
 
-func (m *nvmlDeviceHealthMonitor) Unhealthy() <-chan *AllocatableDevice {
+func (m *nvmlDeviceHealthMonitor) Unhealthy() <-chan *DeviceHealthEvent {
 	return m.unhealthy
 }
 
-func (m *nvmlDeviceHealthMonitor) markAllDevicesUnhealthy() {
+func (m *nvmlDeviceHealthMonitor) sendHealthEventsForAllDevices(eventType DeviceHealthEventType) {
 	for _, giMap := range m.deviceByPlacement {
-		m.markAllMigDevicesUnhealthy(giMap)
+		m.sendHealthEventsForDevices(giMap, eventType)
 	}
 }
 
-// markAllMigDevicesUnhealthy is a helper function to mark every mig device under a parent as unhealthy.
-func (m *nvmlDeviceHealthMonitor) markAllMigDevicesUnhealthy(giMap map[uint32]map[uint32]*AllocatableDevice) {
+// sendHealthEventsForDevices sends a DeviceHealthEvent for every device under
+// the given parent-level map. Uses non-blocking sends to protect the monitor
+// goroutine from deadlocks when the channel is full.
+func (m *nvmlDeviceHealthMonitor) sendHealthEventsForDevices(giMap map[uint32]map[uint32]*AllocatableDevice, eventType DeviceHealthEventType) {
 	for _, ciMap := range giMap {
 		for _, dev := range ciMap {
-			// Non-blocking send to avoid deadlocks if channel is full.
+			event := &DeviceHealthEvent{
+				Device:    dev,
+				EventType: eventType,
+			}
 			select {
-			case m.unhealthy <- dev:
-				klog.V(6).Infof("Marked device %s as unhealthy", dev.UUID())
-			// TODO: The non-blocking send protects the health-monitor goroutine from deadlocks,
-			// but dropping an unhealthy notification means the device's health transition may
-			// never reach the consumer. Consider follow-up improvements:
-			//   - increase the channel buffer beyond len(allocatable) to reduce backpressure;
-			//   - introduce a special "all devices unhealthy" message when bulk updates occur;
-			//   - or revisit whether blocking briefly here is acceptable.
+			case m.unhealthy <- event:
+				klog.V(6).Infof("Sent %s health event for device %s", eventType, dev.UUID())
 			default:
-				klog.Errorf("Unhealthy channel full. Dropping unhealthy notification for device %s", dev.UUID())
+				klog.Errorf("Health event channel full; dropping %s event for device %s", eventType, dev.UUID())
 			}
 		}
 	}
@@ -326,6 +393,7 @@ func getAdditionalXids(input string) []uint64 {
 	return additionalXids
 }
 
+// TODO: this can be simplified now with "EffectNone"
 func xidsToSkip(additionalXids string) map[uint64]bool {
 	// Add the list of hardcoded disabled (ignored) XIDs:
 	// http://docs.nvidia.com/deploy/xid-errors/index.html#topic_4

cmd/gpu-kubelet-plugin/device_state.go

@@ -638,12 +638,6 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 		if !exists {
 			return nil, fmt.Errorf("requested device is not allocatable: %v", result.Device)
 		}
-		// only proceed with config mapping if device is healthy.
-		if featuregates.Enabled(featuregates.NVMLDeviceHealthCheck) {
-			if !device.IsHealthy() {
-				return nil, fmt.Errorf("requested device is not healthy: %v", result.Device)
-			}
-		}
 		for _, c := range slices.Backward(configs) {
 			if slices.Contains(c.Requests, result.Request) {
 				if _, ok := c.Config.(*configapi.GpuConfig); ok && device.Type() != GpuDeviceType {
@@ -1071,31 +1065,6 @@ func GetOpaqueDeviceConfigs(
 	return resultConfigs, nil
 }
 
-// TODO: this needs a code comment.
-func (s *DeviceState) UpdateDeviceHealthStatus(d *AllocatableDevice, hs HealthStatus) {
-	s.Lock()
-	defer s.Unlock()
-
-	switch d.Type() {
-	case GpuDeviceType:
-		d.Gpu.health = hs
-	case MigDynamicDeviceType:
-		// Here we do not have access to a concrete MIG device. The
-		// 'allocatable' MIG device is an abstract representation to a specific
-		// MIG device.
-		klog.Warningf("UpdateDeviceHealthStatus() called for abstract, dynamic MIG device %s: %s", d.MigDynamic.CanonicalName(), hs)
-	case MigStaticDeviceType:
-		// Does it make sense to update the health for a MIG device? Do we
-		// receive health events that are specific to individual MIG devices? If
-		// yes: does this allow for making conclusions about the parent device?
-		d.MigStatic.health = hs
-	default:
-		klog.V(6).Infof("Cannot update health status for unknown device type: %s", d.Type())
-		return
-	}
-	klog.V(4).Infof("Updated device: %s health status to %s", d.UUID(), hs)
-}
-
 // requestedNonAdminDevices returns the set of device names requested by the claim,
 // excluding admin-access allocations.
 func (s *DeviceState) requestedNonAdminDevices(claim *resourceapi.ResourceClaim) map[string]struct{} {
@@ -1182,3 +1151,11 @@ func (s *DeviceState) deleteMigDevIfExistsAndNotUsedByCompletedClaim(ms *MigSpec
 
 	return nil
 }
+
+// AddDeviceTaint adds or updates a DRA device taint on the given device under
+// the DeviceState lock. Returns true if the taint set was actually modified.
+func (s *DeviceState) AddDeviceTaint(d *AllocatableDevice, taint resourceapi.DeviceTaint) bool {
+	s.Lock()
+	defer s.Unlock()
+	return d.AddOrUpdateTaint(taint)
+}

cmd/gpu-kubelet-plugin/deviceinfo.go

@@ -27,15 +27,6 @@ import (
 	"k8s.io/utils/ptr"
 )
 
-// Defined similarly as https://pkg.go.dev/k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1#Healthy.
-type HealthStatus string
-
-const (
-	Healthy HealthStatus = "Healthy"
-	// With NVMLDeviceHealthCheck, Unhealthy means that there are critcal xid errors on the device.
-	Unhealthy HealthStatus = "Unhealthy"
-)
-
 // Represents a specific, full, physical GPU device.
 type GpuInfo struct {
 	UUID                  string `json:"uuid"`
@@ -53,7 +44,6 @@ type GpuInfo struct {
 	pcieRootAttr          *deviceattribute.DeviceAttribute
 	migProfiles           []*MigProfileInfo
 	addressingMode        *string
-	health                HealthStatus
 
 	// The following properties that can only be known after inspecting MIG
 	// profiles.
@@ -89,7 +79,6 @@ type MigDeviceInfo struct {
 	ciProfileInfo *nvml.ComputeInstanceProfileInfo
 	pcieBusID     string
 	pcieRootAttr  *deviceattribute.DeviceAttribute
-	health        HealthStatus
 }
 
 type VfioDeviceInfo struct {