retrieve entire graph then clone (#1728)

add more print

profile time and threads

workers with clone only command

include origin_ds_set and s3 batch write

clean up pr for review

remove wait-only flag because it is no longer used

address comments

fix OSS resume

Author: darinyu

metaflow/cli.py

@@ -473,13 +473,6 @@ def echo_unicode(line, **kwargs):
     help="Pathspec of the origin task for this task to clone. Do "
     "not execute anything.",
 )
-@click.option(
-    "--clone-wait-only/--no-clone-wait-only",
-    default=False,
-    show_default=True,
-    help="If specified, waits for an external process to clone the task",
-    hidden=True,
-)
 @click.option(
     "--clone-run-id",
     default=None,
@@ -519,7 +512,6 @@ def step(
     retry_count=None,
     max_user_code_retries=None,
     clone_only=None,
-    clone_wait_only=False,
     clone_run_id=None,
     decospecs=None,
     ubf_context="none",
@@ -575,7 +567,6 @@ def step(
             task_id,
             clone_only,
             retry_count,
-            wait_only=clone_wait_only,
         )
     else:
         task.run_step(
@@ -802,7 +793,10 @@ def resume(
     write_run_id(run_id_file, runtime.run_id)
     runtime.print_workflow_info()
     runtime.persist_constants()
-    runtime.execute()
+    if clone_only:
+        runtime.clone_original_run()
+    else:
+        runtime.execute()
 
 
 @tracing.cli_entrypoint("cli/run")
@@ -830,7 +824,7 @@ def run(
     decospecs=None,
     run_id_file=None,
     user_namespace=None,
-    **kwargs
+    **kwargs,
 ):
     if user_namespace is not None:
         namespace(user_namespace or None)
@@ -985,7 +979,7 @@ def start(
     pylint=None,
     event_logger=None,
     monitor=None,
-    **deco_options
+    **deco_options,
 ):
     global echo
     if quiet:

metaflow/clone_util.py

@@ -0,0 +1,71 @@
+import time
+from .metadata import MetaDatum
+
+
+def clone_task_helper(
+    flow_name,
+    clone_run_id,
+    run_id,
+    step_name,
+    clone_task_id,
+    task_id,
+    flow_datastore,
+    metadata_service,
+    origin_ds_set=None,
+    attempt_id=0,
+):
+    # 1. initialize output datastore
+    output = flow_datastore.get_task_datastore(
+        run_id, step_name, task_id, attempt=attempt_id, mode="w"
+    )
+    output.init_task()
+
+    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:
+        origin = origin_ds_set.get_with_pathspec(
+            "{}/{}/{}".format(origin_run_id, origin_step_name, origin_task_id)
+        )
+    else:
+        origin = flow_datastore.get_task_datastore(
+            origin_run_id, origin_step_name, origin_task_id
+        )
+    metadata_tags = ["attempt_id:{0}".format(attempt_id)]
+    output.clone(origin)
+    _ = metadata_service.register_task_id(
+        run_id,
+        step_name,
+        task_id,
+        attempt_id,
+    )
+    metadata_service.register_metadata(
+        run_id,
+        step_name,
+        task_id,
+        [
+            MetaDatum(
+                field="origin-task-id",
+                value=str(origin_task_id),
+                type="origin-task-id",
+                tags=metadata_tags,
+            ),
+            MetaDatum(
+                field="origin-run-id",
+                value=str(origin_run_id),
+                type="origin-run-id",
+                tags=metadata_tags,
+            ),
+            MetaDatum(
+                field="attempt",
+                value=str(attempt_id),
+                type="attempt",
+                tags=metadata_tags,
+            ),
+        ],
+    )
+    output.done()

metaflow/runtime.py

@@ -13,6 +13,7 @@
 from datetime import datetime
 from io import BytesIO
 from functools import partial
+from concurrent import futures
 
 from metaflow.datastore.exceptions import DataException
 
@@ -30,6 +31,7 @@
 from .decorators import flow_decorators
 from .mflog import mflog, RUNTIME_LOG_SOURCE
 from .util import to_unicode, compress_list, unicode_type
+from .clone_util import clone_task_helper
 from .unbounded_foreach import (
     CONTROL_TASK_TAG,
     UBF_CONTROL,
@@ -188,7 +190,7 @@ def _new_task(self, step, input_paths=None, **kwargs):
             decos=decos,
             logger=self._logger,
             resume_identifier=self._resume_identifier,
-            **kwargs
+            **kwargs,
         )
 
     @property
@@ -234,6 +236,65 @@ def _should_skip_clone_only_execution(self):
                 )
         return False, None
 
+    def clone_task(self, step_name, task_id):
+        self._logger(
+            "Cloning task from {}/{}/{}/{} to {}/{}/{}/{}".format(
+                self._flow.name,
+                self._clone_run_id,
+                step_name,
+                task_id,
+                self._flow.name,
+                self._run_id,
+                step_name,
+                task_id,
+            ),
+            system_msg=True,
+        )
+        clone_task_helper(
+            self._flow.name,
+            self._clone_run_id,
+            self._run_id,
+            step_name,
+            task_id,  # origin_task_id
+            task_id,
+            self._flow_datastore,
+            self._metadata,
+            origin_ds_set=self._origin_ds_set,
+        )
+
+    def clone_original_run(self):
+        (
+            should_skip_clone_only_execution,
+            skip_reason,
+        ) = self._should_skip_clone_only_execution()
+        if should_skip_clone_only_execution:
+            self._logger(skip_reason, system_msg=True)
+            return
+        self._metadata.start_run_heartbeat(self._flow.name, self._run_id)
+        self._logger(
+            "Start cloning original run: {}/{}".format(
+                self._flow.name, self._clone_run_id
+            ),
+            system_msg=True,
+        )
+
+        inputs = []
+
+        for task_ds in self._origin_ds_set:
+            _, step_name, task_id = task_ds.pathspec.split("/")
+            if task_ds["_task_ok"] and step_name != "_parameters":
+                inputs.append((step_name, task_id))
+
+        with futures.ThreadPoolExecutor(max_workers=self._max_workers) as executor:
+            all_tasks = [
+                executor.submit(self.clone_task, step_name, task_id)
+                for (step_name, task_id) in inputs
+            ]
+            _, _ = futures.wait(all_tasks)
+        self._logger("Cloning original run is done", system_msg=True)
+        self._params_task.mark_resume_done()
+        self._metadata.stop_heartbeat()
+
     def execute(self):
         (
             should_skip_clone_only_execution,
@@ -1361,8 +1422,6 @@ def _launch(self):
             # disabling sidecars for cloned tasks due to perf reasons
             args.top_level_options["event-logger"] = "nullSidecarLogger"
             args.top_level_options["monitor"] = "nullSidecarMonitor"
-            if self.task.should_skip_cloning:
-                args.command_options["clone-wait-only"] = True
         else:
             # decorators may modify the CLIArgs object in-place
             for deco in self.task.decos:

metaflow/task.py

@@ -20,6 +20,7 @@
 )
 from .unbounded_foreach import UBF_CONTROL
 from .util import all_equal, get_username, resolve_identity, unicode_type
+from .clone_util import clone_task_helper
 from .metaflow_current import current
 from metaflow.tracing import get_trace_id
 from metaflow.util import namedtuple_with_defaults
@@ -284,57 +285,43 @@ def clone_only(
         task_id,
         clone_origin_task,
         retry_count,
-        wait_only=False,
     ):
         if not clone_origin_task:
             raise MetaflowInternalError(
                 "task.clone_only needs a valid clone_origin_task value."
             )
-        if wait_only:
-            print("Not cloning anything in wait_only mode.")
-            return
+        origin_run_id, _, origin_task_id = clone_origin_task.split("/")
 
+        msg = {
+            "task_id": task_id,
+            "msg": "Cloning task from {}/{}/{}/{} to {}/{}/{}/{}".format(
+                self.flow.name,
+                origin_run_id,
+                step_name,
+                origin_task_id,
+                self.flow.name,
+                run_id,
+                step_name,
+                task_id,
+            ),
+            "step_name": step_name,
+            "run_id": run_id,
+            "flow_name": self.flow.name,
+            "ts": round(time.time()),
+        }
+        self.event_logger.log(msg)
         # If we actually have to do the clone ourselves, proceed...
-        # 1. initialize output datastore
-        output = self.flow_datastore.get_task_datastore(
-            run_id, step_name, task_id, attempt=0, mode="w"
-        )
-
-        output.init_task()
-
-        origin_run_id, origin_step_name, origin_task_id = clone_origin_task.split("/")
-        # 2. initialize origin datastore
-        origin = self.flow_datastore.get_task_datastore(
-            origin_run_id, origin_step_name, origin_task_id
-        )
-        metadata_tags = ["attempt_id:{0}".format(retry_count)]
-        output.clone(origin)
-        self.metadata.register_metadata(
+        clone_task_helper(
+            self.flow.name,
+            origin_run_id,
             run_id,
             step_name,
+            origin_task_id,
             task_id,
-            [
-                MetaDatum(
-                    field="origin-task-id",
-                    value=str(origin_task_id),
-                    type="origin-task-id",
-                    tags=metadata_tags,
-                ),
-                MetaDatum(
-                    field="origin-run-id",
-                    value=str(origin_run_id),
-                    type="origin-run-id",
-                    tags=metadata_tags,
-                ),
-                MetaDatum(
-                    field="attempt",
-                    value=str(retry_count),
-                    type="attempt",
-                    tags=metadata_tags,
-                ),
-            ],
+            self.flow_datastore,
+            self.metadata,
+            attempt_id=retry_count,
         )
-        output.done()
 
     def _finalize_control_task(self):
         # Update `_transition` which is expected by the NativeRuntime.