fix OSS resume

Netflix · Mar 12, 2024 · 4f9e3c9 · 4f9e3c9
1 parent c97026c
commit 4f9e3c9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/metaflow/clone_util.py b/metaflow/clone_util.py
@@ -7,6 +7,7 @@ def clone_task_helper(
     clone_run_id,
     run_id,
     step_name,
+    clone_task_id,
     task_id,
     flow_datastore,
     metadata_service,
@@ -22,7 +23,11 @@ def clone_task_helper(
     )
     output.init_task()
 
-    origin_run_id, origin_step_name, origin_task_id = clone_run_id, step_name, task_id
+    origin_run_id, origin_step_name, origin_task_id = (
+        clone_run_id,
+        step_name,
+        clone_task_id,
+    )
     # 2. initialize origin datastore
     origin = None
     if origin_ds_set:

diff --git a/metaflow/runtime.py b/metaflow/runtime.py
@@ -246,6 +246,7 @@ def clone_task(self, step_name, task_id):
             self._clone_run_id,
             self._run_id,
             step_name,
+            task_id,  # origin_task_id
             task_id,
             self._flow_datastore,
             self._metadata,

diff --git a/metaflow/task.py b/metaflow/task.py
@@ -290,14 +290,15 @@ def clone_only(
             raise MetaflowInternalError(
                 "task.clone_only needs a valid clone_origin_task value."
             )
-        origin_run_id, _, _ = clone_origin_task.split("/")
+        origin_run_id, _, origin_task_id = clone_origin_task.split("/")
 
         # If we actually have to do the clone ourselves, proceed...
         clone_task_helper(
             self.flow.name,
             origin_run_id,
             run_id,
             step_name,
+            origin_task_id,
             task_id,
             self.flow_datastore,
             self.metadata,