fix: LSDV-5289: Prevent db deadlocks (#4394)

* fix: LSDV-5289: Prevent db deadlocks * Cover by atomic only select_for_update * Overlap and is_labeled update optimization * Fix empty label stream
HumanSignal · Jun 29, 2023 · f0828fd · f0828fd
1 parent c6f72b0
commit f0828fd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/label_studio/projects/models.py b/label_studio/projects/models.py
@@ -412,7 +412,8 @@ def _rearrange_overlap_cohort(self):
         logger.info(f"Required tasks {must_tasks} and left required tasks {left_must_tasks}")
         if left_must_tasks > 0:
             # if there are unfinished tasks update tasks with count(annotations) >= overlap
-            tasks_with_max_annotations.update(overlap=max_annotations, is_labeled=True)
+            ids = tasks_with_max_annotations.values_list('id', flat=True)
+            all_project_tasks.filter(id__in=ids).update(overlap=max_annotations, is_labeled=True)
             # order other tasks by count(annotations)
             tasks_with_min_annotations = tasks_with_min_annotations.annotate(
                 anno=Count('annotations')
@@ -426,8 +427,10 @@ def _rearrange_overlap_cohort(self):
             min_tasks_to_update = all_project_tasks.filter(id__in=ids)
             min_tasks_to_update.update(overlap=1)
         else:
-            tasks_with_max_annotations.update(overlap=max_annotations)
-            tasks_with_min_annotations.update(overlap=1)
+            ids = tasks_with_max_annotations.values_list('id', flat=True)
+            all_project_tasks.filter(id__in=ids).update(overlap=max_annotations)
+            ids = tasks_with_min_annotations.values_list('id', flat=True)
+            all_project_tasks.filter(id__in=ids).update(overlap=1)
         # update is labeled after tasks rearrange overlap
         bulk_update_stats_project_tasks(all_project_tasks, project=self)
 

diff --git a/label_studio/tasks/models.py b/label_studio/tasks/models.py
@@ -884,9 +884,10 @@ def bulk_update_stats_project_tasks(tasks, project=None):
             # finished tasks
             finished_tasks = tasks.filter(Q(total_annotations__gte=maximum_annotations) |
                                           Q(total_annotations__gte=1, overlap=1))
-            finished_tasks.update(is_labeled=True)
+            ids = finished_tasks.values_list('id', flat=True)
+            tasks.filter(id__in=ids).update(is_labeled=True)
             # unfinished tasks
-            tasks.exclude(id__in=finished_tasks).update(is_labeled=False)
+            tasks.exclude(id__in=ids).update(is_labeled=False)
         else:
             # update objects without saving if we can't use overlap
             for task in tasks: