fix: DEV-1694: Make update_tasks_states async (#2014)

* fix: DEV-1694: Make update_tasks_states async

* fix: DEV-1694: Disconnect signals and update state after

* refactor: DEV-1694: Update summary in delete specific tasks

* fix: DEV-1693: Add disconnect a list of signals and refactor signal delete annotations from S3

* fix: DEV-1694: Move project summary recalculation to async call

* fix: DEV-1694: Change delete_annotation_from_s3_storages signal to pre_delete

* docs: DEV-1694: Fix example in temporary_disconnect_list_signal

Author: KonstantinKorotaev

label_studio/core/utils/common.py

@@ -605,3 +605,40 @@ def round_floats(o):
     if isinstance(o, (list, tuple)):
         return [round_floats(x) for x in o]
     return o
+
+
+class temporary_disconnect_list_signal:
+    """ Temporarily disconnect a list of signals
+        Each signal tuple: (signal_type, signal_method, object)
+        Example:
+            with temporary_disconnect_list_signal(
+                [(signals.post_delete, update_is_labeled_after_removing_annotation, Annotation)]
+                ):
+                do_something()
+    """
+    def __init__(self, signals):
+        self.signals = signals
+
+    def __enter__(self):
+        for signal in self.signals:
+            sig = signal[0]
+            receiver = signal[1]
+            sender = signal[2]
+            dispatch_uid = signal[3] if len(signal) > 3 else None
+            sig.disconnect(
+                receiver=receiver,
+                sender=sender,
+                dispatch_uid=dispatch_uid
+            )
+
+    def __exit__(self, type_, value, traceback):
+        for signal in self.signals:
+            sig = signal[0]
+            receiver = signal[1]
+            sender = signal[2]
+            dispatch_uid = signal[3] if len(signal) > 3 else None
+            sig.connect(
+                receiver=receiver,
+                sender=sender,
+                dispatch_uid=dispatch_uid
+            )

label_studio/data_manager/actions/basic.py

@@ -2,15 +2,17 @@
 """
 import logging
 
-from django.db.models import signals
 from datetime import datetime
+from django.db.models.signals import post_delete, pre_delete
 
 from core.permissions import AllPermissions
 from core.redis import start_job_async_or_sync
-from core.utils.common import temporary_disconnect_signal, temporary_disconnect_all_signals
+from core.utils.common import temporary_disconnect_list_signal
+from projects.models import Project
+
 from tasks.models import (
-    Annotation, Prediction, Task, update_is_labeled_after_removing_annotation,
-    bulk_update_stats_project_tasks
+    Annotation, Prediction, Task, bulk_update_stats_project_tasks, update_is_labeled_after_removing_annotation,
+    update_all_task_states_after_deleting_task, remove_data_columns, remove_project_summary_annotations
 )
 from webhooks.utils import emit_webhooks_for_instance
 from webhooks.models import WebhookAction
@@ -38,25 +40,36 @@ def delete_tasks(project, queryset, **kwargs):
     """
     tasks_ids = list(queryset.values('id'))
     count = len(tasks_ids)
+    tasks_ids_list = [task['id'] for task in tasks_ids]
+    # signals to switch off
+    signals = [
+        (post_delete, update_is_labeled_after_removing_annotation, Annotation),
+        (post_delete, update_all_task_states_after_deleting_task, Task),
+        (pre_delete, remove_data_columns, Task),
+        (pre_delete, remove_project_summary_annotations, Annotation)
+    ]
 
     # delete all project tasks
     if count == project.tasks.count():
-        with temporary_disconnect_all_signals():
+        with temporary_disconnect_list_signal(signals):
             queryset.delete()
-
         project.summary.reset()
-        project.update_tasks_states(
-            maximum_annotations_changed=False,
-            overlap_cohort_percentage_changed=False,
-            tasks_number_changed=True
-        )
 
     # delete only specific tasks
     else:
-        # this signal re-save the task back
-        with temporary_disconnect_signal(signals.post_delete, update_is_labeled_after_removing_annotation, Annotation):
+        # update project summary
+        start_job_async_or_sync(async_project_summary_recalculation, tasks_ids_list, project.id)
+
+        with temporary_disconnect_list_signal(signals):
             queryset.delete()
 
+    project.update_tasks_states(
+        maximum_annotations_changed=False,
+        overlap_cohort_percentage_changed=False,
+        tasks_number_changed=True
+    )
+
+    # emit webhooks for project
     emit_webhooks_for_instance(project.organization, project, WebhookAction.TASKS_DELETED, tasks_ids)
 
     # remove all tabs if there are no tasks in project
@@ -103,6 +116,13 @@ def delete_tasks_predictions(project, queryset, **kwargs):
     return {'processed_items': count, 'detail': 'Deleted ' + str(count) + ' predictions'}
 
 
+def async_project_summary_recalculation(tasks_ids_list, project_id):
+    queryset = Task.objects.filter(id__in=tasks_ids_list)
+    project = Project.objects.get(id=project_id)
+    project.summary.remove_created_annotations_and_labels(Annotation.objects.filter(task__in=queryset))
+    project.summary.remove_data_columns(queryset)
+
+
 actions = [
     {
         'entry_point': retrieve_tasks_predictions,

label_studio/io_storages/s3/models.py

@@ -10,7 +10,7 @@
 from django.conf import settings
 from django.utils.translation import gettext_lazy as _
 from django.dispatch import receiver
-from django.db.models.signals import post_save, post_delete
+from django.db.models.signals import post_save, post_delete, pre_delete
 
 from io_storages.base_models import ImportStorage, ImportStorageLink, ExportStorage, ExportStorageLink
 from io_storages.utils import get_uri_via_regex
@@ -209,14 +209,14 @@ def export_annotation_to_s3_storages(sender, instance, **kwargs):
             storage.save_annotation(instance)
 
 
-@receiver(post_delete, sender=Annotation)
+@receiver(pre_delete, sender=Annotation)
 def delete_annotation_from_s3_storages(sender, instance, **kwargs):
-    project = instance.task.project
-    if hasattr(project, 'io_storages_s3exportstorages'):
-        for storage in project.io_storages_s3exportstorages.all():
-            if storage.can_delete_objects:
-                logger.debug(f'Delete {instance} from S3 storage {storage}')
-                storage.delete_annotation(instance)
+    links = S3ExportStorageLink.objects.filter(annotation=instance)
+    for link in links:
+        storage = link.storage
+        if storage.can_delete_objects:
+            logger.debug(f'Delete {instance} from S3 storage {storage}')
+            storage.delete_annotation(instance)
 
 
 class S3ImportStorageLink(ImportStorageLink):

label_studio/projects/models.py

@@ -12,6 +12,7 @@
 from annoying.fields import AutoOneToOneField
 from functools import lru_cache
 
+from core.redis import start_job_async_or_sync
 from tasks.models import Task, Prediction, Annotation, Q_task_finished_annotations, bulk_update_stats_project_tasks
 from core.utils.common import create_hash, get_attr_or_item, load_func
 from core.utils.exceptions import LabelStudioValidationErrorSentryIgnored
@@ -344,10 +345,9 @@ def has_collaborator_enabled(self, user):
         membership = ProjectMember.objects.filter(user=user, project=self)
         return membership.exists() and membership.first().enabled
 
-    def update_tasks_states(
+    def _update_tasks_states(
         self, maximum_annotations_changed, overlap_cohort_percentage_changed, tasks_number_changed
     ):
-
         # if only maximum annotations parameter is tweaked
         if maximum_annotations_changed and not overlap_cohort_percentage_changed:
             tasks_with_overlap = self.tasks.filter(overlap__gt=1)
@@ -371,6 +371,12 @@ def update_tasks_states(
                 self.tasks.filter(Q(annotations__isnull=False) & Q(annotations__ground_truth=False))
             )
 
+    def update_tasks_states(
+        self, maximum_annotations_changed, overlap_cohort_percentage_changed, tasks_number_changed
+    ):
+        start_job_async_or_sync(self._update_tasks_states, maximum_annotations_changed, overlap_cohort_percentage_changed, tasks_number_changed)
+
+
     def _rearrange_overlap_cohort(self):
         """
         Rearrange overlap depending on annotation count in tasks