List API for AWS Batch (meta-pytorch#553)

Summary:
List jobs scheduled in AWS Batch scheduler across all queues.
Might support listing by queueName in the future.

Pull Request resolved: meta-pytorch#553

Test Plan:
Unit tests
```
(venv38) ubuntu@ip-172-31-24-88:~/rsync/torchx$ torchx list -s aws_batch
torchx 2022-07-13 00:33:16 INFO     Found credentials in environment variables.
aws_batch://torchx/torchx:touch-c79fsrptv..
aws_batch://torchx/torchx:sh-qpdf3z7..
aws_batch://torchx/torchx:torchx-torchserve-vcdqjk..
.
.
```

Reviewed By: d4l3k

Differential Revision: D37807402

Pulled By: priyaramani

fbshipit-source-id: ba160f107fb6c3366d75da6cbacb1a43bb01e183

Author: priyaramani

torchx/schedulers/aws_batch_scheduler.py

@@ -577,7 +577,33 @@ def log_iter(
             return iterator
 
     def list(self) -> List[str]:
-        raise NotImplementedError()
+        # TODO: get queue name input instead of iterating over all queues?
+        resp = self._client.describe_job_queues()
+        queue_names = [queue["jobQueueName"] for queue in resp["jobQueues"]]
+        all_app_ids = []
+        for qn in queue_names:
+            all_app_ids += self._list_by_queue(qn)
+        return all_app_ids
+
+    def _list_by_queue(self, queue_name: str) -> List[str]:
+        # 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"
+        jobs = self._client.list_jobs(
+            jobQueue=queue_name,
+            filters=[
+                {
+                    "name": "AFTER_CREATED_AT",
+                    "values": [
+                        milli_seconds_after_epoch,
+                    ],
+                },
+            ],
+        )["jobSummaryList"]
+        app_ids = [f"{queue_name}:{job['jobName']}" for job in jobs]
+        return app_ids
 
     def _stream_events(
         self,

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -514,6 +514,17 @@ def _mock_scheduler(self) -> AWSBatchScheduler:
             ],
         }
 
+        scheduler._client.describe_job_queues.return_value = {
+            "ResponseMetadata": {},
+            "jobQueues": [
+                {
+                    "jobQueueName": "torchx",
+                    "jobQueueArn": "arn:aws:batch:test-region:4000005:job-queue/torchx",
+                    "state": "ENABLED",
+                },
+            ],
+        }
+
         return scheduler
 
     @mock_rand()
@@ -559,6 +570,12 @@ def test_describe(self) -> None:
             ),
         )
 
+    def test_list(self) -> None:
+        scheduler = self._mock_scheduler()
+        expected_app_ids = ["torchx:echo-v1r560pmwn5t3c"]
+        app_ids = scheduler.list()
+        self.assertEqual(app_ids, expected_app_ids)
+
     def test_log_iter(self) -> None:
         scheduler = self._mock_scheduler()
         logs = scheduler.log_iter("testqueue:app-name-42", "echo", k=1, regex="foo.*")