feat: add support for custom service labels (kubeflow#1952)

* feat: add support for custom service labels Signed-off-by: Cian Gallagher <[email protected]> * chore: correctly format files Signed-off-by: Cian Gallagher <[email protected]> * chore: bump chart version to 1.1.30 Signed-off-by: Cian Gallagher <[email protected]> --------- Signed-off-by: Cian Gallagher <[email protected]>
spotinst · Aug 7, 2024 · 022199d · 022199d
1 parent 3c1fc60
commit 022199d
Show file tree

Hide file tree

Showing 13 changed files with 120 additions and 1 deletion.
diff --git a/charts/spark-operator-chart/Chart.yaml b/charts/spark-operator-chart/Chart.yaml
@@ -1,7 +1,7 @@
 apiVersion: v2
 name: spark-operator
 description: A Helm chart for Spark on Kubernetes operator
-version: 1.1.29
+version: 1.1.30
 appVersion: v1beta2-1.3.8-3.1.1
 keywords:
   - spark

diff --git a/charts/spark-operator-chart/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml b/charts/spark-operator-chart/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml
@@ -3750,6 +3750,10 @@ spec:
                           additionalProperties:
                             type: string
                           type: object
+                        serviceLabels:
+                          additionalProperties:
+                            type: string
+                          type: object
                         ingressAnnotations:
                           additionalProperties:
                             type: string

diff --git a/charts/spark-operator-chart/crds/sparkoperator.k8s.io_sparkapplications.yaml b/charts/spark-operator-chart/crds/sparkoperator.k8s.io_sparkapplications.yaml
@@ -3736,6 +3736,10 @@ spec:
                       additionalProperties:
                         type: string
                       type: object
+                    serviceLabels:
+                      additionalProperties:
+                        type: string
+                      type: object
                     ingressAnnotations:
                       additionalProperties:
                         type: string

diff --git a/docs/api-docs.md b/docs/api-docs.md
@@ -3120,6 +3120,18 @@ map[string]string
 </tr>
 <tr>
 <td>
+<code>serviceLabels</code><br/>
+<em>
+map[string]string
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>ServiceLables is a map of key,value pairs of labels that might be added to the service object.</p>
+</td>
+</tr>
+<tr>
+<td>
 <code>ingressAnnotations</code><br/>
 <em>
 map[string]string

diff --git a/examples/spark-pi.yaml b/examples/spark-pi.yaml
@@ -26,6 +26,9 @@ spec:
   mainClass: org.apache.spark.examples.SparkPi
   mainApplicationFile: "local:///opt/spark/examples/jars/spark-examples_2.12-3.1.1.jar"
   sparkVersion: "3.1.1"
+  sparkUIOptions:
+    serviceLabels:
+      test-label/v1: 'true'
   restartPolicy:
     type: Never
   volumes:

diff --git a/go.sum b/go.sum
@@ -2427,6 +2427,7 @@ k8s.io/gengo v0.0.0-20200428234225-8167cfdcfc14/go.mod h1:ezvh/TsK7cY6rbqRK0oQQ8
 k8s.io/gengo v0.0.0-20210813121822-485abfe95c7c/go.mod h1:FiNAH4ZV3gBg2Kwh89tzAEV2be7d5xI0vBa/VySYy3E=
 k8s.io/gengo v0.0.0-20211129171323-c02415ce4185/go.mod h1:FiNAH4ZV3gBg2Kwh89tzAEV2be7d5xI0vBa/VySYy3E=
 k8s.io/heapster v1.2.0-beta.1/go.mod h1:h1uhptVXMwC8xtZBYsPXKVi8fpdlYkTs6k949KozGrM=
+k8s.io/klog v1.0.0 h1:Pt+yjF5aB1xDSVbau4VsWe+dQNzA0qv1LlXdC2dF6Q8=
 k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 k8s.io/klog/v2 v2.0.0/go.mod h1:PBfzABfn139FHAV07az/IF9Wp1bkk3vpT2XSJ76fSDE=
 k8s.io/klog/v2 v2.2.0/go.mod h1:Od+F08eJP+W3HUb4pSrPpgp9DGU4GzlpG/TmITuYh/Y=

diff --git a/manifest/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml b/manifest/crds/sparkoperator.k8s.io_scheduledsparkapplications.yaml
@@ -3750,6 +3750,10 @@ spec:
                           additionalProperties:
                             type: string
                           type: object
+                        serviceLabels:
+                          additionalProperties:
+                            type: string
+                          type: object
                         ingressAnnotations:
                           additionalProperties:
                             type: string

diff --git a/manifest/crds/sparkoperator.k8s.io_sparkapplications.yaml b/manifest/crds/sparkoperator.k8s.io_sparkapplications.yaml
@@ -3736,6 +3736,10 @@ spec:
                       additionalProperties:
                         type: string
                       type: object
+                    serviceLabels:
+                      additionalProperties:
+                        type: string
+                      type: object
                     ingressAnnotations:
                       additionalProperties:
                         type: string

diff --git a/pkg/apis/sparkoperator.k8s.io/v1beta2/types.go b/pkg/apis/sparkoperator.k8s.io/v1beta2/types.go
@@ -317,6 +317,9 @@ type SparkUIConfiguration struct {
 	// ServiceAnnotations is a map of key,value pairs of annotations that might be added to the service object.
 	// +optional
 	ServiceAnnotations map[string]string `json:"serviceAnnotations,omitempty"`
+	// ServiceLables is a map of key,value pairs of labels that might be added to the service object.
+	// +optional
+	ServiceLabels map[string]string `json:"serviceLabels,omitempty"`
 	// IngressAnnotations is a map of key,value pairs of annotations that might be added to the ingress object. i.e. specify nginx as ingress.class
 	// +optional
 	IngressAnnotations map[string]string `json:"ingressAnnotations,omitempty"`

diff --git a/pkg/controller/sparkapplication/sparkapp_util.go b/pkg/controller/sparkapplication/sparkapp_util.go
@@ -94,6 +94,16 @@ func getServiceAnnotations(app *v1beta2.SparkApplication) map[string]string {
 	return serviceAnnotations
 }
 
+func getServiceLabels(app *v1beta2.SparkApplication) map[string]string {
+	serviceLabels := map[string]string{}
+	if app.Spec.SparkUIOptions != nil && app.Spec.SparkUIOptions.ServiceLabels != nil {
+		for key, value := range app.Spec.SparkUIOptions.ServiceLabels {
+			serviceLabels[key] = value
+		}
+	}
+	return serviceLabels
+}
+
 func getIngressResourceAnnotations(app *v1beta2.SparkApplication) map[string]string {
 	ingressAnnotations := map[string]string{}
 	if app.Spec.SparkUIOptions != nil && app.Spec.SparkUIOptions.IngressAnnotations != nil {

diff --git a/pkg/controller/sparkapplication/sparkui.go b/pkg/controller/sparkapplication/sparkui.go
@@ -71,6 +71,7 @@ type SparkService struct {
 	targetPort         intstr.IntOrString
 	serviceIP          string
 	serviceAnnotations map[string]string
+	serviceLabels      map[string]string
 }
 
 // SparkIngress encapsulates information about the driver UI ingress.
@@ -285,6 +286,12 @@ func createSparkUIService(
 		service.ObjectMeta.Annotations = serviceAnnotations
 	}
 
+	serviceLabels := getServiceLabels(app)
+	if len(serviceLabels) != 0 {
+		glog.Infof("Creating a service labels %s for the Spark UI: %v", service.Name, &serviceLabels)
+		service.ObjectMeta.Labels = serviceLabels
+	}
+
 	glog.Infof("Creating a service %s for the Spark UI for application %s", service.Name, app.Name)
 	service, err = kubeClient.CoreV1().Services(app.Namespace).Create(context.TODO(), service, metav1.CreateOptions{})
 	if err != nil {
@@ -299,6 +306,7 @@ func createSparkUIService(
 		targetPort:         service.Spec.Ports[0].TargetPort,
 		serviceIP:          service.Spec.ClusterIP,
 		serviceAnnotations: serviceAnnotations,
+		serviceLabels:      serviceLabels,
 	}, nil
 }
 

diff --git a/pkg/controller/sparkapplication/sparkui_test.go b/pkg/controller/sparkapplication/sparkui_test.go
@@ -87,6 +87,10 @@ func TestCreateSparkUIService(t *testing.T) {
 		if !reflect.DeepEqual(serviceAnnotations, test.expectedService.serviceAnnotations) {
 			t.Errorf("%s: unexpected annotations wanted %s got %s", test.name, test.expectedService.serviceAnnotations, serviceAnnotations)
 		}
+		serviceLabels := service.ObjectMeta.Labels
+		if !reflect.DeepEqual(serviceLabels, test.expectedService.serviceLabels) {
+			t.Errorf("%s: unexpected labels wanted %s got %s", test.name, test.expectedService.serviceLabels, serviceLabels)
+		}
 	}
 	defaultPort := defaultSparkWebUIPort
 	defaultPortName := defaultSparkWebUIPortName
@@ -205,6 +209,25 @@ func TestCreateSparkUIService(t *testing.T) {
 			ExecutionAttempts:  1,
 		},
 	}
+	app8 := &v1beta2.SparkApplication{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "foo8",
+			Namespace: "default",
+			UID:       "foo-123",
+		},
+		Spec: v1beta2.SparkApplicationSpec{
+			SparkUIOptions: &v1beta2.SparkUIConfiguration{
+				ServiceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo8",
+					"key":                           "value",
+				},
+			},
+		},
+		Status: v1beta2.SparkApplicationStatus{
+			SparkApplicationID: "foo-8",
+			ExecutionAttempts:  1,
+		},
+	}
 	testcases := []testcase{
 		{
 			name: "service with custom serviceport and serviceport and target port are same",
@@ -214,6 +237,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceType:     apiv1.ServiceTypeClusterIP,
 				servicePortName: defaultPortName,
 				servicePort:     4041,
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo1",
+				},
 				targetPort: intstr.IntOrString{
 					Type:   intstr.Int,
 					IntVal: int32(4041),
@@ -233,6 +259,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceType:     apiv1.ServiceTypeClusterIP,
 				servicePortName: defaultPortName,
 				servicePort:     int32(defaultPort),
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo2",
+				},
 			},
 			expectedSelector: map[string]string{
 				config.SparkAppNameLabel: "foo2",
@@ -248,6 +277,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceType:     apiv1.ServiceTypeClusterIP,
 				servicePortName: defaultPortName,
 				servicePort:     80,
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo4",
+				},
 				targetPort: intstr.IntOrString{
 					Type:   intstr.Int,
 					IntVal: int32(4041),
@@ -267,6 +299,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceType:     apiv1.ServiceTypeNodePort,
 				servicePortName: defaultPortName,
 				servicePort:     int32(defaultPort),
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo5",
+				},
 			},
 			expectedSelector: map[string]string{
 				config.SparkAppNameLabel: "foo5",
@@ -282,6 +317,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceType:     apiv1.ServiceTypeClusterIP,
 				servicePortName: "http-spark-test",
 				servicePort:     int32(80),
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo6",
+				},
 			},
 			expectedSelector: map[string]string{
 				config.SparkAppNameLabel: "foo6",
@@ -300,6 +338,9 @@ func TestCreateSparkUIService(t *testing.T) {
 				serviceAnnotations: map[string]string{
 					"key": "value",
 				},
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo7",
+				},
 				targetPort: intstr.IntOrString{
 					Type:   intstr.Int,
 					IntVal: int32(4041),
@@ -311,6 +352,29 @@ func TestCreateSparkUIService(t *testing.T) {
 			},
 			expectError: false,
 		},
+		{
+			name: "service with custom labels",
+			app:  app8,
+			expectedService: SparkService{
+				serviceName:     fmt.Sprintf("%s-ui-svc", app8.GetName()),
+				serviceType:     apiv1.ServiceTypeClusterIP,
+				servicePortName: defaultPortName,
+				servicePort:     defaultPort,
+				serviceLabels: map[string]string{
+					"sparkoperator.k8s.io/app-name": "foo8",
+					"key":                           "value",
+				},
+				targetPort: intstr.IntOrString{
+					Type:   intstr.Int,
+					IntVal: int32(4041),
+				},
+			},
+			expectedSelector: map[string]string{
+				config.SparkAppNameLabel: "foo8",
+				config.SparkRoleLabel:    config.SparkDriverRole,
+			},
+			expectError: false,
+		},
 		{
 			name:        "service with bad port configurations",
 			app:         app3,

diff --git a/test/e2e/basic_test.go b/test/e2e/basic_test.go
@@ -18,6 +18,7 @@ package e2e
 
 import (
 	"context"
+	"log"
 	"strings"
 	"testing"
 
@@ -64,6 +65,7 @@ func TestSubmitSparkPiYaml(t *testing.T) {
 
 	app, _ := appFramework.GetSparkApplication(framework.SparkApplicationClient, appFramework.SparkTestNamespace, appName)
 	podName := app.Status.DriverInfo.PodName
+	log.Printf("LABELS: %v", app.ObjectMeta.GetLabels())
 	rawLogs, err := framework.KubeClient.CoreV1().Pods(appFramework.SparkTestNamespace).GetLogs(podName, &v1.PodLogOptions{}).Do(context.TODO()).Raw()
 	assert.Equal(t, nil, err)
 	assert.NotEqual(t, -1, strings.Index(string(rawLogs), "Pi is roughly 3"))