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