Stop large class masking jobs from being cancelled before they report any progress#1376
Merged
Conversation
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 <noreply@anthropic.com>
✅ Deploy Preview for antenna-ssec canceled.
|
✅ Deploy Preview for antenna-preview canceled.
|
Contributor
|
Warning Review limit reached
Next review available in: 22 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughClass masking now counts and retrieves scoped classifications without ChangesClass masking execution
Agent guidance
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ClassMaskingTask
participant make_classifications_filtered_by_taxa_list
participant ProgressStage
ClassMaskingTask->>make_classifications_filtered_by_taxa_list: pass on_setup callback
make_classifications_filtered_by_taxa_list->>make_classifications_filtered_by_taxa_list: compute classification total
make_classifications_filtered_by_taxa_list->>ClassMaskingTask: invoke on_setup(total)
ClassMaskingTask->>ProgressStage: set initial progress and classifications_total
make_classifications_filtered_by_taxa_list-->>ClassMaskingTask: process first batch
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ami/ml/post_processing/class_masking.py`:
- Around line 342-345: Update the _on_setup callback to mark empty runs
complete: when total is 0, report 100% progress instead of 0%, while preserving
the existing 0% initialization for non-empty runs and the classifications_total
metric.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 781a2e37-02bf-41bc-bbec-cfc62c211c27
📒 Files selected for processing (3)
ami/ml/management/commands/run_class_masking.pyami/ml/post_processing/class_masking.pyami/ml/post_processing/tests/test_class_masking.py
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves the responsiveness and observability of the class-masking post-processing task by removing unnecessary .distinct() calls from large classification scope queries (eliminating an expensive SELECT DISTINCT over large array columns) and adding an early “setup” progress callback so operators can see the scope size before the first batch is processed.
Changes:
- Remove
.distinct()from class-masking scope queries (task and management command) to avoid costly de-duplication over largelogits/scoresarrays. - Add an
on_setup(total)callback to report the scope size before iterating classifications, improving early job feedback. - Add tests to pin the scope query’s non-duplicating behavior and verify setup callback ordering.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| ami/ml/post_processing/tests/test_class_masking.py | Adds tests for scope uniqueness, absence of DISTINCT, and setup-callback ordering. |
| ami/ml/post_processing/class_masking.py | Removes .distinct() from scope query; adds on_setup callback and wires it into ClassMaskingTask. |
| ami/ml/management/commands/run_class_masking.py | Removes .distinct() from preflight count to avoid slow startup. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ami/ml/post_processing/base.py`:
- Around line 77-85: Update the stage status handling in the progress-update
method around self.job.progress.update_stage: set JobState.STARTED only when the
stage is in a pre-start state, and preserve terminal or cancelling statuses such
as FAILURE, REVOKED, and CANCELING while still updating progress. Ensure
progress-only heartbeats cannot resurrect a completed, failed, or cancelled
stage.
In `@ami/ml/tests.py`:
- Around line 1677-1683: Update the test around
SmallSizeFilterTask.update_progress to use 0.0 progress instead of 0.5, while
preserving the refresh and assertions that the job stage transitions to STARTED
with the expected status label for the zero-progress heartbeat.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b4758c1c-4d6f-4678-9f8a-af4651090909
📒 Files selected for processing (6)
.agents/AGENTS.mdami/ml/management/commands/run_class_masking.pyami/ml/post_processing/base.pyami/ml/post_processing/class_masking.pyami/ml/post_processing/tests/test_class_masking.pyami/ml/tests.py
🚧 Files skipped from review as they are similar to previous changes (3)
- ami/ml/management/commands/run_class_masking.py
- ami/ml/post_processing/class_masking.py
- ami/ml/post_processing/tests/test_class_masking.py
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 <noreply@anthropic.com>
…beat 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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Large class masking runs were being cancelled roughly eleven minutes in, having reported no progress at all: the job page showed zero classifications checked. The cause was not the masking work itself but the query that decides which classifications are in scope, which asked the database to de-duplicate rows it could not have duplicated. On a scope of 55,530 classifications that took 208 seconds to count and just as long again before the first row could be read, so the run crossed the stalled-job cutoff before it ever committed a batch and reported progress.
Removing the de-duplication brings the same count back in half a second. A masking run now starts working almost immediately and reports the size of the job before it processes the first prediction, so an operator watching a long run can see how much was found rather than a stage sitting at zero.
Checking that claim in review turned up a smaller related gap: post-processing jobs never marked their stage started, so they read as "Waiting to start" until the first progress report. That is fixed here too, and it affects the small size filter as well.
The scope query returns exactly the same classifications, and the masking logic is untouched.
List of Changes
ClassMaskingTask._scoped_classificationsno longer calls.distinct()on either scoperun_class_maskingcommand's preflight count returns promptly instead of appearing to hangon_setupcallback reports the scope size before the first rowPostProcessingJob.runmarks the stage STARTED before the task runs, matching the other job types.agents/AGENTS.mdgains a rule on comment length and linking to the PRDetailed Description
Why de-duplication was both unnecessary and expensive
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. Filtering on a single occurrence or a single collection therefore cannot multiply rows, so there was nothing to de-duplicate. The two new scope tests pin exactly that, including the case a reader would worry about: a capture that belongs to several collections still contributes one row per classification.
The cost comes from what
SELECT DISTINCThas to compare. The selected columns include thelogitsandscoresarrays, and a production classifier can carry tens of thousands of categories, which puts those two arrays at roughly half a megabyte per row. The database was sorting all of them to prove uniqueness of a column set that already contains the primary key.De-duplication is also a blocking step, so it delayed the iteration pass as well as the count — no row could be returned until the whole sort finished. Both delays land before the first batch commits, and the first batch commit is where the run reports its first progress.
Measured, not estimated
Against a live database, on the scope of a run that had been cancelled:
.distinct().distinct().distinct()For context on the remaining setup work, expanding that classifier's category map (29,176 categories) took 10.4 seconds and reading the taxa list took 0.4 seconds. With the de-duplication gone, the work ahead of the first progress report is on the order of ten seconds rather than minutes.
The stage-status fix, and two corrections it went through
A stage that has started but not yet reported progress reads as "Waiting to start", because
get_status_labelrenders aCREATEDstage that way regardless of progress. Every job type is supposed to move its stage toSTARTEDat the top ofrun();PostProcessingJobwas the only one that did not, so its stage sat inCREATEDuntil the task's first progress report, which on a large masking run is minutes in.This took two rounds of review to land correctly, and the audit trail is worth keeping:
CREATEDas evidence of the stall. @copilot pointed out a healthy run looked identical, so that was not diagnostic; the evidence that stands is the counter, still at zero eleven minutes in.BasePostProcessingTask.update_progress, guarding the promotion to fire only fromCREATED. @coderabbitai correctly noted thatJob.update_progressalready promotes on any progress above zero, so my first test at 0.5 was vacuous, and that a per-heartbeat write is the wrong shape for a status transition.PostProcessingJob.run, next toupdate_status(STARTED), exactly asRegroupEventsJobandSourceImageCollectionPopulateJobalready do it. That is one line, it fires before the scope count rather than at the first heartbeat, and it removes the base-class guard entirely.test_post_processing_stage_is_started_before_the_task_runspins it.A note on the one
.distinct()that staysClassMaskingActionForm._algorithms_for_scopestill de-duplicates, and should: it joins fromAlgorithmacross a reverse relation to many classifications, so one algorithm genuinely does appear once per classification. It also selects only algorithm columns, so the comparison is cheap.Testing
ami.ml.post_processing,ami.ml.testsandami.jobspass, including five new tests.black,isortandflake8clean on the touched files.