[History Server] Prefer active RayCluster for RayService sessions

historyserver/config/service_account.yaml

@@ -12,6 +12,9 @@ rules:
 - apiGroups: ["ray.io"]
   resources: ["rayclusters"]
   verbs: ["list", "get"]
+- apiGroups: ["ray.io"]
+  resources: ["rayjobs", "rayservices"]
+  verbs: ["get"]
 ---
 apiVersion: rbac.authorization.k8s.io/v1
 kind: ClusterRoleBinding

historyserver/pkg/historyserver/clientmanager.go

@@ -71,6 +71,22 @@ func (c *ClientManager) GetRayCluster(ctx context.Context, namespace, name strin
 	return &rayCluster, nil
 }
 
+func (c *ClientManager) GetRayJob(ctx context.Context, namespace, name string) (*rayv1.RayJob, error) {
+	var rayJob rayv1.RayJob
+	if err := c.Client().Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, &rayJob); err != nil {
+		return nil, err
+	}
+	return &rayJob, nil
+}
+
+func (c *ClientManager) GetRayService(ctx context.Context, namespace, name string) (*rayv1.RayService, error) {
+	var rayService rayv1.RayService
+	if err := c.Client().Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, &rayService); err != nil {
+		return nil, err
+	}
+	return &rayService, nil
+}
+
 type ClientManagerConfig struct {
 	Kubeconfigs        string
 	UseKubernetesProxy bool

historyserver/pkg/historyserver/enter_cluster_test.go

@@ -562,6 +562,135 @@ func TestEnterClusterRayJobAndRayService(t *testing.T) {
 	})
 }
 
+func TestResolveSessionOwnerClusterSelection(t *testing.T) {
+	const (
+		namespace          = "default"
+		resourceName       = "my-owner"
+		rayJobClusterName  = "rayjob-cluster"
+		activeClusterName  = "active-cluster"
+		pendingClusterName = "pending-cluster"
+		staleClusterName   = "stale-cluster"
+	)
+
+	tests := []struct {
+		name                 string
+		resourceType         string
+		rayJobClusterName    string
+		activeClusterName    string
+		pendingClusterName   string
+		existingClusterNames []string
+		ownerGetErr          error
+		wantClusterName      string
+		wantErrorContains    string
+	}{
+		{
+			name:                 "RayJob selects the cluster named in status",
+			resourceType:         utils.RayJobKind,
+			rayJobClusterName:    rayJobClusterName,
+			existingClusterNames: []string{staleClusterName, rayJobClusterName},
+			wantClusterName:      rayJobClusterName,
+		},
+		{
+			name:              "RayJob with an empty cluster name has no live session",
+			resourceType:      utils.RayJobKind,
+			rayJobClusterName: "",
+		},
+		{
+			name:               "RayService prefers active",
+			resourceType:       utils.RayServiceKind,
+			activeClusterName:  activeClusterName,
+			pendingClusterName: pendingClusterName,
+			existingClusterNames: []string{
+				pendingClusterName,
+				activeClusterName,
+			},
+			wantClusterName: activeClusterName,
+		},
+		{
+			name:               "RayService uses pending when active name is empty",
+			resourceType:       utils.RayServiceKind,
+			pendingClusterName: pendingClusterName,
+			existingClusterNames: []string{
+				pendingClusterName,
+			},
+			wantClusterName: pendingClusterName,
+		},
+		{
+			name:              "RayService owner lookup error is returned",
+			resourceType:      utils.RayServiceKind,
+			activeClusterName: activeClusterName,
+			ownerGetErr:       fmt.Errorf("connection timeout"),
+			wantErrorContains: "failed to get RayService default/my-owner",
+		},
+	}
+
+	scheme := runtime.NewScheme()
+	if err := rayv1.AddToScheme(scheme); err != nil {
+		t.Fatalf("Failed to add Ray types to scheme: %v", err)
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			objects := make([]client.Object, 0, len(tt.existingClusterNames)+1)
+			switch tt.resourceType {
+			case utils.RayJobKind:
+				objects = append(objects, &rayv1.RayJob{
+					ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: resourceName},
+					Status:     rayv1.RayJobStatus{RayClusterName: tt.rayJobClusterName},
+				})
+			case utils.RayServiceKind:
+				objects = append(objects, &rayv1.RayService{
+					ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: resourceName},
+					Status: rayv1.RayServiceStatuses{
+						ActiveServiceStatus:  rayv1.RayServiceStatus{RayClusterName: tt.activeClusterName},
+						PendingServiceStatus: rayv1.RayServiceStatus{RayClusterName: tt.pendingClusterName},
+					},
+				})
+			default:
+				t.Fatalf("unsupported test resource type %q", tt.resourceType)
+			}
+			for _, clusterName := range tt.existingClusterNames {
+				objects = append(objects, &rayv1.RayCluster{
+					ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: clusterName},
+				})
+			}
+			k8sClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(objects...).Build()
+			var testClient client.Client = k8sClient
+			if tt.ownerGetErr != nil {
+				testClient = &errorClient{err: tt.ownerGetErr}
+			}
+			clientManager := &ClientManager{
+				clients: []client.Client{testClient},
+			}
+			handler := &ServerHandler{
+				enableLiveClusters: true,
+				clientManager:      clientManager,
+				reader:             &mockStorageReader{},
+			}
+
+			clusterInfo, _, err := handler.resolveSession(context.Background(), namespace, tt.resourceType, resourceName, "latest")
+			if tt.wantErrorContains != "" {
+				if err == nil || !strings.Contains(err.Error(), tt.wantErrorContains) {
+					t.Fatalf("resolveSession() error = %v, want an error containing %q", err, tt.wantErrorContains)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("resolveSession() returned an unexpected error: %v", err)
+			}
+			if clusterInfo.Name != tt.wantClusterName {
+				t.Errorf("resolveSession() selected RayCluster %q, want %q", clusterInfo.Name, tt.wantClusterName)
+			}
+			if tt.wantClusterName == "" {
+				return
+			}
+			if clusterInfo.OwnerKind != tt.resourceType || clusterInfo.OwnerName != resourceName {
+				t.Errorf("resolveSession() owner = %s/%s, want %s/%s", clusterInfo.OwnerKind, clusterInfo.OwnerName, tt.resourceType, resourceName)
+			}
+		})
+	}
+}
+
 // newDisabledLiveHandler returns a handler with live clusters disabled and a single RayCluster that
 // exists in both Kubernetes and storage. status controls what the fake session loader reports for
 // its stored session.

historyserver/pkg/historyserver/reader.go

@@ -24,8 +24,6 @@ import (
 
 	"github.com/ray-project/kuberay/historyserver/pkg/utils"
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
-	rayutils "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
-	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 const (
@@ -102,19 +100,6 @@ func buildLiveClusterInfo(liveCluster *rayv1.RayCluster) utils.ClusterInfo {
 	}
 }
 
-func crdLabelValueFor(kindLower string) string {
-	switch kindLower {
-	case utils.RayJobKind:
-		return "RayJob"
-	case utils.RayServiceKind:
-		return "RayService"
-	case utils.RayClusterKind:
-		return "RayCluster"
-	default:
-		return kindLower
-	}
-}
-
 func (s *ServerHandler) listClusters(limit int) []utils.ClusterInfo {
 	// Initial continuation marker
 	logrus.Debugf("Prepare to get list clusters info ...")
@@ -143,6 +128,41 @@ func (s *ServerHandler) listClusters(limit int) []utils.ClusterInfo {
 	return clusters
 }
 
+func (s *ServerHandler) ownerRayClusterNames(ctx context.Context, namespace, resourceType, resourceName string) ([]string, error) {
+	switch resourceType {
+	case utils.RayJobKind:
+		rayJob, err := s.clientManager.GetRayJob(ctx, namespace, resourceName)
+		if err != nil {
+			if apierrors.IsNotFound(err) {
+				return nil, nil
+			}
+			return nil, fmt.Errorf("failed to get RayJob %s/%s: %w", namespace, resourceName, err)
+		}
+		if rayJob.Status.RayClusterName == "" {
+			return nil, nil
+		}
+		return []string{rayJob.Status.RayClusterName}, nil
+	case utils.RayServiceKind:
+		rayService, err := s.clientManager.GetRayService(ctx, namespace, resourceName)
+		if err != nil {
+			if apierrors.IsNotFound(err) {
+				return nil, nil
+			}
+			return nil, fmt.Errorf("failed to get RayService %s/%s: %w", namespace, resourceName, err)
+		}
+
+		clusterNames := make([]string, 0, 2)
+		for _, name := range []string{rayService.Status.ActiveServiceStatus.RayClusterName, rayService.Status.PendingServiceStatus.RayClusterName} {
+			if name != "" {
+				clusterNames = append(clusterNames, name)
+			}
+		}
+		return clusterNames, nil
+	default:
+		return nil, fmt.Errorf("unsupported owner resource kind: %q", resourceType)
+	}
+}
+
 // resolveSession maps (namespace, resourceType, resourceName, session) to a concrete ClusterInfo.
 func (s *ServerHandler) resolveSession(ctx context.Context, namespace, resourceType, resourceName, session string) (utils.ClusterInfo, bool, error) {
 	isLatestOrEmpty := session == "latest" || session == ""
@@ -166,21 +186,20 @@ func (s *ServerHandler) resolveSession(ctx context.Context, namespace, resourceT
 				return utils.ClusterInfo{}, false, fmt.Errorf("failed to check live RayCluster %s/%s: %w", namespace, resourceName, err)
 			}
 		} else {
-			// Both labels are needed: name alone is ambiguous since a RayJob and a
-			// RayService in the same namespace can share a name.
-			liveClusters, err := s.clientManager.ListRayClusters(ctx,
-				client.InNamespace(namespace),
-				client.MatchingLabels{
-					rayutils.RayOriginatedFromCRNameLabelKey: resourceName,
-					rayutils.RayOriginatedFromCRDLabelKey:    crdLabelValueFor(resTypeLower),
-				},
-			)
+			clusterNames, err := s.ownerRayClusterNames(ctx, namespace, resTypeLower, resourceName)
 			if err != nil {
-				return utils.ClusterInfo{}, false, fmt.Errorf("failed to list live RayClusters: %w", err)
+				return utils.ClusterInfo{}, false, err
 			}
-			// TODO: A RayService owns both an active and a pending cluster during an upgrade. Needs to decide which one to take.
-			if len(liveClusters) > 0 {
-				info := buildLiveClusterInfo(liveClusters[0])
+			for _, clusterName := range clusterNames {
+				liveCluster, err := s.clientManager.GetRayCluster(ctx, namespace, clusterName)
+				if err != nil {
+					if !apierrors.IsNotFound(err) {
+						return utils.ClusterInfo{}, false, fmt.Errorf("failed to check live RayCluster %s/%s: %w", namespace, clusterName, err)
+					}
+					continue
+				}
+
+				info := buildLiveClusterInfo(liveCluster)
 				if info.OwnerKind == "" {
 					info.OwnerKind = resTypeLower
 					info.OwnerName = resourceName