step functions create cli aws_tags

Netflix · Nov 18, 2024 · e178c81 · e178c81
1 parent d0c847c
commit e178c81
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py
@@ -48,6 +48,7 @@ def __init__(
         event_logger,
         monitor,
         tags=None,
+        aws_tags=None,
         namespace=None,
         username=None,
         max_workers=None,
@@ -67,6 +68,7 @@ def __init__(
         self.event_logger = event_logger
         self.monitor = monitor
         self.tags = tags
+        self.aws_tags = aws_tags
         self.namespace = namespace
         self.username = username
         self.max_workers = max_workers
@@ -933,6 +935,8 @@ def _step_cli(self, node, paths, code_package_url, user_code_retries):
             # Assign tags to run objects.
             if self.tags:
                 params.extend("--tag %s" % tag for tag in self.tags)
+            if self.aws_tags:
+                params.extend("--aws-tags %s" % aws_tag for aws_tag in self.aws_tags)
 
             # If the start step gets retried, we must be careful not to
             # regenerate multiple parameters tasks. Hence, we check first if
@@ -978,6 +982,8 @@ def _step_cli(self, node, paths, code_package_url, user_code_retries):
             step.append("--split-index $METAFLOW_SPLIT_INDEX")
         if self.tags:
             step.extend("--tag %s" % tag for tag in self.tags)
+        if self.aws_tags:
+            step.extend("--aws-tag %s" % aws_tag for aws_tag in self.aws_tags)
         if self.namespace is not None:
             step.append("--namespace=%s" % self.namespace)
         cmds.append(" ".join(entrypoint + top_level + step))

diff --git a/metaflow/plugins/aws/step_functions/step_functions_cli.py b/metaflow/plugins/aws/step_functions/step_functions_cli.py
@@ -97,6 +97,11 @@ def step_functions(obj, name=None):
     "with the given tag. You can specify this option multiple "
     "times to attach multiple tags.",
 )
+@click.option(
+    "--aws-tags",
+    multiple=True,
+    default=None,
+    help="AWS tags.")
 @click.option(
     "--namespace",
     "user_namespace",
@@ -143,6 +148,7 @@ def step_functions(obj, name=None):
 def create(
     obj,
     tags=None,
+    aws_tags=None,
     user_namespace=None,
     only_json=False,
     authorize=None,
@@ -196,6 +202,7 @@ def create(
         token,
         obj.state_machine_name,
         tags,
+        aws_tags,
         user_namespace,
         max_workers,
         workflow_timeout,
@@ -315,6 +322,7 @@ def make_flow(
     token,
     name,
     tags,
+    aws_tags,
     namespace,
     max_workers,
     workflow_timeout,
@@ -350,6 +358,7 @@ def make_flow(
         obj.event_logger,
         obj.monitor,
         tags=tags,
+        aws_tags=aws_tags,
         namespace=namespace,
         max_workers=max_workers,
         username=get_username(),