From 22044cd807fd31b995b1a0ba7f8db0db056d02ac Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 8 Oct 2026 09:01:58 +0200 Subject: [PATCH] feat(ray): resolve RayService names for existing-cluster tasks KubeRay resolves clusterSelector as a RayCluster name, and a RayService names its cluster - and replaces it on every upgrade. When ray.existing_cluster names a RayService in the task namespace, the plugin substitutes the RayService's active RayCluster. Other names pass through. A RayService without an active cluster returns a retryable error. BuildResource gets no Kubernetes client, so the plugin creates one on first use from the in-cluster or remote cluster config. Flyte's service account needs get on rayservices.ray.io. Signed-off-by: Christian --- flyteplugins/go/tasks/plugins/k8s/ray/ray.go | 56 +++++++++++++++++++ .../go/tasks/plugins/k8s/ray/ray_test.go | 55 ++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go index 6e9bca982d..85e83850d6 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go +++ b/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) { diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go index 14ffc800dc..3c52247ff2 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go +++ b/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) + }) +}