metaflow 2.12.8__py2.py3-none-any.whl → 2.12.9__py2.py3-none-any.whl

metaflow/plugins/parallel_decorator.py

@@ -1,11 +1,34 @@
+from collections import namedtuple
 from metaflow.decorators import StepDecorator
-from metaflow.unbounded_foreach import UBF_CONTROL
+from metaflow.unbounded_foreach import UBF_CONTROL, CONTROL_TASK_TAG
 from metaflow.exception import MetaflowException
+from metaflow.metadata import MetaDatum
+from metaflow.metaflow_current import current, Parallel
 import os
 import sys
 
 
 class ParallelDecorator(StepDecorator):
+    """
+    MF Add To Current
+    -----------------
+    parallel -> metaflow.metaflow_current.Parallel
+
+        @@ Returns
+        -------
+        Parallel
+            `namedtuple` with the following fields:
+                - main_ip : str
+                    The IP address of the control task.
+                - num_nodes : int
+                    The total number of tasks created by @parallel
+                - node_index : int
+                    The index of the current task in all the @parallel tasks.
+
+    is_parallel -> bool
+        True if the current step is a @parallel step.
+    """
+
     name = "parallel"
     defaults = {}
     IS_PARALLEL = True
@@ -16,7 +39,6 @@ class ParallelDecorator(StepDecorator):
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context
     ):
-
         if ubf_context == UBF_CONTROL:
             num_parallel = cli_args.task.ubf_iter.num_parallel
             cli_args.command_options["num-parallel"] = str(num_parallel)
@@ -25,6 +47,82 @@ class ParallelDecorator(StepDecorator):
         self, flow, graph, step_name, decorators, environment, flow_datastore, logger
     ):
         self.environment = environment
+        # Previously, the `parallel` property was a hardcoded, static property within `current`.
+        # Whenever `current.parallel` was called, it returned a named tuple with values coming from
+        # environment variables, loaded dynamically at runtime.
+        # Now, many of these environment variables are set by compute-related decorators in `task_pre_step`.
+        # This necessitates ensuring the correct ordering of the `parallel` and compute decorators if we want to
+        # statically set the namedtuple via `current._update_env` in `task_pre_step`. Hence we avoid using
+        # `current._update_env` since:
+        # - it will set a static named tuple, resolving environment variables only once (at the time of calling `current._update_env`).
+        # - we cannot guarantee the order of calling the decorator's `task_pre_step` (calling `current._update_env` may not set
+        #   the named tuple with the correct values).
+        # Therefore, we explicitly set the property in `step_init` to ensure the property can resolve the appropriate values in the named tuple
+        # when accessed at runtime.
+        setattr(
+            current.__class__,
+            "parallel",
+            property(
+                fget=lambda _: Parallel(
+                    main_ip=os.environ.get("MF_PARALLEL_MAIN_IP", "127.0.0.1"),
+                    num_nodes=int(os.environ.get("MF_PARALLEL_NUM_NODES", "1")),
+                    node_index=int(os.environ.get("MF_PARALLEL_NODE_INDEX", "0")),
+                )
+            ),
+        )
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        from metaflow import current
+
+        # Set `is_parallel` to `True` in `current` just like we
+        # with `is_production` in the project decorator.
+        current._update_env(
+            {
+                "is_parallel": True,
+            }
+        )
+
+        self.input_paths = [obj.pathspec for obj in inputs]
+        task_metadata_list = [
+            MetaDatum(
+                field="parallel-world-size",
+                value=flow._parallel_ubf_iter.num_parallel,
+                type="parallel-world-size",
+                tags=["attempt_id:{0}".format(0)],
+            )
+        ]
+        if ubf_context == UBF_CONTROL:
+            # A Task's tags are now those of its ancestral Run, so we are not able
+            # to rely on a task's tags to indicate the presence of a control task
+            # so, on top of adding the tags above, we also add a task metadata
+            # entry indicating that this is a "control task".
+            #
+            # Here we will also add a task metadata entry to indicate "control
+            # task". Within the metaflow repo, the only dependency of such a
+            # "control task" indicator is in the integration test suite (see
+            # Step.control_tasks() in client API).
+            task_metadata_list += [
+                MetaDatum(
+                    field="internal_task_type",
+                    value=CONTROL_TASK_TAG,
+                    type="internal_task_type",
+                    tags=["attempt_id:{0}".format(0)],
+                )
+            ]
+        metadata.register_metadata(run_id, step_name, task_id, task_metadata_list)
 
     def task_decorate(
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
@@ -47,6 +145,7 @@ class ParallelDecorator(StepDecorator):
                 env_to_use,
                 _step_func_with_setup,
                 retry_count,
+                ",".join(self.input_paths),
             )
         else:
             return _step_func_with_setup
@@ -56,7 +155,9 @@ class ParallelDecorator(StepDecorator):
         pass
 
 
-def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_count):
+def _local_multinode_control_task_step_func(
+    flow, env_to_use, step_func, retry_count, input_paths
+):
     """
     Used as multinode UBF control task when run in local mode.
     """
@@ -80,10 +181,7 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
     run_id = current.run_id
     step_name = current.step_name
     control_task_id = current.task_id
-
-    (_, split_step_name, split_task_id) = control_task_id.split("-")[1:]
     # UBF handling for multinode case
-    top_task_id = control_task_id.replace("control-", "")  # chop "-0"
     mapper_task_ids = [control_task_id]
     # If we are running inside Conda, we use the base executable FIRST;
     # the conda environment will then be used when runtime_step_cli is
@@ -93,12 +191,13 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
     script = sys.argv[0]
 
     # start workers
+    # TODO: Logs for worker processes are assigned to control process as of today, which
+    #       should be fixed at some point
     subprocesses = []
     for node_index in range(1, num_parallel):
-        task_id = "%s_node_%d" % (top_task_id, node_index)
+        task_id = "%s_node_%d" % (control_task_id, node_index)
         mapper_task_ids.append(task_id)
         os.environ["MF_PARALLEL_NODE_INDEX"] = str(node_index)
-        input_paths = "%s/%s/%s" % (run_id, split_step_name, split_task_id)
         # Override specific `step` kwargs.
         kwargs = cli_args.step_kwargs
         kwargs["split_index"] = str(node_index)
@@ -109,6 +208,7 @@ def _local_multinode_control_task_step_func(flow, env_to_use, step_func, retry_c
         kwargs["retry_count"] = str(retry_count)
 
         cmd = cli_args.step_command(executable, script, step_name, step_kwargs=kwargs)
+
         p = subprocess.Popen(cmd)
         subprocesses.append(p)
 

metaflow/plugins/secrets/secrets_decorator.py

@@ -4,7 +4,7 @@ import re
 from metaflow.exception import MetaflowException
 from metaflow.decorators import StepDecorator
 from metaflow.metaflow_config import DEFAULT_SECRETS_ROLE
-from metaflow.unbounded_foreach import UBF_CONTROL
+from metaflow.unbounded_foreach import UBF_TASK
 
 from typing import Any, Dict, List, Union
 
@@ -210,8 +210,17 @@ class SecretsDecorator(StepDecorator):
         ubf_context,
         inputs,
     ):
-        if ubf_context == UBF_CONTROL:
-            """control tasks (as used in "unbounded for each") don't need secrets"""
+        if (
+            ubf_context
+            and ubf_context == UBF_TASK
+            and os.environ.get("METAFLOW_RUNTIME_ENVIRONMENT", "local") == "local"
+        ):
+            # We will skip the secret injection for "locally" launched UBF_TASK (worker) tasks
+            # When we "locally" run @parallel tasks, the control task will create the worker tasks and the environment variables
+            # of the control task are inherited by the worker tasks. If we don't skip setting secrets in the worker task then the
+            # worker tasks will try to set the environment variables again which will cause a clash with the control task's env vars,
+            # causing the @secrets' `task_pre_step` to fail. In remote settings, (e.g. AWS Batch/Kubernetes), the worker task and
+            # control task are independently created so there is no chances of an env var clash.
             return
         # List of pairs (secret_spec, env_vars_from_this_spec)
         all_secrets_env_vars = []

metaflow/plugins/test_unbounded_foreach_decorator.py

@@ -8,8 +8,14 @@ import sys
 from metaflow.cli_args import cli_args
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
-from metaflow.unbounded_foreach import UnboundedForeachInput, UBF_CONTROL, UBF_TASK
+from metaflow.unbounded_foreach import (
+    UnboundedForeachInput,
+    UBF_CONTROL,
+    UBF_TASK,
+    CONTROL_TASK_TAG,
+)
 from metaflow.util import to_unicode
+from metaflow.metadata import MetaDatum
 
 
 class InternalTestUnboundedForeachInput(UnboundedForeachInput):
@@ -60,13 +66,42 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
     ):
         self.environment = environment
 
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        if ubf_context == UBF_CONTROL:
+            metadata.register_metadata(
+                run_id,
+                step_name,
+                task_id,
+                [
+                    MetaDatum(
+                        field="internal_task_type",
+                        value=CONTROL_TASK_TAG,
+                        type="internal_task_type",
+                        tags=["attempt_id:{0}".format(0)],
+                    )
+                ],
+            )
+            self.input_paths = [obj.pathspec for obj in inputs]
+
     def control_task_step_func(self, flow, graph, retry_count):
         from metaflow import current
 
         run_id = current.run_id
         step_name = current.step_name
         control_task_id = current.task_id
-        (_, split_step_name, split_task_id) = control_task_id.split("-")[1:]
         # If we are running inside Conda, we use the base executable FIRST;
         # the conda environment will then be used when runtime_step_cli is
         # called. This is so that it can properly set up all the metaflow
@@ -94,10 +129,10 @@ class InternalTestUnboundedForeachDecorator(StepDecorator):
         mapper_tasks = []
 
         for i in range(foreach_num_splits):
-            task_id = "%s-%d" % (control_task_id.replace("control-", "test-ubf-"), i)
+            task_id = "%s-%d" % (control_task_id, i)
             pathspec = "%s/%s/%s" % (run_id, step_name, task_id)
             mapper_tasks.append(to_unicode(pathspec))
-            input_paths = "%s/%s/%s" % (run_id, split_step_name, split_task_id)
+            input_paths = ",".join(self.input_paths)
 
             # Override specific `step` kwargs.
             kwargs = cli_args.step_kwargs

metaflow/runner/deployer.py

@@ -0,0 +1,386 @@
+import os
+import sys
+import json
+import importlib
+import functools
+import tempfile
+from typing import Optional, Dict, ClassVar
+
+from metaflow.exception import MetaflowNotFound
+from metaflow.runner.subprocess_manager import CommandManager, SubprocessManager
+from metaflow.runner.utils import read_from_file_when_ready
+
+
+def handle_timeout(tfp_runner_attribute, command_obj: CommandManager):
+    """
+    Handle the timeout for a running subprocess command that reads a file
+    and raises an error with appropriate logs if a TimeoutError occurs.
+
+    Parameters
+    ----------
+    tfp_runner_attribute : NamedTemporaryFile
+        Temporary file that stores runner attribute data.
+    command_obj : CommandManager
+        Command manager object that encapsulates the running command details.
+
+    Returns
+    -------
+    str
+        Content read from the temporary file.
+
+    Raises
+    ------
+    RuntimeError
+        If a TimeoutError occurs, it raises a RuntimeError with the command's
+        stdout and stderr logs.
+    """
+    try:
+        content = read_from_file_when_ready(tfp_runner_attribute.name, timeout=10)
+        return content
+    except TimeoutError as e:
+        stdout_log = open(command_obj.log_files["stdout"]).read()
+        stderr_log = open(command_obj.log_files["stderr"]).read()
+        command = " ".join(command_obj.command)
+        error_message = "Error executing: '%s':\n" % command
+        if stdout_log.strip():
+            error_message += "\nStdout:\n%s\n" % stdout_log
+        if stderr_log.strip():
+            error_message += "\nStderr:\n%s\n" % stderr_log
+        raise RuntimeError(error_message) from e
+
+
+def get_lower_level_group(
+    api, top_level_kwargs: Dict, _type: Optional[str], deployer_kwargs: Dict
+):
+    """
+    Retrieve a lower-level group from the API based on the type and provided arguments.
+
+    Parameters
+    ----------
+    api : MetaflowAPI
+        Metaflow API instance.
+    top_level_kwargs : Dict
+        Top-level keyword arguments to pass to the API.
+    _type : str
+        Type of the deployer implementation to target.
+    deployer_kwargs : Dict
+        Keyword arguments specific to the deployer.
+
+    Returns
+    -------
+    Any
+        The lower-level group object retrieved from the API.
+
+    Raises
+    ------
+    ValueError
+        If the `_type` is None.
+    """
+    if _type is None:
+        raise ValueError(
+            "DeployerImpl doesn't have a 'TYPE' to target. Please use a sub-class of DeployerImpl."
+        )
+    return getattr(api(**top_level_kwargs), _type)(**deployer_kwargs)
+
+
+class Deployer(object):
+    """
+    Use the `Deployer` class to configure and access one of the production
+    orchestrators supported by Metaflow.
+
+    Parameters
+    ----------
+    flow_file : str
+        Path to the flow file to deploy.
+    show_output : bool, default True
+        Show the 'stdout' and 'stderr' to the console by default.
+    profile : Optional[str], default None
+        Metaflow profile to use for the deployment. If not specified, the default
+        profile is used.
+    env : Optional[Dict[str, str]], default None
+        Additional environment variables to set for the deployment.
+    cwd : Optional[str], default None
+        The directory to run the subprocess in; if not specified, the current
+        directory is used.
+    **kwargs : Any
+        Additional arguments that you would pass to `python myflow.py` before
+        the deployment command.
+    """
+
+    def __init__(
+        self,
+        flow_file: str,
+        show_output: bool = True,
+        profile: Optional[str] = None,
+        env: Optional[Dict] = None,
+        cwd: Optional[str] = None,
+        **kwargs
+    ):
+        self.flow_file = flow_file
+        self.show_output = show_output
+        self.profile = profile
+        self.env = env
+        self.cwd = cwd
+        self.top_level_kwargs = kwargs
+
+        from metaflow.plugins import DEPLOYER_IMPL_PROVIDERS
+
+        for provider_class in DEPLOYER_IMPL_PROVIDERS:
+            # TYPE is the name of the CLI groups i.e.
+            # `argo-workflows` instead of `argo_workflows`
+            # The injected method names replace '-' by '_' though.
+            method_name = provider_class.TYPE.replace("-", "_")
+            setattr(Deployer, method_name, self.__make_function(provider_class))
+
+    def __make_function(self, deployer_class):
+        """
+        Create a function for the given deployer class.
+
+        Parameters
+        ----------
+        deployer_class : Type[DeployerImpl]
+            Deployer implementation class.
+
+        Returns
+        -------
+        Callable
+            Function that initializes and returns an instance of the deployer class.
+        """
+
+        def f(self, **deployer_kwargs):
+            return deployer_class(
+                deployer_kwargs=deployer_kwargs,
+                flow_file=self.flow_file,
+                show_output=self.show_output,
+                profile=self.profile,
+                env=self.env,
+                cwd=self.cwd,
+                **self.top_level_kwargs
+            )
+
+        return f
+
+
+class TriggeredRun(object):
+    """
+    TriggeredRun class represents a run that has been triggered on a production orchestrator.
+
+    Only when the `start` task starts running, the `run` object corresponding to the run
+    becomes available.
+    """
+
+    def __init__(
+        self,
+        deployer: "DeployerImpl",
+        content: str,
+    ):
+        self.deployer = deployer
+        content_json = json.loads(content)
+        self.metadata_for_flow = content_json.get("metadata")
+        self.pathspec = content_json.get("pathspec")
+        self.name = content_json.get("name")
+
+    def _enrich_object(self, env):
+        """
+        Enrich the TriggeredRun object with additional properties and methods.
+
+        Parameters
+        ----------
+        env : dict
+            Environment dictionary containing properties and methods to add.
+        """
+        for k, v in env.items():
+            if isinstance(v, property):
+                setattr(self.__class__, k, v)
+            elif callable(v):
+                setattr(self, k, functools.partial(v, self))
+            else:
+                setattr(self.__class__, k, property(fget=lambda _, v=v: v))
+
+    @property
+    def run(self):
+        """
+        Retrieve the `Run` object for the triggered run.
+
+        Note that Metaflow `Run` becomes available only when the `start` task
+        has started executing.
+
+        Returns
+        -------
+        Run, optional
+            Metaflow Run object if the `start` step has started executing, otherwise None.
+        """
+        from metaflow import Run
+
+        try:
+            return Run(self.pathspec, _namespace_check=False)
+        except MetaflowNotFound:
+            return None
+
+
+class DeployedFlow(object):
+    """
+    DeployedFlow class represents a flow that has been deployed.
+
+    Parameters
+    ----------
+    deployer : DeployerImpl
+        Instance of the deployer implementation.
+    """
+
+    def __init__(self, deployer: "DeployerImpl"):
+        self.deployer = deployer
+
+    def _enrich_object(self, env):
+        """
+        Enrich the DeployedFlow object with additional properties and methods.
+
+        Parameters
+        ----------
+        env : dict
+            Environment dictionary containing properties and methods to add.
+        """
+        for k, v in env.items():
+            if isinstance(v, property):
+                setattr(self.__class__, k, v)
+            elif callable(v):
+                setattr(self, k, functools.partial(v, self))
+            else:
+                setattr(self.__class__, k, property(fget=lambda _, v=v: v))
+
+
+class DeployerImpl(object):
+    """
+    Base class for deployer implementations. Each implementation should define a TYPE
+    class variable that matches the name of the CLI group.
+
+    Parameters
+    ----------
+    flow_file : str
+        Path to the flow file to deploy.
+    show_output : bool, default True
+        Show the 'stdout' and 'stderr' to the console by default.
+    profile : Optional[str], default None
+        Metaflow profile to use for the deployment. If not specified, the default
+        profile is used.
+    env : Optional[Dict], default None
+        Additional environment variables to set for the deployment.
+    cwd : Optional[str], default None
+        The directory to run the subprocess in; if not specified, the current
+        directory is used.
+    **kwargs : Any
+        Additional arguments that you would pass to `python myflow.py` before
+        the deployment command.
+    """
+
+    TYPE: ClassVar[Optional[str]] = None
+
+    def __init__(
+        self,
+        flow_file: str,
+        show_output: bool = True,
+        profile: Optional[str] = None,
+        env: Optional[Dict] = None,
+        cwd: Optional[str] = None,
+        **kwargs
+    ):
+        if self.TYPE is None:
+            raise ValueError(
+                "DeployerImpl doesn't have a 'TYPE' to target. Please use a sub-class of DeployerImpl."
+            )
+
+        if "metaflow.cli" in sys.modules:
+            importlib.reload(sys.modules["metaflow.cli"])
+        from metaflow.cli import start
+        from metaflow.runner.click_api import MetaflowAPI
+
+        self.flow_file = flow_file
+        self.show_output = show_output
+        self.profile = profile
+        self.env = env
+        self.cwd = cwd
+
+        self.env_vars = os.environ.copy()
+        self.env_vars.update(self.env or {})
+        if self.profile:
+            self.env_vars["METAFLOW_PROFILE"] = profile
+
+        self.spm = SubprocessManager()
+        self.top_level_kwargs = kwargs
+        self.api = MetaflowAPI.from_cli(self.flow_file, start)
+
+    def __enter__(self) -> "DeployerImpl":
+        return self
+
+    def create(self, **kwargs) -> DeployedFlow:
+        """
+        Create a deployed flow using the deployer implementation.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional arguments to pass to `create` corresponding to the
+            command line arguments of `create`
+
+        Returns
+        -------
+        DeployedFlow
+            DeployedFlow object representing the deployed flow.
+
+        Raises
+        ------
+        Exception
+            If there is an error during deployment.
+        """
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tfp_runner_attribute = tempfile.NamedTemporaryFile(
+                dir=temp_dir, delete=False
+            )
+            # every subclass needs to have `self.deployer_kwargs`
+            command = get_lower_level_group(
+                self.api, self.top_level_kwargs, self.TYPE, self.deployer_kwargs
+            ).create(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
+
+            pid = self.spm.run_command(
+                [sys.executable, *command],
+                env=self.env_vars,
+                cwd=self.cwd,
+                show_output=self.show_output,
+            )
+
+            command_obj = self.spm.get(pid)
+            content = handle_timeout(tfp_runner_attribute, command_obj)
+            content = json.loads(content)
+            self.name = content.get("name")
+            self.flow_name = content.get("flow_name")
+            self.metadata = content.get("metadata")
+
+            if command_obj.process.returncode == 0:
+                deployed_flow = DeployedFlow(deployer=self)
+                self._enrich_deployed_flow(deployed_flow)
+                return deployed_flow
+
+        raise Exception("Error deploying %s to %s" % (self.flow_file, self.TYPE))
+
+    def _enrich_deployed_flow(self, deployed_flow: DeployedFlow):
+        """
+        Enrich the DeployedFlow object with additional properties and methods.
+
+        Parameters
+        ----------
+        deployed_flow : DeployedFlow
+            The DeployedFlow object to enrich.
+        """
+        raise NotImplementedError
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Cleanup resources on exit.
+        """
+        self.cleanup()
+
+    def cleanup(self):
+        """
+        Cleanup resources.
+        """
+        self.spm.cleanup()

metaflow/runner/metaflow_runner.py

@@ -2,33 +2,14 @@ import importlib
 import os
 import sys
 import tempfile
-import time
 from typing import Dict, Iterator, Optional, Tuple
 
 from metaflow import Run, metadata
 
+from .utils import clear_and_set_os_environ, read_from_file_when_ready
 from .subprocess_manager import CommandManager, SubprocessManager
 
 
-def clear_and_set_os_environ(env: Dict):
-    os.environ.clear()
-    os.environ.update(env)
-
-
-def read_from_file_when_ready(file_path: str, timeout: float = 5):
-    start_time = time.time()
-    with open(file_path, "r", encoding="utf-8") as file_pointer:
-        content = file_pointer.read()
-        while not content:
-            if time.time() - start_time > timeout:
-                raise TimeoutError(
-                    "Timeout while waiting for file content from '%s'" % file_path
-                )
-            time.sleep(0.1)
-            content = file_pointer.read()
-        return content
-
-
 class ExecutingRun(object):
     """
     This class contains a reference to a `metaflow.Run` object representing