metaflow 2.18.2__py2.py3-none-any.whl → 2.18.3__py2.py3-none-any.whl

metaflow/client/core.py

@@ -303,7 +303,7 @@ class MetaflowObject(object):
             # distinguish between "attempt will happen" and "no such
             # attempt exists".
 
-        if pathspec:
+        if pathspec and _object is None:
             ids = pathspec.split("/")
 
             if self._NAME == "flow" and len(ids) != 1:

metaflow/metaflow_config.py

@@ -344,6 +344,8 @@ SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH = from_conf(
         else None
     ),
 )
+# Toggle for step command being part of the Step Function payload, or if it should be offloaded to S3
+SFN_COMPRESS_STATE_MACHINE = from_conf("SFN_COMPRESS_STATE_MACHINE", False)
 ###
 # Kubernetes configuration
 ###
@@ -409,6 +411,9 @@ ARGO_EVENTS_INTERNAL_WEBHOOK_URL = from_conf(
     "ARGO_EVENTS_INTERNAL_WEBHOOK_URL", ARGO_EVENTS_WEBHOOK_URL
 )
 ARGO_EVENTS_WEBHOOK_AUTH = from_conf("ARGO_EVENTS_WEBHOOK_AUTH", "none")
+ARGO_EVENTS_SENSOR_NAMESPACE = from_conf(
+    "ARGO_EVENTS_SENSOR_NAMESPACE", KUBERNETES_NAMESPACE
+)
 
 ARGO_WORKFLOWS_UI_URL = from_conf("ARGO_WORKFLOWS_UI_URL")
 

metaflow/plugins/argo/argo_client.py

@@ -1,5 +1,6 @@
 import json
 
+from metaflow.metaflow_config import ARGO_EVENTS_SENSOR_NAMESPACE
 from metaflow.exception import MetaflowException
 from metaflow.plugins.kubernetes.kubernetes_client import KubernetesClient
 
@@ -377,12 +378,15 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def register_sensor(self, name, sensor=None):
+    def register_sensor(
+        self, name, sensor=None, sensor_namespace=ARGO_EVENTS_SENSOR_NAMESPACE
+    ):
         if sensor is None:
             sensor = {}
         # Unfortunately, Kubernetes client does not handle optimistic
         # concurrency control by itself unlike kubectl
         client = self._client.get()
+
         if not sensor:
             sensor["metadata"] = {}
 
@@ -392,7 +396,7 @@ class ArgoClient(object):
             ] = client.CustomObjectsApi().get_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )[
@@ -407,7 +411,7 @@ class ArgoClient(object):
                     return client.CustomObjectsApi().create_namespaced_custom_object(
                         group=self._group,
                         version=self._version,
-                        namespace=self._namespace,
+                        namespace=sensor_namespace,
                         plural="sensors",
                         body=sensor,
                     )
@@ -425,7 +429,7 @@ class ArgoClient(object):
             return client.CustomObjectsApi().replace_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 body=sensor,
                 name=name,
@@ -435,7 +439,7 @@ class ArgoClient(object):
                 json.loads(e.body)["message"] if e.body is not None else e.reason
             )
 
-    def delete_sensor(self, name):
+    def delete_sensor(self, name, sensor_namespace):
         """
         Issues an API call for deleting a sensor
 
@@ -447,7 +451,7 @@ class ArgoClient(object):
             return client.CustomObjectsApi().delete_namespaced_custom_object(
                 group=self._group,
                 version=self._version,
-                namespace=self._namespace,
+                namespace=sensor_namespace,
                 plural="sensors",
                 name=name,
             )

metaflow/plugins/argo/argo_workflows.py

@@ -19,6 +19,7 @@ from metaflow.metaflow_config import (
     ARGO_EVENTS_EVENT_BUS,
     ARGO_EVENTS_EVENT_SOURCE,
     ARGO_EVENTS_INTERNAL_WEBHOOK_URL,
+    ARGO_EVENTS_SENSOR_NAMESPACE,
     ARGO_EVENTS_SERVICE_ACCOUNT,
     ARGO_EVENTS_WEBHOOK_AUTH,
     ARGO_WORKFLOWS_CAPTURE_ERROR_SCRIPT,
@@ -73,6 +74,10 @@ class ArgoWorkflowsException(MetaflowException):
     headline = "Argo Workflows error"
 
 
+class ArgoWorkflowsSensorCleanupException(MetaflowException):
+    headline = "Argo Workflows sensor clean up error"
+
+
 class ArgoWorkflowsSchedulingException(MetaflowException):
     headline = "Argo Workflows scheduling error"
 
@@ -186,6 +191,7 @@ class ArgoWorkflows(object):
         return str(self._workflow_template)
 
     def deploy(self):
+        self.cleanup_previous_sensors()
         try:
             # Register workflow template.
             ArgoClient(namespace=KUBERNETES_NAMESPACE).register_workflow_template(
@@ -194,6 +200,37 @@ class ArgoWorkflows(object):
         except Exception as e:
             raise ArgoWorkflowsException(str(e))
 
+    def cleanup_previous_sensors(self):
+        try:
+            client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
+            # Check for existing deployment and do cleanup
+            old_template = client.get_workflow_template(self.name)
+            if not old_template:
+                return None
+            # Clean up old sensors
+            old_sensor_namespace = old_template["metadata"]["annotations"].get(
+                "metaflow/sensor_namespace"
+            )
+
+            if old_sensor_namespace is None:
+                # This workflow was created before sensor annotations
+                # and may have a sensor in the default namespace
+                # we will delete it and it'll get recreated if need be
+                old_sensor_name = ArgoWorkflows._sensor_name(self.name)
+                client.delete_sensor(old_sensor_name, client._namespace)
+            else:
+                # delete old sensor only if it was somewhere else, otherwise it'll get replaced
+                old_sensor_name = old_template["metadata"]["annotations"][
+                    "metaflow/sensor_name"
+                ]
+                if (
+                    not self._sensor
+                    or old_sensor_namespace != ARGO_EVENTS_SENSOR_NAMESPACE
+                ):
+                    client.delete_sensor(old_sensor_name, old_sensor_namespace)
+        except Exception as e:
+            raise ArgoWorkflowsSensorCleanupException(str(e))
+
     @staticmethod
     def _sanitize(name):
         # Metaflow allows underscores in node names, which are disallowed in Argo
@@ -221,6 +258,17 @@ class ArgoWorkflows(object):
     def delete(name):
         client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
 
+        workflow_template = client.get_workflow_template(name)
+        sensor_name = ArgoWorkflows._sensor_name(
+            workflow_template["metadata"]["annotations"].get(
+                "metaflow/sensor_name", name
+            )
+        )
+        # if below is missing then it was deployed before custom sensor namespaces
+        sensor_namespace = workflow_template["metadata"]["annotations"].get(
+            "metaflow/sensor_namespace", KUBERNETES_NAMESPACE
+        )
+
         # Always try to delete the schedule. Failure in deleting the schedule should not
         # be treated as an error, due to any of the following reasons
         # - there might not have been a schedule, or it was deleted by some other means
@@ -230,7 +278,7 @@ class ArgoWorkflows(object):
 
         # The workflow might have sensors attached to it, which consume actual resources.
         # Try to delete these as well.
-        sensor_deleted = client.delete_sensor(ArgoWorkflows._sensor_name(name))
+        sensor_deleted = client.delete_sensor(sensor_name, sensor_namespace)
 
         # After cleaning up related resources, delete the workflow in question.
         # Failure in deleting is treated as critical and will be made visible to the user
@@ -399,11 +447,10 @@ class ArgoWorkflows(object):
             # Metaflow will overwrite any existing sensor.
             sensor_name = ArgoWorkflows._sensor_name(self.name)
             if self._sensor:
-                argo_client.register_sensor(sensor_name, self._sensor.to_json())
-            else:
-                # Since sensors occupy real resources, delete existing sensor if needed
-                # Deregister sensors that might have existed before this deployment
-                argo_client.delete_sensor(sensor_name)
+                # The new sensor will go into the sensor namespace specified
+                ArgoClient(namespace=ARGO_EVENTS_SENSOR_NAMESPACE).register_sensor(
+                    sensor_name, self._sensor.to_json(), ARGO_EVENTS_SENSOR_NAMESPACE
+                )
         except Exception as e:
             raise ArgoWorkflowsSchedulingException(str(e))
 
@@ -730,6 +777,7 @@ class ArgoWorkflows(object):
         # references to them within the DAGTask.
 
         annotations = {}
+
         if self._schedule is not None:
             # timezone is an optional field and json dumps on None will result in null
             # hence configuring it to an empty string
@@ -752,7 +800,9 @@ class ArgoWorkflows(object):
                             {key: trigger.get(key) for key in ["name", "type"]}
                             for trigger in self.triggers
                         ]
-                    )
+                    ),
+                    "metaflow/sensor_name": ArgoWorkflows._sensor_name(self.name),
+                    "metaflow/sensor_namespace": ARGO_EVENTS_SENSOR_NAMESPACE,
                 }
             )
         if self.notify_on_error:
@@ -931,7 +981,7 @@ class ArgoWorkflows(object):
         node_conditional_parents = {}
         node_conditional_branches = {}
 
-        def _visit(node, seen, conditional_branch, conditional_parents=None):
+        def _visit(node, conditional_branch, conditional_parents=None):
             if not node.type == "split-switch" and not (
                 conditional_branch and conditional_parents
             ):
@@ -940,7 +990,10 @@ class ArgoWorkflows(object):
 
             if node.type == "split-switch":
                 conditional_branch = conditional_branch + [node.name]
-                node_conditional_branches[node.name] = conditional_branch
+                c_br = node_conditional_branches.get(node.name, [])
+                node_conditional_branches[node.name] = c_br + [
+                    b for b in conditional_branch if b not in c_br
+                ]
 
                 conditional_parents = (
                     [node.name]
@@ -958,21 +1011,36 @@ class ArgoWorkflows(object):
             if conditional_parents and not node.type == "split-switch":
                 node_conditional_parents[node.name] = conditional_parents
                 conditional_branch = conditional_branch + [node.name]
-                node_conditional_branches[node.name] = conditional_branch
+                c_br = node_conditional_branches.get(node.name, [])
+                node_conditional_branches[node.name] = c_br + [
+                    b for b in conditional_branch if b not in c_br
+                ]
 
                 self.conditional_nodes.add(node.name)
 
             if conditional_branch and conditional_parents:
                 for n in node.out_funcs:
                     child = self.graph[n]
-                    if n not in seen:
-                        _visit(
-                            child, seen + [n], conditional_branch, conditional_parents
-                        )
+                    if child.name == node.name:
+                        continue
+                    _visit(child, conditional_branch, conditional_parents)
 
         # First we visit all nodes to determine conditional parents and branches
         for n in self.graph:
-            _visit(n, [], [])
+            _visit(n, [])
+
+        # helper to clean up conditional info for all children of a node, until a new split-switch is encountered.
+        def _cleanup_conditional_status(node_name, seen):
+            if self.graph[node_name].type == "split-switch":
+                # stop recursive cleanup if we hit a new split-switch
+                return
+            if node_name in self.conditional_nodes:
+                self.conditional_nodes.remove(node_name)
+            node_conditional_parents[node_name] = []
+            node_conditional_branches[node_name] = []
+            for p in self.graph[node_name].out_funcs:
+                if p not in seen:
+                    _cleanup_conditional_status(p, seen + [p])
 
         # Then we traverse again in order to determine conditional join nodes, and matching conditional join info
         for node in self.graph:
@@ -1005,14 +1073,44 @@ class ArgoWorkflows(object):
                     last_conditional_split_nodes = self.graph[
                         last_split_switch
                     ].out_funcs
-                    # p needs to be in at least one conditional_branch for it to be closed.
-                    if all(
-                        any(
-                            p in node_conditional_branches.get(in_func, [])
-                            for in_func in conditional_in_funcs
+                    # NOTE: How do we define a conditional join step?
+                    # The idea here is that we check if the conditional branches(e.g. chains of conditional steps leading to) of all the in_funcs
+                    # manage to tick off every step name that follows a split-switch
+                    # For example, consider the following structure
+                    # switch_step -> A, B, C
+                    # A -> A2 -> A3 -> A4 -> B2
+                    # B -> B2 -> B3 -> C3
+                    # C -> C2 -> C3 -> end
+                    #
+                    # if we look at the in_funcs for C3, they are (C2, B3)
+                    # B3 closes off branches started by A and B
+                    # C3 closes off branches started by C
+                    # therefore C3 is a conditional join step for the 'switch_step'
+                    # NOTE: Then what about a skip step?
+                    # some switch cases might not introduce any distinct steps of their own, opting to instead skip ahead to a later common step.
+                    # Example:
+                    # switch_step -> A, B, C
+                    # A -> A1 -> B2 -> C
+                    # B -> B1 -> B2 -> C
+                    #
+                    # In this case, C is a skip step as it does not add any conditional branching of its own.
+                    # C is also a conditional join, as it closes all branches started by 'switch_step'
+
+                    closes_branches = all(
+                        (
+                            # branch_root_node_name needs to be in at least one conditional_branch for it to be closed.
+                            any(
+                                branch_root_node_name
+                                in node_conditional_branches.get(in_func, [])
+                                for in_func in conditional_in_funcs
+                            )
+                            # need to account for a switch case skipping completely, not having a conditional-branch of its own.
+                            if branch_root_node_name != node.name
+                            else True
                         )
-                        for p in last_conditional_split_nodes
-                    ):
+                        for branch_root_node_name in last_conditional_split_nodes
+                    )
+                    if closes_branches:
                         closed_conditional_parents.append(last_split_switch)
 
                         self.conditional_join_nodes.add(node.name)
@@ -1026,25 +1124,45 @@ class ArgoWorkflows(object):
                     for p in node_conditional_parents.get(node.name, [])
                     if p not in closed_conditional_parents
                 ]:
-                    if node.name in self.conditional_nodes:
-                        self.conditional_nodes.remove(node.name)
-                    node_conditional_parents[node.name] = []
-                    for p in node.out_funcs:
-                        if p in self.conditional_nodes:
-                            self.conditional_nodes.remove(p)
-                        node_conditional_parents[p] = []
+                    _cleanup_conditional_status(node.name, [])
 
     def _is_conditional_node(self, node):
         return node.name in self.conditional_nodes
 
+    def _is_conditional_skip_node(self, node):
+        return (
+            self._is_conditional_node(node)
+            and any(
+                self.graph[in_func].type == "split-switch" for in_func in node.in_funcs
+            )
+            and len(
+                [
+                    in_func
+                    for in_func in node.in_funcs
+                    if self._is_conditional_node(self.graph[in_func])
+                    or self.graph[in_func].type == "split-switch"
+                ]
+            )
+            > 1
+        )
+
     def _is_conditional_join_node(self, node):
         return node.name in self.conditional_join_nodes
 
+    def _many_in_funcs_all_conditional(self, node):
+        cond_in_funcs = [
+            in_func
+            for in_func in node.in_funcs
+            if self._is_conditional_node(self.graph[in_func])
+        ]
+        return len(cond_in_funcs) > 1 and len(cond_in_funcs) == len(node.in_funcs)
+
     def _is_recursive_node(self, node):
         return node.name in self.recursive_nodes
 
     def _matching_conditional_join(self, node):
-        return self.matching_conditional_join_dict.get(node.name, None)
+        # If no earlier conditional join step is found during parsing, then 'end' is always one.
+        return self.matching_conditional_join_dict.get(node.name, "end")
 
     # Visit every node and yield the uber DAGTemplate(s).
     def _dag_templates(self):
@@ -1224,12 +1342,24 @@ class ArgoWorkflows(object):
                     "%s.Succeeded" % self._sanitize(in_func)
                     for in_func in node.in_funcs
                     if self._is_conditional_node(self.graph[in_func])
+                    or self.graph[in_func].type == "split-switch"
                 ]
                 required_deps = [
                     "%s.Succeeded" % self._sanitize(in_func)
                     for in_func in node.in_funcs
                     if not self._is_conditional_node(self.graph[in_func])
+                    and self.graph[in_func].type != "split-switch"
                 ]
+                if self._is_conditional_skip_node(
+                    node
+                ) or self._many_in_funcs_all_conditional(node):
+                    # skip nodes need unique condition handling
+                    conditional_deps = [
+                        "%s.Succeeded" % self._sanitize(in_func)
+                        for in_func in node.in_funcs
+                    ]
+                    required_deps = []
+
                 both_conditions = required_deps and conditional_deps
 
                 depends_str = "{required}{_and}{conditional}".format(
@@ -1247,16 +1377,46 @@ class ArgoWorkflows(object):
                 )
 
                 # Add conditional if this is the first step in a conditional branch
+                switch_in_funcs = [
+                    in_func
+                    for in_func in node.in_funcs
+                    if self.graph[in_func].type == "split-switch"
+                ]
                 if (
                     self._is_conditional_node(node)
-                    and self.graph[node.in_funcs[0]].type == "split-switch"
-                ):
-                    in_func = node.in_funcs[0]
-                    dag_task.when(
-                        "{{tasks.%s.outputs.parameters.switch-step}}==%s"
-                        % (self._sanitize(in_func), node.name)
+                    or self._is_conditional_skip_node(node)
+                    or self._is_conditional_join_node(node)
+                ) and switch_in_funcs:
+                    conditional_when = "||".join(
+                        [
+                            "{{tasks.%s.outputs.parameters.switch-step}}==%s"
+                            % (self._sanitize(switch_in_func), node.name)
+                            for switch_in_func in switch_in_funcs
+                        ]
                     )
 
+                    non_switch_in_funcs = [
+                        in_func
+                        for in_func in node.in_funcs
+                        if in_func not in switch_in_funcs
+                    ]
+                    status_when = ""
+                    if non_switch_in_funcs:
+                        status_when = "||".join(
+                            [
+                                "{{tasks.%s.status}}==Succeeded"
+                                % self._sanitize(in_func)
+                                for in_func in non_switch_in_funcs
+                            ]
+                        )
+
+                    total_when = (
+                        f"({status_when}) || ({conditional_when})"
+                        if status_when
+                        else conditional_when
+                    )
+                    dag_task.when(total_when)
+
             dag_tasks.append(dag_task)
             # End the workflow if we have reached the end of the flow
             if node.type == "end":
@@ -1699,7 +1859,11 @@ class ArgoWorkflows(object):
                 input_paths_expr = (
                     "export INPUT_PATHS={{inputs.parameters.input-paths}}"
                 )
-                if self._is_conditional_join_node(node):
+                if (
+                    self._is_conditional_join_node(node)
+                    or self._many_in_funcs_all_conditional(node)
+                    or self._is_conditional_skip_node(node)
+                ):
                     # NOTE: Argo template expressions that fail to resolve, output the expression itself as a value.
                     # With conditional steps, some of the input-paths are therefore 'broken' due to containing a nil expression
                     # e.g. "{{ tasks['A'].outputs.parameters.task-id }}" when task A never executed.
@@ -1879,20 +2043,33 @@ class ArgoWorkflows(object):
                 )
                 input_paths = "%s/_parameters/%s" % (run_id, task_id_params)
             # Only for static joins and conditional_joins
-            elif self._is_conditional_join_node(node) and not (
+            elif (
+                self._is_conditional_join_node(node)
+                or self._many_in_funcs_all_conditional(node)
+                or self._is_conditional_skip_node(node)
+            ) and not (
                 node.type == "join"
                 and self.graph[node.split_parents[-1]].type == "foreach"
             ):
+                # we need to pass in the set of conditional in_funcs to the pathspec generating script as in the case of split-switch skipping cases,
+                # non-conditional input-paths need to be ignored in favour of conditional ones when they have executed.
+                skippable_input_steps = ",".join(
+                    [
+                        in_func
+                        for in_func in node.in_funcs
+                        if self.graph[in_func].type == "split-switch"
+                    ]
+                )
                 input_paths = (
-                    "$(python -m metaflow.plugins.argo.conditional_input_paths %s)"
-                    % input_paths
+                    "$(python -m metaflow.plugins.argo.conditional_input_paths %s %s)"
+                    % (input_paths, skippable_input_steps)
                 )
             elif (
                 node.type == "join"
                 and self.graph[node.split_parents[-1]].type == "foreach"
             ):
                 # foreach-joins straight out of conditional branches are not yet supported
-                if self._is_conditional_join_node(node):
+                if self._is_conditional_join_node(node) and len(node.in_funcs) > 1:
                     raise ArgoWorkflowsException(
                         "Conditional steps inside a foreach that transition directly into a join step are not currently supported.\n"
                         "As a workaround, add a common step after the conditional steps %s "
@@ -3521,7 +3698,7 @@ class ArgoWorkflows(object):
                 # Sensor metadata.
                 ObjectMeta()
                 .name(ArgoWorkflows._sensor_name(self.name))
-                .namespace(KUBERNETES_NAMESPACE)
+                .namespace(ARGO_EVENTS_SENSOR_NAMESPACE)
                 .labels(self._base_labels)
                 .label("app.kubernetes.io/name", "metaflow-sensor")
                 .annotations(self._base_annotations)

metaflow/plugins/argo/conditional_input_paths.py

@@ -4,7 +4,7 @@ from metaflow.util import decompress_list, compress_list
 import base64
 
 
-def generate_input_paths(input_paths):
+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")
@@ -13,9 +13,23 @@ def generate_input_paths(input_paths):
     # 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]
-    return compress_list(trimmed, zlibmin=inf)
+
+    # 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__":
-    print(generate_input_paths(sys.argv[1]))
+    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/aws/batch/batch.py

@@ -53,9 +53,10 @@ 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)
 
@@ -67,6 +68,7 @@ class Batch(object):
         step_name,
         step_cmds,
         task_spec,
+        offload_command_to_s3,
     ):
         mflog_expr = export_mflog_env_vars(
             datastore_type="s3",
@@ -104,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:
@@ -207,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"),
@@ -228,6 +267,7 @@ class Batch(object):
                     step_name,
                     [step_cli],
                     task_spec,
+                    offload_command_to_s3,
                 )
             )
             .image(image)

metaflow/plugins/aws/step_functions/step_functions.py

@@ -57,6 +57,7 @@ class StepFunctions(object):
         workflow_timeout=None,
         is_project=False,
         use_distributed_map=False,
+        compress_state_machine=False,
     ):
         self.name = name
         self.graph = graph
@@ -81,6 +82,9 @@ class StepFunctions(object):
         # https://aws.amazon.com/blogs/aws/step-functions-distributed-map-a-serverless-solution-for-large-scale-parallel-data-processing/
         self.use_distributed_map = use_distributed_map
 
+        # S3 command upload configuration
+        self.compress_state_machine = compress_state_machine
+
         self._client = StepFunctionsClient()
         self._workflow = self._compile()
         self._cron = self._cron()
@@ -858,7 +862,7 @@ class StepFunctions(object):
         # merge batch tags supplied through step-fuctions CLI and ones defined in decorator
         batch_tags = {**self.aws_batch_tags, **resources["aws_batch_tags"]}
         return (
-            Batch(self.metadata, self.environment)
+            Batch(self.metadata, self.environment, self.flow_datastore)
             .create_job(
                 step_name=node.name,
                 step_cli=self._step_cli(
@@ -894,6 +898,7 @@ class StepFunctions(object):
                 ephemeral_storage=resources["ephemeral_storage"],
                 log_driver=resources["log_driver"],
                 log_options=resources["log_options"],
+                offload_command_to_s3=self.compress_state_machine,
             )
             .attempts(total_retries + 1)
         )

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -10,6 +10,7 @@ from metaflow.metaflow_config import (
     FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     SERVICE_VERSION_CHECK,
     SFN_STATE_MACHINE_PREFIX,
+    SFN_COMPRESS_STATE_MACHINE,
     UI_URL,
 )
 from metaflow.package import MetaflowPackage
@@ -140,6 +141,12 @@ def step_functions(obj, name=None):
     help="Use AWS Step Functions Distributed Map instead of Inline Map for "
     "defining foreach tasks in Amazon State Language.",
 )
+@click.option(
+    "--compress-state-machine/--no-compress-state-machine",
+    is_flag=True,
+    default=SFN_COMPRESS_STATE_MACHINE,
+    help="Compress AWS Step Functions state machine to fit within the 8K limit.",
+)
 @click.option(
     "--deployer-attribute-file",
     default=None,
@@ -162,6 +169,7 @@ def create(
     workflow_timeout=None,
     log_execution_history=False,
     use_distributed_map=False,
+    compress_state_machine=False,
     deployer_attribute_file=None,
 ):
     for node in obj.graph:
@@ -212,6 +220,7 @@ def create(
         workflow_timeout,
         obj.is_project,
         use_distributed_map,
+        compress_state_machine,
     )
 
     if only_json:
@@ -332,6 +341,7 @@ def make_flow(
     workflow_timeout,
     is_project,
     use_distributed_map,
+    compress_state_machine=False,
 ):
     if obj.flow_datastore.TYPE != "s3":
         raise MetaflowException("AWS Step Functions requires --datastore=s3.")
@@ -390,6 +400,7 @@ def make_flow(
         workflow_timeout=workflow_timeout,
         is_project=is_project,
         use_distributed_map=use_distributed_map,
+        compress_state_machine=compress_state_machine,
     )
 
 

metaflow/plugins/aws/step_functions/step_functions_deployer.py

@@ -76,6 +76,9 @@ class StepFunctionsDeployer(DeployerImpl):
         use_distributed_map : bool, optional, default False
             Use AWS Step Functions Distributed Map instead of Inline Map for defining foreach
             tasks in Amazon State Language.
+        compress_state_machine : bool, optional, default False
+            Compress AWS Step Functions state machine to fit within the 8K limit.
+
         deployer_attribute_file : str, optional, default None
             Write the workflow name to the specified file. Used internally for Metaflow's Deployer API.
 

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.18.2"
+metaflow_version = "2.18.3"

{metaflow-2.18.2.dist-info → metaflow-2.18.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow
-Version: 2.18.2
+Version: 2.18.3
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.18.2; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.18.3; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{metaflow-2.18.2.dist-info → metaflow-2.18.3.dist-info}/RECORD

@@ -16,7 +16,7 @@ metaflow/includefile.py,sha256=RtISGl1V48qjkJBakUZ9yPpHV102h7pOIFiKP8PLHpc,20927
 metaflow/integrations.py,sha256=LlsaoePRg03DjENnmLxZDYto3NwWc9z_PtU6nJxLldg,1480
 metaflow/lint.py,sha256=A2NdUq_MnQal_RUCMC8ZOSR0VYZGyi2mSgwPQB0UzQo,15343
 metaflow/meta_files.py,sha256=vlgJHI8GJUKzXoxdrVoH8yyCF5bhFgwYemUgnyd1wgM,342
-metaflow/metaflow_config.py,sha256=v0bGKQ3EMGsRyd-UkHZxwrISV_5jsnxFLn8PG2_dyzs,24272
+metaflow/metaflow_config.py,sha256=Abn_4SiqpdrmMOSGYlFnEUUzVJqwfTFpvYwXlpXwcMw,24551
 metaflow/metaflow_config_funcs.py,sha256=5GlvoafV6SxykwfL8D12WXSfwjBN_NsyuKE_Q3gjGVE,6738
 metaflow/metaflow_current.py,sha256=pfkXmkyHeMJhxIs6HBJNBEaBDpcl5kz9Wx5mW6F_3qo,7164
 metaflow/metaflow_environment.py,sha256=20PIhA5R_rJneNj8f8UaWRmznGRPcEd6hP7goj_rc1s,11477
@@ -36,7 +36,7 @@ metaflow/tuple_util.py,sha256=_G5YIEhuugwJ_f6rrZoelMFak3DqAR2tt_5CapS1XTY,830
 metaflow/unbounded_foreach.py,sha256=p184WMbrMJ3xKYHwewj27ZhRUsSj_kw1jlye5gA9xJk,387
 metaflow/util.py,sha256=g2SOU_CRzJLgDM_UGF9QDMANMAIHAsDRXE6S76_YzsY,14594
 metaflow/vendor.py,sha256=A82CGHfStZGDP5pQ5XzRjFkbN1ZC-vFmghXIrzMDDNg,5868
-metaflow/version.py,sha256=TtcVTU2Og_O5qZDagNFBeiE246eQcw_mTgBScy5TV7U,28
+metaflow/version.py,sha256=tMCDMFSk5MlOlJbUQrrKKLwLAp53oxtJaV317s14sEM,28
 metaflow/_vendor/__init__.py,sha256=y_CiwUD3l4eAKvTVDZeqgVujMy31cAM1qjAB-HfI-9s,353
 metaflow/_vendor/typing_extensions.py,sha256=q9zxWa6p6CzF1zZvSkygSlklduHf_b3K7MCxGz7MJRc,134519
 metaflow/_vendor/zipp.py,sha256=ajztOH-9I7KA_4wqDYygtHa6xUBVZgFpmZ8FE74HHHI,8425
@@ -157,7 +157,7 @@ metaflow/cli_components/run_cmds.py,sha256=_xk2asy3SkqsJfZVhbgYSJ2rkkJe7cvLik6b0
 metaflow/cli_components/step_cmd.py,sha256=zGJgTv7wxrv34nWDi__CHaC2eS6kItR95EdVGJX803w,4766
 metaflow/cli_components/utils.py,sha256=gpoDociadjnJD7MuiJup_MDR02ZJjjleejr0jPBu29c,6057
 metaflow/client/__init__.py,sha256=1GtQB4Y_CBkzaxg32L1syNQSlfj762wmLrfrDxGi1b8,226
-metaflow/client/core.py,sha256=tj2PuqQt1RXg8GuyLQ_WRuxYEvTaDyi_lW18YGwWvAQ,83714
+metaflow/client/core.py,sha256=6h4oohNXZkP2QC0M4dKdrwz5Zx5s6j1_DIWB8VYmbN0,83734
 metaflow/client/filecache.py,sha256=Wy0yhhCqC1JZgebqi7z52GCwXYnkAqMZHTtxThvwBgM,15229
 metaflow/cmd/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 metaflow/cmd/configure_cmd.py,sha256=o-DKnUf2FBo_HiMVyoyzQaGBSMtpbEPEdFTQZ0hkU-k,33396
@@ -228,15 +228,15 @@ metaflow/plugins/airflow/sensors/base_sensor.py,sha256=s-OQBfPWZ_T3wn96Ua59CCEj1
 metaflow/plugins/airflow/sensors/external_task_sensor.py,sha256=zhYlrZnXT20KW8-fVk0fCNtTyNiKJB5PMVASacu30r0,6034
 metaflow/plugins/airflow/sensors/s3_sensor.py,sha256=iDReG-7FKnumrtQg-HY6cCUAAqNA90nARrjjjEEk_x4,3275
 metaflow/plugins/argo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/argo/argo_client.py,sha256=jLz0FjCTBvFLZt-8lZcMQhDcInhgEcGdPrU2Gvh67zA,17080
+metaflow/plugins/argo/argo_client.py,sha256=oT4ZrCyE7CYEbqNN0SfoZfSHd5fYW9XtuOrQEiUd1co,17230
 metaflow/plugins/argo/argo_events.py,sha256=_C1KWztVqgi3zuH57pInaE9OzABc2NnncC-zdwOMZ-w,5909
-metaflow/plugins/argo/argo_workflows.py,sha256=URF89Inj4gXQJ5p4_nFPTQihj-LlQZP0yq38kIlmhKU,207599
+metaflow/plugins/argo/argo_workflows.py,sha256=t132jaRyVWRLNM8IF0wupyp0cW-_6VDO-woIPrWG_B4,215964
 metaflow/plugins/argo/argo_workflows_cli.py,sha256=-blfZp-kAS8oWFTarfou9gRyE4QCnnJwa-0g8QuE0zk,52280
 metaflow/plugins/argo/argo_workflows_decorator.py,sha256=CLSjPqFTGucZ2_dSQGAYkoWWUZBQ9TCBXul4rxhDj3w,8282
 metaflow/plugins/argo/argo_workflows_deployer.py,sha256=6kHxEnYXJwzNCM9swI8-0AckxtPWqwhZLerYkX8fxUM,4444
 metaflow/plugins/argo/argo_workflows_deployer_objects.py,sha256=ydBE-lP42eNKvep36nQdUBPS3rQQErvoA7rCgyp5M6I,14949
 metaflow/plugins/argo/capture_error.py,sha256=Ys9dscGrTpW-ZCirLBU0gD9qBM0BjxyxGlUMKcwewQc,1852
-metaflow/plugins/argo/conditional_input_paths.py,sha256=kpzwyjXjO7PjeAJMapYX_ajUulFQJvCj-2DhH7OHzy0,645
+metaflow/plugins/argo/conditional_input_paths.py,sha256=Vtca74XbhnqAXgJJXKasLEa28jZbKBZPC5w4NAIOURc,1251
 metaflow/plugins/argo/exit_hooks.py,sha256=nh8IEkzAtQnbKVnh3N9CVnVKZB39Bjm3e0LFrACsLz8,6109
 metaflow/plugins/argo/generate_input_paths.py,sha256=loYsI6RFX9LlFsHb7Fe-mzlTTtRdySoOu7sYDy-uXK0,881
 metaflow/plugins/argo/jobset_input_paths.py,sha256=-h0E_e0w6FMiBUod9Rf_XOSCtZv_C0exacw4q1SfIfg,501
@@ -244,7 +244,7 @@ metaflow/plugins/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 metaflow/plugins/aws/aws_client.py,sha256=BTiLMXa1agjja-N73oWinaOZHs-lGPbfKJG8CqdRgaU,4287
 metaflow/plugins/aws/aws_utils.py,sha256=5mGZLu6wwTo2KzIke_MFqiomM3sYjAkU7Fx55dIMLfg,8561
 metaflow/plugins/aws/batch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow/plugins/aws/batch/batch.py,sha256=sdhpQVG3QAWD6KlpZH1MRCQcdwO0S4D-qXxgRK2A4bE,18466
+metaflow/plugins/aws/batch/batch.py,sha256=ugQ4YOjEs_PS5VtilMF0u1AArAZJTZPHCrCoVHPSqHc,20395
 metaflow/plugins/aws/batch/batch_cli.py,sha256=yZAy6WqZPsCqxjMnRhnTellCjLo27tD_OcoBLHezH8E,12508
 metaflow/plugins/aws/batch/batch_client.py,sha256=J50RMEXeEXFe5RqNUM1HN22BuDQFYFVQ4FSMOK55VWY,28798
 metaflow/plugins/aws/batch/batch_decorator.py,sha256=apKBywlseCGxoKxpLLYQSjK3rz8QEJom4ekKf6OvzNU,20104
@@ -256,11 +256,11 @@ metaflow/plugins/aws/step_functions/event_bridge_client.py,sha256=U9-tqKdih4KR-Z
 metaflow/plugins/aws/step_functions/production_token.py,sha256=rREx9djJzKYDiGhPCZ919pSpfrBCYuhSL5WlwnAojNM,1890
 metaflow/plugins/aws/step_functions/schedule_decorator.py,sha256=Ab1rW8O_no4HNZm4__iBmFDCDW0Z8-TgK4lnxHHA6HI,1940
 metaflow/plugins/aws/step_functions/set_batch_environment.py,sha256=ibiGWFHDjKcLfprH3OsX-g2M9lUsh6J-bp7v2cdLhD4,1294
-metaflow/plugins/aws/step_functions/step_functions.py,sha256=dAdJVo2UMsUe2StQak9VlUJK3-GtJ_50jlG8jk5QWLc,54110
-metaflow/plugins/aws/step_functions/step_functions_cli.py,sha256=MF3UPdZvRgVfZrcD4UCu0Fq-n-EUMM8m7RK57iUvOQs,27152
+metaflow/plugins/aws/step_functions/step_functions.py,sha256=jj21WBIvsNULnIjBDFwbn6h__GPRF1sxH6X-DjLGCY4,54340
+metaflow/plugins/aws/step_functions/step_functions_cli.py,sha256=dgCA3RbRfYf48wqzM_JJYiIH-eqcMjAXGb_jdXCfsp4,27556
 metaflow/plugins/aws/step_functions/step_functions_client.py,sha256=DKpNwAIWElvWjFANs5Ku3rgzjxFoqAD6k-EF8Xhkg3Q,4754
 metaflow/plugins/aws/step_functions/step_functions_decorator.py,sha256=jzDHYmgU_XvLffZDazR_1viow_1qQFblx9UKyjtoM_0,3788
-metaflow/plugins/aws/step_functions/step_functions_deployer.py,sha256=JKYtDhKivtXUWPklprZFzkqezh14loGDmk8mNk6QtpI,3714
+metaflow/plugins/aws/step_functions/step_functions_deployer.py,sha256=MOQ6H42szNHJw0ii3RJm7eyVGCmHL2j8kuZp7AuoiA8,3860
 metaflow/plugins/aws/step_functions/step_functions_deployer_objects.py,sha256=n7AEPs3uULXEuG3TVf2ZlTNq1LFd2n7x1IPVO2T5Ekk,8174
 metaflow/plugins/azure/__init__.py,sha256=GuuhTVC-zSdyAf79a1wiERMq0Zts7fwVT7t9fAf234A,100
 metaflow/plugins/azure/azure_credential.py,sha256=JmdGEbVzgxy8ucqnQDdTTI_atyMX9WSZUw3qYOo7RhE,2174
@@ -428,12 +428,12 @@ metaflow/user_decorators/mutable_flow.py,sha256=EywKTN3cnXPQF_s62wQaC4a4aH14j8oe
 metaflow/user_decorators/mutable_step.py,sha256=-BY0UDXf_RCAEnC5JlLzEXGdiw1KD9oSrSxS_SWaB9Y,16791
 metaflow/user_decorators/user_flow_decorator.py,sha256=2yDwZq9QGv9W-7kEuKwa8o4ZkTvuHJ5ESz7VVrGViAI,9890
 metaflow/user_decorators/user_step_decorator.py,sha256=4558NR8RJtN22OyTwCXO80bAMhMTaRGMoX12b1GMcPc,27232
-metaflow-2.18.2.data/data/share/metaflow/devtools/Makefile,sha256=TT4TCq8ALSfqYyGqDPocN5oPcZe2FqoCZxmGO1LmyCc,13760
-metaflow-2.18.2.data/data/share/metaflow/devtools/Tiltfile,sha256=Ty5p6AD3MwJAcAnOGv4yMz8fExAsnNQ11r8whK6uzzw,21381
-metaflow-2.18.2.data/data/share/metaflow/devtools/pick_services.sh,sha256=DCnrMXwtApfx3B4S-YiZESMyAFHbXa3VuNL0MxPLyiE,2196
-metaflow-2.18.2.dist-info/licenses/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
-metaflow-2.18.2.dist-info/METADATA,sha256=XxsMQK9awCthO2EntKCQ3nHRN8T4SOL2ODQlsGoqtJk,6740
-metaflow-2.18.2.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-metaflow-2.18.2.dist-info/entry_points.txt,sha256=RvEq8VFlgGe_FfqGOZi0D7ze1hLD0pAtXeNyGfzc_Yc,103
-metaflow-2.18.2.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
-metaflow-2.18.2.dist-info/RECORD,,
+metaflow-2.18.3.data/data/share/metaflow/devtools/Makefile,sha256=TT4TCq8ALSfqYyGqDPocN5oPcZe2FqoCZxmGO1LmyCc,13760
+metaflow-2.18.3.data/data/share/metaflow/devtools/Tiltfile,sha256=Ty5p6AD3MwJAcAnOGv4yMz8fExAsnNQ11r8whK6uzzw,21381
+metaflow-2.18.3.data/data/share/metaflow/devtools/pick_services.sh,sha256=DCnrMXwtApfx3B4S-YiZESMyAFHbXa3VuNL0MxPLyiE,2196
+metaflow-2.18.3.dist-info/licenses/LICENSE,sha256=nl_Lt5v9VvJ-5lWJDT4ddKAG-VZ-2IaLmbzpgYDz2hU,11343
+metaflow-2.18.3.dist-info/METADATA,sha256=yH6g8ewcnbbincQ9dCZO6YAifTpoA8DzJyIyEbq1Ojw,6740
+metaflow-2.18.3.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+metaflow-2.18.3.dist-info/entry_points.txt,sha256=RvEq8VFlgGe_FfqGOZi0D7ze1hLD0pAtXeNyGfzc_Yc,103
+metaflow-2.18.3.dist-info/top_level.txt,sha256=v1pDHoWaSaKeuc5fKTRSfsXCKSdW1zvNVmvA-i0if3o,9
+metaflow-2.18.3.dist-info/RECORD,,