Refactor enrich_images_unsorted_sync/enrich_images_unsorted_async

Author: tcompa

fractal_server/app/routes/api/v2/history.py

@@ -34,7 +34,7 @@
 from fractal_server.app.schemas.v2 import HistoryUnitStatus
 from fractal_server.app.schemas.v2 import HistoryUnitStatusWithUnset
 from fractal_server.app.schemas.v2 import ImageLogsRequest
-from fractal_server.images.status_tools import enrich_images_async
+from fractal_server.images.status_tools import enrich_images_unsorted_async
 from fractal_server.images.status_tools import IMAGE_STATUS_KEY
 from fractal_server.images.tools import aggregate_attributes
 from fractal_server.images.tools import aggregate_types
@@ -334,7 +334,7 @@ async def get_history_images(
     types = aggregate_types(type_filtered_images)
 
     # (3) Enrich images with status attribute
-    type_filtered_images_with_status = await enrich_images_async(
+    type_filtered_images_with_status = await enrich_images_unsorted_async(
         dataset_id=dataset_id,
         workflowtask_id=workflowtask_id,
         images=type_filtered_images,

fractal_server/app/routes/api/v2/pre_submission_checks.py

@@ -14,7 +14,7 @@
 from fractal_server.app.routes.auth import current_active_user
 from fractal_server.app.schemas.v2 import HistoryUnitStatus
 from fractal_server.app.schemas.v2 import TaskType
-from fractal_server.images.status_tools import enrich_images_async
+from fractal_server.images.status_tools import enrich_images_unsorted_async
 from fractal_server.images.status_tools import IMAGE_STATUS_KEY
 from fractal_server.images.tools import aggregate_types
 from fractal_server.images.tools import filter_image_list
@@ -46,7 +46,7 @@ async def verify_unique_types(
         filtered_images = dataset.images
     else:
         if IMAGE_STATUS_KEY in query.attribute_filters.keys():
-            images = await enrich_images_async(
+            images = await enrich_images_unsorted_async(
                 dataset_id=dataset_id,
                 workflowtask_id=workflowtask_id,
                 images=dataset.images,
@@ -134,7 +134,7 @@ async def check_non_processed_images(
         attribute_filters=filters.attribute_filters,
     )
 
-    filtered_images_with_status = await enrich_images_async(
+    filtered_images_with_status = await enrich_images_unsorted_async(
         dataset_id=dataset_id,
         workflowtask_id=previous_wft.id,
         images=filtered_images,

fractal_server/app/runner/v2/runner.py

@@ -32,7 +32,7 @@
 from fractal_server.app.schemas.v2 import TaskGroupDumpV2
 from fractal_server.app.schemas.v2 import TaskType
 from fractal_server.images import SingleImage
-from fractal_server.images.status_tools import enrich_images_sync
+from fractal_server.images.status_tools import enrich_images_unsorted_sync
 from fractal_server.images.status_tools import IMAGE_STATUS_KEY
 from fractal_server.images.tools import filter_image_list
 from fractal_server.images.tools import find_image_by_zarr_url
@@ -147,7 +147,7 @@ def execute_tasks_v2(
 
             if ind_wftask == 0 and ENRICH_IMAGES_WITH_STATUS:
                 # FIXME: Could this be done on `type_filtered_images`?
-                tmp_images = enrich_images_sync(
+                tmp_images = enrich_images_unsorted_sync(
                     images=tmp_images,
                     dataset_id=dataset.id,
                     workflowtask_id=wftask.id,

fractal_server/images/status_tools.py

@@ -10,15 +10,18 @@
 from fractal_server.app.models.v2 import HistoryUnit
 from fractal_server.app.schemas.v2 import HistoryUnitStatusWithUnset
 from fractal_server.logger import set_logger
-from fractal_server.types import ImageAttributeValue
 
 logger = set_logger(__name__)
 
 
 IMAGE_STATUS_KEY = "__wftask_dataset_image_status__"
 
 
-def _enriched_image(*, img: dict[str, Any], status: str) -> dict[str, Any]:
+def _enriched_image(
+    *,
+    img: dict[str, Any],
+    status: str,
+) -> dict[str, Any]:
     return img | {
         "attributes": (img["attributes"] | {IMAGE_STATUS_KEY: status})
     }
@@ -29,6 +32,9 @@ def _prepare_query(
     dataset_id: int,
     workflowtask_id: int,
 ) -> Select:
+    """
+    Note: the query does not include `.order_by`.
+    """
     stm = (
         select(HistoryImageCache.zarr_url, HistoryUnit.status)
         .join(HistoryUnit)
@@ -39,13 +45,56 @@ def _prepare_query(
     return stm
 
 
-async def enrich_images_async(
+def _postprocess_image_lists(
+    target_images: list[dict[str, Any]],
+    list_query_url_status: list[tuple[str, str]],
+) -> list[dict[str, Any]]:
+    """ """
+    t_1 = time.perf_counter()
+
+    # Select only processed images that are part of the target image set
+    zarr_url_to_image = {img["zarr_url"]: img for img in target_images}
+    target_zarr_urls = zarr_url_to_image.keys()
+    list_processed_url_status = [
+        url_status
+        for url_status in list_query_url_status
+        if url_status[0] in target_zarr_urls
+    ]
+
+    set_processed_urls = set(
+        url_status[0] for url_status in list_processed_url_status
+    )
+    processed_images_with_status = [
+        _enriched_image(
+            img=zarr_url_to_image[item[0]],
+            status=item[1],
+        )
+        for item in list_processed_url_status
+    ]
+
+    non_processed_urls = target_zarr_urls - set_processed_urls
+    non_processed_images_with_status = [
+        _enriched_image(
+            img=zarr_url_to_image[zarr_url],
+            status=HistoryUnitStatusWithUnset.UNSET,
+        )
+        for zarr_url in non_processed_urls
+    ]
+    t_2 = time.perf_counter()
+    logger.debug(
+        f"[enrich_images_async] post-processing, elapsed={t_2 - t_1:.5f} s"
+    )
+
+    return processed_images_with_status + non_processed_images_with_status
+
+
+async def enrich_images_unsorted_async(
     *,
     images: list[dict[str, Any]],
     dataset_id: int,
     workflowtask_id: int,
     db: AsyncSession,
-) -> list[dict[str, ImageAttributeValue]]:
+) -> list[dict[str, Any]]:
     """
     Enrich images with a status-related attribute.
 
@@ -56,124 +105,75 @@ async def enrich_images_async(
         db: An async db session
 
     Returns:
-        The list of enriched images
+        The list of enriched images, not necessarily in the same order as
+        the input.
     """
     t_0 = time.perf_counter()
     logger.info(
         f"[enrich_images_async] START, {dataset_id=}, {workflowtask_id=}"
     )
 
-    zarr_url_to_image = {img["zarr_url"]: img for img in images}
-
+    # Get `(zarr_url, status)` for _all_ processed images (including those that
+    # are not part of the target image set)
     res = await db.execute(
         _prepare_query(
             dataset_id=dataset_id,
             workflowtask_id=workflowtask_id,
         )
     )
-    zarr_urls = zarr_url_to_image.keys()
-    list_processed_url_status = [
-        item for item in res.all() if item[0] in zarr_urls
-    ]
+    list_query_url_status = res.all()
     t_1 = time.perf_counter()
-    logger.debug(f"[enrich_images_async] db-query, elapsed={t_1 - t_0:.4f} s")
-
-    set_processed_urls = set(item[0] for item in list_processed_url_status)
-    processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[item[0]],
-            status=item[1],
-        )
-        for item in list_processed_url_status
-    ]
-    t_2 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] processed-images, " f"elapsed={t_2 - t_1:.4f} s"
-    )
+    logger.debug(f"[enrich_images_async] query, elapsed={t_1 - t_0:.5f} s")
 
-    non_processed_urls = zarr_url_to_image.keys() - set_processed_urls
-    non_processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[zarr_url],
-            status=HistoryUnitStatusWithUnset.UNSET,
-        )
-        for zarr_url in non_processed_urls
-    ]
-    t_3 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] non-processed-images, "
-        f"elapsed={t_3 - t_2:.4f} s"
+    output = _postprocess_image_lists(
+        target_images=images,
+        list_query_url_status=list_query_url_status,
     )
 
-    return processed_images_with_status + non_processed_images_with_status
+    return output
 
 
-def enrich_images_sync(
+def enrich_images_unsorted_sync(
     *,
     images: list[dict[str, Any]],
     dataset_id: int,
     workflowtask_id: int,
-) -> list[dict[str, ImageAttributeValue]]:
+) -> list[dict[str, Any]]:
     """
     Enrich images with a status-related attribute.
 
+
     Args:
         images: The input image list
         dataset_id: The dataset ID
         workflowtask_id: The workflow-task ID
 
     Returns:
-        The list of enriched images
+        The list of enriched images, not necessarily in the same order as
+        the input.
     """
+
     t_0 = time.perf_counter()
     logger.info(
         f"[enrich_images_async] START, {dataset_id=}, {workflowtask_id=}"
     )
 
-    zarr_url_to_image = {img["zarr_url"]: img for img in images}
-
-    t_1 = time.perf_counter()
-    logger.debug(f"[enrich_images_async] deep-copy, elapsed={t_1 - t_0:.4f} s")
-
+    # Get `(zarr_url, status)` for _all_ processed images (including those that
+    # are not part of the target image set)
     with next(get_sync_db()) as db:
         res = db.execute(
             _prepare_query(
                 dataset_id=dataset_id,
                 workflowtask_id=workflowtask_id,
             )
         )
-    zarr_urls = zarr_url_to_image.keys()
-    list_processed_url_status = [
-        item for item in res.all() if item[0] in zarr_urls
-    ]
-    t_2 = time.perf_counter()
-    logger.debug(f"[enrich_images_async] db-query, elapsed={t_2 - t_1:.4f} s")
-
-    set_processed_urls = set(item[0] for item in list_processed_url_status)
-    processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[item[0]],
-            status=item[1],
-        )
-        for item in list_processed_url_status
-    ]
-    t_3 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] processed-images, " f"elapsed={t_3 - t_2:.4f} s"
-    )
+        list_query_url_status = res.all()
+    t_1 = time.perf_counter()
+    logger.debug(f"[enrich_images_async] query, elapsed={t_1 - t_0:.5f} s")
 
-    non_processed_urls = zarr_url_to_image.keys() - set_processed_urls
-    non_processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[zarr_url],
-            status=HistoryUnitStatusWithUnset.UNSET,
-        )
-        for zarr_url in non_processed_urls
-    ]
-    t_4 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] non-processed-images, "
-        f"elapsed={t_4 - t_3:.4f} s"
+    output = _postprocess_image_lists(
+        target_images=images,
+        list_query_url_status=list_query_url_status,
     )
 
-    return processed_images_with_status + non_processed_images_with_status
+    return output