publish health status

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

Author: guptaNswati

cmd/gpu-kubelet-plugin/allocatable.go

@@ -81,6 +81,16 @@ func (d *AllocatableDevice) IsHealthy() bool {
 	return d.Health == Healthy
 }
 
+func (d *AllocatableDevice) GetUUID() string {
+	if d.Gpu != nil {
+		return d.Gpu.UUID
+	}
+	if d.Mig != nil {
+		return d.Mig.UUID
+	}
+	return ""
+}
+
 func (d AllocatableDevices) GpuUUIDs() []string {
 	var uuids []string
 	for _, device := range d {

cmd/gpu-kubelet-plugin/device_health.go

@@ -35,23 +35,20 @@ type deviceHealthMonitor 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)
-	}
-
+func newDeviceHealthMonitor(ctx context.Context, config *Config, allocatable AllocatableDevices, nvdevlib *deviceLib) (*deviceHealthMonitor, error) {
+	klog.Info("[SWATI DEBUG] initializing NVML..")
 	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)
-	}
+	//klog.Info("[SWATI DEBUG] getting all devices..")
+	//allocatable, err := nvdevlib.enumerateAllPossibleDevices(config)
+	//if err != nil {
+	//	return nil, fmt.Errorf("error enumerating all possible devices: %w", err)
+	//}
 
+	klog.Info("[SWATI DEBUG] creating NVML events")
 	eventSet, err := nvdevlib.nvmllib.EventSetCreate()
 	if err != nvml.SUCCESS {
 		return nil, fmt.Errorf("failed to create event set: %w", err)
@@ -65,6 +62,7 @@ func newDeviceHealthMonitor(ctx context.Context, config *Config) (*deviceHealthM
 		stop:        make(chan struct{}),
 	}
 
+	klog.Info("[SWATI DEBUG] registering NVML events")
 	if err := monitor.registerDevicesForEvents(); err != nil {
 		monitor.eventSet.Free()
 		return nil, fmt.Errorf("failed to register devices for health monitoring: %w", err)
@@ -93,6 +91,7 @@ func (m *deviceHealthMonitor) registerDevicesForEvents() error {
 }
 
 func (m *deviceHealthMonitor) start() {
+	klog.Info("[SWATI DEBUG] starting health monitor")
 	m.wg.Add(1)
 	go m.run()
 }
@@ -101,6 +100,7 @@ func (m *deviceHealthMonitor) Stop() {
 	if m == nil {
 		return
 	}
+	klog.Info("[SWATI DEBUG] stopping health monitor")
 	close(m.stop)
 	m.wg.Wait()
 	close(m.unhealthy)
@@ -116,7 +116,7 @@ func (m *deviceHealthMonitor) run() {
 
 	uuidToDeviceMap := make(map[string]*AllocatableDevice)
 	for _, device := range m.allocatable {
-		uuid := device.getUUID()
+		uuid := device.GetUUID()
 		if uuid != "" {
 			uuidToDeviceMap[uuid] = device
 		}
@@ -132,6 +132,7 @@ func (m *deviceHealthMonitor) run() {
 		default:
 			event, err := m.eventSet.Wait(5000)
 			if err == nvml.ERROR_TIMEOUT {
+				klog.Info("[SWATI DEBUG] timedout")
 				continue
 			}
 			if err != nvml.SUCCESS {
@@ -145,11 +146,11 @@ func (m *deviceHealthMonitor) run() {
 			// Process health events
 			switch event.EventType {
 			case nvml.EventTypeXidCriticalError:
-				klog.Warningf("Critical XID error detected on device")
+				klog.Warningf("Critical XID error detected on device: %+v", event)
 			case nvml.EventTypeDoubleBitEccError:
-				klog.Warningf("Double-bit ECC error detected on device")
+				klog.Warningf("Double-bit ECC error detected on device: %+v", event)
 			case nvml.EventTypeSingleBitEccError:
-				klog.Infof("Single-bit ECC error detected on device")
+				klog.Infof("Single-bit ECC error detected on device:%+v", event)
 			default:
 				continue
 			}
@@ -184,13 +185,3 @@ func (m *deviceHealthMonitor) run() {
 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/device_state.go

@@ -23,7 +23,10 @@ import (
 	"sync"
 
 	resourceapi "k8s.io/api/resource/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	metav1apply "k8s.io/client-go/applyconfigurations/meta/v1"
+	resourceapply "k8s.io/client-go/applyconfigurations/resource/v1"
 	"k8s.io/dynamic-resource-allocation/kubeletplugin"
 	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/kubelet/checkpointmanager"
@@ -285,8 +288,11 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 		Config:   configapi.DefaultMigDeviceConfig(),
 	})
 
+	// Swati: Add resourceclaim status update
 	// Look through the configs and figure out which one will be applied to
 	// each device allocation result based on their order of precedence and type.
+	resourceClaimStatus := resourceapply.ResourceClaimStatus()
+	var deviceStatuses []*resourceapply.AllocatedDeviceStatusApplyConfiguration
 	configResultsMap := make(map[runtime.Object][]*resourceapi.DeviceRequestAllocationResult)
 	for _, result := range claim.Status.Allocation.Devices.Results {
 		if result.Driver != DriverName {
@@ -296,6 +302,41 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 		if !exists {
 			return nil, fmt.Errorf("requested device is not allocatable: %v", result.Device)
 		}
+		// Swati add health check
+		klog.Info("[SWATI DEBUG] adding device status")
+		deviceStatus := resourceapply.AllocatedDeviceStatus().
+			WithDevice(result.Device).
+			WithDriver(result.Driver).
+			WithPool(result.Pool)
+
+		if device.Health == Unhealthy {
+			deviceStatus = deviceStatus.WithConditions(
+				metav1apply.Condition().
+					WithType("Ready").
+					WithStatus(metav1.ConditionFalse).
+					WithReason("Unhealthy").
+					WithMessage(fmt.Sprintf("Device %s is not healthy", result.Device)).
+					WithLastTransitionTime(metav1.Now()),
+			)
+			klog.Warningf("Device %s is unhealthy, marking as not ready", result.Device)
+		} else {
+			deviceStatus = deviceStatus.WithConditions(
+				metav1apply.Condition().
+					WithType("Ready").
+					WithStatus(metav1.ConditionTrue).
+					WithReason("Healthy").
+					WithMessage("Device is healthy and ready").
+					WithLastTransitionTime(metav1.Now()),
+			)
+			klog.Infof("Device %s is healthy, marking as ready", result.Device)
+		}
+		deviceStatuses = append(deviceStatuses, deviceStatus)
+
+		// Only proceed with config mapping for healthy devices
+		if device.Health == Unhealthy {
+			continue
+		}
+
 		for _, c := range slices.Backward(configs) {
 			if slices.Contains(c.Requests, result.Request) {
 				if _, ok := c.Config.(*configapi.GpuConfig); ok && device.Type() != GpuDeviceType {
@@ -319,7 +360,25 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 			}
 		}
 	}
+	resourceClaimStatus = resourceClaimStatus.WithDevices(deviceStatuses...)
+
+	// Update the resource claim status
+	resourceClaimApply := resourceapply.ResourceClaim(claim.Name, claim.Namespace).WithStatus(resourceClaimStatus)
+	_, err = s.config.clientsets.Resource.ResourceClaims(claim.Namespace).ApplyStatus(ctx,
+		resourceClaimApply,
+		metav1.ApplyOptions{FieldManager: DriverName, Force: true},
+	)
+
+	if err != nil {
+		klog.Infof("failed to update status for claim %s/%s : %v", claim.Namespace, claim.Name, err)
+	} else {
+		klog.Infof("update status for claim %s/%s", claim.Namespace, claim.Name)
+	}
 
+	// If no healthy devices are available for configuration, return
+	if len(configResultsMap) == 0 {
+		return nil, fmt.Errorf("no healthy devices available for allocation")
+	}
 	// Normalize, validate, and apply all configs associated with devices that
 	// need to be prepared. Track device group configs generated from applying the
 	// config to the set of device allocation results.
@@ -550,6 +609,21 @@ func GetOpaqueDeviceConfigs(
 	return resultConfigs, nil
 }
 
+func (s *DeviceState) MarkDeviceUnhealthy(unhealthyDevice *AllocatableDevice) {
+	s.Lock()
+	defer s.Unlock()
+
+	uuid := unhealthyDevice.GetUUID()
+	device, ok := s.allocatable[uuid]
+	if !ok {
+		klog.Warningf("Attempted to mark unknown device as unhealthy: %s", uuid)
+		return
+	}
+
+	device.Health = Unhealthy
+	klog.Infof("Marked device:%s unhealthy", uuid)
+}
+
 // TODO: Dynamic MIG is not yet supported with structured parameters.
 // Refactor this to allow for the allocation of statically partitioned MIG
 // devices.

cmd/gpu-kubelet-plugin/driver.go

@@ -94,10 +94,11 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 	}
 	driver.healthcheck = healthcheck
 
-	deviceHealthMonitor, err := newDeviceHealthMonitor(ctx, config)
+	deviceHealthMonitor, err := newDeviceHealthMonitor(ctx, config, state.allocatable, state.nvdevlib)
 	if err != nil {
 		return nil, fmt.Errorf("start deviceHealthMonitor: %w", err)
 	}
+	klog.Info("[SWATI DEBUGS] Started device health monitor")
 	driver.deviceHealthMonitor = deviceHealthMonitor
 
 	if err := driver.pluginhelper.PublishResources(ctx, resources); err != nil {
@@ -191,6 +192,7 @@ func (d *driver) nodeUnprepareResource(ctx context.Context, claimNs kubeletplugi
 }
 
 func (d *driver) handleHealthNotifications(ctx context.Context, nodeName string) {
+	klog.Info("[SWATI DEBUG] handling Health Notifications")
 	for {
 		select {
 		case <-ctx.Done():
@@ -202,23 +204,25 @@ func (d *driver) handleHealthNotifications(ctx context.Context, nodeName string)
 				return
 			}
 
-			uuid := device.getUUID()
+			uuid := device.GetUUID()
 			klog.Warningf("Received unhealthy notification for device: %s", uuid)
 
 			// Mark device as unhealthy in state
-			//d.state.MarkUnhealthy(device)
+			d.state.MarkDeviceUnhealthy(device)
 
 			// Rebuild resource slice with only healthy devices
 			var resourceSlice resourceslice.Slice
 			for _, dev := range d.state.allocatable {
 				if dev.IsHealthy() {
+					klog.Infof("[SWATI DEBUG] device is healthy, added to resoureslice: %v", dev)
 					resourceSlice.Devices = append(resourceSlice.Devices, dev.GetDevice())
 				} else {
 					klog.Errorf("device:%v with uuid:%s is unhealthy", uuid, dev)
 				}
 			}
 
 			// Republish updated resources
+			klog.Info("[SWATI DEBUG] rebulishing resourceslice with healthy devices")
 			resources := resourceslice.DriverResources{
 				Pools: map[string]resourceslice.Pool{
 					nodeName: {Slices: []resourceslice.Slice{resourceSlice}},