step function tags option for batch create job

Netflix · Nov 15, 2024 · d0c847c · d0c847c
1 parent 57a1115
commit d0c847c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/metaflow/plugins/aws/batch/batch.py b/metaflow/plugins/aws/batch/batch.py
@@ -188,6 +188,7 @@ def create_job(
         efs_volumes=None,
         use_tmpfs=None,
         tags=None,
+        step_function_tags=None,
         tmpfs_tempdir=None,
         tmpfs_size=None,
         tmpfs_path=None,
@@ -342,6 +343,11 @@ def create_job(
                     raise BatchException("tags must be a dictionary of key-value tags.")
                 for name, value in tags.items():
                     job.tag(name, value)
+
+            if step_function_tags is not None:
+                for tag in step_function_tags:
+                    job.tag(tag['key'], tag['value'])
+
         return job
 
     def launch_job(

diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py
@@ -841,7 +841,7 @@ def _batch(self, node):
                 swappiness=resources["swappiness"],
                 efa=resources["efa"],
                 use_tmpfs=resources["use_tmpfs"],
-                tags=step_function_tags,
+                step_function_tags=step_function_tags,
                 tmpfs_tempdir=resources["tmpfs_tempdir"],
                 tmpfs_size=resources["tmpfs_size"],
                 tmpfs_path=resources["tmpfs_path"],