metaflow 2.15.18__py2.py3-none-any.whl → 2.15.19__py2.py3-none-any.whl

metaflow/plugins/argo/exit_hooks.py

@@ -0,0 +1,209 @@
+from collections import defaultdict
+import json
+from typing import Dict, List
+
+
+class JsonSerializable(object):
+    def to_json(self):
+        return self.payload
+
+    def __str__(self):
+        return json.dumps(self.payload, indent=4)
+
+
+class _LifecycleHook(JsonSerializable):
+    # https://argoproj.github.io/argo-workflows/fields/#lifecyclehook
+
+    def __init__(self, name):
+        tree = lambda: defaultdict(tree)
+        self.name = name
+        self.payload = tree()
+
+    def expression(self, expression):
+        self.payload["expression"] = str(expression)
+        return self
+
+    def template(self, template):
+        self.payload["template"] = template
+        return self
+
+
+class _Template(JsonSerializable):
+    # https://argoproj.github.io/argo-workflows/fields/#template
+
+    def __init__(self, name):
+        tree = lambda: defaultdict(tree)
+        self.name = name
+        self.payload = tree()
+        self.payload["name"] = name
+
+    def http(self, http):
+        self.payload["http"] = http.to_json()
+        return self
+
+    def script(self, script):
+        self.payload["script"] = script.to_json()
+        return self
+
+    def container(self, container):
+        self.payload["container"] = container
+        return self
+
+    def service_account_name(self, service_account_name):
+        self.payload["serviceAccountName"] = service_account_name
+        return self
+
+
+class Hook(object):
+    """
+    Abstraction for Argo Workflows exit hooks.
+    A hook consists of a Template, and one or more LifecycleHooks that trigger the template
+    """
+
+    template: "_Template"
+    lifecycle_hooks: List["_LifecycleHook"]
+
+
+class _HttpSpec(JsonSerializable):
+    # https://argoproj.github.io/argo-workflows/fields/#http
+
+    def __init__(self, method):
+        tree = lambda: defaultdict(tree)
+        self.payload = tree()
+        self.payload["method"] = method
+        self.payload["headers"] = []
+
+    def header(self, header, value):
+        self.payload["headers"].append({"name": header, "value": value})
+        return self
+
+    def body(self, body):
+        self.payload["body"] = str(body)
+        return self
+
+    def url(self, url):
+        self.payload["url"] = url
+        return self
+
+    def success_condition(self, success_condition):
+        self.payload["successCondition"] = success_condition
+        return self
+
+
+# HTTP hook
+class HttpExitHook(Hook):
+    def __init__(
+        self,
+        name,
+        url,
+        method="GET",
+        headers=None,
+        body=None,
+        on_success=False,
+        on_error=False,
+    ):
+        self.template = _Template(name)
+        http = _HttpSpec(method).url(url)
+        if headers is not None:
+            for header, value in headers.items():
+                http.header(header, value)
+
+        if body is not None:
+            http.body(json.dumps(body))
+
+        self.template.http(http)
+
+        self.lifecycle_hooks = []
+
+        if on_success and on_error:
+            raise Exception("Set only one of the on_success/on_error at a time.")
+
+        if on_success:
+            self.lifecycle_hooks.append(
+                _LifecycleHook(name)
+                .expression("workflow.status == 'Succeeded'")
+                .template(self.template.name)
+            )
+
+        if on_error:
+            self.lifecycle_hooks.append(
+                _LifecycleHook(name)
+                .expression("workflow.status == 'Error' || workflow.status == 'Failed'")
+                .template(self.template.name)
+            )
+
+        if not on_success and not on_error:
+            # add an expressionless lifecycle hook
+            self.lifecycle_hooks.append(_LifecycleHook(name).template(name))
+
+
+class ExitHookHack(Hook):
+    # Warning: terrible hack to workaround a bug in Argo Workflow where the
+    #          templates listed above do not execute unless there is an
+    #          explicit exit hook. as and when this bug is patched, we should
+    #          remove this effectively no-op template.
+    # Note: We use the Http template because changing this to an actual no-op container had the side-effect of
+    # leaving LifecycleHooks in a pending state even when they have finished execution.
+    def __init__(
+        self,
+        url,
+        headers=None,
+        body=None,
+    ):
+        self.template = _Template("exit-hook-hack")
+        http = _HttpSpec("GET").url(url)
+        if headers is not None:
+            for header, value in headers.items():
+                http.header(header, value)
+
+        if body is not None:
+            http.body(json.dumps(body))
+
+        http.success_condition("true == true")
+
+        self.template.http(http)
+
+        self.lifecycle_hooks = []
+
+        # add an expressionless lifecycle hook
+        self.lifecycle_hooks.append(_LifecycleHook("exit").template("exit-hook-hack"))
+
+
+class ContainerHook(Hook):
+    def __init__(
+        self,
+        name: str,
+        container: Dict,
+        service_account_name: str = None,
+        on_success: bool = False,
+        on_error: bool = False,
+    ):
+        self.template = _Template(name)
+
+        if service_account_name is not None:
+            self.template.service_account_name(service_account_name)
+
+        self.template.container(container)
+
+        self.lifecycle_hooks = []
+
+        if on_success and on_error:
+            raise Exception("Set only one of the on_success/on_error at a time.")
+
+        if on_success:
+            self.lifecycle_hooks.append(
+                _LifecycleHook(name)
+                .expression("workflow.status == 'Succeeded'")
+                .template(self.template.name)
+            )
+
+        if on_error:
+            self.lifecycle_hooks.append(
+                _LifecycleHook(name)
+                .expression("workflow.status == 'Error' || workflow.status == 'Failed'")
+                .template(self.template.name)
+            )
+
+        if not on_success and not on_error:
+            # add an expressionless lifecycle hook
+            self.lifecycle_hooks.append(_LifecycleHook(name).template(name))

metaflow/plugins/aws/aws_utils.py

@@ -48,7 +48,7 @@ def get_ec2_instance_metadata():
         # Try to get an IMDSv2 token.
         token = requests.put(
             url="http://169.254.169.254/latest/api/token",
-            headers={"X-aws-ec2-metadata-token-ttl-seconds": 100},
+            headers={"X-aws-ec2-metadata-token-ttl-seconds": "100"},
             timeout=timeout,
         ).text
     except:

metaflow/plugins/aws/step_functions/step_functions.py

@@ -301,6 +301,12 @@ class StepFunctions(object):
                 "to AWS Step Functions is not supported currently."
             )
 
+        if self.flow._flow_decorators.get("exit_hook"):
+            raise StepFunctionsException(
+                "Deploying flows with the @exit_hook decorator "
+                "to AWS Step Functions is not currently supported."
+            )
+
         # Visit every node of the flow and recursively build the state machine.
         def _visit(node, workflow, exit_node=None):
             if node.parallel_foreach:

metaflow/plugins/cards/card_cli.py

@@ -30,7 +30,7 @@ from .exception import (
 )
 import traceback
 from collections import namedtuple
-
+from .metadata import _save_metadata
 from .card_resolver import resolve_paths_from_task, resumed_info
 
 id_func = id
@@ -613,6 +613,14 @@ def update_card(mf_card, mode, task, data, timeout_value=None):
     hidden=True,
     help="Delete data-file and component-file after reading. (internal)",
 )
+@click.option(
+    "--save-metadata",
+    default=None,
+    show_default=True,
+    type=JSONTypeClass(),
+    hidden=True,
+    help="JSON string containing metadata to be saved. (internal)",
+)
 @click.pass_context
 def create(
     ctx,
@@ -627,6 +635,7 @@ def create(
     card_uuid=None,
     delete_input_files=None,
     id=None,
+    save_metadata=None,
 ):
     card_id = id
     rendered_info = None  # Variable holding all the information which will be rendered
@@ -824,6 +833,16 @@ def create(
                 % (card_info.type, card_info.hash[:NUM_SHORT_HASH_CHARS]),
                 fg="green",
             )
+            if save_metadata:
+                _save_metadata(
+                    ctx.obj.metadata,
+                    task.parent.parent.id,
+                    task.parent.id,
+                    task.id,
+                    task.current_attempt,
+                    card_uuid,
+                    save_metadata,
+                )
 
 
 @card.command()

metaflow/plugins/cards/card_creator.py

@@ -5,6 +5,8 @@ import json
 import sys
 import os
 from metaflow import current
+from typing import Callable, Tuple, Dict
+
 
 ASYNC_TIMEOUT = 30
 
@@ -44,8 +46,18 @@ class CardProcessManager:
 
 
 class CardCreator:
-    def __init__(self, top_level_options):
+    def __init__(
+        self,
+        top_level_options,
+        should_save_metadata_lambda: Callable[[str], Tuple[bool, Dict]],
+    ):
+        # should_save_metadata_lambda is a lambda that provides a flag to indicate if
+        # card metadata should be written to the metadata store.
+        # It gets called only once when the card is created inside the subprocess.
+        # The intent is that this is a stateful lambda that will ensure that we only end
+        # up writing to the metadata store once.
         self._top_level_options = top_level_options
+        self._should_save_metadata = should_save_metadata_lambda
 
     def create(
         self,
@@ -62,6 +74,8 @@ class CardCreator:
         # Setting `final` will affect the Reload token set during the card refresh
         # data creation along with synchronous execution of subprocess.
         # Setting `sync` will only cause synchronous execution of subprocess.
+        save_metadata = False
+        metadata_dict = {}
         if mode != "render" and not runtime_card:
             # silently ignore runtime updates for cards that don't support them
             return
@@ -71,6 +85,8 @@ class CardCreator:
             component_strings = []
         else:
             component_strings = current.card._serialize_components(card_uuid)
+            # Since the mode is a render, we can check if we need to write to the metadata store.
+            save_metadata, metadata_dict = self._should_save_metadata(card_uuid)
         data = current.card._get_latest_data(card_uuid, final=final, mode=mode)
         runspec = "/".join([current.run_id, current.step_name, current.task_id])
         self._run_cards_subprocess(
@@ -85,6 +101,8 @@ class CardCreator:
             data,
             final=final,
             sync=sync,
+            save_metadata=save_metadata,
+            metadata_dict=metadata_dict,
         )
 
     def _run_cards_subprocess(
@@ -100,6 +118,8 @@ class CardCreator:
         data=None,
         final=False,
         sync=False,
+        save_metadata=False,
+        metadata_dict=None,
     ):
         components_file = data_file = None
         wait = final or sync
@@ -156,6 +176,9 @@ class CardCreator:
         if data_file is not None:
             cmd += ["--data-file", data_file.name]
 
+        if save_metadata:
+            cmd += ["--save-metadata", json.dumps(metadata_dict)]
+
         response, fail = self._run_command(
             cmd,
             card_uuid,

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.

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/secrets/__init__.py

@@ -9,3 +9,6 @@ class SecretsProvider(abc.ABC):
     def get_secret_as_dict(self, secret_id, options={}, role=None) -> Dict[str, str]:
         """Retrieve the secret from secrets backend, and return a dictionary of
         environment variables."""
+
+
+from .secrets_func import get_secrets