ob-metaflow 2.15.18.1__py2.py3-none-any.whl → 2.16.0.1__py2.py3-none-any.whl

metaflow/plugins/cards/card_decorator.py

@@ -2,8 +2,10 @@ import json
 import os
 import re
 import tempfile
+from typing import Tuple, Dict
 
 from metaflow.decorators import StepDecorator
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_current import current
 from metaflow.user_configs.config_options import ConfigInput
 from metaflow.user_configs.config_parameters import dump_config_values
@@ -22,6 +24,24 @@ def warning_message(message, logger=None, ts=False):
         logger(msg, timestamp=ts, bad=True)
 
 
+class MetadataStateManager(object):
+    def __init__(self, info_func):
+        self._info_func = info_func
+        self._metadata_registered = {}
+
+    def register_metadata(self, card_uuid) -> Tuple[bool, Dict]:
+        info = self._info_func()
+        # Check that metadata was not written yet. We only want to write once.
+        if (
+            info is None
+            or info.get(card_uuid) is None
+            or self._metadata_registered.get(card_uuid)
+        ):
+            return False, {}
+        self._metadata_registered[card_uuid] = True
+        return True, info.get(card_uuid)
+
+
 class CardDecorator(StepDecorator):
     """
     Creates a human-readable report, a Metaflow Card, after this step completes.
@@ -55,11 +75,14 @@ class CardDecorator(StepDecorator):
             The or one of the cards attached to this step.
     """
 
+    _GLOBAL_CARD_INFO = {}
+
     name = "card"
     defaults = {
         "type": "default",
         "options": {},
         "scope": "task",
+        "rank": None,  # Can be one of "high", "medium", "low". Can help derive ordering on the UI.
         "timeout": 45,
         "id": None,
         "save_errors": True,
@@ -91,6 +114,7 @@ class CardDecorator(StepDecorator):
         self._is_editable = False
         self._card_uuid = None
         self._user_set_card_id = None
+        self._metadata_registered = False
 
     @classmethod
     def _set_card_creator(cls, card_creator):
@@ -131,6 +155,16 @@ class CardDecorator(StepDecorator):
                 json.dump(config_value, config_file)
                 cls._config_file_name = config_file.name
 
+    @classmethod
+    def _register_card_info(cls, **kwargs):
+        if not kwargs.get("card_uuid"):
+            raise ValueError("card_uuid is required")
+        cls._GLOBAL_CARD_INFO[kwargs["card_uuid"]] = kwargs
+
+    @classmethod
+    def all_cards_info(cls):
+        return cls._GLOBAL_CARD_INFO.copy()
+
     def step_init(
         self, flow, graph, step_name, decorators, environment, flow_datastore, logger
     ):
@@ -191,6 +225,11 @@ class CardDecorator(StepDecorator):
         # 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)
+        # The MetadataStateManager is used to track the state of the metadata registration.
+        # It is there to ensure that we only register metadata for the card once. This is so that we
+        # avoid any un-necessary metadata writes because the create command can be called multiple times during the
+        # card creation process.
+        self._metadata_state_manager = MetadataStateManager(self.all_cards_info)
 
         card_type = self.attributes["type"]
         card_class = get_card_class(card_type)
@@ -225,7 +264,12 @@ 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(flow)))
+            self._set_card_creator(
+                CardCreator(
+                    self._create_top_level_args(flow),
+                    self._metadata_state_manager.register_metadata,
+                )
+            )
 
             current._update_env(
                 {"card": CardComponentCollector(self._logger, self.card_creator)}
@@ -248,6 +292,18 @@ class CardDecorator(StepDecorator):
         )
         self._card_uuid = card_metadata["uuid"]
 
+        self._register_card_info(
+            card_uuid=self._card_uuid,
+            rank=self.attributes["rank"],
+            type=self.attributes["type"],
+            options=self.card_options,
+            is_editable=self._is_editable,
+            is_runtime_card=self._is_runtime_card,
+            refresh_interval=self.attributes["refresh_interval"],
+            customize=customize,
+            id=self._user_set_card_id,
+        )
+
         # This means that we are calling `task_pre_step` on the last card decorator.
         # We can now `finalize` method in the CardComponentCollector object.
         # This will set up the `current.card` object for usage inside `@step` code.
@@ -305,11 +361,6 @@ class CardDecorator(StepDecorator):
 
         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]:

metaflow/plugins/cards/card_modules/convert_to_native_type.py

@@ -143,7 +143,10 @@ class TaskToDict:
         obj_type_name = self._get_object_type(data_object)
         if obj_type_name == "bytes":
             # Works for python 3.1+
-            import imghdr
+            # Python 3.13 removes the standard ``imghdr`` module. Metaflow
+            # vendors a copy so we can keep using ``what`` to detect image
+            # formats irrespective of the Python version.
+            from metaflow._vendor import imghdr
 
             resp = imghdr.what(None, h=data_object)
             # Only accept types supported on the web
@@ -157,7 +160,7 @@ class TaskToDict:
         obj_type_name = self._get_object_type(data_object)
         if obj_type_name == "bytes":
             # Works for python 3.1+
-            import imghdr
+            from metaflow._vendor import imghdr
 
             resp = imghdr.what(None, h=data_object)
             # Only accept types supported on the web

metaflow/plugins/cards/card_modules/test_cards.py

@@ -213,3 +213,19 @@ class TestRefreshComponentCard(MetaflowCard):
         if task.finished:
             return "final"
         return "runtime-%s" % _component_values_to_hash(data["components"])
+
+
+class TestImageCard(MetaflowCard):
+    """Card that renders a tiny PNG using ``TaskToDict.parse_image``."""
+
+    type = "test_image_card"
+
+    def render(self, task):
+        from .convert_to_native_type import TaskToDict
+        import base64
+
+        png_bytes = base64.b64decode(
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABRDE8UwAAAABJRU5ErkJggg=="
+        )
+        img_src = TaskToDict().parse_image(png_bytes)
+        return f"<html><img src='{img_src}' /></html>"

metaflow/plugins/cards/metadata.py

@@ -0,0 +1,22 @@
+import json
+from metaflow.metadata_provider import MetaDatum
+
+
+def _save_metadata(
+    metadata_provider,
+    run_id,
+    step_name,
+    task_id,
+    attempt_id,
+    card_uuid,
+    save_metadata,
+):
+    entries = [
+        MetaDatum(
+            field=card_uuid,
+            value=json.dumps(save_metadata),
+            type="card-info",
+            tags=["attempt_id:{0}".format(attempt_id)],
+        )
+    ]
+    metadata_provider.register_metadata(run_id, step_name, task_id, entries)

metaflow/plugins/exit_hook/exit_hook_decorator.py

@@ -0,0 +1,46 @@
+from metaflow.decorators import FlowDecorator
+from metaflow.exception import MetaflowException
+
+
+class ExitHookDecorator(FlowDecorator):
+    name = "exit_hook"
+    allow_multiple = True
+
+    defaults = {
+        "on_success": [],
+        "on_error": [],
+        "options": {},
+    }
+
+    def flow_init(
+        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
+    ):
+        on_success = self.attributes["on_success"]
+        on_error = self.attributes["on_error"]
+
+        if not on_success and not on_error:
+            raise MetaflowException(
+                "Choose at least one of the options on_success/on_error"
+            )
+
+        self.success_hooks = []
+        self.error_hooks = []
+        for success_fn in on_success:
+            if isinstance(success_fn, str):
+                self.success_hooks.append(success_fn)
+            elif callable(success_fn):
+                self.success_hooks.append(success_fn.__name__)
+            else:
+                raise ValueError(
+                    "Exit hooks inside 'on_success' must be a function or a string referring to the function"
+                )
+
+        for error_fn in on_error:
+            if isinstance(error_fn, str):
+                self.error_hooks.append(error_fn)
+            elif callable(error_fn):
+                self.error_hooks.append(error_fn.__name__)
+            else:
+                raise ValueError(
+                    "Exit hooks inside 'on_error' must be a function or a string referring to the function"
+                )

metaflow/plugins/exit_hook/exit_hook_script.py

@@ -0,0 +1,52 @@
+import os
+import inspect
+import importlib
+import sys
+
+
+def main(flow_file, fn_name_or_path, run_pathspec):
+    hook_fn = None
+
+    try:
+        module_path, function_name = fn_name_or_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        hook_fn = getattr(module, function_name)
+    except (ImportError, AttributeError, ValueError):
+        try:
+            module_name = os.path.splitext(os.path.basename(flow_file))[0]
+            spec = importlib.util.spec_from_file_location(module_name, flow_file)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            hook_fn = getattr(module, fn_name_or_path)
+        except (AttributeError, IOError) as e:
+            print(
+                f"[exit_hook] Could not load function '{fn_name_or_path}' "
+                f"as an import path or from '{flow_file}': {e}"
+            )
+            sys.exit(1)
+
+    argspec = inspect.getfullargspec(hook_fn)
+
+    # Check if fn expects a run object as an arg.
+    if "run" in argspec.args or argspec.varkw is not None:
+        from metaflow import Run
+
+        try:
+            _run = Run(run_pathspec, _namespace_check=False)
+        except Exception as ex:
+            print(ex)
+            _run = None
+
+        hook_fn(run=_run)
+    else:
+        hook_fn()
+
+
+if __name__ == "__main__":
+    try:
+        flow_file, fn_name, run_pathspec = sys.argv[1:4]
+    except Exception:
+        print("Usage: exit_hook_script.py <flow_file> <function_name> <run_pathspec>")
+        sys.exit(1)
+
+    main(flow_file, fn_name, run_pathspec)

metaflow/plugins/kubernetes/kubernetes.py

@@ -90,6 +90,7 @@ class Kubernetes(object):
         step_name,
         task_id,
         attempt,
+        code_package_metadata,
         code_package_url,
         step_cmds,
     ):
@@ -104,7 +105,7 @@ class Kubernetes(object):
             stderr_path=STDERR_PATH,
         )
         init_cmds = self._environment.get_package_commands(
-            code_package_url, self._datastore.TYPE
+            code_package_url, self._datastore.TYPE, code_package_metadata
         )
         init_expr = " && ".join(init_cmds)
         step_expr = bash_capture_logs(
@@ -165,6 +166,7 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -232,6 +234,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)
@@ -431,6 +434,7 @@ class Kubernetes(object):
             step_name=step_name,
             task_id=_tskid,
             attempt=attempt,
+            code_package_metadata=code_package_metadata,
             code_package_url=code_package_url,
             step_cmds=[
                 step_cli.replace(
@@ -479,6 +483,7 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -527,6 +532,7 @@ class Kubernetes(object):
                     step_name=step_name,
                     task_id=task_id,
                     attempt=attempt,
+                    code_package_metadata=code_package_metadata,
                     code_package_url=code_package_url,
                     step_cmds=[step_cli],
                 ),
@@ -555,6 +561,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -41,6 +41,7 @@ def kubernetes():
 )
 @tracing.cli("kubernetes/step")
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option(
@@ -161,6 +162,7 @@ def kubernetes():
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
@@ -304,6 +306,7 @@ def step(
                 task_id=task_id,
                 attempt=str(retry_count),
                 user=util.get_username(),
+                code_package_metadata=code_package_metadata,
                 code_package_sha=code_package_sha,
                 code_package_url=code_package_url,
                 code_package_ds=ctx.obj.flow_datastore.TYPE,

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -11,6 +11,7 @@ from metaflow.metadata_provider import MetaDatum
 from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     KUBERNETES_CONTAINER_IMAGE,
     KUBERNETES_CONTAINER_REGISTRY,
     KUBERNETES_CPU,
@@ -170,6 +171,7 @@ class KubernetesDecorator(StepDecorator):
         "qos": KUBERNETES_QOS,
         "security_context": None,
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
@@ -179,8 +181,6 @@ class KubernetesDecorator(StepDecorator):
     target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
-        super(KubernetesDecorator, self).init()
-
         if not self.attributes["namespace"]:
             self.attributes["namespace"] = KUBERNETES_NAMESPACE
         if not self.attributes["service_account"]:
@@ -486,6 +486,7 @@ class KubernetesDecorator(StepDecorator):
             # to execute on Kubernetes anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["kubernetes", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
 
@@ -657,9 +658,16 @@ class KubernetesDecorator(StepDecorator):
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
 
 
 # TODO: Unify this method with the multi-node setup in @batch

metaflow/plugins/package_cli.py

@@ -9,35 +9,30 @@ def cli():
 
 
 @cli.group(help="Commands related to code packages.")
+@click.option(
+    "--timeout", default=60, help="Timeout for package operations in seconds."
+)
 @click.pass_obj
-def package(obj):
+def package(obj, timeout):
     # Prepare the package before any of the sub-commands are invoked.
+    # We explicitly will *not* upload it to the datastore.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=None,
     )
+    obj.package_op_timeout = timeout
 
 
-@package.command(help="Output information about the current code package.")
+@package.command(help="Output information about the code package.")
 @click.pass_obj
 def info(obj):
-    obj.echo("Status of the current working directory:", fg="magenta", bold=False)
-    obj.echo_always(
-        "Hash: *%s*" % sha1(obj.package.blob).hexdigest(),
-        highlight="green",
-        highlight_bold=False,
-    )
-    obj.echo_always(
-        "Package size: *%d* KB" % (len(obj.package.blob) / 1024),
-        highlight="green",
-        highlight_bold=False,
-    )
-    num = sum(1 for _ in obj.package.path_tuples())
-    obj.echo_always(
-        "Number of files: *%d*" % num, highlight="green", highlight_bold=False
-    )
+    obj.echo_always(obj.package.show())
 
 
-@package.command(help="List files included in the code package.")
+@package.command(help="List all files included in the code package.")
 @click.option(
     "--archive/--no-archive",
     default=False,
@@ -47,8 +42,10 @@ def info(obj):
 )
 @click.pass_obj
 def list(obj, archive=False):
+    _ = obj.package.blob_with_timeout(timeout=obj.package_op_timeout)
+    # We now have all the information about the blob
     obj.echo(
-        "Files included in the code package " "(change with --package-suffixes):",
+        "Files included in the code package (change with --package-suffixes):",
         fg="magenta",
         bold=False,
     )
@@ -58,10 +55,15 @@ def list(obj, archive=False):
         obj.echo_always("\n".join(path for path, _ in obj.package.path_tuples()))
 
 
-@package.command(help="Save the current code package in a tar file")
+@package.command(help="Save the current code package to a file.")
 @click.argument("path")
 @click.pass_obj
 def save(obj, path):
     with open(path, "wb") as f:
-        f.write(obj.package.blob)
-    obj.echo("Code package saved in *%s*." % path, fg="magenta", bold=False)
+        f.write(obj.package.blob())
+    obj.echo(
+        "Code package saved in *%s* with metadata: %s"
+        % (path, obj.package.package_metadata),
+        fg="magenta",
+        bold=False,
+    )

metaflow/plugins/parallel_decorator.py

@@ -36,8 +36,10 @@ class ParallelDecorator(StepDecorator):
     defaults = {}
     IS_PARALLEL = True
 
-    def __init__(self, attributes=None, statically_defined=False):
-        super(ParallelDecorator, self).__init__(attributes, statically_defined)
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
+        super(ParallelDecorator, self).__init__(
+            attributes, statically_defined, inserted_by
+        )
 
     def runtime_step_cli(
         self, cli_args, retry_count, max_user_code_retries, ubf_context

metaflow/plugins/pypi/bootstrap.py

@@ -12,9 +12,10 @@ import platform
 from urllib.error import URLError
 from urllib.request import urlopen
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, CONDA_USE_FAST_INIT
+from metaflow.packaging_sys import MetaflowCodeContent, ContentType
 from metaflow.plugins import DATASTORES
 from metaflow.plugins.pypi.utils import MICROMAMBA_MIRROR_URL, MICROMAMBA_URL
-from metaflow.util import which, get_metaflow_root
+from metaflow.util import which
 from urllib.request import Request
 import warnings
 
@@ -365,8 +366,13 @@ if __name__ == "__main__":
 
         # Move MAGIC_FILE inside local datastore.
         os.makedirs(manifest_dir, exist_ok=True)
+        path_to_manifest = MetaflowCodeContent.get_filename(
+            MAGIC_FILE, ContentType.OTHER_CONTENT
+        )
+        if path_to_manifest is None:
+            raise RuntimeError(f"Cannot find {MAGIC_FILE} in the package")
         shutil.move(
-            os.path.join(get_metaflow_root(), MAGIC_FILE),
+            path_to_manifest,
             os.path.join(manifest_dir, MAGIC_FILE),
         )
         with open(os.path.join(manifest_dir, MAGIC_FILE)) as f: