(torchx/aws_batch) Add user tag and runopt. Make list API only return torchx jobs by filtering on job tag (#672)

Author: kiukchung

torchx/schedulers/aws_batch_scheduler.py

@@ -34,7 +34,7 @@
 https://docs.aws.amazon.com/AmazonECR/latest/userguide/getting-started-cli.html#cli-create-repository
 for how to create a image repository.
 """
-
+import getpass
 import re
 import threading
 from dataclasses import dataclass
@@ -79,6 +79,11 @@
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
 from typing_extensions import TypedDict
 
+TAG_TORCHX_VER = "torchx.pytorch.org/version"
+TAG_TORCHX_APPNAME = "torchx.pytorch.org/app-name"
+TAG_TORCHX_USER = "torchx.pytorch.org/user"
+
+
 if TYPE_CHECKING:
     from docker import DockerClient
 
@@ -244,6 +249,7 @@ def _local_session() -> "boto3.session.Session":
 
 class AWSBatchOpts(TypedDict, total=False):
     queue: str
+    user: str
     image_repo: Optional[str]
     share_id: Optional[str]
     priority: Optional[int]
@@ -417,8 +423,9 @@ def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJ
                     ],
                 },
                 "tags": {
-                    "torchx.pytorch.org/version": torchx.__version__,
-                    "torchx.pytorch.org/app-name": app.name,
+                    TAG_TORCHX_VER: torchx.__version__,
+                    TAG_TORCHX_APPNAME: app.name,
+                    TAG_TORCHX_USER: cfg.get("user"),
                 },
             },
             **(
@@ -455,6 +462,12 @@ def _cancel_existing(self, app_id: str) -> None:
     def _run_opts(self) -> runopts:
         opts = runopts()
         opts.add("queue", type_=str, help="queue to schedule job in", required=True)
+        opts.add(
+            "user",
+            type_=str,
+            default=getpass.getuser(),
+            help="The username to tag the job with. `getpass.getuser()` if not specified.",
+        )
         opts.add(
             "share_id",
             type_=str,
@@ -582,36 +595,50 @@ def list(self) -> List[ListAppResponse]:
         for resp in self._client.get_paginator("describe_job_queues").paginate():
             queue_names = [queue["jobQueueName"] for queue in resp["jobQueues"]]
             for qn in queue_names:
-                apps_in_queue = self._list_by_queue(qn)
-                all_apps += [
-                    ListAppResponse(
-                        app_id=f"{qn}:{app['jobName']}",
-                        state=JOB_STATE[app["status"]],
-                    )
-                    for app in apps_in_queue
-                ]
+                all_apps.extend(self._list_by_queue(qn))
         return all_apps
 
-    def _list_by_queue(self, queue_name: str) -> List[Dict[str, Any]]:
-        # By default only running jobs are listed by batch/boto client's list_jobs API
+    def _list_by_queue(self, queue_name: str) -> List[ListAppResponse]:
+        # By default, only running jobs are listed by batch/boto client's list_jobs API
         # When 'filters' parameter is specified, jobs with all statuses are listed
         # So use AFTER_CREATED_AT filter to list jobs in all statuses
         # milli_seconds_after_epoch can later be used to list jobs by timeframe
-        milli_seconds_after_epoch = "1"
-        job_summary_list = []
+        MS_AFTER_EPOCH = "1"
+        EVERY_STATUS = {"name": "AFTER_CREATED_AT", "values": [MS_AFTER_EPOCH]}
+
+        jobs = []
         for resp in self._client.get_paginator("list_jobs").paginate(
             jobQueue=queue_name,
-            filters=[
-                {
-                    "name": "AFTER_CREATED_AT",
-                    "values": [
-                        milli_seconds_after_epoch,
-                    ],
-                },
-            ],
+            filters=[EVERY_STATUS],
+            # describe-jobs API can take up to 100 jobIds
+            PaginationConfig={"MaxItems": 100},
         ):
-            job_summary_list.extend(resp["jobSummaryList"])
-        return job_summary_list
+
+            # torchx.pytorch.org/version tag is used to filter torchx jobs
+            # list_jobs() API only returns a job summary which does not include the job's tag
+            # so we need to call the describe_jobs API.
+            # Ideally batch lets us pass tags as a filter to list_jobs API
+            # but this is currently not supported
+            job_ids = [js["jobId"] for js in resp["jobSummaryList"]]
+            for jobdesc in self._get_torchx_submitted_jobs(job_ids):
+                jobs.append(
+                    ListAppResponse(
+                        app_id=f"{queue_name}:{jobdesc['jobName']}",
+                        state=JOB_STATE[jobdesc["status"]],
+                    )
+                )
+
+        return jobs
+
+    def _get_torchx_submitted_jobs(self, job_ids: List[str]) -> List[Dict[str, Any]]:
+        if not job_ids:
+            return []
+
+        return [
+            jobdesc
+            for jobdesc in self._client.describe_jobs(jobs=job_ids)["jobs"]
+            if TAG_TORCHX_VER in jobdesc["tags"]
+        ]
 
     def _stream_events(
         self,

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -94,41 +94,50 @@ def test_create_scheduler(self) -> None:
         self.assertIsInstance(scheduler, AWSBatchScheduler)
 
     def test_submit_dryrun_with_share_id(self) -> None:
-        scheduler = create_scheduler("test")
         app = _test_app()
         cfg = AWSBatchOpts({"queue": "testqueue", "share_id": "fooshare"})
-        info = scheduler._submit_dryrun(app, cfg)
+        info = create_scheduler("test").submit_dryrun(app, cfg)
 
         req = info.request
         job_def = req.job_def
         self.assertEqual(req.share_id, "fooshare")
         # must be set for jobs submitted to a queue with scheduling policy
         self.assertEqual(job_def["schedulingPriority"], 0)
 
-    def test_submit_dryrun_with_priority(self) -> None:
-        scheduler = create_scheduler("test")
-        app = _test_app()
-
+    def test_submit_dryrun_with_priority_but_not_share_id(self) -> None:
         with self.assertRaisesRegex(ValueError, "config value.*priority.*share_id"):
             cfg = AWSBatchOpts({"queue": "testqueue", "priority": 42})
-            info = scheduler._submit_dryrun(app, cfg)
+            create_scheduler("test").submit_dryrun(_test_app(), cfg)
 
-        cfg = AWSBatchOpts(
-            {"queue": "testqueue", "share_id": "fooshare", "priority": 42}
-        )
-        info = scheduler._submit_dryrun(app, cfg)
+    def test_submit_dryrun_with_priority(self) -> None:
+        cfg = AWSBatchOpts({"queue": "testqueue", "share_id": "foo", "priority": 42})
+        info = create_scheduler("test").submit_dryrun(_test_app(), cfg)
 
         req = info.request
         job_def = req.job_def
-        self.assertEqual(req.share_id, "fooshare")
+        self.assertEqual(req.share_id, "foo")
         self.assertEqual(job_def["schedulingPriority"], 42)
 
+    @patch(
+        "torchx.schedulers.aws_batch_scheduler.getpass.getuser", return_value="testuser"
+    )
+    def test_submit_dryrun_tags(self, _) -> None:
+        # intentionally not specifying user in cfg to test default
+        cfg = AWSBatchOpts({"queue": "ignored_in_test"})
+        info = create_scheduler("test").submit_dryrun(_test_app(), cfg)
+        self.assertEqual(
+            {
+                "torchx.pytorch.org/version": torchx.__version__,
+                "torchx.pytorch.org/app-name": "test",
+                "torchx.pytorch.org/user": "testuser",
+            },
+            info.request.job_def["tags"],
+        )
+
     @mock_rand()
     def test_submit_dryrun(self) -> None:
-        scheduler = create_scheduler("test")
-        app = _test_app()
-        cfg = AWSBatchOpts({"queue": "testqueue"})
-        info = scheduler._submit_dryrun(app, cfg)
+        cfg = AWSBatchOpts({"queue": "testqueue", "user": "testuser"})
+        info = create_scheduler("test").submit_dryrun(_test_app(), cfg)
 
         req = info.request
         self.assertEqual(req.share_id, None)
@@ -248,6 +257,7 @@ def test_submit_dryrun(self) -> None:
                 "tags": {
                     "torchx.pytorch.org/version": torchx.__version__,
                     "torchx.pytorch.org/app-name": "test",
+                    "torchx.pytorch.org/user": "testuser",
                 },
             },
         )
@@ -440,13 +450,13 @@ def _mock_scheduler(self) -> AWSBatchScheduler:
                         {
                             "jobArn": "arn:aws:batch:us-west-2:495572122715:job/6afc27d7-3559-43ca-89fd-1007b6bf2546",
                             "jobId": "6afc27d7-3559-43ca-89fd-1007b6bf2546",
-                            "jobName": "echo-v1r560pmwn5t3c",
+                            "jobName": "app-name-42",
                             "createdAt": 1643949940162,
                             "status": "SUCCEEDED",
                             "stoppedAt": 1643950324125,
                             "container": {"exitCode": 0},
                             "nodeProperties": {"numNodes": 2},
-                            "jobDefinition": "arn:aws:batch:us-west-2:495572122715:job-definition/echo-v1r560pmwn5t3c:1",
+                            "jobDefinition": "arn:aws:batch:us-west-2:495572122715:job-definition/app-name-42:1",
                         }
                     ]
                 }
@@ -568,11 +578,7 @@ def _mock_scheduler(self) -> AWSBatchScheduler:
     def test_submit(self) -> None:
         scheduler = self._mock_scheduler()
         app = _test_app()
-        cfg = AWSBatchOpts(
-            {
-                "queue": "testqueue",
-            }
-        )
+        cfg = AWSBatchOpts({"queue": "testqueue"})
 
         info = scheduler._submit_dryrun(app, cfg)
         id = scheduler.schedule(info)
@@ -610,34 +616,35 @@ def test_describe(self) -> None:
     def test_list(self) -> None:
         scheduler = self._mock_scheduler()
         expected_apps = [
-            ListAppResponse(
-                app_id="torchx:echo-v1r560pmwn5t3c", state=AppState.SUCCEEDED
-            )
+            ListAppResponse(app_id="torchx:app-name-42", state=AppState.SUCCEEDED)
         ]
         apps = scheduler.list()
-        self.assertEqual(apps, expected_apps)
+        self.assertEqual(expected_apps, apps)
+
+    def test_list_no_jobs(self) -> None:
+        scheduler = AWSBatchScheduler("test", client=MagicMock())
+        scheduler._client.get_paginator.side_effect = MockPaginator(
+            describe_job_queues=[
+                {
+                    "jobQueues": [
+                        {"jobQueueName": "torchx", "state": "ENABLED"},
+                    ],
+                }
+            ],
+            list_jobs=[{"jobSummaryList": []}],
+        )
+
+        self.assertEqual([], scheduler.list())
 
     def test_log_iter(self) -> None:
         scheduler = self._mock_scheduler()
         logs = scheduler.log_iter("testqueue:app-name-42", "echo", k=1, regex="foo.*")
-        self.assertEqual(
-            list(logs),
-            [
-                "foo\n",
-                "foobar\n",
-            ],
-        )
+        self.assertEqual(list(logs), ["foo\n", "foobar\n"])
 
     def test_log_iter_running_job(self) -> None:
         scheduler = self._mock_scheduler_running_job()
         logs = scheduler.log_iter("testqueue:app-name-42", "echo", k=1, regex="foo.*")
-        self.assertEqual(
-            [
-                "foo\n",
-                "foobar\n",
-            ],
-            list(logs),
-        )
+        self.assertEqual(["foo\n", "foobar\n"], list(logs))
 
     def test_local_session(self) -> None:
         a: object = _local_session()

torchx/specs/api.py

@@ -794,7 +794,7 @@ def resolve(self, cfg: Mapping[str, CfgVal]) -> Dict[str, CfgVal]:
 
     def cfg_from_str(self, cfg_str: str) -> Dict[str, CfgVal]:
         """
-        Parses scheduler ``runcfg`` from a string literal and returns
+        Parses scheduler ``cfg`` from a string literal and returns
         a cfg map where the cfg values have been cast into the appropriate
         types as specified by this runopts object. Unknown keys are ignored
         and not returned in the resulting map.