ob-metaflow 2.17.3.1__tar.gz → 2.18.1.1__tar.gz

{ob_metaflow-2.17.3.1/ob_metaflow.egg-info → ob_metaflow-2.18.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ob-metaflow
-Version: 2.17.3.1
+Version: 2.18.1.1
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co
@@ -12,7 +12,7 @@ Requires-Dist: boto3
 Requires-Dist: pylint
 Requires-Dist: kubernetes
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.17.3.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.18.1.1; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/datastore/task_datastore.py

@@ -222,6 +222,9 @@ class TaskDataStore(object):
     @property
     def pathspec_index(self):
         idxstr = ",".join(map(str, (f.index for f in self["_foreach_stack"])))
+        if "_iteration_stack" in self:
+            itrstr = ",".join(map(str, (f for f in self["_iteration_stack"])))
+            return "%s/%s[%s][%s]" % (self._run_id, self._step_name, idxstr, itrstr)
         return "%s/%s[%s]" % (self._run_id, self._step_name, idxstr)
 
     @property

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/graph.py

@@ -478,7 +478,9 @@ class FlowGraph(object):
                     cur_name = cur_node.matching_join
                 elif node_type == "split-switch":
                     all_paths = [
-                        populate_block(s, end_name) for s in cur_node.out_funcs
+                        populate_block(s, end_name)
+                        for s in cur_node.out_funcs
+                        if s != cur_name
                     ]
                     resulting_list.append(all_paths)
                     cur_name = end_name

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/lint.py

@@ -175,6 +175,8 @@ def check_for_acyclicity(graph):
 
     def check_path(node, seen):
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             if n in seen:
                 path = "->".join(seen + [n])
                 raise LintWarn(
@@ -241,6 +243,8 @@ def check_split_join_balance(graph):
         elif node.type == "split-switch":
             # For a switch, continue traversal down each path with the same stack
             for n in node.out_funcs:
+                if node.type == "split-switch" and n == node.name:
+                    continue
                 traverse(graph[n], split_stack)
             return
         elif node.type == "end":
@@ -293,6 +297,8 @@ def check_split_join_balance(graph):
             new_stack = split_stack
 
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             traverse(graph[n], new_stack)
 
     traverse(graph["start"], [])
@@ -410,3 +416,25 @@ def check_nested_foreach(graph):
         if node.type == "foreach":
             if any(graph[p].type == "foreach" for p in node.split_parents):
                 raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_ambiguous_joins(graph):
+    for node in graph:
+        if node.type == "join":
+            problematic_parents = [
+                p_name
+                for p_name in node.in_funcs
+                if graph[p_name].type == "split-switch"
+            ]
+            if problematic_parents:
+                msg = (
+                    "A conditional path cannot lead directly to a join step.\n"
+                    "In your conditional step(s) {parents}, one or more of the possible paths transition directly to the join step {join_name}.\n"
+                    "As a workaround, please introduce an intermediate, unconditional step on that specific path before joining."
+                ).format(
+                    parents=", ".join("*%s*" % p for p in problematic_parents),
+                    join_name="*%s*" % node.name,
+                )
+                raise LintWarn(msg, node.func_lineno, node.source_file)

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/metaflow_config.py

@@ -315,6 +315,8 @@ SERVICE_INTERNAL_URL = from_conf("SERVICE_INTERNAL_URL", SERVICE_URL)
 # in all Metaflow deployments. Hopefully, some day we can flip the
 # default to True.
 BATCH_EMIT_TAGS = from_conf("BATCH_EMIT_TAGS", False)
+# Default tags to add to AWS Batch jobs. These are in addition to the defaults set when BATCH_EMIT_TAGS is true.
+BATCH_DEFAULT_TAGS = from_conf("BATCH_DEFAULT_TAGS", {})
 
 ###
 # AWS Step Functions configuration

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/argo/argo_workflows.py

@@ -935,6 +935,7 @@ class ArgoWorkflows(object):
         self.conditional_nodes = set()
         self.conditional_join_nodes = set()
         self.matching_conditional_join_dict = {}
+        self.recursive_nodes = set()
 
         node_conditional_parents = {}
         node_conditional_branches = {}
@@ -957,6 +958,12 @@ class ArgoWorkflows(object):
                 )
                 node_conditional_parents[node.name] = conditional_parents
 
+                # check for recursion. this split is recursive if any of its out functions are itself.
+                if any(
+                    out_func for out_func in node.out_funcs if out_func == node.name
+                ):
+                    self.recursive_nodes.add(node.name)
+
             if conditional_parents and not node.type == "split-switch":
                 node_conditional_parents[node.name] = conditional_parents
                 conditional_branch = conditional_branch + [node.name]
@@ -1042,6 +1049,9 @@ class ArgoWorkflows(object):
     def _is_conditional_join_node(self, node):
         return node.name in self.conditional_join_nodes
 
+    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)
 
@@ -1053,6 +1063,7 @@ class ArgoWorkflows(object):
             templates=None,
             dag_tasks=None,
             parent_foreach=None,
+            seen=None,
         ):  # Returns Tuple[List[Template], List[DAGTask]]
             """ """
             # Every for-each node results in a separate subDAG and an equivalent
@@ -1062,6 +1073,8 @@ class ArgoWorkflows(object):
             # of the for-each node.
 
             # Emit if we have reached the end of the sub workflow
+            if seen is None:
+                seen = []
             if dag_tasks is None:
                 dag_tasks = []
             if templates is None:
@@ -1069,6 +1082,13 @@ class ArgoWorkflows(object):
 
             if exit_node is not None and exit_node is node.name:
                 return templates, dag_tasks
+            if node.name in seen:
+                return templates, dag_tasks
+
+            seen.append(node.name)
+
+            # helper variable for recursive conditional inputs
+            has_foreach_inputs = False
             if node.name == "start":
                 # Start node has no dependencies.
                 dag_task = DAGTask(self._sanitize(node.name)).template(
@@ -1082,9 +1102,10 @@ class ArgoWorkflows(object):
                 # vs what is a "num_parallel" based foreach (i.e. something that follows gang semantics.)
                 # A `regular` foreach is basically any arbitrary kind of foreach.
             ):
+                # helper variable for recursive conditional inputs
+                has_foreach_inputs = True
                 # Child of a foreach node needs input-paths as well as split-index
                 # This child is the first node of the sub workflow and has no dependency
-
                 parameters = [
                     Parameter("input-paths").value("{{inputs.parameters.input-paths}}"),
                     Parameter("split-index").value("{{inputs.parameters.split-index}}"),
@@ -1262,6 +1283,7 @@ class ArgoWorkflows(object):
                         templates,
                         dag_tasks,
                         parent_foreach,
+                        seen,
                     )
                 return _visit(
                     self.graph[node.matching_join],
@@ -1269,8 +1291,102 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             elif node.type == "split-switch":
+                if self._is_recursive_node(node):
+                    # we need an additional recursive template if the step is recursive
+                    # NOTE: in the recursive case, the original step is renamed in the container templates to 'recursive-<step_name>'
+                    # so that we do not have to touch the step references in the DAG.
+                    #
+                    # NOTE: The way that recursion in Argo Workflows is achieved is with the following structure:
+                    # - the usual 'example-step' template which would match example_step in flow code is renamed to 'recursive-example-step'
+                    # - templates has another template with the original task name: 'example-step'
+                    # - the template 'example-step' in turn has steps
+                    #   - 'example-step-internal' which uses the metaflow step executing template 'recursive-example-step'
+                    #   - 'example-step-recursion' which calls the parent template 'example-step' if switch-step output from 'example-step-internal' matches the condition.
+                    sanitized_name = self._sanitize(node.name)
+                    templates.append(
+                        Template(sanitized_name)
+                        .steps(
+                            [
+                                WorkflowStep()
+                                .name("%s-internal" % sanitized_name)
+                                .template("recursive-%s" % sanitized_name)
+                                .arguments(
+                                    Arguments().parameters(
+                                        [
+                                            Parameter("input-paths").value(
+                                                "{{inputs.parameters.input-paths}}"
+                                            )
+                                        ]
+                                        # Add the additional inputs required by specific node types.
+                                        # We do not need to cover joins or @parallel, as a split-switch step can not be either one of these.
+                                        + (
+                                            [
+                                                Parameter("split-index").value(
+                                                    "{{inputs.parameters.split-index}}"
+                                                )
+                                            ]
+                                            if has_foreach_inputs
+                                            else []
+                                        )
+                                    )
+                                )
+                            ]
+                        )
+                        .steps(
+                            [
+                                WorkflowStep()
+                                .name("%s-recursion" % sanitized_name)
+                                .template(sanitized_name)
+                                .when(
+                                    "{{steps.%s-internal.outputs.parameters.switch-step}}==%s"
+                                    % (sanitized_name, node.name)
+                                )
+                                .arguments(
+                                    Arguments().parameters(
+                                        [
+                                            Parameter("input-paths").value(
+                                                "argo-{{workflow.name}}/%s/{{steps.%s-internal.outputs.parameters.task-id}}"
+                                                % (node.name, sanitized_name)
+                                            )
+                                        ]
+                                        + (
+                                            [
+                                                Parameter("split-index").value(
+                                                    "{{inputs.parameters.split-index}}"
+                                                )
+                                            ]
+                                            if has_foreach_inputs
+                                            else []
+                                        )
+                                    )
+                                ),
+                            ]
+                        )
+                        .inputs(Inputs().parameters(parameters))
+                        .outputs(
+                            # NOTE: We try to read the output parameters from the recursive template call first (<step>-recursion), and the internal step second (<step>-internal).
+                            # This guarantees that we always get the output parameters of the last recursive step that executed.
+                            Outputs().parameters(
+                                [
+                                    Parameter("task-id").valueFrom(
+                                        {
+                                            "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['task-id']"
+                                            % (sanitized_name, sanitized_name)
+                                        }
+                                    ),
+                                    Parameter("switch-step").valueFrom(
+                                        {
+                                            "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['switch-step']"
+                                            % (sanitized_name, sanitized_name)
+                                        }
+                                    ),
+                                ]
+                            )
+                        )
+                    )
                 for n in node.out_funcs:
                     _visit(
                         self.graph[n],
@@ -1278,6 +1394,7 @@ class ArgoWorkflows(object):
                         templates,
                         dag_tasks,
                         parent_foreach,
+                        seen,
                     )
 
                 return _visit(
@@ -1286,6 +1403,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             # For foreach nodes generate a new sub DAGTemplate
             # We do this for "regular" foreaches (ie. `self.next(self.a, foreach=)`)
@@ -1376,6 +1494,7 @@ class ArgoWorkflows(object):
                     templates,
                     [],
                     node.name,
+                    seen,
                 )
 
                 # How do foreach's work on Argo:
@@ -1509,6 +1628,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             # For linear nodes continue traversing to the next node
             if node.type in ("linear", "join", "start"):
@@ -1518,6 +1638,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             else:
                 raise ArgoWorkflowsException(
@@ -1782,8 +1903,10 @@ class ArgoWorkflows(object):
                 # foreach-joins straight out of conditional branches are not yet supported
                 if self._is_conditional_join_node(node):
                     raise ArgoWorkflowsException(
-                        "Conditionals steps that transition directly into a join step are not currently supported. "
-                        "As a workaround, you can add a normal step after the conditional steps that transitions to a join step."
+                        "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 "
+                        "that will transition to a join."
+                        % ", ".join("*%s*" % f for f in node.in_funcs)
                     )
                 # Set aggregated input-paths for a for-each join
                 foreach_step = next(
@@ -2307,8 +2430,13 @@ class ArgoWorkflows(object):
                     )
                 )
             else:
+                template_name = self._sanitize(node.name)
+                if self._is_recursive_node(node):
+                    # The recursive template has the original step name,
+                    # this becomes a template within the recursive ones 'steps'
+                    template_name = self._sanitize("recursive-%s" % node.name)
                 yield (
-                    Template(self._sanitize(node.name))
+                    Template(template_name)
                     # Set @timeout values
                     .active_deadline_seconds(run_time_limit)
                     # Set service account
@@ -3801,6 +3929,10 @@ class WorkflowStep(object):
         self.payload["template"] = str(template)
         return self
 
+    def arguments(self, arguments):
+        self.payload["arguments"] = arguments.to_json()
+        return self
+
     def when(self, condition):
         self.payload["when"] = str(condition)
         return self

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/argo/argo_workflows_cli.py

@@ -1369,7 +1369,7 @@ def sanitize_for_argo(text):
     Sanitizes a string so it does not contain characters that are not permitted in
     Argo Workflow resource names.
     """
-    return (
+    sanitized = (
         re.compile(r"^[^A-Za-z0-9]+")
         .sub("", text)
         .replace("_", "")
@@ -1377,6 +1377,12 @@ def sanitize_for_argo(text):
         .replace("+", "")
         .lower()
     )
+    # This is added in order to get sanitized and truncated project branch names to adhere to RFC 1123 subdomain requirements
+    # f.ex. after truncation a project flow name might be project.branch-cut-short-.flowname
+    # sanitize around the . separators by removing any non-alphanumeric characters
+    sanitized = re.compile(r"[^a-z0-9]*\.[^a-z0-9]*").sub(".", sanitized)
+
+    return sanitized
 
 
 def remap_status(status):

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/aws/aws_utils.py

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

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/aws/batch/batch.py

@@ -199,6 +199,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,
@@ -344,6 +345,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(
@@ -371,6 +377,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,
@@ -414,6 +421,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,

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/aws/batch/batch_cli.py

@@ -12,6 +12,7 @@ from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from .batch import Batch, BatchKilledException
+from ..aws_utils import validate_aws_tag
 
 
 @click.group()
@@ -47,7 +48,7 @@ def _execute_cmd(func, flow_name, run_id, user, my_runs, echo):
     func(flow_name, run_id, user, echo)
 
 
-@batch.command(help="List unfinished AWS Batch tasks of this flow")
+@batch.command("list", help="List unfinished AWS Batch tasks of this flow")
 @click.option(
     "--my-runs",
     default=False,
@@ -61,7 +62,7 @@ def _execute_cmd(func, flow_name, run_id, user, my_runs, echo):
     help="List unfinished tasks corresponding to the run id.",
 )
 @click.pass_context
-def list(ctx, run_id, user, my_runs):
+def _list(ctx, run_id, user, my_runs):
     batch = Batch(ctx.obj.metadata, ctx.obj.environment)
     _execute_cmd(
         batch.list_jobs, ctx.obj.flow.name, run_id, user, my_runs, ctx.obj.echo
@@ -147,6 +148,13 @@ def kill(ctx, run_id, user, my_runs):
     help="Activate designated number of elastic fabric adapter devices. "
     "EFA driver must be installed and instance type compatible with EFA",
 )
+@click.option(
+    "--aws-batch-tag",
+    "aws_batch_tags",
+    multiple=True,
+    default=None,
+    help="AWS tags. Format: key=value, multiple allowed",
+)
 @click.option("--use-tmpfs", is_flag=True, help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-tempdir", is_flag=True, help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-size", help="tmpfs requirement for AWS Batch.")
@@ -203,6 +211,7 @@ def step(
     swappiness=None,
     inferentia=None,
     efa=None,
+    aws_batch_tags=None,
     use_tmpfs=None,
     tmpfs_tempdir=None,
     tmpfs_size=None,
@@ -213,7 +222,7 @@ def step(
     log_driver=None,
     log_options=None,
     num_parallel=None,
-    **kwargs,
+    **kwargs
 ):
     def echo(msg, stream="stderr", batch_id=None, **kwargs):
         msg = util.to_unicode(msg)
@@ -277,6 +286,14 @@ def step(
 
     env = {"METAFLOW_FLOW_FILENAME": os.path.basename(sys.argv[0])}
 
+    if aws_batch_tags is not None:
+        # We do not need to validate these again,
+        # as they come supplied by the batch decorator which already performed validation.
+        batch_tags = {}
+        for item in list(aws_batch_tags):
+            key, value = item.split("=")
+            batch_tags[key] = value
+
     env_deco = [deco for deco in node.decorators if deco.name == "environment"]
     if env_deco:
         env.update(env_deco[0].attributes["vars"])
@@ -341,6 +358,7 @@ def step(
                 host_volumes=host_volumes,
                 efs_volumes=efs_volumes,
                 use_tmpfs=use_tmpfs,
+                aws_batch_tags=batch_tags,
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/aws/batch/batch_decorator.py

@@ -10,6 +10,7 @@ from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     BATCH_CONTAINER_IMAGE,
     BATCH_CONTAINER_REGISTRY,
+    BATCH_DEFAULT_TAGS,
     BATCH_JOB_QUEUE,
     DATASTORE_LOCAL_DIR,
     ECS_FARGATE_EXECUTION_ROLE,
@@ -24,6 +25,7 @@ from ..aws_utils import (
     compute_resource_attributes,
     get_docker_registry,
     get_ec2_instance_metadata,
+    validate_aws_tag,
 )
 from .batch import BatchException
 
@@ -68,6 +70,9 @@ class BatchDecorator(StepDecorator):
         A swappiness value of 0 causes swapping not to happen unless absolutely
         necessary. A swappiness value of 100 causes pages to be swapped very
         aggressively. Accepted values are whole numbers between 0 and 100.
+    aws_batch_tags: Dict[str, str], optional, default None
+        Sets arbitrary AWS tags on the AWS Batch compute environment.
+        Set as string key-value pairs.
     use_tmpfs : bool, default False
         This enables an explicit tmpfs mount for this step. Note that tmpfs is
         not available on Fargate compute environments
@@ -114,6 +119,7 @@ class BatchDecorator(StepDecorator):
         "host_volumes": None,
         "efs_volumes": None,
         "use_tmpfs": False,
+        "aws_batch_tags": None,
         "tmpfs_tempdir": True,
         "tmpfs_size": None,
         "tmpfs_path": "/metaflow_temp",
@@ -175,6 +181,29 @@ class BatchDecorator(StepDecorator):
         if self.attributes["trainium"] is not None:
             self.attributes["inferentia"] = self.attributes["trainium"]
 
+        if not isinstance(BATCH_DEFAULT_TAGS, dict) and not all(
+            isinstance(k, str) and isinstance(v, str)
+            for k, v in BATCH_DEFAULT_TAGS.items()
+        ):
+            raise BatchException(
+                "BATCH_DEFAULT_TAGS environment variable must be Dict[str, str]"
+            )
+
+        if self.attributes["aws_batch_tags"] is not None:
+            if not isinstance(self.attributes["aws_batch_tags"], dict) and not all(
+                isinstance(k, str) and isinstance(v, str)
+                for k, v in self.attributes["aws_batch_tags"].items()
+            ):
+                raise BatchException("aws_batch_tags must be Dict[str, str]")
+        else:
+            self.attributes["aws_batch_tags"] = {}
+
+        if BATCH_DEFAULT_TAGS:
+            self.attributes["aws_batch_tags"] = {
+                **BATCH_DEFAULT_TAGS,
+                **self.attributes["aws_batch_tags"],
+            }
+
         # clean up the alias attribute so it is not passed on.
         self.attributes.pop("trainium", None)
 
@@ -207,6 +236,11 @@ class BatchDecorator(StepDecorator):
         if self.attributes["tmpfs_path"] and self.attributes["tmpfs_path"][0] != "/":
             raise BatchException("'tmpfs_path' needs to be an absolute path")
 
+        # Validate Batch tags
+        if self.attributes["aws_batch_tags"]:
+            for key, val in self.attributes["aws_batch_tags"].items():
+                validate_aws_tag(key, val)
+
     def runtime_init(self, flow, graph, package, run_id):
         # Set some more internal state.
         self.flow = flow
@@ -231,8 +265,17 @@ class BatchDecorator(StepDecorator):
             cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
-            cli_args.command_options.update(self.attributes)
+            # skip certain keys as CLI arguments
+            _skip_keys = ["aws_batch_tags"]
+            cli_args.command_options.update(
+                {k: v for k, v in self.attributes.items() if k not in _skip_keys}
+            )
             cli_args.command_options["run-time-limit"] = self.run_time_limit
+
+            # Pass the supplied AWS batch tags to the step CLI cmd
+            cli_args.command_options["aws-batch-tag"] = [
+                "%s=%s" % (k, v) for k, v in self.attributes["aws_batch_tags"].items()
+            ]
             if not R.use_r():
                 cli_args.entrypoint[0] = sys.executable
 

{ob_metaflow-2.17.3.1 → ob_metaflow-2.18.1.1}/metaflow/plugins/aws/step_functions/step_functions.py

@@ -50,6 +50,7 @@ class StepFunctions(object):
         event_logger,
         monitor,
         tags=None,
+        aws_batch_tags=None,
         namespace=None,
         username=None,
         max_workers=None,
@@ -70,6 +71,7 @@ class StepFunctions(object):
         self.event_logger = event_logger
         self.monitor = monitor
         self.tags = tags
+        self.aws_batch_tags = aws_batch_tags or {}
         self.namespace = namespace
         self.username = username
         self.max_workers = max_workers
@@ -194,6 +196,7 @@ class StepFunctions(object):
                 "on AWS Step Functions. Please "
                 "deploy your flow first." % name
             )
+
         # Dump parameters into `Parameters` input field.
         input = json.dumps({"Parameters": json.dumps(parameters)})
         # AWS Step Functions limits input to be 32KiB, but AWS Batch
@@ -852,7 +855,8 @@ class StepFunctions(object):
             # AWS_BATCH_JOB_ATTEMPT as the job counter.
             "retry_count": "$((AWS_BATCH_JOB_ATTEMPT-1))",
         }
-
+        # 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)
             .create_job(
@@ -878,6 +882,7 @@ class StepFunctions(object):
                 swappiness=resources["swappiness"],
                 efa=resources["efa"],
                 use_tmpfs=resources["use_tmpfs"],
+                aws_batch_tags=batch_tags,
                 tmpfs_tempdir=resources["tmpfs_tempdir"],
                 tmpfs_size=resources["tmpfs_size"],
                 tmpfs_path=resources["tmpfs_path"],