[no-relnote] refactor use p as receiver for nvmlHealthProvider

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>

Author: ArangoGutierrez

internal/rm/health.go

@@ -114,69 +114,69 @@ func NewNVMLHealthProvider(
 
 // Start initializes NVML, registers event handlers, and starts the
 // monitoring goroutine. Blocks until initialization completes.
-func (r *nvmlHealthProvider) Start(ctx context.Context) error {
-	r.Lock()
-	if r.started {
-		r.Unlock()
+func (p *nvmlHealthProvider) Start(ctx context.Context) error {
+	p.Lock()
+	if p.started {
+		p.Unlock()
 		return fmt.Errorf("health provider already started")
 	}
-	r.started = true
-	r.Unlock()
+	p.started = true
+	p.Unlock()
 
 	// Get XID filter configuration
-	r.xidsDisabled = getDisabledHealthCheckXids()
+	p.xidsDisabled = getDisabledHealthCheckXids()
 
 	// Initialize NVML
-	ret := r.nvml.Init()
+	ret := p.nvml.Init()
 	if ret != nvml.SUCCESS {
-		if *r.config.Flags.FailOnInitError {
+		if *p.config.Flags.FailOnInitError {
 			return fmt.Errorf("failed to initialize NVML: %v", ret)
 		}
 		klog.Warningf("NVML init failed: %v; health checks disabled", ret)
 		return nil
 	}
 
 	// Create event set
-	eventSet, ret := r.nvml.EventSetCreate()
+	eventSet, ret := p.nvml.EventSetCreate()
 	if ret != nvml.SUCCESS {
-		if shutdownRet := r.nvml.Shutdown(); shutdownRet != nvml.SUCCESS {
+		if shutdownRet := p.nvml.Shutdown(); shutdownRet != nvml.SUCCESS {
 			klog.Warningf("Failed to shutdown NVML: %v", shutdownRet)
 		}
 		return fmt.Errorf("failed to create event set: %v", ret)
 	}
-	r.eventSet = eventSet
+	p.eventSet = eventSet
 
 	// Register devices
-	if err := r.registerDevices(); err != nil {
-		r.cleanup()
+	if err := p.registerDevices(); err != nil {
+		p.cleanup()
 		return fmt.Errorf("failed to register devices: %w", err)
 	}
 
-	klog.Infof("Health monitoring started for %d devices", len(r.devices))
+	klog.Infof("Health monitoring started for %d devices", len(p.devices))
 	klog.Infof("Ignoring the following XIDs for health checks: %v",
-		r.xidsDisabled)
+		p.xidsDisabled)
 
 	// Create child context
-	r.ctx, r.cancel = context.WithCancel(ctx)
+	p.ctx, p.cancel = context.WithCancel(ctx)
 
 	// Start monitoring goroutine
-	r.wg.Add(1)
-	go r.runEventMonitor()
+	p.wg.Add(1)
+	go p.runEventMonitor()
 
 	return nil
 }
 
 // Stop gracefully shuts down health monitoring and waits for the
 // monitoring goroutine to complete.
-func (r *nvmlHealthProvider) Stop() {
-	r.Lock()
-	if r.stopped {
-		r.Unlock()
+func (p *nvmlHealthProvider) Stop() {
+	p.Lock()
+	if p.stopped {
+		p.Unlock()
 		return
 	}
-	r.stopped = true
-	alreadyStarted := r.started
-	r.Unlock()
+	p.stopped = true
+	alreadyStarted := p.started
+	p.Unlock()
 
 	if !alreadyStarted {
 		return
@@ -185,61 +185,61 @@ func (r *nvmlHealthProvider) Stop() {
 	klog.V(2).Info("Stopping health provider...")
 
 	// Signal goroutine to stop
-	if r.cancel != nil {
-		r.cancel()
+	if p.cancel != nil {
+		p.cancel()
 	}
 
 	// Wait for goroutine to finish
-	r.wg.Wait()
+	p.wg.Wait()
 
 	// Cleanup NVML resources
-	r.cleanup()
+	p.cleanup()
 
 	// Close channel
-	close(r.healthChan)
+	close(p.healthChan)
 
 	klog.Info("Health provider stopped")
 }
 
 // Health returns a read-only channel that receives devices that have
 // become unhealthy.
-func (r *nvmlHealthProvider) Health() <-chan *Device {
-	return r.healthChan
+func (p *nvmlHealthProvider) Health() <-chan *Device {
+	return p.healthChan
 }
 
 // cleanup releases NVML resources.
-func (r *nvmlHealthProvider) cleanup() {
-	if r.eventSet != nil {
-		ret := r.eventSet.Free()
+func (p *nvmlHealthProvider) cleanup() {
+	if p.eventSet != nil {
+		ret := p.eventSet.Free()
 		if ret != nvml.SUCCESS {
 			klog.Warningf("Failed to free event set: %v", ret)
 		}
-		r.eventSet = nil
+		p.eventSet = nil
 	}
 
-	if ret := r.nvml.Shutdown(); ret != nvml.SUCCESS {
+	if ret := p.nvml.Shutdown(); ret != nvml.SUCCESS {
 		klog.Warningf("NVML shutdown failed: %v", ret)
 	}
 }
 
 // runEventMonitor monitors NVML events and reports unhealthy devices.
 // This is the existing checkHealth logic refactored into a goroutine.
-func (r *nvmlHealthProvider) runEventMonitor() {
-	defer r.wg.Done()
+func (p *nvmlHealthProvider) runEventMonitor() {
+	defer p.wg.Done()
 
 	klog.V(2).Info("Health check: event monitor started")
 	defer klog.V(2).Info("Health check: event monitor stopped")
 
 	for {
 		// Check for context cancellation
 		select {
-		case <-r.ctx.Done():
+		case <-p.ctx.Done():
 			return
 		default:
 		}
 
 		// Wait for NVML event (5 second timeout)
-		event, ret := r.eventSet.Wait(5000)
+		event, ret := p.eventSet.Wait(5000)
 
 		if ret == nvml.ERROR_TIMEOUT {
 			continue
@@ -248,8 +248,8 @@ func (r *nvmlHealthProvider) runEventMonitor() {
 		if ret != nvml.SUCCESS {
 			klog.Infof("Error waiting for event: %v; marking all "+
 				"devices as unhealthy", ret)
-			for _, device := range r.devices {
-				r.sendUnhealthy(device)
+			for _, device := range p.devices {
+				p.sendUnhealthy(device)
 			}
 			continue
 		}
@@ -262,7 +262,7 @@ func (r *nvmlHealthProvider) runEventMonitor() {
 		}
 
 		// Check if XID is disabled
-		if r.xidsDisabled.IsDisabled(event.EventData) {
+		if p.xidsDisabled.IsDisabled(event.EventData) {
 			klog.Infof("Skipping event %+v", event)
 			continue
 		}
@@ -274,13 +274,13 @@ func (r *nvmlHealthProvider) runEventMonitor() {
 		if ret != nvml.SUCCESS {
 			klog.Infof("Failed to determine uuid for event %v: %v; "+
 				"marking all devices as unhealthy.", event, ret)
-			for _, device := range r.devices {
-				r.sendUnhealthy(device)
+			for _, device := range p.devices {
+				p.sendUnhealthy(device)
 			}
 			continue
 		}
 
-		device, exists := r.parentToDeviceMap[eventUUID]
+		device, exists := p.parentToDeviceMap[eventUUID]
 		if !exists {
 			klog.Infof("Ignoring event for unexpected device: %v",
 				eventUUID)
@@ -291,8 +291,8 @@ func (r *nvmlHealthProvider) runEventMonitor() {
 		if device.IsMigDevice() &&
 			event.GpuInstanceId != 0xFFFFFFFF &&
 			event.ComputeInstanceId != 0xFFFFFFFF {
-			gi := r.deviceIDToGiMap[device.ID]
-			ci := r.deviceIDToCiMap[device.ID]
+			gi := p.deviceIDToGiMap[device.ID]
+			ci := p.deviceIDToCiMap[device.ID]
 
 			if gi != event.GpuInstanceId || ci != event.ComputeInstanceId {
 				continue
@@ -307,14 +307,71 @@ func (r *nvmlHealthProvider) runEventMonitor() {
 			"device as unhealthy.", event.EventData, device.ID)
 
 		device.Health = pluginapi.Unhealthy
-		r.sendUnhealthy(device)
+		p.sendUnhealthy(device)
 	}
 }
 
+<<<<<<< HEAD
+=======
+// getDevicePlacement returns the placement of the specified device.
+// 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 0xFFFFFFFF for the other two elements.
+func (p *nvmlHealthProvider) getDevicePlacement(
+	d *Device,
+) (string, uint32, uint32, error) {
+	if !d.IsMigDevice() {
+		return d.GetUUID(), 0xFFFFFFFF, 0xFFFFFFFF, nil
+	}
+	return p.getMigDeviceParts(d)
+}
+
+// getMigDeviceParts returns the parent GI and CI ids of the MIG device.
+func (p *nvmlHealthProvider) getMigDeviceParts(
+	d *Device,
+) (string, uint32, uint32, error) {
+	if !d.IsMigDevice() {
+		return "", 0, 0, fmt.Errorf("cannot get GI and CI of full device")
+	}
+
+	uuid := d.GetUUID()
+	// For older driver versions, the call to DeviceGetHandleByUUID will
+	// fail for MIG devices.
+	mig, ret := p.nvml.DeviceGetHandleByUUID(uuid)
+	if ret == nvml.SUCCESS {
+		parentHandle, ret := mig.GetDeviceHandleFromMigDeviceHandle()
+		if ret != nvml.SUCCESS {
+			return "", 0, 0, fmt.Errorf("failed to get parent "+
+				"device handle: %v", ret)
+		}
+
+		parentUUID, ret := parentHandle.GetUUID()
+		if ret != nvml.SUCCESS {
+			return "", 0, 0, fmt.Errorf("failed to get parent "+
+				"uuid: %v", ret)
+		}
+		gi, ret := mig.GetGpuInstanceId()
+		if ret != nvml.SUCCESS {
+			return "", 0, 0, fmt.Errorf("failed to get GPU "+
+				"Instance ID: %v", ret)
+		}
+
+		ci, ret := mig.GetComputeInstanceId()
+		if ret != nvml.SUCCESS {
+			return "", 0, 0, fmt.Errorf("failed to get Compute "+
+				"Instance ID: %v", ret)
+		}
+		//nolint:gosec  // We know that the values returned from Get*InstanceId are within the valid uint32 range.
+		return parentUUID, uint32(gi), uint32(ci), nil
+	}
+	return parseMigDeviceUUID(uuid)
+}
+
+>>>>>>> 3f8110390 ([no-relnote] refactor use p as receiver for nvmlHealthProvider)
 // sendUnhealthy sends device to unhealthy channel (non-blocking).
-func (r *nvmlHealthProvider) sendUnhealthy(device *Device) {
+func (p *nvmlHealthProvider) sendUnhealthy(device *Device) {
 	select {
-	case r.healthChan <- device:
+	case p.healthChan <- device:
 		// Sent successfully
 	default:
 		// Channel full
@@ -326,35 +383,35 @@ func (r *nvmlHealthProvider) sendUnhealthy(device *Device) {
 
 // registerDevices registers all devices with the NVML event set.
 // This is the existing logic from checkHealth().
-func (r *nvmlHealthProvider) registerDevices() error {
-	r.parentToDeviceMap = make(map[string]*Device)
-	r.deviceIDToGiMap = make(map[string]uint32)
-	r.deviceIDToCiMap = make(map[string]uint32)
+func (p *nvmlHealthProvider) registerDevices() error {
+	p.parentToDeviceMap = make(map[string]*Device)
+	p.deviceIDToGiMap = make(map[string]uint32)
+	p.deviceIDToCiMap = make(map[string]uint32)
 
 	eventMask := uint64(nvml.EventTypeXidCriticalError |
 		nvml.EventTypeDoubleBitEccError |
 		nvml.EventTypeSingleBitEccError)
 
-	for _, device := range r.devices {
-		uuid, gi, ci, err := r.getDevicePlacement(device)
+	for _, device := range p.devices {
+		uuid, gi, ci, err := p.getDevicePlacement(device)
 		if err != nil {
 			klog.Warningf("Could not determine device placement for "+
 				"%v: %v; marking it unhealthy.", device.ID, err)
 			device.Health = pluginapi.Unhealthy
-			r.sendUnhealthy(device)
+			p.sendUnhealthy(device)
 			continue
 		}
 
-		r.deviceIDToGiMap[device.ID] = gi
-		r.deviceIDToCiMap[device.ID] = ci
-		r.parentToDeviceMap[uuid] = device
+		p.deviceIDToGiMap[device.ID] = gi
+		p.deviceIDToCiMap[device.ID] = ci
+		p.parentToDeviceMap[uuid] = device
 
-		gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid)
+		gpu, ret := p.nvml.DeviceGetHandleByUUID(uuid)
 		if ret != nvml.SUCCESS {
 			klog.Infof("unable to get device handle from UUID: %v; "+
 				"marking it as unhealthy", ret)
 			device.Health = pluginapi.Unhealthy
-			r.sendUnhealthy(device)
+			p.sendUnhealthy(device)
 			continue
 		}
 
@@ -363,11 +420,11 @@ func (r *nvmlHealthProvider) registerDevices() error {
 			klog.Infof("unable to determine the supported events for "+
 				"%v: %v; marking it as unhealthy", device.ID, ret)
 			device.Health = pluginapi.Unhealthy
-			r.sendUnhealthy(device)
+			p.sendUnhealthy(device)
 			continue
 		}
 
-		ret = gpu.RegisterEvents(eventMask&supportedEvents, r.eventSet)
+		ret = gpu.RegisterEvents(eventMask&supportedEvents, p.eventSet)
 		if ret == nvml.ERROR_NOT_SUPPORTED {
 			klog.Warningf("Device %v is too old to support "+
 				"healthchecking.", device.ID)
@@ -376,7 +433,7 @@ func (r *nvmlHealthProvider) registerDevices() error {
 			klog.Infof("Marking device %v as unhealthy: %v",
 				device.ID, ret)
 			device.Health = pluginapi.Unhealthy
-			r.sendUnhealthy(device)
+			p.sendUnhealthy(device)
 		}
 	}