[Refactor][RayService] Add GetRayClusterWithRayServiceAssociationOpti…

…ons (#2070)
ray-project · Apr 22, 2024 · b8212e5 · b8212e5
1 parent 7af8988
commit b8212e5
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 10 deletions.
diff --git a/ray-operator/controllers/ray/common/association.go b/ray-operator/controllers/ray/common/association.go
@@ -93,6 +93,16 @@ func RayClusterAllPodsAssociationOptions(instance *rayv1.RayCluster) Association
 	}
 }
 
+func RayServiceRayClustersAssociationOptions(rayService *rayv1.RayService) AssociationOptions {
+	return AssociationOptions{
+		client.InNamespace(rayService.Namespace),
+		client.MatchingLabels{
+			utils.RayOriginatedFromCRNameLabelKey: rayService.Name,
+			utils.RayOriginatedFromCRDLabelKey:    utils.RayOriginatedFromCRDLabelValue(utils.RayServiceCRD),
+		},
+	}
+}
+
 func RayServiceServeServiceNamespacedName(rayService *rayv1.RayService) types.NamespacedName {
 	if rayService.Spec.ServeService != nil && rayService.Spec.ServeService.Name != "" {
 		return types.NamespacedName{

diff --git a/ray-operator/controllers/ray/rayservice_controller.go b/ray-operator/controllers/ray/rayservice_controller.go
@@ -435,12 +435,9 @@ func (r *RayServiceReconciler) reconcileRayCluster(ctx context.Context, rayServi
 func (r *RayServiceReconciler) cleanUpRayClusterInstance(ctx context.Context, rayServiceInstance *rayv1.RayService) error {
 	logger := ctrl.LoggerFrom(ctx)
 	rayClusterList := rayv1.RayClusterList{}
-	filterLabels := client.MatchingLabels{
-		utils.RayOriginatedFromCRNameLabelKey: rayServiceInstance.Name,
-		utils.RayOriginatedFromCRDLabelKey:    utils.RayOriginatedFromCRDLabelValue(utils.RayServiceCRD),
-	}
+
 	var err error
-	if err = r.List(ctx, &rayClusterList, client.InNamespace(rayServiceInstance.Namespace), filterLabels); err != nil {
+	if err = r.List(ctx, &rayClusterList, common.RayServiceRayClustersAssociationOptions(rayServiceInstance).ToListOptions()...); err != nil {
 		logger.Error(err, "Fail to list RayCluster for "+rayServiceInstance.Name)
 		return err
 	}
@@ -1237,9 +1234,7 @@ func isServeAppUnhealthyOrDeployedFailed(appStatus string) bool {
 // TODO: Move this function to util.go and always use this function to retrieve the head Pod.
 func (r *RayServiceReconciler) getHeadPod(ctx context.Context, instance *rayv1.RayCluster) (*corev1.Pod, error) {
 	podList := corev1.PodList{}
-	filterLabels := client.MatchingLabels{utils.RayClusterLabelKey: instance.Name, utils.RayNodeTypeLabelKey: string(rayv1.HeadNode)}
-
-	if err := r.List(ctx, &podList, client.InNamespace(instance.Namespace), filterLabels); err != nil {
+	if err := r.List(ctx, &podList, common.RayClusterHeadPodsAssociationOptions(instance).ToListOptions()...); err != nil {
 		return nil, err
 	}
 

diff --git a/ray-operator/controllers/ray/rayservice_controller_test.go b/ray-operator/controllers/ray/rayservice_controller_test.go
@@ -21,6 +21,8 @@ import (
 	"os"
 	"time"
 
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray/common"
+
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
 
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -247,9 +249,8 @@ var _ = Context("Inside the default namespace", func() {
 		})
 
 		It("should create more than 1 worker", func() {
-			filterLabels := client.MatchingLabels{utils.RayClusterLabelKey: myRayService.Status.ActiveServiceStatus.RayClusterName, utils.RayNodeGroupLabelKey: "small-group"}
 			Eventually(
-				listResourceFunc(ctx, &workerPods, filterLabels, &client.ListOptions{Namespace: "default"}),
+				listResourceFunc(ctx, &workerPods, common.RayClusterGroupPodsAssociationOptions(myRayCluster, "small-group").ToListOptions()...),
 				time.Second*15, time.Millisecond*500).Should(Equal(3), fmt.Sprintf("workerGroup %v", workerPods.Items))
 			if len(workerPods.Items) > 0 {
 				Expect(workerPods.Items[0].Status.Phase).Should(Or(Equal(corev1.PodRunning), Equal(corev1.PodPending)))