feat(ray): resolve RayService names for existing-cluster tasks

flyteplugins/go/tasks/plugins/k8s/ray/ray.go

@@ -9,17 +9,21 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 	"github.com/samber/lo"
 	"gopkg.in/yaml.v2"
 	v1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/validation"
 	"k8s.io/client-go/kubernetes/scheme"
+	"k8s.io/client-go/rest"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	ctrlconfig "sigs.k8s.io/controller-runtime/pkg/client/config"
 
 	flyteerr "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/errors"
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/logs"
@@ -54,6 +58,10 @@ const (
 	rayClusterSelectorKey = "ray.io/cluster"
 )
 
+// rayServiceReader resolves RayService names for existing-cluster jobs. BuildResource gets no
+// Kubernetes client, so the plugin creates its own on first use.
+var rayServiceReader = sync.OnceValues(newRayServiceReader)
+
 var logTemplateRegexes = struct {
 	RayClusterName *regexp.Regexp
 	RayJobID       *regexp.Regexp
@@ -131,6 +139,10 @@ func (rayJobResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsC
 		if !cfg.AllowExistingCluster {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "existing Ray clusters are disabled by the ray plugin config (allowExistingCluster)")
 		}
+		clusterSelector, err = resolveRayService(ctx, taskCtx.TaskExecutionMetadata().GetNamespace(), clusterSelector)
+		if err != nil {
+			return nil, err
+		}
 		return constructSelectorRayJob(taskCtx, &rayJob, objectMeta, *primaryContainer, clusterSelector)
 	}
 
@@ -419,6 +431,50 @@ func existingClusterSelector(taskTemplate *core.TaskTemplate) (map[string]string
 	return map[string]string{rayClusterSelectorKey: name}, nil
 }
 
+// resolveRayService replaces a RayService name in the selector with the RayService's active RayCluster.
+// KubeRay resolves ray.io/cluster as a RayCluster name, and a RayService generates a new one on every
+// upgrade. Names that are not a RayService pass through.
+func resolveRayService(ctx context.Context, namespace string, clusterSelector map[string]string) (map[string]string, error) {
+	name := clusterSelector[rayClusterSelectorKey]
+	if name == "" {
+		return clusterSelector, nil
+	}
+	reader, err := rayServiceReader()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create client for RayService lookup: %w", err)
+	}
+
+	rayService := &rayv1.RayService{}
+	err = reader.Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, rayService)
+	if apierrors.IsNotFound(err) {
+		return clusterSelector, nil
+	} else if err != nil {
+		return nil, fmt.Errorf("failed to get RayService %s/%s: %w", namespace, name, err)
+	}
+
+	active := rayService.Status.ActiveServiceStatus.RayClusterName
+	if active == "" {
+		return nil, fmt.Errorf("RayService %s/%s has no active RayCluster yet", namespace, name)
+	}
+	resolved := maps.Clone(clusterSelector)
+	resolved[rayClusterSelectorKey] = active
+	return resolved, nil
+}
+
+func newRayServiceReader() (client.Reader, error) {
+	var restConfig *rest.Config
+	var err error
+	if remote := GetConfig().RemoteClusterConfig; remote.Enabled {
+		restConfig, err = k8s.KubeClientConfig(remote.Endpoint, remote.Auth)
+	} else {
+		restConfig, err = ctrlconfig.GetConfig()
+	}
+	if err != nil {
+		return nil, err
+	}
+	return client.New(restConfig, client.Options{Scheme: scheme.Scheme})
+}
+
 // constructSelectorRayJob builds a RayJob that runs the task as the driver on the existing RayCluster
 // matched by clusterSelector. KubeRay never deletes a selected cluster, so shutdown and TTL stay unset.
 func constructSelectorRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob *plugins.RayJob, objectMeta *metav1.ObjectMeta, primaryContainer v1.Container, clusterSelector map[string]string) (*rayv1.RayJob, error) {

flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go

@@ -18,6 +18,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/logs"
@@ -2634,7 +2635,15 @@ func dummyExistingClusterTaskTemplate(cluster string, rayJob *plugins.RayJob) *c
 	return taskTemplate
 }
 
+func stubRayServiceReader(t *testing.T, objs ...client.Object) {
+	original := rayServiceReader
+	t.Cleanup(func() { rayServiceReader = original })
+	reader := fake.NewClientBuilder().WithObjects(objs...).Build()
+	rayServiceReader = func() (client.Reader, error) { return reader, nil }
+}
+
 func TestBuildResourceRayExistingCluster(t *testing.T) {
+	stubRayServiceReader(t)
 	originalConfig := GetConfig()
 	defer func() { assert.NoError(t, SetConfig(originalConfig)) }()
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{}))
@@ -2735,3 +2744,49 @@ func TestBuildResourceRayExistingCluster(t *testing.T) {
 func TestAllowExistingClusterDefaultsToFalse(t *testing.T) {
 	assert.False(t, defaultConfig.AllowExistingCluster)
 }
+
+func TestBuildResourceRayExistingClusterRayService(t *testing.T) {
+	originalConfig := GetConfig()
+	defer func() { assert.NoError(t, SetConfig(originalConfig)) }()
+	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{}))
+	cfg := *originalConfig
+	cfg.AllowExistingCluster = true
+	assert.NoError(t, SetConfig(&cfg))
+
+	stubRayServiceReader(t,
+		&rayv1.RayService{
+			ObjectMeta: metav1.ObjectMeta{Name: "ray-inference", Namespace: "test-namespace"},
+			Status: rayv1.RayServiceStatuses{
+				ActiveServiceStatus: rayv1.RayServiceStatus{RayClusterName: "ray-inference-abcde"},
+			},
+		},
+		&rayv1.RayService{ObjectMeta: metav1.ObjectMeta{Name: "starting", Namespace: "test-namespace"}},
+	)
+
+	build := func(cluster string) (*rayv1.RayJob, error) {
+		taskCtx := dummyRayTaskContext(dummyExistingClusterTaskTemplate(cluster, &plugins.RayJob{}), resourceRequirements, nil, "", serviceAccount)
+		obj, err := rayJobResourceHandler{}.BuildResource(context.TODO(), taskCtx)
+		if err != nil {
+			return nil, err
+		}
+		return obj.(*rayv1.RayJob), nil
+	}
+
+	t.Run("rayservice name resolves to its active cluster", func(t *testing.T) {
+		job, err := build("ray-inference")
+		require.NoError(t, err)
+		assert.Equal(t, map[string]string{"ray.io/cluster": "ray-inference-abcde"}, job.Spec.ClusterSelector)
+	})
+
+	t.Run("rayservice without active cluster is retried", func(t *testing.T) {
+		_, err := build("starting")
+		assert.ErrorContains(t, err, "no active RayCluster")
+		assert.NotContains(t, err.Error(), "BadTaskSpecification")
+	})
+
+	t.Run("plain cluster name passes through", func(t *testing.T) {
+		job, err := build("shared")
+		require.NoError(t, err)
+		assert.Equal(t, map[string]string{"ray.io/cluster": "shared"}, job.Spec.ClusterSelector)
+	})
+}