Skip to content

Commit 35df30d

Browse files
committed
add section for setting k8s_pod for ray head & worker pods
1 parent 02e0dc9 commit 35df30d

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

examples/ray_plugin/ray_plugin/ray_example.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,9 +85,21 @@ def ray_task(n: int) -> typing.List[int]:
8585

8686
# %% [markdown]
8787
# :::{note}
88-
# The `Resources` section here is utilized to specify the resources allocated to the worker nodes.
89-
# :::
90-
#
88+
# The `Resources` section here is utilized to specify the resources allocated to all Ray components like head, worker & driver pods.
89+
# %% [markdown]
90+
# For more fine grain resource allocation, you are able to set the resources for Ray head & worker pods by setting the K8s pod argument for the respective config.
91+
# %%
92+
from flytekit.models.task import K8sPod
93+
from flytekit.core.resources import pod_spec_from_resources
94+
95+
ray_config = RayJobConfig(
96+
head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}, k8s_pod=K8sPod(pod_spec=pod_spec_from_resources(k8s_pod_name="ray-head", requests=Resources(cpu="4",mem="5Gi")))),
97+
worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=1, k8s_pod=K8sPod(pod_spec=pod_spec_from_resources(k8s_pod_name="ray-worker", requests=Resources(cpu="1",mem="1Gi"))))],
98+
runtime_env={"pip": ["numpy", "pandas"]}, # or runtime_env="./requirements.txt"
99+
enable_autoscaling=True,
100+
shutdown_after_job_finishes=True,
101+
ttl_seconds_after_finished=3600,
102+
)
91103
# Lastly, define a workflow to call the Ray task.
92104
# %%
93105
@workflow

0 commit comments

Comments
 (0)