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

metaflow/plugins/argo/argo_workflows_cli.py

@@ -15,6 +15,7 @@ from metaflow.exception import (
 )
 from metaflow.metaflow_config import (
     ARGO_WORKFLOWS_UI_URL,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     KUBERNETES_NAMESPACE,
     SERVICE_VERSION_CHECK,
     UI_URL,
@@ -518,16 +519,27 @@ def make_flow(
     # Save the code package in the flow datastore so that both user code and
     # metaflow package can be retrieved during workflow execution.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
-    package_url, package_sha = obj.flow_datastore.save_data(
-        [obj.package.blob], len_hint=1
-    )[0]
+
+    # This blocks until the package is created
+    if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+        package_url = obj.package.package_url()
+        package_sha = obj.package.package_sha()
+    else:
+        package_url, package_sha = obj.flow_datastore.save_data(
+            [obj.package.blob], len_hint=1
+        )[0]
 
     return ArgoWorkflows(
         name,
         obj.graph,
         obj.flow,
+        obj.package.package_metadata,
         package_sha,
         package_url,
         token,

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

@@ -49,7 +49,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/batch/batch.py

@@ -59,14 +59,24 @@ class Batch(object):
         self._client = BatchClient()
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
-    def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
+    def _command(
+        self,
+        environment,
+        code_package_metadata,
+        code_package_url,
+        step_name,
+        step_cmds,
+        task_spec,
+    ):
         mflog_expr = export_mflog_env_vars(
             datastore_type="s3",
             stdout_path=STDOUT_PATH,
             stderr_path=STDERR_PATH,
             **task_spec
         )
-        init_cmds = environment.get_package_commands(code_package_url, "s3")
+        init_cmds = environment.get_package_commands(
+            code_package_url, "s3", code_package_metadata
+        )
         init_expr = " && ".join(init_cmds)
         step_expr = bash_capture_logs(
             " && ".join(environment.bootstrap_commands(step_name, "s3") + step_cmds)
@@ -167,6 +177,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -210,7 +221,12 @@ class Batch(object):
             .job_queue(queue)
             .command(
                 self._command(
-                    self.environment, code_package_url, step_name, [step_cli], task_spec
+                    self.environment,
+                    code_package_metadata,
+                    code_package_url,
+                    step_name,
+                    [step_cli],
+                    task_spec,
                 )
             )
             .image(image)
@@ -249,6 +265,7 @@ class Batch(object):
             )
             .task_id(attrs.get("metaflow.task_id"))
             .environment_variable("AWS_DEFAULT_REGION", self._client.region())
+            .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)
@@ -334,6 +351,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -374,6 +392,7 @@ class Batch(object):
             step_name,
             step_cli,
             task_spec,
+            code_package_metadata,
             code_package_sha,
             code_package_url,
             code_package_ds,

metaflow/plugins/aws/batch/batch_cli.py

@@ -100,6 +100,7 @@ def kill(ctx, run_id, user, my_runs):
     "Metaflow."
 )
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option("--executable", help="Executable requirement for AWS Batch.")
@@ -185,6 +186,7 @@ def kill(ctx, run_id, user, my_runs):
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
@@ -317,6 +319,7 @@ def step(
                 step_name,
                 step_cli,
                 task_spec,
+                code_package_metadata,
                 code_package_sha,
                 code_package_url,
                 ctx.obj.flow_datastore.TYPE,

metaflow/plugins/aws/batch/batch_decorator.py

@@ -14,6 +14,7 @@ from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
     ECS_FARGATE_EXECUTION_ROLE,
     ECS_S3_ACCESS_IAM_ROLE,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
 )
 from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
 from metaflow.sidecar import Sidecar
@@ -126,6 +127,7 @@ class BatchDecorator(StepDecorator):
         "gpu": "0",
         "memory": "4096",
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
@@ -135,8 +137,6 @@ class BatchDecorator(StepDecorator):
     target_platform = "linux-64"
 
     def init(self):
-        super(BatchDecorator, self).init()
-
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
             # If metaflow-config specifies a docker image, just use that.
@@ -228,6 +228,7 @@ class BatchDecorator(StepDecorator):
             # to execute on AWS Batch anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["batch", "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)
             cli_args.command_options.update(self.attributes)
@@ -403,9 +404,16 @@ class BatchDecorator(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
 
 
 def _setup_multinode_environment():

metaflow/plugins/aws/step_functions/step_functions.py

@@ -40,6 +40,7 @@ class StepFunctions(object):
         name,
         graph,
         flow,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         production_token,
@@ -59,6 +60,7 @@ class StepFunctions(object):
         self.name = name
         self.graph = graph
         self.flow = flow
+        self.code_package_metadata = code_package_metadata
         self.code_package_sha = code_package_sha
         self.code_package_url = code_package_url
         self.production_token = production_token
@@ -301,6 +303,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:
@@ -847,6 +855,7 @@ class StepFunctions(object):
                     node, input_paths, self.code_package_url, user_code_retries
                 ),
                 task_spec=task_spec,
+                code_package_metadata=self.code_package_metadata,
                 code_package_sha=self.code_package_sha,
                 code_package_url=self.code_package_url,
                 code_package_ds=self.flow_datastore.TYPE,
@@ -907,7 +916,7 @@ class StepFunctions(object):
             "with": [
                 decorator.make_decorator_spec()
                 for decorator in node.decorators
-                if not decorator.statically_defined
+                if not decorator.statically_defined and decorator.inserted_by is None
             ]
         }
         # FlowDecorators can define their own top-level options. They are

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -7,6 +7,7 @@ from metaflow import JSONType, current, decorators, parameters
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     SERVICE_VERSION_CHECK,
     SFN_STATE_MACHINE_PREFIX,
     UI_URL,
@@ -331,16 +332,26 @@ def make_flow(
     )
 
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
-    package_url, package_sha = obj.flow_datastore.save_data(
-        [obj.package.blob], len_hint=1
-    )[0]
+    # This blocks until the package is created
+    if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+        package_url = obj.package.package_url()
+        package_sha = obj.package.package_sha()
+    else:
+        package_url, package_sha = obj.flow_datastore.save_data(
+            [obj.package.blob], len_hint=1
+        )[0]
 
     return StepFunctions(
         name,
         obj.graph,
         obj.flow,
+        obj.package.package_metadata,
         package_sha,
         package_url,
         token,

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,