ob-metaflow 2.14.0.2rc0__tar.gz → 2.14.3.1__tar.gz

{ob_metaflow-2.14.0.2rc0/ob_metaflow.egg-info → ob_metaflow-2.14.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ob-metaflow
-Version: 2.14.0.2rc0
+Version: 2.14.3.1
 Summary: Metaflow: More Data Science, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co
@@ -12,7 +12,7 @@ Requires-Dist: boto3
 Requires-Dist: pylint
 Requires-Dist: kubernetes
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.14.0.2rc0; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.14.3.1; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/metaflow/cli.py

@@ -16,7 +16,6 @@ from .exception import CommandException, MetaflowException
 from .flowspec import _FlowState
 from .graph import FlowGraph
 from .metaflow_config import (
-    DECOSPECS,
     DEFAULT_DATASTORE,
     DEFAULT_ENVIRONMENT,
     DEFAULT_EVENT_LOGGER,
@@ -106,26 +105,6 @@ def logger(body="", system_msg=False, head="", bad=False, timestamp=True, nl=Tru
     click.secho(body, bold=system_msg, fg=LOGGER_BAD_COLOR if bad else None, nl=nl)
 
 
-def config_merge_cb(ctx, param, value):
-    # Callback to:
-    #  - read  the Click auto_envvar variable from both the
-    #    environment AND the configuration
-    #  - merge that value with the value passed in the command line (value)
-    #  - return the value as a tuple
-    # Note that this function gets called even if there is no option passed on the
-    # command line.
-    # NOTE: Assumes that ctx.auto_envvar_prefix is set to METAFLOW (same as in
-    # from_conf)
-
-    # Special case where DECOSPECS and value are the same. This happens
-    # when there is no --with option at the TL and DECOSPECS is read from
-    # the env var. In this case, click also passes it as value
-    splits = DECOSPECS.split()
-    if len(splits) == len(value) and all([a == b for (a, b) in zip(splits, value)]):
-        return value
-    return tuple(list(value) + DECOSPECS.split())
-
-
 @click.group(
     cls=LazyGroup,
     lazy_subcommands={
@@ -238,7 +217,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.
@@ -285,7 +263,6 @@ def version(obj):
     multiple=True,
     help="Add a decorator to all steps. You can specify this option "
     "multiple times to attach multiple decorators in steps.",
-    callback=config_merge_cb,
 )
 @click.option(
     "--pylint/--no-pylint",

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/metaflow/cli_components/run_cmds.py

@@ -8,6 +8,7 @@ from .. import decorators, namespace, parameters, tracing
 from ..exception import CommandException
 from ..graph import FlowGraph
 from ..metaflow_current import current
+from ..metaflow_config import DEFAULT_DECOSPECS
 from ..package import MetaflowPackage
 from ..runtime import NativeRuntime
 from ..system import _system_logger
@@ -70,6 +71,23 @@ def write_file(file_path, content):
             f.write(str(content))
 
 
+def config_merge_cb(ctx, param, value):
+    # Callback to:
+    #  - read  the Click auto_envvar variable from both the
+    #    environment AND the configuration
+    #  - merge that value with the value passed in the command line (value)
+    #  - return the value as a tuple
+    # Note that this function gets called even if there is no option passed on the
+    # command line.
+    # NOTE: Assumes that ctx.auto_envvar_prefix is set to METAFLOW (same as in
+    # from_conf)
+
+    # Read decospecs options from the environment (METAFLOW_DEFAULT_DECOSPECS=...)
+    # and merge them with the one provided as --with.
+    splits = DEFAULT_DECOSPECS.split()
+    return tuple(list(value) + splits)
+
+
 def common_run_options(func):
     @click.option(
         "--tag",
@@ -109,6 +127,7 @@ def common_run_options(func):
         help="Add a decorator to all steps. You can specify this "
         "option multiple times to attach multiple decorators "
         "in steps.",
+        callback=config_merge_cb,
     )
     @click.option(
         "--run-id-file",
@@ -166,6 +185,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 +303,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 +365,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()

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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",

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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):
     """

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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
 

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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)
 

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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:

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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.
@@ -286,7 +285,7 @@ CONTACT_INFO = from_conf(
 ###
 # Format is a space separated string of decospecs (what is passed
 # using --with)
-DECOSPECS = from_conf("DECOSPECS", "")
+DEFAULT_DECOSPECS = from_conf("DEFAULT_DECOSPECS", "")
 
 ###
 # AWS Batch configuration
@@ -584,9 +583,9 @@ try:
                 _TOGGLE_DECOSPECS.extend(o)
             elif not n.startswith("__") and not isinstance(o, types.ModuleType):
                 globals()[n] = o
-    # If DECOSPECS is set, use that, else extrapolate from extensions
-    if not DECOSPECS:
-        DECOSPECS = " ".join(_TOGGLE_DECOSPECS)
+    # If DEFAULT_DECOSPECS is set, use that, else extrapolate from extensions
+    if not DEFAULT_DECOSPECS:
+        DEFAULT_DECOSPECS = " ".join(_TOGGLE_DECOSPECS)
 
 finally:
     # Erase all temporary names to avoid leaking things

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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:

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/metaflow/plugins/argo/argo_workflows.py

@@ -78,13 +78,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):
@@ -743,6 +739,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),
                         }
                     )
                 }
@@ -754,6 +751,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),
                         }
                     )
                 }
@@ -1966,11 +1964,6 @@ class ArgoWorkflows(object):
                 resources["disk"],
             )
 
-            extended_resources = resources.get("extended_resources", {})
-
-            qos_requests = {**qos_requests, **extended_resources}
-            qos_limits = {**qos_limits, **extended_resources}
-
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -2027,7 +2020,6 @@ class ArgoWorkflows(object):
                     shared_memory=shared_memory,
                     port=port,
                     qos=resources["qos"],
-                    extended_resources=extended_resources,
                 )
 
                 for k, v in env.items():
@@ -2043,6 +2035,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(
                     {
@@ -2334,6 +2328,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
@@ -3682,7 +3680,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

{ob_metaflow-2.14.0.2rc0 → ob_metaflow-2.14.3.1}/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(