ob-metaflow 2.12.10.1rc2__py2.py3-none-any.whl → 2.12.11.0__py2.py3-none-any.whl

metaflow/plugins/argo/argo_workflows_cli.py

@@ -5,11 +5,14 @@ import re
 import sys
 from hashlib import sha1
 
-from metaflow import Run, JSONType, current, decorators, parameters
-from metaflow.client.core import get_metadata
-from metaflow.exception import MetaflowNotFound
+from metaflow import JSONType, Run, current, decorators, parameters
 from metaflow._vendor import click
-from metaflow.exception import MetaflowException, MetaflowInternalError
+from metaflow.client.core import get_metadata
+from metaflow.exception import (
+    MetaflowException,
+    MetaflowInternalError,
+    MetaflowNotFound,
+)
 from metaflow.metaflow_config import (
     ARGO_WORKFLOWS_UI_URL,
     KUBERNETES_NAMESPACE,
@@ -167,6 +170,12 @@ def argo_workflows(obj, name=None):
     default="",
     help="PagerDuty Events API V2 Integration key for workflow success/failure notifications.",
 )
+@click.option(
+    "--enable-heartbeat-daemon/--no-enable-heartbeat-daemon",
+    default=False,
+    show_default=True,
+    help="Use a daemon container to broadcast heartbeats.",
+)
 @click.option(
     "--deployer-attribute-file",
     default=None,
@@ -175,6 +184,12 @@ def argo_workflows(obj, name=None):
     help="Write the workflow name to the file specified. Used internally for Metaflow's Deployer API.",
     hidden=True,
 )
+@click.option(
+    "--enable-error-msg-capture/--no-enable-error-msg-capture",
+    default=False,
+    show_default=True,
+    help="Capture stack trace of first failed task in exit hook.",
+)
 @click.pass_obj
 def create(
     obj,
@@ -192,7 +207,9 @@ def create(
     notify_on_success=False,
     notify_slack_webhook_url=None,
     notify_pager_duty_integration_key=None,
+    enable_heartbeat_daemon=True,
     deployer_attribute_file=None,
+    enable_error_msg_capture=False,
 ):
     validate_tags(tags)
 
@@ -240,6 +257,8 @@ def create(
         notify_on_success,
         notify_slack_webhook_url,
         notify_pager_duty_integration_key,
+        enable_heartbeat_daemon,
+        enable_error_msg_capture,
     )
 
     if only_json:
@@ -412,6 +431,8 @@ def make_flow(
     notify_on_success,
     notify_slack_webhook_url,
     notify_pager_duty_integration_key,
+    enable_heartbeat_daemon,
+    enable_error_msg_capture,
 ):
     # TODO: Make this check less specific to Amazon S3 as we introduce
     #       support for more cloud object stores.
@@ -474,6 +495,8 @@ def make_flow(
         notify_on_success=notify_on_success,
         notify_slack_webhook_url=notify_slack_webhook_url,
         notify_pager_duty_integration_key=notify_pager_duty_integration_key,
+        enable_heartbeat_daemon=enable_heartbeat_daemon,
+        enable_error_msg_capture=enable_error_msg_capture,
     )
 
 

metaflow/plugins/argo/argo_workflows_decorator.py

@@ -54,7 +54,7 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
                                 "_", 1
                             )[
                                 0
-                            ]  # infer type from env var key
+                            ],  # infer type from env var key
                             # Add more event metadata here in the future
                         }
                     )
@@ -108,18 +108,12 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
         # we run pods with a security context. We work around this constraint by
         # mounting an emptyDir volume.
         if graph[step_name].type == "foreach":
-            # A DAGNode is considered a `parallel_step` if it is annotated by the @parallel decorator.
-            # A DAGNode is considered a `parallel_foreach` if it contains a `num_parallel` kwarg provided to the
-            # `next` method of that DAGNode.
-            # At this moment in the code we care if a node is marked as a `parallel_foreach` so that we can pass down the
-            # value of `num_parallel` to the subsequent steps.
-            # For @parallel, the implmentation uses 1 jobset object. That one jobset
-            # object internally creates 'num_parallel' jobs. So, we set foreach_num_splits
-            # to 1 here for @parallel. The parallelism of jobset is handled in
-            # kubernetes_job.py.
             if graph[step_name].parallel_foreach:
+                # If a node is marked as a `parallel_foreach`, pass down the value of
+                # `num_parallel` to the subsequent steps.
                 with open("/mnt/out/num_parallel", "w") as f:
                     json.dump(flow._parallel_ubf_iter.num_parallel, f)
+                # Set splits to 1 since parallelism is handled by JobSet.
                 flow._foreach_num_splits = 1
                 with open("/mnt/out/task_id_entropy", "w") as file:
                     import uuid
@@ -131,10 +125,9 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
             with open("/mnt/out/split_cardinality", "w") as file:
                 json.dump(flow._foreach_num_splits, file)
 
-        # for steps that have a `@parallel` decorator set to them, we will be relying on Jobsets
+        # For steps that have a `@parallel` decorator set to them, we will be relying on Jobsets
         # to run the task. In this case, we cannot set anything in the
-        # `/mnt/out` directory, since such form of output mounts are not available to jobset execution as
-        # argo just treats it like A K8s resource that it throws in the cluster.
+        # `/mnt/out` directory, since such form of output mounts are not available to Jobset executions.
         if not graph[step_name].parallel_step:
             # Unfortunately, we can't always use pod names as task-ids since the pod names
             # are not static across retries. We write the task-id to a file that is read

metaflow/plugins/argo/capture_error.py

@@ -0,0 +1,70 @@
+import json
+import os
+from datetime import datetime
+
+###
+# Algorithm to determine 1st error:
+#   ignore the failures where message = ""
+#   group the failures via templateName
+#     sort each group by finishedAt
+#   find the group for which the last finishedAt is earliest
+#   if the earliest message is "No more retries left" then
+#     get the n-1th message from that group
+#   else
+#     return the last message.
+###
+
+
+def parse_workflow_failures():
+    failures = json.loads(
+        json.loads(os.getenv("METAFLOW_ARGO_WORKFLOW_FAILURES", "[]"), strict=False),
+        strict=False,
+    )
+    return [wf for wf in failures if wf.get("message")]
+
+
+def group_failures_by_template(failures):
+    groups = {}
+    for failure in failures:
+        groups.setdefault(failure["templateName"], []).append(failure)
+    return groups
+
+
+def sort_by_finished_at(items):
+    return sorted(
+        items, key=lambda x: datetime.strptime(x["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
+    )
+
+
+def find_earliest_last_finished_group(groups):
+    return min(
+        groups,
+        key=lambda k: datetime.strptime(
+            groups[k][-1]["finishedAt"], "%Y-%m-%dT%H:%M:%SZ"
+        ),
+    )
+
+
+def determine_first_error():
+    failures = parse_workflow_failures()
+    if not failures:
+        return None
+
+    grouped_failures = group_failures_by_template(failures)
+    for group in grouped_failures.values():
+        group.sort(
+            key=lambda x: datetime.strptime(x["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
+        )
+
+    earliest_group = grouped_failures[
+        find_earliest_last_finished_group(grouped_failures)
+    ]
+
+    if earliest_group[-1]["message"] == "No more retries left":
+        return earliest_group[-2]
+    return earliest_group[-1]
+
+
+if __name__ == "__main__":
+    first_err = determine_first_error()
+    print(json.dumps(first_err, indent=2))

metaflow/plugins/argo/daemon.py

@@ -0,0 +1,59 @@
+from collections import namedtuple
+from time import sleep
+from metaflow.metaflow_config import DEFAULT_METADATA
+from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.plugins import METADATA_PROVIDERS
+from metaflow._vendor import click
+
+
+class CliState:
+    pass
+
+
+@click.group()
+@click.option("--flow_name", required=True)
+@click.option("--run_id", required=True)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    default=None,
+    help="Annotate all objects produced by Argo Workflows runs "
+    "with the given tag. You can specify this option multiple "
+    "times to attach multiple tags.",
+)
+@click.pass_context
+def cli(ctx, flow_name, run_id, tags=None):
+    ctx.obj = CliState()
+    ctx.obj.flow_name = flow_name
+    ctx.obj.run_id = run_id
+    ctx.obj.tags = tags
+    # Use a dummy flow to initialize the environment and metadata service,
+    # as we only need a name for the flow object.
+    flow = namedtuple("DummyFlow", "name")
+    dummyflow = flow(flow_name)
+
+    # Initialize a proper metadata service instance
+    environment = MetaflowEnvironment(dummyflow)
+
+    ctx.obj.metadata = [m for m in METADATA_PROVIDERS if m.TYPE == DEFAULT_METADATA][0](
+        environment, dummyflow, None, None
+    )
+
+
+@cli.command(help="start heartbeat process for a run")
+@click.pass_obj
+def heartbeat(obj):
+    # Try to register a run in case the start task has not taken care of it yet.
+    obj.metadata.register_run_id(obj.run_id, obj.tags)
+    # Start run heartbeat
+    obj.metadata.start_run_heartbeat(obj.flow_name, obj.run_id)
+    # Keepalive loop
+    while True:
+        # Do not pollute daemon logs with anything unnecessary,
+        # as they might be extremely long running.
+        sleep(10)
+
+
+if __name__ == "__main__":
+    cli()

metaflow/plugins/aws/step_functions/step_functions.py

@@ -664,9 +664,9 @@ class StepFunctions(object):
                 # input to those descendent tasks. We set and propagate the
                 # task ids pointing to split_parents through every state.
                 if any(self.graph[n].type == "foreach" for n in node.in_funcs):
-                    attrs[
-                        "split_parent_task_id_%s.$" % node.split_parents[-1]
-                    ] = "$.SplitParentTaskId"
+                    attrs["split_parent_task_id_%s.$" % node.split_parents[-1]] = (
+                        "$.SplitParentTaskId"
+                    )
                     for parent in node.split_parents[:-1]:
                         if self.graph[parent].type == "foreach":
                             attrs["split_parent_task_id_%s.$" % parent] = (

metaflow/plugins/cards/card_modules/basic.py

@@ -26,9 +26,11 @@ def transform_flow_graph(step_info):
         graph_dict[stepname] = {
             "type": node_to_type(step_info[stepname]["type"]),
             "box_next": step_info[stepname]["type"] not in ("linear", "join"),
-            "box_ends": None
-            if "matching_join" not in step_info[stepname]
-            else step_info[stepname]["matching_join"],
+            "box_ends": (
+                None
+                if "matching_join" not in step_info[stepname]
+                else step_info[stepname]["matching_join"]
+            ),
             "next": step_info[stepname]["next"],
             "doc": step_info[stepname]["doc"],
         }

metaflow/plugins/cards/card_modules/convert_to_native_type.py

@@ -314,8 +314,8 @@ class TaskToDict:
         # If there is any form of TypeError or ValueError we set the column value to "Unsupported Type"
         # We also set columns which are have null values to "null" strings
         time_format = "%Y-%m-%dT%H:%M:%S%Z"
-        truncate_long_objects = (
-            lambda x: x.astype("string").str.slice(0, 30) + "..."
+        truncate_long_objects = lambda x: (
+            x.astype("string").str.slice(0, 30) + "..."
             if len(x) > 0 and x.astype("string").str.len().max() > 30
             else x.astype("string")
         )

metaflow/plugins/cards/card_modules/renderer_tools.py

@@ -40,6 +40,7 @@ def render_safely(func):
     This is a decorator that can be added to any `MetaflowCardComponent.render`
     The goal is to render subcomponents safely and ensure that they are JSON serializable.
     """
+
     # expects a renderer func
     def ret_func(self, *args, **kwargs):
         return _render_component_safely(self, func, True, *args, **kwargs)

metaflow/plugins/cards/card_modules/test_cards.py

@@ -138,7 +138,6 @@ class TestJSONComponent(MetaflowCardComponent):
 
 
 class TestRefreshCard(MetaflowCard):
-
     """
     This card takes no components and helps test the `current.card.refresh(data)` interface.
     """
@@ -178,7 +177,6 @@ def _component_values_to_hash(components):
 
 
 class TestRefreshComponentCard(MetaflowCard):
-
     """
     This card takes components and helps test the `current.card.components["A"].update()`
     interface

metaflow/plugins/datastores/gs_storage.py

@@ -8,7 +8,6 @@ from tempfile import mkdtemp
 
 from metaflow.datastore.datastore_storage import DataStoreStorage, CloseAfterUse
 from metaflow.exception import MetaflowInternalError
-from metaflow.tracing import traced, tracing
 from metaflow.metaflow_config import (
     DATASTORE_SYSROOT_GS,
     ARTIFACT_LOCALROOT,
@@ -71,14 +70,12 @@ class _GSRootClient(object):
         """Drives GSStorage.is_file()"""
         try:
             blob = self.get_blob_client(path)
-            with traced("exists", dict(path=path)):
-                result = blob.exists()
+            result = blob.exists()
 
             return result
         except Exception as e:
             process_gs_exception(e)
 
-    @tracing
     def list_content_single(self, path):
         """Drives GSStorage.list_content()"""
 
@@ -107,7 +104,6 @@ class _GSRootClient(object):
         except Exception as e:
             process_gs_exception(e)
 
-    @tracing
     def save_bytes_single(
         self,
         path_tmpfile_metadata_triple,
@@ -123,12 +119,10 @@ class _GSRootClient(object):
                 blob.metadata = {"metaflow-user-attributes": json.dumps(metadata)}
             from google.cloud.storage.retry import DEFAULT_RETRY
 
-            with traced("upload_blob", dict(path=path)):
-                blob.upload_from_filename(tmpfile, retry=DEFAULT_RETRY)
+            blob.upload_from_filename(tmpfile, retry=DEFAULT_RETRY)
         except Exception as e:
             process_gs_exception(e)
 
-    @tracing
     def load_bytes_single(self, tmpdir, key):
         """Drives GSStorage.load_bytes()"""
         tmp_filename = os.path.join(tmpdir, str(uuid.uuid4()))
@@ -142,8 +136,7 @@ class _GSRootClient(object):
                 metaflow_user_attributes = json.loads(
                     blob.metadata["metaflow-user-attributes"]
                 )
-            with traced("download_blob", dict(path=key)):
-                blob.download_to_filename(tmp_filename)
+            blob.download_to_filename(tmp_filename)
         except google.api_core.exceptions.NotFound:
             tmp_filename = None
         return key, tmp_filename, metaflow_user_attributes

metaflow/plugins/datatools/s3/s3op.py

@@ -1119,9 +1119,11 @@ def get(
                         str(url.idx),
                         url_quote(url.prefix).decode(encoding="utf-8"),
                         url_quote(url.url).decode(encoding="utf-8"),
-                        url_quote(url.range).decode(encoding="utf-8")
-                        if url.range
-                        else "<norange>",
+                        (
+                            url_quote(url.range).decode(encoding="utf-8")
+                            if url.range
+                            else "<norange>"
+                        ),
                     ]
                 )
                 + "\n"

metaflow/plugins/kubernetes/kubernetes.py

@@ -299,6 +299,7 @@ class Kubernetes(object):
 
         jobset.environment_variables_from_selectors(
             {
+                "METAFLOW_KUBERNETES_NAMESPACE": "metadata.namespace",
                 "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
                 "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
                 "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -5,7 +5,6 @@ import random
 import sys
 import time
 
-from metaflow.tracing import inject_tracing_vars, tracing, traced
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_SECRETS
 from metaflow.tracing import inject_tracing_vars
@@ -99,13 +98,16 @@ class KubernetesJob(object):
                     containers=[
                         client.V1Container(
                             command=self._kwargs["command"],
-                            ports=[]
-                            if self._kwargs["port"] is None
-                            else [
-                                client.V1ContainerPort(
-                                    container_port=int(self._kwargs["port"])
-                                )
-                            ],
+                            termination_message_policy="FallbackToLogsOnError",
+                            ports=(
+                                []
+                                if self._kwargs["port"] is None
+                                else [
+                                    client.V1ContainerPort(
+                                        container_port=int(self._kwargs["port"])
+                                    )
+                                ]
+                            ),
                             env=[
                                 client.V1EnvVar(name=k, value=str(v))
                                 for k, v in self._kwargs.get(
@@ -125,6 +127,7 @@ class KubernetesJob(object):
                                     ),
                                 )
                                 for k, v in {
+                                    "METAFLOW_KUBERNETES_NAMESPACE": "metadata.namespace",
                                     "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
                                     "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
                                     "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
@@ -257,7 +260,6 @@ class KubernetesJob(object):
                         if self._kwargs["persistent_volume_claims"] is not None
                         else []
                     ),
-                    # TODO (savin): Set termination_message_policy
                 ),
             ),
         )
@@ -300,19 +302,13 @@ class KubernetesJob(object):
             #       achieve the guarantees that we are seeking.
             #       https://github.com/kubernetes/enhancements/issues/1040
             #       Hopefully, we will be able to get creative with kube-batch
-
-            with traced(
-                "submit_kubernetes_job",
-                {"job": "%s/%s" % (self._job.metadata.name, self._kwargs["namespace"])},
-            ):
-
-                response = (
-                    client.BatchV1Api()
-                    .create_namespaced_job(
-                        body=self._job, namespace=self._kwargs["namespace"]
-                    )
-                    .to_dict()
+            response = (
+                client.BatchV1Api()
+                .create_namespaced_job(
+                    body=self._job, namespace=self._kwargs["namespace"]
                 )
+                .to_dict()
+            )
             return RunningJob(
                 client=self._client,
                 name=response["metadata"]["name"],
@@ -441,14 +437,11 @@ class RunningJob(object):
     def _fetch_job(self):
         client = self._client.get()
         try:
-            with traced(
-                "fetch_kubernetes_job", {"job": "%s/%s" % (self._name, self._namespace)}
-            ):
-                return (
-                    client.BatchV1Api()
-                    .read_namespaced_job(name=self._name, namespace=self._namespace)
-                    .to_dict()
-                )
+            return (
+                client.BatchV1Api()
+                .read_namespaced_job(name=self._name, namespace=self._namespace)
+                .to_dict()
+            )
         except client.rest.ApiException as e:
             if e.status == 404:
                 raise KubernetesJobException(
@@ -460,20 +453,17 @@ class RunningJob(object):
     def _fetch_pod(self):
         # Fetch pod metadata.
         client = self._client.get()
-        with traced(
-            "fetch_kubernetes_pod", {"job": "%s/%s" % (self._name, self._namespace)}
-        ):
-            pods = (
-                client.CoreV1Api()
-                .list_namespaced_pod(
-                    namespace=self._namespace,
-                    label_selector="job-name={}".format(self._name),
-                )
-                .to_dict()["items"]
+        pods = (
+            client.CoreV1Api()
+            .list_namespaced_pod(
+                namespace=self._namespace,
+                label_selector="job-name={}".format(self._name),
             )
-            if pods:
-                return pods[0]
-            return {}
+            .to_dict()["items"]
+        )
+        if pods:
+            return pods[0]
+        return {}
 
     def kill(self):
         # Terminating a Kubernetes job is a bit tricky. Issuing a

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -52,8 +52,6 @@ def k8s_retry(deadline_seconds=60, max_backoff=32):
     return decorator
 
 
-CONTROL_JOB_NAME = "control"
-
 JobsetStatus = namedtuple(
     "JobsetStatus",
     [
@@ -586,13 +584,18 @@ class JobSetSpec(object):
                                 containers=[
                                     client.V1Container(
                                         command=self._kwargs["command"],
-                                        ports=[]
-                                        if self._kwargs["port"] is None
-                                        else [
-                                            client.V1ContainerPort(
-                                                container_port=int(self._kwargs["port"])
-                                            )
-                                        ],
+                                        termination_message_policy="FallbackToLogsOnError",
+                                        ports=(
+                                            []
+                                            if self._kwargs["port"] is None
+                                            else [
+                                                client.V1ContainerPort(
+                                                    container_port=int(
+                                                        self._kwargs["port"]
+                                                    )
+                                                )
+                                            ]
+                                        ),
                                         env=[
                                             client.V1EnvVar(name=k, value=str(v))
                                             for k, v in self._kwargs.get(
@@ -756,7 +759,6 @@ class JobSetSpec(object):
                                     is not None
                                     else []
                                 ),
-                                # TODO (savin): Set termination_message_policy
                             ),
                         ),
                     ),
@@ -790,14 +792,14 @@ class KubernetesJobSet(object):
 
         self._jobset_control_addr = _make_domain_name(
             name,
-            CONTROL_JOB_NAME,
+            "control",
             0,
             0,
             namespace,
         )
 
         self._control_spec = JobSetSpec(
-            client.get(), name=CONTROL_JOB_NAME, namespace=namespace, **kwargs
+            client.get(), name="control", namespace=namespace, **kwargs
         )
         self._worker_spec = JobSetSpec(
             client.get(), name="worker", namespace=namespace, **kwargs
@@ -918,14 +920,14 @@ class KubernetesArgoJobSet(object):
 
         self._jobset_control_addr = _make_domain_name(
             name,
-            CONTROL_JOB_NAME,
+            "control",
             0,
             0,
             namespace,
         )
 
         self._control_spec = JobSetSpec(
-            kubernetes_sdk, name=CONTROL_JOB_NAME, namespace=namespace, **kwargs
+            kubernetes_sdk, name="control", namespace=namespace, **kwargs
         )
         self._worker_spec = JobSetSpec(
             kubernetes_sdk, name="worker", namespace=namespace, **kwargs

metaflow/plugins/logs_cli.py

@@ -7,6 +7,7 @@ from ..datastore import TaskDataStoreSet, TaskDataStore
 
 from ..mflog import mflog, LOG_SOURCES
 
+
 # main motivation from https://github.com/pallets/click/issues/430
 # in order to support a default command being called for a Click group.
 #

metaflow/plugins/pypi/conda_environment.py

@@ -301,9 +301,7 @@ class CondaEnvironment(MetaflowEnvironment):
             lambda f: lambda obj: (
                 {k: f(f)(v) for k, v in sorted(obj.items())}
                 if isinstance(obj, dict)
-                else sorted([f(f)(e) for e in obj])
-                if isinstance(obj, list)
-                else obj
+                else sorted([f(f)(e) for e in obj]) if isinstance(obj, list) else obj
             )
         )
 

metaflow/plugins/pypi/pip.py

@@ -121,9 +121,9 @@ class Pip(object):
                     res["url"] = "{vcs}+{url}@{commit_id}{subdir_str}".format(
                         **vcs_info,
                         **res,
-                        subdir_str="#subdirectory=%s" % subdirectory
-                        if subdirectory
-                        else ""
+                        subdir_str=(
+                            "#subdirectory=%s" % subdirectory if subdirectory else ""
+                        )
                     )
                     # used to deduplicate the storage location in case wheel does not
                     # build with enough unique identifiers.

metaflow/plugins/storage_executor.py

@@ -6,7 +6,6 @@ import sys
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 
 from metaflow.exception import MetaflowException
-from metaflow.tracing import TracedThreadPoolExecutor
 
 if sys.version_info[:2] < (3, 7):
     # in 3.6, Only BrokenProcessPool exists (there is no BrokenThreadPool)
@@ -133,10 +132,7 @@ class StorageExecutor(object):
                     msg="Cannot use ProcessPoolExecutor because Python version is older than 3.7 and multiprocess start method has been set to something other than 'spawn'"
                 )
         else:
-            self._executor = TracedThreadPoolExecutor(
-                "StorageExecutor", max_workers=threadpool_max_workers
-            )
-            # self._executor = ThreadPoolExecutor(max_workers=threadpool_max_workers)
+            self._executor = ThreadPoolExecutor(max_workers=threadpool_max_workers)
 
     def warm_up(self):
         # warm up at least one process or thread in the pool.