Merge pull request #153 from unicef/hotfix/report-long-execution-in-encode_chunk

Hotfix: Report long execution in encode_chunk

Author: sergey-misuk-valor

docker/bin/docker-entrypoint.sh

@@ -31,7 +31,7 @@ case "$1" in
       ;;
     worker)
 	    set -- tini -- "$@"
-      set -- gosu hope:unicef celery -A hope_dedup_engine.config.celery worker -E --loglevel=ERROR --concurrency=4
+      set -- gosu hope:unicef celery -A hope_dedup_engine.config.celery worker -E --loglevel=ERROR --concurrency=2
       ;;
     beat)
 	    set -- tini -- "$@"

src/hope_dedup_engine/apps/faces/celery_tasks.py

@@ -15,6 +15,7 @@
 from hope_dedup_engine.apps.api.utils.notification import send_notification
 from hope_dedup_engine.apps.faces.managers import FileSyncManager
 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
 
@@ -72,11 +73,14 @@ def encode_chunk(
     config: dict[str, Any],
 ) -> tuple[EncodingType, int, int]:
     """Encode faces in a chunk of files."""
-    ds = DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))
+    with report_long_execution('DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))'):
+        ds = DeduplicationSet.objects.get(pk=config.get("deduplication_set_id"))
     try:
         callback = partial(notify_status, task=self, dedup_job_id=ds.dedupjob.pk)
-        pre_encodings = ds.get_encodings()
-        return encode_faces(files, config.get("encoding"), pre_encodings, progress=callback)
+        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)
     except Exception as e:
         sentry_sdk.capture_exception(e)
         handle_error(ds)

src/hope_dedup_engine/apps/faces/services/facial.py

@@ -6,7 +6,7 @@
 
 from hope_dedup_engine.apps.api.models import Image
 from hope_dedup_engine.apps.faces.managers import ImagesStorageManager
-from hope_dedup_engine.apps.faces.utils import is_facial_error
+from hope_dedup_engine.apps.faces.utils import is_facial_error, report_long_execution
 from hope_dedup_engine.types import EncodingType, FindingType, IgnoredPairType
 
 logger = logging.getLogger(__name__)
@@ -25,24 +25,30 @@ def encode_faces(
     if not callable(progress):
         progress = default_progress
 
-    storage = ImagesStorageManager()
-    images = storage.get_files()
+    with report_long_execution("ImagesStorageManager()"):
+        storage = ImagesStorageManager()
+
+    with report_long_execution("storage.get_files()"):
+        images = storage.get_files()
 
     encoded = {}
     if pre_encodings:
-        encoded.update(pre_encodings)
+        with report_long_execution("encoded.update(pre_encodings)"):
+            encoded.update(pre_encodings)
     added_cnt = existing_cnt = 0
     existing_cnt = 1000
     for file in files:
-        progress()
+        with report_long_execution("progress()"):
+            progress()
         if file not in images:
             encoded[file] = Image.StatusCode.NO_FILE_FOUND.name
             continue
         if file in encoded:
             existing_cnt += 1
             continue
         try:
-            result = DeepFace.represent(storage.load_image(file), **(options or {}))
+            with report_long_execution("DeepFace.represent(storage.load_image(file), **(options or {}))"):
+                result = DeepFace.represent(storage.load_image(file), **(options or {}))
             if len(result) > 1:
                 encoded[file] = Image.StatusCode.MULTIPLE_FACES_DETECTED.name
             else:

src/hope_dedup_engine/apps/faces/utils.py

@@ -1,3 +1,9 @@
+import time
+from contextlib import contextmanager
+from typing import Generator
+
+import sentry_sdk
+
 from hope_dedup_engine.apps.api.models import Image
 
 
@@ -11,3 +17,16 @@ def is_facial_error(value):
             Image.StatusCode.values + Image.StatusCode.names + [choice.label for choice in Image.StatusCode]
         )
     return False
+
+
+DEFAULT_THRESHOLD_SECONDS = 60
+
+
+@contextmanager
+def report_long_execution(
+    message: str, threshold_seconds: int = DEFAULT_THRESHOLD_SECONDS
+) -> Generator[None, None, None]:
+    start = time.time()
+    yield
+    if (total := time.time() - start) > threshold_seconds:
+        sentry_sdk.capture_message(f"Execution took {total} seconds: {message}")