handle mig devices

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

Author: guptaNswati

cmd/gpu-kubelet-plugin/device_health.go

@@ -1,6 +1,5 @@
 /*
- * Copyright 2025 The Kubernetes Authors.
- * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022, 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.
@@ -26,124 +25,146 @@ import (
 	"k8s.io/klog/v2"
 )
 
+const (
+	InstanceID uint32 = 0xFFFFFFFF
+)
+
 type deviceHealthMonitor struct {
-	nvdevlib    *deviceLib
-	allocatable AllocatableDevices
-	eventSet    nvml.EventSet
-	unhealthy   chan *AllocatableDevice
-	stop        chan struct{}
-	wg          sync.WaitGroup
+	nvmllib         nvml.Interface
+	eventSet        nvml.EventSet
+	unhealthy       chan *AllocatableDevice
+	stop            chan struct{}
+	uuidToDeviceMap map[string]*AllocatableDevice
+	wg              sync.WaitGroup
 }
 
 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)
+	if nvdevlib.nvmllib == nil {
+		return nil, fmt.Errorf("nvml library is nil")
 	}
-	//defer nvdevlib.alwaysShutdown()
 
-	//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)
-	//}
+	m := &deviceHealthMonitor{
+		nvmllib:   nvdevlib.nvmllib,
+		unhealthy: make(chan *AllocatableDevice, len(allocatable)),
+		stop:      make(chan struct{}),
+	}
 
-	klog.Info("[SWATI DEBUG] creating NVML events")
-	eventSet, err := nvdevlib.nvmllib.EventSetCreate()
+	if r := m.nvmllib.Init(); r != nvml.SUCCESS {
+		return nil, fmt.Errorf("failed to initialize NVML: %v", r)
+	}
+
+	klog.V(6).Info("creating NVML events for device health monitor")
+	eventSet, err := m.nvmllib.EventSetCreate()
 	if err != nvml.SUCCESS {
+		_ = m.nvmllib.Shutdown()
 		return nil, fmt.Errorf("failed to create event set: %w", err)
 	}
+	m.eventSet = eventSet
 
-	monitor := &deviceHealthMonitor{
-		nvdevlib:    nvdevlib,
-		allocatable: allocatable,
-		eventSet:    eventSet,
-		unhealthy:   make(chan *AllocatableDevice, len(allocatable)),
-		stop:        make(chan struct{}),
-	}
+	m.uuidToDeviceMap = getUUIDToDeviceMap(allocatable)
 
-	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)
-	}
+	klog.V(6).Info("registering NVML events for device health monitor")
+	m.registerDevicesForEvents()
 
-	monitor.start()
-	return monitor, nil
+	klog.V(6).Info("started device health monitoring")
+	m.wg.Add(1)
+	go m.run()
+
+	return m, nil
 }
 
-func (m *deviceHealthMonitor) registerDevicesForEvents() error {
-	nvmllib := m.nvdevlib.nvmllib
+func (m *deviceHealthMonitor) registerDevicesForEvents() {
+	// TODO: add a list of xids to ignore
 	eventMask := uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError)
 
-	for _, uuid := range m.allocatable.UUIDs() {
-		gpu, err := nvmllib.DeviceGetHandleByUUID(uuid)
+	processedUUIDs := make(map[string]bool)
+
+	for uuid, dev := range m.uuidToDeviceMap {
+		var u string
+		if dev.Type() == MigDeviceType {
+			u = dev.Mig.parent.UUID
+		} else {
+			u = uuid
+		}
+
+		if processedUUIDs[u] {
+			continue
+		}
+		gpu, err := m.nvmllib.DeviceGetHandleByUUID(u)
+		if err != nvml.SUCCESS {
+			klog.Infof("Unable to get device handle from UUID[%s]: %v; marking it as unhealthy", u, err)
+			m.unhealthy <- dev
+			continue
+		}
+
+		supportedEvents, err := gpu.GetSupportedEventTypes()
 		if err != nvml.SUCCESS {
-			klog.Infof("Unable to get NVML handle for UUID %s: %v; skipping health check for this device", uuid, err)
+			klog.Infof("unable to determine the supported events for %s: %v; marking it as unhealthy", u, err)
+			m.unhealthy <- dev
 			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)
+		err = gpu.RegisterEvents(eventMask&supportedEvents, m.eventSet)
+		if err == nvml.ERROR_NOT_SUPPORTED {
+			klog.Warningf("Device %v is too old to support healthchecking.", u)
+		}
+		if err != nvml.SUCCESS {
+			klog.Infof("unable to register events for %s: %v; marking it as unhealthy", u, err)
+			m.unhealthy <- dev
 		}
+		processedUUIDs[u] = true
 	}
-	return nil
-}
-
-func (m *deviceHealthMonitor) start() {
-	klog.Info("[SWATI DEBUG] starting health monitor")
-	m.wg.Add(1)
-	go m.run()
 }
 
 func (m *deviceHealthMonitor) Stop() {
 	if m == nil {
 		return
 	}
-	klog.Info("[SWATI DEBUG] stopping health monitor")
+	klog.V(6).Info("stopping health monitor")
+
 	close(m.stop)
 	m.wg.Wait()
-	close(m.unhealthy)
+
 	m.eventSet.Free()
 
-	if m.nvdevlib != nil {
-		m.nvdevlib.alwaysShutdown()
+	if r := m.nvmllib.Shutdown(); r != nvml.SUCCESS {
+		klog.Warningf("failed to shutdown NVML: %v", r)
 	}
+	close(m.unhealthy)
 }
 
-func (m *deviceHealthMonitor) run() {
-	defer m.wg.Done()
-
+func getUUIDToDeviceMap(allocatable AllocatableDevices) map[string]*AllocatableDevice {
 	uuidToDeviceMap := make(map[string]*AllocatableDevice)
-	for _, device := range m.allocatable {
-		uuid := device.GetUUID()
-		if uuid != "" {
-			uuidToDeviceMap[uuid] = device
+
+	for _, d := range allocatable {
+		if u := d.GetUUID(); u != "" {
+			uuidToDeviceMap[u] = d
 		}
 	}
+	return uuidToDeviceMap
+}
 
-	klog.Info("Starting event-driven GPU health monitor...")
-
+func (m *deviceHealthMonitor) run() {
+	defer m.wg.Done()
 	for {
 		select {
 		case <-m.stop:
-			klog.Info("Stopping event-driven GPU health monitor...")
+			klog.V(6).Info("Stopping event-driven GPU health monitor...")
 			return
 		default:
 			event, err := m.eventSet.Wait(5000)
 			if err == nvml.ERROR_TIMEOUT {
-				klog.Info("[SWATI DEBUG] timedout")
 				continue
 			}
 			if err != nvml.SUCCESS {
 				klog.Infof("Error waiting for event: %v; Marking all devices as unhealthy", err)
-				for _, dev := range m.allocatable {
+				for _, dev := range m.uuidToDeviceMap {
 					m.unhealthy <- dev
 				}
 				continue
 			}
 
-			// Process health events
+			klog.Infof("Processing event %+v", event)
 			switch event.EventType {
 			case nvml.EventTypeXidCriticalError:
 				klog.Warningf("Critical XID error detected on device: %+v", event)
@@ -158,30 +179,53 @@ func (m *deviceHealthMonitor) run() {
 			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 {
+				for _, dev := range m.uuidToDeviceMap {
 					m.unhealthy <- dev
 				}
 				continue
 			}
 
-			device, exists := uuidToDeviceMap[eventUUID]
-			if !exists {
-				continue
+			var affectedDevice *AllocatableDevice
+			if event.GpuInstanceId != InstanceID && event.ComputeInstanceId != InstanceID {
+				affectedDevice = m.findMigDevice(eventUUID, event.GpuInstanceId, event.ComputeInstanceId)
+				klog.Infof("Event for mig device: %v", affectedDevice)
+			} else {
+				affectedDevice = m.findGpuDevice(eventUUID)
 			}
 
-			// 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)
+			if affectedDevice == nil {
+				klog.Infof("Ignoring event for unexpected device (UUID: %s, GI: %d, CI: %d)", eventUUID, event.GpuInstanceId, event.ComputeInstanceId)
+				continue
 			}
+			klog.Infof("Sending unhealthy notification for device %s due to event type %v", eventUUID, event.EventType)
+			m.unhealthy <- affectedDevice
 		}
 	}
 }
 
 func (m *deviceHealthMonitor) 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 {
+			continue
+		}
+
+		if device.Mig.parent.UUID == parentUUID &&
+			device.Mig.giInfo.Id == giID &&
+			device.Mig.ciInfo.Id == ciID {
+			return device
+		}
+	}
+	return nil
+}
+
+func (m *deviceHealthMonitor) findGpuDevice(uuid string) *AllocatableDevice {
+	device, exists := m.uuidToDeviceMap[uuid]
+	if exists && device.Type() == GpuDeviceType {
+		return device
+	}
+	return nil
+}

cmd/gpu-kubelet-plugin/driver.go

@@ -105,7 +105,7 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
 		return nil, err
 	}
 
-	go driver.handleHealthNotifications(ctx, config.flags.nodeName)
+	go driver.deviceHealthEvents(ctx, config.flags.nodeName)
 
 	return driver, nil
 }
@@ -191,12 +191,12 @@ func (d *driver) nodeUnprepareResource(ctx context.Context, claimNs kubeletplugi
 	return nil
 }
 
-func (d *driver) handleHealthNotifications(ctx context.Context, nodeName string) {
-	klog.Info("[SWATI DEBUG] handling Health Notifications")
+func (d *driver) deviceHealthEvents(ctx context.Context, nodeName string) {
+	klog.Info("Processing device health notifications")
 	for {
 		select {
 		case <-ctx.Done():
-			klog.Info("Stopping health notification handler")
+			klog.Info("Stop processing device health notifications")
 			return
 		case device, ok := <-d.deviceHealthMonitor.Unhealthy():
 			if !ok {
@@ -222,6 +222,8 @@ func (d *driver) handleHealthNotifications(ctx context.Context, nodeName string)
 			}
 
 			// Republish updated resources
+			// TODO: 1. remove this.
+			// 2. Add device taints
 			klog.Info("[SWATI DEBUG] rebulishing resourceslice with healthy devices")
 			resources := resourceslice.DriverResources{
 				Pools: map[string]resourceslice.Pool{