Repository: pytorch/torchx

Issue:
metadata support for aws_batch_scheduler
## Description
`AppDef.metadata` is currently not implemented for `aws_batch_scheduler`. This feature request suggests adding it via AWS Batch Tags.

## Motivation/Background
This allows adding metadata that is not available to the training jobs itself (unlike `env` variables). It could be used to filter/group AWs Batch Jobs but won't be pollute `env` variables for the training job.


## Detailed Proposal
Currently there's no way to pass/change AWS Batch job tags for `aws_batch_scheduler` (unlike `env` variables). We can add `AppDef.metadata` key-value pairs as tags in [here](https://github.com/pytorch/torchx/blob/13a4787e0b107cb96d97adf8f9f04e921b58257a/torchx/schedulers/aws_batch_scheduler.py#L515).
This way a custom component can pass the required key-value pairs via `AppDef.metadata` and they will be set on the AWS Batch job as tags.


## Alternatives
- Using `env` variables could result in name clash and, more importantly, leak metadata into the training process itself (e.g. the flag can be used by a dependency).
- Implementing a custom scheduler on top of the standard `aws_batch_scheduler` just to pass custom tags seems like an overkill


## Additional context/links
* https://github.com/pytorch/torchx/blob/13a4787e0b107cb96d97adf8f9f04e921b58257a/torchx/schedulers/aws_batch_scheduler.py#L515



Discussion:
Seems reasonable, would you be open to contributing a PR with this change? Feel free to assign me as the reviewer. Thanks! 
Thanks, @kiukchung I'll contribute a PR shortly.

Resolving patch:
diff --git a/torchx/schedulers/aws_batch_scheduler.py b/torchx/schedulers/aws_batch_scheduler.py
--- 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 {}),


Test patch:
diff --git a/torchx/schedulers/test/aws_batch_scheduler_test.py b/torchx/schedulers/test/aws_batch_scheduler_test.py
--- 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",
                 },
             },
         )