fix: mark all devices sharing a physical GPU unhealthy on Xid errors

[tryuuu]

Fixes #1929

checkHealth kept only one *Device per physical GPU UUID in parentToDeviceMap, so devices sharing a physical GPU overwrote each other. As a result, an Xid event marked at most one device as unhealthy.

With replication (time-slicing or MPS), the other replicas of the affected GPU remained schedulable. With MIG, an event could be dropped entirely if the single device retained in the map did not match the event's GI and CI.

This PR groups all logical devices by their parent GPU UUID and processes each device associated with the event UUID. Replicated devices are all reported as unhealthy, while MIG devices are filtered using each available event ID independently. If either the GI or CI is unavailable, the remaining available ID is still used to identify the affected devices.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tariq1890]

Thank you @tryuuu for your contribution. On an initial glance, this looks fine to me. Let me run some more tests against this branch and get back to you sometime next week

[kvalliyurnatt]

NVML I believe reports GPU instance Id and Compute Instance Id independently. I believe we can have an Xid event like
GI=0, CI=FFFFFFFF
GI=1, CI=0

In this case GI 1 should remain healthy, but currently we end up skipping the entire GPU and mark all instances of that GPU unhealthy, is that intended ?. I think we should individually verify the GIs and CIs

for _, d := range ds {
	if d.IsMigDevice() {
		gi := deviceIDToGiMap[d.ID]
		ci := deviceIDToCiMap[d.ID]
		if e.GpuInstanceId != 0xFFFFFFFF && gi != e.GpuInstanceId {
			continue
		}
		if e.ComputeInstanceId != 0xFFFFFFFF && ci != e.ComputeInstanceId {
			continue
		}
	}
	klog.Infof("XidCriticalError: Xid=%d on Device=%s; marking device as unhealthy.", e.EventData, d.ID)
	unhealthy <- d
}

We should also add a unit test for this case.

[aryangorwade]

@tryuuu About the testing changes: this small map fan-out fix adds five NVML fakes and ~266 lines. Existing tests isolate deterministic data transformations instead of simulating raw NVML (eg. testing DeviceMap's setEntry/insert methods). Could we extract and directly test the parent-UUID-to-logical-devices grouping, instead of testing the whole NVML event loop end to end?

I'm thinking we could create a type in health.go like this:

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
}

You would call placedDevices := make([]placedDevice, 0, len(devices)), placeDevices.append, and then ultimately set parentToDeviceMap := groupByParent(placedDevices).

This allows for easy unit testing:

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])
}

[aryangorwade]

Duplicate gpu registration has been an issue before this PR as well (and still is). Just something to note for a future fix.

[tryuuu]

@kvalliyurnatt @aryangorwade
Thank you both for the helpful and insightful comments. I agreed with your suggestions and have updated the PR accordingly. Could you please take another look when you have a chance?
Please feel free to let me know if any further changes are needed.