From 566b772292473eb00772ef79f71dee9ef8e07f2c Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 22 Jul 2026 09:57:57 -0700 Subject: [PATCH 1/4] fix: stop de-duplicating class-masking scope rows Class masking resolved its scope with `.distinct()`, which made Postgres de-duplicate whole classification rows. Those rows carry the `logits` and `scores` arrays, roughly half a megabyte each for a classifier with tens of thousands of categories, so the database had to sort all of them to prove uniqueness of a column set that already contains the primary key. Neither scope can produce duplicate rows in the first place: a classification reaches an occurrence through a foreign key, and it reaches a collection through a membership table holding one row per (capture, collection) pair. 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 the iteration pass could not yield its first row until the same sort finished. Both delays land before the first batch commits, which is where the run reports its first progress, so large masking jobs crossed the stalled-job cutoff and were revoked having reported nothing. Also report the scope size through a new `on_setup` callback before the first row is processed, so a long run shows how much work it found rather than an untouched stage, and apply the same fix to the management command's preflight count. Co-Authored-By: Claude --- .../management/commands/run_class_masking.py | 23 ++-- ami/ml/post_processing/class_masking.py | 33 ++++- .../tests/test_class_masking.py | 130 ++++++++++++++++++ 3 files changed, 173 insertions(+), 13 deletions(-) diff --git a/ami/ml/management/commands/run_class_masking.py b/ami/ml/management/commands/run_class_masking.py index 940956c6e..eb31c6208 100644 --- a/ami/ml/management/commands/run_class_masking.py +++ b/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() diff --git a/ami/ml/post_processing/class_masking.py b/ami/ml/post_processing/class_masking.py index ffc77f346..ab75e9321 100644 --- a/ami/ml/post_processing/class_masking.py +++ b/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) diff --git a/ami/ml/post_processing/tests/test_class_masking.py b/ami/ml/post_processing/tests/test_class_masking.py index b0578272c..b1b61cd11 100644 --- a/ami/ml/post_processing/tests/test_class_masking.py +++ b/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"]) From 41e5e05f6ead2105f827cec97791c0fefd71dbd2 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 22 Jul 2026 13:48:03 -0700 Subject: [PATCH 2/4] fix: mark post-processing stages started, and trim the comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback on #1376. @copilot pointed out that the new setup callback does not move a stage out of its created state, and the cause turned out to be broader than the callback: nothing in the post-processing framework ever set a stage status. The job wrapper creates the stage and later marks it SUCCESS, and `get_status_label` reports "Waiting to start" for a created stage whatever its progress, so every post-processing run — the small size filter included — was labelled as not yet begun for its whole duration. `update_progress` now marks the stage STARTED. This does not disturb `JobProgress.is_complete()`, which gates on a final status, and STARTED is no more final than CREATED. Also on @copilot's suggestion, the de-duplication test now reads `query.distinct` instead of matching the word DISTINCT in generated SQL. The comments and test docstrings added in the previous commit were too long for this repository's norms: the measurements and the reasoning belong in the pull request, where they are dated and can be argued with, and the comment should carry the rule a future editor must not break plus a link. Trimmed accordingly, and the norm is now written down in the agent instructions so the next change does not repeat the round. Co-Authored-By: Claude --- .agents/AGENTS.md | 7 ++++ .../management/commands/run_class_masking.py | 7 +--- ami/ml/post_processing/base.py | 10 ++++- ami/ml/post_processing/class_masking.py | 25 ++++------- .../tests/test_class_masking.py | 42 +++++-------------- ami/ml/tests.py | 31 ++++++++++++++ 6 files changed, 66 insertions(+), 56 deletions(-) diff --git a/.agents/AGENTS.md b/.agents/AGENTS.md index 458ae513b..1499fc174 100644 --- a/.agents/AGENTS.md +++ b/.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 diff --git a/ami/ml/management/commands/run_class_masking.py b/ami/ml/management/commands/run_class_masking.py index eb31c6208..c0fa88b65 100644 --- a/ami/ml/management/commands/run_class_masking.py +++ b/ami/ml/management/commands/run_class_masking.py @@ -46,11 +46,8 @@ def handle(self, *args, **options): from ami.main.models import Classification - # 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. + # Mirrors ClassMaskingTask._scoped_classifications; do not add .distinct(), + # see #1376. classification_count = Classification.objects.filter( detection__source_image__collections=collection, terminal=True, diff --git a/ami/ml/post_processing/base.py b/ami/ml/post_processing/base.py index 8f197f192..b3cfab8f8 100644 --- a/ami/ml/post_processing/base.py +++ b/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 diff --git a/ami/ml/post_processing/class_masking.py b/ami/ml/post_processing/class_masking.py index ab75e9321..ab191a16d 100644 --- a/ami/ml/post_processing/class_masking.py +++ b/ami/ml/post_processing/class_masking.py @@ -68,10 +68,8 @@ 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. + ``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. @@ -119,9 +117,8 @@ 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. + # 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) @@ -282,17 +279,9 @@ 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. + 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, diff --git a/ami/ml/post_processing/tests/test_class_masking.py b/ami/ml/post_processing/tests/test_class_masking.py index b1b61cd11..7f2275282 100644 --- a/ami/ml/post_processing/tests/test_class_masking.py +++ b/ami/ml/post_processing/tests/test_class_masking.py @@ -490,14 +490,8 @@ def test_reweight_false_winner_identical_scores_differ(self): # ----- 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. - """ + """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", @@ -552,21 +546,11 @@ def test_occurrence_scope_returns_each_classification_once(self): 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. + """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]) @@ -579,17 +563,11 @@ def test_scope_query_does_not_deduplicate_rows(self): 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()) + 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. - - 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. - """ + """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]) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index cb88ee6fd..6cecb0d34 100644 --- a/ami/ml/tests.py +++ b/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. From 3c1d54288812cb670b9cbea976bee2a4bad54df5 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 22 Jul 2026 14:18:35 -0700 Subject: [PATCH 3/4] fix: promote a post-processing stage only from its created state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback on #1376, both points from @coderabbitai. My previous commit message claimed nothing ever set a stage status. That was wrong: `Job.update_progress` promotes a stage from CREATED to STARTED as soon as its progress exceeds zero. The gap is narrower than I described — progress of exactly zero does not clear that bar, and exactly zero is what the new setup report means, so it would have left the stage looking untouched. Two consequences. The promotion now happens only from CREATED, mirroring the condition `Job.update_progress` already applies, so a late heartbeat cannot pull a finished, failed or cancelled stage back to STARTED and hide the outcome. And the test that was meant to pin this used a progress of 0.5, which the existing promotion already handled, so it passed with or without the change; it now uses zero, and a second test covers the stage that must not be resurrected. Co-Authored-By: Claude --- ami/ml/post_processing/base.py | 13 +++++---- ami/ml/tests.py | 50 ++++++++++++++++++++++++++-------- 2 files changed, 46 insertions(+), 17 deletions(-) diff --git a/ami/ml/post_processing/base.py b/ami/ml/post_processing/base.py index b3cfab8f8..8f402faa8 100644 --- a/ami/ml/post_processing/base.py +++ b/ami/ml/post_processing/base.py @@ -78,11 +78,14 @@ def update_progress(self, progress: float): # 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) + # Job.update_progress only promotes a stage out of CREATED once its + # progress exceeds zero, so a stage reporting exactly 0% stays labelled + # "Waiting to start". Promote it here, and only from CREATED, so a + # heartbeat can never pull a finished or cancelled stage backwards. + # See #1376. + stage = self.job.progress.get_stage(self.job.job_type_key) + status = {"status": JobState.STARTED} if stage.status == JobState.CREATED else {} + self.job.progress.update_stage(self.job.job_type_key, progress=progress, **status) # 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 diff --git a/ami/ml/tests.py b/ami/ml/tests.py index 6cecb0d34..3723f9c15 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -1651,19 +1651,12 @@ 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 + def _post_processing_job(self, name: str): + from ami.jobs.models import Job job = Job.objects.create( project=self.project, - name="stage status test", + name=name, job_type_key="post_processing", params={ "task": "small_size_filter", @@ -1672,15 +1665,48 @@ def test_progress_marks_the_stage_started(self): ) job.progress.add_stage("Post Processing", key="post_processing") job.save() + return job + + def test_zero_progress_heartbeat_marks_the_stage_started(self): + """A heartbeat at exactly 0% still moves the stage out of CREATED. + + ``Job.update_progress`` only promotes a stage once its progress exceeds + zero, so the setup heartbeat would otherwise leave a running job labelled + "Waiting to start". See #1376. + """ + from ami.jobs.models import JobState + + job = self._post_processing_job("stage status test") 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) + task.update_progress(0.0) job.refresh_from_db() stage = job.progress.stages[0] self.assertEqual(stage.status, JobState.STARTED) - self.assertEqual(stage.status_label, "50% complete") + self.assertEqual(stage.status_label, "0% complete") + + def test_heartbeat_does_not_resurrect_a_finished_stage(self): + """Promotion applies only to a stage that has not started. + + A late heartbeat must not pull a stage that already reached a final state + back to STARTED, which would hide a failure or a cancellation. See #1376. + """ + from ami.jobs.models import Job, JobState + + for final in (JobState.SUCCESS, JobState.FAILURE, JobState.REVOKED): + with self.subTest(final=final): + job = self._post_processing_job(f"stage status {final} test") + job.progress.update_stage("post_processing", status=final, progress=1.0) + job.save() + + task = SmallSizeFilterTask(job=job, source_image_collection_id=self.collection.pk, size_threshold=0.01) + task.update_progress(0.5) + + job.refresh_from_db() + self.assertEqual(job.progress.stages[0].status, final) + Job.objects.filter(pk=job.pk).delete() def test_occurrences_updated_counts_only_changed_determinations(self): """``occurrences_updated`` counts occurrences whose determination actually From d44e00b26aa4036fde88fe06d6bb3a30e14db10d Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 22 Jul 2026 15:30:37 -0700 Subject: [PATCH 4/4] fix: mark the post-processing stage started in the job, not the heartbeat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Takeaway review feedback on #1376. The previous two commits marked the stage STARTED from inside BasePostProcessingTask.update_progress. That was the wrong layer. Every other job type sets its stage started at the top of run() — RegroupEventsJob and SourceImageCollectionPopulateJob both do — and PostProcessingJob was simply missing that line. Doing it there fires before the scope is even counted, rather than at the first heartbeat, so the stage reads as running through the whole setup delay this PR is about. It also lets update_progress return to main's version, dropping the per-heartbeat status write and the guard that came with it. Pinned by test_post_processing_stage_is_started_before_the_task_runs. Also from the review: - The comment on _scoped_classifications now says when de-duplication would become necessary again — broadening either scope to match several collections or a whole project would fan out — rather than only that it is currently forbidden. - Dropped test_occurrence_scope_returns_each_classification_once: an occurrence filters a foreign key, which structurally cannot duplicate a row, so the test passed under any change. The collection variant, where a capture belongs to two collections, is the one that earns its place. Co-Authored-By: Claude --- ami/jobs/models.py | 4 ++ ami/ml/post_processing/base.py | 13 +--- ami/ml/post_processing/class_masking.py | 8 ++- .../tests/test_class_masking.py | 26 +------- ami/ml/tests.py | 62 ++++++------------- 5 files changed, 32 insertions(+), 81 deletions(-) diff --git a/ami/jobs/models.py b/ami/jobs/models.py index bced66575..ff65f31f2 100644 --- a/ami/jobs/models.py +++ b/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 {} diff --git a/ami/ml/post_processing/base.py b/ami/ml/post_processing/base.py index 8f402faa8..8f197f192 100644 --- a/ami/ml/post_processing/base.py +++ b/ami/ml/post_processing/base.py @@ -74,18 +74,7 @@ def update_progress(self, progress: float): """ if self.job: - # 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 - - # Job.update_progress only promotes a stage out of CREATED once its - # progress exceeds zero, so a stage reporting exactly 0% stays labelled - # "Waiting to start". Promote it here, and only from CREATED, so a - # heartbeat can never pull a finished or cancelled stage backwards. - # See #1376. - stage = self.job.progress.get_stage(self.job.job_type_key) - status = {"status": JobState.STARTED} if stage.status == JobState.CREATED else {} - self.job.progress.update_stage(self.job.job_type_key, progress=progress, **status) + self.job.progress.update_stage(self.job.job_type_key, 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 diff --git a/ami/ml/post_processing/class_masking.py b/ami/ml/post_processing/class_masking.py index ab191a16d..fc72b836a 100644 --- a/ami/ml/post_processing/class_masking.py +++ b/ami/ml/post_processing/class_masking.py @@ -279,9 +279,11 @@ 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. + 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, diff --git a/ami/ml/post_processing/tests/test_class_masking.py b/ami/ml/post_processing/tests/test_class_masking.py index 7f2275282..0e9b28efa 100644 --- a/ami/ml/post_processing/tests/test_class_masking.py +++ b/ami/ml/post_processing/tests/test_class_masking.py @@ -525,32 +525,12 @@ def test_collection_scope_returns_each_classification_once(self): "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. + 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]) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index 3723f9c15..9366eef47 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -1651,62 +1651,38 @@ 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 _post_processing_job(self, name: str): - from ami.jobs.models import Job + 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=name, + 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() - return job - - def test_zero_progress_heartbeat_marks_the_stage_started(self): - """A heartbeat at exactly 0% still moves the stage out of CREATED. - ``Job.update_progress`` only promotes a stage once its progress exceeds - zero, so the setup heartbeat would otherwise leave a running job labelled - "Waiting to start". See #1376. - """ - from ami.jobs.models import JobState - - job = self._post_processing_job("stage status test") - 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.0) - - job.refresh_from_db() - stage = job.progress.stages[0] - self.assertEqual(stage.status, JobState.STARTED) - self.assertEqual(stage.status_label, "0% complete") - - def test_heartbeat_does_not_resurrect_a_finished_stage(self): - """Promotion applies only to a stage that has not started. - - A late heartbeat must not pull a stage that already reached a final state - back to STARTED, which would hide a failure or a cancellation. See #1376. - """ - from ami.jobs.models import Job, JobState + observed = {} - for final in (JobState.SUCCESS, JobState.FAILURE, JobState.REVOKED): - with self.subTest(final=final): - job = self._post_processing_job(f"stage status {final} test") - job.progress.update_stage("post_processing", status=final, progress=1.0) - job.save() + def _capture(self_task): + stage = self_task.job.progress.get_stage("post_processing") + observed["status"] = stage.status + observed["label"] = stage.status_label - task = SmallSizeFilterTask(job=job, source_image_collection_id=self.collection.pk, size_threshold=0.01) - task.update_progress(0.5) + with patch.object(SmallSizeFilterTask, "run", _capture): + PostProcessingJob.run(job) - job.refresh_from_db() - self.assertEqual(job.progress.stages[0].status, final) - Job.objects.filter(pk=job.pk).delete() + 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