Improve get_status performance

This commit is contained in:
allegroai 2023-02-28 17:06:38 +02:00
parent dd5036a0a0
commit 890f4cfa77
2 changed files with 13 additions and 10 deletions

View File

@ -30,7 +30,7 @@ class TaskStopSignal(object):
try: try:
# we use internal status read, so we do not need to constantly pull the entire task object, # we use internal status read, so we do not need to constantly pull the entire task object,
# it might be large, and we want to also avoid the edit lock on it. # it might be large, and we want to also avoid the edit lock on it.
status, message = self.task._get_status() status, message, _ = self.task._get_status([self.task.id])[0]
# if we did not get a proper status, return and recheck later # if we did not get a proper status, return and recheck later
if status is None and message is None: if status is None and message is None:
return None return None

View File

@ -14,6 +14,7 @@ from operator import itemgetter
from tempfile import gettempdir from tempfile import gettempdir
from threading import Thread from threading import Thread
from typing import Optional, Any, Sequence, Callable, Mapping, Union, List, Set, Dict from typing import Optional, Any, Sequence, Callable, Mapping, Union, List, Set, Dict
from collections import namedtuple
from uuid import uuid4 from uuid import uuid4
from pathlib2 import Path from pathlib2 import Path
@ -1881,7 +1882,7 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin):
:return: str: Task status as string (TaskStatusEnum) :return: str: Task status as string (TaskStatusEnum)
""" """
status, status_message = self._get_status() status, status_message, _ = self._get_status([self.id])[0]
if self._data: if self._data:
self._data.status = status self._data.status = status
self._data.status_message = str(status_message) self._data.status_message = str(status_message)
@ -2276,19 +2277,21 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin):
self._files_server = Session.get_files_server_host() self._files_server = Session.get_files_server_host()
return self._files_server return self._files_server
def _get_status(self): @classmethod
# type: () -> (Optional[str], Optional[str]) def _get_status(cls, ids):
if self._offline_mode: # type: (List[str]) -> List[namedtuple[Optional[str], Optional[str], Optional[str]]]
return tasks.TaskStatusEnum.created, 'offline' if cls._offline_mode:
return [(tasks.TaskStatusEnum.created, "offline") for _ in ids]
Status = namedtuple("Status", ["status", "status_message", "id"])
# noinspection PyBroadException # noinspection PyBroadException
try: try:
all_tasks = self.send( all_tasks = cls._get_default_session().send(
tasks.GetAllRequest(id=[self.id], only_fields=['status', 'status_message']), tasks.GetAllRequest(id=ids, only_fields=["status", "status_message", "id"]),
).response.tasks ).response.tasks
return all_tasks[0].status, all_tasks[0].status_message return [Status(task.status, task.status_message, task.id) for task in all_tasks]
except Exception: except Exception:
return None, None return [Status(None, None, None) for _ in ids]
def _get_last_update(self): def _get_last_update(self):
# type: () -> (Optional[datetime]) # type: () -> (Optional[datetime])