Add bulk cancellation for active jobs - #685
Conversation
|
Warning Review limit reached
Next review available in: 4 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: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughBulk cancellation was added to the admin jobs list. The shared ChangesTask cancellation and bulk selection
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant JobListPage
participant DataTable
participant ConfirmDialog
participant AdminCancellationAPI
participant WorkerDispatcher
participant TaskStateManager
JobListPage->>DataTable: select active task rows
JobListPage->>ConfirmDialog: request bulk cancellation confirmation
ConfirmDialog-->>JobListPage: confirm cancellation
JobListPage->>AdminCancellationAPI: cancel selected task IDs
AdminCancellationAPI->>WorkerDispatcher: cancel task
WorkerDispatcher->>TaskStateManager: set_cancelled_if_active
TaskStateManager-->>WorkerDispatcher: conditional transition result
AdminCancellationAPI-->>JobListPage: cancellation response
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2973302e73
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
Reviewed. No correctness, security, or test-coverage issues — approving.
A few minor, non-blocking things worth a look (not required for merge):
bulkCancelMutation'sonErrorinlist.tsxis effectively dead code —mutationFnawaitsPromise.allSettled, so per-task failures are already captured into{ ok, failed }and the mutation itself won't reject in normal operation.- Styling inconsistency: "Cancel selected" uses
variant="outline", while the equivalent single-task cancel ondetail.tsxusesvariant="destructive"for the same class of action. - Bulk "select all" only selects the current page (TanStack default), and a partial cancel failure clears the whole selection rather than leaving the failed ones selected for retry. Both reasonable trade-offs for this PR's scope, just flagging for awareness.
…-104451 # Conflicts: # ui/src/pages/admin/jobs/list.test.tsx # ui/src/pages/admin/jobs/list.tsx
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
ui/src/pages/admin/jobs/list.test.tsx (1)
24-30: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
vi.importActualfor pure utility functions.Mocking pure functions like
isActiveStateduplicates business logic in the test file, which can lead to false positives if the actual implementation changes (e.g., if a new active status is introduced). Consider importing the actual module and only mocking the API calls.♻️ Proposed refactor
-vi.mock("`@/lib/api/jobs`", () => ({ - cancelTask: vi.fn(), - getQueueInfo: vi.fn(), - isActiveState: (state: string) => ["QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"].includes(state), - isTerminalState: (state: string) => ["COMPLETED", "FAILED", "CANCELLED"].includes(state), - listTasks: vi.fn(), -})); +vi.mock("`@/lib/api/jobs`", async () => { + const actual = await vi.importActual<typeof import("`@/lib/api/jobs`")>("`@/lib/api/jobs`"); + return { + ...actual, + cancelTask: vi.fn(), + getQueueInfo: vi.fn(), + listTasks: vi.fn(), + }; +});🤖 Prompt for 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. In `@ui/src/pages/admin/jobs/list.test.tsx` around lines 24 - 30, Update the jobs module mock in the test to use vi.importActual and retain the real isActiveState and isTerminalState implementations, while continuing to mock only cancelTask, getQueueInfo, and listTasks.ui/src/components/shared/data-table.tsx (1)
77-86: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider disabling the "Select visible rows" checkbox when no rows are selectable.
If a page only contains rows that cannot be selected (e.g., all tasks are completed), the header checkbox will have no effect when clicked. Disabling it in this state improves the user experience.
✨ Proposed UX improvement
- header: ({ table }) => ( + header: ({ table }) => { + const hasSelectableRows = table.getRowModel().rows.some((row) => row.getCanSelect()); + return ( <Checkbox checked={ table.getIsAllPageRowsSelected() || (table.getIsSomePageRowsSelected() && "indeterminate") } onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)} + disabled={!hasSelectableRows} aria-label="Select visible rows" /> - ), + ); + },🤖 Prompt for 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. In `@ui/src/components/shared/data-table.tsx` around lines 77 - 86, Update the header Checkbox in the table column definition to be disabled when the current page has no selectable rows, using the table selection APIs or row-selection predicate already used by the table. Preserve the existing checked state, toggle behavior, and aria label when at least one visible row is selectable.
🤖 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.
Nitpick comments:
In `@ui/src/components/shared/data-table.tsx`:
- Around line 77-86: Update the header Checkbox in the table column definition
to be disabled when the current page has no selectable rows, using the table
selection APIs or row-selection predicate already used by the table. Preserve
the existing checked state, toggle behavior, and aria label when at least one
visible row is selectable.
In `@ui/src/pages/admin/jobs/list.test.tsx`:
- Around line 24-30: Update the jobs module mock in the test to use
vi.importActual and retain the real isActiveState and isTerminalState
implementations, while continuing to mock only cancelTask, getQueueInfo, and
listTasks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: acc68d46-efc2-4841-bf24-e99dc7516783
📒 Files selected for processing (3)
ui/src/components/shared/data-table.tsxui/src/pages/admin/jobs/list.test.tsxui/src/pages/admin/jobs/list.tsx
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7d10ae650b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
Terminal states are duplicated across three files with no shared source of truth This PR adds a terminal-state check in three separate places, each spelled differently:
Same three strings, copy-pasted three times, nothing keeping them in sync. If a state is ever added/renamed, it's easy to update two of the three and end up with the route guard and the actor's own guard disagreeing about what's cancellable. This also isn't the first independent formulation of this boundary: Suggest hoisting one Possible redundant Ray call in state = await self._call(self._tsm.get_state.remote(task_id), ...)
if state in self._TERMINAL_STATES:
return False
obj_ref = await self._call(self._tsm.get_object_ref.remote(task_id), ...)
if obj_ref is None:
return False
state = await self._call(self._tsm.get_state.remote(task_id), ...)
if state in self._TERMINAL_STATES:
return False |
Terminal-state guards were spelled out independently in the route, dispatcher, and TaskStateManager. Consolidate into core.models.catalog.TERMINAL_TASK_STATES so the three stay in sync, and drop the redundant pre-obj_ref state check in WorkerDispatcher.cancel_task (the state check right before ray.cancel, plus the atomic set_cancelled_if_active guard, already cover it).
|
Good catch on both points. Fixed in commit 3d55607. Duplicated terminal states: hoisted a single Redundant Ray call: confirmed it was genuinely redundant — the state check right before |
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
Re-reviewing after the latest push (3d55607).
Status of everything raised in this PR's discussion:
- codex P1 (stale-state race on bulk cancel): fixed in b4bce26 — client re-fetches and re-filters by
isActiveStateright before cancelling, server-sidecancel_task/set_cancelled_if_activeatomically reject terminal tasks. - codex P2 (hidden selections surviving a search change): fixed in b4bce26 — search
onChangenow clearsrowSelection. - @aditykris — duplicated terminal-state literals across 3 files: fixed in 3d55607 — hoisted
TERMINAL_TASK_STATESintocore/models/catalog.py, all three call sites import it. - @aditykris — redundant Ray
get_statecall: fixed in the same commit — dropped the pre-get_object_refstate check; the remaining check + the atomicset_cancelled_if_activeguard are sufficient.
My earlier review's non-blocking notes still stand as optional/out-of-scope for this PR (not required for merge): the mutation's onError handler is effectively unreachable given Promise.allSettled, "Cancel selected" uses variant="outline" vs. detail.tsx's variant="destructive" for the single-task cancel, and a partial cancel failure clears the whole selection rather than leaving failed rows selected for retry.
No correctness, security, or test-coverage issues. Approving.
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 `@openrag/services/workers/dispatcher.py`:
- Around line 241-255: In openrag/services/workers/dispatcher.py lines 241-255,
update the cancellation flow to call set_cancelled_if_active atomically before
ray.cancel, remove the preceding get_state RPC and terminal-state check, and
return False without cancelling when the atomic update reports the task is not
active. In tests/unit/services/workers/test_dispatcher.py lines 218-239, remove
the get_state mock and assertions, configure set_cancelled_if_active to control
whether cancellation proceeds, and assert the updated call order and abort
behavior.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4e841814-3826-4f40-920b-b5027900a711
📒 Files selected for processing (9)
openrag/api/routers/admin/indexing.pyopenrag/core/indexing/dispatcher.pyopenrag/core/models/__init__.pyopenrag/core/models/catalog.pyopenrag/services/workers/dispatcher.pyopenrag/services/workers/task_state.pytests/unit/services/workers/test_dispatcher.pyui/src/pages/admin/jobs/list.test.tsxui/src/pages/admin/jobs/list.tsx
🚧 Files skipped from review as they are similar to previous changes (1)
- ui/src/pages/admin/jobs/list.tsx
Calling ray.cancel() before the set_cancelled_if_active RPC left a window where a killed worker could never report back if that RPC then failed, stranding the task active forever. Atomically claim the cancellation first and only send ray.cancel() once that succeeds, so a failed cancel signal just leaves a task marked CANCELLED that later self-corrects instead of a permanent zombie. Also drops the now fully-redundant get_state precheck.
# Conflicts: # ui/src/pages/admin/jobs/list.test.tsx
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
Re-reviewing after commit ded3534.
CodeRabbit flagged a real bug in my previous fix (ded3534 supersedes 3d55607): in WorkerDispatcher.cancel_task, calling ray.cancel() before the atomic set_cancelled_if_active RPC meant that if that RPC failed/timed out after the worker was killed, the task would be stuck active forever with no worker left to report completion — a permanent zombie. Fixed by claiming the cancellation atomically first and only sending ray.cancel() once that succeeds; a failed ray.cancel() now just leaves a task marked CANCELLED that self-corrects later instead of hanging forever. Also drops the now fully-redundant get_state precheck, so cancel_task is down to 2 remote calls. Tests updated to match, full unit suite (1721 tests) and all CI checks green.
Everything raised across this PR's review history is now addressed:
- codex P1 (stale-state race on bulk cancel) — fixed upstream (b4bce26)
- codex P2 (hidden selections surviving a search change) — fixed upstream (b4bce26)
- duplicated terminal-state literals across 3 files — fixed (3d55607)
- redundant
get_stateRPC — fixed, twice over (3d55607, then fully in ded3534) - zombie-task race on RPC failure after
ray.cancel()— fixed (ded3534)
No remaining correctness, security, or test-coverage issues. Approving.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 800910a411
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
Re-approving after commit 92d8a9a.
Codex found one more real race in my previous fix (ded3534): claiming cancellation before ray.cancel() closed the RPC-failure zombie risk, but opened a window where a QUEUED task's worker could start process_file and unconditionally overwrite the CANCELLED claim via set_state("SERIALIZING"/"COMPLETED") before the cancel signal took effect. hedhoud fixed this at the right layer in 92d8a9a: TaskStateManager.set_state now refuses to leave CANCELLED for any other state, so any worker write in that window is a no-op. Covered by a new regression test, verified locally alongside the existing dispatcher suite.
All CI checks green: lint, tests, api-tests, milvus-integration, layer-import-guard, CodeRabbit. No outstanding comments from CodeRabbit, codex, or aditykris. Approving.
`cancel_task` used to write CANCELLED to the jobs row unconditionally. Since
`get_object_ref` still answers for a finished task — the ref is only dropped
when the entry is evicted — a late `DELETE /task/{id}` reached that path for
work that had already succeeded, and durably recorded a COMPLETED file as
cancelled.
Rebasing onto #685 supplies the fix: `set_cancelled_if_active` refuses the
claim on a terminal task, and returning on that answer now short-circuits
before the durable write rather than after it. This pins the durable half,
which #685 alone does not cover.
`set_cancelled_if_active` is the only state write that does not call `_mark_terminal`. Every sibling does — `set_state` and `set_failed_if_not_cancelled` both register the transition — and eviction is driven entirely off `terminal_at`, so an entry that never lands there is retained for the lifetime of the detached actor: neither the 2000 cap nor the 1h TTL can reclaim it. Nothing recovers it afterwards. `ray.cancel` raises `CancelledError`, a `BaseException` that `process_file`'s `except Exception` sails straight past, so no later write ever re-registers the task. Every user-initiated cancel therefore leaks one `TaskInfo` — its details, its `user_index` entry, and its pinned `object_ref` — which is verbatim the unbounded growth this branch exists to fix. The method predates the eviction machinery this branch adds (it arrived with the bulk-cancel work in #685), and the two only met on the rebase; the gap is the merge, not either change on its own. `#685`'s cancel-all UI fires one cancel per selected task, so the leak scales with a single click. `test_cancelled_tasks_are_evictable` did not catch it because it drives `set_state(..., "CANCELLED")`, a path no cancellation actually takes. The new test uses the entry point `cancel_task` really calls.
Cancelling a task that had not started yet leaked its reserved file slot outright. After dispatch the worker owns the slot and only releases it in `IndexerWorker.process_file`'s `finally`; when `ray.cancel` retires the task before that body runs, the `finally` never executes and nothing else gives the slot back. The CANCELLED row is terminal, so the loss is indistinguishable from a clean cancel, and the uploader's quota is permanently one narrower. `cancel_task` could not simply release, because a task cancelled *mid-flight* does run its `finally` — releasing in both places would drive `file_count` below reality and hand out free quota. The two outcomes were only distinguishable by how far the worker happened to get. `TaskStateManager` settles it. It is a single-threaded Ray actor and already arbitrates this way (`set_cancelled_if_active`, `set_failed_if_not_cancelled`), so a one-shot `claim_quota_release` compare-and-set is the idiomatic primitive here: the worker and the canceller both ask, exactly one is told to release. The dispatcher records `quota_reserved` on the task at admission so the actor knows whether a slot is outstanding at all — `put_file` reuses an existing row and reserves nothing. The setup-failure release in `indexer_pool` goes through the same claim: a cancellation *during* setup reaches both it and `cancel_task`. This also closes the dispatch-partial-success double-release noted in the PR description — the worker and the request teardown can no longer both consume the same reservation. Failure handling stays biased the way the rest of the branch is: an unreachable arbiter releases anyway. An undercount is recoverable and self-heals under the #676 reconciliation; a leak is permanent and silent. `_release_cancelled_slot` is best-effort throughout, because a cancellation the user already asked for and already saw must not fail on bookkeeping. Worth doing now rather than deferring to #676: #685 shipped a cancel-all UI that fires one cancel per selected task, so the leak had become reachable at scale from a single click.
`cancel_task` used to write CANCELLED to the jobs row unconditionally. Since
`get_object_ref` still answers for a finished task — the ref is only dropped
when the entry is evicted — a late `DELETE /task/{id}` reached that path for
work that had already succeeded, and durably recorded a COMPLETED file as
cancelled.
Rebasing onto #685 supplies the fix: `set_cancelled_if_active` refuses the
claim on a terminal task, and returning on that answer now short-circuits
before the durable write rather than after it. This pins the durable half,
which #685 alone does not cover.
`set_cancelled_if_active` is the only state write that does not call `_mark_terminal`. Every sibling does — `set_state` and `set_failed_if_not_cancelled` both register the transition — and eviction is driven entirely off `terminal_at`, so an entry that never lands there is retained for the lifetime of the detached actor: neither the 2000 cap nor the 1h TTL can reclaim it. Nothing recovers it afterwards. `ray.cancel` raises `CancelledError`, a `BaseException` that `process_file`'s `except Exception` sails straight past, so no later write ever re-registers the task. Every user-initiated cancel therefore leaks one `TaskInfo` — its details, its `user_index` entry, and its pinned `object_ref` — which is verbatim the unbounded growth this branch exists to fix. The method predates the eviction machinery this branch adds (it arrived with the bulk-cancel work in #685), and the two only met on the rebase; the gap is the merge, not either change on its own. `#685`'s cancel-all UI fires one cancel per selected task, so the leak scales with a single click. `test_cancelled_tasks_are_evictable` did not catch it because it drives `set_state(..., "CANCELLED")`, a path no cancellation actually takes. The new test uses the entry point `cancel_task` really calls.
Cancelling a task that had not started yet leaked its reserved file slot outright. After dispatch the worker owns the slot and only releases it in `IndexerWorker.process_file`'s `finally`; when `ray.cancel` retires the task before that body runs, the `finally` never executes and nothing else gives the slot back. The CANCELLED row is terminal, so the loss is indistinguishable from a clean cancel, and the uploader's quota is permanently one narrower. `cancel_task` could not simply release, because a task cancelled *mid-flight* does run its `finally` — releasing in both places would drive `file_count` below reality and hand out free quota. The two outcomes were only distinguishable by how far the worker happened to get. `TaskStateManager` settles it. It is a single-threaded Ray actor and already arbitrates this way (`set_cancelled_if_active`, `set_failed_if_not_cancelled`), so a one-shot `claim_quota_release` compare-and-set is the idiomatic primitive here: the worker and the canceller both ask, exactly one is told to release. The dispatcher records `quota_reserved` on the task at admission so the actor knows whether a slot is outstanding at all — `put_file` reuses an existing row and reserves nothing. The setup-failure release in `indexer_pool` goes through the same claim: a cancellation *during* setup reaches both it and `cancel_task`. This also closes the dispatch-partial-success double-release noted in the PR description — the worker and the request teardown can no longer both consume the same reservation. Failure handling stays biased the way the rest of the branch is: an unreachable arbiter releases anyway. An undercount is recoverable and self-heals under the is best-effort throughout, because a cancellation the user already asked for and already saw must not fail on bookkeeping. Worth doing now rather than deferring to #676: #685 shipped a cancel-all UI that fires one cancel per selected task, so the leak had become reachable at scale from a single click.
3 participants
Context
Admins can cancel one active task from the detail page, but cancelling several active jobs still required API calls or scripts.
Change
Add safe bulk cancellation to the existing Jobs table:
Partially addresses #545.
Limitation
Retry is not implemented here. The current API does not expose a safe retry operation or enough durable submission data to recreate the original indexing request without risking duplicate or incorrect submissions. That should be a separate backend/API design before adding a retry button.
Validation
Summary by CodeRabbit