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

.agents/AGENTS.md

@@ -69,6 +69,13 @@ A confident but wrong comment is worse than no comment — future developers and
 - **Keep the load-bearing fact; cut unverified mechanism.** A reader changing a constant should not have to trust a paragraph of speculation. State what the code does and the one reason that matters.
 - **Label guesses as guesses.** Use "best-guess", "not measured", "appears to" for anything you have not tested or profiled. Never state a hunch in the authoritative voice.
 - **Short, then link.** Point to the PR or issue for the long reasoning (`See #1231`) instead of embedding it as truth. A PR thread is dated and discussable; an inline comment reads as eternal fact.
+- **Keep comments to a few lines; put the investigation in the PR.** A rationale comment states the rule a future editor must not break, the one reason it exists, and a `See #NNNN` link — normally two or three lines, rarely more than four. Benchmark tables, before/after measurements, the history of what was tried, and anything a reviewer would want to argue with belong in the PR description, where they are dated and can be replied to. The same limit applies to docstrings on tests and helpers: say what the test protects, then link out for why it matters. Long explanatory blocks also age badly, because the next person to change the code updates the line and leaves the essay.
+
+  ```python
+  # Do not add .distinct(): neither scope can duplicate a row, and de-duplicating
+  # sorts the logits and scores arrays — 208s versus 0.5s on a 55,530-row scope.
+  # See #1376.
+  ```
 - **In PR descriptions, distinguish measured from inferred.** Numbers from a profiler, logs, or `ps` are measurements; numbers from reading code are estimates. Say which.
 
 ## Project Overview

ami/ml/management/commands/run_class_masking.py

@@ -46,17 +46,15 @@ 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; do not add .distinct(),
+        # see #1376.
+        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/base.py

@@ -74,7 +74,15 @@ def update_progress(self, progress: float):
         """
 
         if self.job:
-            self.job.progress.update_stage(self.job.job_type_key, progress=progress)
+            # Imported here because ami.jobs.models imports the post-processing
+            # registry, so a module-level import would be circular.
+            from ami.jobs.models import JobState
+
+            # The stage is marked STARTED here because nothing else does it: the job
+            # wrapper creates the stage and later marks it SUCCESS, so a stage
+            # reporting progress would otherwise still be labelled "Waiting to
+            # start" for the whole run. See #1376.
+            self.job.progress.update_stage(self.job.job_type_key, status=JobState.STARTED, progress=progress)
             # Bump updated_at alongside progress: the stale-job reaper
             # (check_stale_jobs) revokes running jobs whose updated_at is older
             # than STALLED_JOBS_MAX_MINUTES. A long post-processing run that only

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,9 @@ 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, before
+    the first row is processed, so a long run reports what it found up front.
+
     ``occurrences_updated`` counts only occurrences whose determination actually
     changed (not just any occurrence touched), matching the size-filter convention.
 
@@ -113,6 +117,11 @@ def make_classifications_filtered_by_taxa_list(
     timestamp = timezone.now()
     masked_count = 0
 
+    # Sizing the scope and expanding the category map both take time on a large
+    # classifier, so report before touching a row.
+    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 +278,10 @@ def _scoped_classifications(
 
         ``config_schema`` guarantees exactly one scope id is set, so the single
         ``else`` branch is sound.
+
+        Do not add ``.distinct()``: neither scope can duplicate a row, and
+        de-duplicating sorts the ``logits`` and ``scores`` arrays, which measured
+        208s versus 0.5s on a 55,530-row scope. See #1376.
         """
         base = Classification.objects.filter(
             terminal=True,
@@ -281,7 +294,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 +303,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 +328,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 +350,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,111 @@ 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 in several collections still contributes one row per
+        classification, which is why the scope needs no de-duplication."""
+        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 de-duplicates rows.
+
+        This costs 208s versus 0.5s on a production-sized scope and is invisible in
+        the output, so the query shape is the only thing that can guard it. The two
+        tests above cover the results being equivalent. See #1376.
+        """
+        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.assertFalse(scoped.query.distinct)
+
+    def test_scope_size_is_reported_before_the_first_batch(self):
+        """The scope size reaches the job before any row is processed, so a large
+        run reports what it found rather than sitting at zero while it starts."""
+        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"])

ami/ml/tests.py

@@ -1651,6 +1651,37 @@ def test_progress_save_bumps_updated_at_for_reaper(self):
         job.refresh_from_db()
         self.assertGreater(job.updated_at, stale, "report_stage_metrics must bump updated_at")
 
+    def test_progress_marks_the_stage_started(self):
+        """A progress heartbeat moves the stage to STARTED.
+
+        The job wrapper creates the stage and later marks it SUCCESS, and nothing
+        in between sets a status. ``get_status_label`` labels CREATED as "Waiting
+        to start" whatever the progress is, so without this a running job reads as
+        not yet begun for its whole duration. See #1376.
+        """
+        from ami.jobs.models import Job, JobState
+
+        job = Job.objects.create(
+            project=self.project,
+            name="stage status test",
+            job_type_key="post_processing",
+            params={
+                "task": "small_size_filter",
+                "config": {"source_image_collection_id": self.collection.pk, "size_threshold": 0.01},
+            },
+        )
+        job.progress.add_stage("Post Processing", key="post_processing")
+        job.save()
+        self.assertEqual(job.progress.stages[0].status, JobState.CREATED)
+
+        task = SmallSizeFilterTask(job=job, source_image_collection_id=self.collection.pk, size_threshold=0.01)
+        task.update_progress(0.5)
+
+        job.refresh_from_db()
+        stage = job.progress.stages[0]
+        self.assertEqual(stage.status, JobState.STARTED)
+        self.assertEqual(stage.status_label, "50% complete")
+
     def test_occurrences_updated_counts_only_changed_determinations(self):
         """``occurrences_updated`` counts occurrences whose determination actually
         changed, not every occurrence the filter re-saved.