Add pod labels and annotations for DCGMExporter

Signed-off-by: maciej-tatarski <mta@corti.ai>

Author: maciej-tatarski

api/nvidia/v1/clusterpolicy_types.go

@@ -884,6 +884,16 @@ type DCGMExporterSpec struct {
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.x-descriptors="urn:alm:descriptor:io.kubernetes:Secret"
 	ImagePullSecrets []string `json:"imagePullSecrets,omitempty"`
 
+	// Optional Pod Annotations for the DCGM Exporter deployment
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="Annotations for the DCGM Exporter deployment"
+	PodAnnotations map[string]string `json:"podAnnotations,omitempty"`
+
+	// Optional Pod Labels for the DCGM Exporter deployment
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
+	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="Labels for the DCGM Exporter deployment"
+	PodLabels map[string]string `json:"podLabels,omitempty"`
+
 	// Optional: Define resources requests and limits for each pod
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors=true
 	// +operator-sdk:gen-csv:customresourcedefinitions.specDescriptors.displayName="Resource Requirements"

controllers/object_controls.go

@@ -23,13 +23,12 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"path/filepath"
 	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 
-	"path/filepath"
-
 	apiconfigv1 "github.com/openshift/api/config/v1"
 	apiimagev1 "github.com/openshift/api/image/v1"
 	secv1 "github.com/openshift/api/security/v1"
@@ -1621,6 +1620,26 @@ func TransformDCGMExporter(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpe
 		obj.Spec.Template.Spec.Containers[0].Args = config.DCGMExporter.Args
 	}
 
+	// set labels if specified for exporter container
+	if obj.Spec.Template.ObjectMeta.Labels == nil {
+		obj.Spec.Template.ObjectMeta.Labels = make(map[string]string)
+	}
+	for labelKey, labelValue := range config.DCGMExporter.PodLabels {
+		// disallow setting app label
+		if labelKey == "app" {
+			continue
+		}
+		obj.Spec.Template.ObjectMeta.Labels[labelKey] = labelValue
+	}
+
+	// set annotations if specified for exporter container
+	if obj.Spec.Template.ObjectMeta.Annotations == nil {
+		obj.Spec.Template.ObjectMeta.Annotations = make(map[string]string)
+	}
+	for annKey, annValue := range config.DCGMExporter.PodAnnotations {
+		obj.Spec.Template.ObjectMeta.Annotations[annKey] = annValue
+	}
+
 	// check if DCGM hostengine is enabled as a separate Pod and setup env accordingly
 	if config.DCGM.IsEnabled() {
 		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), DCGMRemoteEngineEnvName, fmt.Sprintf("nvidia-dcgm:%d", DCGMDefaultPort))

controllers/object_controls_test.go

@@ -974,6 +974,14 @@ func getDCGMExporterTestInput(testCase string) *gpuv1.ClusterPolicy {
 	case "standalone-dcgm":
 		dcgmEnabled := true
 		cp.Spec.DCGM.Enabled = &dcgmEnabled
+	case "podLabelsAndAnnotations":
+		cp.Spec.DCGMExporter.PodAnnotations = map[string]string{
+			"prometheus.io/scrape": "true",
+			"prometheus.io/port":   "9400",
+		}
+		cp.Spec.DCGMExporter.PodLabels = map[string]string{
+			"team": "test",
+		}
 	default:
 		return nil
 	}
@@ -989,15 +997,25 @@ func getDCGMExporterTestOutput(testCase string) map[string]interface{} {
 		"numDaemonsets":     1,
 		"dcgmExporterImage": "nvcr.io/nvidia/k8s/dcgm-exporter:3.3.0-3.2.0-ubuntu22.04",
 		"imagePullSecret":   "ngc-secret",
+		"env":               map[string]string{},
+		"podLabels":         map[string]string{},
+		"podAnnotations":    map[string]string{},
 	}
 
 	switch testCase {
 	case "default":
-		output["env"] = map[string]string{}
 	case "standalone-dcgm":
 		output["env"] = map[string]string{
 			"DCGM_REMOTE_HOSTENGINE_INFO": "nvidia-dcgm:5555",
 		}
+	case "podLabelsAndAnnotations":
+		output["podLabels"] = map[string]string{
+			"team": "test",
+		}
+		output["podAnnotations"] = map[string]string{
+			"prometheus.io/scrape": "true",
+			"prometheus.io/port":   "9400",
+		}
 	default:
 		return nil
 	}
@@ -1023,6 +1041,11 @@ func TestDCGMExporter(t *testing.T) {
 			getDCGMExporterTestInput("standalone-dcgm"),
 			getDCGMExporterTestOutput("standalone-dcgm"),
 		},
+		{
+			"podLabelsAndAnnotations",
+			getDCGMExporterTestInput("podLabelsAndAnnotations"),
+			getDCGMExporterTestOutput("podLabelsAndAnnotations"),
+		},
 	}
 
 	for _, tc := range testCases {
@@ -1053,6 +1076,18 @@ func TestDCGMExporter(t *testing.T) {
 					t.Fatalf("Expected env is not set for daemonset nvidia-dcgm-exporter %s->%s", key, value)
 				}
 			}
+			for key, value := range tc.output["podLabels"].(map[string]string) {
+				objLabelvalue, ok := ds.Spec.Template.Labels[key]
+				if !ok || objLabelvalue != value {
+					t.Fatalf("Expected label is not set for daemonset nvidia-dcgm-exporter %s->%s", key, value)
+				}
+			}
+			for key, value := range tc.output["podAnnotations"].(map[string]string) {
+				objAnnotationValue, ok := ds.Spec.Template.Annotations[key]
+				if !ok || objAnnotationValue != value {
+					t.Fatalf("Expected annotation is not set for daemonset nvidia-dcgm-exporter %s->%s", key, value)
+				}
+			}
 
 			require.Equal(t, tc.output["dcgmExporterImage"], dcgmExporterImage, "Unexpected configuration for dcgm-exporter image")
 

deployments/gpu-operator/templates/clusterpolicy.yaml

@@ -502,6 +502,12 @@ spec:
     {{- if .Values.dcgmExporter.resources }}
     resources: {{ toYaml .Values.dcgmExporter.resources | nindent 6 }}
     {{- end }}
+    {{- if .Values.dcgmExporter.podAnnotations }}
+    podAnnotations: {{ toYaml .Values.dcgmExporter.podAnnotations | nindent 6 }}
+    {{- end }}
+    {{- if .Values.dcgmExporter.podLabels }}
+    podLabels: {{ toYaml .Values.dcgmExporter.podLabels | nindent 6 }}
+    {{- end }}
     {{- if .Values.dcgmExporter.env }}
     env: {{ toYaml .Values.dcgmExporter.env | nindent 6 }}
     {{- end }}

deployments/gpu-operator/values.yaml

@@ -315,6 +315,8 @@ dcgmExporter:
   image: dcgm-exporter
   version: 4.2.3-4.1.3-ubuntu22.04
   imagePullPolicy: IfNotPresent
+  podAnnotations: {}
+  podLabels: {}
   env:
     - name: DCGM_EXPORTER_LISTEN
       value: ":9400"