Stop large class masking jobs from being cancelled before they report any progress

ami/ml/management/commands/run_class_masking.py

@@ -46,17 +46,18 @@ def handle(self, *args, **options):
 
         from ami.main.models import Classification
 
-        classification_count = (
-            Classification.objects.filter(
-                detection__source_image__collections=collection,
-                terminal=True,
-                algorithm=algorithm,
-                scores__isnull=False,
-                logits__isnull=False,
-            )
-            .distinct()
-            .count()
-        )
+        # Mirrors ClassMaskingTask._scoped_classifications, including its reason for
+        # not de-duplicating rows: collection membership holds one row per
+        # (capture, collection) pair, so this cannot double-count, and de-duplicating
+        # rows that carry the logits and scores arrays is expensive enough to dominate
+        # the command's runtime.
+        classification_count = Classification.objects.filter(
+            detection__source_image__collections=collection,
+            terminal=True,
+            algorithm=algorithm,
+            scores__isnull=False,
+            logits__isnull=False,
+        ).count()
 
         taxa_count = taxa_list.taxa.count()
 

ami/ml/post_processing/class_masking.py

@@ -50,6 +50,7 @@ def make_classifications_filtered_by_taxa_list(
     batch_size: int = 200,
     reweight: bool = True,
     task_logger: logging.Logger = logger,
+    on_setup: Callable[[int], None] | None = None,
     on_batch: Callable[[dict], None] | None = None,
 ) -> dict[str, int]:
     """Re-score ``classifications`` by masking out classes absent from ``taxa_list``.
@@ -67,6 +68,11 @@ def make_classifications_filtered_by_taxa_list(
     ``{"classifications_checked": i, "classifications_total": total,
        "classifications_masked": masked_count, "occurrences_updated": n_changed}``.
 
+    ``on_setup`` is called once with the number of classifications in scope, after
+    the scope has been sized and the category map expanded but before the first row
+    is processed. It gives an operator the size of the run up front and marks the
+    job as alive before any masking work begins.
+
     ``occurrences_updated`` counts only occurrences whose determination actually
     changed (not just any occurrence touched), matching the size-filter convention.
 
@@ -113,6 +119,12 @@ def make_classifications_filtered_by_taxa_list(
     timestamp = timezone.now()
     masked_count = 0
 
+    # Report the size of the run before touching a row: sizing the scope and
+    # expanding the category map both take time on a large classifier, and until
+    # this fires the job looks untouched.
+    if on_setup is not None:
+        on_setup(total)
+
     for i, classification in enumerate(classifications.iterator(chunk_size=batch_size), start=1):
         logits = classification.logits
         if not isinstance(logits, list) or not all(isinstance(x, (int, float)) for x in logits):
@@ -269,6 +281,18 @@ def _scoped_classifications(
 
         ``config_schema`` guarantees exactly one scope id is set, so the single
         ``else`` branch is sound.
+
+        Neither branch de-duplicates rows, and neither needs to: a classification
+        reaches an occurrence through a plain foreign key, and it reaches a
+        collection through a membership table that holds one row per
+        (capture, collection) pair, so filtering on a single occurrence or a single
+        collection cannot multiply rows. De-duplicating is also costly here, because
+        the selected columns include the ``logits`` and ``scores`` arrays, which run
+        to roughly half a megabyte per row for a classifier with tens of thousands
+        of categories. Measured on a 55,530-row scope: counting took 208 seconds
+        with ``.distinct()`` and 0.5 seconds without, returning the same number.
+        De-duplication is also a blocking step, so it delayed the first row of the
+        iteration pass as well as the count.
         """
         base = Classification.objects.filter(
             terminal=True,
@@ -281,7 +305,7 @@ def _scoped_classifications(
             if not Occurrence.objects.filter(pk=config.occurrence_id).exists():
                 raise ValueError(f"Occurrence {config.occurrence_id} not found")
             return (
-                base.filter(detection__occurrence_id=config.occurrence_id).distinct(),
+                base.filter(detection__occurrence_id=config.occurrence_id),
                 f"occurrence {config.occurrence_id}",
             )
 
@@ -290,7 +314,7 @@ def _scoped_classifications(
         except SourceImageCollection.DoesNotExist:
             raise ValueError(f"SourceImageCollection {config.source_image_collection_id} not found")
         return (
-            base.filter(detection__source_image__collections=collection).distinct(),
+            base.filter(detection__source_image__collections=collection),
             f"collection {collection.pk}",
         )
 
@@ -315,6 +339,10 @@ def run(self) -> None:
         classifications, scope_desc = self._scoped_classifications(config, source_algorithm)
         self.logger.info(f"Applying class masking on {scope_desc} using taxa list {taxa_list.pk}")
 
+        def _on_setup(total: int) -> None:
+            self.update_progress(0.0)
+            self.report_stage_metrics({"classifications_total": total})
+
         def _on_batch(m: dict) -> None:
             total = m["classifications_total"]
             self.update_progress(m["classifications_checked"] / total if total else 1.0)
@@ -333,6 +361,7 @@ def _on_batch(m: dict) -> None:
             new_algorithm=masking_algorithm,
             reweight=config.reweight,
             task_logger=self.logger,
+            on_setup=_on_setup,
             on_batch=_on_batch,
         )
         self.report_stage_metrics(metrics)

ami/ml/post_processing/tests/test_class_masking.py

@@ -486,3 +486,133 @@ def test_reweight_false_winner_identical_scores_differ(self):
         # Stored confidence: reweight=False uses the original pre-mask probability.
         self.assertAlmostEqual(new_clf_f.score, scores[0], places=6)
         self.assertNotAlmostEqual(new_clf_t.score, scores[0], places=2)
+
+    # ----- scope query shape ----------------------------------------------
+
+    def test_collection_scope_returns_each_classification_once(self):
+        """A capture that belongs to several collections still contributes one row per
+        classification when the scope is filtered to one of those collections.
+
+        This is the property that makes de-duplication unnecessary: the collection
+        membership table holds at most one row per (capture, collection) pair, and
+        detection -> capture is a plain foreign key, so filtering on a single
+        collection cannot multiply rows.
+        """
+        capture = self.collection.images.first()
+        other_collection = SourceImageCollection.objects.create(
+            name="Second collection containing the same capture",
+            project=self.project,
+            method="manual",
+            kwargs={"image_ids": [capture.pk]},
+        )
+        other_collection.images.add(capture)
+
+        detection = Detection.objects.create(source_image=capture, bbox=[0, 0, 200, 200])
+        occurrence = Occurrence.objects.create(project=self.project, event=self.deployment.events.first())
+        occurrence.detections.add(detection)
+        logits = [2.0, 1.0, 5.0]
+        classification = self._create_classification_with_logits(
+            detection, self.species_taxa[2], _softmax(logits), logits
+        )
+
+        taxa_list = TaxaList.objects.create(name="Scope shape list")
+        taxa_list.taxa.set(self.species_taxa[:2])
+        task = ClassMaskingTask(
+            source_image_collection_id=self.collection.pk,
+            taxa_list_id=taxa_list.pk,
+            algorithm_id=self.algorithm.pk,
+        )
+        scoped, _ = task._scoped_classifications(task.config, self.algorithm)
+
+        self.assertEqual(list(scoped.filter(pk=classification.pk)), [classification])
+        self.assertEqual(
+            scoped.filter(pk=classification.pk).count(),
+            1,
+            "A capture in two collections must not yield the classification twice",
+        )
+
+    def test_occurrence_scope_returns_each_classification_once(self):
+        """An occurrence with several detections yields one row per classification."""
+        taxa_list = TaxaList.objects.create(name="Occurrence scope shape list")
+        taxa_list.taxa.set(self.species_taxa[:2])
+
+        occurrence = Occurrence.objects.create(project=self.project, event=self.deployment.events.first())
+        logits = [2.0, 1.0, 5.0]
+        for capture in self.collection.images.all()[:3]:
+            detection = Detection.objects.create(source_image=capture, bbox=[0, 0, 200, 200])
+            occurrence.detections.add(detection)
+            self._create_classification_with_logits(detection, self.species_taxa[2], _softmax(logits), logits)
+
+        task = ClassMaskingTask(
+            occurrence_id=occurrence.pk,
+            taxa_list_id=taxa_list.pk,
+            algorithm_id=self.algorithm.pk,
+        )
+        scoped, _ = task._scoped_classifications(task.config, self.algorithm)
+        self.assertEqual(scoped.count(), 3)
+
+    def test_scope_query_does_not_deduplicate_rows(self):
+        """Neither scope may emit ``SELECT DISTINCT``, because the selected columns
+        include the ``logits`` and ``scores`` arrays.
+
+        A production classifier carries 29,176 categories, so those two arrays are
+        roughly half a megabyte per row and de-duplicating whole rows forces the
+        database to sort all of them. Measured against a 55,530-row scope: counting
+        the scope took 208 seconds with de-duplication and 0.5 seconds without,
+        returning the same number. The de-duplication is also blocking, so the first
+        row of the iteration pass cannot arrive until the whole sort completes,
+        which is what pushed masking runs past the stalled-job cutoff before any
+        progress was reported.
+
+        The equivalent-results half of this claim is covered by the two scope tests
+        above; de-duplication cannot be observed in the output, so the query shape
+        is the only thing left to pin.
+        """
+        taxa_list = TaxaList.objects.create(name="Query shape list")
+        taxa_list.taxa.set(self.species_taxa[:2])
+        occurrence = Occurrence.objects.create(project=self.project, event=self.deployment.events.first())
+
+        for kwargs in (
+            {"source_image_collection_id": self.collection.pk},
+            {"occurrence_id": occurrence.pk},
+        ):
+            with self.subTest(**kwargs):
+                task = ClassMaskingTask(taxa_list_id=taxa_list.pk, algorithm_id=self.algorithm.pk, **kwargs)
+                scoped, _ = task._scoped_classifications(task.config, self.algorithm)
+                self.assertNotIn("DISTINCT", str(scoped.query).upper())
+
+    def test_scope_size_is_reported_before_the_first_batch(self):
+        """The scope size reaches the job before any row is processed.
+
+        Resolving the scope reads the taxa list and expands the classifier's
+        category map, so an operator watching a large run would otherwise see a
+        stage sitting at zero with no indication of how much work was found. The
+        same callback moves the stage out of its created state, which is what
+        proves to the stalled-job check that the run is alive.
+        """
+        taxa_list = TaxaList.objects.create(name="Setup callback list")
+        taxa_list.taxa.set(self.species_taxa[:2])
+
+        det, _ = self._detection_with_occurrence()
+        logits = [2.0, 1.0, 5.0]
+        self._create_classification_with_logits(det, self.species_taxa[2], _softmax(logits), logits)
+
+        new_algorithm = Algorithm.objects.create(
+            name="masked_setup",
+            key="masked_setup_test",
+            task_type=AlgorithmTaskType.CLASSIFICATION.value,
+            category_map=self.algorithm.category_map,
+        )
+
+        events: list[tuple[str, int]] = []
+        make_classifications_filtered_by_taxa_list(
+            classifications=Classification.objects.filter(detection=det, terminal=True),
+            taxa_list=taxa_list,
+            algorithm=self.algorithm,
+            new_algorithm=new_algorithm,
+            on_setup=lambda total: events.append(("setup", total)),
+            on_batch=lambda m: events.append(("batch", m["classifications_checked"])),
+        )
+
+        self.assertEqual(events[0], ("setup", 1), "Scope size is reported once, before the first batch")
+        self.assertEqual([name for name, _ in events], ["setup", "batch"])