metaflow 2.18.2__tar.gz → 2.18.4__tar.gz

{metaflow-2.18.2/metaflow.egg-info → metaflow-2.18.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow
-Version: 2.18.2
+Version: 2.18.4
 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.4; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{metaflow-2.18.2 → metaflow-2.18.4}/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-2.18.2 → metaflow-2.18.4}/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-2.18.2 → metaflow-2.18.4}/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-2.18.2 → metaflow-2.18.4}/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,20 @@ class ArgoWorkflows(object):
     def delete(name):
         client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
 
+        # the workflow template might not exist, but we still want to try clean up associated sensors and schedules.
+        workflow_template = client.get_workflow_template(name) or {}
+        workflow_annotations = workflow_template.get("metadata", {}).get(
+            "annotations", {}
+        )
+
+        sensor_name = ArgoWorkflows._sensor_name(
+            workflow_annotations.get("metaflow/sensor_name", name)
+        )
+        # if below is missing then it was deployed before custom sensor namespaces
+        sensor_namespace = workflow_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 +281,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 +450,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 +780,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 +803,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 +984,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 +993,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 +1014,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 +1076,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 +1127,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 +1345,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,15 +1380,45 @@ 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
@@ -1699,7 +1862,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 +2046,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 +3701,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-2.18.4/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-2.18.2 → metaflow-2.18.4}/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)