Skip to content

Commit

Permalink
add section for setting k8s_pod for ray head & worker pods
Browse files Browse the repository at this point in the history
  • Loading branch information
fiedlerNr9 committed Jan 3, 2025
1 parent 02e0dc9 commit 35df30d
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions examples/ray_plugin/ray_plugin/ray_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,21 @@ def ray_task(n: int) -> typing.List[int]:

# %% [markdown]
# :::{note}
# The `Resources` section here is utilized to specify the resources allocated to the worker nodes.
# :::
#
# The `Resources` section here is utilized to specify the resources allocated to all Ray components like head, worker & driver pods.
# %% [markdown]
# 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.
# %%
from flytekit.models.task import K8sPod
from flytekit.core.resources import pod_spec_from_resources

ray_config = RayJobConfig(
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")))),
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"))))],
runtime_env={"pip": ["numpy", "pandas"]}, # or runtime_env="./requirements.txt"
enable_autoscaling=True,
shutdown_after_job_finishes=True,
ttl_seconds_after_finished=3600,
)
# Lastly, define a workflow to call the Ray task.
# %%
@workflow
Expand Down

0 comments on commit 35df30d

Please sign in to comment.