From e5c1bed9d41af38789a395c75f1032ac2735e602 Mon Sep 17 00:00:00 2001 From: Kornel Csernai <69173677+cskornel-doordash@users.noreply.github.com> Date: Thu, 15 Dec 2022 18:47:36 -0800 Subject: [PATCH] [RayCluster controller] Add headServiceAnnotations field to RayCluster CR (#841) Adds headServiceAnnotations field to RayCluster CR. Signed-off-by: Kornel Csernai <69173677+cskornel-doordash@users.noreply.github.com> --- .../crds/ray.io_rayclusters.yaml | 4 ++++ .../kuberay-operator/crds/ray.io_rayjobs.yaml | 4 ++++ .../crds/ray.io_rayservices.yaml | 4 ++++ .../apis/ray/v1alpha1/raycluster_types.go | 3 ++- .../apis/ray/v1alpha1/zz_generated.deepcopy.go | 7 +++++++ .../config/crd/bases/ray.io_rayclusters.yaml | 4 ++++ .../config/crd/bases/ray.io_rayjobs.yaml | 4 ++++ .../config/crd/bases/ray.io_rayservices.yaml | 4 ++++ .../config/samples/ray_v1alpha1_rayservice.yaml | 3 +++ ray-operator/controllers/ray/common/service.go | 15 ++++++++++----- .../controllers/ray/common/service_test.go | 17 +++++++++++++++-- .../controllers/ray/raycluster_controller.go | 6 +++++- .../ray/raycluster_controller_fake_test.go | 4 ++-- 13 files changed, 68 insertions(+), 11 deletions(-) diff --git a/helm-chart/kuberay-operator/crds/ray.io_rayclusters.yaml b/helm-chart/kuberay-operator/crds/ray.io_rayclusters.yaml index 1a4a2f1a3f9..3898a48460b 100644 --- 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 index 72efe71a79f..6513284a973 100644 --- 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 index 24e5466a1cc..1f0fb579a2a 100644 --- 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 index 713928d961d..a3ea2a49191 100644 --- 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 index e84cb10216d..a5451b378d1 100644 --- a/ray-operator/apis/ray/v1alpha1/zz_generated.deepcopy.go +++ b/ray-operator/apis/ray/v1alpha1/zz_generated.deepcopy.go @@ -278,6 +278,13 @@ func (in *RayClusterSpec) DeepCopyInto(out *RayClusterSpec) { *out = new(AutoscalerOptions) (*in).DeepCopyInto(*out) } + if in.HeadServiceAnnotations != nil { + in, out := &in.HeadServiceAnnotations, &out.HeadServiceAnnotations + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new RayClusterSpec. diff --git a/ray-operator/config/crd/bases/ray.io_rayclusters.yaml b/ray-operator/config/crd/bases/ray.io_rayclusters.yaml index 1a4a2f1a3f9..3898a48460b 100644 --- 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 index 72efe71a79f..6513284a973 100644 --- 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 index 24e5466a1cc..1f0fb579a2a 100644 --- 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 index b5c8a02684f..40276933c95 100644 --- 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 index 7344b990d27..1db4143395d 100644 --- 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 index 55ec9af6f25..e0879077c60 100644 --- 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 index f2f8c30a181..34875b3f66f 100644 --- 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 index 5ac52826e8b..7fbd6603636 100644 --- 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) }