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/jobs/models.py

@@ -897,6 +897,10 @@ def run(cls, job: "Job"):
         job.progress.add_stage(cls.name, key=cls.key)
         job.update_status(JobState.STARTED)
         job.started_at = datetime.datetime.now()
+        # Mark the stage started before the task runs, as the other job types do.
+        # A stage left at CREATED reads as "Waiting to start" however long the task
+        # takes, and a task's first progress report can be minutes in. See #1376.
+        job.progress.update_stage(cls.key, status=JobState.STARTED, progress=0)
         job.save()
 
         params = job.params or {}

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/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,12 @@ def _scoped_classifications(
 
         ``config_schema`` guarantees exactly one scope id is set, so the single
         ``else`` branch is sound.
+
+        Do not add ``.distinct()`` while both scopes filter a to-one path: neither
+        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. Broadening
+        either scope to match several collections or a whole project would fan out
+        and need de-duplication again — restrict the columns first. See #1376.
         """
         base = Classification.objects.filter(
             terminal=True,
@@ -281,7 +296,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 +305,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 +330,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 +352,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,91 @@ 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_scope_query_does_not_deduplicate_rows(self):
+        """Neither scope de-duplicates rows.
+
+        De-duplication is invisible in the output but very costly here, so the
+        query shape is the only thing that can guard it; the test above covers 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,39 @@ 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_post_processing_stage_is_started_before_the_task_runs(self):
+        """The stage reads as running from the moment the job starts.
+
+        A task's first progress report can be minutes into a large run, and a stage
+        left at CREATED renders as "Waiting to start" until then. See #1376.
+        """
+        from unittest.mock import patch
+
+        from ami.jobs.models import Job, JobState, PostProcessingJob
+
+        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},
+            },
+        )
+
+        observed = {}
+
+        def _capture(self_task):
+            stage = self_task.job.progress.get_stage("post_processing")
+            observed["status"] = stage.status
+            observed["label"] = stage.status_label
+
+        with patch.object(SmallSizeFilterTask, "run", _capture):
+            PostProcessingJob.run(job)
+
+        self.assertEqual(observed["status"], JobState.STARTED)
+        self.assertEqual(observed["label"], "0% 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.