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

metaflow/plugins/argo/argo_workflows_deployer_objects.py

@@ -1,5 +1,6 @@
 import sys
 import json
+import time
 import tempfile
 from typing import ClassVar, Optional
 
@@ -170,6 +171,47 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
         command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
+    def wait_for_completion(self, timeout: Optional[int] = None):
+        """
+        Wait for the workflow to complete or timeout.
+
+        Parameters
+        ----------
+        timeout : int, optional, default None
+            Maximum time in seconds to wait for workflow completion.
+            If None, waits indefinitely.
+
+        Raises
+        ------
+        TimeoutError
+            If the workflow does not complete within the specified timeout period.
+        """
+        start_time = time.time()
+        check_interval = 5
+        while self.is_running:
+            if timeout is not None and (time.time() - start_time) > timeout:
+                raise TimeoutError(
+                    "Workflow did not complete within specified timeout."
+                )
+            time.sleep(check_interval)
+
+    @property
+    def is_running(self):
+        """
+        Check if the workflow is currently running.
+
+        Returns
+        -------
+        bool
+            True if the workflow status is either 'Pending' or 'Running',
+            False otherwise.
+        """
+        workflow_status = self.status
+        # full list of all states present here:
+        # https://github.com/argoproj/argo-workflows/blob/main/pkg/apis/workflow/v1alpha1/workflow_types.go#L54
+        # we only consider non-terminal states to determine if the workflow has not finished
+        return workflow_status is not None and workflow_status in ["Pending", "Running"]
+
     @property
     def status(self) -> Optional[str]:
         """

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/secrets_manager/aws_secrets_manager_secrets_provider.py

@@ -50,24 +50,27 @@ class AwsSecretsManagerSecretsProvider(SecretsProvider):
         The secret payload from AWS is EITHER a string OR a binary blob.
 
         If the secret contains a string payload ("SecretString"):
-        - if the `parse_secret_string_as_json` option is True (default):
+        - if the `json` option is True (default):
             {SecretString} will be parsed as a JSON. If successfully parsed, AND the JSON contains a
             top-level object, each entry K/V in the object will also be converted to an entry in the result. V will
             always be casted to a string (if not already a string).
-        - If `parse_secret_string_as_json` option is False:
-            {SecretString} will be returned as a single entry in the result, with the key being the secret_id.
+        - If `json` option is False:
+            {SecretString} will be returned as a single entry in the result, where the key is either:
+                - the `secret_id`, OR
+                - the value set by `options={"env_var_name": custom_env_var_name}`.
 
-        Otherwise, the secret contains a binary blob payload ("SecretBinary"). In this case
-        - The result dic contains '{SecretName}': '{SecretBinary}', where {SecretBinary} is a base64-encoded string
+        Otherwise, if the secret contains a binary blob payload ("SecretBinary"):
+        - The result dict contains '{SecretName}': '{SecretBinary}', where {SecretBinary} is a base64-encoded string.
 
-        All keys in the result are sanitized to be more valid environment variable names. This is done on a best effort
+        All keys in the result are sanitized to be more valid environment variable names. This is done on a best-effort
         basis. Further validation is expected to be done by the invoking @secrets decorator itself.
 
-        :param secret_id: ARN or friendly name of the secret
-        :param options: unused
-        :param role: AWS IAM Role ARN to assume before reading the secret
-        :return: dict of environment variables. All keys and values are strings.
+        :param secret_id: ARN or friendly name of the secret.
+        :param options: Dictionary of additional options. E.g., `options={"env_var_name": custom_env_var_name}`.
+        :param role: AWS IAM Role ARN to assume before reading the secret.
+        :return: Dictionary of environment variables. All keys and values are strings.
         """
+
         import botocore
         from metaflow.plugins.aws.aws_client import get_aws_client
 

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

metaflow/plugins/cards/card_creator.py

@@ -122,6 +122,7 @@ class CardCreator:
             executable,
             sys.argv[0],
         ]
+
         cmd += self._top_level_options + [
             "card",
             "create",

metaflow/plugins/cards/card_decorator.py

@@ -1,13 +1,16 @@
+import json
+import os
+import re
+import tempfile
+
 from metaflow.decorators import StepDecorator
 from metaflow.metaflow_current import current
+from metaflow.user_configs.config_options import ConfigInput
+from metaflow.user_configs.config_parameters import dump_config_values
 from metaflow.util import to_unicode
+
 from .component_serializer import CardComponentCollector, get_card_class
 from .card_creator import CardCreator
-
-
-# from metaflow import get_metadata
-import re
-
 from .exception import CARD_ID_PATTERN, TYPE_CHECK_REGEX
 
 ASYNC_TIMEOUT = 30
@@ -111,6 +114,14 @@ class CardDecorator(StepDecorator):
         self._logger = logger
         self.card_options = None
 
+        # We check for configuration options. We do this here before they are
+        # converted to properties.
+        self._config_values = [
+            (config.name, ConfigInput.make_key_name(config.name))
+            for _, config in flow._get_parameters()
+            if config.IS_CONFIG_PARAMETER
+        ]
+
         self.card_options = self.attributes["options"]
 
         evt_name = "step-init"
@@ -146,6 +157,18 @@ class CardDecorator(StepDecorator):
         self._task_datastore = task_datastore
         self._metadata = metadata
 
+        # If we have configs, we need to dump them to a file so we can re-use them
+        # when calling the card creation subprocess.
+        if self._config_values:
+            with tempfile.NamedTemporaryFile(
+                mode="w", encoding="utf-8", delete=False
+            ) as config_file:
+                config_value = dump_config_values(flow)
+                json.dump(config_value, config_file)
+                self._config_file_name = config_file.name
+        else:
+            self._config_file_name = None
+
         card_type = self.attributes["type"]
         card_class = get_card_class(card_type)
 
@@ -179,7 +202,7 @@ class CardDecorator(StepDecorator):
         # we need to ensure that `current.card` has `CardComponentCollector` instantiated only once.
         if not self._is_event_registered("pre-step"):
             self._register_event("pre-step")
-            self._set_card_creator(CardCreator(self._create_top_level_args()))
+            self._set_card_creator(CardCreator(self._create_top_level_args(flow)))
 
             current._update_env(
                 {"card": CardComponentCollector(self._logger, self.card_creator)}
@@ -223,6 +246,13 @@ class CardDecorator(StepDecorator):
             self.card_creator.create(mode="render", final=True, **create_options)
             self.card_creator.create(mode="refresh", final=True, **create_options)
 
+        # Unlink the config file if it exists
+        if self._config_file_name:
+            try:
+                os.unlink(self._config_file_name)
+            except Exception as e:
+                pass
+
     @staticmethod
     def _options(mapping):
         for k, v in mapping.items():
@@ -232,9 +262,13 @@ class CardDecorator(StepDecorator):
                 for value in v:
                     yield "--%s" % k
                     if not isinstance(value, bool):
-                        yield to_unicode(value)
+                        if isinstance(value, tuple):
+                            for val in value:
+                                yield to_unicode(val)
+                        else:
+                            yield to_unicode(value)
 
-    def _create_top_level_args(self):
+    def _create_top_level_args(self, flow):
         top_level_options = {
             "quiet": True,
             "metadata": self._metadata.TYPE,
@@ -247,4 +281,8 @@ class CardDecorator(StepDecorator):
             # We don't provide --with as all execution is taking place in
             # the context of the main process
         }
+        if self._config_values:
+            top_level_options["config-value"] = self._config_values
+            top_level_options["local-config-file"] = self._config_file_name
+
         return list(self._options(top_level_options))

metaflow/plugins/kubernetes/kube_utils.py

@@ -1,7 +1,14 @@
-from metaflow.exception import CommandException
+import re
+from typing import Dict, List, Optional
+from metaflow.exception import CommandException, MetaflowException
 from metaflow.util import get_username, get_latest_run_id
 
 
+# avoid circular import by having the exception class contained here
+class KubernetesException(MetaflowException):
+    headline = "Kubernetes error"
+
+
 def parse_cli_options(flow_name, run_id, user, my_runs, echo):
     if user and my_runs:
         raise CommandException("--user and --my-runs are mutually exclusive.")
@@ -52,3 +59,50 @@ def qos_requests_and_limits(qos: str, cpu: int, memory: int, storage: int):
     # TODO: Add support for BestEffort once there is a use case for it.
     # BestEffort - no limit or requests for cpu/memory
     return qos_requests, qos_limits
+
+
+def validate_kube_labels(
+    labels: Optional[Dict[str, Optional[str]]],
+) -> bool:
+    """Validate label values.
+
+    This validates the kubernetes label values.  It does not validate the keys.
+    Ideally, keys should be static and also the validation rules for keys are
+    more complex than those for values.  For full validation rules, see:
+
+    https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
+    """
+
+    def validate_label(s: Optional[str]):
+        regex_match = r"^(([A-Za-z0-9][-A-Za-z0-9_.]{0,61})?[A-Za-z0-9])?$"
+        if not s:
+            # allow empty label
+            return True
+        if not re.search(regex_match, s):
+            raise KubernetesException(
+                'Invalid value: "%s"\n'
+                "A valid label must be an empty string or one that\n"
+                "  - Consist of alphanumeric, '-', '_' or '.' characters\n"
+                "  - Begins and ends with an alphanumeric character\n"
+                "  - Is at most 63 characters" % s
+            )
+        return True
+
+    return all([validate_label(v) for v in labels.values()]) if labels else True
+
+
+def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):
+    try:
+        ret = {}
+        for item_str in items:
+            item = item_str.split("=", 1)
+            if requires_both:
+                item[1]  # raise IndexError
+            if str(item[0]) in ret:
+                raise KubernetesException("Duplicate key found: %s" % str(item[0]))
+            ret[str(item[0])] = str(item[1]) if len(item) > 1 else None
+        return ret
+    except KubernetesException as e:
+        raise e
+    except (AttributeError, IndexError):
+        raise KubernetesException("Unable to parse kubernetes list: %s" % items)

metaflow/plugins/kubernetes/kubernetes.py

@@ -1,11 +1,8 @@
-import copy
 import json
 import math
 import os
-import re
 import shlex
 import time
-from typing import Dict, List, Optional
 from uuid import uuid4
 
 from metaflow import current, util
@@ -35,7 +32,6 @@ from metaflow.metaflow_config import (
     DEFAULT_SECRETS_BACKEND_TYPE,
     GCP_SECRET_MANAGER_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
-    KUBERNETES_LABELS,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
     OTEL_ENDPOINT,
     S3_ENDPOINT_URL,
@@ -193,6 +189,7 @@ class Kubernetes(object):
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
+        annotations=None,
         shared_memory=None,
         port=None,
         num_parallel=None,
@@ -304,10 +301,6 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
-        _labels = self._get_labels(labels)
-        for k, v in _labels.items():
-            jobset.label(k, v)
-
         for k in list(
             [] if not secrets else [secrets] if isinstance(secrets, str) else secrets
         ) + KUBERNETES_SECRETS.split(","):
@@ -379,13 +372,16 @@ class Kubernetes(object):
         for name, value in env.items():
             jobset.environment_variable(name, value)
 
-        annotations = {
+        system_annotations = {
             "metaflow/user": user,
             "metaflow/flow_name": flow_name,
             "metaflow/control-task-id": task_id,
+            "metaflow/run_id": run_id,
+            "metaflow/step_name": step_name,
+            "metaflow/attempt": attempt,
         }
         if current.get("project_name"):
-            annotations.update(
+            system_annotations.update(
                 {
                     "metaflow/project_name": current.project_name,
                     "metaflow/branch_name": current.branch_name,
@@ -393,15 +389,15 @@ class Kubernetes(object):
                 }
             )
 
-        for name, value in annotations.items():
-            jobset.annotation(name, value)
+        system_labels = {
+            "app.kubernetes.io/name": "metaflow-task",
+            "app.kubernetes.io/part-of": "metaflow",
+        }
 
-        (
-            jobset.annotation("metaflow/run_id", run_id)
-            .annotation("metaflow/step_name", step_name)
-            .annotation("metaflow/attempt", attempt)
-            .label("app.kubernetes.io/name", "metaflow-task")
-            .label("app.kubernetes.io/part-of", "metaflow")
+        jobset.labels({**({} if not labels else labels), **system_labels})
+
+        jobset.annotations(
+            {**({} if not annotations else annotations), **system_annotations}
         )
         # We need this task-id set so that all the nodes are aware of the control
         # task's task-id. These "MF_" variables populate the `current.parallel` namedtuple
@@ -491,6 +487,7 @@ class Kubernetes(object):
         port=None,
         name_pattern=None,
         qos=None,
+        annotations=None,
     ):
         if env is None:
             env = {}
@@ -523,7 +520,8 @@ class Kubernetes(object):
                 retries=0,
                 step_name=step_name,
                 tolerations=tolerations,
-                labels=self._get_labels(labels),
+                labels=labels,
+                annotations=annotations,
                 use_tmpfs=use_tmpfs,
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,
@@ -642,13 +640,25 @@ class Kubernetes(object):
 
         for name, value in env.items():
             job.environment_variable(name, value)
+        # Add job specific labels
+        system_labels = {
+            "app.kubernetes.io/name": "metaflow-task",
+            "app.kubernetes.io/part-of": "metaflow",
+        }
+        for name, value in system_labels.items():
+            job.label(name, value)
 
-        annotations = {
-            "metaflow/user": user,
+        # Add job specific annotations not set in the decorator.
+        system_annotations = {
             "metaflow/flow_name": flow_name,
+            "metaflow/run_id": run_id,
+            "metaflow/step_name": step_name,
+            "metaflow/task_id": task_id,
+            "metaflow/attempt": attempt,
+            "metaflow/user": user,
         }
         if current.get("project_name"):
-            annotations.update(
+            system_annotations.update(
                 {
                     "metaflow/project_name": current.project_name,
                     "metaflow/branch_name": current.branch_name,
@@ -656,7 +666,7 @@ class Kubernetes(object):
                 }
             )
 
-        for name, value in annotations.items():
+        for name, value in system_annotations.items():
             job.annotation(name, value)
 
         (
@@ -775,60 +785,3 @@ class Kubernetes(object):
             "stderr",
             job_id=self._job.id,
         )
-
-    @staticmethod
-    def _get_labels(extra_labels=None):
-        if extra_labels is None:
-            extra_labels = {}
-        env_labels = KUBERNETES_LABELS.split(",") if KUBERNETES_LABELS else []
-        env_labels = parse_kube_keyvalue_list(env_labels, False)
-        labels = {**env_labels, **extra_labels}
-        validate_kube_labels(labels)
-        return labels
-
-
-def validate_kube_labels(
-    labels: Optional[Dict[str, Optional[str]]],
-) -> bool:
-    """Validate label values.
-
-    This validates the kubernetes label values.  It does not validate the keys.
-    Ideally, keys should be static and also the validation rules for keys are
-    more complex than those for values.  For full validation rules, see:
-
-    https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
-    """
-
-    def validate_label(s: Optional[str]):
-        regex_match = r"^(([A-Za-z0-9][-A-Za-z0-9_.]{0,61})?[A-Za-z0-9])?$"
-        if not s:
-            # allow empty label
-            return True
-        if not re.search(regex_match, s):
-            raise KubernetesException(
-                'Invalid value: "%s"\n'
-                "A valid label must be an empty string or one that\n"
-                "  - Consist of alphanumeric, '-', '_' or '.' characters\n"
-                "  - Begins and ends with an alphanumeric character\n"
-                "  - Is at most 63 characters" % s
-            )
-        return True
-
-    return all([validate_label(v) for v in labels.values()]) if labels else True
-
-
-def parse_kube_keyvalue_list(items: List[str], requires_both: bool = True):
-    try:
-        ret = {}
-        for item_str in items:
-            item = item_str.split("=", 1)
-            if requires_both:
-                item[1]  # raise IndexError
-            if str(item[0]) in ret:
-                raise KubernetesException("Duplicate key found: %s" % str(item[0]))
-            ret[str(item[0])] = str(item[1]) if len(item) > 1 else None
-        return ret
-    except KubernetesException as e:
-        raise e
-    except (AttributeError, IndexError):
-        raise KubernetesException("Unable to parse kubernetes list: %s" % items)