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

Merging #1093 (444bd49) into master (a968a5a) will decrease coverage by 18.42%.
The diff coverage is n/a.

@@             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     

Impacted Files	Coverage Δ
flytekit/core/python_function_task.py	52.63% <ø> (-35.15%)	⬇️
plugins/setup.py	0.00% <ø> (ø)
flytekit/deck/__init__.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/core/constants.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/core/resources.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/core/tracked_abc.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/types/file/__init__.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/types/numpy/__init__.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/types/pickle/__init__.py	0.00% <0.00%> (-100.00%)	⬇️
flytekit/types/schema/__init__.py	0.00% <0.00%> (-100.00%)	⬇️
... and 132 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? 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.