From 17d38cc66b8205db21ab39e7cb2667a98c44d4b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Abra=C3=A3o=20Teixeira?= Date: Sun, 30 Aug 2026 12:36:45 -0300 Subject: [PATCH] Fix Xid healthcheck for replicated GPUs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a physical GPU experiences a critical Xid error, the device plugin needs to mark all associated logical devices (Time-Slicing, MPS) as unhealthy. Previously, the parentToDeviceMap was overwritten by replicated logical devices, causing only one of them to be marked unhealthy on an Xid event. This patch changes the mapping to map[string][]*Device so all logical devices bound to the failing physical UUID are properly marked. Fixes #1929 Signed-off-by: Abraão Teixeira --- internal/rm/health.go | 67 +++-- internal/rm/health_test.go | 579 +++++++++++++++++++++++++++++++++++++ 2 files changed, 615 insertions(+), 31 deletions(-) diff --git a/internal/rm/health.go b/internal/rm/health.go index 1ff0e9b9b..580f2c3fc 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,30 +85,33 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic } deviceIDToGiMap[d.ID] = gi deviceIDToCiMap[d.ID] = ci - parentToDeviceMap[uuid] = d - gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid) - if ret != nvml.SUCCESS { - klog.Infof("unable to get device handle from UUID: %v; marking it as unhealthy", ret) - unhealthy <- d - continue - } + if _, exists := parentToDeviceMap[uuid]; !exists { + gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid) + if ret != nvml.SUCCESS { + klog.Infof("unable to get device handle from UUID: %v; marking it as unhealthy", ret) + 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 - 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 + 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) - case ret != nvml.SUCCESS: - klog.Infof("Marking device %v as unhealthy: %v", d.ID, ret) - unhealthy <- d + 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) + case ret != nvml.SUCCESS: + klog.Infof("Marking device %v as unhealthy: %v", d.ID, ret) + unhealthy <- d + } } + + parentToDeviceMap[uuid] = append(parentToDeviceMap[uuid], d) } for { @@ -151,23 +154,25 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic continue } - d, exists := parentToDeviceMap[eventUUID] + devs, 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 devs { + 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..4ebfce57b 100644 --- a/internal/rm/health_test.go +++ b/internal/rm/health_test.go @@ -19,11 +19,16 @@ package rm import ( "fmt" "strings" + "sync" "testing" + "time" "github.com/NVIDIA/go-nvml/pkg/nvml" "github.com/stretchr/testify/require" pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + "k8s.io/utils/ptr" + + spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) func TestNewHealthCheckXIDs(t *testing.T) { @@ -409,3 +414,577 @@ func TestGetMigDeviceParts(t *testing.T) { }) } } + +// fakeHealthEventSet is a test double for nvml.EventSet that yields mock events sequentially. +type fakeHealthEventSet struct { + nvml.EventSet + events []nvml.EventData + index int + mu sync.Mutex +} + +func (f *fakeHealthEventSet) Wait(timeout uint32) (nvml.EventData, nvml.Return) { + f.mu.Lock() + defer f.mu.Unlock() + if f.index < len(f.events) { + e := f.events[f.index] + f.index++ + return e, nvml.SUCCESS + } + time.Sleep(10 * time.Millisecond) + return nvml.EventData{}, nvml.ERROR_TIMEOUT +} + +func (f *fakeHealthEventSet) Free() nvml.Return { + return nvml.SUCCESS +} + +// fakeHealthDevice is a test double for nvml.Device supporting UUID, supported events, and MIG hierarchy. +type fakeHealthDevice struct { + nvml.Device + uuid string + parentUUID string + gi int + ci int + supportedEvents uint64 + registerCallCount int + mu sync.Mutex +} + +func (f *fakeHealthDevice) GetUUID() (string, nvml.Return) { + return f.uuid, nvml.SUCCESS +} + +func (f *fakeHealthDevice) GetSupportedEventTypes() (uint64, nvml.Return) { + events := f.supportedEvents + if events == 0 { + events = uint64(nvml.EventTypeXidCriticalError | nvml.EventTypeDoubleBitEccError | nvml.EventTypeSingleBitEccError) + } + return events, nvml.SUCCESS +} + +func (f *fakeHealthDevice) RegisterEvents(mask uint64, set nvml.EventSet) nvml.Return { + f.mu.Lock() + defer f.mu.Unlock() + f.registerCallCount++ + return nvml.SUCCESS +} + +func (f *fakeHealthDevice) GetDeviceHandleFromMigDeviceHandle() (nvml.Device, nvml.Return) { + return &fakeHealthDevice{uuid: f.parentUUID}, nvml.SUCCESS +} + +func (f *fakeHealthDevice) GetGpuInstanceId() (int, nvml.Return) { + return f.gi, nvml.SUCCESS +} + +func (f *fakeHealthDevice) GetComputeInstanceId() (int, nvml.Return) { + return f.ci, nvml.SUCCESS +} + +// fakeHealthNvmlLib is a test double for nvml.Interface providing handle lookup and event set creation. +type fakeHealthNvmlLib struct { + nvml.Interface + devices map[string]nvml.Device + eventSet nvml.EventSet +} + +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(uuid string) (nvml.Device, nvml.Return) { + if d, ok := f.devices[uuid]; ok { + return d, nvml.SUCCESS + } + return nil, nvml.ERROR_NOT_FOUND +} + +func TestCheckHealthReplicatedDevicesAllMarkedUnhealthy(t *testing.T) { + parentUUID := "GPU-12345678-1234-1234-1234-123456789abc" + dev0 := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::0", Health: pluginapi.Healthy}, + Index: "0", + } + dev1 := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::1", Health: pluginapi.Healthy}, + Index: "0", + } + dev2 := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::2", Health: pluginapi.Healthy}, + Index: "0", + } + dev3 := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::3", Health: pluginapi.Healthy}, + Index: "0", + } + + mockDevice := &fakeHealthDevice{uuid: parentUUID} + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockDevice, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, // Critical Xid error + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockDevice, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + devices := Devices{ + dev0.ID: dev0, + dev1.ID: dev1, + dev2.ID: dev2, + dev3.ID: dev3, + } + + go func() { + _ = r.checkHealth(stop, devices, unhealthy) + }() + + var received []string + timeout := time.After(2 * time.Second) + for len(received) < 4 { + select { + case d := <-unhealthy: + received = append(received, d.ID) + case <-timeout: + t.Fatalf("timed out waiting for unhealthy devices, received %d of 4: %v", len(received), received) + } + } + + require.ElementsMatch(t, []string{dev0.ID, dev1.ID, dev2.ID, dev3.ID}, received) + require.Equal(t, 1, mockDevice.registerCallCount, "RegisterEvents should be called exactly once per unique parent UUID") +} + +func TestCheckHealthMultiplePhysicalGPUs(t *testing.T) { + parentUUID1 := "GPU-aaaa-1111" + parentUUID2 := "GPU-bbbb-2222" + + dev1a := &Device{ + Device: pluginapi.Device{ID: parentUUID1 + "::0", Health: pluginapi.Healthy}, + Index: "0", + } + dev1b := &Device{ + Device: pluginapi.Device{ID: parentUUID1 + "::1", Health: pluginapi.Healthy}, + Index: "0", + } + dev2a := &Device{ + Device: pluginapi.Device{ID: parentUUID2 + "::0", Health: pluginapi.Healthy}, + Index: "1", + } + dev2b := &Device{ + Device: pluginapi.Device{ID: parentUUID2 + "::1", Health: pluginapi.Healthy}, + Index: "1", + } + + mockDevice1 := &fakeHealthDevice{uuid: parentUUID1} + mockDevice2 := &fakeHealthDevice{uuid: parentUUID2} + + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockDevice1, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID1: mockDevice1, + parentUUID2: mockDevice2, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + devices := Devices{ + dev1a.ID: dev1a, + dev1b.ID: dev1b, + dev2a.ID: dev2a, + dev2b.ID: dev2b, + } + + go func() { + _ = r.checkHealth(stop, devices, unhealthy) + }() + + var received []string + timeout := time.After(2 * time.Second) + for len(received) < 2 { + select { + case d := <-unhealthy: + received = append(received, d.ID) + case <-timeout: + t.Fatalf("timed out waiting for unhealthy devices, received %d of 2: %v", len(received), received) + } + } + + require.ElementsMatch(t, []string{dev1a.ID, dev1b.ID}, received) + + // Ensure no unexpected devices from GPU-2 are sent + select { + case unexpected := <-unhealthy: + t.Fatalf("unexpected device received on unhealthy channel: %v", unexpected.ID) + case <-time.After(100 * time.Millisecond): + // Success + } +} + +func TestCheckHealthMigPlacementAndGlobalError(t *testing.T) { + parentUUID := "GPU-MIG-PARENT-1234" + mig1 := &Device{ + Device: pluginapi.Device{ID: "MIG-GPU-MIG-PARENT-1234/1/0", Health: pluginapi.Healthy}, + Index: "0:0", + } + mig2 := &Device{ + Device: pluginapi.Device{ID: "MIG-GPU-MIG-PARENT-1234/2/0", Health: pluginapi.Healthy}, + Index: "0:1", + } + + mockParent := &fakeHealthDevice{uuid: parentUUID} + + t.Run("GI-specific error marks only affected MIG instance unhealthy", func(t *testing.T) { + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockParent, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: 1, + ComputeInstanceId: 0, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockParent, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + devices := Devices{ + mig1.ID: mig1, + mig2.ID: mig2, + } + + go func() { + _ = r.checkHealth(stop, devices, unhealthy) + }() + + select { + case d := <-unhealthy: + require.Equal(t, mig1.ID, d.ID) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for unhealthy MIG device") + } + + // Ensure mig2 was not marked unhealthy + select { + case unexpected := <-unhealthy: + t.Fatalf("unexpected device marked unhealthy: %v", unexpected.ID) + case <-time.After(100 * time.Millisecond): + } + }) + + t.Run("Global GPU error marks all MIG instances on parent GPU unhealthy", func(t *testing.T) { + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockParent, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockParent, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + devices := Devices{ + mig1.ID: mig1, + mig2.ID: mig2, + } + + go func() { + _ = r.checkHealth(stop, devices, unhealthy) + }() + + var received []string + timeout := time.After(2 * time.Second) + for len(received) < 2 { + select { + case d := <-unhealthy: + received = append(received, d.ID) + case <-timeout: + t.Fatalf("timed out waiting for unhealthy MIG devices, received %d of 2: %v", len(received), received) + } + } + + require.ElementsMatch(t, []string{mig1.ID, mig2.ID}, received) + }) +} + +func TestCheckHealthIgnoredXids(t *testing.T) { + parentUUID := "GPU-IGNORED-TEST" + dev := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::0", Health: pluginapi.Healthy}, + Index: "0", + } + + mockDevice := &fakeHealthDevice{uuid: parentUUID} + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockDevice, + EventType: nvml.EventTypeXidCriticalError, + EventData: 31, // Application error (memory page fault) - default ignored + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + { + Device: mockDevice, + EventType: nvml.EventTypeXidCriticalError, + EventData: 43, // Application error (GPU stopped processing) - default ignored + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockDevice, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + go func() { + _ = r.checkHealth(stop, Devices{dev.ID: dev}, unhealthy) + }() + + select { + case unexpected := <-unhealthy: + t.Fatalf("ignored Xid should not mark device unhealthy: %v", unexpected.ID) + case <-time.After(200 * time.Millisecond): + // Success + } +} + +func TestCheckHealthUnexpectedDevice(t *testing.T) { + parentUUID := "GPU-EXPECTED" + dev := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::0", Health: pluginapi.Healthy}, + Index: "0", + } + + mockDevice := &fakeHealthDevice{uuid: parentUUID} + mockUnknownDevice := &fakeHealthDevice{uuid: "GPU-UNKNOWN"} + + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockUnknownDevice, + EventType: nvml.EventTypeXidCriticalError, + EventData: 79, + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockDevice, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + go func() { + _ = r.checkHealth(stop, Devices{dev.ID: dev}, unhealthy) + }() + + select { + case unexpected := <-unhealthy: + t.Fatalf("unexpected device should not mark device unhealthy: %v", unexpected.ID) + case <-time.After(200 * time.Millisecond): + // Success + } +} + +func TestCheckHealthNonCriticalEventType(t *testing.T) { + parentUUID := "GPU-NON-CRITICAL" + dev := &Device{ + Device: pluginapi.Device{ID: parentUUID + "::0", Health: pluginapi.Healthy}, + Index: "0", + } + + mockDevice := &fakeHealthDevice{uuid: parentUUID} + eventSet := &fakeHealthEventSet{ + events: []nvml.EventData{ + { + Device: mockDevice, + EventType: nvml.EventTypeSingleBitEccError, + EventData: 0, + GpuInstanceId: 0xFFFFFFFF, + ComputeInstanceId: 0xFFFFFFFF, + }, + }, + } + + nvmllib := &fakeHealthNvmlLib{ + devices: map[string]nvml.Device{ + parentUUID: mockDevice, + }, + eventSet: eventSet, + } + + r := &nvmlResourceManager{ + resourceManager: resourceManager{ + config: &spec.Config{ + Flags: spec.Flags{ + CommandLineFlags: spec.CommandLineFlags{ + FailOnInitError: ptr.To(true), + }, + }, + }, + }, + nvml: nvmllib, + } + + stop := make(chan interface{}) + unhealthy := make(chan *Device, 10) + defer close(stop) + + go func() { + _ = r.checkHealth(stop, Devices{dev.ID: dev}, unhealthy) + }() + + select { + case unexpected := <-unhealthy: + t.Fatalf("non-critical event type should not mark device unhealthy: %v", unexpected.ID) + case <-time.After(200 * time.Millisecond): + // Success + } +}