From 25772ed616ed0eb60c65278eaa6796f97be5ad9f Mon Sep 17 00:00:00 2001 From: tryuuu Date: Fri, 17 Jul 2026 20:09:49 +0900 Subject: [PATCH 1/4] fix: mark all devices sharing a physical GPU unhealthy on Xid errors Signed-off-by: tryuuu --- internal/rm/health.go | 26 ++-- internal/rm/health_test.go | 266 +++++++++++++++++++++++++++++++++++++ 2 files changed, 280 insertions(+), 12 deletions(-) diff --git a/internal/rm/health.go b/internal/rm/health.go index 1ff0e9b9b..49779dcc4 100644 --- a/internal/rm/health.go +++ b/internal/rm/health.go @@ -71,7 +71,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic _ = eventSet.Free() }() - parentToDeviceMap := make(map[string]*Device) + parentToDeviceMap := make(map[string][]*Device) deviceIDToGiMap := make(map[string]uint32) deviceIDToCiMap := make(map[string]uint32) @@ -85,7 +85,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic } deviceIDToGiMap[d.ID] = gi deviceIDToCiMap[d.ID] = ci - parentToDeviceMap[uuid] = d + parentToDeviceMap[uuid] = append(parentToDeviceMap[uuid], d) gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid) if ret != nvml.SUCCESS { @@ -151,23 +151,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() && e.GpuInstanceId != 0xFFFFFFFF && e.ComputeInstanceId != 0xFFFFFFFF { + gi := deviceIDToGiMap[d.ID] + ci := 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("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 + } } } diff --git a/internal/rm/health_test.go b/internal/rm/health_test.go index 602c69815..6c93fa841 100644 --- a/internal/rm/health_test.go +++ b/internal/rm/health_test.go @@ -20,6 +20,7 @@ import ( "fmt" "strings" "testing" + "time" "github.com/NVIDIA/go-nvml/pkg/nvml" "github.com/stretchr/testify/require" @@ -409,3 +410,268 @@ func TestGetMigDeviceParts(t *testing.T) { }) } } + +// fakeHealthNvmlLib is a minimal nvml.Interface test double for driving checkHealth +// without a real GPU; only the calls made by checkHealth are implemented. +type fakeHealthNvmlLib struct { + nvml.Interface + eventSet *fakeEventSet +} + +func (f *fakeHealthNvmlLib) Init() nvml.Return { + return nvml.SUCCESS +} + +func (f *fakeHealthNvmlLib) Shutdown() nvml.Return { + return nvml.SUCCESS +} + +func (f *fakeHealthNvmlLib) EventSetCreate() (nvml.EventSet, nvml.Return) { + return f.eventSet, nvml.SUCCESS +} + +func (f *fakeHealthNvmlLib) DeviceGetHandleByUUID(string) (nvml.Device, nvml.Return) { + return &fakeGpuHandle{}, nvml.SUCCESS +} + +// fakeGpuHandle is a minimal nvml.Device test double for the handle returned by +// DeviceGetHandleByUUID; only the event-registration calls are implemented. +type fakeGpuHandle struct { + nvml.Device +} + +func (f *fakeGpuHandle) GetSupportedEventTypes() (uint64, nvml.Return) { + return ^uint64(0), nvml.SUCCESS +} + +func (f *fakeGpuHandle) RegisterEvents(uint64, nvml.EventSet) nvml.Return { + return nvml.SUCCESS +} + +// fakeEventDevice is a minimal nvml.Device test double for the device embedded in +// a scripted nvml.EventData; only GetUUID is used by checkHealth. +type fakeEventDevice struct { + nvml.Device + uuid string +} + +func (f *fakeEventDevice) GetUUID() (string, nvml.Return) { + return f.uuid, nvml.SUCCESS +} + +// fakeEventSet is a scripted nvml.EventSet test double: it returns the configured +// events in order, and ERROR_TIMEOUT once they have all been consumed. +type fakeEventSet struct { + nvml.EventSet + events []nvml.EventData + idx int +} + +func (f *fakeEventSet) Wait(uint32) (nvml.EventData, nvml.Return) { + if f.idx < len(f.events) { + e := f.events[f.idx] + f.idx++ + return e, nvml.SUCCESS + } + return nvml.EventData{}, nvml.ERROR_TIMEOUT +} + +func (f *fakeEventSet) Free() nvml.Return { + return nvml.SUCCESS +} + +// fakeMigHealthNvmlLib is a fakeHealthNvmlLib whose MIG UUID lookups fail, so +// that getMigDeviceParts falls back to parsing the legacy MIG UUID format. +type fakeMigHealthNvmlLib struct { + fakeHealthNvmlLib +} + +func (f *fakeMigHealthNvmlLib) DeviceGetHandleByUUID(uuid string) (nvml.Device, nvml.Return) { + if strings.HasPrefix(uuid, "MIG-") { + return nil, nvml.ERROR_NOT_SUPPORTED + } + return &fakeGpuHandle{}, nvml.SUCCESS +} + +// TestCheckHealthMigXidFanout verifies that an Xid event on a parent GPU is +// routed to the MIG devices it hosts: an instance-specific event marks only +// the matching MIG device unhealthy, while a GPU-wide event marks all of them. +func TestCheckHealthMigXidFanout(t *testing.T) { + t.Setenv(envDisableHealthChecks, "") + t.Setenv(envEnableHealthChecks, "") + + parentUUID := "GPU-5c89852c-d268-c3f3-1b07-005d5ae1dc3f" + migGi3 := "MIG-" + parentUUID + "/3/0" + migGi5 := "MIG-" + parentUUID + "/5/0" + + testCases := []struct { + description string + gi uint32 + ci uint32 + expected []string + }{ + { + description: "instance-specific event marks only the matching MIG device unhealthy", + gi: 3, + ci: 0, + expected: []string{migGi3}, + }, + { + description: "GPU-wide event marks every MIG device on the parent unhealthy", + gi: 0xFFFFFFFF, + ci: 0xFFFFFFFF, + expected: []string{migGi3, migGi5}, + }, + { + description: "event for an unknown instance marks nothing unhealthy", + gi: 9, + ci: 0, + expected: nil, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices := make(Devices) + for _, id := range []string{migGi3, migGi5} { + devices[id] = &Device{ + Device: pluginapi.Device{ID: id, Health: pluginapi.Healthy}, + Index: "0:0", + } + } + + eventSet := &fakeEventSet{ + events: []nvml.EventData{ + { + Device: &fakeEventDevice{uuid: parentUUID}, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: tc.gi, + ComputeInstanceId: tc.ci, + }, + }, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{devices: devices}, + nvml: &fakeMigHealthNvmlLib{fakeHealthNvmlLib{eventSet: eventSet}}, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, len(devices)) + done := make(chan error, 1) + go func() { + done <- r.checkHealth(stop, devices, unhealthy) + }() + + received := make(map[string]bool) + timeout := time.After(2 * time.Second) + collect: + for len(received) < len(tc.expected) { + select { + case d := <-unhealthy: + received[d.ID] = true + case <-timeout: + break collect + } + } + // Grace period to catch devices that should not have been reported. + select { + case d := <-unhealthy: + received[d.ID] = true + case <-time.After(100 * time.Millisecond): + } + + close(stop) + require.NoError(t, <-done) + + expected := make(map[string]bool) + for _, id := range tc.expected { + expected[id] = true + } + require.Equal(t, expected, received) + }) + } +} + +func TestCheckHealthReplicatedXidFanout(t *testing.T) { + t.Setenv(envDisableHealthChecks, "") + t.Setenv(envEnableHealthChecks, "") + + uuid := "GPU-3a1f2c4e-8b2d-41a9-9c3f-1a2b3c4d5e6f" + + newReplicatedDevices := func(replicas int) Devices { + devices := make(Devices) + for i := 0; i < replicas; i++ { + id := string(NewAnnotatedID(uuid, i)) + devices[id] = &Device{ + Device: pluginapi.Device{ID: id, Health: pluginapi.Healthy}, + Index: "0", + } + } + return devices + } + + testCases := []struct { + description string + replicas int + }{ + { + description: "non-replicated device: the Xid event marks the single device unhealthy", + replicas: 1, + }, + { + description: "replicated device: an Xid event on the shared physical GPU marks every replica unhealthy", + replicas: 4, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices := newReplicatedDevices(tc.replicas) + + eventSet := &fakeEventSet{ + events: []nvml.EventData{ + { + Device: &fakeEventDevice{uuid: uuid}, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{devices: devices}, + nvml: &fakeHealthNvmlLib{eventSet: eventSet}, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, len(devices)) + done := make(chan error, 1) + go func() { + done <- r.checkHealth(stop, devices, unhealthy) + }() + + received := make(map[string]bool) + timeout := time.After(2 * time.Second) + for len(received) < len(devices) { + select { + case d := <-unhealthy: + received[d.ID] = true + case <-timeout: + t.Fatalf("timed out waiting for %d unhealthy devices; received %d", len(devices), len(received)) + } + } + + close(stop) + require.NoError(t, <-done) + + require.Len(t, received, len(devices)) + for id := range devices { + require.True(t, received[id], "expected device %v to be marked unhealthy", id) + } + }) + } +} From 5397a2b5388a78321f8b88518944862d577eb85b Mon Sep 17 00:00:00 2001 From: tryuuu Date: Sun, 20 Sep 2026 15:45:05 +0900 Subject: [PATCH 2/4] fix: handle partial MIG placement in Xid events Signed-off-by: tryuuu --- internal/rm/health.go | 7 +++++-- internal/rm/health_test.go | 14 +++++++++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/internal/rm/health.go b/internal/rm/health.go index 49779dcc4..4c5a3f762 100644 --- a/internal/rm/health.go +++ b/internal/rm/health.go @@ -158,10 +158,13 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic } for _, d := range ds { - if d.IsMigDevice() && e.GpuInstanceId != 0xFFFFFFFF && e.ComputeInstanceId != 0xFFFFFFFF { + if d.IsMigDevice() { gi := deviceIDToGiMap[d.ID] ci := deviceIDToCiMap[d.ID] - if gi != e.GpuInstanceId || ci != e.ComputeInstanceId { + if e.GpuInstanceId != 0xFFFFFFFF && gi != e.GpuInstanceId { + continue + } + if e.ComputeInstanceId != 0xFFFFFFFF && ci != e.ComputeInstanceId { continue } klog.Infof("Event for mig device %v (gi=%v, ci=%v)", d.ID, gi, ci) diff --git a/internal/rm/health_test.go b/internal/rm/health_test.go index 6c93fa841..5ffaf0f85 100644 --- a/internal/rm/health_test.go +++ b/internal/rm/health_test.go @@ -502,7 +502,7 @@ func TestCheckHealthMigXidFanout(t *testing.T) { parentUUID := "GPU-5c89852c-d268-c3f3-1b07-005d5ae1dc3f" migGi3 := "MIG-" + parentUUID + "/3/0" - migGi5 := "MIG-" + parentUUID + "/5/0" + migGi5 := "MIG-" + parentUUID + "/5/1" testCases := []struct { description string @@ -516,6 +516,18 @@ func TestCheckHealthMigXidFanout(t *testing.T) { ci: 0, expected: []string{migGi3}, }, + { + description: "event with only a GI ID marks only devices in the matching GI unhealthy", + gi: 3, + ci: 0xFFFFFFFF, + expected: []string{migGi3}, + }, + { + description: "event with only a CI ID marks only devices with the matching CI unhealthy", + gi: 0xFFFFFFFF, + ci: 0, + expected: []string{migGi3}, + }, { description: "GPU-wide event marks every MIG device on the parent unhealthy", gi: 0xFFFFFFFF, From 59ea5f1cfb8b7577ae1ff50f556e6a55d142f720 Mon Sep 17 00:00:00 2001 From: tryuuu Date: Sun, 20 Sep 2026 17:29:26 +0900 Subject: [PATCH 3/4] refactor: isolate Xid health check logic for testing Signed-off-by: tryuuu --- internal/rm/health.go | 33 +++- internal/rm/health_test.go | 306 +++++++------------------------------ 2 files changed, 83 insertions(+), 256 deletions(-) diff --git a/internal/rm/health.go b/internal/rm/health.go index 4c5a3f762..7edac3da1 100644 --- a/internal/rm/health.go +++ b/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,7 +108,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic } deviceIDToGiMap[d.ID] = gi deviceIDToCiMap[d.ID] = ci - parentToDeviceMap[uuid] = append(parentToDeviceMap[uuid], d) + placedDevices = append(placedDevices, placedDevice{parentUUID: uuid, device: d}) gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid) if ret != nvml.SUCCESS { @@ -110,6 +133,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic unhealthy <- d } } + parentToDeviceMap := groupByParent(placedDevices) for { select { @@ -161,10 +185,7 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic 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 { + if !matchesMigEvent(gi, ci, e.GpuInstanceId, e.ComputeInstanceId) { continue } klog.Infof("Event for mig device %v (gi=%v, ci=%v)", d.ID, gi, ci) diff --git a/internal/rm/health_test.go b/internal/rm/health_test.go index 5ffaf0f85..9cdde73f9 100644 --- a/internal/rm/health_test.go +++ b/internal/rm/health_test.go @@ -20,7 +20,6 @@ import ( "fmt" "strings" "testing" - "time" "github.com/NVIDIA/go-nvml/pkg/nvml" "github.com/stretchr/testify/require" @@ -411,279 +410,86 @@ func TestGetMigDeviceParts(t *testing.T) { } } -// fakeHealthNvmlLib is a minimal nvml.Interface test double for driving checkHealth -// without a real GPU; only the calls made by checkHealth are implemented. -type fakeHealthNvmlLib struct { - nvml.Interface - eventSet *fakeEventSet -} - -func (f *fakeHealthNvmlLib) Init() nvml.Return { - return nvml.SUCCESS -} - -func (f *fakeHealthNvmlLib) Shutdown() nvml.Return { - return nvml.SUCCESS -} - -func (f *fakeHealthNvmlLib) EventSetCreate() (nvml.EventSet, nvml.Return) { - return f.eventSet, nvml.SUCCESS -} - -func (f *fakeHealthNvmlLib) DeviceGetHandleByUUID(string) (nvml.Device, nvml.Return) { - return &fakeGpuHandle{}, nvml.SUCCESS -} - -// fakeGpuHandle is a minimal nvml.Device test double for the handle returned by -// DeviceGetHandleByUUID; only the event-registration calls are implemented. -type fakeGpuHandle struct { - nvml.Device -} - -func (f *fakeGpuHandle) GetSupportedEventTypes() (uint64, nvml.Return) { - return ^uint64(0), nvml.SUCCESS -} +func TestGroupByParent(t *testing.T) { + parentA := "GPU-A" + parentB := "GPU-B" -func (f *fakeGpuHandle) RegisterEvents(uint64, nvml.EventSet) nvml.Return { - return nvml.SUCCESS -} + 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"}} -// fakeEventDevice is a minimal nvml.Device test double for the device embedded in -// a scripted nvml.EventData; only GetUUID is used by checkHealth. -type fakeEventDevice struct { - nvml.Device - uuid string -} - -func (f *fakeEventDevice) GetUUID() (string, nvml.Return) { - return f.uuid, nvml.SUCCESS -} + grouped := groupByParent([]placedDevice{ + {parentUUID: parentA, device: deviceA0}, + {parentUUID: parentA, device: deviceA1}, + {parentUUID: parentB, device: deviceB0}, + }) -// fakeEventSet is a scripted nvml.EventSet test double: it returns the configured -// events in order, and ERROR_TIMEOUT once they have all been consumed. -type fakeEventSet struct { - nvml.EventSet - events []nvml.EventData - idx int + require.Equal(t, []*Device{deviceA0, deviceA1}, grouped[parentA]) + require.Equal(t, []*Device{deviceB0}, grouped[parentB]) } -func (f *fakeEventSet) Wait(uint32) (nvml.EventData, nvml.Return) { - if f.idx < len(f.events) { - e := f.events[f.idx] - f.idx++ - return e, nvml.SUCCESS - } - return nvml.EventData{}, nvml.ERROR_TIMEOUT -} - -func (f *fakeEventSet) Free() nvml.Return { - return nvml.SUCCESS -} - -// fakeMigHealthNvmlLib is a fakeHealthNvmlLib whose MIG UUID lookups fail, so -// that getMigDeviceParts falls back to parsing the legacy MIG UUID format. -type fakeMigHealthNvmlLib struct { - fakeHealthNvmlLib -} - -func (f *fakeMigHealthNvmlLib) DeviceGetHandleByUUID(uuid string) (nvml.Device, nvml.Return) { - if strings.HasPrefix(uuid, "MIG-") { - return nil, nvml.ERROR_NOT_SUPPORTED - } - return &fakeGpuHandle{}, nvml.SUCCESS -} - -// TestCheckHealthMigXidFanout verifies that an Xid event on a parent GPU is -// routed to the MIG devices it hosts: an instance-specific event marks only -// the matching MIG device unhealthy, while a GPU-wide event marks all of them. -func TestCheckHealthMigXidFanout(t *testing.T) { - t.Setenv(envDisableHealthChecks, "") - t.Setenv(envEnableHealthChecks, "") - - parentUUID := "GPU-5c89852c-d268-c3f3-1b07-005d5ae1dc3f" - migGi3 := "MIG-" + parentUUID + "/3/0" - migGi5 := "MIG-" + parentUUID + "/5/1" - +func TestMatchesMigEvent(t *testing.T) { testCases := []struct { description string - gi uint32 - ci uint32 - expected []string + deviceGI uint32 + deviceCI uint32 + eventGI uint32 + eventCI uint32 + expected bool }{ { - description: "instance-specific event marks only the matching MIG device unhealthy", - gi: 3, - ci: 0, - expected: []string{migGi3}, + description: "GI and CI match", + deviceGI: 3, + deviceCI: 0, + eventGI: 3, + eventCI: 0, + expected: true, }, { - description: "event with only a GI ID marks only devices in the matching GI unhealthy", - gi: 3, - ci: 0xFFFFFFFF, - expected: []string{migGi3}, + description: "only GI is specified and matches", + deviceGI: 3, + deviceCI: 0, + eventGI: 3, + eventCI: 0xFFFFFFFF, + expected: true, }, { - description: "event with only a CI ID marks only devices with the matching CI unhealthy", - gi: 0xFFFFFFFF, - ci: 0, - expected: []string{migGi3}, + description: "only GI is specified and does not match", + deviceGI: 5, + deviceCI: 0, + eventGI: 3, + eventCI: 0xFFFFFFFF, + expected: false, }, { - description: "GPU-wide event marks every MIG device on the parent unhealthy", - gi: 0xFFFFFFFF, - ci: 0xFFFFFFFF, - expected: []string{migGi3, migGi5}, + description: "only CI is specified and matches", + deviceGI: 3, + deviceCI: 0, + eventGI: 0xFFFFFFFF, + eventCI: 0, + expected: true, }, { - description: "event for an unknown instance marks nothing unhealthy", - gi: 9, - ci: 0, - expected: nil, + description: "only CI is specified and does not match", + deviceGI: 3, + deviceCI: 1, + eventGI: 0xFFFFFFFF, + eventCI: 0, + expected: false, }, - } - - for _, tc := range testCases { - t.Run(tc.description, func(t *testing.T) { - devices := make(Devices) - for _, id := range []string{migGi3, migGi5} { - devices[id] = &Device{ - Device: pluginapi.Device{ID: id, Health: pluginapi.Healthy}, - Index: "0:0", - } - } - - eventSet := &fakeEventSet{ - events: []nvml.EventData{ - { - Device: &fakeEventDevice{uuid: parentUUID}, - EventType: nvml.EventTypeXidCriticalError, - EventData: 79, - GpuInstanceId: tc.gi, - ComputeInstanceId: tc.ci, - }, - }, - } - - r := &nvmlResourceManager{ - resourceManager: resourceManager{devices: devices}, - nvml: &fakeMigHealthNvmlLib{fakeHealthNvmlLib{eventSet: eventSet}}, - } - - stop := make(chan interface{}) - unhealthy := make(chan *Device, len(devices)) - done := make(chan error, 1) - go func() { - done <- r.checkHealth(stop, devices, unhealthy) - }() - - received := make(map[string]bool) - timeout := time.After(2 * time.Second) - collect: - for len(received) < len(tc.expected) { - select { - case d := <-unhealthy: - received[d.ID] = true - case <-timeout: - break collect - } - } - // Grace period to catch devices that should not have been reported. - select { - case d := <-unhealthy: - received[d.ID] = true - case <-time.After(100 * time.Millisecond): - } - - close(stop) - require.NoError(t, <-done) - - expected := make(map[string]bool) - for _, id := range tc.expected { - expected[id] = true - } - require.Equal(t, expected, received) - }) - } -} - -func TestCheckHealthReplicatedXidFanout(t *testing.T) { - t.Setenv(envDisableHealthChecks, "") - t.Setenv(envEnableHealthChecks, "") - - uuid := "GPU-3a1f2c4e-8b2d-41a9-9c3f-1a2b3c4d5e6f" - - newReplicatedDevices := func(replicas int) Devices { - devices := make(Devices) - for i := 0; i < replicas; i++ { - id := string(NewAnnotatedID(uuid, i)) - devices[id] = &Device{ - Device: pluginapi.Device{ID: id, Health: pluginapi.Healthy}, - Index: "0", - } - } - return devices - } - - testCases := []struct { - description string - replicas int - }{ { - description: "non-replicated device: the Xid event marks the single device unhealthy", - replicas: 1, - }, - { - description: "replicated device: an Xid event on the shared physical GPU marks every replica unhealthy", - replicas: 4, + 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) { - devices := newReplicatedDevices(tc.replicas) - - eventSet := &fakeEventSet{ - events: []nvml.EventData{ - { - Device: &fakeEventDevice{uuid: uuid}, - EventType: nvml.EventTypeXidCriticalError, - EventData: 79, - GpuInstanceId: 0xFFFFFFFF, - ComputeInstanceId: 0xFFFFFFFF, - }, - }, - } - - r := &nvmlResourceManager{ - resourceManager: resourceManager{devices: devices}, - nvml: &fakeHealthNvmlLib{eventSet: eventSet}, - } - - stop := make(chan interface{}) - unhealthy := make(chan *Device, len(devices)) - done := make(chan error, 1) - go func() { - done <- r.checkHealth(stop, devices, unhealthy) - }() - - received := make(map[string]bool) - timeout := time.After(2 * time.Second) - for len(received) < len(devices) { - select { - case d := <-unhealthy: - received[d.ID] = true - case <-timeout: - t.Fatalf("timed out waiting for %d unhealthy devices; received %d", len(devices), len(received)) - } - } - - close(stop) - require.NoError(t, <-done) - - require.Len(t, received, len(devices)) - for id := range devices { - require.True(t, received[id], "expected device %v to be marked unhealthy", id) - } + require.Equal(t, tc.expected, matchesMigEvent(tc.deviceGI, tc.deviceCI, tc.eventGI, tc.eventCI)) }) } } From 144e037e17b4d798155763aedfafc9475f5aefb5 Mon Sep 17 00:00:00 2001 From: Aryan Gorwade Date: Mon, 21 Sep 2026 12:25:57 -0700 Subject: [PATCH 4/4] Prevent duplicate GPU registration Signed-off-by: Aryan Gorwade --- internal/rm/health.go | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/internal/rm/health.go b/internal/rm/health.go index 7edac3da1..1f3e68d67 100644 --- a/internal/rm/health.go +++ b/internal/rm/health.go @@ -109,31 +109,39 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic deviceIDToGiMap[d.ID] = gi deviceIDToCiMap[d.ID] = ci 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 + } } } - parentToDeviceMap := groupByParent(placedDevices) for { select {