From a4ae806ba879391578f128dfcf89800f6c10f7ef Mon Sep 17 00:00:00 2001 From: Alexander Jipa Date: Tue, 17 Oct 2023 13:56:53 -0400 Subject: [PATCH] fixes #775 passing AppDef metadata to job tags for AWS Batch scheduler --- torchx/schedulers/aws_batch_scheduler.py | 1 + torchx/schedulers/test/aws_batch_scheduler_test.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/torchx/schedulers/aws_batch_scheduler.py b/torchx/schedulers/aws_batch_scheduler.py index 80d57815d..1b88b8835 100644 --- 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 index 44d490485..c23505533 100644 --- 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", }, }, )