diff --git a/examples/plugins/ray_reusable_cluster.py b/examples/plugins/ray_reusable_cluster.py new file mode 100644 index 000000000..2b045db5d --- /dev/null +++ b/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") diff --git a/plugins/ray/src/flyteplugins/ray/task.py b/plugins/ray/src/flyteplugins/ray/task.py index c4c70ad1a..729eb9d8a 100644 --- a/plugins/ray/src/flyteplugins/ray/task.py +++ b/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) diff --git a/plugins/ray/tests/test_task.py b/plugins/ray/tests/test_task.py index f416ba668..7567edb83 100644 --- a/plugins/ray/tests/test_task.py +++ b/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) diff --git a/src/flyte/_internal/runtime/task_serde.py b/src/flyte/_internal/runtime/task_serde.py index 254fd8632..b48c8b823 100644 --- a/src/flyte/_internal/runtime/task_serde.py +++ b/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 diff --git a/src/flyte/_task_environment.py b/src/flyte/_task_environment.py index f648094d6..d5f4378f3 100644 --- a/src/flyte/_task_environment.py +++ b/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)):