add section for setting k8s_pod for ray head & worker pods

Signed-off-by: Jan Fiedler <[email protected]>

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