diff --git a/examples/ray_plugin/ray_plugin/ray_example.py b/examples/ray_plugin/ray_plugin/ray_example.py index ec836cb91..d35700260 100644 --- a/examples/ray_plugin/ray_plugin/ray_example.py +++ b/examples/ray_plugin/ray_plugin/ray_example.py @@ -93,13 +93,28 @@ def ray_task(n: int) -> typing.List[int]: 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"))))], + 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