simplify code

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

Author: guptaNswati

cmd/gpu-kubelet-plugin/device_health.go

@@ -20,130 +20,138 @@ package main
 import (
 	"context"
 	"fmt"
-	"sync"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 	"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
 }
 
 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")
+	}
+
+	m := &deviceHealthMonitor{
+		nvmllib:   nvdevlib.nvmllib,
+		unhealthy: make(chan *AllocatableDevice, len(allocatable)),
+		stop:      make(chan struct{}),
 	}
-	//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)
-	//}
+	if r := m.nvmllib.Init(); r != nvml.SUCCESS {
+		return nil, fmt.Errorf("failed to initialize NVML: %v", r)
+	}
 
-	klog.Info("[SWATI DEBUG] creating NVML events")
-	eventSet, err := nvdevlib.nvmllib.EventSetCreate()
+	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")
+	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)
+	for uuid, dev := range m.uuidToDeviceMap {
+		// if its a mig device, get its Parent UUID
+		if dev.Type() == MigDeviceType {
+			uuid = dev.Mig.parent.UUID
+		}
+		gpu, err := m.nvmllib.DeviceGetHandleByUUID(uuid)
 		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 get device handle from UUID[%s]: %v; marking it as unhealthy", uuid, 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)
+		supportedEvents, err := gpu.GetSupportedEventTypes()
+		if err != nvml.SUCCESS {
+			klog.Infof("unable to determine the supported events for %s: %v; marking it as unhealthy", uuid, err)
+			m.unhealthy <- dev
+			continue
 		}
-	}
-	return nil
-}
 
-func (m *deviceHealthMonitor) start() {
-	klog.Info("[SWATI DEBUG] starting health monitor")
-	m.wg.Add(1)
-	go m.run()
+		err = gpu.RegisterEvents(eventMask&supportedEvents, m.eventSet)
+		if err == nvml.ERROR_NOT_SUPPORTED {
+			klog.Warningf("Device %v is too old to support healthchecking.", uuid)
+		}
+		if err != nvml.SUCCESS {
+			klog.Infof("unable to register events for %s: %v; marking it as unhealthy", uuid, err)
+			m.unhealthy <- dev
+		}
+
+	}
 }
 
 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 {
+	for _, device := range allocatable {
 		uuid := device.GetUUID()
 		if uuid != "" {
 			uuidToDeviceMap[uuid] = device
 		}
 	}
+	return uuidToDeviceMap
+}
 
-	klog.Info("Starting event-driven GPU health monitor...")
-
+func (m *deviceHealthMonitor) run() {
 	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,16 +166,25 @@ 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]
+			device, exists := m.uuidToDeviceMap[eventUUID]
 			if !exists {
+				klog.Infof("Ignoring event for unexpected device: %v", eventUUID)
 				continue
 			}
+			if device.Type() == MigDeviceType && event.GpuInstanceId != InstanceID && event.ComputeInstanceId != InstanceID {
+				giID := device.Mig.giInfo.Id
+				ciID := device.Mig.ciInfo.Id
+				if giID != event.GpuInstanceId || ciID != event.ComputeInstanceId {
+					continue
+				}
+				klog.Infof("Event for mig device %v (giID=%v, ciID=%v)", device.Mig.UUID, giID, ciID)
+			}
 
 			// Send notification to driver
 			klog.Infof("Sending unhealthy notification for device %s due to event type %v", eventUUID, event.EventType)

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{