feat: add ClusterPlugin interface + migrate Ray to shared clusters by pingsutw · Pull Request #7460 · flyteorg/flyte

pingsutw · 2026-05-30T02:19:29Z

Tracking issue

Related to shared-cluster execution for Ray (and future cluster-backed plugins).

Why are the changes needed?

Today every Ray task produces a RayJob with an embedded RayClusterSpec. KubeRay spins up a brand-new Ray cluster per task, runs the job, and tears the cluster down. Two Ray tasks with identical cluster configuration cannot share a cluster, so we pay full cluster cold-start cost on every task.

This PR introduces a generic ClusterPlugin model so a plugin can run a job on top of a shared, long-lived cluster keyed by the hash of the job-spec proto. Two tasks whose spec hashes to the same name reuse one cluster; the job is created only after the cluster is ready. Ray is migrated to this model as the first consumer.

What changes were proposed in this pull request?

1. New `ClusterPlugin` interface (`flyteplugins/.../pluginmachinery/k8s/plugin.go`)

Sits alongside the existing Plugin interface. Instead of one BuildResource, it declares two resources plus a readiness check:

type ClusterPlugin interface {
    GetClusterName(ctx, taskCtx) (string, error)               // deterministic name = hash(spec proto)
    BuildClusterResource(ctx, taskCtx) (client.Object, error)  // the shared cluster
    BuildClusterIdentityResource(ctx, taskMeta) (client.Object, error)
    IsClusterReady(ctx, pluginContext, cluster) (bool, error)
    BuildJobResource(ctx, taskCtx, clusterName) (client.Object, error)  // job bound to the cluster
    BuildJobIdentityResource(ctx, taskMeta) (client.Object, error)
    GetJobPhase(ctx, pluginContext, job) (pluginsCore.PhaseInfo, error)
    GetProperties() PluginProperties
}

PluginEntry gains ClusterPlugin + ClusterResourceToWatch; RegisterK8sPlugin validates exactly one of Plugin/ClusterPlugin is set.

2. New `ClusterPluginManager` (`executor/pkg/plugin/k8s/cluster_plugin_manager.go`)

A sibling to PluginManager that implements pluginsCore.Plugin and drives a 3-phase state machine:

NotStarted   ── create-or-reuse shared cluster (IsAlreadyExists tolerated) ─▶ ClusterWait
ClusterWait  ── IsClusterReady? ──no──▶ stay Initializing
                               ──yes─▶ submit job (ClusterSelector) ────────▶ JobStarted
JobStarted   ── GetJobPhase ──▶ Running / Success(+read outputs) / Failure

Key invariants:

The cluster is created without owner reference or finalizer so it outlives any single task execution (cleanup delegated to KubeRay's idle autoscaler / external GC).
The job is owned normally (owner ref + finalizer); Abort/Finalize act only on the job, never the shared cluster.
ClusterName is persisted in plugin state so it round-trips across reconcile rounds.

3. Architecture

                         ┌───────────────────────────────────────────────┐
                         │              ClusterPluginManager              │
 TaskAction (ray) ─────▶ │  Handle() state machine                       │
                         │                                               │
                         │  ┌─ NotStarted ─────────────────────────────┐ │
                         │  │ name = GetClusterName(spec)  (proto hash) │ │
                         │  │ GET cluster; if absent:                   │ │
                         │  │   BuildClusterResource → Create           │ │
                         │  │   (no ownerRef / no finalizer)            │ │
                         │  └──────────────────────────────────────────┘ │
                         │  ┌─ ClusterWait ────────────────────────────┐ │
                         │  │ IsClusterReady(cluster)?                  │ │
                         │  │   ready → BuildJobResource(name) → Create │ │
                         │  │          (ownerRef + finalizer)           │ │
                         │  └──────────────────────────────────────────┘ │
                         │  ┌─ JobStarted ─────────────────────────────┐ │
                         │  │ GetJobPhase(job) → phase + outputs        │ │
                         │  └──────────────────────────────────────────┘ │
                         └───────────────────────────────────────────────┘
        run A (spec X) ┐
                       ├─ hash(X) = ray-<h> ──▶ ONE RayCluster ray-<h>
        run B (spec X) ┘                          ▲   ▲
                                       RayJob A ───┘   └─── RayJob B
                                       (ClusterSelector: ray.io/cluster=ray-<h>)

4. Ray migration (`flyteplugins/.../plugins/k8s/ray/ray.go`)

GetClusterName: "ray-" + truncate(pbhash(RayJob proto)), DNS1123-safe. Identical configs → identical name.
BuildClusterResource: standalone RayCluster labeled ray.io/cluster=<name>, autoscaling enabled.
BuildJobResource: RayJob with RayClusterSpec=nil, ClusterSelector={ray.io/cluster: <name>}, ShutdownAfterJobFinishes=false (cluster must outlive the job).
GetJobPhase: the former GetTaskPhase body, unchanged.
Old Plugin methods (BuildResource, BuildIdentityResource, IsTerminal, GetCompletionTime) removed; registered via ClusterPlugin.

5. Two correctness fixes discovered during live testing

Because the Ray driver runs on the shared head pod, run identity could not come from the cluster pod's env (that env belongs to whichever task created the cluster). Two fixes make each job report under its own run:

Strip run-scoped env from the shared cluster pods — RUN_NAME, ACTION_NAME, _U_RUN_BASE, FLYTE_INTERNAL_* are removed from the head/worker containers so the cluster is run-agnostic (user-supplied env preserved).
Inject per-job identity into the RayJob runtime_env.env_vars — Ray applies runtime_env to the driver process, so each job's driver sees its own RUN_NAME/_U_RUN_BASE/FLYTE_INTERNAL_* and reports under its own run. User-supplied runtime_env values are never overwritten.

Dataflow (per run)

flyte run ──▶ RunService creates TaskAction(ray)
   └─▶ ClusterPluginManager.Handle
        ├─ GetClusterName(spec)  ────────────▶ ray-<hash>   (same for identical specs)
        ├─ ensure RayCluster ray-<hash>      ─▶ KubeRay provisions head+workers (shared)
        ├─ IsClusterReady → ready
        └─ submit RayJob  (ClusterSelector=ray-<hash>,
                           runtime_env.env_vars = THIS run's identity)
                 └─▶ KubeRay binds job to existing cluster
                       └─▶ driver on head pod reads its OWN identity from runtime_env
                             └─▶ job runs → Complete → outputs read → TaskAction Completed

How was this patch tested?

Unit tests (go test ./flyteplugins/.../ray/... ./executor/pkg/plugin/...): registry validation, cluster/job split, readiness, deterministic naming, env stripping, and runtime_env injection (including "does not overwrite user-supplied env"). All green; go vet clean.

End-to-end in the local devbox (k3s + KubeRay v1.6.1):

Example (reusable_ray.py):

import asyncio, typing, ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
import flyte

@ray.remote
def f(x): return x * x

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=1)],
    runtime_env={"pip": ["numpy", "ty"]},
    enable_autoscaling=False,
)
image = (flyte.Image.from_debian_base(name="flyte")
         .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray"))
ray_env = flyte.TaskEnvironment(name="ray_env", plugin_config=ray_config, image=image,
                                resources=flyte.Resources(cpu=(2, 4), memory=("3000Mi", "5000Mi")))

@ray_env.task
async def hello_ray_nested(n: int = 3) -> typing.List[int]:
    futures = [f.remote(i) for i in range(n)]
    return ray.get(futures)

Ran it multiple times concurrently against the devbox. All runs succeeded on a single shared cluster, each with its own identity:

$ kubectl -n flyte get rayclusters,rayjobs
NAME                                     STATUS   AGE
raycluster.ray.io/ray-jw97ftk5ayrzerhm   ready    95s          # ← ONE shared cluster

NAME                                      JOB STATUS   RAY CLUSTER NAME       END TIME
rayjob.ray.io/rb9t6k2vcvgqlv4cc42n-a0-0   SUCCEEDED    ray-jw97ftk5ayrzerhm   ...:51Z
rayjob.ray.io/rgnnxfd7swdks6wh8kjz-a0-0   SUCCEEDED    ray-jw97ftk5ayrzerhm   ...:58Z
rayjob.ray.io/rkgt2ls7kdrtqkrsmx7b-a0-0   SUCCEEDED    ray-jw97ftk5ayrzerhm   ...:08Z

# Each RayJob's runtime_env carries its OWN run identity (the fix):
rb9t6k2vcvgqlv4cc42n-a0-0 → RUN_NAME=rb9t6k2vcvgqlv4cc42n  cluster=ray-jw97ftk5ayrzerhm
rgnnxfd7swdks6wh8kjz-a0-0 → RUN_NAME=rgnnxfd7swdks6wh8kjz  cluster=ray-jw97ftk5ayrzerhm
rkgt2ls7kdrtqkrsmx7b-a0-0 → RUN_NAME=rkgt2ls7kdrtqkrsmx7b  cluster=ray-jw97ftk5ayrzerhm

All TaskActions (and nested python children) reached Completed.

Labels

added (ClusterPlugin interface + manager, Ray shared-cluster support)
changed (Ray plugin migrated from embedded-cluster to shared-cluster model)