diff --git a/pkg/addons/assets/nvidia-device-plugin.yaml b/pkg/addons/assets/nvidia-device-plugin.yaml index a6b452f078..135de8ac2b 100644 --- a/pkg/addons/assets/nvidia-device-plugin.yaml +++ b/pkg/addons/assets/nvidia-device-plugin.yaml @@ -37,6 +37,8 @@ spec: # be rescheduled after a failure. # See https://kubernetes.io/docs/tasks/administer-cluster/guaranteed-scheduling-critical-addon-pods/ priorityClassName: "system-node-critical" + nodeSelector: + nvidia.com/gpu.present: "true" containers: - image: nvcr.io/nvidia/k8s-device-plugin:v0.20.0 name: nvidia-device-plugin-ctr diff --git a/pkg/addons/device_plugin.go b/pkg/addons/device_plugin.go index 5a72fe7799..acbe5e964c 100644 --- a/pkg/addons/device_plugin.go +++ b/pkg/addons/device_plugin.go @@ -81,6 +81,7 @@ type DevicePlugin interface { Manifest() []byte SetImage(t *corev1.PodTemplateSpec) error SetTolerations(t *corev1.PodTemplateSpec) error + SetNodeSelector(t *corev1.PodTemplateSpec) error Deploy() error } @@ -108,6 +109,9 @@ func applyDevicePlugin(dp DevicePlugin) error { if err := dp.SetTolerations(&daemonSet.Spec.Template); err != nil { return fmt.Errorf("adding tolerations to device plugin daemonset: %w", err) } + if err := dp.SetNodeSelector(&daemonSet.Spec.Template); err != nil { + return fmt.Errorf("adding nodeSelector to device plugin daemonset: %w", err) + } msg, err := rawResource.CreateOrReplace(dp.PlanMode()) if err != nil { return fmt.Errorf("calling create or replace on raw device plugin daemonset: %w", err) @@ -163,6 +167,10 @@ func (n *NeuronDevicePlugin) SetTolerations(t *corev1.PodTemplateSpec) error { return nil } +func (n *NeuronDevicePlugin) SetNodeSelector(t *corev1.PodTemplateSpec) error { + return nil +} + // Deploy deploys the Neuron device plugin to the specified cluster func (n *NeuronDevicePlugin) Deploy() error { return applyDevicePlugin(n) @@ -254,6 +262,15 @@ func (n *NvidiaDevicePlugin) SetTolerations(spec *corev1.PodTemplateSpec) error return nil } +// SetNodeSelector sets a nodeSelector to ensure the DaemonSet only runs on nodes with NVIDIA GPUs. +func (n *NvidiaDevicePlugin) SetNodeSelector(spec *corev1.PodTemplateSpec) error { + if spec.Spec.NodeSelector == nil { + spec.Spec.NodeSelector = make(map[string]string) + } + spec.Spec.NodeSelector["nvidia.com/gpu.present"] = "true" + return nil +} + // A EFADevicePlugin deploys the EFA Device Plugin to a cluster type EFADevicePlugin struct { rawClient kubernetes.RawClientInterface @@ -282,6 +299,10 @@ func (n *EFADevicePlugin) SetTolerations(spec *corev1.PodTemplateSpec) error { return nil } +func (n *EFADevicePlugin) SetNodeSelector(spec *corev1.PodTemplateSpec) error { + return nil +} + // NewEFADevicePlugin creates a new EFADevicePlugin func NewEFADevicePlugin(rawClient kubernetes.RawClientInterface, region string, planMode bool, spec *api.ClusterConfig) DevicePlugin { return &EFADevicePlugin{ diff --git a/pkg/addons/device_plugin_test.go b/pkg/addons/device_plugin_test.go index 5f1bd67fc4..8be56c6b42 100644 --- a/pkg/addons/device_plugin_test.go +++ b/pkg/addons/device_plugin_test.go @@ -250,4 +250,45 @@ var _ = Describe("NvidiaDevicePlugin", func() { }) }) }) + + Describe("SetNodeSelector", func() { + var ( + plugin *addons.NvidiaDevicePlugin + spec *corev1.PodTemplateSpec + config *api.ClusterConfig + ) + + BeforeEach(func() { + spec = &corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{}, + } + config = &api.ClusterConfig{} + plugin = addons.NewNvidiaDevicePlugin(nil, "us-west-2", false, config).(*addons.NvidiaDevicePlugin) + }) + + It("should set nvidia.com/gpu.present nodeSelector", func() { + err := plugin.SetNodeSelector(spec) + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Spec.NodeSelector).To(HaveKeyWithValue("nvidia.com/gpu.present", "true")) + }) + + It("should not overwrite existing nodeSelector keys", func() { + spec.Spec.NodeSelector = map[string]string{ + "existing-key": "existing-value", + } + err := plugin.SetNodeSelector(spec) + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Spec.NodeSelector).To(HaveKeyWithValue("existing-key", "existing-value")) + Expect(spec.Spec.NodeSelector).To(HaveKeyWithValue("nvidia.com/gpu.present", "true")) + }) + + It("should overwrite the nodeSelector value if already set", func() { + spec.Spec.NodeSelector = map[string]string{ + "nvidia.com/gpu.present": "false", + } + err := plugin.SetNodeSelector(spec) + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Spec.NodeSelector).To(HaveKeyWithValue("nvidia.com/gpu.present", "true")) + }) + }) })