[RayCluster controller] Add headServiceAnnotations field to RayCluste…

…r CR (ray-project#841) Adds headServiceAnnotations field to RayCluster CR. Signed-off-by: Kornel Csernai <[email protected]>
lowang-bh · Dec 16, 2022 · e5c1bed · e5c1bed
1 parent 1ec19a9
commit e5c1bed
Show file tree

Hide file tree

Showing 13 changed files with 68 additions and 11 deletions.
diff --git a/helm-chart/kuberay-operator/crds/ray.io_rayclusters.yaml b/helm-chart/kuberay-operator/crds/ray.io_rayclusters.yaml
@@ -5709,6 +5709,10 @@ spec:
                 - serviceType
                 - template
                 type: object
+              headServiceAnnotations:
+                additionalProperties:
+                  type: string
+                type: object
               rayVersion:
                 description: RayVersion is the version of ray being used. this affects
                   the command used to start ray

diff --git a/helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml b/helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml
@@ -5973,6 +5973,10 @@ spec:
                     - serviceType
                     - template
                     type: object
+                  headServiceAnnotations:
+                    additionalProperties:
+                      type: string
+                    type: object
                   rayVersion:
                     description: RayVersion is the version of ray being used. this
                       affects the command used to start ray

diff --git a/helm-chart/kuberay-operator/crds/ray.io_rayservices.yaml b/helm-chart/kuberay-operator/crds/ray.io_rayservices.yaml
@@ -5959,6 +5959,10 @@ spec:
                     - serviceType
                     - template
                     type: object
+                  headServiceAnnotations:
+                    additionalProperties:
+                      type: string
+                    type: object
                   rayVersion:
                     description: RayVersion is the version of ray being used. this
                       affects the command used to start ray

diff --git a/ray-operator/apis/ray/v1alpha1/raycluster_types.go b/ray-operator/apis/ray/v1alpha1/raycluster_types.go
@@ -21,7 +21,8 @@ type RayClusterSpec struct {
 	// EnableInTreeAutoscaling indicates whether operator should create in tree autoscaling configs
 	EnableInTreeAutoscaling *bool `json:"enableInTreeAutoscaling,omitempty"`
 	// AutoscalerOptions specifies optional configuration for the Ray autoscaler.
-	AutoscalerOptions *AutoscalerOptions `json:"autoscalerOptions,omitempty"`
+	AutoscalerOptions      *AutoscalerOptions `json:"autoscalerOptions,omitempty"`
+	HeadServiceAnnotations map[string]string  `json:"headServiceAnnotations,omitempty"`
 }
 
 // HeadGroupSpec are the spec for the head pod

diff --git a/ray-operator/apis/ray/v1alpha1/zz_generated.deepcopy.go b/ray-operator/apis/ray/v1alpha1/zz_generated.deepcopy.go
diff --git a/ray-operator/config/crd/bases/ray.io_rayclusters.yaml b/ray-operator/config/crd/bases/ray.io_rayclusters.yaml
@@ -5709,6 +5709,10 @@ spec:
                 - serviceType
                 - template
                 type: object
+              headServiceAnnotations:
+                additionalProperties:
+                  type: string
+                type: object
               rayVersion:
                 description: RayVersion is the version of ray being used. this affects
                   the command used to start ray

diff --git a/ray-operator/config/crd/bases/ray.io_rayjobs.yaml b/ray-operator/config/crd/bases/ray.io_rayjobs.yaml
@@ -5973,6 +5973,10 @@ spec:
                     - serviceType
                     - template
                     type: object
+                  headServiceAnnotations:
+                    additionalProperties:
+                      type: string
+                    type: object
                   rayVersion:
                     description: RayVersion is the version of ray being used. this
                       affects the command used to start ray

diff --git a/ray-operator/config/crd/bases/ray.io_rayservices.yaml b/ray-operator/config/crd/bases/ray.io_rayservices.yaml
@@ -5959,6 +5959,10 @@ spec:
                     - serviceType
                     - template
                     type: object
+                  headServiceAnnotations:
+                    additionalProperties:
+                      type: string
+                    type: object
                   rayVersion:
                     description: RayVersion is the version of ray being used. this
                       affects the command used to start ray

diff --git a/ray-operator/config/samples/ray_v1alpha1_rayservice.yaml b/ray-operator/config/samples/ray_v1alpha1_rayservice.yaml
@@ -109,3 +109,6 @@ spec:
                   requests:
                     cpu: "500m"
                     memory: "2Gi"
+    headServiceAnnotations: {}
+      # annotations passed on for the Head Service
+      # service_key: "service_value"
diff --git a/ray-operator/controllers/ray/common/service.go b/ray-operator/controllers/ray/common/service.go
@@ -21,7 +21,7 @@ func HeadServiceLabels(cluster rayiov1alpha1.RayCluster) map[string]string {
 
 // BuildServiceForHeadPod Builds the service for a pod. Currently, there is only one service that allows
 // the worker nodes to connect to the head node.
-func BuildServiceForHeadPod(cluster rayiov1alpha1.RayCluster, labels map[string]string) (*corev1.Service, error) {
+func BuildServiceForHeadPod(cluster rayiov1alpha1.RayCluster, labels map[string]string, annotations map[string]string) (*corev1.Service, error) {
 	if labels == nil {
 		labels = make(map[string]string)
 	}
@@ -34,11 +34,16 @@ func BuildServiceForHeadPod(cluster rayiov1alpha1.RayCluster, labels map[string]
 		}
 	}
 
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+
 	service := &corev1.Service{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      utils.GenerateServiceName(cluster.Name),
-			Namespace: cluster.Namespace,
-			Labels:    labels,
+			Name:        utils.GenerateServiceName(cluster.Name),
+			Namespace:   cluster.Namespace,
+			Labels:      labels,
+			Annotations: annotations,
 		},
 		Spec: corev1.ServiceSpec{
 			Selector: labels,
@@ -61,7 +66,7 @@ func BuildServiceForHeadPod(cluster rayiov1alpha1.RayCluster, labels map[string]
 // the worker nodes to connect to the head node.
 // RayService controller updates the service whenever a new RayCluster serves the traffic.
 func BuildHeadServiceForRayService(rayService rayiov1alpha1.RayService, rayCluster rayiov1alpha1.RayCluster) (*corev1.Service, error) {
-	service, err := BuildServiceForHeadPod(rayCluster, nil)
+	service, err := BuildServiceForHeadPod(rayCluster, nil, nil)
 	if err != nil {
 		return nil, err
 	}

diff --git a/ray-operator/controllers/ray/common/service_test.go b/ray-operator/controllers/ray/common/service_test.go
@@ -63,7 +63,7 @@ var instanceWithWrongSvc = &rayiov1alpha1.RayCluster{
 }
 
 func TestBuildServiceForHeadPod(t *testing.T) {
-	svc, err := BuildServiceForHeadPod(*instanceWithWrongSvc, nil)
+	svc, err := BuildServiceForHeadPod(*instanceWithWrongSvc, nil, nil)
 	assert.Nil(t, err)
 
 	actualResult := svc.Spec.Selector[RayClusterLabelKey]
@@ -96,7 +96,8 @@ func TestBuildServiceForHeadPod(t *testing.T) {
 func TestBuildServiceForHeadPodWithAppNameLabel(t *testing.T) {
 	labels := make(map[string]string)
 	labels[KubernetesApplicationNameLabelKey] = "testname"
-	svc, err := BuildServiceForHeadPod(*instanceWithWrongSvc, labels)
+
+	svc, err := BuildServiceForHeadPod(*instanceWithWrongSvc, labels, nil)
 	assert.Nil(t, err)
 
 	actualResult := svc.Spec.Selector[KubernetesApplicationNameLabelKey]
@@ -113,3 +114,15 @@ func TestBuildServiceForHeadPodWithAppNameLabel(t *testing.T) {
 		t.Fatalf("Expected `%v` but got `%v`", expectedLength, actualLength)
 	}
 }
+
+func TestBuildServiceForHeadPodWithAnnotations(t *testing.T) {
+	annotations := make(map[string]string)
+	annotations["key1"] = "testvalue1"
+	annotations["key2"] = "testvalue2"
+	svc, err := BuildServiceForHeadPod(*instanceWithWrongSvc, nil, annotations)
+	assert.Nil(t, err)
+
+	if !reflect.DeepEqual(svc.ObjectMeta.Annotations, annotations) {
+		t.Fatalf("Expected `%v` but got `%v`", annotations, svc.ObjectMeta.Annotations)
+	}
+}
diff --git a/ray-operator/controllers/ray/raycluster_controller.go b/ray-operator/controllers/ray/raycluster_controller.go
@@ -310,7 +310,11 @@ func (r *RayClusterReconciler) reconcileServices(instance *rayiov1alpha1.RayClus
 			if val, ok := instance.Spec.HeadGroupSpec.Template.ObjectMeta.Labels[common.KubernetesApplicationNameLabelKey]; ok {
 				labels[common.KubernetesApplicationNameLabelKey] = val
 			}
-			raySvc, err = common.BuildServiceForHeadPod(*instance, labels)
+			annotations := make(map[string]string)
+			for k, v := range instance.Spec.HeadServiceAnnotations {
+				annotations[k] = v
+			}
+			raySvc, err = common.BuildServiceForHeadPod(*instance, labels, annotations)
 		} else if serviceType == common.AgentService {
 			raySvc, err = common.BuildDashboardService(*instance)
 		}

diff --git a/ray-operator/controllers/ray/raycluster_controller_fake_test.go b/ray-operator/controllers/ray/raycluster_controller_fake_test.go
@@ -319,7 +319,7 @@ func setupTest(t *testing.T) {
 		},
 	}
 
-	headService, err := common.BuildServiceForHeadPod(*testRayCluster, nil)
+	headService, err := common.BuildServiceForHeadPod(*testRayCluster, nil, nil)
 	if err != nil {
 		t.Errorf("failed to build head service: %v", err)
 	}
@@ -904,7 +904,7 @@ func TestGetHeadServiceIP(t *testing.T) {
 	defer tearDown(t)
 
 	headServiceIP := "1.2.3.4"
-	headService, err := common.BuildServiceForHeadPod(*testRayCluster, nil)
+	headService, err := common.BuildServiceForHeadPod(*testRayCluster, nil, nil)
 	if err != nil {
 		t.Errorf("failed to build head service: %v", err)
 	}