Merge branch 'master' into sendToGitHub/fix-scheduling-reports-auth

Author: mauriceyap

client/python/armada_client/asyncio_client.py

@@ -105,6 +105,7 @@ def __init__(
         event_timeout: timedelta = timedelta(minutes=15),
     ) -> None:
         self.submit_stub = submit_pb2_grpc.SubmitStub(channel)
+        self.queue_stub = submit_pb2_grpc.QueueServiceStub(channel)
         self.event_stub = event_pb2_grpc.EventStub(channel)
         self.job_stub = job_pb2_grpc.JobsStub(channel)
         self.event_timeout = event_timeout
@@ -390,7 +391,7 @@ async def create_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
         :param queue: A queue to create.
         """
 
-        response = await self.submit_stub.CreateQueue(queue)
+        response = await self.queue_stub.CreateQueue(queue)
         return response
 
     async def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
@@ -400,7 +401,7 @@ async def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
         :param queue: A queue to update.
         """
 
-        response = await self.submit_stub.UpdateQueue(queue)
+        response = await self.queue_stub.UpdateQueue(queue)
         return response
 
     async def create_queues(
@@ -413,7 +414,7 @@ async def create_queues(
         """
 
         queue_list = submit_pb2.QueueList(queues=queues)
-        response = await self.submit_stub.CreateQueues(queue_list)
+        response = await self.queue_stub.CreateQueues(queue_list)
         return response
 
     async def update_queues(
@@ -426,7 +427,7 @@ async def update_queues(
         """
 
         queue_list = submit_pb2.QueueList(queues=queues)
-        response = await self.submit_stub.UpdateQueues(queue_list)
+        response = await self.queue_stub.UpdateQueues(queue_list)
         return response
 
     async def delete_queue(self, name: str) -> None:
@@ -438,7 +439,7 @@ async def delete_queue(self, name: str) -> None:
         :return: None
         """
         request = submit_pb2.QueueDeleteRequest(name=name)
-        await self.submit_stub.DeleteQueue(request)
+        await self.queue_stub.DeleteQueue(request)
 
     async def get_queue(self, name: str) -> submit_pb2.Queue:
         """Get the queue by name.
@@ -449,9 +450,23 @@ async def get_queue(self, name: str) -> submit_pb2.Queue:
         :return: A queue object. See the api definition.
         """
         request = submit_pb2.QueueGetRequest(name=name)
-        response = await self.submit_stub.GetQueue(request)
+        response = await self.queue_stub.GetQueue(request)
         return response
 
+    async def get_queues(self) -> list:
+        """Retrieves all queues
+        :return: List containing all queues.
+        """
+        queues = []
+        request = submit_pb2.StreamingQueueGetRequest()
+        async for message in self.queue_stub.GetQueues(request):
+            event_type = message.WhichOneof("event")
+            if event_type == "queue":
+                queues.append(message.queue)
+            elif event_type == "end":
+                break
+        return queues
+
     @staticmethod
     def unwatch_events(event_stream) -> None:
         """Closes gRPC event streams

client/python/armada_client/client.py

@@ -103,8 +103,9 @@ class ArmadaClient:
     def __init__(self, channel, event_timeout: timedelta = timedelta(minutes=15)):
         self.submit_stub = submit_pb2_grpc.SubmitStub(channel)
         self.event_stub = event_pb2_grpc.EventStub(channel)
-        self.event_timeout = event_timeout
         self.job_stub = job_pb2_grpc.JobsStub(channel)
+        self.queue_stub = submit_pb2_grpc.QueueServiceStub(channel)
+        self.event_timeout = event_timeout
 
     def get_job_events_stream(
         self,
@@ -290,6 +291,7 @@ def cancel_jobset(
     ) -> empty_pb2.Empty:
         """Cancel jobs in a given queue.
 
+
         Uses the CancelJobSet RPC to cancel jobs.
         A filter is used to only cancel jobs in certain states.
 
@@ -378,7 +380,7 @@ def create_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
         :param queue: A queue to create.
         """
 
-        response = self.submit_stub.CreateQueue(queue)
+        response = self.queue_stub.CreateQueue(queue)
         return response
 
     def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
@@ -388,7 +390,7 @@ def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
         :param queue: A queue to update.
         """
 
-        response = self.submit_stub.UpdateQueue(queue)
+        response = self.queue_stub.UpdateQueue(queue)
         return response
 
     def create_queues(
@@ -401,7 +403,7 @@ def create_queues(
         """
 
         queue_list = submit_pb2.QueueList(queues=queues)
-        response = self.submit_stub.CreateQueues(queue_list)
+        response = self.queue_stub.CreateQueues(queue_list)
         return response
 
     def update_queues(
@@ -414,7 +416,7 @@ def update_queues(
         """
 
         queue_list = submit_pb2.QueueList(queues=queues)
-        response = self.submit_stub.UpdateQueues(queue_list)
+        response = self.queue_stub.UpdateQueues(queue_list)
         return response
 
     def delete_queue(self, name: str) -> None:
@@ -426,7 +428,7 @@ def delete_queue(self, name: str) -> None:
         :return: None
         """
         request = submit_pb2.QueueDeleteRequest(name=name)
-        self.submit_stub.DeleteQueue(request)
+        self.queue_stub.DeleteQueue(request)
 
     def get_queue(self, name: str) -> submit_pb2.Queue:
         """Get the queue by name.
@@ -437,9 +439,27 @@ def get_queue(self, name: str) -> submit_pb2.Queue:
         :return: A queue object. See the api definition.
         """
         request = submit_pb2.QueueGetRequest(name=name)
-        response = self.submit_stub.GetQueue(request)
+        response = self.queue_stub.GetQueue(request)
         return response
 
+    def get_queues(self) -> List[submit_pb2.Queue]:
+        """Get all queues.
+
+        Uses the GetQueues RPC to get the queues.
+
+        :return: list containing all queues
+        """
+        queues = []
+
+        request = submit_pb2.StreamingQueueGetRequest()
+
+        for message in self.queue_stub.GetQueues(request):
+            if message.HasField("queue"):
+                queues.append(message.queue)
+            elif message.HasField("end"):
+                break
+        return queues
+
     @staticmethod
     def unwatch_events(event_stream) -> None:
         """Closes gRPC event streams

client/python/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "armada_client"
-version = "0.4.10"
+version = "0.4.11"
 description = "Armada gRPC API python client"
 readme = "README.md"
 requires-python = ">=3.9"

client/python/tests/unit/server_mock.py

@@ -15,7 +15,7 @@
 from armada_client.armada.submit_pb2 import JobState
 
 
-class SubmitService(submit_pb2_grpc.SubmitServicer):
+class QueueService(submit_pb2_grpc.QueueServiceServicer):
     def CreateQueue(self, request, context):
         return empty_pb2.Empty()
 
@@ -25,6 +25,40 @@ def DeleteQueue(self, request, context):
     def GetQueue(self, request, context):
         return submit_pb2.Queue(name=request.name)
 
+    def GetQueues(self, request, context):
+        queue_names = ["test_queue1", "test_queue2", "test_queue3"]
+        for name in queue_names:
+            queue_message = submit_pb2.StreamingQueueMessage(
+                queue=submit_pb2.Queue(name=name)
+            )
+            yield queue_message
+
+        yield submit_pb2.StreamingQueueMessage(end=submit_pb2.EndMarker())
+
+    def GetQueueInfo(self, request, context):
+        return submit_pb2.QueueInfo(name=request.name)
+
+    def CreateQueues(self, request, context):
+        return submit_pb2.BatchQueueCreateResponse(
+            failed_queues=[
+                submit_pb2.QueueCreateResponse(queue=submit_pb2.Queue(name=queue.name))
+                for queue in request.queues
+            ]
+        )
+
+    def UpdateQueues(self, request, context):
+        return submit_pb2.BatchQueueUpdateResponse(
+            failed_queues=[
+                submit_pb2.QueueUpdateResponse(queue=submit_pb2.Queue(name=queue.name))
+                for queue in request.queues
+            ]
+        )
+
+    def UpdateQueue(self, request, context):
+        return empty_pb2.Empty()
+
+
+class SubmitService(submit_pb2_grpc.SubmitServicer):
     def SubmitJobs(self, request, context):
         # read job_ids from request.job_request_items
         job_ids = [f"job-{i}" for i in range(1, len(request.job_request_items) + 1)]
@@ -35,9 +69,6 @@ def SubmitJobs(self, request, context):
 
         return submit_pb2.JobSubmitResponse(job_response_items=job_response_items)
 
-    def GetQueueInfo(self, request, context):
-        return submit_pb2.QueueInfo(name=request.name)
-
     def CancelJobs(self, request, context):
         return submit_pb2.CancellationResult(
             cancelled_ids=["job-1"],
@@ -72,25 +103,6 @@ def ReprioritizeJobs(self, request, context):
 
         return submit_pb2.JobReprioritizeResponse(reprioritization_results=results)
 
-    def UpdateQueue(self, request, context):
-        return empty_pb2.Empty()
-
-    def CreateQueues(self, request, context):
-        return submit_pb2.BatchQueueCreateResponse(
-            failed_queues=[
-                submit_pb2.QueueCreateResponse(queue=submit_pb2.Queue(name=queue.name))
-                for queue in request.queues
-            ]
-        )
-
-    def UpdateQueues(self, request, context):
-        return submit_pb2.BatchQueueUpdateResponse(
-            failed_queues=[
-                submit_pb2.QueueUpdateResponse(queue=submit_pb2.Queue(name=queue.name))
-                for queue in request.queues
-            ]
-        )
-
     def Health(self, request, context):
         return health_pb2.HealthCheckResponse(
             status=health_pb2.HealthCheckResponse.SERVING

client/python/tests/unit/test_asyncio_client.py

@@ -6,7 +6,7 @@
 
 from armada_client.typings import JobState
 from armada_client.armada.job_pb2 import JobRunState
-from server_mock import EventService, SubmitService, QueryAPIService
+from server_mock import EventService, SubmitService, QueueService, QueryAPIService
 
 from armada_client.armada import (
     event_pb2_grpc,
@@ -28,6 +28,7 @@
 def server_mock():
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
     submit_pb2_grpc.add_SubmitServicer_to_server(SubmitService(), server)
+    submit_pb2_grpc.add_QueueServiceServicer_to_server(QueueService(), server)
     event_pb2_grpc.add_EventServicer_to_server(EventService(), server)
     job_pb2_grpc.add_JobsServicer_to_server(QueryAPIService(), server)
     server.add_insecure_port("[::]:50051")
@@ -175,6 +176,13 @@ async def test_get_queue(aio_client):
     assert queue.name == "test"
 
 
+@pytest.mark.asyncio
+async def test_get_queues(aio_client):
+    queues = await aio_client.get_queues()
+    queue_names = [q.name for q in queues]
+    assert queue_names == ["test_queue1", "test_queue2", "test_queue3"]
+
+
 @pytest.mark.asyncio
 async def test_delete_queue(aio_client):
     await aio_client.delete_queue("test")

client/python/tests/unit/test_client.py

@@ -5,7 +5,7 @@
 
 from armada_client.typings import JobState
 from armada_client.armada.job_pb2 import JobRunState
-from server_mock import EventService, SubmitService, QueryAPIService
+from server_mock import EventService, SubmitService, QueryAPIService, QueueService
 
 from armada_client.armada import (
     event_pb2_grpc,
@@ -27,6 +27,7 @@
 def server_mock():
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
     submit_pb2_grpc.add_SubmitServicer_to_server(SubmitService(), server)
+    submit_pb2_grpc.add_QueueServiceServicer_to_server(QueueService(), server)
     event_pb2_grpc.add_EventServicer_to_server(EventService(), server)
     job_pb2_grpc.add_JobsServicer_to_server(QueryAPIService(), server)
     server.add_insecure_port("[::]:50051")
@@ -165,6 +166,12 @@ def test_get_queue():
     assert tester.get_queue("test").name == "test"
 
 
+def test_get_queues():
+    queues = tester.get_queues()
+    queue_names = [q.name for q in queues]
+    assert queue_names == ["test_queue1", "test_queue2", "test_queue3"]
+
+
 def test_delete_queue():
     tester.delete_queue("test")
 

…r/templates/scheduler-pruner-secret.yaml …emplates/scheduler-pruner-configmap.yamldeployment/scheduler/templates/scheduler-pruner-secret.yaml renamed to deployment/scheduler/templates/scheduler-pruner-configmap.yaml deployment/scheduler/templates/scheduler-pruner-secret.yaml renamed to deployment/scheduler/templates/scheduler-pruner-configmap.yaml

@@ -1,13 +1,12 @@
 apiVersion: v1
-kind: Secret
+kind: ConfigMap
 metadata:
   name: {{ include "armada-scheduler-pruner.config.name" . }}
   namespace: {{ .Release.Namespace }}
   labels:
     {{- include "armada-scheduler-pruner.labels.all" . | nindent 4 }}
-type: Opaque
 data:
   {{ include "armada-scheduler-pruner.config.filename" . }}: |
 {{- if .Values.pruner.applicationConfig }}
-{{ toYaml .Values.pruner.applicationConfig | b64enc | indent 4 }}
+{{ toYaml .Values.pruner.applicationConfig | indent 4 }}
 {{- end }}

docs/python_armada_client.md

@@ -447,6 +447,24 @@ Uses the GetQueue RPC to get the queue.
 
 
 
+#### get_queues()
+Get all queues.
+
+Uses the GetQueues RPC to get the queues.
+
+
+* **Returns**
+
+    list containing all queues
+
+
+
+* **Return type**
+
+    *List*[armada.submit_pb2.Queue]
+
+
+
 #### preempt_jobs(queue, job_set_id, job_id)
 Preempt jobs in a given queue.
 

internal/scheduler/database/migrations/019_add_runs_executor_index.sql

@@ -0,0 +1 @@
+CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_runs_filtered ON runs (executor, succeeded, failed, cancelled, serial);