Add Ray Autoscaler to the Flyte-Ray plugin

[Yicheng-Lu-llll]

TL;DR

NOTE：The ray CI test failed because we need to first merge flyteorg/flyte#4363 to update flytecli.

Currently, the Flyte-Ray plugin utilizes Rayjob. However, there are cases where Rayjob may require an autoscaler.

1. After completing a workload with Rayjob, a user might want to retain all the information, logs, past tasks, and actor execution history for a period. As of now, Ray lacks a mechanism to persist these data, necessitating the continuous operation of the Ray cluster even after workload completion. With an autoscaler, the Ray cluster will maintain only the head pod while scaling down all worker pods.
2. User does not need to pre estimate the need of the resource. Autoscaler will care everything.

So, this PR adds Ray Autoscaler config to the Flyte-Ray plugin. Also see flyteorg/flyte#4363.

btw, This PR adds shutdown_after_job_finishes and ttl_seconds_after_finished.

• ttl_seconds_after_finished specifies the number of seconds after which the RayCluster will be deleted after the RayJob finishes.
• shutdown_after_job_finishes specifies whether the RayCluster should be deleted after the RayJob finishes.

Below is an example:

import typing
import ray
from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
from flytekit import Resources, task, workflow

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

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}),
    # The behavior will be:
    # 1. Create a head node and 0 worker node.
    # 2. The worker node will be scaled to 2.
    # 3. The worker node will be scaled to 0 and will be terminated later(120s).
    worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=0, min_replicas=0, max_replicas=2)],
    enable_autoscaling=True,
    shutdown_after_job_finishes=True,
    # ttl_seconds_after_finished specifies the number of seconds after which the RayCluster will be deleted after the RayJob finishes.
    ttl_seconds_after_finished=120,
)

@task(
    task_config=ray_config,
    requests=Resources(mem="1Gi", cpu="2"),
)
def ray_task() -> int:
    import time

    # Import placement group APIs.
    from ray.util.placement_group import (
        placement_group,
        placement_group_table,
        remove_placement_group,
    )
    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

    # use the `placement_group` to trigger autoscaling.
    # each node have two cpus, so this will create 2 workers when executing.
    # Once the task is done, the two workers will be removed.
    pg = placement_group([{"CPU": 2} for i in range(3)])
    ray.get(pg.ready(), timeout=100)

    return 1

@workflow
def ray_workflow() -> int:
    return ray_task()

init(replicas=0), only head:

(flytekit) ubuntu@ip-172-31-2-249:~/flyte/flytekit$ kubectl get pod -n flytesnacks-development
NAME                                                    READY   STATUS    RESTARTS   AGE
f50cd4d1a5ceb40bc9aa-n0-0-raycluster-pbvx9-head-szdgn   2/2     Running   0          3s

task executing(max_replicas=2):

(flytekit) ubuntu@ip-172-31-2-249:~/flyte/flytekit$ kubectl get pod -n flytesnacks-development
NAME                                                      READY   STATUS    RESTARTS   AGE
f50cd4d1a5ceb40bc9aa-n0-0-raycluster-pbvx9-head-szdgn     2/2     Running   0          54s
ceb40bc9aa-n0-0-raycluster-pbvx9-worker-ray-group-llzp7   1/1     Running   0          14s
ceb40bc9aa-n0-0-raycluster-pbvx9-worker-ray-group-mqrb6   1/1     Running   0          14s

finsh(min_replicas=0):

(flytekit) ubuntu@ip-172-31-2-249:~/flyte/flytekit$ kubectl get pod -n flytesnacks-development
NAME                                                      READY   STATUS        RESTARTS   AGE
f50cd4d1a5ceb40bc9aa-n0-0-raycluster-pbvx9-head-szdgn     2/2     Running       0          2m16s
ceb40bc9aa-n0-0-raycluster-pbvx9-worker-ray-group-llzp7   1/1     Terminating   0          96s
ceb40bc9aa-n0-0-raycluster-pbvx9-worker-ray-group-mqrb6   1/1     Terminating   0          96s

After TTL 120s:

(flytekit) ubuntu@ip-172-31-2-249:~/flyte/flytekit$ kubectl get pod -n flytesnacks-development
No resources found in flytesnacks-development namespace.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

How did you fix the bug, make the feature etc. Link to any design docs etc

Tracking Issue

flyteorg/flyte#4187

Follow-up issue

NA
OR
https://github.com/flyteorg/flyte/issues/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (7794d3a) 85.53% compared to head (150ab4d) 85.53%.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1937   +/-   ##
=======================================
  Coverage   85.53%   85.53%           
=======================================
  Files         309      309           
  Lines       23460    23475   +15     
  Branches     3630     3630           
=======================================
+ Hits        20066    20079   +13     
- Misses       2752     2753    +1     
- Partials      642      643    +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[pingsutw]

LGTM, could we add a small test here?

[pingsutw]

qq: why in_tree? Does Ray have other autoscaling strategies?

[Yicheng-Lu-llll]

Only have one autoscaling strategy. I have changed name to enable_autoscaling. Thank you!

[eapolinario]

does this mean that by default the rayjob will not be reclaimed by kuberay once the job finishes?

[Yicheng-Lu-llll]

According to https://github.com/ray-project/kuberay/blob/master/ray-operator/config/samples/ray_v1alpha1_rayjob.yaml#L7, Yes.

[eapolinario]

this is a no-op if shutdown_after_job_finishes is set to False, right?

[Yicheng-Lu-llll]

Yes.

[eapolinario]

This needs a new release of flyteidl.