Prevent duplicate GPU registration for health checks

internal/rm/health.go

@@ -40,6 +40,29 @@ const (
 	envEnableHealthChecks = "DP_ENABLE_HEALTHCHECKS"
 )
 
+type placedDevice struct {
+	parentUUID string
+	device     *Device
+}
+
+func groupByParent(devices []placedDevice) map[string][]*Device {
+	grouped := make(map[string][]*Device)
+	for _, d := range devices {
+		grouped[d.parentUUID] = append(grouped[d.parentUUID], d.device)
+	}
+	return grouped
+}
+
+func matchesMigEvent(deviceGI, deviceCI, eventGI, eventCI uint32) bool {
+	if eventGI != 0xFFFFFFFF && deviceGI != eventGI {
+		return false
+	}
+	if eventCI != 0xFFFFFFFF && deviceCI != eventCI {
+		return false
+	}
+	return true
+}
+
 // 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 {
 	xids := getDisabledHealthCheckXids()
@@ -71,7 +94,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		_ = eventSet.Free()
 	}()
 
-	parentToDeviceMap := make(map[string]*Device)
+	placedDevices := make([]placedDevice, 0, len(devices))
 	deviceIDToGiMap := make(map[string]uint32)
 	deviceIDToCiMap := make(map[string]uint32)
 
@@ -85,29 +108,38 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		}
 		deviceIDToGiMap[d.ID] = gi
 		deviceIDToCiMap[d.ID] = ci
-		parentToDeviceMap[uuid] = d
+		placedDevices = append(placedDevices, placedDevice{parentUUID: uuid, device: d})
+	}
+	parentToDeviceMap := groupByParent(placedDevices)
 
-		gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid)
+	for parentUUID, d := range parentToDeviceMap {
+		gpu, ret := r.nvml.DeviceGetHandleByUUID(parentUUID)
 		if ret != nvml.SUCCESS {
 			klog.Infof("unable to get device handle from UUID: %v; marking it as unhealthy", ret)
-			unhealthy <- d
+			for _, d := range d {
+				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
+			klog.Infof("unable to determine the supported events for %v: %v; marking it as unhealthy", parentUUID, ret)
+			for _, d := range d {
+				unhealthy <- d
+			}
 			continue
 		}
 
 		ret = gpu.RegisterEvents(eventMask&supportedEvents, eventSet)
 		switch {
 		case ret == nvml.ERROR_NOT_SUPPORTED:
-			klog.Warningf("Device %v is too old to support healthchecking.", d.ID)
+			klog.Warningf("Device %v is too old to support healthchecking.", parentUUID)
 		case ret != nvml.SUCCESS:
-			klog.Infof("Marking device %v as unhealthy: %v", d.ID, ret)
-			unhealthy <- d
+			klog.Infof("Marking device %v as unhealthy: %v", parentUUID, ret)
+			for _, d := range d {
+				unhealthy <- d
+			}
 		}
 	}
 
@@ -151,23 +183,25 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 			continue
 		}
 
-		d, exists := parentToDeviceMap[eventUUID]
+		ds, exists := 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]
-			if gi != e.GpuInstanceId || ci != e.ComputeInstanceId {
-				continue
+		for _, d := range ds {
+			if d.IsMigDevice() {
+				gi := deviceIDToGiMap[d.ID]
+				ci := deviceIDToCiMap[d.ID]
+				if !matchesMigEvent(gi, ci, e.GpuInstanceId, e.ComputeInstanceId) {
+					continue
+				}
+				klog.Infof("Event for mig device %v (gi=%v, ci=%v)", d.ID, gi, ci)
 			}
-			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
+			klog.Infof("XidCriticalError: Xid=%d on Device=%s; marking device as unhealthy.", e.EventData, d.ID)
+			unhealthy <- d
+		}
 	}
 }
 

internal/rm/health_test.go

@@ -409,3 +409,87 @@ func TestGetMigDeviceParts(t *testing.T) {
 		})
 	}
 }
+
+func TestGroupByParent(t *testing.T) {
+	parentA := "GPU-A"
+	parentB := "GPU-B"
+
+	deviceA0 := &Device{Device: pluginapi.Device{ID: "GPU-A::0"}}
+	deviceA1 := &Device{Device: pluginapi.Device{ID: "GPU-A::1"}}
+	deviceB0 := &Device{Device: pluginapi.Device{ID: "GPU-B::0"}}
+
+	grouped := groupByParent([]placedDevice{
+		{parentUUID: parentA, device: deviceA0},
+		{parentUUID: parentA, device: deviceA1},
+		{parentUUID: parentB, device: deviceB0},
+	})
+
+	require.Equal(t, []*Device{deviceA0, deviceA1}, grouped[parentA])
+	require.Equal(t, []*Device{deviceB0}, grouped[parentB])
+}
+
+func TestMatchesMigEvent(t *testing.T) {
+	testCases := []struct {
+		description string
+		deviceGI    uint32
+		deviceCI    uint32
+		eventGI     uint32
+		eventCI     uint32
+		expected    bool
+	}{
+		{
+			description: "GI and CI match",
+			deviceGI:    3,
+			deviceCI:    0,
+			eventGI:     3,
+			eventCI:     0,
+			expected:    true,
+		},
+		{
+			description: "only GI is specified and matches",
+			deviceGI:    3,
+			deviceCI:    0,
+			eventGI:     3,
+			eventCI:     0xFFFFFFFF,
+			expected:    true,
+		},
+		{
+			description: "only GI is specified and does not match",
+			deviceGI:    5,
+			deviceCI:    0,
+			eventGI:     3,
+			eventCI:     0xFFFFFFFF,
+			expected:    false,
+		},
+		{
+			description: "only CI is specified and matches",
+			deviceGI:    3,
+			deviceCI:    0,
+			eventGI:     0xFFFFFFFF,
+			eventCI:     0,
+			expected:    true,
+		},
+		{
+			description: "only CI is specified and does not match",
+			deviceGI:    3,
+			deviceCI:    1,
+			eventGI:     0xFFFFFFFF,
+			eventCI:     0,
+			expected:    false,
+		},
+		{
+			description: "neither GI nor CI is specified",
+			deviceGI:    3,
+			deviceCI:    0,
+			eventGI:     0xFFFFFFFF,
+			eventCI:     0xFFFFFFFF,
+			expected:    true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			require.Equal(t, tc.expected, matchesMigEvent(tc.deviceGI, tc.deviceCI, tc.eventGI, tc.eventCI))
+		})
+	}
+}