diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index 4640af7a71..5e04f3f833 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -454,7 +454,11 @@ func (plugin *NvidiaDevicePlugin) GetDevicePluginOptions(context.Context, *kubel // ListAndWatch lists devices and update that list according to the health status func (plugin *NvidiaDevicePlugin) ListAndWatch(e *kubeletdevicepluginv1beta1.Empty, s kubeletdevicepluginv1beta1.DevicePlugin_ListAndWatchServer) error { - s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()}) + err := s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()}) + if err != nil { + klog.Errorf("Failed to send ListAndWatch response: %v", err) + return err + } for { select { @@ -464,7 +468,10 @@ func (plugin *NvidiaDevicePlugin) ListAndWatch(e *kubeletdevicepluginv1beta1.Emp // 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) - s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()}) + if err := s.Send(&kubeletdevicepluginv1beta1.ListAndWatchResponse{Devices: plugin.apiDevices()}); err != nil { + klog.Errorf("Failed to send health-update ListAndWatch response: %v", err) + return nil + } } } } diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index b69e7cafac..07c511c17b 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -48,6 +48,7 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/util" "github.com/stretchr/testify/require" + "google.golang.org/grpc" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" @@ -1308,3 +1309,66 @@ func TestMigCurrentConfigsNeverNil(t *testing.T) { require.False(t, cfg.MigEnabled) } } + +type mockListAndWatchServer struct { + grpc.ServerStream + sendErrs []error + sent []*kubeletdevicepluginv1beta1.ListAndWatchResponse +} + +func (m *mockListAndWatchServer) Send(resp *kubeletdevicepluginv1beta1.ListAndWatchResponse) error { + m.sent = append(m.sent, resp) + if len(m.sendErrs) > 0 { + err := m.sendErrs[0] + m.sendErrs = m.sendErrs[1:] + return err + } + return nil +} + +func TestListAndWatch_SendError(t *testing.T) { + mockRM := &rm.ResourceManagerMock{ + DevicesFunc: func() rm.Devices { + return rm.Devices{} + }, + ResourceFunc: func() v1.ResourceName { + return v1.ResourceName("nvidia.com/gpu") + }, + } + + t.Run("InitialSendFails", func(t *testing.T) { + expectedErr := fmt.Errorf("initial send failed") + server := &mockListAndWatchServer{ + sendErrs: []error{expectedErr}, + } + plugin := &NvidiaDevicePlugin{ + rm: mockRM, + stop: make(chan any), + health: make(chan *rm.Device, 1), + schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}}, + } + + err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server) + require.ErrorIs(t, err, expectedErr) + require.Len(t, server.sent, 1) + }) + + t.Run("UpdateSendFails", func(t *testing.T) { + expectedErr := fmt.Errorf("update send failed") + server := &mockListAndWatchServer{ + sendErrs: []error{nil, expectedErr}, + } + plugin := &NvidiaDevicePlugin{ + rm: mockRM, + stop: make(chan any), + health: make(chan *rm.Device, 1), + schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}}, + } + + plugin.health <- &rm.Device{Device: kubeletdevicepluginv1beta1.Device{ID: "gpu-1"}} + err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server) + require.NoError(t, err) + require.Len(t, server.sent, 2) + }) +} +