Add nodeSelector to nvidia-device-plugin DaemonSet

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

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{

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"))
+		})
+	})
 })