feat(dcgm-exporter): naive implementation of adding extra annotations to the daemonset and it's underlying pods

Author: MadJlzz

controllers/object_controls.go

@@ -1755,6 +1755,12 @@ func TransformDCGMExporter(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpe
 	if len(config.DCGMExporter.ImagePullSecrets) > 0 {
 		addPullSecrets(&obj.Spec.Template.Spec, config.DCGMExporter.ImagePullSecrets)
 	}
+
+	// set extra annotations for the DS
+	if len(config.DCGMExporter.Annotations) > 0 {
+		addExtraAnnotations(obj, config.DCGMExporter.Annotations)
+	}
+
 	// set resource limits
 	if config.DCGMExporter.Resources != nil {
 		// apply resource limits to all containers
@@ -1851,6 +1857,19 @@ func TransformDCGMExporter(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpe
 	return nil
 }
 
+func addExtraAnnotations(obj *appsv1.DaemonSet, annotations map[string]string) {
+	if obj.Annotations == nil {
+		obj.Annotations = make(map[string]string)
+	}
+	if obj.Spec.Template.Annotations == nil {
+		obj.Spec.Template.Annotations = make(map[string]string)
+	}
+	for k, v := range annotations {
+		obj.Annotations[k] = v
+		obj.Spec.Template.Annotations[k] = v
+	}
+}
+
 // TransformDCGM transforms dcgm daemonset with required config as per ClusterPolicy
 func TransformDCGM(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpec, n ClusterPolicyController) error {
 	// update validation container

controllers/transforms_test.go

@@ -88,6 +88,11 @@ func NewDaemonset() Daemonset {
 	return Daemonset{ds}
 }
 
+func (d Daemonset) WithAnnotations(annotations map[string]string) Daemonset {
+	d.Annotations = annotations
+	return d
+}
+
 func (d Daemonset) WithHostPathVolume(name string, path string, hostPathType *corev1.HostPathType) Daemonset {
 	volume := corev1.Volume{
 		Name: name,
@@ -1512,6 +1517,78 @@ func TestTransformDCGMExporter(t *testing.T) {
 				WithRuntimeClassName("nvidia").
 				WithHostPathVolume("hpc-job-mapping", "/run/nvidia/dcgm-job-mapping", ptr.To(corev1.HostPathDirectoryOrCreate)),
 		},
+		{
+			description: "transform dcgm exporter with extra annotations",
+			ds: NewDaemonset().
+				WithContainer(corev1.Container{Name: "dcgm-exporter"}),
+			cpSpec: &gpuv1.ClusterPolicySpec{
+				DCGMExporter: gpuv1.DCGMExporterSpec{
+					Repository: "nvcr.io/nvidia/k8s",
+					Image:      "dcgm-exporter",
+					Version:    "v1.0.0",
+					Annotations: map[string]string{
+						"prometheus.io/scrape": "true",
+						"prometheus.io/port":   "8080",
+						"prometheus.io/path":   "/metrics",
+					},
+				},
+			},
+			expectedDs: NewDaemonset().
+				WithAnnotations(map[string]string{
+					"prometheus.io/scrape": "true",
+					"prometheus.io/port":   "8080",
+					"prometheus.io/path":   "/metrics",
+				}).
+				WithPodAnnotations(map[string]string{
+					"prometheus.io/scrape": "true",
+					"prometheus.io/port":   "8080",
+					"prometheus.io/path":   "/metrics",
+				}).
+				WithContainer(corev1.Container{
+					Name:            "dcgm-exporter",
+					Image:           "nvcr.io/nvidia/k8s/dcgm-exporter:v1.0.0",
+					ImagePullPolicy: corev1.PullIfNotPresent,
+					Env: []corev1.EnvVar{
+						{Name: "DCGM_REMOTE_HOSTENGINE_INFO", Value: "nvidia-dcgm:5555"},
+					},
+				}).
+				WithRuntimeClassName("nvidia"),
+		},
+		{
+			description: "transform dcgm exporter appends extra annotations to existing ones",
+			ds: NewDaemonset().
+				WithAnnotations(map[string]string{
+					"existing-key": "existing-value",
+				}).
+				WithContainer(corev1.Container{Name: "dcgm-exporter"}),
+			cpSpec: &gpuv1.ClusterPolicySpec{
+				DCGMExporter: gpuv1.DCGMExporterSpec{
+					Repository: "nvcr.io/nvidia/k8s",
+					Image:      "dcgm-exporter",
+					Version:    "v1.0.0",
+					Annotations: map[string]string{
+						"prometheus.io/scrape": "true",
+					},
+				},
+			},
+			expectedDs: NewDaemonset().
+				WithAnnotations(map[string]string{
+					"existing-key":         "existing-value",
+					"prometheus.io/scrape": "true",
+				}).
+				WithPodAnnotations(map[string]string{
+					"prometheus.io/scrape": "true",
+				}).
+				WithContainer(corev1.Container{
+					Name:            "dcgm-exporter",
+					Image:           "nvcr.io/nvidia/k8s/dcgm-exporter:v1.0.0",
+					ImagePullPolicy: corev1.PullIfNotPresent,
+					Env: []corev1.EnvVar{
+						{Name: "DCGM_REMOTE_HOSTENGINE_INFO", Value: "nvidia-dcgm:5555"},
+					},
+				}).
+				WithRuntimeClassName("nvidia"),
+		},
 	}
 
 	for _, tc := range testCases {