Memory optimization for image chunk preparation

changelog.d/20241022_174329_andrey_optimize_chunk_preparation.md

@@ -0,0 +1,4 @@
+### Fixed
+
+- Memory consumption during preparation of image chunks
+  (<https://github.com/cvat-ai/cvat/pull/8581>)

cvat/apps/engine/cache.py

@@ -54,7 +54,7 @@
     ZipChunkWriter,
     ZipCompressedChunkWriter,
 )
-from cvat.apps.engine.utils import md5_hash, preload_images
+from cvat.apps.engine.utils import md5_hash, load_image
 from utils.dataset_manifest import ImageManifestManager
 
 slogger = ServerLogManager(__name__)
@@ -321,15 +321,13 @@ def _read_raw_images(
                 cloud_storage_instance.bulk_download_to_dir(
                     files=files_to_download, upload_dir=tmp_dir
                 )
-                media = preload_images(media)
 
-                for checksum, (_, fs_filename, _) in zip(checksums, media):
-                    if checksum and not md5_hash(fs_filename) == checksum:
+                for checksum, media_item in zip(checksums, media):
+                    if checksum and not md5_hash(media_item[1]) == checksum:
                         slogger.cloud_storage[db_cloud_storage.id].warning(
                             "Hash sums of files {} do not match".format(file_name)
                         )
-
-                yield from media
+                    yield load_image(media_item)
         else:
             requested_frame_iter = iter(frame_ids)
             next_requested_frame_id = next(requested_frame_iter, None)
@@ -359,7 +357,7 @@ def _read_raw_images(
             assert next_requested_frame_id is None
 
             if db_task.dimension == models.DimensionType.DIM_2D:
-                media = preload_images(media)
+                media = map(load_image, media)
 
             yield from media
 

cvat/apps/engine/task.py

@@ -38,7 +38,7 @@
 from cvat.apps.engine.models import RequestAction, RequestTarget
 from cvat.apps.engine.utils import (
     av_scan_paths, format_list,get_rq_job_meta,
-    define_dependent_job, get_rq_lock_by_user, preload_images
+    define_dependent_job, get_rq_lock_by_user, load_image
 )
 from cvat.apps.engine.rq_job_handler import RQId
 from cvat.utils.http import make_requests_session, PROXIES_FOR_UNTRUSTED_URLS
@@ -1537,7 +1537,7 @@ def save_chunks(
                 MEDIA_TYPES['archive']['extractor'],
             ))
         ):
-            chunk_data = preload_images(chunk_data)
+            chunk_data = list(map(load_image, chunk_data))
 
         # TODO: extract into a class
 

cvat/apps/engine/utils.py

@@ -11,7 +11,7 @@
 import sys
 import traceback
 from contextlib import suppress, nullcontext
-from typing import Any, Dict, Optional, Callable, Sequence, Union, Iterable
+from typing import Any, Dict, Optional, Callable, Sequence, Union
 import subprocess
 import os
 import urllib.parse
@@ -363,14 +363,11 @@ def sendfile(
 
     return _sendfile(request, filename, attachment, attachment_filename, mimetype, encoding)
 
-def preload_image(image: tuple[str, str, str])-> tuple[Image.Image, str, str]:
+def load_image(image: tuple[str, str, str])-> tuple[Image.Image, str, str]:
     pil_img = Image.open(image[0])
     pil_img.load()
     return pil_img, image[1], image[2]
 
-def preload_images(images: Iterable[tuple[str, str, str]]) -> list[tuple[Image.Image, str, str]]:
-    return list(map(preload_image, images))
-
 def build_backup_file_name(
     *,
     class_name: str,