ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/airflow/airflow_cli.py

@@ -7,6 +7,7 @@ from hashlib import sha1
 from metaflow import current, decorators
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
+from metaflow.metaflow_config import FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
 from metaflow.package import MetaflowPackage
 from metaflow.plugins.aws.step_functions.production_token import (
     load_token,
@@ -283,24 +284,35 @@ def make_flow(
 ):
     # Attach @kubernetes.
     decorators._attach_decorators(obj.flow, [KubernetesDecorator.name])
+    decorators._init(obj.flow)
 
     decorators._init_step_decorators(
         obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger
     )
-
+    obj.graph = obj.flow._graph
     # Save the code package in the flow datastore so that both user code and
     # metaflow package can be retrieved during workflow execution.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
-    package_url, package_sha = obj.flow_datastore.save_data(
-        [obj.package.blob], len_hint=1
-    )[0]
+    # This blocks until the package is created
+    if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+        package_url = obj.package.package_url()
+        package_sha = obj.package.package_sha()
+    else:
+        package_url, package_sha = obj.flow_datastore.save_data(
+            [obj.package.blob], len_hint=1
+        )[0]
 
     return Airflow(
         dag_name,
         obj.graph,
         obj.flow,
+        obj.package.package_metadata,
         package_sha,
         package_url,
         obj.metadata,
@@ -389,6 +401,11 @@ def _validate_workflow(flow, graph, flow_datastore, metadata, workflow_timeout):
                 "Step *%s* is marked for execution on AWS Batch with Airflow which isn't currently supported."
                 % node.name
             )
+        if any([d.name == "slurm" for d in node.decorators]):
+            raise NotSupportedException(
+                "Step *%s* is marked for execution on Slurm with Airflow which isn't currently supported."
+                % node.name
+            )
     SUPPORTED_DATASTORES = ("azure", "s3", "gs")
     if flow_datastore.TYPE not in SUPPORTED_DATASTORES:
         raise AirflowException(

metaflow/plugins/airflow/airflow_decorator.py

@@ -1,7 +1,7 @@
 import json
 import os
 from metaflow.decorators import StepDecorator
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 
 from .airflow_utils import (
     TASK_ID_XCOM_KEY,

metaflow/plugins/airflow/airflow_utils.py

@@ -567,9 +567,11 @@ class AirflowTask(object):
         return cls(
             task_dict["name"],
             is_mapper_node=is_mapper_node,
-            operator_type=task_dict["operator_type"]
-            if "operator_type" in task_dict
-            else "kubernetes",
+            operator_type=(
+                task_dict["operator_type"]
+                if "operator_type" in task_dict
+                else "kubernetes"
+            ),
             flow_name=flow_name,
             flow_contains_foreach=flow_contains_foreach,
         ).set_operator_args(**op_args)

metaflow/plugins/airflow/sensors/base_sensor.py

@@ -1,5 +1,5 @@
 import uuid
-from metaflow.decorators import FlowDecorator
+from metaflow.decorators import FlowDecorator, flow_decorators
 from ..exception import AirflowException
 from ..airflow_utils import AirflowTask, id_creator, TASK_ID_HASH_LEN
 
@@ -49,7 +49,7 @@ class AirflowSensorDecorator(FlowDecorator):
             operator_type=self.operator_type,
         ).set_operator_args(**{k: v for k, v in task_args.items() if v is not None})
 
-    def validate(self):
+    def validate(self, flow):
         """
         Validate if the arguments for the sensor are correct.
         """
@@ -58,7 +58,7 @@ class AirflowSensorDecorator(FlowDecorator):
         if self.attributes["name"] is None:
             deco_index = [
                 d._id
-                for d in self._flow_decorators
+                for d in flow_decorators(flow)
                 if issubclass(d.__class__, AirflowSensorDecorator)
             ].index(self._id)
             self._airflow_task_name = "%s-%s" % (
@@ -71,4 +71,4 @@ class AirflowSensorDecorator(FlowDecorator):
     def flow_init(
         self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
     ):
-        self.validate()
+        self.validate(flow)

metaflow/plugins/airflow/sensors/external_task_sensor.py

@@ -85,7 +85,7 @@ class ExternalTaskSensorDecorator(AirflowSensorDecorator):
             )
         return task_args
 
-    def validate(self):
+    def validate(self, flow):
         if self.attributes["external_dag_id"] is None:
             raise AirflowException(
                 "`%s` argument of `@%s`cannot be `None`."
@@ -131,4 +131,4 @@ class ExternalTaskSensorDecorator(AirflowSensorDecorator):
                         self.name,
                     )
                 )
-        super().validate()
+        super().validate(flow)

metaflow/plugins/airflow/sensors/s3_sensor.py

@@ -58,9 +58,9 @@ class S3KeySensorDecorator(AirflowSensorDecorator):
         #  `verify` is a airflow variable.
     )
 
-    def validate(self):
+    def validate(self, flow):
         if self.attributes["bucket_key"] is None:
             raise AirflowException(
                 "`bucket_key` for `@%s`cannot be empty." % (self.name)
             )
-        super().validate()
+        super().validate(flow)

metaflow/plugins/argo/argo_client.py

@@ -1,7 +1,6 @@
 import json
-import os
-import sys
 
+from metaflow.metaflow_config import ARGO_EVENTS_SENSOR_NAMESPACE
 from metaflow.exception import MetaflowException
 from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
 
@@ -10,6 +9,14 @@ class ArgoClientException(MetaflowException):
     headline = "Argo Client error"
 
 
+class ArgoResourceNotFound(MetaflowException):
+    headline = "Resource not found"
+
+
+class ArgoNotPermitted(MetaflowException):
+    headline = "Operation not permitted"
+
+
 class ArgoClient(object):
     def __init__(self, namespace=None):
         self._client = KubernetesClient()
@@ -52,21 +59,38 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def get_workflow_templates(self):
+    def get_workflow_templates(self, page_size=100):
         client = self._client.get()
-        try:
-            return client.CustomObjectsApi().list_namespaced_custom_object(
-                group=self._group,
-                version=self._version,
-                namespace=self._namespace,
-                plural="workflowtemplates",
-            )["items"]
-        except client.rest.ApiException as e:
-            if e.status == 404:
-                return None
-            raise ArgoClientException(
-                json.loads(e.body)["message"] if e.body is not None else e.reason
-            )
+        continue_token = None
+
+        while True:
+            try:
+                params = {"limit": page_size}
+                if continue_token:
+                    params["_continue"] = continue_token
+
+                response = client.CustomObjectsApi().list_namespaced_custom_object(
+                    group=self._group,
+                    version=self._version,
+                    namespace=self._namespace,
+                    plural="workflowtemplates",
+                    **params,
+                )
+
+                for item in response.get("items", []):
+                    yield item
+
+                metadata = response.get("metadata", {})
+                continue_token = metadata.get("continue")
+
+                if not continue_token:
+                    break
+            except client.rest.ApiException as e:
+                if e.status == 404:
+                    return None
+                raise ArgoClientException(
+                    json.loads(e.body)["message"] if e.body is not None else e.reason
+                )
 
     def register_workflow_template(self, name, workflow_template):
         # Unfortunately, Kubernetes client does not handle optimistic
@@ -140,9 +164,7 @@ class ArgoClient(object):
             if e.status == 404:
                 return None
             else:
-                raise ArgoClientException(
-                    json.loads(e.body)["message"] if e.body is not None else e.reason
-                )
+                raise wrap_api_error(e)
 
     def delete_workflow_template(self, name):
         """
@@ -164,9 +186,7 @@ class ArgoClient(object):
             if e.status == 404:
                 return None
             else:
-                raise ArgoClientException(
-                    json.loads(e.body)["message"] if e.body is not None else e.reason
-                )
+                raise wrap_api_error(e)
 
     def terminate_workflow(self, name):
         client = self._client.get()
@@ -254,12 +274,19 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def trigger_workflow_template(self, name, parameters={}):
+    def trigger_workflow_template(self, name, usertype, username, parameters={}):
         client = self._client.get()
         body = {
             "apiVersion": "argoproj.io/v1alpha1",
             "kind": "Workflow",
-            "metadata": {"generateName": name + "-"},
+            "metadata": {
+                "generateName": name + "-",
+                "annotations": {
+                    "metaflow/triggered_by_user": json.dumps(
+                        {"type": usertype, "name": username}
+                    )
+                },
+            },
             "spec": {
                 "workflowTemplateRef": {"name": name},
                 "arguments": {
@@ -295,7 +322,10 @@ class ArgoClient(object):
                 "suspend": schedule is None,
                 "schedule": schedule,
                 "timezone": timezone,
+                "failedJobsHistoryLimit": 10000,  # default is unfortunately 1
+                "successfulJobsHistoryLimit": 10000,  # default is unfortunately 3
                 "workflowSpec": {"workflowTemplateRef": {"name": name}},
+                "startingDeadlineSeconds": 3540,  # configuring this to 59 minutes so a failed trigger of cron workflow can succeed at most 59 mins after scheduled execution
             },
         }
         try:
@@ -349,12 +379,15 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def register_sensor(self, name, sensor=None):
+    def register_sensor(
+        self, name, sensor=None, sensor_namespace=ARGO_EVENTS_SENSOR_NAMESPACE
+    ):
         if sensor is None:
             sensor = {}
         # Unfortunately, Kubernetes client does not handle optimistic
         # concurrency control by itself unlike kubectl
         client = self._client.get()
+
         if not sensor:
             sensor["metadata"] = {}
 
@@ -364,7 +397,7 @@ class ArgoClient(object):
             ] = client.CustomObjectsApi().get_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )[
@@ -379,7 +412,7 @@ class ArgoClient(object):
                     return client.CustomObjectsApi().create_namespaced_custom_object(
                         group=self._group,
                         version=self._version,
-                        namespace=self._namespace,
+                        namespace=sensor_namespace,
                         plural="sensors",
                         body=sensor,
                     )
@@ -397,7 +430,7 @@ class ArgoClient(object):
             return client.CustomObjectsApi().replace_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 body=sensor,
                 name=name,
@@ -407,7 +440,7 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def delete_sensor(self, name):
+    def delete_sensor(self, name, sensor_namespace):
         """
         Issues an API call for deleting a sensor
 
@@ -419,13 +452,25 @@ class ArgoClient(object):
             return client.CustomObjectsApi().delete_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )
         except client.rest.ApiException as e:
             if e.status == 404:
                 return None
-            raise ArgoClientException(
-                json.loads(e.body)["message"] if e.body is not None else e.reason
-            )
+            raise wrap_api_error(e)
+
+
+def wrap_api_error(error):
+    message = (
+        json.loads(error.body)["message"] if error.body is not None else error.reason
+    )
+    # catch all
+    ex = ArgoClientException(message)
+    if error.status == 404:
+        # usually handled outside this function as most cases want to return None instead.
+        ex = ArgoResourceNotFound(message)
+    if error.status == 403:
+        ex = ArgoNotPermitted(message)
+    return ex

metaflow/plugins/argo/argo_events.py

@@ -11,6 +11,7 @@ from metaflow.metaflow_config import (
     ARGO_EVENTS_WEBHOOK_AUTH,
     ARGO_EVENTS_WEBHOOK_URL,
     SERVICE_HEADERS,
+    SERVICE_RETRY_COUNT,
 )
 
 
@@ -127,12 +128,11 @@ class ArgoEvent(object):
                     headers={"Content-Type": "application/json", **headers},
                     data=json.dumps(data).encode("utf-8"),
                 )
-                retries = 3
-                backoff_factor = 2
 
-                for i in range(retries):
+                for i in range(SERVICE_RETRY_COUNT):
                     try:
-                        urllib.request.urlopen(request, timeout=10.0)
+                        # we do not want to wait indefinitely for a response on the event broadcast, as this will keep the task running.
+                        urllib.request.urlopen(request, timeout=60)
                         print(
                             "Argo Event (%s) published." % self._name, file=sys.stderr
                         )
@@ -141,10 +141,10 @@ class ArgoEvent(object):
                         # TODO: Retry retryable HTTP error codes
                         raise e
                     except urllib.error.URLError as e:
-                        if i == retries - 1:
+                        if i == SERVICE_RETRY_COUNT - 1:
                             raise e
                         else:
-                            time.sleep(backoff_factor**i)
+                            time.sleep(2**i)
             except Exception as e:
                 msg = "Unable to publish Argo Event (%s): %s" % (self._name, e)
                 if ignore_errors: