Don't pass encodings through Celery

Author: sergey-misuk-valor

src/hope_dedup_engine/apps/api/models/deduplication.py

@@ -3,7 +3,7 @@
 
 from django.conf import settings
 from django.core.validators import MaxValueValidator, MinValueValidator
-from django.db import models
+from django.db import models, transaction
 
 from hope_dedup_engine.apps.security.models import ExternalSystem
 from hope_dedup_engine.types import EncodingType, FindingType, IgnoredPairType
@@ -70,8 +70,10 @@ def get_ignored_pairs(self) -> IgnoredPairType:
         )
 
     def update_encodings(self, encodings: EncodingType) -> None:
-        self.encodings.update(encodings)
-        self.save()
+        with transaction.atomic():
+            fresh_self: DeduplicationSet = DeduplicationSet.objects.select_for_update().get(pk=self.pk)
+            fresh_self.encodings.update(encodings)
+            fresh_self.save()
 
     def update_findings(self, findings: FindingType) -> None:
         images = Image.objects.filter(deduplication_set=self).values("filename", "reference_pk")

src/hope_dedup_engine/apps/faces/celery_tasks.py

@@ -1,5 +1,4 @@
 import traceback
-from collections import ChainMap
 from functools import partial
 from typing import Any, Final
 
@@ -17,7 +16,7 @@
 from hope_dedup_engine.apps.faces.services.facial import dedupe_images, encode_faces
 from hope_dedup_engine.apps.faces.utils import report_long_execution
 from hope_dedup_engine.config.celery import DedupeTask, app
-from hope_dedup_engine.types import EncodingType, FindingType
+from hope_dedup_engine.types import FindingType
 
 CHUNK_SIZE: Final[int] = 25
 
@@ -71,7 +70,7 @@ def encode_chunk(
     self: DedupeTask,
     files: list[str],
     config: dict[str, Any],
-) -> tuple[EncodingType, int, int]:
+) -> None:
     """Encode faces in a chunk of files."""
     with report_long_execution('DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))'):
         ds = DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))
@@ -80,7 +79,9 @@ def encode_chunk(
         with report_long_execution("ds.get_encodings()"):
             pre_encodings = ds.get_encodings()
         with report_long_execution('encode_faces(files, config.get("encoding"), pre_encodings, progress=callback)'):
-            return encode_faces(files, config.get("encoding"), pre_encodings, progress=callback)
+            results = encode_faces(files, config.get("encoding"), pre_encodings, progress=callback)
+        with report_long_execution('ds.update_encodings(results[0])'):
+            ds.update_encodings(results[0])
     except Exception as e:
         sentry_sdk.capture_exception(e)
         handle_error(ds)
@@ -149,17 +150,15 @@ def callback_findings(
 @app.task(bind=True, base=DedupeTask)
 def callback_encodings(
     self: Task,
-    results: tuple[EncodingType, int, int],
+    results: list[None],
     config: dict[str, Any],
 ) -> dict[str, Any]:
     """Aggregate and save encodings."""
     ds = DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))
     try:
-        encodings = dict(ChainMap(*[result[0] for result in results]))
-        ds.update_encodings(encodings)
         deduplicate_dataset.delay(config)
         return {
-            "Encoded": len(encodings),
+            "Encoded": True,
         }
     except Exception as e:
         sentry_sdk.capture_exception(e)