diff --git a/ray-operator/controllers/ray/common/service.go b/ray-operator/controllers/ray/common/service.go index 7ffb8f58eeb..274d3e7034e 100644 --- a/ray-operator/controllers/ray/common/service.go +++ b/ray-operator/controllers/ray/common/service.go @@ -5,6 +5,8 @@ import ( "fmt" "sort" + "sigs.k8s.io/controller-runtime/pkg/client" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ctrl "sigs.k8s.io/controller-runtime" @@ -13,17 +15,6 @@ import ( "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils" ) -// HeadServiceLabels returns the default labels for a cluster's head service. -func HeadServiceLabels(cluster rayv1.RayCluster) map[string]string { - return map[string]string{ - utils.RayClusterLabelKey: cluster.Name, - utils.RayNodeTypeLabelKey: string(rayv1.HeadNode), - utils.RayIDLabelKey: utils.CheckLabel(utils.GenerateIdentifier(cluster.Name, rayv1.HeadNode)), - utils.KubernetesApplicationNameLabelKey: utils.ApplicationName, - utils.KubernetesCreatedByLabelKey: utils.ComponentName, - } -} - // 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(ctx context.Context, cluster rayv1.RayCluster, labels map[string]string, annotations map[string]string) (*corev1.Service, error) { @@ -33,7 +24,7 @@ func BuildServiceForHeadPod(ctx context.Context, cluster rayv1.RayCluster, label labels = make(map[string]string) } - default_labels := HeadServiceLabels(cluster) + default_labels := HeadServiceAssociationOptions(&cluster)[0].(client.MatchingLabels) // selector consists of *only* the keys in default_labels, updated with the values in labels if they exist selector := make(map[string]string) diff --git a/ray-operator/controllers/ray/common/service_test.go b/ray-operator/controllers/ray/common/service_test.go index 148826158b3..8f8044f3c86 100644 --- a/ray-operator/controllers/ray/common/service_test.go +++ b/ray-operator/controllers/ray/common/service_test.go @@ -7,6 +7,8 @@ import ( "reflect" "testing" + "sigs.k8s.io/controller-runtime/pkg/client" + rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils" @@ -293,7 +295,7 @@ func TestUserSpecifiedHeadService(t *testing.T) { // The selector field should only use the keys from the five default labels. The values should be updated with the values from the template labels. // The user-provided HeadService labels should be ignored for the purposes of the selector field. The user-provided Selector field should be ignored. - default_labels := HeadServiceLabels(*testRayClusterWithHeadService) + default_labels := HeadServiceAssociationOptions(testRayClusterWithHeadService)[0].(client.MatchingLabels) // Make sure this test isn't spuriously passing. Check that RayClusterLabelKey is in the default labels. if _, ok := default_labels[utils.RayClusterLabelKey]; !ok { t.Errorf("utils.RayClusterLabelKey=%s should be in the default labels", utils.RayClusterLabelKey) diff --git a/ray-operator/controllers/ray/raycluster_controller_fake_test.go b/ray-operator/controllers/ray/raycluster_controller_fake_test.go index 5aef753ea5c..e4d0bac7761 100644 --- a/ray-operator/controllers/ray/raycluster_controller_fake_test.go +++ b/ray-operator/controllers/ray/raycluster_controller_fake_test.go @@ -1372,7 +1372,7 @@ func TestGetHeadServiceIP(t *testing.T) { ObjectMeta: metav1.ObjectMeta{ Name: "unexpectedExtraHeadService", Namespace: namespaceStr, - Labels: common.HeadServiceLabels(*testRayCluster), + Labels: common.HeadServiceAssociationOptions(testRayCluster)[0].(client.MatchingLabels), }, }