ob-metaflow 2.12.30.2__py2.py3-none-any.whl → 2.13.6.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
 
@@ -10,7 +11,7 @@ from metaflow.metaflow_config import KUBERNETES_NAMESPACE
 from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
 from metaflow.runner.deployer import Deployer, DeployedFlow, TriggeredRun
 
-from metaflow.runner.utils import get_lower_level_group, handle_timeout
+from metaflow.runner.utils import get_lower_level_group, handle_timeout, temporary_fifo
 
 
 def generate_fake_flow_file_contents(
@@ -97,6 +98,7 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
     def unsuspend(self, **kwargs) -> bool:
@@ -131,6 +133,7 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
     def terminate(self, **kwargs) -> bool:
@@ -165,8 +168,50 @@ class ArgoWorkflowsTriggeredRun(TriggeredRun):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        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]:
         """
@@ -319,6 +364,7 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
     def trigger(self, **kwargs) -> ArgoWorkflowsTriggeredRun:
@@ -341,18 +387,14 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
         Exception
             If there is an error during the trigger process.
         """
-        with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_runner_attribute = tempfile.NamedTemporaryFile(
-                dir=temp_dir, delete=False
-            )
-
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
             # every subclass needs to have `self.deployer_kwargs`
             command = get_lower_level_group(
                 self.deployer.api,
                 self.deployer.top_level_kwargs,
                 self.deployer.TYPE,
                 self.deployer.deployer_kwargs,
-            ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+            ).trigger(deployer_attribute_file=attribute_file_path, **kwargs)
 
             pid = self.deployer.spm.run_command(
                 [sys.executable, *command],
@@ -363,9 +405,9 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
 
             command_obj = self.deployer.spm.get(pid)
             content = handle_timeout(
-                tfp_runner_attribute, command_obj, self.deployer.file_read_timeout
+                attribute_file_fd, command_obj, self.deployer.file_read_timeout
             )
-
+            command_obj.sync_wait()
             if command_obj.process.returncode == 0:
                 return ArgoWorkflowsTriggeredRun(
                     deployer=self.deployer, content=content

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
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
@@ -31,6 +30,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
@@ -160,6 +163,8 @@ def compute_resource_attributes(decos, compute_deco, step_name, resource_default
                         # 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):
@@ -138,8 +134,8 @@ class BatchDecorator(StepDecorator):
     supports_conda_environment = True
     target_platform = "linux-64"
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(BatchDecorator, self).__init__(attributes, statically_defined)
+    def init(self):
+        super(BatchDecorator, self).init()
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
@@ -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

@@ -18,6 +18,7 @@ from metaflow.metaflow_config import (
     SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH,
 )
 from metaflow.parameters import deploy_time_eval
+from metaflow.user_configs.config_options import ConfigInput
 from metaflow.util import dict_to_cli_options, to_pascalcase
 
 from ..batch.batch import Batch
@@ -71,6 +72,7 @@ class StepFunctions(object):
         self.username = username
         self.max_workers = max_workers
         self.workflow_timeout = workflow_timeout
+        self.config_parameters = self._process_config_parameters()
 
         # https://aws.amazon.com/blogs/aws/step-functions-distributed-map-a-serverless-solution-for-large-scale-parallel-data-processing/
         self.use_distributed_map = use_distributed_map
@@ -234,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 "
@@ -485,6 +487,10 @@ class StepFunctions(object):
                     "case-insensitive." % param.name
                 )
             seen.add(norm)
+            # NOTE: We skip config parameters as these do not have dynamic values,
+            # and need to be treated differently.
+            if param.IS_CONFIG_PARAMETER:
+                continue
 
             is_required = param.kwargs.get("required", False)
             # Throw an exception if a schedule is set for a flow with required
@@ -501,6 +507,27 @@ class StepFunctions(object):
             parameters.append(dict(name=param.name, value=value))
         return parameters
 
+    def _process_config_parameters(self):
+        parameters = []
+        seen = set()
+        for var, param in self.flow._get_parameters():
+            if not param.IS_CONFIG_PARAMETER:
+                continue
+            # Throw an exception if the parameter is specified twice.
+            norm = param.name.lower()
+            if norm in seen:
+                raise MetaflowException(
+                    "Parameter *%s* is specified twice. "
+                    "Note that parameter names are "
+                    "case-insensitive." % param.name
+                )
+            seen.add(norm)
+
+            parameters.append(
+                dict(name=param.name, kv_name=ConfigInput.make_key_name(param.name))
+            )
+        return parameters
+
     def _batch(self, node):
         attrs = {
             # metaflow.user is only used for setting the AWS Job Name.
@@ -747,6 +774,11 @@ class StepFunctions(object):
         metaflow_version["production_token"] = self.production_token
         env["METAFLOW_VERSION"] = json.dumps(metaflow_version)
 
+        # map config values
+        cfg_env = {param["name"]: param["kv_name"] for param in self.config_parameters}
+        if cfg_env:
+            env["METAFLOW_FLOW_CONFIG_VALUE"] = json.dumps(cfg_env)
+
         # Set AWS DynamoDb Table Name for state tracking for for-eaches.
         # There are three instances when metaflow runtime directly interacts
         # with AWS DynamoDB.

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 (
@@ -326,6 +325,7 @@ def make_flow(
 
     # Attach AWS Batch decorator to the flow
     decorators._attach_decorators(obj.flow, [BatchDecorator.name])
+    decorators._init(obj.flow)
     decorators._init_step_decorators(
         obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger
     )

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,12 +1,11 @@
 import sys
 import json
-import tempfile
 from typing import ClassVar, Optional, List
 
 from metaflow.plugins.aws.step_functions.step_functions import StepFunctions
 from metaflow.runner.deployer import DeployedFlow, TriggeredRun
 
-from metaflow.runner.utils import get_lower_level_group, handle_timeout
+from metaflow.runner.utils import get_lower_level_group, handle_timeout, temporary_fifo
 
 
 class StepFunctionsTriggeredRun(TriggeredRun):
@@ -46,6 +45,7 @@ class StepFunctionsTriggeredRun(TriggeredRun):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
 
@@ -174,6 +174,7 @@ class StepFunctionsDeployedFlow(DeployedFlow):
         )
 
         command_obj = self.deployer.spm.get(pid)
+        command_obj.sync_wait()
         return command_obj.process.returncode == 0
 
     def trigger(self, **kwargs) -> StepFunctionsTriggeredRun:
@@ -196,18 +197,14 @@ class StepFunctionsDeployedFlow(DeployedFlow):
         Exception
             If there is an error during the trigger process.
         """
-        with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_runner_attribute = tempfile.NamedTemporaryFile(
-                dir=temp_dir, delete=False
-            )
-
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
             # every subclass needs to have `self.deployer_kwargs`
             command = get_lower_level_group(
                 self.deployer.api,
                 self.deployer.top_level_kwargs,
                 self.deployer.TYPE,
                 self.deployer.deployer_kwargs,
-            ).trigger(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+            ).trigger(deployer_attribute_file=attribute_file_path, **kwargs)
 
             pid = self.deployer.spm.run_command(
                 [sys.executable, *command],
@@ -218,9 +215,10 @@ class StepFunctionsDeployedFlow(DeployedFlow):
 
             command_obj = self.deployer.spm.get(pid)
             content = handle_timeout(
-                tfp_runner_attribute, command_obj, self.deployer.file_read_timeout
+                attribute_file_fd, command_obj, self.deployer.file_read_timeout
             )
 
+            command_obj.sync_wait()
             if command_obj.process.returncode == 0:
                 return StepFunctionsTriggeredRun(
                     deployer=self.deployer, content=content

metaflow/plugins/cards/card_cli.py

@@ -691,10 +691,15 @@ def create(
         try:
             if options is not None:
                 mf_card = filtered_card(
-                    options=options, components=component_arr, graph=graph_dict
+                    options=options,
+                    components=component_arr,
+                    graph=graph_dict,
+                    flow=ctx.obj.flow,
                 )
             else:
-                mf_card = filtered_card(components=component_arr, graph=graph_dict)
+                mf_card = filtered_card(
+                    components=component_arr, graph=graph_dict, flow=ctx.obj.flow
+                )
         except TypeError as e:
             if render_error_card:
                 mf_card = None

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
@@ -73,6 +76,12 @@ class CardDecorator(StepDecorator):
 
     card_creator = None
 
+    _config_values = None
+
+    _config_file_name = None
+
+    task_finished_decos = 0
+
     def __init__(self, *args, **kwargs):
         super(CardDecorator, self).__init__(*args, **kwargs)
         self._task_datastore = None
@@ -103,6 +112,25 @@ class CardDecorator(StepDecorator):
     def _increment_step_counter(cls):
         cls.step_counter += 1
 
+    @classmethod
+    def _increment_completed_counter(cls):
+        cls.task_finished_decos += 1
+
+    @classmethod
+    def _set_config_values(cls, config_values):
+        cls._config_values = config_values
+
+    @classmethod
+    def _set_config_file_name(cls, flow):
+        # Only create a config file from the very first card decorator.
+        if cls._config_values and not cls._config_file_name:
+            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)
+                cls._config_file_name = config_file.name
+
     def step_init(
         self, flow, graph, step_name, decorators, environment, flow_datastore, logger
     ):
@@ -111,6 +139,16 @@ 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._set_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 +184,14 @@ 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.
+        # Since a step can contain multiple card decorators, and all the card creation processes
+        # will reference the same config file (because of how the CardCreator is created (only single class instance)),
+        # we need to ensure that a single config file is being referenced for all card create commands.
+        # This config file will be removed when the last card decorator has finished creating its card.
+        self._set_config_file_name(flow)
+
         card_type = self.attributes["type"]
         card_class = get_card_class(card_type)
 
@@ -179,7 +225,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 +269,8 @@ class CardDecorator(StepDecorator):
             self.card_creator.create(mode="render", final=True, **create_options)
             self.card_creator.create(mode="refresh", final=True, **create_options)
 
+        self._cleanup(step_name)
+
     @staticmethod
     def _options(mapping):
         for k, v in mapping.items():
@@ -232,9 +280,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 +299,23 @@ 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))
+
+    def task_exception(
+        self, exception, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        self._cleanup(step_name)
+
+    def _cleanup(self, step_name):
+        self._increment_completed_counter()
+        if self.task_finished_decos == self.total_decos_on_step[step_name]:
+            # Unlink the config file if it exists
+            if self._config_file_name:
+                try:
+                    os.unlink(self._config_file_name)
+                except Exception as e:
+                    pass