fixes #775 (#777)

passing AppDef metadata to job tags for AWS Batch scheduler Co-authored-by: Alexander Jipa <[email protected]>
pytorch · Oct 18, 2023 · 9421dfa · 9421dfa
1 parent 6153315
commit 9421dfa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/torchx/schedulers/aws_batch_scheduler.py b/torchx/schedulers/aws_batch_scheduler.py
@@ -516,6 +516,7 @@ def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJ
                     TAG_TORCHX_VER: torchx.__version__,
                     TAG_TORCHX_APPNAME: app.name,
                     TAG_TORCHX_USER: cfg.get("user"),
+                    **app.metadata,
                 },
             },
             **({"schedulingPriority": priority} if share_id is not None else {}),

diff --git a/torchx/schedulers/test/aws_batch_scheduler_test.py b/torchx/schedulers/test/aws_batch_scheduler_test.py
@@ -54,7 +54,7 @@ def _test_app() -> specs.AppDef:
         ],
     )
 
-    return specs.AppDef("test", roles=[trainer_role])
+    return specs.AppDef("test", roles=[trainer_role], metadata={"FIZZ": "buzz"})
 
 
 @contextmanager
@@ -144,6 +144,7 @@ def test_submit_dryrun_tags(self, _) -> None:
                 "torchx.pytorch.org/version": torchx.__version__,
                 "torchx.pytorch.org/app-name": "test",
                 "torchx.pytorch.org/user": "testuser",
+                "FIZZ": "buzz",
             },
             info.request.job_def["tags"],
         )
@@ -231,6 +232,7 @@ def test_submit_dryrun(self) -> None:
                     "torchx.pytorch.org/version": torchx.__version__,
                     "torchx.pytorch.org/app-name": "test",
                     "torchx.pytorch.org/user": "testuser",
+                    "FIZZ": "buzz",
                 },
             },
         )