From 3f58b60523cb2b528e65b702912194c68bdc18b4 Mon Sep 17 00:00:00 2001 From: Jan Fiedler Date: Fri, 3 Jan 2025 12:11:11 +0100 Subject: [PATCH] add section for setting k8s_pod for ray head & worker pods Signed-off-by: Jan Fiedler --- examples/ray_plugin/ray_plugin/ray_example.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/examples/ray_plugin/ray_plugin/ray_example.py b/examples/ray_plugin/ray_plugin/ray_example.py index 185652b0d..ec836cb91 100644 --- a/examples/ray_plugin/ray_plugin/ray_example.py +++ b/examples/ray_plugin/ray_plugin/ray_example.py @@ -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