refactor(health): extract nvmlHealthProvider and add tests

Refactor the health monitoring implementation:
- Extract nvmlHealthProvider struct to encapsulate health check state
- Add registerDeviceEvents() and runEventMonitor() methods
- Rename xids field to xidsDisabled for clarity
- Migrate from stop channel to context.Context for cancellation

Add comprehensive tests:
- TestCheckHealth validates XID event handling with mocks
- TestRegisterDeviceEventsNotSupported ensures old GPUs returning
  ERROR_NOT_SUPPORTED are not incorrectly marked unhealthy

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

Author: elezar

internal/rm/health.go

@@ -17,6 +17,7 @@
 package rm
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strconv"
@@ -40,8 +41,19 @@ const (
 	envEnableHealthChecks = "DP_ENABLE_HEALTHCHECKS"
 )
 
+type nvmlHealthProvider struct {
+	nvmllib           nvml.Interface
+	devices           Devices
+	parentToDeviceMap map[string]*Device
+	deviceIDToGiMap   map[string]uint32
+	deviceIDToCiMap   map[string]uint32
+
+	xidsDisabled disabledXIDs
+	unhealthy    chan<- *Device
+}
+
 // CheckHealth performs health checks on a set of devices, writing to the 'unhealthy' channel with any unhealthy devices
-func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devices, unhealthy chan<- *Device) error {
+func (r *nvmlResourceManager) checkHealth(ctx context.Context, devices Devices, unhealthy chan<- *Device) error {
 	xids := getDisabledHealthCheckXids()
 	if xids.IsAllDisabled() {
 		return nil
@@ -71,13 +83,14 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		_ = eventSet.Free()
 	}()
 
+	// Construct the device maps.
+	// TODO: This should be factored out. The main issue is marking the devices
+	// unhealthy as part of this loop.
 	parentToDeviceMap := make(map[string]*Device)
 	deviceIDToGiMap := make(map[string]uint32)
 	deviceIDToCiMap := make(map[string]uint32)
-
-	eventMask := uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError)
 	for _, d := range devices {
-		uuid, gi, ci, err := r.getDevicePlacement(d)
+		uuid, gi, ci, err := (&withDevicePlacements{r.nvml}).getDevicePlacement(d)
 		if err != nil {
 			klog.Warningf("Could not determine device placement for %v: %v; Marking it unhealthy.", d.ID, err)
 			unhealthy <- d
@@ -86,18 +99,35 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		deviceIDToGiMap[d.ID] = gi
 		deviceIDToCiMap[d.ID] = ci
 		parentToDeviceMap[uuid] = d
+	}
+
+	p := nvmlHealthProvider{
+		nvmllib:           r.nvml,
+		devices:           devices,
+		unhealthy:         unhealthy,
+		parentToDeviceMap: parentToDeviceMap,
+		deviceIDToGiMap:   deviceIDToGiMap,
+		deviceIDToCiMap:   deviceIDToCiMap,
+		xidsDisabled:      xids,
+	}
+	p.registerDeviceEvents(eventSet)
+
+	return p.runEventMonitor(ctx, eventSet)
+}
 
-		gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid)
+func (r *nvmlHealthProvider) registerDeviceEvents(eventSet nvml.EventSet) {
+	eventMask := uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError)
+	for uuid, d := range r.parentToDeviceMap {
+		gpu, ret := r.nvmllib.DeviceGetHandleByUUID(uuid)
 		if ret != nvml.SUCCESS {
 			klog.Infof("unable to get device handle from UUID: %v; marking it as unhealthy", ret)
-			unhealthy <- d
+			r.unhealthy <- d
 			continue
 		}
-
 		supportedEvents, ret := gpu.GetSupportedEventTypes()
 		if ret != nvml.SUCCESS {
 			klog.Infof("unable to determine the supported events for %v: %v; marking it as unhealthy", d.ID, ret)
-			unhealthy <- d
+			r.unhealthy <- d
 			continue
 		}
 
@@ -107,14 +137,16 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 			klog.Warningf("Device %v is too old to support healthchecking.", d.ID)
 		case ret != nvml.SUCCESS:
 			klog.Infof("Marking device %v as unhealthy: %v", d.ID, ret)
-			unhealthy <- d
+			r.unhealthy <- d
 		}
 	}
+}
 
+func (r *nvmlHealthProvider) runEventMonitor(ctx context.Context, eventSet nvml.EventSet) error {
 	for {
 		select {
-		case <-stop:
-			return nil
+		case <-ctx.Done():
+			return ctx.Err()
 		default:
 		}
 
@@ -124,18 +156,20 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		}
 		if ret != nvml.SUCCESS {
 			klog.Infof("Error waiting for event: %v; Marking all devices as unhealthy", ret)
-			for _, d := range devices {
-				unhealthy <- d
+			for _, d := range r.devices {
+				r.unhealthy <- d
 			}
 			continue
 		}
 
+		// TODO: We create an event mask for other event types but don't handle
+		// them here.
 		if e.EventType != nvml.EventTypeXidCriticalError {
 			klog.Infof("Skipping non-nvmlEventTypeXidCriticalError event: %+v", e)
 			continue
 		}
 
-		if xids.IsDisabled(e.EventData) {
+		if r.xidsDisabled.IsDisabled(e.EventData) {
 			klog.Infof("Skipping event %+v", e)
 			continue
 		}
@@ -145,29 +179,29 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		if ret != nvml.SUCCESS {
 			// If we cannot reliably determine the device UUID, we mark all devices as unhealthy.
 			klog.Infof("Failed to determine uuid for event %v: %v; Marking all devices as unhealthy.", e, ret)
-			for _, d := range devices {
-				unhealthy <- d
+			for _, d := range r.devices {
+				r.unhealthy <- d
 			}
 			continue
 		}
 
-		d, exists := parentToDeviceMap[eventUUID]
+		d, exists := r.parentToDeviceMap[eventUUID]
 		if !exists {
 			klog.Infof("Ignoring event for unexpected device: %v", eventUUID)
 			continue
 		}
 
 		if d.IsMigDevice() && e.GpuInstanceId != 0xFFFFFFFF && e.ComputeInstanceId != 0xFFFFFFFF {
-			gi := deviceIDToGiMap[d.ID]
-			ci := deviceIDToCiMap[d.ID]
+			gi := r.deviceIDToGiMap[d.ID]
+			ci := r.deviceIDToCiMap[d.ID]
 			if gi != e.GpuInstanceId || ci != e.ComputeInstanceId {
 				continue
 			}
 			klog.Infof("Event for mig device %v (gi=%v, ci=%v)", d.ID, gi, ci)
 		}
 
 		klog.Infof("XidCriticalError: Xid=%d on Device=%s; marking device as unhealthy.", e.EventData, d.ID)
-		unhealthy <- d
+		r.unhealthy <- d
 	}
 }
 
@@ -276,25 +310,29 @@ func newHealthCheckXIDs(xids ...string) disabledXIDs {
 	return output
 }
 
+type withDevicePlacements struct {
+	nvml.Interface
+}
+
 // 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 (r *nvmlResourceManager) getDevicePlacement(d *Device) (string, uint32, uint32, error) {
+func (r *withDevicePlacements) getDevicePlacement(d *Device) (string, uint32, uint32, error) {
 	if !d.IsMigDevice() {
 		return d.GetUUID(), 0xFFFFFFFF, 0xFFFFFFFF, nil
 	}
 	return r.getMigDeviceParts(d)
 }
 
 // getMigDeviceParts returns the parent GI and CI ids of the MIG device.
-func (r *nvmlResourceManager) getMigDeviceParts(d *Device) (string, uint32, uint32, error) {
+func (r *withDevicePlacements) 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 := r.nvml.DeviceGetHandleByUUID(uuid)
+	mig, ret := r.DeviceGetHandleByUUID(uuid)
 	if ret == nvml.SUCCESS {
 		parentHandle, ret := mig.GetDeviceHandleFromMigDeviceHandle()
 		if ret != nvml.SUCCESS {

internal/rm/health_test.go

@@ -17,10 +17,15 @@
 package rm
 
 import (
+	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
 
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/NVIDIA/go-nvml/pkg/nvml/mock"
+	"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
 	"github.com/stretchr/testify/require"
 )
 
@@ -221,3 +226,147 @@ func TestGetDisabledHealthCheckXids(t *testing.T) {
 		})
 	}
 }
+
+func TestCheckHealth(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	unhealthy := make(chan *Device)
+
+	server := dgxa100.New()
+
+	deviceMock := server.Devices[0].(*dgxa100.Device)
+	deviceMock.GetSupportedEventTypesFunc = func() (uint64, nvml.Return) {
+		return nvml.EventTypeXidCriticalError, nvml.SUCCESS
+	}
+	// TODO: Should this be more dynamic?
+	deviceMock.RegisterEventsFunc = func(v uint64, eventSet nvml.EventSet) nvml.Return {
+		return nvml.SUCCESS
+	}
+
+	var count int
+	eventData := []nvml.EventData{
+		{
+			EventData: 109,
+			EventType: nvml.EventTypeXidCriticalError,
+			Device:    server.Devices[0],
+		},
+		{
+			EventData: 48,
+			EventType: nvml.EventTypeXidCriticalError,
+			Device:    server.Devices[0],
+		},
+	}
+
+	server.EventSetCreateFunc = func() (nvml.EventSet, nvml.Return) {
+		es := &mock.EventSet{
+			WaitFunc: func(v uint32) (nvml.EventData, nvml.Return) {
+				ed := eventData[count%len(eventData)]
+				count++
+				if count == len(eventData) {
+					// Cancel the context to signal the health checker to terminate
+					// after the predefined events have been triggered.
+					cancel()
+				}
+				return ed, nvml.SUCCESS
+			},
+			FreeFunc: func() nvml.Return {
+				return nvml.SUCCESS
+			},
+		}
+		return es, nvml.SUCCESS
+	}
+
+	r := &nvmlResourceManager{
+		nvml: server,
+	}
+
+	var unhealthyDevices []*Device
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+		for d := range unhealthy {
+			unhealthyDevices = append(unhealthyDevices, d)
+		}
+	}()
+
+	var expectedDevices []*Device
+
+	devices := make(Devices)
+	for i, d := range server.Devices {
+		device, err := BuildDevice(newNvmlGPUDevice(i, d))
+		require.NoError(t, err)
+		devices[device.GetUUID()] = device
+		expectedDevices = append(expectedDevices, device)
+		// TODO: We only expect a single unhealthy event for the first device.
+		break
+	}
+
+	err := r.checkHealth(ctx, devices, unhealthy)
+	require.ErrorIs(t, err, context.Canceled)
+
+	// Close the channel and wait for the goroutine to finish collecting unhealthy devices
+	close(unhealthy)
+	wg.Wait()
+
+	require.EqualValues(t, expectedDevices, unhealthyDevices)
+}
+
+func TestRegisterDeviceEventsNotSupported(t *testing.T) {
+	// This test verifies that devices which return ERROR_NOT_SUPPORTED from
+	// RegisterEvents are NOT marked as unhealthy. This is the expected behavior
+	// for old GPUs that don't support event registration.
+	// See: commit 8cd14472a "Fix healthchecking on old devices"
+
+	unhealthy := make(chan *Device, 10)
+
+	server := dgxa100.New()
+
+	deviceMock := server.Devices[0].(*dgxa100.Device)
+	deviceMock.GetSupportedEventTypesFunc = func() (uint64, nvml.Return) {
+		return nvml.EventTypeXidCriticalError, nvml.SUCCESS
+	}
+	// Simulate an old device that doesn't support event registration
+	deviceMock.RegisterEventsFunc = func(v uint64, eventSet nvml.EventSet) nvml.Return {
+		return nvml.ERROR_NOT_SUPPORTED
+	}
+
+	eventSet := &mock.EventSet{
+		FreeFunc: func() nvml.Return {
+			return nvml.SUCCESS
+		},
+	}
+
+	devices := make(Devices)
+	parentToDeviceMap := make(map[string]*Device)
+
+	for i, d := range server.Devices {
+		device, err := BuildDevice(newNvmlGPUDevice(i, d))
+		require.NoError(t, err)
+		devices[device.GetUUID()] = device
+		parentToDeviceMap[device.GetUUID()] = device
+		break
+	}
+
+	p := nvmlHealthProvider{
+		nvmllib:           server,
+		devices:           devices,
+		unhealthy:         unhealthy,
+		parentToDeviceMap: parentToDeviceMap,
+		deviceIDToGiMap:   make(map[string]uint32),
+		deviceIDToCiMap:   make(map[string]uint32),
+	}
+
+	p.registerDeviceEvents(eventSet)
+
+	// Close the channel so we can drain it
+	close(unhealthy)
+
+	// Verify that NO devices were marked as unhealthy
+	var unhealthyDevices []*Device
+	for d := range unhealthy {
+		unhealthyDevices = append(unhealthyDevices, d)
+	}
+
+	require.Empty(t, unhealthyDevices, "Devices returning ERROR_NOT_SUPPORTED should NOT be marked as unhealthy")
+}