Fix device plugin recovery from Unhealthy state

pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go

@@ -461,9 +461,7 @@ func (plugin *NvidiaDevicePlugin) ListAndWatch(e *kubeletdevicepluginv1beta1.Emp
 		case <-plugin.stop:
 			return nil
 		case d := <-plugin.health:
-			// FIXME: there is no way to recover from the Unhealthy state.
-			d.Health = kubeletdevicepluginv1beta1.Unhealthy
-			klog.Infof("'%s' device marked unhealthy: %s", plugin.rm.Resource(), d.ID)
+			klog.Infof("'%s' device health changed to %s: %s", plugin.rm.Resource(), d.Health, d.ID)
 			s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()})
 		}
 	}

pkg/device-plugin/nvidiadevice/nvinternal/rm/health.go

@@ -40,6 +40,7 @@ import (
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 	"k8s.io/klog/v2"
+	kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
 )
 
 const (
@@ -56,8 +57,8 @@ const (
 	envEnableHealthChecks = "DP_ENABLE_HEALTHCHECKS"
 )
 
-// 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, disableNVML <-chan bool) error {
+// CheckHealth performs health checks on a set of devices, writing to the 'health' channel with any unhealthy devices
+func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devices, health chan<- *Device, disableNVML <-chan bool) error {
 	xids := getHealthCheckXids()
 	if xids.IsAllDisabled() {
 		return nil
@@ -123,10 +124,13 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		}
 		if ret != nvml.SUCCESS {
 			klog.Infof("Marking device %v as unhealthy: %v", d.ID, ret)
-			unhealthy <- d
+			d.Health = kubeletdevicepluginv1beta1.Unhealthy
+			health <- d
 		}
 	}
 
+	unhealthyDevices := make(map[string]*Device)
+
 	for {
 		select {
 		case <-stop:
@@ -141,12 +145,31 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 
 		e, ret := eventSet.Wait(5000)
 		if ret == nvml.ERROR_TIMEOUT {
+			for id, d := range unhealthyDevices {
+				uuid, _, _, err := r.getDevicePlacement(d)
+				if err != nil {
+					continue
+				}
+				gpu, ret := r.nvml.DeviceGetHandleByUUID(uuid)
+				if ret != nvml.SUCCESS {
+					continue
+				}
+				_, ret = gpu.GetMemoryInfo()
+				if ret == nvml.SUCCESS {
+					klog.Infof("Device %s has recovered", d.ID)
+					d.Health = kubeletdevicepluginv1beta1.Healthy
+					health <- d
+					delete(unhealthyDevices, id)
+				}
+			}
 			continue
 		}
 		if ret != nvml.SUCCESS {
 			klog.Infof("Error waiting for event: %v; Marking all devices as unhealthy", ret)
 			for _, d := range devices {
-				unhealthy <- d
+				d.Health = kubeletdevicepluginv1beta1.Unhealthy
+				unhealthyDevices[d.ID] = d
+				health <- d
 			}
 			continue
 		}
@@ -167,7 +190,9 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 			// 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
+				d.Health = kubeletdevicepluginv1beta1.Unhealthy
+				unhealthyDevices[d.ID] = d
+				health <- d
 			}
 			continue
 		}
@@ -188,7 +213,9 @@ func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devic
 		}
 
 		klog.Infof("XidCriticalError: Xid=%d on Device=%s; marking device as unhealthy.", e.EventData, d.ID)
-		unhealthy <- d
+		d.Health = kubeletdevicepluginv1beta1.Unhealthy
+		unhealthyDevices[d.ID] = d
+		health <- d
 	}
 }
 

pkg/device-plugin/nvidiadevice/nvinternal/rm/nvml_manager.go

@@ -113,12 +113,12 @@ func (r *nvmlResourceManager) GetDevicePaths(ids []string) []string {
 	return append(paths, r.Devices().Subset(ids).GetPaths()...)
 }
 
-// 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{}, unhealthy chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
+// CheckHealth performs health checks on a set of devices, writing to the 'health' channel with any unhealthy devices
+func (r *nvmlResourceManager) CheckHealth(stop <-chan interface{}, health chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
 	for {
 		// first check if disableNVML channel signal is pass close into checkHealth function
 		// if signal is pass close, return error "close signal received"
-		err := r.checkHealth(stop, r.devices, unhealthy, disableNVML)
+		err := r.checkHealth(stop, r.devices, health, disableNVML)
 		// checkHealth returns nil when health checks are disabled or on shutdown.
 		if err != nil && err.Error() == "close signal received" {
 			ackDisableHealthChecks <- true

pkg/device-plugin/nvidiadevice/nvinternal/rm/rm.go

@@ -58,7 +58,7 @@ type ResourceManager interface {
 	Devices() Devices
 	GetDevicePaths([]string) []string
 	GetPreferredAllocation(available, required []string, size int) ([]string, error)
-	CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error
+	CheckHealth(stop <-chan interface{}, health chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error
 	ValidateRequest(AnnotatedIDs) error
 }
 

pkg/device-plugin/nvidiadevice/nvinternal/rm/rm_mock.go

@@ -135,13 +135,13 @@ type ResourceManagerMock struct {
 }
 
 // CheckHealth calls CheckHealthFunc.
-func (mock *ResourceManagerMock) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
+func (mock *ResourceManagerMock) CheckHealth(stop <-chan interface{}, health chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
 	callInfo := struct {
-		Stop      <-chan interface{}
-		Unhealthy chan<- *Device
+		Stop   <-chan interface{}
+		Health chan<- *Device
 	}{
-		Stop:      stop,
-		Unhealthy: unhealthy,
+		Stop:   stop,
+		Health: health,
 	}
 	mock.lockCheckHealth.Lock()
 	mock.calls.CheckHealth = append(mock.calls.CheckHealth, callInfo)
@@ -152,7 +152,7 @@ func (mock *ResourceManagerMock) CheckHealth(stop <-chan interface{}, unhealthy
 		)
 		return errOut
 	}
-	return mock.CheckHealthFunc(stop, unhealthy)
+	return mock.CheckHealthFunc(stop, health)
 }
 
 // CheckHealthCalls gets all the calls that were made to CheckHealth.

pkg/device-plugin/nvidiadevice/nvinternal/rm/tegra_manager.go

@@ -87,6 +87,6 @@ func (r *tegraResourceManager) GetDevicePaths(ids []string) []string {
 }
 
 // CheckHealth is disabled for the tegraResourceManager
-func (r *tegraResourceManager) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
+func (r *tegraResourceManager) CheckHealth(stop <-chan interface{}, health chan<- *Device, disableNVML <-chan bool, ackDisableHealthChecks chan<- bool) error {
 	return nil
 }