From f82bc2c70ca044e9d5ba443f7994a5b909b89fb9 Mon Sep 17 00:00:00 2001
From: Jan Fiedler <jan@union.ai>
Date: Tue, 16 Jul 2024 14:38:06 -0700
Subject: [PATCH] test

Signed-off-by: Jan Fiedler <jan@union.ai>
---
 .../flytekit-ray/flytekitplugins/ray/task.py  | 33 ++++++++++++++++---
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/plugins/flytekit-ray/flytekitplugins/ray/task.py b/plugins/flytekit-ray/flytekitplugins/ray/task.py
index f0ee542b32..eeb166592a 100644
--- a/plugins/flytekit-ray/flytekitplugins/ray/task.py
+++ b/plugins/flytekit-ray/flytekitplugins/ray/task.py
@@ -4,7 +4,12 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Optional
 
-from flytekitplugins.ray.models import HeadGroupSpec, RayCluster, RayJob, WorkerGroupSpec
+from flytekitplugins.ray.models import (
+    HeadGroupSpec,
+    RayCluster,
+    RayJob,
+    WorkerGroupSpec,
+)
 from google.protobuf.json_format import MessageToDict
 
 from flytekit import lazy_module
@@ -46,11 +51,19 @@ class RayFunctionTask(PythonFunctionTask):
     _RAY_TASK_TYPE = "ray"
 
     def __init__(self, task_config: RayJobConfig, task_function: Callable, **kwargs):
-        super().__init__(task_config=task_config, task_type=self._RAY_TASK_TYPE, task_function=task_function, **kwargs)
+        super().__init__(
+            task_config=task_config,
+            task_type=self._RAY_TASK_TYPE,
+            task_function=task_function,
+            **kwargs,
+        )
         self._task_config = task_config
 
     def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
-        ray.init(address=self._task_config.address)
+        ray.init(
+            address=self._task_config.address,
+            runtime_env={"working_dir": "/root"},
+        )
         return user_params
 
     def post_execute(self, user_params: ExecutionParameters, rval: Any) -> Any:
@@ -62,9 +75,19 @@ def get_custom(self, settings: SerializationSettings) -> Optional[Dict[str, Any]
 
         ray_job = RayJob(
             ray_cluster=RayCluster(
-                head_group_spec=HeadGroupSpec(cfg.head_node_config.ray_start_params) if cfg.head_node_config else None,
+                head_group_spec=(
+                    HeadGroupSpec(cfg.head_node_config.ray_start_params)
+                    if cfg.head_node_config
+                    else None
+                ),
                 worker_group_spec=[
-                    WorkerGroupSpec(c.group_name, c.replicas, c.min_replicas, c.max_replicas, c.ray_start_params)
+                    WorkerGroupSpec(
+                        c.group_name,
+                        c.replicas,
+                        c.min_replicas,
+                        c.max_replicas,
+                        c.ray_start_params,
+                    )
                     for c in cfg.worker_node_config
                 ],
             ),