flyteorg / flyteorg/flyte

Support Ray AutoscalerOptions in RayCluster spec

Open
#7,005 4 comments 1 reaction 1 assignee Claimed by @0yukali0 View on GitHub
flyte flyte2
Dominant language
Go
Stars
7.5k
Forks
886
Avg merge
1d 14h
Merged PRs (30d)
120

Description

## Motivation

Users who need to configure Ray autoscaler behavior (e.g., environment variables, upscaling mode, resource limits) currently cannot do so because the Flyte Ray integration doesn't pass `AutoscalerOptions` through to the KubeRay `RayCluster` CR.

Currently, only `enable_autoscaling` (bool) and per-worker `min_replicas`/`max_replicas` are supported. The KubeRay `RayCluster` CRD has a rich [`autoscalerOptions`](https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/configuring-autoscaling.html) field that allows configuring:

- **`upscalingMode`** — `Default`, `Aggressive`, or `Conservative`
- **`env`** — environment variables for the autoscaler sidecar container
- **`resources`** — CPU/memory requests and limits for the autoscaler container
- **`image`** — custom autoscaler container image
- **`idleTimeoutSeconds`** — how long idle workers are kept before scale-down

## Proposal

Expose these options end-to-end across the Flyte stack:

### 1. flyteidl (protobuf definitions)

Add an `AutoscalerOptions` message to [`flyteidl/plugins/ray.proto`](https://github.com/flyteorg/flyte/blob/master/flyteidl/protos/flyteidl/plugins/ray.proto):

```protobuf
message AutoscalerOptions {
string upscaling_mode = 1; // "Default", "Aggressive", "Conservative"
int32 idle_timeout_seconds = 2;
repeated EnvVar env = 3; // autoscaler sidecar env vars
string image = 4; // custom autoscaler image
Resources resources = 5; // autoscaler container resources
}
```

Add a field to the existing `RayCluster` message:

```protobuf
message RayCluster {
...
AutoscalerOptions autoscaler_options = N;
}
```

### 2. flytepropeller (backend Ray plugin)

Update the [Ray plugin handler](https://github.com/flyteorg/flyte/tree/master/flyteplugins/go/tasks/plugins/k8s/ray) to:
- Read `AutoscalerOptions` from the task's protobuf `RayCluster` message
- Map it onto the KubeRay `RayCluster` CR's `.spec.autoscalerOptions`

### 3. flytekit / flyte-sdk (Python SDK)

In the Ray plugin (`plugins/ray/src/flyteplugins/ray/task.py` in [flyte-sdk](https://github.com/flyteorg/flyte-sdk)):

- Add an `AutoscalerOptions` dataclass:
```python
@dataclass
class AutoscalerOptions:
upscaling_mode: Optional[str] = None # "Default", "Aggressive", "Conservative"
idle_timeout_seconds: Optional[int] = None
env: Optional[Dict[str, str]] = None
image: Optional[str] = None
resources: Optional[Resources] = None
```

- Add `autoscaler_options: Optional[AutoscalerOptions] = None` to `RayJobConfig`
- Update `RayFunctionTask.custom_config()` to serialize autoscaler options into the `RayCluster` protobuf

Currently `RayJobConfig` only has:
```python
enable_autoscaling: bool = False
```

The new field would sit alongside it:
```python
enable_autoscaling: bool = False
autoscaler_options: Optional[AutoscalerOptions] = None
```

## Goal

Allow users to write:
```python
RayJobConfig(
enable_autoscaling=True,
autoscaler_options=AutoscalerOptions(
upscaling_mode="Conservative",
idle_timeout_seconds=120,
image="rayproject/ray:2.9.0",
),
worker_node_config=[...],
)
```

And have KubeRay receive the full `autoscalerOptions` on the resulting `RayCluster` CR.

## Are you sure this issue hasn't been raised already?

Yes

## Have you read the Code of Conduct?

Yes

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.