diff --git a/changelog.d/20241113_130658_dmitrii.lavrukhin_no_queryset_cache.md b/changelog.d/20241113_130658_dmitrii.lavrukhin_no_queryset_cache.md new file mode 100644 index 000000000000..512c4b1cd5a6 --- /dev/null +++ b/changelog.d/20241113_130658_dmitrii.lavrukhin_no_queryset_cache.md @@ -0,0 +1,4 @@ +### Fixed + +- Optimized memory consumption when importing annotations to a task with a lot of jobs and images + () diff --git a/cvat/apps/dataset_manager/task.py b/cvat/apps/dataset_manager/task.py index 5b72f92a1ebc..0ad01b06129e 100644 --- a/cvat/apps/dataset_manager/task.py +++ b/cvat/apps/dataset_manager/task.py @@ -13,7 +13,7 @@ from datumaro.components.errors import DatasetError, DatasetImportError, DatasetNotFoundError from django.db import transaction -from django.db.models.query import Prefetch +from django.db.models.query import Prefetch, QuerySet from django.conf import settings from rest_framework.exceptions import ValidationError @@ -81,9 +81,10 @@ def merge_table_rows(rows, keys_for_merge, field_id): return list(merged_rows.values()) + class JobAnnotation: @classmethod - def add_prefetch_info(cls, queryset): + def add_prefetch_info(cls, queryset: QuerySet, prefetch_images: bool = True): assert issubclass(queryset.model, models.Job) label_qs = add_prefetch_fields(models.Label.objects.all(), [ @@ -93,6 +94,12 @@ def add_prefetch_info(cls, queryset): ]) label_qs = JobData.add_prefetch_info(label_qs) + task_data_queryset = models.Data.objects.select_related('video') + if prefetch_images: + task_data_queryset = task_data_queryset.prefetch_related( + Prefetch('images', queryset=models.Image.objects.order_by('frame')) + ) + return queryset.select_related( 'segment', 'segment__task', @@ -103,18 +110,15 @@ def add_prefetch_info(cls, queryset): 'segment__task__project__owner', 'segment__task__project__assignee', - Prefetch('segment__task__data', - queryset=models.Data.objects.select_related('video').prefetch_related( - Prefetch('images', queryset=models.Image.objects.order_by('frame')) - )), + Prefetch('segment__task__data', queryset=task_data_queryset), Prefetch('segment__task__label_set', queryset=label_qs), Prefetch('segment__task__project__label_set', queryset=label_qs), ) - def __init__(self, pk, *, is_prefetched=False, queryset=None): + def __init__(self, pk, *, is_prefetched: bool = False, queryset: QuerySet = None, prefetch_images: bool = True): if queryset is None: - queryset = self.add_prefetch_info(models.Job.objects) + queryset = self.add_prefetch_info(models.Job.objects, prefetch_images=prefetch_images) if is_prefetched: self.db_job: models.Job = queryset.select_related( @@ -1018,7 +1022,7 @@ def put_job_data(pk, data): @plugin_decorator @transaction.atomic def patch_job_data(pk, data, action): - annotation = JobAnnotation(pk) + annotation = JobAnnotation(pk, prefetch_images=False) if action == PatchAction.CREATE: annotation.create(data) elif action == PatchAction.UPDATE: @@ -1031,7 +1035,7 @@ def patch_job_data(pk, data, action): @silk_profile(name="DELETE job data") @transaction.atomic def delete_job_data(pk): - annotation = JobAnnotation(pk) + annotation = JobAnnotation(pk, prefetch_images=False) annotation.delete() def export_job(job_id, dst_file, format_name, server_url=None, save_images=False):