ob-metaflow 2.11.4.9__py2.py3-none-any.whl → 2.11.8.1__py2.py3-none-any.whl

metaflow/plugins/env_escape/utils.py

@@ -13,12 +13,12 @@ def get_methods(class_object):
     for base_class in mros:
         all_attributes.update(base_class.__dict__)
     for name, attribute in all_attributes.items():
-        if hasattr(attribute, "__call__"):
-            all_methods[name] = inspect.getdoc(attribute)
-        elif isinstance(attribute, staticmethod):
+        if isinstance(attribute, staticmethod):
             all_methods["___s___%s" % name] = inspect.getdoc(attribute)
         elif isinstance(attribute, classmethod):
             all_methods["___c___%s" % name] = inspect.getdoc(attribute)
+        elif hasattr(attribute, "__call__"):
+            all_methods[name] = inspect.getdoc(attribute)
     return all_methods
 
 

metaflow/plugins/gcp/gcp_secret_manager_secrets_provider.py

@@ -90,7 +90,7 @@ class GcpSecretManagerSecretsProvider(SecretsProvider):
         # The latter two forms require METAFLOW_GCP_SECRET_MANAGER_PREFIX to be set.
 
         match_full = re.match(r"^projects/\d+/secrets/([\w\-]+)(/versions/([\w\-]+))?$", secret_id)
-        match_partial = re.match(r"^[\w\-]+(/versions/[\w\-]+)?$", secret_id)
+        match_partial = re.match(r"^([\w\-]+)(/versions/[\w\-]+)?$", secret_id)
         if match_full:
             # Full path
             env_var_name = match_full.group(1)
@@ -107,9 +107,10 @@ class GcpSecretManagerSecretsProvider(SecretsProvider):
                 raise ValueError(
                     f"Cannot use simple secret_id without setting METAFLOW_GCP_SECRET_MANAGER_PREFIX. {GCP_SECRET_MANAGER_PREFIX}"
                 )
-            if match_partial.group(1):
+            if match_partial.group(2):
                 # With version specified
                 full_secret_name = f"{GCP_SECRET_MANAGER_PREFIX}{secret_id}"
+                env_var_name = match_partial.group(1)
             else:
                 # No version specified, use latest
                 full_secret_name = (

metaflow/plugins/pypi/conda_environment.py

@@ -282,6 +282,15 @@ class CondaEnvironment(MetaflowEnvironment):
             # Match PyPI and Conda python versions with the resolved environment Python.
             environment["pypi"]["python"] = environment["conda"]["python"] = env_python
 
+            # When using `Application Default Credentials` for private GCP
+            # PyPI registries, the usage of environment variable `GOOGLE_APPLICATION_CREDENTIALS`
+            # demands that `keyrings.google-artifactregistry-auth` has to be installed
+            # and available in the underlying python environment.
+            if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
+                environment["conda"]["packages"][
+                    "keyrings.google-artifactregistry-auth"
+                ] = ">=1.1.1"
+
         # Z combinator for a recursive lambda
         deep_sort = (lambda f: f(f))(
             lambda f: lambda obj: (

metaflow/plugins/pypi/pip.py

@@ -92,7 +92,14 @@ class Pip(object):
                 # so using @branch as a version acts as expected.
                 vcs_info = dl_info.get("vcs_info")
                 if vcs_info:
-                    res["url"] = "{vcs}+{url}@{commit_id}".format(**vcs_info, **res)
+                    subdirectory = dl_info.get("subdirectory")
+                    res["url"] = "{vcs}+{url}@{commit_id}{subdir_str}".format(
+                        **vcs_info,
+                        **res,
+                        subdir_str="#subdirectory=%s" % subdirectory
+                        if subdirectory
+                        else ""
+                    )
                     # used to deduplicate the storage location in case wheel does not
                     # build with enough unique identifiers.
                     res["hash"] = vcs_info["commit_id"]
@@ -270,9 +277,17 @@ class Pip(object):
                         prefix,
                         "pip3",
                         "--disable-pip-version-check",
-                        "--no-input",
                         "--no-color",
                     ]
+                    # credentials are being determined from the JSON file referenced by
+                    # the GOOGLE_APPLICATION_CREDENTIALS environment variable and are
+                    # probably injected dynamically via `keyrings.google-artifactregistry-auth`
+                    # Thus, we avoid passing `--no-input` in this case.
+                    + (
+                        ["--no-input"]
+                        if os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is None
+                        else []
+                    )
                     + (["--isolated"] if isolated else [])
                     + args,
                     stderr=subprocess.PIPE,

metaflow/runtime.py

@@ -13,6 +13,7 @@ import subprocess
 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 .debug import debug
 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,
@@ -46,6 +48,7 @@ PROGRESS_INTERVAL = 300  # s
 # executing a flow. These are prefetched during the resume operation by
 # leveraging the TaskDataStoreSet.
 PREFETCH_DATA_ARTIFACTS = ["_foreach_stack", "_task_ok", "_transition"]
+RESUME_POLL_SECONDS = 60
 
 # Runtime must use logsource=RUNTIME_LOG_SOURCE for all loglines that it
 # formats according to mflog. See a comment in mflog.__init__
@@ -75,8 +78,8 @@ class NativeRuntime(object):
         max_workers=MAX_WORKERS,
         max_num_splits=MAX_NUM_SPLITS,
         max_log_size=MAX_LOG_SIZE,
+        resume_identifier=None,
     ):
-
         if run_id is None:
             self._run_id = metadata.new_run_id()
         else:
@@ -101,11 +104,13 @@ class NativeRuntime(object):
         self._entrypoint = entrypoint
         self.event_logger = event_logger
         self._monitor = monitor
+        self._resume_identifier = resume_identifier
 
         self._clone_run_id = clone_run_id
         self._clone_only = clone_only
         self._clone_steps = {} if clone_steps is None else clone_steps
         self._reentrant = reentrant
+        self._run_url = None
 
         self._origin_ds_set = None
         if clone_run_id:
@@ -184,7 +189,8 @@ class NativeRuntime(object):
             origin_ds_set=self._origin_ds_set,
             decos=decos,
             logger=self._logger,
-            **kwargs
+            resume_identifier=self._resume_identifier,
+            **kwargs,
         )
 
     @property
@@ -192,25 +198,25 @@ class NativeRuntime(object):
         return self._run_id
 
     def persist_constants(self, task_id=None):
-        task = self._new_task("_parameters", task_id=task_id)
-        if not task.is_cloned:
-            task.persist(self._flow)
-        self._params_task = task.path
-        self._is_cloned[task.path] = task.is_cloned
+        self._params_task = self._new_task("_parameters", task_id=task_id)
+        if not self._params_task.is_cloned:
+            self._params_task.persist(self._flow)
 
-    def execute(self):
-        run_url = (
+        self._is_cloned[self._params_task.path] = self._params_task.is_cloned
+
+    def print_workflow_info(self):
+        self._run_url = (
             "%s/%s/%s" % (UI_URL.rstrip("/"), self._flow.name, self._run_id)
             if UI_URL
             else None
         )
 
-        if run_url:
+        if self._run_url:
             self._logger(
                 "Workflow starting (run-id %s), see it in the UI at %s"
                 % (
                     self._run_id,
-                    run_url,
+                    self._run_url,
                 ),
                 system_msg=True,
             )
@@ -219,10 +225,88 @@ class NativeRuntime(object):
                 "Workflow starting (run-id %s):" % self._run_id, system_msg=True
             )
 
+    def _should_skip_clone_only_execution(self):
+        if self._clone_only and self._params_task:
+            if self._params_task.resume_done():
+                return True, "Resume already complete. Skip clone-only execution."
+            if not self._params_task.is_resume_leader():
+                return (
+                    True,
+                    "Not resume leader under resume execution. Skip clone-only execution.",
+                )
+        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,
+            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)
 
         if self._params_task:
-            self._queue_push("start", {"input_paths": [self._params_task]})
+            self._queue_push("start", {"input_paths": [self._params_task.path]})
         else:
             self._queue_push("start", {})
 
@@ -303,9 +387,9 @@ class NativeRuntime(object):
 
         # assert that end was executed and it was successful
         if ("end", ()) in self._finished:
-            if run_url:
+            if self._run_url:
                 self._logger(
-                    "Done! See the run in the UI at %s" % run_url,
+                    "Done! See the run in the UI at %s" % self._run_url,
                     system_msg=True,
                 )
             else:
@@ -316,6 +400,7 @@ class NativeRuntime(object):
                 "cloned; no new tasks executed!",
                 system_msg=True,
             )
+            self._params_task.mark_resume_done()
         else:
             raise MetaflowInternalError(
                 "The *end* step was not successful by the end of flow."
@@ -711,6 +796,7 @@ class Task(object):
         ubf_iter=None,
         join_type=None,
         task_id=None,
+        resume_identifier=None,
     ):
 
         self.step = step
@@ -749,32 +835,38 @@ class Task(object):
         self.datastore_sysroot = flow_datastore.datastore_root
         self._results_ds = None
 
+        # Only used in clone-only resume.
+        self._is_resume_leader = None
+        self._resume_done = None
+        self._resume_identifier = resume_identifier
+
         origin = None
         if clone_run_id and may_clone:
             origin = self._find_origin_task(clone_run_id, join_type)
         if origin and origin["_task_ok"]:
             # At this point, we know we are going to clone
             self._is_cloned = True
+
+            task_id_exists_already = False
+            task_completed = False
             if reentrant:
                 # A re-entrant clone basically allows multiple concurrent processes
                 # to perform the clone at the same time to the same new run id. Let's
                 # assume two processes A and B both simultaneously calling
                 # `resume --reentrant --run-id XX`.
-                # For each task that is cloned, we want to guarantee that:
-                #   - one and only one of A or B will do the actual cloning
-                #   - the other process (or other processes) will block until the cloning
-                #     is complete.
-                # This ensures that the rest of the clone algorithm can proceed as normal
-                # and also guarantees that we only write once to the datastore and
-                # metadata.
+                # We want to guarantee that:
+                #   - All incomplete tasks are cloned exactly once.
+                # To achieve this, we will select a resume leader and let it clone the
+                # entire execution graph. This ensures that we only write once to the
+                # datastore and metadata.
                 #
-                # To accomplish this, we use the cloned task's task-id as the "key" to
-                # synchronize on. We then try to "register" this new task-id (or rather
-                # the full pathspec <run>/<step>/<taskid>) with the metadata service
-                # which will indicate if we actually registered it or if it existed
-                # already. If we did manage to register it, we are the "elected cloner"
+                # We use the cloned _parameter task's task-id as the "key" to synchronize
+                # on. We try to "register" this new task-id (or rather the full pathspec
+                # <run>/<step>/<taskid>) with the metadata service which will indicate
+                # if we actually registered it or if it existed already. If we did manage
+                # to register it (_parameter task), we are the "elected resume leader"
                 # in essence and proceed to clone. If we didn't, we just wait to make
-                # sure the task is fully done (ie: the clone is finished).
+                # sure the entire clone execution is fully done (ie: the clone is finished).
                 if task_id is not None:
                     # Sanity check -- this should never happen. We cannot allow
                     # for explicit task-ids because in the reentrant case, we use the
@@ -783,6 +875,15 @@ class Task(object):
                         "Reentrant clone-only resume does not allow for explicit task-id"
                     )
 
+                if resume_identifier:
+                    self.log(
+                        "Resume identifier is %s." % resume_identifier,
+                        system_msg=True,
+                    )
+                else:
+                    raise MetaflowInternalError(
+                        "Reentrant clone-only resume needs a resume identifier."
+                    )
                 # We will use the same task_id as the original task
                 # to use it effectively as a synchronization key
                 clone_task_id = origin.task_id
@@ -798,56 +899,124 @@ class Task(object):
 
                 # If _get_task_id returns True it means the task already existed, so
                 # we wait for it.
-                self._wait_for_clone = self._get_task_id(clone_task_id)
+                task_id_exists_already = self._get_task_id(clone_task_id)
+
+                # We may not have access to task datastore on first resume attempt, but
+                # on later resume attempt, we should check if the resume task is complete
+                # or not. This is to fix the issue where the resume leader was killed
+                # unexpectedly during cloning and never mark task complete.
+                try:
+                    task_completed = self.results["_task_ok"]
+                except DataException as e:
+                    pass
             else:
-                self._wait_for_clone = False
                 self._get_task_id(task_id)
 
             # Store the mapping from current_pathspec -> origin_pathspec which
             # will be useful for looking up origin_ds_set in find_origin_task.
             self.clone_pathspec_mapping[self._path] = origin.pathspec
             if self.step == "_parameters":
-                # We don't put _parameters on the queue so we either clone it or wait
-                # for it.
-                if not self._wait_for_clone:
+                # In the _parameters task, we need to resolve who is the resume leader.
+                self._is_resume_leader = False
+                resume_leader = None
+
+                if task_id_exists_already:
+                    # If the task id already exists, we need to check if current task is the resume leader in previous attempt.
+                    ds = self._flow_datastore.get_task_datastore(
+                        self.run_id, self.step, self.task_id
+                    )
+                    if not ds["_task_ok"]:
+                        raise MetaflowInternalError(
+                            "Externally cloned _parameters task did not succeed"
+                        )
+
+                    # Check if we should be the resume leader (maybe from previous attempt).
+                    # To avoid the edge case where the resume leader is selected but has not
+                    # yet written the _resume_leader metadata, we will wait for a few seconds.
+                    # We will wait for resume leader for at most 3 times.
+                    for resume_leader_wait_retry in range(3):
+
+                        if ds.has_metadata("_resume_leader", add_attempt=False):
+                            resume_leader = ds.load_metadata(
+                                ["_resume_leader"], add_attempt=False
+                            )["_resume_leader"]
+                            self._is_resume_leader = resume_leader == resume_identifier
+                        else:
+                            self.log(
+                                "Waiting for resume leader to be selected. Sleeping ...",
+                                system_msg=True,
+                            )
+                            time.sleep(3)
+                else:
+                    # If the task id does not exist, current task is the resume leader.
+                    resume_leader = resume_identifier
+                    self._is_resume_leader = True
+
+                if reentrant:
+                    if resume_leader:
+                        self.log(
+                            "Resume leader is %s." % resume_leader,
+                            system_msg=True,
+                        )
+                    else:
+                        raise MetaflowInternalError(
+                            "Can not determine the resume leader in distributed resume mode."
+                        )
+
+                if self._is_resume_leader:
+                    self.log(
+                        "Selected as the reentrant clone leader.",
+                        system_msg=True,
+                    )
                     # Clone in place without relying on run_queue.
                     self.new_attempt()
                     self._ds.clone(origin)
+                    # Set the resume leader be the task that calls the resume (first task to clone _parameters task).
+                    # We will always set resume leader regardless whether we are in distributed resume case or not.
+                    if resume_identifier:
+                        self._ds.save_metadata(
+                            {"_resume_leader": resume_identifier}, add_attempt=False
+                        )
+
                     self._ds.done()
                 else:
-                    # TODO: There is a bit of a duplication with the task.py
-                    # clone_only function here
-                    self.log(
-                        "Waiting for clone of _parameters step to occur...",
-                        system_msg=True,
-                    )
+                    # Wait for the resume leader to complete
                     while True:
-                        try:
-                            ds = self._flow_datastore.get_task_datastore(
-                                self.run_id, self.step, self.task_id
-                            )
-                            if not ds["_task_ok"]:
-                                raise MetaflowInternalError(
-                                    "Externally cloned _parameters task did not succeed"
-                                )
+                        ds = self._flow_datastore.get_task_datastore(
+                            self.run_id, self.step, self.task_id
+                        )
+
+                        # Check if resume is complete. Resume leader will write the done file.
+                        self._resume_done = ds.has_metadata(
+                            "_resume_done", add_attempt=False
+                        )
+
+                        if self._resume_done:
                             break
-                        except DataException:
-                            self.log("Sleeping for 5s...", system_msg=True)
-                            # No need to get fancy with the sleep here.
-                            time.sleep(5)
-                    self.log("_parameters clone successful", system_msg=True)
+
+                        self.log(
+                            "Waiting for resume leader to complete. Sleeping for %ds..."
+                            % RESUME_POLL_SECONDS,
+                            system_msg=True,
+                        )
+                        time.sleep(RESUME_POLL_SECONDS)
+                    self.log(
+                        "_parameters clone completed by resume leader", system_msg=True
+                    )
             else:
-                # For non parameter steps
+                # Only leader can reach non-parameter steps in resume.
+
                 # Store the origin pathspec in clone_origin so this can be run
                 # as a task by the runtime.
                 self.clone_origin = origin.pathspec
                 # Save a call to creating the results_ds since its same as origin.
                 self._results_ds = origin
-                if self._wait_for_clone:
+
+                # If the task is already completed in new run, we don't need to clone it.
+                self._should_skip_cloning = task_completed
+                if self._should_skip_cloning:
                     self.log(
-                        "Waiting for the successful cloning of results "
-                        "of a previously run task %s (this may take some time)"
-                        % self.clone_origin,
+                        "Skip cloning of previously run task %s" % self.clone_origin,
                         system_msg=True,
                     )
                 else:
@@ -916,6 +1085,31 @@ class Task(object):
         self._logger(msg, head=prefix, system_msg=system_msg, timestamp=timestamp)
         sys.stdout.flush()
 
+    def is_resume_leader(self):
+        assert (
+            self.step == "_parameters"
+        ), "Only _parameters step can check resume leader."
+        return self._is_resume_leader
+
+    def resume_done(self):
+        assert (
+            self.step == "_parameters"
+        ), "Only _parameters step can check wheather resume is complete."
+        return self._resume_done
+
+    def mark_resume_done(self):
+        assert (
+            self.step == "_parameters"
+        ), "Only _parameters step can mark resume as done."
+        assert self.is_resume_leader(), "Only resume leader can mark resume as done."
+
+        # Mark the resume as done. This is called at the end of the resume flow and after
+        # the _parameters step was successfully cloned, so we need to 'dangerously' save
+        # this done file, but the risk should be minimal.
+        self._ds._dangerous_save_metadata_post_done(
+            {"_resume_done": True}, add_attempt=False
+        )
+
     def _get_task_id(self, task_id):
         already_existed = True
         if self.ubf_context == UBF_CONTROL:
@@ -1040,8 +1234,8 @@ class Task(object):
         return self._is_cloned
 
     @property
-    def wait_for_clone(self):
-        return self._wait_for_clone
+    def should_skip_cloning(self):
+        return self._should_skip_cloning
 
     def persist(self, flow):
         # this is used to persist parameters before the start step
@@ -1185,7 +1379,6 @@ class CLIArgs(object):
 
 class Worker(object):
     def __init__(self, task, max_logs_size):
-
         self.task = task
         self._proc = self._launch()
 
@@ -1229,8 +1422,6 @@ class Worker(object):
             # 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.wait_for_clone:
-                args.command_options["clone-wait-only"] = True
         else:
             # decorators may modify the CLIArgs object in-place
             for deco in self.task.decos:

metaflow/sidecar/sidecar.py

@@ -13,14 +13,24 @@ class Sidecar(object):
         if t is not None and t.get_worker() is not None:
             self._has_valid_worker = True
         self.sidecar_process = None
+        # Whether to send msg in a thread-safe fashion.
+        self._threadsafe_send_enabled = False
 
     def start(self):
         if not self.is_active and self._has_valid_worker:
             self.sidecar_process = SidecarSubProcess(self._sidecar_type)
 
+    def enable_threadsafe_send(self):
+        self._threadsafe_send_enabled = True
+
+    def disable_threadsafe_send(self):
+        self._threadsafe_send_enabled = False
+
     def send(self, msg):
         if self.is_active:
-            self.sidecar_process.send(msg)
+            self.sidecar_process.send(
+                msg, thread_safe_send=self._threadsafe_send_enabled
+            )
 
     def terminate(self):
         if self.is_active: