metaflow 2.12.9__py2.py3-none-any.whl → 2.12.11__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/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

@@ -98,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(
@@ -124,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",
@@ -256,7 +260,6 @@ class KubernetesJob(object):
                         if self._kwargs["persistent_volume_claims"] is not None
                         else []
                     ),
-                    # TODO (savin): Set termination_message_policy
                 ),
             ),
         )

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

@@ -298,9 +298,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/tag_cli.py

@@ -507,9 +507,9 @@ def tag_list(
 
     if not group_by_run and not group_by_tag:
         # We list all the runs that match to print them out if needed.
-        system_tags_by_some_grouping[
-            ",".join(pathspecs)
-        ] = system_tags_by_some_grouping.get("_", set())
+        system_tags_by_some_grouping[",".join(pathspecs)] = (
+            system_tags_by_some_grouping.get("_", set())
+        )
         all_tags_by_some_grouping[",".join(pathspecs)] = all_tags_by_some_grouping.get(
             "_", set()
         )

metaflow/procpoll.py

@@ -31,7 +31,7 @@ class LinuxProcPoll(ProcPoll):
         self._poll.unregister(fd)
 
     def poll(self, timeout):
-        for (fd, event) in self._poll.poll(timeout):
+        for fd, event in self._poll.poll(timeout):
             yield ProcPollEvent(
                 fd=fd,
                 can_read=bool(event & select.POLLIN),

metaflow/runtime.py

@@ -4,6 +4,7 @@ Local backend
 Execute the flow with a native runtime
 using local / remote processes
 """
+
 from __future__ import print_function
 import os
 import sys

metaflow/util.py

@@ -382,9 +382,9 @@ def to_camelcase(obj):
     if isinstance(obj, dict):
         res = obj.__class__()
         for k in obj:
-            res[
-                re.sub(r"(?!^)_([a-zA-Z])", lambda x: x.group(1).upper(), k)
-            ] = to_camelcase(obj[k])
+            res[re.sub(r"(?!^)_([a-zA-Z])", lambda x: x.group(1).upper(), k)] = (
+                to_camelcase(obj[k])
+            )
     elif isinstance(obj, (list, set, tuple)):
         res = obj.__class__(to_camelcase(v) for v in obj)
     else:
@@ -401,9 +401,9 @@ def to_pascalcase(obj):
     if isinstance(obj, dict):
         res = obj.__class__()
         for k in obj:
-            res[
-                re.sub("([a-zA-Z])", lambda x: x.groups()[0].upper(), k, 1)
-            ] = to_pascalcase(obj[k])
+            res[re.sub("([a-zA-Z])", lambda x: x.groups()[0].upper(), k, 1)] = (
+                to_pascalcase(obj[k])
+            )
     elif isinstance(obj, (list, set, tuple)):
         res = obj.__class__(to_pascalcase(v) for v in obj)
     else:

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.12.9"
+metaflow_version = "2.12.11"

{metaflow-2.12.9.dist-info → metaflow-2.12.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: metaflow
-Version: 2.12.9
+Version: 2.12.11
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs ==2.12.9 ; extra == 'stubs'
+Requires-Dist: metaflow-stubs==2.12.11; extra == "stubs"
 
 ![Metaflow_Logo_Horizontal_FullColor_Ribbon_Dark_RGB](https://user-images.githubusercontent.com/763451/89453116-96a57e00-d713-11ea-9fa6-82b29d4d6eff.png)