From 893ba48eda3ad3d7880a61a55e25f95b4e3ae036 Mon Sep 17 00:00:00 2001 From: clearml <> Date: Wed, 18 Dec 2024 17:47:51 +0200 Subject: [PATCH] Add support for queue display name --- apiserver/apimodels/queues.py | 2 + apiserver/apimodels/workers.py | 1 + apiserver/bll/queue/queue_bll.py | 2 + apiserver/bll/workers/__init__.py | 2 + apiserver/database/model/queue.py | 1 + apiserver/schema/services/queues.conf | 20 ++++++++- apiserver/schema/services/workers.conf | 4 ++ apiserver/services/queues.py | 1 + apiserver/tests/automated/test_queues.py | 56 +++++++++++++++++++++--- 9 files changed, 81 insertions(+), 8 deletions(-) diff --git a/apiserver/apimodels/queues.py b/apiserver/apimodels/queues.py index de1b69a..d9f87e1 100644 --- a/apiserver/apimodels/queues.py +++ b/apiserver/apimodels/queues.py @@ -17,6 +17,7 @@ class GetDefaultResp(Base): class CreateRequest(Base): name = StringField(required=True) + display_name = StringField() tags = ListField(items_types=[str]) system_tags = ListField(items_types=[str]) metadata = DictField(value_types=[MetadataItem]) @@ -47,6 +48,7 @@ class DeleteRequest(QueueRequest): class UpdateRequest(QueueRequest): name = StringField() + display_name = StringField() tags = ListField(items_types=[str]) system_tags = ListField(items_types=[str]) metadata = DictField(value_types=[MetadataItem]) diff --git a/apiserver/apimodels/workers.py b/apiserver/apimodels/workers.py index ba98503..d1f631d 100644 --- a/apiserver/apimodels/workers.py +++ b/apiserver/apimodels/workers.py @@ -86,6 +86,7 @@ class CurrentTaskEntry(IdNameEntry): class QueueEntry(IdNameEntry): + display_name = StringField() next_task = EmbeddedField(IdNameEntry) num_tasks = IntField() diff --git a/apiserver/bll/queue/queue_bll.py b/apiserver/bll/queue/queue_bll.py index fc5aac6..097aa44 100644 --- a/apiserver/bll/queue/queue_bll.py +++ b/apiserver/bll/queue/queue_bll.py @@ -34,6 +34,7 @@ class QueueBLL(object): def create( company_id: str, name: str, + display_name: str = None, tags: Optional[Sequence[str]] = None, system_tags: Optional[Sequence[str]] = None, metadata: Optional[dict] = None, @@ -46,6 +47,7 @@ class QueueBLL(object): company=company_id, created=now, name=name, + display_name=display_name, tags=tags or [], system_tags=system_tags or [], metadata=metadata, diff --git a/apiserver/bll/workers/__init__.py b/apiserver/bll/workers/__init__.py index 1caec6c..4eec5a0 100644 --- a/apiserver/bll/workers/__init__.py +++ b/apiserver/bll/workers/__init__.py @@ -297,6 +297,7 @@ class WorkerBLL: { "$project": { "name": 1, + "display_name": 1, "next_entry": {"$arrayElemAt": ["$entries", 0]}, "num_entries": {"$size": "$entries"}, } @@ -330,6 +331,7 @@ class WorkerBLL: if not info: continue entry.name = info.get("name", None) + entry.display_name = info.get("display_name", None) entry.num_tasks = info.get("num_entries", 0) task_id = nested_get(info, ("next_entry", "task")) if task_id: diff --git a/apiserver/database/model/queue.py b/apiserver/database/model/queue.py index aab7b49..f8d66db 100644 --- a/apiserver/database/model/queue.py +++ b/apiserver/database/model/queue.py @@ -47,6 +47,7 @@ class Queue(DbModelMixin, Document): name = StrippedStringField( required=True, unique_with="company", min_length=3, user_set_allowed=True ) + display_name = StringField(user_set_allowed=True) company = StringField(required=True, reference_field=Company) created = DateTimeField(required=True) tags = SafeSortedListField( diff --git a/apiserver/schema/services/queues.conf b/apiserver/schema/services/queues.conf index 775063e..96d5a62 100644 --- a/apiserver/schema/services/queues.conf +++ b/apiserver/schema/services/queues.conf @@ -50,6 +50,10 @@ _definitions { description: "Queue name" type: string } + display_name { + description: "Display name" + type: string + } user { description: "Associated user id" type: string @@ -324,7 +328,7 @@ create { } } "2.13": ${create."2.4"} { - metadata { + request.properties.metadata { description: "Queue metadata" type: object additionalProperties { @@ -332,6 +336,12 @@ create { } } } + "2.31": ${create."2.13"} { + request.properties.display_name { + description: "Display name" + type: string + } + } } update { "2.4" { @@ -377,7 +387,7 @@ update { } } "2.13": ${update."2.4"} { - metadata { + request.properties.metadata { description: "Queue metadata" type: object additionalProperties { @@ -385,6 +395,12 @@ update { } } } + "2.31": ${update."2.13"} { + request.properties.display_name { + description: "Display name" + type: string + } + } } delete { "2.4" { diff --git a/apiserver/schema/services/workers.conf b/apiserver/schema/services/workers.conf index 1e6d310..9eeca7f 100644 --- a/apiserver/schema/services/workers.conf +++ b/apiserver/schema/services/workers.conf @@ -194,6 +194,10 @@ _definitions { queue_entry = ${_definitions.id_name_entry} { properties { + display_name { + description: "Display name for the queue (if defined)" + type: string + } next_task { description: "Next task in the queue" "$ref": "#/definitions/id_name_entry" diff --git a/apiserver/services/queues.py b/apiserver/services/queues.py index 225a141..73c291f 100644 --- a/apiserver/services/queues.py +++ b/apiserver/services/queues.py @@ -121,6 +121,7 @@ def create(call: APICall, company_id, request: CreateRequest): queue = queue_bll.create( company_id=company_id, name=request.name, + display_name=request.display_name, tags=tags, system_tags=system_tags, metadata=Metadata.metadata_from_api(request.metadata), diff --git a/apiserver/tests/automated/test_queues.py b/apiserver/tests/automated/test_queues.py index fb402e4..9ba30aa 100644 --- a/apiserver/tests/automated/test_queues.py +++ b/apiserver/tests/automated/test_queues.py @@ -259,12 +259,19 @@ class TestQueues(TestService): def test_get_all_ex(self): queue_name = "TestTempQueue1" + queue_display_name = "Test display name" queue_tags = ["Test1", "Test2"] - queue = self._temp_queue(queue_name, tags=queue_tags) + queue = self._temp_queue(queue_name, display_name=queue_display_name, tags=queue_tags) res = self.api.queues.get_all_ex(name="TestTempQueue*").queues self.assertQueue( - res, queue_id=queue, name=queue_name, tags=queue_tags, tasks=[], workers=[] + res, + queue_id=queue, + display_name=queue_display_name, + name=queue_name, + tags=queue_tags, + tasks=[], + workers=[], ) tasks = [ @@ -279,6 +286,7 @@ class TestQueues(TestService): res, queue_id=queue, name=queue_name, + display_name=queue_display_name, tags=queue_tags, tasks=tasks, workers=workers, @@ -306,6 +314,7 @@ class TestQueues(TestService): queues: Sequence[AttrDict], queue_id: str, name: str, + display_name: str, tags: Sequence[str], tasks: Sequence[dict], workers: Sequence[dict], @@ -314,15 +323,33 @@ class TestQueues(TestService): assert queue.last_update self.assertEqualNoOrder(queue.tags, tags) self.assertEqual(queue.name, name) - self.assertQueueTasks(queue, tasks) - self.assertQueueWorkers(queue, workers) + self.assertEqual(queue.display_name, display_name) + self.assertQueueTasks(queue, tasks, name, display_name) + self.assertQueueWorkers(queue, workers, name, display_name) def assertTaskTags(self, task, system_tags): res = self.api.tasks.get_by_id(task=task) self.assertSequenceEqual(res.task.system_tags, system_tags) - def assertQueueTasks(self, queue: AttrDict, tasks: Sequence): + def assertQueueTasks( + self, + queue: AttrDict, + tasks: Sequence, + queue_name: str = None, + display_queue_name: str = None, + ): self.assertEqual([e.task for e in queue.entries], tasks) + if queue_name: + for task in tasks: + execution = self.api.tasks.get_by_id_ex( + id=[task["id"]], + only_fields=[ + "execution.queue.name", + "execution.queue.display_name", + ], + ).tasks[0].execution + self.assertEqual(execution.queue.name, queue_name) + self.assertEqual(execution.queue.display_name, display_queue_name) def assertGetNextTasks(self, queue, tasks): for task_id in tasks: @@ -330,11 +357,28 @@ class TestQueues(TestService): self.assertEqual(res.entry.task, task_id) assert not self.api.queues.get_next_task(queue=queue) - def assertQueueWorkers(self, queue: AttrDict, workers: Sequence[dict]): + def assertQueueWorkers( + self, + queue: AttrDict, + workers: Sequence[dict], + queue_name: str = None, + display_queue_name: str = None, + ): sort_key = itemgetter("name") self.assertEqual( sorted(queue.workers, key=sort_key), sorted(workers, key=sort_key) ) + if not workers: + return + + res = self.api.workers.get_all() + worker_ids = {w["key"] for w in workers} + found = [w for w in res.workers if w.key in worker_ids] + self.assertEqual(len(found), len(worker_ids)) + for worker in found: + for queue in worker.queues: + self.assertEqual(queue.name, queue_name) + self.assertEqual(queue.display_name, display_queue_name) def _temp_queue(self, queue_name, **kwargs): return self.create_temp("queues", name=queue_name, **kwargs)