metaflow 2.13__py2.py3-none-any.whl → 2.13.1__py2.py3-none-any.whl

metaflow/metadata_provider/heartbeat.py

@@ -1,11 +1,12 @@
+import json
 import time
+from threading import Thread
+
 import requests
-import json
 
-from threading import Thread
-from metaflow.sidecar import MessageTypes, Message
-from metaflow.metaflow_config import SERVICE_HEADERS
 from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import SERVICE_HEADERS
+from metaflow.sidecar import Message, MessageTypes
 
 HB_URL_KEY = "hb_url"
 
@@ -52,13 +53,27 @@ class MetadataHeartBeat(object):
                 retry_counter = 0
             except HeartBeatException as e:
                 retry_counter = retry_counter + 1
-                time.sleep(4**retry_counter)
+                time.sleep(1.5**retry_counter)
 
     def _heartbeat(self):
         if self.hb_url is not None:
-            response = requests.post(
-                url=self.hb_url, data="{}", headers=self.headers.copy()
-            )
+            try:
+                response = requests.post(
+                    url=self.hb_url, data="{}", headers=self.headers.copy()
+                )
+            except requests.exceptions.ConnectionError as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (ConnectionError)" % (self.hb_url)
+                )
+            except requests.exceptions.Timeout as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed" " (Timeout)" % (self.hb_url)
+                )
+            except requests.exceptions.RequestException as e:
+                raise HeartBeatException(
+                    "HeartBeat request (%s) failed"
+                    " (RequestException) %s" % (self.hb_url, str(e))
+                )
             # Unfortunately, response.json() returns a string that we need
             # to cast to json; however when the request encounters an error
             # the return type is a json blob :/

metaflow/metaflow_config.py

@@ -355,6 +355,8 @@ KUBERNETES_PERSISTENT_VOLUME_CLAIMS = from_conf(
 KUBERNETES_SECRETS = from_conf("KUBERNETES_SECRETS", "")
 # Default labels for kubernetes pods
 KUBERNETES_LABELS = from_conf("KUBERNETES_LABELS", "")
+# Default annotations for kubernetes pods
+KUBERNETES_ANNOTATIONS = from_conf("KUBERNETES_ANNOTATIONS", "")
 # Default GPU vendor to use by K8S jobs created by Metaflow (supports nvidia, amd)
 KUBERNETES_GPU_VENDOR = from_conf("KUBERNETES_GPU_VENDOR", "nvidia")
 # Default container image for K8S

metaflow/plugins/argo/argo_client.py

@@ -1,6 +1,4 @@
 import json
-import os
-import sys
 
 from metaflow.exception import MetaflowException
 from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient

metaflow/plugins/argo/argo_workflows.py

@@ -7,12 +7,11 @@ import sys
 from collections import defaultdict
 from hashlib import sha1
 from math import inf
-from typing import List, Tuple
 
 from metaflow import JSONType, current
 from metaflow.decorators import flow_decorators
 from metaflow.exception import MetaflowException
-from metaflow.graph import DAGNode, FlowGraph
+from metaflow.graph import FlowGraph
 from metaflow.includefile import FilePathClass
 from metaflow.metaflow_config import (
     ARGO_EVENTS_EVENT,
@@ -39,9 +38,7 @@ from metaflow.metaflow_config import (
     DEFAULT_SECRETS_BACKEND_TYPE,
     GCP_SECRET_MANAGER_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
-    KUBERNETES_LABELS,
     KUBERNETES_NAMESPACE,
-    KUBERNETES_NODE_SELECTOR,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
     KUBERNETES_SECRETS,
     S3_ENDPOINT_URL,
@@ -54,10 +51,7 @@ from metaflow.metaflow_config_funcs import config_values
 from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
 from metaflow.parameters import deploy_time_eval
 from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
-from metaflow.plugins.kubernetes.kubernetes import (
-    parse_kube_keyvalue_list,
-    validate_kube_labels,
-)
+
 from metaflow.plugins.kubernetes.kubernetes_jobsets import KubernetesArgoJobSet
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.user_configs.config_options import ConfigInput
@@ -173,7 +167,8 @@ class ArgoWorkflows(object):
         self.triggers, self.trigger_options = self._process_triggers()
         self._schedule, self._timezone = self._get_schedule()
 
-        self.kubernetes_labels = self._get_kubernetes_labels()
+        self._base_labels = self._base_kubernetes_labels()
+        self._base_annotations = self._base_kubernetes_annotations()
         self._workflow_template = self._compile_workflow_template()
         self._sensor = self._compile_sensor()
 
@@ -310,7 +305,7 @@ class ArgoWorkflows(object):
             try:
                 # Check that the workflow was deployed through Metaflow
                 workflow_template["metadata"]["annotations"]["metaflow/owner"]
-            except KeyError as e:
+            except KeyError:
                 raise ArgoWorkflowsException(
                     "An existing non-metaflow workflow with the same name as "
                     "*%s* already exists in Argo Workflows. \nPlease modify the "
@@ -324,18 +319,42 @@ class ArgoWorkflows(object):
         except Exception as e:
             raise ArgoWorkflowsException(str(e))
 
-    @staticmethod
-    def _get_kubernetes_labels():
+    def _base_kubernetes_labels(self):
         """
-        Get Kubernetes labels from environment variable.
-        Parses the string into a dict and validates that values adhere to Kubernetes restrictions.
+        Get shared Kubernetes labels for Argo resources.
         """
-        if not KUBERNETES_LABELS:
-            return {}
-        env_labels = KUBERNETES_LABELS.split(",")
-        env_labels = parse_kube_keyvalue_list(env_labels, False)
-        validate_kube_labels(env_labels)
-        return env_labels
+        # TODO: Add configuration through an environment variable or Metaflow config in the future if required.
+        labels = {"app.kubernetes.io/part-of": "metaflow"}
+
+        return labels
+
+    def _base_kubernetes_annotations(self):
+        """
+        Get shared Kubernetes annotations for Argo resources.
+        """
+        from datetime import datetime, timezone
+
+        # TODO: Add configuration through an environment variable or Metaflow config in the future if required.
+        # base annotations
+        annotations = {
+            "metaflow/production_token": self.production_token,
+            "metaflow/owner": self.username,
+            "metaflow/user": "argo-workflows",
+            "metaflow/flow_name": self.flow.name,
+            "metaflow/deployment_timestamp": str(
+                datetime.now(timezone.utc).isoformat()
+            ),
+        }
+
+        if current.get("project_name"):
+            annotations.update(
+                {
+                    "metaflow/project_name": current.project_name,
+                    "metaflow/branch_name": current.branch_name,
+                    "metaflow/project_flow_name": current.project_flow_name,
+                }
+            )
+        return annotations
 
     def _get_schedule(self):
         schedule = self.flow._flow_decorators.get("schedule")
@@ -411,7 +430,7 @@ class ArgoWorkflows(object):
                         "metaflow/production_token"
                     ],
                 )
-            except KeyError as e:
+            except KeyError:
                 raise ArgoWorkflowsException(
                     "An existing non-metaflow workflow with the same name as "
                     "*%s* already exists in Argo Workflows. \nPlease modify the "
@@ -676,18 +695,7 @@ class ArgoWorkflows(object):
         # generate container templates at the top level (in WorkflowSpec) and maintain
         # references to them within the DAGTask.
 
-        from datetime import datetime, timezone
-
-        annotations = {
-            "metaflow/production_token": self.production_token,
-            "metaflow/owner": self.username,
-            "metaflow/user": "argo-workflows",
-            "metaflow/flow_name": self.flow.name,
-            "metaflow/deployment_timestamp": str(
-                datetime.now(timezone.utc).isoformat()
-            ),
-        }
-
+        annotations = {}
         if self._schedule is not None:
             # timezone is an optional field and json dumps on None will result in null
             # hence configuring it to an empty string
@@ -699,15 +707,6 @@ class ArgoWorkflows(object):
         if self.parameters:
             annotations.update({"metaflow/parameters": json.dumps(self.parameters)})
 
-        if current.get("project_name"):
-            annotations.update(
-                {
-                    "metaflow/project_name": current.project_name,
-                    "metaflow/branch_name": current.branch_name,
-                    "metaflow/project_flow_name": current.project_flow_name,
-                }
-            )
-
         # Some more annotations to populate the Argo UI nicely
         if self.tags:
             annotations.update({"metaflow/tags": json.dumps(self.tags)})
@@ -755,9 +754,10 @@ class ArgoWorkflows(object):
                 # is released, we should be able to support multi-namespace /
                 # multi-cluster scheduling.
                 .namespace(KUBERNETES_NAMESPACE)
-                .label("app.kubernetes.io/name", "metaflow-flow")
-                .label("app.kubernetes.io/part-of", "metaflow")
                 .annotations(annotations)
+                .annotations(self._base_annotations)
+                .labels(self._base_labels)
+                .label("app.kubernetes.io/name", "metaflow-flow")
             )
             .spec(
                 WorkflowSpec()
@@ -787,10 +787,14 @@ class ArgoWorkflows(object):
                 # Set workflow metadata
                 .workflow_metadata(
                     Metadata()
+                    .labels(self._base_labels)
                     .label("app.kubernetes.io/name", "metaflow-run")
-                    .label("app.kubernetes.io/part-of", "metaflow")
                     .annotations(
-                        {**annotations, **{"metaflow/run_id": "argo-{{workflow.name}}"}}
+                        {
+                            **annotations,
+                            **self._base_annotations,
+                            **{"metaflow/run_id": "argo-{{workflow.name}}"},
+                        }
                     )
                     # TODO: Set dynamic labels using labels_from. Ideally, we would
                     #       want to expose run_id as a label. It's easy to add labels,
@@ -823,10 +827,10 @@ class ArgoWorkflows(object):
                 # Set common pod metadata.
                 .pod_metadata(
                     Metadata()
+                    .labels(self._base_labels)
                     .label("app.kubernetes.io/name", "metaflow-task")
-                    .label("app.kubernetes.io/part-of", "metaflow")
                     .annotations(annotations)
-                    .labels(self.kubernetes_labels)
+                    .annotations(self._base_annotations)
                 )
                 # Set the entrypoint to flow name
                 .entrypoint(self.flow.name)
@@ -1895,15 +1899,7 @@ class ArgoWorkflows(object):
             # twice, but due to issues with variable substitution, we will have to
             # live with this routine.
             if node.parallel_step:
-                # Explicitly add the task-id-hint label. This is important because this label
-                # is returned as an Output parameter of this step and is used subsequently as an
-                # an input in the join step.
-                kubernetes_labels = self.kubernetes_labels.copy()
                 jobset_name = "{{inputs.parameters.jobset-name}}"
-                kubernetes_labels["task_id_entropy"] = (
-                    "{{inputs.parameters.task-id-entropy}}"
-                )
-                kubernetes_labels["num_parallel"] = "{{inputs.parameters.num-parallel}}"
                 jobset = KubernetesArgoJobSet(
                     kubernetes_sdk=kubernetes_sdk,
                     name=jobset_name,
@@ -1959,8 +1955,22 @@ class ArgoWorkflows(object):
                 for k, v in env.items():
                     jobset.environment_variable(k, v)
 
-                for k, v in kubernetes_labels.items():
-                    jobset.label(k, v)
+                # Set labels. Do not allow user-specified task labels to override internal ones.
+                #
+                # Explicitly add the task-id-hint label. This is important because this label
+                # is returned as an Output parameter of this step and is used subsequently as an
+                # an input in the join step.
+                kubernetes_labels = {
+                    "task_id_entropy": "{{inputs.parameters.task-id-entropy}}",
+                    "num_parallel": "{{inputs.parameters.num-parallel}}",
+                }
+                jobset.labels(
+                    {
+                        **resources["labels"],
+                        **self._base_labels,
+                        **kubernetes_labels,
+                    }
+                )
 
                 jobset.environment_variable(
                     "MF_MASTER_ADDR", jobset.jobset_control_addr
@@ -1989,27 +1999,23 @@ class ArgoWorkflows(object):
                         "TASK_ID_SUFFIX": "metadata.annotations['jobset.sigs.k8s.io/job-index']",
                     }
                 )
+
+                # Set annotations. Do not allow user-specified task-specific annotations to override internal ones.
                 annotations = {
                     # setting annotations explicitly as they wont be
                     # passed down from WorkflowTemplate level
                     "metaflow/step_name": node.name,
                     "metaflow/attempt": str(retry_count),
                     "metaflow/run_id": run_id,
-                    "metaflow/production_token": self.production_token,
-                    "metaflow/owner": self.username,
-                    "metaflow/user": "argo-workflows",
-                    "metaflow/flow_name": self.flow.name,
                 }
-                if current.get("project_name"):
-                    annotations.update(
-                        {
-                            "metaflow/project_name": current.project_name,
-                            "metaflow/branch_name": current.branch_name,
-                            "metaflow/project_flow_name": current.project_flow_name,
-                        }
-                    )
-                for k, v in annotations.items():
-                    jobset.annotation(k, v)
+
+                jobset.annotations(
+                    {
+                        **resources["annotations"],
+                        **self._base_annotations,
+                        **annotations,
+                    }
+                )
 
                 jobset.control.replicas(1)
                 jobset.worker.replicas("{{=asInt(inputs.parameters.workerCount)}}")
@@ -2066,13 +2072,16 @@ class ArgoWorkflows(object):
                         minutes_between_retries=minutes_between_retries,
                     )
                     .metadata(
-                        ObjectMeta().annotation("metaflow/step_name", node.name)
+                        ObjectMeta()
+                        .annotation("metaflow/step_name", node.name)
                         # Unfortunately, we can't set the task_id since it is generated
                         # inside the pod. However, it can be inferred from the annotation
                         # set by argo-workflows - `workflows.argoproj.io/outputs` - refer
                         # the field 'task-id' in 'parameters'
                         # .annotation("metaflow/task_id", ...)
                         .annotation("metaflow/attempt", retry_count)
+                        .annotations(resources["annotations"])
+                        .labels(resources["labels"])
                     )
                     # Set emptyDir volume for state management
                     .empty_dir_volume("out")
@@ -2840,33 +2849,6 @@ class ArgoWorkflows(object):
                 "sdk (https://pypi.org/project/kubernetes/) first."
             )
 
-        labels = {"app.kubernetes.io/part-of": "metaflow"}
-
-        annotations = {
-            "metaflow/production_token": self.production_token,
-            "metaflow/owner": self.username,
-            "metaflow/user": "argo-workflows",
-            "metaflow/flow_name": self.flow.name,
-        }
-        if current.get("project_name"):
-            annotations.update(
-                {
-                    "metaflow/project_name": current.project_name,
-                    "metaflow/branch_name": current.branch_name,
-                    "metaflow/project_flow_name": current.project_flow_name,
-                }
-            )
-
-        # Useful to paint the UI
-        trigger_annotations = {
-            "metaflow/triggered_by": json.dumps(
-                [
-                    {key: trigger.get(key) for key in ["name", "type"]}
-                    for trigger in self.triggers
-                ]
-            )
-        }
-
         return (
             Sensor()
             .metadata(
@@ -2874,10 +2856,9 @@ class ArgoWorkflows(object):
                 ObjectMeta()
                 .name(ArgoWorkflows._sensor_name(self.name))
                 .namespace(KUBERNETES_NAMESPACE)
+                .labels(self._base_labels)
                 .label("app.kubernetes.io/name", "metaflow-sensor")
-                .label("app.kubernetes.io/part-of", "metaflow")
-                .labels(self.kubernetes_labels)
-                .annotations(annotations)
+                .annotations(self._base_annotations)
             )
             .spec(
                 SensorSpec().template(
@@ -2887,7 +2868,7 @@ class ArgoWorkflows(object):
                         ObjectMeta()
                         .label("app.kubernetes.io/name", "metaflow-sensor")
                         .label("app.kubernetes.io/part-of", "metaflow")
-                        .annotations(annotations)
+                        .annotations(self._base_annotations)
                     )
                     .container(
                         # Run sensor in guaranteed QoS. The sensor isn't doing a lot
@@ -2934,6 +2915,7 @@ class ArgoWorkflows(object):
                                         "metadata": {
                                             "generateName": "%s-" % self.name,
                                             "namespace": KUBERNETES_NAMESPACE,
+                                            # Useful to paint the UI
                                             "annotations": {
                                                 "metaflow/triggered_by": json.dumps(
                                                     [

metaflow/plugins/argo/argo_workflows_cli.py

@@ -8,7 +8,6 @@ from time import sleep
 
 from metaflow import JSONType, Run, current, decorators, parameters
 from metaflow._vendor import click
-from metaflow.client.core import get_metadata
 from metaflow.exception import (
     MetaflowException,
     MetaflowInternalError,

metaflow/plugins/argo/argo_workflows_decorator.py

@@ -1,14 +1,12 @@
 import json
 import os
-import time
 
 
 from metaflow import current
 from metaflow.decorators import StepDecorator
 from metaflow.events import Trigger
 from metaflow.metadata_provider import MetaDatum
-from metaflow.metaflow_config import ARGO_EVENTS_WEBHOOK_URL
-from metaflow.graph import DAGNode, FlowGraph
+from metaflow.graph import FlowGraph
 from metaflow.flowspec import FlowSpec
 from .argo_events import ArgoEvent
 
@@ -42,7 +40,7 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
                 if payload != "null":  # Argo-Workflow's None
                     try:
                         payload = json.loads(payload)
-                    except (TypeError, ValueError) as e:
+                    except (TypeError, ValueError):
                         # There could be arbitrary events that Metaflow doesn't know of
                         payload = {}
                     triggers.append(

metaflow/plugins/argo/jobset_input_paths.py

@@ -1,5 +1,4 @@
 import sys
-from hashlib import md5
 
 
 def generate_input_paths(run_id, step_name, task_id_entropy, num_parallel):

metaflow/plugins/aws/aws_utils.py

@@ -1,5 +1,4 @@
 import re
-import requests
 
 from metaflow.exception import MetaflowException
 
@@ -30,6 +29,10 @@ def get_ec2_instance_metadata():
         - ec2-region
         - ec2-availability-zone
     """
+
+    # TODO: Remove dependency on requests
+    import requests
+
     meta = {}
     # Capture AWS instance identity metadata. This is best-effort only since
     # access to this end-point might be blocked on AWS and not available
@@ -159,6 +162,8 @@ def compute_resource_attributes(decos, compute_deco, resource_defaults):
                         # Here we don't have ints, so we compare the value and raise
                         # an exception if not equal
                         if my_val != v:
+                            # TODO: Throw a better exception since the user has no
+                            #       knowledge of 'compute' decorator
                             raise MetaflowException(
                                 "'resources' and compute decorator have conflicting "
                                 "values for '%s'. Please use consistent values or "

metaflow/plugins/aws/batch/batch_client.py

@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
-from collections import defaultdict, deque
+from collections import defaultdict
 import copy
 import random
-import select
-import sys
 import time
 import hashlib
 

metaflow/plugins/aws/batch/batch_decorator.py

@@ -1,34 +1,30 @@
 import os
-import sys
 import platform
-import requests
+import sys
 import time
 
-from metaflow import util
 from metaflow import R, current
-
 from metaflow.decorators import StepDecorator
-from metaflow.plugins.resources_decorator import ResourcesDecorator
-from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.metadata_provider import MetaDatum
 from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
-    ECS_S3_ACCESS_IAM_ROLE,
-    BATCH_JOB_QUEUE,
     BATCH_CONTAINER_IMAGE,
     BATCH_CONTAINER_REGISTRY,
-    ECS_FARGATE_EXECUTION_ROLE,
+    BATCH_JOB_QUEUE,
     DATASTORE_LOCAL_DIR,
+    ECS_FARGATE_EXECUTION_ROLE,
+    ECS_S3_ACCESS_IAM_ROLE,
 )
+from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
 from metaflow.unbounded_foreach import UBF_CONTROL
 
-from .batch import BatchException
 from ..aws_utils import (
     compute_resource_attributes,
     get_docker_registry,
     get_ec2_instance_metadata,
 )
+from .batch import BatchException
 
 
 class BatchDecorator(StepDecorator):
@@ -280,6 +276,10 @@ class BatchDecorator(StepDecorator):
             # Metaflow would be running the container agent compatible with
             # version V4.
             # https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint.html
+
+            # TODO: Remove dependency on requests
+            import requests
+
             try:
                 logs_meta = (
                     requests.get(url=os.environ["ECS_CONTAINER_METADATA_URI_V4"])
@@ -386,7 +386,7 @@ class BatchDecorator(StepDecorator):
                             len(flow._control_mapper_tasks),
                         )
                     )
-            except Exception as e:
+            except Exception:
                 pass
         raise Exception(
             "Batch secondary workers did not finish in %s seconds" % TIMEOUT

metaflow/plugins/aws/step_functions/dynamo_db_client.py

@@ -1,8 +1,5 @@
-import os
 import time
 
-import requests
-
 from metaflow.metaflow_config import SFN_DYNAMO_DB_TABLE
 
 

metaflow/plugins/aws/step_functions/production_token.py

@@ -3,7 +3,7 @@ import os
 import random
 import string
 import zlib
-from itertools import dropwhile, islice
+from itertools import dropwhile
 
 from metaflow.util import to_bytes
 

metaflow/plugins/aws/step_functions/step_functions.py

@@ -236,7 +236,7 @@ class StepFunctions(object):
                 return parameters.get("metaflow.owner"), parameters.get(
                     "metaflow.production_token"
                 )
-            except KeyError as e:
+            except KeyError:
                 raise StepFunctionsException(
                     "An existing non-metaflow "
                     "workflow with the same name as "

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -4,7 +4,6 @@ import re
 from hashlib import sha1
 
 from metaflow import JSONType, current, decorators, parameters
-from metaflow.client.core import get_metadata
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (

metaflow/plugins/aws/step_functions/step_functions_decorator.py

@@ -1,4 +1,3 @@
-import json
 import os
 import time
 

metaflow/plugins/aws/step_functions/step_functions_deployer_objects.py

@@ -1,6 +1,5 @@
 import sys
 import json
-import tempfile
 from typing import ClassVar, Optional, List
 
 from metaflow.plugins.aws.step_functions.step_functions import StepFunctions