feat(ray): support ReusePolicy — run Ray tasks on shared, reusable RayClusters

examples/plugins/ray_reusable_cluster.py

@@ -0,0 +1,81 @@
+"""Reusable shared Ray cluster.
+
+A driver task spawns two Ray tasks concurrently. Because both Ray tasks belong to the same
+environment declared with ``flyte.ReusePolicy``, they run on ONE long-lived, shared RayCluster
+instead of each paying a full cluster cold-start: the first task to arrive creates the cluster,
+the second reuses it. The cluster is shut down automatically once it has been idle for ``idle_ttl``
+seconds.
+"""
+
+import asyncio
+import typing
+
+import ray
+from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
+
+import flyte
+
+
+@ray.remote
+def square(x: int) -> int:
+    return x * x
+
+
+ray_config = RayJobConfig(
+    head_node_config=HeadNodeConfig(),
+    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
+    enable_autoscaling=False,
+)
+
+image = (
+    flyte.Image.from_debian_base(name="flyte")
+    .with_apt_packages("wget")
+    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray")
+)
+
+# The reusable Ray environment. Every task in this environment with an identical configuration
+# shares one long-lived RayCluster (keyed by the environment's identity). `idle_ttl` shuts the
+# cluster down after it sits idle with no jobs for that many seconds.
+ray_env = flyte.TaskEnvironment(
+    name="reusable_ray_env",
+    plugin_config=ray_config,
+    image=image,
+    resources=flyte.Resources(cpu=(1, 2), memory=("2000Mi", "4000Mi")),
+    reusable=flyte.ReusePolicy(replicas=1, idle_ttl=600),
+)
+
+# A plain environment for the driver that orchestrates the Ray tasks. The driver itself does not
+# need Ray — it just fans out to the Ray tasks. `depends_on=[ray_env]` deploys the Ray environment
+# alongside the driver so the driver can spawn tasks in it.
+driver_env = flyte.TaskEnvironment(name="flyte_driver_env", image=image, depends_on=[ray_env])
+
+
+@ray_env.task
+async def sum_of_squares(n: int) -> int:
+    """Runs on the shared RayCluster and fans the work out to Ray workers via `@ray.remote`..."""
+    print(f"running Ray task on the shared cluster (n={n})")
+    results = ray.get([square.remote(i) for i in range(n)])
+    print(f"partial results: {results}")
+    return sum(results)
+
+
+@driver_env.task
+async def driver(a: int = 5, b: int = 8) -> typing.List[int]:
+    """Spawns two Ray tasks concurrently; both bind to the same reusable RayCluster.
+
+    The first task to reach the cluster creates it; the second reuses it (no second cold start).
+    """
+    first, second = await asyncio.gather(
+        sum_of_squares(a),
+        sum_of_squares(b),
+    )
+    print(f"driver results: sum_of_squares({a})={first}, sum_of_squares({b})={second}")
+    return [first, second]
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(driver, a=5, b=8)
+    print("run url:", run.url)
+    run.wait()
+    print("phase:", run.action.phase if run.action else "unknown")

plugins/ray/src/flyteplugins/ray/task.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, Optional
 
 import flyte
+import flyte.errors
 import yaml
 from flyte import PodTemplate, Resources
 from flyte.extend import (
@@ -131,6 +132,7 @@ class RayFunctionTask(AsyncFunctionTaskTemplate):
     task_type: str = "ray"
     plugin_config: RayJobConfig
     debuggable: bool = True
+    supports_reuse_policy: typing.ClassVar[bool] = True
 
     async def pre(self, *args, **kwargs) -> Dict[str, Any]:
         init_params = {"address": self.plugin_config.address}
@@ -200,7 +202,27 @@ def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]:
             shutdown_after_job_finishes=cfg.shutdown_after_job_finishes,
         )
 
-        return MessageToDict(ray_job)
+        custom = MessageToDict(ray_job)
+
+        if self.reusable is not None:
+            # `replicas` is the number of shared clusters; only 1 is supported for now.
+            if self.reusable.max_replicas != 1:
+                raise flyte.errors.RuntimeUserError(
+                    "BadConfiguration",
+                    f"Reusable Ray tasks currently support exactly 1 replica (one shared RayCluster); "
+                    f"got replicas={self.reusable.replicas}. Use ReusePolicy(replicas=1).",
+                )
+            idle_ttl = self.reusable.idle_ttl
+            scaledown_ttl = self.reusable.get_scaledown_ttl()
+            custom["reusePolicy"] = {
+                "parallelism": self.reusable.concurrency,
+                "min_replica_count": self.reusable.min_replicas,
+                "replica_count": self.reusable.max_replicas,
+                "ttl_seconds": idle_ttl.total_seconds() if idle_ttl else None,  # type: ignore[union-attr]
+                "scaledown_ttl_seconds": scaledown_ttl.total_seconds() if scaledown_ttl else None,
+            }
+
+        return custom
 
 
 TaskPluginRegistry.register(config_type=RayJobConfig, plugin=RayFunctionTask)

plugins/ray/tests/test_task.py

@@ -9,7 +9,7 @@
 from google.protobuf.json_format import MessageToDict, ParseDict
 from kubernetes.client import V1Container, V1PodSpec, V1ResourceRequirements
 
-from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
+from flyteplugins.ray.task import HeadNodeConfig, RayFunctionTask, RayJobConfig, WorkerNodeConfig
 
 
 @pytest.fixture
@@ -209,3 +209,38 @@ def test_pod_template_without_resources_is_unchanged(sctx):
     container = _primary_container(ray_job.ray_cluster.worker_group_spec[0].k8s_pod)
     assert container["args"] == ["wut update-aws-credentials-file default"]
     assert container["resources"]["requests"]["cpu"] == "15000m"
+
+
+def test_custom_config_records_reuse_policy(sctx):
+    task = RayFunctionTask(
+        name="t",
+        interface=None,
+        func=lambda: None,
+        plugin_config=RayJobConfig(worker_node_config=[]),
+        reusable=flyte.ReusePolicy(replicas=1, idle_ttl=600),
+    )
+    custom = task.custom_config(sctx)
+    assert custom["reusePolicy"] == {
+        "parallelism": 1,
+        "min_replica_count": 1,
+        "replica_count": 1,
+        "ttl_seconds": 600,
+        "scaledown_ttl_seconds": 30,
+    }
+    # The rest of the spec still parses as a RayJob (the extra field is ignored by the proto).
+    assert "rayCluster" in custom
+
+
+def test_custom_config_rejects_multiple_reuse_replicas(sctx):
+    import flyte.errors
+
+    task = RayFunctionTask(
+        name="t",
+        interface=None,
+        func=lambda: None,
+        plugin_config=RayJobConfig(worker_node_config=[]),
+        reusable=flyte.ReusePolicy(replicas=(1, 3)),
+    )
+    # `replicas` is the number of shared clusters; only 1 is supported for now.
+    with pytest.raises(flyte.errors.RuntimeUserError, match="exactly 1 replica"):
+        task.custom_config(sctx)

src/flyte/_internal/runtime/task_serde.py

@@ -288,6 +288,8 @@ def get_proto_task(
             env = task.parent_env()
             if env is not None:
                 env_name = env.name
+        if getattr(task, "supports_reuse_policy", False):
+            return task_template
         return add_reusable(task_template, task.reusable, serialize_context.code_bundle, env_name)
 
     return task_template

src/flyte/_task_environment.py

@@ -145,7 +145,11 @@ async def my_task():
     def __post_init__(self) -> None:
         super().__post_init__()
         if self.reusable is not None and self.plugin_config is not None:
-            raise ValueError("Cannot set plugin_config when environment is reusable.")
+            from flyte.extend import TaskPluginRegistry
+
+            plugin_cls: Optional[type] = TaskPluginRegistry.find(type(self.plugin_config))
+            if plugin_cls is None or not getattr(plugin_cls, "supports_reuse_policy", False):
+                raise ValueError("Cannot set plugin_config when environment is reusable.")
         if self.reusable and not isinstance(self.reusable, ReusePolicy):
             raise TypeError(f"Expected reusable to be of type ReusePolicy, got {type(self.reusable)}")
         if self.cache and not isinstance(self.cache, (str, Cache)):