feat: add affinity and nodeSelector to common daemonset config

api/nvidia/v1/clusterpolicy_types.go

@@ -309,6 +309,18 @@ type DaemonsetsSpec struct {
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.x-descriptors="urn:alm:descriptor:com.tectonic.ui:advanced,urn:alm:descriptor:io.kubernetes:Tolerations"
 	Tolerations []corev1.Toleration `json:"tolerations,omitempty"`
 
+	// Optional: Set nodeSelector merged with operand nodeSelector labels
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="NodeSelector for all Daemonsets"
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.x-descriptors="urn:alm:descriptor:com.tectonic.ui:advanced,urn:alm:descriptor:com.tectonic.ui:nodeSelector"
+	NodeSelector map[string]string `json:"nodeSelector,omitempty"`
+
+	// Optional: Set affinity for all Daemonsets
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="Affinity for all Daemonsets"
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.x-descriptors="urn:alm:descriptor:com.tectonic.ui:advanced,urn:alm:descriptor:com.tectonic.ui:affinity"
+	Affinity *corev1.Affinity `json:"affinity,omitempty"`
+
 	// +kubebuilder:validation:Optional
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="PriorityClassName"

controllers/object_controls.go

@@ -829,13 +829,129 @@ func applyCommonDaemonsetConfig(obj *appsv1.DaemonSet, config *gpuv1.ClusterPoli
 		obj.Spec.Template.Spec.Tolerations = config.Daemonsets.Tolerations
 	}
 
+	if len(config.Daemonsets.NodeSelector) > 0 {
+		if obj.Spec.Template.Spec.NodeSelector == nil {
+			obj.Spec.Template.Spec.NodeSelector = make(map[string]string)
+		}
+		for key, value := range config.Daemonsets.NodeSelector {
+			if _, exists := obj.Spec.Template.Spec.NodeSelector[key]; exists {
+				continue
+			}
+			obj.Spec.Template.Spec.NodeSelector[key] = value
+		}
+	}
+
+	if config.Daemonsets.Affinity != nil {
+		obj.Spec.Template.Spec.Affinity = mergeDaemonsetAffinity(obj.Spec.Template.Spec.Affinity, config.Daemonsets.Affinity)
+	}
+
 	// set pod-level security context if specified (applies as defaults to all containers in the pod)
 	if config.Daemonsets.PodSecurityContext != nil {
 		obj.Spec.Template.Spec.SecurityContext = config.Daemonsets.PodSecurityContext
 	}
 	return nil
 }
 
+func mergeDaemonsetAffinity(dst, src *corev1.Affinity) *corev1.Affinity {
+	if src == nil {
+		return dst
+	}
+	if dst == nil {
+		return src.DeepCopy()
+	}
+	merged := dst.DeepCopy()
+	if src.NodeAffinity != nil {
+		if merged.NodeAffinity == nil {
+			merged.NodeAffinity = src.NodeAffinity.DeepCopy()
+		} else {
+			merged.NodeAffinity = mergeNodeAffinity(merged.NodeAffinity, src.NodeAffinity)
+		}
+	}
+	if src.PodAffinity != nil {
+		if merged.PodAffinity == nil {
+			merged.PodAffinity = src.PodAffinity.DeepCopy()
+		} else {
+			if src.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution != nil {
+				merged.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+					merged.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+					src.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution...,
+				)
+			}
+			if src.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution != nil {
+				merged.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+					merged.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+					src.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution...,
+				)
+			}
+		}
+	}
+	if src.PodAntiAffinity != nil {
+		if merged.PodAntiAffinity == nil {
+			merged.PodAntiAffinity = src.PodAntiAffinity.DeepCopy()
+		} else {
+			if src.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution != nil {
+				merged.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+					merged.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+					src.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution...,
+				)
+			}
+			if src.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution != nil {
+				merged.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+					merged.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+					src.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution...,
+				)
+			}
+		}
+	}
+	return merged
+}
+
+func mergeNodeAffinity(dst, src *corev1.NodeAffinity) *corev1.NodeAffinity {
+	if src == nil {
+		return dst
+	}
+	if dst == nil {
+		return src.DeepCopy()
+	}
+	merged := dst.DeepCopy()
+	if src.RequiredDuringSchedulingIgnoredDuringExecution != nil {
+		if merged.RequiredDuringSchedulingIgnoredDuringExecution == nil {
+			merged.RequiredDuringSchedulingIgnoredDuringExecution = src.RequiredDuringSchedulingIgnoredDuringExecution.DeepCopy()
+		} else {
+			merged.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = mergeNodeSelectorTerms(
+				merged.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
+				src.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
+			)
+		}
+	}
+	if src.PreferredDuringSchedulingIgnoredDuringExecution != nil {
+		merged.PreferredDuringSchedulingIgnoredDuringExecution = append(
+			merged.PreferredDuringSchedulingIgnoredDuringExecution,
+			src.PreferredDuringSchedulingIgnoredDuringExecution...,
+		)
+	}
+	return merged
+}
+
+func mergeNodeSelectorTerms(dst, src []corev1.NodeSelectorTerm) []corev1.NodeSelectorTerm {
+	if len(dst) == 0 {
+		return append([]corev1.NodeSelectorTerm(nil), src...)
+	}
+	if len(src) == 0 {
+		return dst
+	}
+	merged := make([]corev1.NodeSelectorTerm, 0, len(dst)*len(src))
+	for _, dstTerm := range dst {
+		for _, srcTerm := range src {
+			merged = append(merged, corev1.NodeSelectorTerm{
+				MatchExpressions: append(append([]corev1.NodeSelectorRequirement(nil), dstTerm.MatchExpressions...), srcTerm.MatchExpressions...),
+				MatchFields:      append(append([]corev1.NodeSelectorRequirement(nil), dstTerm.MatchFields...), srcTerm.MatchFields...),
+			})
+		}
+	}
+	return merged
+}
+
 // apply necessary transforms if a custom host root path is configured
 func transformForHostRoot(obj *appsv1.DaemonSet, hostRoot string) {
 	if hostRoot == "" || hostRoot == "/" {

controllers/transforms_test.go

@@ -149,6 +149,16 @@ func (d Daemonset) WithTolerations(tolerations []corev1.Toleration) Daemonset {
 	return d
 }
 
+func (d Daemonset) WithNodeSelector(nodeSelector map[string]string) Daemonset {
+	d.Spec.Template.Spec.NodeSelector = nodeSelector
+	return d
+}
+
+func (d Daemonset) WithAffinity(affinity *corev1.Affinity) Daemonset {
+	d.Spec.Template.Spec.Affinity = affinity
+	return d
+}
+
 func (d Daemonset) WithPodSecurityContext(psc *corev1.PodSecurityContext) Daemonset {
 	d.Spec.Template.Spec.SecurityContext = psc
 	return d
@@ -621,7 +631,8 @@ func TestApplyUpdateStrategyConfig(t *testing.T) {
 				UpdateStrategy: "RollingUpdate",
 				RollingUpdate: &gpuv1.RollingUpdateSpec{
 					MaxUnavailable: "10%abc",
-				}},
+				},
+			},
 			errorExpected: true,
 		},
 		{
@@ -707,14 +718,163 @@ func TestApplyCommonDaemonSetConfig(t *testing.T) {
 				},
 			}),
 		},
+		{
+			description: "nodeSelector configured",
+			ds: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/gpu.deploy.device-plugin": "true",
+			}),
+			dsSpec: gpuv1.DaemonsetsSpec{
+				NodeSelector: map[string]string{
+					"karpenter.sh/nodepool": "gpu",
+				},
+			},
+			expectedDs: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/gpu.deploy.device-plugin": "true",
+				"karpenter.sh/nodepool":               "gpu",
+			}),
+		},
+		{
+			description: "existing nodeSelector keys are not overridden",
+			ds: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/gpu.deploy.device-plugin": "true",
+			}),
+			dsSpec: gpuv1.DaemonsetsSpec{
+				NodeSelector: map[string]string{
+					"nvidia.com/gpu.deploy.device-plugin": "false",
+					"karpenter.sh/nodepool":               "gpu",
+				},
+			},
+			expectedDs: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/gpu.deploy.device-plugin": "true",
+				"karpenter.sh/nodepool":               "gpu",
+			}),
+		},
+		{
+			description: "operand nodeSelector keys such as mps.capable are not overridden",
+			ds: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/mps.capable": "true",
+			}),
+			dsSpec: gpuv1.DaemonsetsSpec{
+				NodeSelector: map[string]string{
+					"nvidia.com/mps.capable": "false",
+					"karpenter.sh/nodepool":  "gpu",
+				},
+			},
+			expectedDs: NewDaemonset().WithNodeSelector(map[string]string{
+				"nvidia.com/mps.capable": "true",
+				"karpenter.sh/nodepool":  "gpu",
+			}),
+		},
+		{
+			description: "affinity configured",
+			ds:          NewDaemonset(),
+			dsSpec: gpuv1.DaemonsetsSpec{
+				Affinity: &corev1.Affinity{
+					NodeAffinity: &corev1.NodeAffinity{
+						RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+							NodeSelectorTerms: []corev1.NodeSelectorTerm{
+								{
+									MatchExpressions: []corev1.NodeSelectorRequirement{
+										{
+											Key:      "karpenter.sh/nodepool",
+											Operator: corev1.NodeSelectorOpIn,
+											Values:   []string{"gpu"},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedDs: NewDaemonset().WithAffinity(&corev1.Affinity{
+				NodeAffinity: &corev1.NodeAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+						NodeSelectorTerms: []corev1.NodeSelectorTerm{
+							{
+								MatchExpressions: []corev1.NodeSelectorRequirement{
+									{
+										Key:      "karpenter.sh/nodepool",
+										Operator: corev1.NodeSelectorOpIn,
+										Values:   []string{"gpu"},
+									},
+								},
+							},
+						},
+					},
+				},
+			}),
+		},
+		{
+			description: "node affinity merged additively with operand",
+			ds: NewDaemonset().WithAffinity(&corev1.Affinity{
+				NodeAffinity: &corev1.NodeAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+						NodeSelectorTerms: []corev1.NodeSelectorTerm{
+							{
+								MatchExpressions: []corev1.NodeSelectorRequirement{
+									{
+										Key:      "nvidia.com/gpu.deploy.device-plugin",
+										Operator: corev1.NodeSelectorOpIn,
+										Values:   []string{"true"},
+									},
+								},
+							},
+						},
+					},
+				},
+			}),
+			dsSpec: gpuv1.DaemonsetsSpec{
+				Affinity: &corev1.Affinity{
+					NodeAffinity: &corev1.NodeAffinity{
+						RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+							NodeSelectorTerms: []corev1.NodeSelectorTerm{
+								{
+									MatchExpressions: []corev1.NodeSelectorRequirement{
+										{
+											Key:      "karpenter.sh/nodepool",
+											Operator: corev1.NodeSelectorOpIn,
+											Values:   []string{"gpu"},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expectedDs: NewDaemonset().WithAffinity(&corev1.Affinity{
+				NodeAffinity: &corev1.NodeAffinity{
+					RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+						NodeSelectorTerms: []corev1.NodeSelectorTerm{
+							{
+								MatchExpressions: []corev1.NodeSelectorRequirement{
+									{
+										Key:      "nvidia.com/gpu.deploy.device-plugin",
+										Operator: corev1.NodeSelectorOpIn,
+										Values:   []string{"true"},
+									},
+									{
+										Key:      "karpenter.sh/nodepool",
+										Operator: corev1.NodeSelectorOpIn,
+										Values:   []string{"gpu"},
+									},
+								},
+							},
+						},
+					},
+				},
+			}),
+		},
 		{
 			description: "invalid updatestrategy configured",
 			ds:          NewDaemonset(),
 			dsSpec: gpuv1.DaemonsetsSpec{
 				UpdateStrategy: "RollingUpdate",
 				RollingUpdate: &gpuv1.RollingUpdateSpec{
 					MaxUnavailable: "10%abc",
-				}},
+				},
+			},
 			errorExpected: true,
 		},
 		{

deployments/gpu-operator/crds/nvidia.com_clusterpolicies.yaml

@@ -159,6 +159,10 @@ spec:
               daemonsets:
                 description: Daemonset defines common configuration for all Daemonsets
                 properties:
+                  affinity:
+                    description: 'Optional: Set affinity for all Daemonsets'
+                    type: object
+                    x-kubernetes-preserve-unknown-fields: true
                   annotations:
                     additionalProperties:
                       type: string
@@ -175,6 +179,12 @@ spec:
                       (scope and select) objects. May match selectors of replication controllers
                       and services.
                     type: object
+                  nodeSelector:
+                    additionalProperties:
+                      type: string
+                    description: 'Optional: Set nodeSelector merged with operand nodeSelector
+                      labels'
+                    type: object
                   podSecurityContext:
                     description: 'Optional: Set pod-level security context for all
                       DaemonSet pods (applies as defaults to all containers)'