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

metaflow/plugins/argo/capture_error.py

@@ -0,0 +1,73 @@
+import json
+import os
+from datetime import datetime, timezone
+
+###
+# Algorithm to determine 1st error:
+#   ignore the failures where message = ""
+#   group the failures via templateName
+#     sort each group by finishedAt
+#   find the group for which the last finishedAt is earliest
+#   if the earliest message is "No more retries left" then
+#     get the n-1th message from that group
+#   else
+#     return the last message.
+###
+
+
+def parse_workflow_failures():
+    failures = json.loads(
+        json.loads(os.getenv("METAFLOW_ARGO_WORKFLOW_FAILURES", "[]"), strict=False),
+        strict=False,
+    )
+    return [wf for wf in failures if wf.get("message")]
+
+
+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
+
+
+def sort_by_finished_at(items):
+    return sorted(
+        items, key=lambda x: datetime.strptime(x["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
+    )
+
+
+def find_earliest_last_finished_group(groups):
+    return min(
+        groups,
+        key=lambda k: datetime.strptime(
+            groups[k][-1]["finishedAt"], "%Y-%m-%dT%H:%M:%SZ"
+        ),
+    )
+
+
+def determine_first_error():
+    failures = parse_workflow_failures()
+    if not failures:
+        return None
+
+    grouped_failures = group_failures_by_template(failures)
+    for group in grouped_failures.values():
+        group.sort(
+            key=lambda g: datetime.strptime(g["finishedAt"], "%Y-%m-%dT%H:%M:%SZ")
+        )
+
+    earliest_group = grouped_failures[
+        find_earliest_last_finished_group(grouped_failures)
+    ]
+
+    if earliest_group[-1]["message"] == "No more retries left":
+        return earliest_group[-2]
+    return earliest_group[-1]
+
+
+if __name__ == "__main__":
+    first_err = determine_first_error()
+    print(json.dumps(first_err, indent=2))

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/jobset_input_paths.py

@@ -0,0 +1,15 @@
+import sys
+
+
+def generate_input_paths(run_id, step_name, task_id_entropy, num_parallel):
+    # => run_id/step/:foo,bar
+    control_id = "control-{}-0".format(task_id_entropy)
+    worker_ids = [
+        "worker-{}-{}".format(task_id_entropy, i) for i in range(int(num_parallel) - 1)
+    ]
+    ids = [control_id] + worker_ids
+    return "{}/{}/:{}".format(run_id, step_name, ",".join(ids))
+
+
+if __name__ == "__main__":
+    print(generate_input_paths(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]))

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

@@ -14,6 +14,7 @@ class Boto3ClientProvider(object):
             AWS_SANDBOX_ENABLED,
             AWS_SANDBOX_STS_ENDPOINT_URL,
             AWS_SANDBOX_API_KEY,
+            S3_CLIENT_RETRY_CONFIG,
         )
 
         if session_vars is None:
@@ -34,13 +35,19 @@ 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
         ):
-            # Use the adaptive retry strategy by default -- do not set anything if
-            # the user has already set something
+            # do not set anything if the user has already set something
             config = client_params.get("config", Config())
-            config.retries = {"max_attempts": 10, "mode": "adaptive"}
+            config.retries = S3_CLIENT_RETRY_CONFIG
+            client_params["config"] = config
 
         if AWS_SANDBOX_ENABLED:
             # role is ignored in the sandbox

metaflow/plugins/aws/aws_utils.py

@@ -1,10 +1,25 @@
 import re
-import requests
 
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import MAX_MEMORY_PER_TASK, MAX_CPU_PER_TASK
 
 
+def parse_s3_full_path(s3_uri):
+    from urllib.parse import urlparse
+
+    #  <scheme>://<netloc>/<path>;<params>?<query>#<fragment>
+    scheme, netloc, path, _, _, _ = urlparse(s3_uri)
+    assert scheme == "s3"
+    assert netloc is not None
+
+    bucket = netloc
+    path = path.lstrip("/").rstrip("/")
+    if path == "":
+        path = None
+
+    return bucket, path
+
+
 def get_ec2_instance_metadata():
     """
     Fetches the EC2 instance metadata through AWS instance metadata service
@@ -15,6 +30,10 @@ def get_ec2_instance_metadata():
         - ec2-region
         - ec2-availability-zone
     """
+
+    # TODO: Remove dependency on requests
+    import requests
+
     meta = {}
     # Capture AWS instance identity metadata. This is best-effort only since
     # access to this end-point might be blocked on AWS and not available
@@ -30,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:
@@ -144,6 +163,8 @@ def compute_resource_attributes(decos, compute_deco, step_name, resource_default
                         # Here we don't have ints, so we compare the value and raise
                         # an exception if not equal
                         if my_val != v:
+                            # TODO: Throw a better exception since the user has no
+                            #       knowledge of 'compute' decorator
                             raise MetaflowException(
                                 "'resources' and compute decorator have conflicting "
                                 "values for '%s'. Please use consistent values or "
@@ -187,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,