Ray Task Support

[pingsutw]

TL;DR

• Design Doc
• Slide

Flyte ray task contains Ray Cluster Config that will be used by propeller to create a Ray cluster.

There are three ways to run Ray job

1. submit to the existing cluster

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

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

@task(task_config=RayJobConfig(
    address=<RAY_CLUSTER_ADDRESS>
    runtime_env={"pip": ["numpy", "pandas"]})
)
def ray_task() -> typing.List[int]:
    futures = [f.remote(i) for i in range(5)]
    return ray.get(futures)

2. Running the below task locally, and flyte will create a local Ray cluster
3. Running the below task in Flyte Cluster, Flyte will create a RayJob CR in Kubernetes.

@task(task_config=RayJobConfig(
    worker_node_config=[WorkerNodeConfig(group_name="test-group", replicas=10)],
)
def ray_task() -> typing.List[int]:
    futures = [f.remote(i) for i in range(5)]
    return ray.get(futures)

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#2641

Follow-up issue

NA

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 68.25%. Comparing base (a968a5a) to head (444bd49).
Report is 1018 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1093       +/-   ##
===========================================
- Coverage   86.67%   68.25%   -18.43%     
===========================================
  Files         269      287       +18     
  Lines       25074    25814      +740     
  Branches     2826     2884       +58     
===========================================
- Hits        21734    17620     -4114     
- Misses       2871     7717     +4846     
- Partials      469      477        +8     

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

[kumare3]

looks pretty good

[wild-endeavor]

trying to be better about pr comments - since this is a new feature/integration, can you write a couple examples and put them in the pr description? we'll need the examples for the documentation in any case so it's not wasted work.