Fix task cleanup

This commit is contained in:
allegroai 2022-07-08 17:50:49 +03:00
parent 0f6c4e75b7
commit ec0436d0da

View File

@ -8,7 +8,6 @@ from boltons.iterutils import partition
from apiserver.apierrors import errors from apiserver.apierrors import errors
from apiserver.bll.event import EventBLL from apiserver.bll.event import EventBLL
from apiserver.bll.event.event_bll import PlotFields from apiserver.bll.event.event_bll import PlotFields
from apiserver.bll.event.event_common import EventType
from apiserver.bll.task.utils import deleted_prefix from apiserver.bll.task.utils import deleted_prefix
from apiserver.database.model.model import Model from apiserver.database.model.model import Model
from apiserver.database.model.task.task import Task, TaskStatus, ArtifactModes from apiserver.database.model.task.task import Task, TaskStatus, ArtifactModes
@ -80,32 +79,18 @@ def collect_debug_image_urls(company: str, task: str) -> Set[str]:
Return the set of unique image urls Return the set of unique image urls
Uses DebugImagesIterator to make sure that we do not retrieve recycled urls Uses DebugImagesIterator to make sure that we do not retrieve recycled urls
""" """
metrics = event_bll.get_metrics_and_variants( after_key = None
company_id=company, task_id=task, event_type=EventType.metrics_image
)
if not metrics:
return set()
task_metrics = {task: {m: [] for m in metrics}}
scroll_id = None
urls = set() urls = set()
while True: while True:
res = event_bll.debug_images_iterator.get_task_events( res, after_key = event_bll.get_debug_image_urls(
company_id=company, company_id=company,
task_metrics=task_metrics, task_id=task,
iter_count=10, after_key=after_key,
state_id=scroll_id,
) )
if not res.metric_events or not any( urls.update(res)
iterations for _, iterations in res.metric_events if not after_key:
):
break break
scroll_id = res.next_scroll_id
for task, iterations in res.metric_events:
urls.update(ev.get("url") for it in iterations for ev in it["events"])
urls.discard({None})
return urls return urls