ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/argo/argo_workflows_decorator.py

@@ -123,6 +123,15 @@ class ArgoWorkflowsInternalDecorator(StepDecorator):
             with open("/mnt/out/split_cardinality", "w") as file:
                 json.dump(flow._foreach_num_splits, file)
 
+        # For conditional branches we need to record the value of the switch to disk, in order to pass it as an
+        # output from the switching step to be used further down the DAG
+        if graph[step_name].type == "split-switch":
+            # TODO: A nicer way to access the chosen step?
+            _out_funcs, _ = flow._transition
+            chosen_step = _out_funcs[0]
+            with open("/mnt/out/switch_step", "w") as file:
+                file.write(chosen_step)
+
         # For steps that have a `@parallel` decorator set to them, we will be relying on Jobsets
         # to run the task. In this case, we cannot set anything in the
         # `/mnt/out` directory, since such form of output mounts are not available to Jobset executions.

metaflow/plugins/argo/argo_workflows_deployer_objects.py

@@ -9,59 +9,16 @@ from metaflow.exception import MetaflowException
 from metaflow.plugins.argo.argo_client import ArgoClient
 from metaflow.metaflow_config import KUBERNETES_NAMESPACE
 from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
-from metaflow.runner.deployer import Deployer, DeployedFlow, TriggeredRun
+from metaflow.runner.deployer import (
+    Deployer,
+    DeployedFlow,
+    TriggeredRun,
+    generate_fake_flow_file_contents,
+)
 
 from metaflow.runner.utils import get_lower_level_group, handle_timeout, temporary_fifo
 
 
-def generate_fake_flow_file_contents(
-    flow_name: str, param_info: dict, project_name: Optional[str] = None
-):
-    params_code = ""
-    for _, param_details in param_info.items():
-        param_python_var_name = param_details["python_var_name"]
-        param_name = param_details["name"]
-        param_type = param_details["type"]
-        param_help = param_details["description"]
-        param_required = param_details["is_required"]
-
-        if param_type == "JSON":
-            params_code += (
-                f"    {param_python_var_name} = Parameter('{param_name}', "
-                f"type=JSONType, help='''{param_help}''', required={param_required})\n"
-            )
-        elif param_type == "FilePath":
-            is_text = param_details.get("is_text", True)
-            encoding = param_details.get("encoding", "utf-8")
-            params_code += (
-                f"    {param_python_var_name} = IncludeFile('{param_name}', "
-                f"is_text={is_text}, encoding='{encoding}', help='''{param_help}''', "
-                f"required={param_required})\n"
-            )
-        else:
-            params_code += (
-                f"    {param_python_var_name} = Parameter('{param_name}', "
-                f"type={param_type}, help='''{param_help}''', required={param_required})\n"
-            )
-
-    project_decorator = f"@project(name='{project_name}')\n" if project_name else ""
-
-    contents = f"""\
-from metaflow import FlowSpec, Parameter, IncludeFile, JSONType, step, project
-{project_decorator}class {flow_name}(FlowSpec):
-{params_code}
-    @step
-    def start(self):
-        self.next(self.end)
-    @step
-    def end(self):
-        pass
-if __name__ == '__main__':
-    {flow_name}()
-"""
-    return contents
-
-
 class ArgoWorkflowsTriggeredRun(TriggeredRun):
     """
     A class representing a triggered Argo Workflow execution.
@@ -246,6 +203,38 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
 
     TYPE: ClassVar[Optional[str]] = "argo-workflows"
 
+    @classmethod
+    def list_deployed_flows(cls, flow_name: Optional[str] = None):
+        """
+        List all deployed Argo Workflow templates.
+
+        Parameters
+        ----------
+        flow_name : str, optional, default None
+            If specified, only list deployed flows for this specific flow name.
+            If None, list all deployed flows.
+
+        Yields
+        ------
+        ArgoWorkflowsDeployedFlow
+            `ArgoWorkflowsDeployedFlow` objects representing deployed
+            workflow templates on Argo Workflows.
+        """
+        from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
+
+        # When flow_name is None, use all=True to get all templates
+        # When flow_name is specified, use all=False to filter by flow_name
+        all_templates = flow_name is None
+        for template_name in ArgoWorkflows.list_templates(
+            flow_name=flow_name, all=all_templates
+        ):
+            try:
+                deployed_flow = cls.from_deployment(template_name)
+                yield deployed_flow
+            except Exception:
+                # Skip templates that can't be converted to DeployedFlow objects
+                continue
+
     @classmethod
     def from_deployment(cls, identifier: str, metadata: Optional[str] = None):
         """
@@ -321,6 +310,43 @@ class ArgoWorkflowsDeployedFlow(DeployedFlow):
 
         return cls(deployer=d)
 
+    @classmethod
+    def get_triggered_run(
+        cls, identifier: str, run_id: str, metadata: Optional[str] = None
+    ):
+        """
+        Retrieves a `ArgoWorkflowsTriggeredRun` object from an identifier, a run id and
+        optional metadata.
+
+        Parameters
+        ----------
+        identifier : str
+            Deployer specific identifier for the workflow to retrieve
+        run_id : str
+            Run ID for the which to fetch the triggered run object
+        metadata : str, optional, default None
+            Optional deployer specific metadata.
+
+        Returns
+        -------
+        ArgoWorkflowsTriggeredRun
+            A `ArgoWorkflowsTriggeredRun` object representing the
+            triggered run on argo workflows.
+        """
+        deployed_flow_obj = cls.from_deployment(identifier, metadata)
+        return ArgoWorkflowsTriggeredRun(
+            deployer=deployed_flow_obj.deployer,
+            content=json.dumps(
+                {
+                    "metadata": deployed_flow_obj.deployer.metadata,
+                    "pathspec": "/".join(
+                        (deployed_flow_obj.deployer.flow_name, run_id)
+                    ),
+                    "name": run_id,
+                }
+            ),
+        )
+
     @property
     def production_token(self) -> Optional[str]:
         """

metaflow/plugins/argo/capture_error.py

@@ -1,6 +1,6 @@
 import json
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 
 ###
 # Algorithm to determine 1st error:
@@ -26,6 +26,9 @@ def parse_workflow_failures():
 def group_failures_by_template(failures):
     groups = {}
     for failure in failures:
+        if failure.get("finishedAt", None) is None:
+            timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+            failure["finishedAt"] = timestamp
         groups.setdefault(failure["templateName"], []).append(failure)
     return groups
 
@@ -53,7 +56,7 @@ def determine_first_error():
     grouped_failures = group_failures_by_template(failures)
     for group in grouped_failures.values():
         group.sort(
-            key=lambda x: datetime.strptime(x["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
+            key=lambda g: datetime.strptime(g["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
         )
 
     earliest_group = grouped_failures[

metaflow/plugins/argo/conditional_input_paths.py

@@ -0,0 +1,35 @@
+from math import inf
+import sys
+from metaflow.util import decompress_list, compress_list
+import base64
+
+
+def generate_input_paths(input_paths, skippable_steps):
+    # => run_id/step/:foo,bar
+    # input_paths are base64 encoded due to Argo shenanigans
+    decoded = base64.b64decode(input_paths).decode("utf-8")
+    paths = decompress_list(decoded)
+
+    # some of the paths are going to be malformed due to never having executed per conditional.
+    # strip these out of the list.
+
+    # all pathspecs of leading steps that executed.
+    trimmed = [path for path in paths if not "{{" in path]
+
+    # pathspecs of leading steps that are conditional, and should be used instead of non-conditional ones
+    # e.g. the case of skipping switches: start -> case_step -> conditional_a or end
+    conditionals = [
+        path for path in trimmed if not any(step in path for step in skippable_steps)
+    ]
+    pathspecs_to_use = conditionals if conditionals else trimmed
+    return compress_list(pathspecs_to_use, zlibmin=inf)
+
+
+if __name__ == "__main__":
+    input_paths = sys.argv[1]
+    try:
+        skippable_steps = sys.argv[2].split(",")
+    except IndexError:
+        skippable_steps = []
+
+    print(generate_input_paths(input_paths, skippable_steps))

metaflow/plugins/argo/exit_hooks.py

@@ -0,0 +1,209 @@
+from collections import defaultdict
+import json
+from typing import Dict, List, Optional
+
+
+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: str,
+        url: str,
+        method: str = "GET",
+        headers: Optional[Dict] = None,
+        body: Optional[str] = None,
+        on_success: bool = False,
+        on_error: bool = 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(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/argo/param_val.py

@@ -0,0 +1,19 @@
+import sys
+import base64
+import json
+
+
+def parse_parameter_value(base64_value):
+    val = base64.b64decode(base64_value).decode("utf-8")
+
+    try:
+        return json.loads(val)
+    except json.decoder.JSONDecodeError:
+        # fallback to using the original value.
+        return val
+
+
+if __name__ == "__main__":
+    base64_val = sys.argv[1]
+
+    print(parse_parameter_value(base64_val))

metaflow/plugins/aws/aws_client.py

@@ -35,6 +35,12 @@ class Boto3ClientProvider(object):
                 "Could not import module 'boto3'. Install boto3 first."
             )
 
+        # Convert dictionary config to Config object if needed
+        if "config" in client_params and not isinstance(
+            client_params["config"], Config
+        ):
+            client_params["config"] = Config(**client_params["config"])
+
         if module == "s3" and (
             "config" not in client_params or client_params["config"].retries is None
         ):

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:
@@ -208,3 +208,35 @@ def sanitize_batch_tag(key, value):
     _value = re.sub(RE_NOT_PERMITTED, "", value)[:256]
 
     return _key, _value
+
+
+def validate_aws_tag(key: str, value: str):
+    PERMITTED = r"[A-Za-z0-9\s\+\-\=\.\_\:\/\@]"
+
+    AWS_PREFIX = r"^aws\:"  # case-insensitive.
+    if re.match(AWS_PREFIX, key, re.IGNORECASE) or re.match(
+        AWS_PREFIX, value, re.IGNORECASE
+    ):
+        raise MetaflowException(
+            "'aws:' is not an allowed prefix for either tag keys or values"
+        )
+
+    if len(key) > 128:
+        raise MetaflowException(
+            "Tag key *%s* is too long. Maximum allowed tag key length is 128." % key
+        )
+    if len(value) > 256:
+        raise MetaflowException(
+            "Tag value *%s* is too long. Maximum allowed tag value length is 256."
+            % value
+        )
+
+    if not re.match(PERMITTED, key):
+        raise MetaflowException(
+            "Key *s* is not permitted. Tags must match pattern: %s" % (key, PERMITTED)
+        )
+    if not re.match(PERMITTED, value):
+        raise MetaflowException(
+            "Value *%s* is not permitted. Tags must match pattern: %s"
+            % (value, PERMITTED)
+        )

metaflow/plugins/aws/batch/batch.py

@@ -53,20 +53,32 @@ class BatchKilledException(MetaflowException):
 
 
 class Batch(object):
-    def __init__(self, metadata, environment):
+    def __init__(self, metadata, environment, flow_datastore=None):
         self.metadata = metadata
         self.environment = environment
+        self.flow_datastore = flow_datastore
         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,
+        offload_command_to_s3,
+    ):
         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)
@@ -94,7 +106,43 @@ class Batch(object):
         # We lose the last logs in this scenario (although they are visible
         # still through AWS CloudWatch console).
         cmd_str += "c=$?; %s; exit $c" % BASH_SAVE_LOGS
-        return shlex.split('bash -c "%s"' % cmd_str)
+        command = shlex.split('bash -c "%s"' % cmd_str)
+
+        if not offload_command_to_s3:
+            return command
+
+        # If S3 upload is enabled, we need to modify the command after it's created
+        if self.flow_datastore is None:
+            raise MetaflowException(
+                "Can not offload Batch command to S3 without a datastore configured."
+            )
+
+        from metaflow.plugins.aws.aws_utils import parse_s3_full_path
+
+        # Get the command that was created
+        # Upload the command to S3 during deployment
+        try:
+            command_bytes = cmd_str.encode("utf-8")
+            result_paths = self.flow_datastore.save_data([command_bytes], len_hint=1)
+            s3_path, _key = result_paths[0]
+
+            bucket, s3_object = parse_s3_full_path(s3_path)
+            download_script = "{python} -c '{script}'".format(
+                python=self.environment._python(),
+                script='import boto3, os; ep=os.getenv(\\"METAFLOW_S3_ENDPOINT_URL\\"); boto3.client(\\"s3\\", **({\\"endpoint_url\\":ep} if ep else {})).download_file(\\"%s\\", \\"%s\\", \\"/tmp/step_command.sh\\")'
+                % (bucket, s3_object),
+            )
+            download_cmd = (
+                f"{self.environment._get_install_dependencies_cmd('s3')} && "  # required for boto3 due to the original dependencies cmd getting packaged, and not being downloaded in time.
+                f"{download_script} && "
+                f"chmod +x /tmp/step_command.sh && "
+                f"bash /tmp/step_command.sh"
+            )
+            new_cmd = shlex.split('bash -c "%s"' % download_cmd)
+            return new_cmd
+        except Exception as e:
+            print(f"Warning: Failed to upload command to S3: {e}")
+            print("Falling back to inline command")
 
     def _search_jobs(self, flow_name, run_id, user):
         if user is None:
@@ -167,6 +215,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -188,6 +237,7 @@ class Batch(object):
         host_volumes=None,
         efs_volumes=None,
         use_tmpfs=None,
+        aws_batch_tags=None,
         tmpfs_tempdir=None,
         tmpfs_size=None,
         tmpfs_path=None,
@@ -195,6 +245,7 @@ class Batch(object):
         ephemeral_storage=None,
         log_driver=None,
         log_options=None,
+        offload_command_to_s3=False,
     ):
         job_name = self._job_name(
             attrs.get("metaflow.user"),
@@ -210,7 +261,13 @@ 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,
+                    offload_command_to_s3,
                 )
             )
             .image(image)
@@ -249,6 +306,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)
@@ -327,6 +385,11 @@ class Batch(object):
                 if key in attrs:
                     k, v = sanitize_batch_tag(key, attrs.get(key))
                     job.tag(k, v)
+
+            if aws_batch_tags is not None:
+                for key, value in aws_batch_tags.items():
+                    job.tag(key, value)
+
         return job
 
     def launch_job(
@@ -334,6 +397,7 @@ class Batch(object):
         step_name,
         step_cli,
         task_spec,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
@@ -353,6 +417,7 @@ class Batch(object):
         host_volumes=None,
         efs_volumes=None,
         use_tmpfs=None,
+        aws_batch_tags=None,
         tmpfs_tempdir=None,
         tmpfs_size=None,
         tmpfs_path=None,
@@ -374,6 +439,7 @@ class Batch(object):
             step_name,
             step_cli,
             task_spec,
+            code_package_metadata,
             code_package_sha,
             code_package_url,
             code_package_ds,
@@ -395,6 +461,7 @@ class Batch(object):
             host_volumes=host_volumes,
             efs_volumes=efs_volumes,
             use_tmpfs=use_tmpfs,
+            aws_batch_tags=aws_batch_tags,
             tmpfs_tempdir=tmpfs_tempdir,
             tmpfs_size=tmpfs_size,
             tmpfs_path=tmpfs_path,