Ray

The Ray plugin lets you run Ray jobs natively on Kubernetes. Flyte provisions a transient Ray cluster for each task execution using KubeRay and tears it down on completion.

When to use this plugin

Distributed Python workloads (parallel computation, data processing)
ML training with Ray Train or hyperparameter tuning with Ray Tune
Ray Serve inference workloads
Any workload that benefits from Ray’s actor model or task parallelism

KubeRay’s readiness and liveness probes for the head and worker pods shell out to wget to poll the raylet and GCS health endpoints. If the image has no wget, both probes fail permanently with wget: command not found, the head pod never reports Ready, the workers stay parked in their wait-gcs-ready init container, and the job is never submitted. Install it with .with_apt_packages("wget"), or use a base image that already ships it.

Configuration

Create a RayJobConfig and pass it as plugin_config to a TaskEnvironment:

        
    
import flyte
from flyteplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
)

`RayJobConfig` parameters

Parameter	Type	Description
`worker_node_config`	`List[WorkerNodeConfig]`	Required. List of worker group configurations
`head_node_config`	`HeadNodeConfig`	Head node configuration (optional)
`enable_autoscaling`	`bool`	Enable Ray autoscaler (default: `False`)
`autoscaler_options`	`AutoscalerOptionsConfig`	Tune the autoscaler sidecar. Has no effect unless `enable_autoscaling` is `True`
`runtime_env`	`dict`	Ray runtime environment (pip packages, env vars, etc.)
`address`	`str`	Connect to an existing Ray cluster instead of provisioning one
`shutdown_after_job_finishes`	`bool`	Shut down the cluster after the job completes (default: `False`)
`ttl_seconds_after_finished`	`int`	Seconds to keep the cluster after completion before cleanup

`WorkerNodeConfig` parameters

Parameter	Type	Description
`group_name`	`str`	Required. Name of this worker group
`replicas`	`int`	Required. Number of worker replicas
`min_replicas`	`int`	Minimum replicas (for autoscaling)
`max_replicas`	`int`	Maximum replicas (for autoscaling)
`ray_start_params`	`Dict[str, str]`	Ray start parameters for workers
`requests`	`Resources`	Resource requests per worker
`limits`	`Resources`	Resource limits per worker
`pod_template`	`PodTemplate`	Full pod template (mutually exclusive with `requests`/`limits`)

`HeadNodeConfig` parameters

Parameter	Type	Description
`ray_start_params`	`Dict[str, str]`	Ray start parameters for the head node
`requests`	`Resources`	Resource requests for the head node
`limits`	`Resources`	Resource limits for the head node
`pod_template`	`PodTemplate`	Full pod template (mutually exclusive with `requests`/`limits`)

The head node runs the Ray dashboard, which starts nine subprocess modules on top of GCS and the raylet. Give it 2 CPU and 4Gi of memory:

        
    
ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(requests=flyte.Resources(cpu=2, memory="4Gi")),
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
)

Under-provisioning the head node is a common cause of a cluster that never becomes ready: if the dashboard cannot start, ray start --head fails and KubeRay recycles the head pod in a loop. A head pod at 1 CPU and 1000Mi has been observed failing this way.

`AutoscalerOptionsConfig` parameters

Setting enable_autoscaling=True runs the Ray autoscaler with KubeRay’s defaults. Pass autoscaler_options to tune it:

        
    
import flyte
from flyteplugins.ray import AutoscalerOptionsConfig, RayJobConfig, WorkerNodeConfig

ray_config = RayJobConfig(
    worker_node_config=[
        WorkerNodeConfig(group_name="ray-group", replicas=1, min_replicas=1, max_replicas=5)
    ],
    enable_autoscaling=True,
    autoscaler_options=AutoscalerOptionsConfig(
        upscaling_mode=AutoscalerOptionsConfig.UpscalingMode.CONSERVATIVE,
        idle_timeout_seconds=120,
        resources=flyte.Resources(cpu=("500m", "1"), memory=("512Mi", "1Gi")),
    ),
)

Parameter	Type	Description
`upscaling_mode`	`AutoscalerOptionsConfig.UpscalingMode`	Rate limiting on adding nodes. `CONSERVATIVE` holds the number of pending worker pods to at most the current cluster size. `DEFAULT` and `AGGRESSIVE` are the same setting: no rate limit. Leaving the field unset, or passing `UNSPECIFIED`, also means no rate limit
`idle_timeout_seconds`	`int`	Seconds a node may sit idle before the autoscaler removes it (default: 60). An explicit `0` is dropped and the default applies
`image`	`str`	Container image for the autoscaler sidecar
`env`	`Dict[str, str]`	Environment variables for the autoscaler container
`resources`	`Resources`	Requests and limits for the autoscaler sidecar

Every field is optional. Leaving upscaling_mode, idle_timeout_seconds, image, or env unset keeps the KubeRay default. resources is the exception: whenever you pass autoscaler_options, whatever you give for resources replaces the sidecar’s default 500m CPU and 512Mi memory requests and limits outright. Omit it and the sidecar runs with no requests or limits at all; set only a request and it also loses the default limits. Set resources explicitly, on both sides. It accepts tuples to set a request and a limit together, as in flyte.Resources(cpu=("500m", "1")).

The options do not switch autoscaling on

autoscaler_options only configures the autoscaler sidecar, and the sidecar is created only when enable_autoscaling is True. Passing options on their own changes nothing.

Connecting to an existing cluster

To connect to an existing Ray cluster instead of provisioning a new one, set the address parameter:

        
    
ray_config = RayJobConfig(
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
    address="ray://existing-cluster:10001",
)

Examples

The following example shows how to configure Ray in a TaskEnvironment. Flyte automatically provisions a Ray cluster for each task using this configuration:

ray_example.py
                
                    
                
            
                
            
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-ray",
#    "ray[default]==2.46.0"
# ]
# main = "hello_ray_nested"
# params = "3"
# ///

import asyncio
import typing

import ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

import flyte.remote
import flyte.storage

@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=2)],
    runtime_env={"pip": ["numpy", "pandas"]},
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=300,
)

image = (
    flyte.Image.from_debian_base(name="ray")
    .with_apt_packages("wget")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray", "pip", "mypy")
)

task_env = flyte.TaskEnvironment(
    name="hello_ray", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
ray_env = flyte.TaskEnvironment(
    name="ray_env",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")),
    depends_on=[task_env],
)

@task_env.task()
async def hello_ray():
    await asyncio.sleep(20)
    print("Hello from the Ray task!")

@ray_env.task
async def hello_ray_nested(n: int = 3) -> typing.List[int]:
    print("running ray task")
    t = asyncio.create_task(hello_ray())
    futures = [f.remote(i) for i in range(n)]
    res = ray.get(futures)
    await t
    return res

if __name__ == "__main__":
    flyte.init_from_config()
    r = flyte.run(hello_ray_nested)
    print(r.name)
    print(r.url)
    r.wait()

The next example demonstrates how Flyte can create ephemeral Ray clusters and run a subtask that connects to an existing Ray cluster:

ray_existing_example.py
                
                    
                
            
                
            
# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.0.0b52",
#    "flyteplugins-ray",
#    "ray[default]==2.46.0"
# ]
# main = "create_ray_cluster"
# params = ""
# ///

import os
import typing

import ray
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

import flyte.storage

@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=2)],
    enable_autoscaling=False,
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=3600,
)

image = (
    flyte.Image.from_debian_base(name="ray")
    .with_apt_packages("wget")
    .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray")
)

task_env = flyte.TaskEnvironment(
    name="ray_client", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image
)
ray_env = flyte.TaskEnvironment(
    name="ray_cluster",
    plugin_config=ray_config,
    image=image,
    resources=flyte.Resources(cpu=(2, 4), memory=("2000Mi", "4000Mi")),
    depends_on=[task_env],
)

@task_env.task()
async def hello_ray(cluster_ip: str) -> typing.List[int]:
    """
    Run a simple Ray task that connects to an existing Ray cluster.
    """
    ray.init(address=f"ray://{cluster_ip}:10001")
    futures = [f.remote(i) for i in range(5)]
    res = ray.get(futures)
    return res

@ray_env.task
async def create_ray_cluster() -> str:
    """
    Create a Ray cluster and return the head node IP address.
    """
    print("creating ray cluster")
    cluster_ip = os.getenv("MY_POD_IP")
    if cluster_ip is None:
        raise ValueError("MY_POD_IP environment variable is not set")
    return f"{cluster_ip}"

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(create_ray_cluster)
    run.wait()
    print("run url:", run.url)
    print("cluster created, running ray task")
    print("ray address:", run.outputs()[0])
    run = flyte.run(hello_ray, cluster_ip=run.outputs()[0])
    print("run url:", run.url)

API reference

See the Ray API reference for full details.