metaflow 2.14.0__py2.py3-none-any.whl → 2.14.1__py2.py3-none-any.whl

metaflow/cli.py

@@ -238,7 +238,6 @@ def version(obj):
     lazy_sources=plugins.get_plugin_cli_path(),
     invoke_without_command=True,
 )
-@tracing.cli_entrypoint("cli/start")
 # Quiet is eager to make sure it is available when processing --config options since
 # we need it to construct a context to pass to any DeployTimeField for the default
 # value.

metaflow/cli_components/run_cmds.py

@@ -166,6 +166,7 @@ def common_run_options(func):
 )
 @click.argument("step-to-rerun", required=False)
 @click.command(help="Resume execution of a previous run of this flow.")
+@tracing.cli("cli/resume")
 @common_run_options
 @click.pass_obj
 def resume(
@@ -283,7 +284,7 @@ def resume(
 
 @parameters.add_custom_parameters(deploy_mode=True)
 @click.command(help="Run the workflow locally.")
-@tracing.cli_entrypoint("cli/run")
+@tracing.cli("cli/run")
 @common_run_options
 @click.option(
     "--namespace",
@@ -345,18 +346,18 @@ def run(
             "msg": "Starting run",
         },
     )
+
+    runtime.print_workflow_info()
+    runtime.persist_constants()
+    if runner_attribute_file:
+        with open(runner_attribute_file, "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "run_id": runtime.run_id,
+                    "flow_name": obj.flow.name,
+                    "metadata": obj.metadata.metadata_str(),
+                },
+                f,
+            )
     with runtime.run_heartbeat():
-        runtime.print_workflow_info()
-        runtime.persist_constants()
-
-        if runner_attribute_file:
-            with open(runner_attribute_file, "w", encoding="utf-8") as f:
-                json.dump(
-                    {
-                        "run_id": runtime.run_id,
-                        "flow_name": obj.flow.name,
-                        "metadata": obj.metadata.metadata_str(),
-                    },
-                    f,
-                )
         runtime.execute()

metaflow/cli_components/step_cmd.py

@@ -7,9 +7,11 @@ from ..exception import CommandException
 from ..task import MetaflowTask
 from ..unbounded_foreach import UBF_CONTROL, UBF_TASK
 from ..util import decompress_list
+import metaflow.tracing as tracing
 
 
 @click.command(help="Internal command to execute a single task.", hidden=True)
+@tracing.cli("cli/step")
 @click.argument("step-name")
 @click.option(
     "--run-id",

metaflow/client/core.py

@@ -5,6 +5,7 @@ import os
 import tarfile
 from collections import namedtuple
 from datetime import datetime
+from tempfile import TemporaryDirectory
 from io import BytesIO
 from itertools import chain
 from typing import (
@@ -379,7 +380,7 @@ class MetaflowObject(object):
             _CLASSES[self._CHILD_CLASS]._NAME,
             query_filter,
             self._attempt,
-            *self.path_components
+            *self.path_components,
         )
         unfiltered_children = unfiltered_children if unfiltered_children else []
         children = filter(
@@ -878,6 +879,73 @@ class MetaflowCode(object):
         """
         return self._tar
 
+    def extract(self) -> TemporaryDirectory:
+        """
+        Extracts the code package to a temporary directory.
+
+        This creates a temporary directory containing all user code
+        files from the code package. The temporary directory is
+        automatically deleted when the returned TemporaryDirectory
+        object is garbage collected or when its cleanup() is called.
+
+        To preserve the contents to a permanent location, use
+        os.replace() which performs a zero-copy move on the same
+        filesystem:
+
+        ```python
+        with task.code.extract() as tmp_dir:
+            # Move contents to permanent location
+            for item in os.listdir(tmp_dir):
+                src = os.path.join(tmp_dir, item)
+                dst = os.path.join('/path/to/permanent/dir', item)
+                os.makedirs(os.path.dirname(dst), exist_ok=True)
+                os.replace(src, dst)  # Atomic move operation
+        ```
+        Returns
+        -------
+        TemporaryDirectory
+            A temporary directory containing the extracted code files.
+            The directory and its contents are automatically deleted when
+            this object is garbage collected.
+        """
+        exclusions = [
+            "metaflow/",
+            "metaflow_extensions/",
+            "INFO",
+            "CONFIG_PARAMETERS",
+            "conda.manifest",
+            # This file is created when using the conda/pypi features available in
+            # nflx-metaflow-extensions: https://github.com/Netflix/metaflow-nflx-extensions
+            "condav2-1.cnd",
+        ]
+        members = [
+            m
+            for m in self.tarball.getmembers()
+            if not any(
+                (x.endswith("/") and m.name.startswith(x)) or (m.name == x)
+                for x in exclusions
+            )
+        ]
+
+        tmp = TemporaryDirectory()
+        self.tarball.extractall(tmp.name, members)
+        return tmp
+
+    @property
+    def script_name(self) -> str:
+        """
+        Returns the filename of the Python script containing the FlowSpec.
+
+        This is the main Python file that was used to execute the flow. For example,
+        if your flow is defined in 'myflow.py', this property will return 'myflow.py'.
+
+        Returns
+        -------
+        str
+            Name of the Python file containing the FlowSpec
+        """
+        return self._info["script"]
+
     def __str__(self):
         return "<MetaflowCode: %s>" % self._info["script"]
 
@@ -1123,6 +1191,143 @@ class Task(MetaflowObject):
         # exclude private data artifacts
         return x.id[0] != "_"
 
+    def _iter_matching_tasks(self, steps, metadata_key, metadata_pattern):
+        """
+        Yield tasks from specified steps matching a foreach path pattern.
+
+        Parameters
+        ----------
+        steps : List[str]
+            List of step names to search for tasks
+        pattern : str
+            Regex pattern to match foreach-indices metadata
+
+        Returns
+        -------
+        Iterator[Task]
+            Tasks matching the foreach path pattern
+        """
+        flow_id, run_id, _, _ = self.path_components
+
+        for step in steps:
+            task_pathspecs = self._metaflow.metadata.filter_tasks_by_metadata(
+                flow_id, run_id, step.id, metadata_key, metadata_pattern
+            )
+            for task_pathspec in task_pathspecs:
+                yield Task(pathspec=task_pathspec, _namespace_check=False)
+
+    @property
+    def parent_tasks(self) -> Iterator["Task"]:
+        """
+        Yields all parent tasks of the current task if one exists.
+
+        Yields
+        ------
+        Task
+            Parent task of the current task
+
+        """
+        flow_id, run_id, _, _ = self.path_components
+
+        steps = list(self.parent.parent_steps)
+        if not steps:
+            return []
+
+        current_path = self.metadata_dict.get("foreach-execution-path", "")
+
+        if len(steps) > 1:
+            # Static join - use exact path matching
+            pattern = current_path or ".*"
+            yield from self._iter_matching_tasks(
+                steps, "foreach-execution-path", pattern
+            )
+            return
+
+        # Handle single step case
+        target_task = Step(
+            f"{flow_id}/{run_id}/{steps[0].id}", _namespace_check=False
+        ).task
+        target_path = target_task.metadata_dict.get("foreach-execution-path")
+
+        if not target_path or not current_path:
+            # (Current task, "A:10") and (Parent task, "")
+            # Pattern: ".*"
+            pattern = ".*"
+        else:
+            current_depth = len(current_path.split(","))
+            target_depth = len(target_path.split(","))
+
+            if current_depth < target_depth:
+                # Foreach join
+                # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13,C:21")
+                # Pattern: "A:10,B:13,.*"
+                pattern = f"{current_path},.*"
+            else:
+                # Foreach split or linear step
+                # Option 1:
+                # (Current task, "A:10,B:13,C:21") and (Parent task, "A:10,B:13")
+                # Option 2:
+                # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13")
+                # Pattern: "A:10,B:13"
+                pattern = ",".join(current_path.split(",")[:target_depth])
+
+        yield from self._iter_matching_tasks(steps, "foreach-execution-path", pattern)
+
+    @property
+    def child_tasks(self) -> Iterator["Task"]:
+        """
+        Yield all child tasks of the current task if one exists.
+
+        Yields
+        ------
+        Task
+            Child task of the current task
+        """
+        flow_id, run_id, _, _ = self.path_components
+        steps = list(self.parent.child_steps)
+        if not steps:
+            return []
+
+        current_path = self.metadata_dict.get("foreach-execution-path", "")
+
+        if len(steps) > 1:
+            # Static split - use exact path matching
+            pattern = current_path or ".*"
+            yield from self._iter_matching_tasks(
+                steps, "foreach-execution-path", pattern
+            )
+            return
+
+        # Handle single step case
+        target_task = Step(
+            f"{flow_id}/{run_id}/{steps[0].id}", _namespace_check=False
+        ).task
+        target_path = target_task.metadata_dict.get("foreach-execution-path")
+
+        if not target_path or not current_path:
+            # (Current task, "A:10") and (Child task, "")
+            # Pattern: ".*"
+            pattern = ".*"
+        else:
+            current_depth = len(current_path.split(","))
+            target_depth = len(target_path.split(","))
+
+            if current_depth < target_depth:
+                # Foreach split
+                # (Current task, "A:10,B:13") and (Child task, "A:10,B:13,C:21")
+                # Pattern: "A:10,B:13,.*"
+                pattern = f"{current_path},.*"
+            else:
+                # Foreach join or linear step
+                # Option 1:
+                # (Current task, "A:10,B:13,C:21") and (Child task, "A:10,B:13")
+                # Option 2:
+                # (Current task, "A:10,B:13") and (Child task, "A:10,B:13")
+                # Pattern: "A:10,B:13"
+                pattern = ",".join(current_path.split(",")[:target_depth])
+
+        yield from self._iter_matching_tasks(steps, "foreach-execution-path", pattern)
+
     @property
     def metadata(self) -> List[Metadata]:
         """
@@ -1837,6 +2042,41 @@ class Step(MetaflowObject):
         for t in self:
             return t.environment_info
 
+    @property
+    def parent_steps(self) -> Iterator["Step"]:
+        """
+        Yields parent steps for the current step.
+
+        Yields
+        ------
+        Step
+            Parent step
+        """
+        graph_info = self.task["_graph_info"].data
+
+        if self.id != "start":
+            flow, run, _ = self.path_components
+            for node_name, attributes in graph_info["steps"].items():
+                if self.id in attributes["next"]:
+                    yield Step(f"{flow}/{run}/{node_name}", _namespace_check=False)
+
+    @property
+    def child_steps(self) -> Iterator["Step"]:
+        """
+        Yields child steps for the current step.
+
+        Yields
+        ------
+        Step
+            Child step
+        """
+        graph_info = self.task["_graph_info"].data
+
+        if self.id != "end":
+            flow, run, _ = self.path_components
+            for next_step in graph_info["steps"][self.id]["next"]:
+                yield Step(f"{flow}/{run}/{next_step}", _namespace_check=False)
+
 
 class Run(MetaflowObject):
     """

metaflow/cmd/main_cli.py

@@ -12,7 +12,7 @@ import metaflow.tracing as tracing
 
 
 @click.group()
-@tracing.cli_entrypoint("cli/main")
+@tracing.cli("cli/main")
 def main():
     pass
 

metaflow/metadata_provider/heartbeat.py

@@ -52,6 +52,7 @@ class MetadataHeartBeat(object):
                 time.sleep(frequency_secs)
                 retry_counter = 0
             except HeartBeatException as e:
+                print(e)
                 retry_counter = retry_counter + 1
                 time.sleep(1.5**retry_counter)
 

metaflow/metadata_provider/metadata.py

@@ -5,6 +5,7 @@ import time
 from collections import namedtuple
 from itertools import chain
 
+from typing import List
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
 from metaflow.tagging_util import validate_tag
 from metaflow.util import get_username, resolve_identity_as_tuple, is_stringish
@@ -672,6 +673,38 @@ class MetadataProvider(object):
         if metadata:
             self.register_metadata(run_id, step_name, task_id, metadata)
 
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        field_name: str,
+        pattern: str,
+    ) -> List[str]:
+        """
+        Filter tasks by metadata field and pattern, returning task pathspecs that match criteria.
+
+        Parameters
+        ----------
+        flow_name : str
+            Flow name, that the run belongs to.
+        run_id: str
+            Run id, together with flow_id, that identifies the specific Run whose tasks to query
+        step_name: str
+            Step name to query tasks from
+        field_name: str
+            Metadata field name to query
+        pattern: str
+            Pattern to match in metadata field value
+
+        Returns
+        -------
+        List[str]
+            List of task pathspecs that satisfy the query
+        """
+        raise NotImplementedError()
+
     @staticmethod
     def _apply_filter(elts, filters):
         if filters is None:

metaflow/metaflow_config.py

@@ -248,8 +248,7 @@ DEFAULT_CONTAINER_IMAGE = from_conf("DEFAULT_CONTAINER_IMAGE")
 # Default container registry
 DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 # Controls whether to include foreach stack information in metadata.
-# TODO(Darin, 05/01/24): Remove this flag once we are confident with this feature.
-INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", False)
+INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", True)
 # Maximum length of the foreach value string to be stored in each ForeachFrame.
 MAXIMUM_FOREACH_VALUE_CHARS = from_conf("MAXIMUM_FOREACH_VALUE_CHARS", 30)
 # The default runtime limit (In seconds) of jobs launched by any compute provider. Default of 5 days.
@@ -433,9 +432,6 @@ CONDA_PACKAGE_GSROOT = from_conf("CONDA_PACKAGE_GSROOT")
 # should result in an appreciable speedup in flow environment initialization.
 CONDA_DEPENDENCY_RESOLVER = from_conf("CONDA_DEPENDENCY_RESOLVER", "conda")
 
-# Default to not using fast init binary.
-CONDA_USE_FAST_INIT = from_conf("CONDA_USE_FAST_INIT", False)
-
 ###
 # Escape hatch configuration
 ###

metaflow/mflog/save_logs.py

@@ -8,12 +8,12 @@ from metaflow.plugins import DATASTORES
 from metaflow.util import Path
 from . import TASK_LOG_SOURCE
 
-from metaflow.tracing import cli_entrypoint
+from metaflow.tracing import cli
 
 SMALL_FILE_LIMIT = 1024 * 1024
 
 
-@cli_entrypoint("save_logs")
+@cli("save_logs")
 def save_logs():
     def _read_file(path):
         with open(path, "rb") as f:

metaflow/plugins/argo/argo_workflows.py

@@ -77,13 +77,9 @@ class ArgoWorkflowsSchedulingException(MetaflowException):
 # List of future enhancements -
 #     1. Configure Argo metrics.
 #     2. Support resuming failed workflows within Argo Workflows.
-#     3. Support gang-scheduled clusters for distributed PyTorch/TF - One option is to
-#        use volcano - https://github.com/volcano-sh/volcano/tree/master/example/integrations/argo
-#     4. Support GitOps workflows.
-#     5. Add Metaflow tags to labels/annotations.
-#     6. Support Multi-cluster scheduling - https://github.com/argoproj/argo-workflows/issues/3523#issuecomment-792307297
-#     7. Support R lang.
-#     8. Ping @savin at slack.outerbounds.co for any feature request
+#     3. Add Metaflow tags to labels/annotations.
+#     4. Support R lang.
+#     5. Ping @savin at slack.outerbounds.co for any feature request
 
 
 class ArgoWorkflows(object):
@@ -742,6 +738,7 @@ class ArgoWorkflows(object):
                         {
                             "slack": bool(self.notify_slack_webhook_url),
                             "pager_duty": bool(self.notify_pager_duty_integration_key),
+                            "incident_io": bool(self.notify_incident_io_api_key),
                         }
                     )
                 }
@@ -753,6 +750,7 @@ class ArgoWorkflows(object):
                         {
                             "slack": bool(self.notify_slack_webhook_url),
                             "pager_duty": bool(self.notify_pager_duty_integration_key),
+                            "incident_io": bool(self.notify_incident_io_api_key),
                         }
                     )
                 }
@@ -2021,6 +2019,8 @@ class ArgoWorkflows(object):
                 kubernetes_labels = {
                     "task_id_entropy": "{{inputs.parameters.task-id-entropy}}",
                     "num_parallel": "{{inputs.parameters.num-parallel}}",
+                    "metaflow/argo-workflows-name": "{{workflow.name}}",
+                    "workflows.argoproj.io/workflow": "{{workflow.name}}",
                 }
                 jobset.labels(
                     {
@@ -2305,6 +2305,10 @@ class ArgoWorkflows(object):
             templates.append(self._slack_success_template())
             templates.append(self._pager_duty_change_template())
             templates.append(self._incident_io_change_template())
+
+        # Clean up None values from templates.
+        templates = list(filter(None, templates))
+
         if self.notify_on_error or self.notify_on_success:
             # Warning: terrible hack to workaround a bug in Argo Workflow where the
             #          templates listed above do not execute unless there is an
@@ -3635,7 +3639,7 @@ class Template(object):
     def resource(self, action, manifest, success_criteria, failure_criteria):
         self.payload["resource"] = {}
         self.payload["resource"]["action"] = action
-        self.payload["setOwnerReference"] = True
+        self.payload["resource"]["setOwnerReference"] = True
         self.payload["resource"]["successCondition"] = success_criteria
         self.payload["resource"]["failureCondition"] = failure_criteria
         self.payload["resource"]["manifest"] = manifest

metaflow/plugins/argo/argo_workflows_cli.py

@@ -168,12 +168,12 @@ def argo_workflows(obj, name=None):
 )
 @click.option(
     "--notify-slack-webhook-url",
-    default="",
+    default=None,
     help="Slack incoming webhook url for workflow success/failure notifications.",
 )
 @click.option(
     "--notify-pager-duty-integration-key",
-    default="",
+    default=None,
     help="PagerDuty Events API V2 Integration key for workflow success/failure notifications.",
 )
 @click.option(

metaflow/plugins/datatools/s3/s3op.py

@@ -155,7 +155,7 @@ def normalize_client_error(err):
 # S3 worker pool
 
 
-@tracing.cli_entrypoint("s3op/worker")
+@tracing.cli("s3op/worker")
 def worker(result_file_name, queue, mode, s3config):
     # Interpret mode, it can either be a single op or something like
     # info_download or info_upload which implies:
@@ -723,7 +723,7 @@ def cli():
 
 
 @cli.command("list", help="List S3 objects")
-@tracing.cli_entrypoint("s3op/list")
+@tracing.cli("s3op/list")
 @click.option(
     "--recursive/--no-recursive",
     default=False,
@@ -783,7 +783,7 @@ def lst(
 
 
 @cli.command(help="Upload files to S3")
-@tracing.cli_entrypoint("s3op/put")
+@tracing.cli("s3op/put")
 @click.option(
     "--file",
     "files",
@@ -978,7 +978,7 @@ def _populate_prefixes(prefixes, inputs):
 
 
 @cli.command(help="Download files from S3")
-@tracing.cli_entrypoint("s3op/get")
+@tracing.cli("s3op/get")
 @click.option(
     "--recursive/--no-recursive",
     default=False,

metaflow/plugins/env_escape/server.py

@@ -264,6 +264,7 @@ class Server(object):
     def serve(self, path=None, port=None):
         # Open up a connection
         if path is not None:
+            # Keep the print line to facilitate debugging
             # print("SERVER: Starting at %s" % path)
             sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
             __try_op__("bind", sock.bind, BIND_RETRY, path)
@@ -503,6 +504,12 @@ class Server(object):
         class_type = self._known_classes.get(class_name)
         if class_type is None:
             raise ValueError("Unknown class %s" % class_name)
+        # Check if __init__ is overridden
+        override_mapping = self._overrides.get(class_type)
+        if override_mapping:
+            override_func = override_mapping.get("__init__")
+            if override_func:
+                return override_func(None, class_type, *args, **kwargs)
         return class_type(*args, **kwargs)
 
     def _handle_subclasscheck(self, target, class_name, otherclass_name, reverse=False):

metaflow/plugins/env_escape/stub.py

@@ -276,9 +276,22 @@ class MetaWithConnection(StubMetaClass):
         if len(args) > 0 and id(args[0]) == id(cls.___class_connection___):
             return super(MetaWithConnection, cls).__call__(*args, **kwargs)
         else:
-            return cls.___class_connection___.stub_request(
-                None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
-            )
+            if hasattr(cls, "__overriden_init__"):
+                return cls.__overriden_init__(
+                    None,
+                    functools.partial(
+                        cls.___class_connection___.stub_request,
+                        None,
+                        OP_INIT,
+                        cls.___class_remote_class_name___,
+                    ),
+                    *args,
+                    **kwargs
+                )
+            else:
+                return cls.___class_connection___.stub_request(
+                    None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
+                )
 
     def __subclasscheck__(cls, subclass):
         subclass_name = "%s.%s" % (subclass.__module__, subclass.__name__)
@@ -381,7 +394,10 @@ def create_class(
             name = name[7:]
             method_type = CLASS_METHOD
         if name in overriden_methods:
-            if method_type == NORMAL_METHOD:
+            if name == "__init__":
+                class_dict["__overriden_init__"] = overriden_methods["__init__"]
+
+            elif method_type == NORMAL_METHOD:
                 class_dict[name] = (
                     lambda override, orig_method: lambda obj, *args, **kwargs: override(
                         obj, functools.partial(orig_method, obj), *args, **kwargs
@@ -412,6 +428,7 @@ def create_class(
             class_dict[name] = _make_method(
                 method_type, connection, class_name, name, doc
             )
+
     # Check for any getattr/setattr overrides
     special_attributes = set(getattr_overrides.keys())
     special_attributes.update(set(setattr_overrides.keys()))

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -39,7 +39,7 @@ def kubernetes():
     "command inside a Kubernetes pod with the given options. Typically you do not call "
     "this command directly; it is used internally by Metaflow."
 )
-@tracing.cli_entrypoint("kubernetes/step")
+@tracing.cli("kubernetes/step")
 @click.argument("step-name")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")