ClearML/clearml-server
1
0
Fork 0
mirror of https://github.com/clearml/clearml-server synced 2025-01-31 10:56:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for specifying a specific task ID in queues.get_next_task

Browse Source
This commit is contained in:
allegroai 2022-09-29 19:27:42 +03:00
parent 0531612bf4
commit 04e7076628
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
apiserver/bll/queue/queue_bll.py
View File

@ -260,16 +260,22 @@ class QueueBLL(object):
return res
def get_next_task(self, company_id: str, queue_id: str) -> Optional[Entry]:
def get_next_task(
self, company_id: str, queue_id: str, task_id: str = None
) -> Optional[Entry]:
"""
Atomically pop and return the first task from the queue (or None)
:raise errors.bad_request.InvalidQueueId: if the queue does not exist
"""
with translate_errors_context():
query = dict(id=queue_id, company=company_id)
queue = Queue.objects(**query).modify(pop__entries=-1, upsert=False)
queue = Queue.objects(
**query, **({"entries__0__task": task_id} if task_id else {})
).modify(pop__entries=-1, upsert=False)
if not queue:
raise errors.bad_request.InvalidQueueId(**query)
if not task_id or not Queue.objects(**query).first():
raise errors.bad_request.InvalidQueueId(**query)
return
self.metrics.log_queue_metrics_to_es(company_id, queues=[queue])