fix(gpu-arbiter): op kinds, shadow queue, cancel_op, position, + 4 pre-deploy defect fixes (#1864)#1984
fix(gpu-arbiter): op kinds, shadow queue, cancel_op, position, + 4 pre-deploy defect fixes (#1864)#1984hognek wants to merge 8 commits into
Conversation
|
Warning Review limit reached
Next review available in: 13 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 Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughGpuArbiter now carries operation metadata through submission and queueing, computes operation-aware queue positions, exposes non-destructive snapshots, and supports cancellation of queued or running tasks with VRAM cleanup. ChangesGPU queue operations
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant GpuArbiter
participant VramReservationManager
participant TaskFuture
Caller->>GpuArbiter: submit_gpu(task, op, model, backend_name)
GpuArbiter->>VramReservationManager: reserve required VRAM
VramReservationManager-->>GpuArbiter: admission result
GpuArbiter->>TaskFuture: queue or run task
Caller->>GpuArbiter: cancel_op(task_id)
GpuArbiter->>TaskFuture: cancel queued future
GpuArbiter->>VramReservationManager: release running-task reservation
Possibly related PRs
🚥 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 |
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
|
|
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| if entry.task.id in self._cancelled_ids: | ||
| self._cancelled_ids.discard(entry.task.id) | ||
| continue | ||
| self._queued_entries.pop(entry.task.id, None) |
There was a problem hiding this comment.
WARNING: Shadow dict _queued_entries is dropped on line 579 but only re-added at line 624 (in the retry re-queue), and _drain_queue awaits asyncio.to_thread at line 580 between them. During that await window the task is still logically queued in self._queue but ABSENT from _queued_entries. A concurrent cancel_op(task_id) will find entry is None, fall through to the running check, find the task is not running, and return False even though the task is genuinely queued. queue_position/queue_snapshot also report it as not-queued mid-drain.
This makes cancel_op and the position/snapshot reads racy against the 2s drain cycle. Consider keeping _queued_entries authoritative and only mutating it atomically with self._queue (e.g. re-add the entry to _queued_entries when it goes into the local retry list, or guard the whole drain step under the same lock used elsewhere), so reads/cancellation stay consistent.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| try: | ||
| return await done | ||
| except asyncio.CancelledError: | ||
| self._queued_entries.pop(task.id, None) |
There was a problem hiding this comment.
SUGGESTION: Submitter-side cancellation (the await done raising CancelledError) pops _queued_entries here but leaves the entry in the real self._queue. Nothing re-checks _cancelled_ids for an entry cancelled this way, so on the next drain that entry is admitted and _run_gpu_task starts — running a task whose submitter already gave up (its _arbiter_future is discarded via _propagate's ct.cancelled() early-return). Consider also recording the id in _cancelled_ids in this branch (or removing the entry from self._queue) so the orphaned queue entry is skipped rather than executed.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| self, task: Task, required_vram_mb: int = 0, | ||
| evictable: bool = False, resource_id: str | None = None, | ||
| required_gpu_arch: str | None = None, | ||
| op: str = "inference", |
There was a problem hiding this comment.
SUGGESTION: op is accepted unchecked. Any value other than "inference" silently falls into the else (global) branch of queue_position (line 664), so a typo like op="infere" or op="LOAD" is treated as a load and counted globally rather than flagged. Consider validating against the allowed set {"inference", "load"} in submit_gpu (and/or normalizing case) to fail fast on misuse. Also note: inference tasks with model=None are counted as the same model (line 665), so two model=None inference tasks occupy positions 1 and 2 behind each other — verify that is the intended per-model semantics when model is unset.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 1 Issue Found | Recommendation: Address before merge OverviewIncremental review of commit
Issue Details (click to expand)SUGGESTION
Files Reviewed (1 file)
Fix these issues in Kilo Cloud Previous Review Summaries (7 snapshots, latest commit 9680011)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 9680011)Status: 1 Issue Found | Recommendation: Address before merge OverviewIncremental review of commit
Issue Details (click to expand)WARNING
Files Reviewed (1 file)
Fix these issues in Kilo Cloud Previous review (commit 1331ea0)Status: 2 Issues Found | Recommendation: Address before merge OverviewIncremental review of commit
Issue Details (click to expand)SUGGESTION
Files Reviewed (2 files)
Fix these issues in Kilo Cloud Previous review (commit c41ac97)Status: No Issues Found | Recommendation: Merge OverviewIncremental review of commit
Both changes mirror the existing admitted-path handling (lines 603-606) and the retry-branch cancel re-check (lines 641-643), giving symmetric coverage of the cancellation race across all branches. All 5 previously-active findings (lines 313, 376, 631, 712, N/A) are outside the incremental change scope and remain unchanged; none are reintroduced or affected by the new code. Files Reviewed (2 files)
Previous review (commit 53a7f53)Status: No Issues Found | Recommendation: Merge OverviewIncremental review of commit
Notes on the incremental change
Files Reviewed (1 file)
Previous review (commit dc0a631)Status: No Issues Found | Recommendation: Merge OverviewIncremental review of commit
Resolved in this incremental diff
Resolved by earlier commits in this PR (carried from prior review)
Files Reviewed (2 files)
Previous review (commit 4374588)Status: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously reported issues — resolved by this incremental diff
Files Reviewed (2 files)
Fix these issues in Kilo Cloud Previous review (commit d441658)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (2 files)
Reviewed by hy3:free · Input: 35K · Output: 2.6K · Cached: 111.1K |
…date op
- _drain_queue: keep entry in _queued_entries during admission check so
concurrent cancel_op finds it (was popped too early before asyncio.to_thread).
Pop from shadow dict only after admission is confirmed (line 621).
- submit_gpu CancelledError handler: add task_id to _cancelled_ids so the
drain loop skips the orphaned queue entry instead of admitting it.
- submit_gpu: validate op against {'inference','load'}, raising ValueError
for unknown values so typos don't silently fall through to global branch.
Kilo review items on PR jaylfc#1984 (GPU queue A2).
Targeted: 29/29 pass.
Kilo review fixes (commit 4374588)Three issues from Kilo at 09:58:
Targeted tests: 29/29 pass. |
|
|
||
| t.add_done_callback(_propagate) | ||
|
|
||
| self._queued_entries.pop(entry.task.id, None) # promoted → _running |
There was a problem hiding this comment.
SUGGESTION: cancel_op (line 678) pops the task from _queued_entries and cancels its future, but _drain_queue already passed the _cancelled_ids check at line 583 and still holds the entry object locally. If cancel_op runs in the window between line 583 and here, the entry is gone from _queued_entries so the check at 583 won't skip it on the next drain, but on the current drain the task is still admitted and _run_gpu_task is spawned (line 600) and runs to completion — despite cancellation. The future is cancelled, yet a background task executes a workload the caller abandoned.
Consider re-checking self._cancelled_ids (or re-confirming presence in _queued_entries) after the admission await at line 586 and before spawning _run_gpu_task, or have cancel_op also signal running/in-flight drain intent under self._running_lock, so an in-progress admission honors a late cancel.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…ancelled_ids after await cancel_op pops from _queued_entries and cancels the future while _drain_queue is suspended on _reserve_and_check — the entry is still admitted despite cancellation. Add a second _cancelled_ids check after the await window so a concurrent cancel_op is caught before the task is spawned. Kilo SUGGESTION on PR jaylfc#1984 — task t_27501d36
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_gpu_arbiter_queue_ops.py (1)
82-92: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover cancellation consistency across both queue representations.
Add regressions that:
- cancel with
max_queue_size=1while processing is paused/not started, then successfully enqueue another task;- cancel during a failed asynchronous admission check and verify snapshots remain empty after the next drain.
🤖 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 `@tests/test_gpu_arbiter_queue_ops.py` around lines 82 - 92, Extend cancellation regression coverage around GpuArbiter.cancel_op: add a max_queue_size=1 case with processing paused or not started, cancel the queued task, and verify a replacement task can be enqueued; also add a case cancelling during a failed asynchronous admission check, then drain and assert both queue representations and snapshots remain empty. Reuse existing arbiter submission, cancellation, drain, and snapshot helpers.
🤖 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 `@tinyagentos/scheduler/gpu_arbiter.py`:
- Around line 688-695: Make cancellation atomic across the physical queue,
shadow entries, and tombstones: update cancel_op to reclaim or compact the
physical queued entry immediately so capacity is reusable, and update the
admission retry/reinsertion flow around the failed-await path to skip
reinsertion when cancellation occurred. In tests/test_gpu_arbiter_queue_ops.py
lines 82-92, add coverage for immediate capacity reuse and cancellation during
failed admission.
---
Nitpick comments:
In `@tests/test_gpu_arbiter_queue_ops.py`:
- Around line 82-92: Extend cancellation regression coverage around
GpuArbiter.cancel_op: add a max_queue_size=1 case with processing paused or not
started, cancel the queued task, and verify a replacement task can be enqueued;
also add a case cancelling during a failed asynchronous admission check, then
drain and assert both queue representations and snapshots remain empty. Reuse
existing arbiter submission, cancellation, drain, and snapshot helpers.
🪄 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 Plus
Run ID: 71342dc9-67e6-40cd-9dd5-e061233a8bbc
📒 Files selected for processing (2)
tests/test_gpu_arbiter_queue_ops.pytinyagentos/scheduler/gpu_arbiter.py
Same race as the admission-path fix: cancel_op may fire during the admission-await window (failed-admission path). The retry list re-queues cancelled entries because it never checks _cancelled_ids. Add the guard so cancelled tasks aren't resurrected. CodeRabbit finding on PR jaylfc#1984
|
Applied the Kilo SUGGESTION (task t_27501d36): two fixes for the cancel_op/_drain_queue race. Fix 1 — Admission path (line 603): Re-check Fix 2 — Retry/re-queue path (line 641): Same Both fixes follow the existing Tests: 22/22 targeted (test_gpu_arbiter_queue_ops + test_gpu_arbiter_894) pass. |
…dow queue, tombstones - cancel_op: release VRAM reservation immediately so capacity is freed rather than leaking until _drain_queue eventually skips the entry. - retry/re-queue loop: re-check _cancelled_ids after put_nowait+dict insert to prevent TOCTOU resurrection of cancelled shadow entries. CodeRabbit finding on PR jaylfc#1984 (GPU queue A2).
…uard, resource_id threading, lease renewal - _process_queue: wrap _drain_queue in try/except so one bad task cannot permanently kill the queue processor loop (defect 1). - _check_cluster_admission: skip workers whose free_vram_mb is None (non-NVIDIA workers) to prevent TypeErrors (defect 2). - _drain_queue: thread resource_id through _QueuedGpuTask so drained tasks hold a cluster lease on the correct resource (defect 3). - _run_gpu_task: add _renew_lease_loop background coroutine that periodically calls renew_lease every 200 s so the lease never expires mid-run (defect 4). Fixes: jaylfc#1864
hognek
changed the title
| renewed = await self._cluster_manager.renew_lease( | ||
| lease_id, ttl_seconds=300, | ||
| ) | ||
| if renewed is None: |
There was a problem hiding this comment.
SUGGESTION: When renew_lease returns None the loop returns and the renew task ends, but _run_gpu_task keeps executing the GPU workload with no active lease. The cluster manager may have already revoked/reclaimed the VRAM or device for that lease, so the still-running task could silently lose its resource mid-inference and fail or corrupt state. Consider signaling the running task (e.g. set a flag the run loop checks, or cancel the task) when the lease cannot be renewed, rather than only logging and returning.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| await self._drain_queue() | ||
| try: | ||
| await self._drain_queue() | ||
| except Exception: |
There was a problem hiding this comment.
SUGGESTION: The bare except Exception: around _drain_queue() swallows every error — including programming bugs like KeyError/AttributeError/TypeError — and the 2s loop then retries forever. A persistent bug would spam logs every 2 seconds and stay masked indefinitely. Consider narrowing the catch (or at least logging exc_info/repr(e) and adding a backoff / consecutive-failure threshold) so real defects surface instead of being hidden behind logger.exception.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…rror handling Two Kilo SUGGESTIONS from PR jaylfc#1984 review: 1. gpu_arbiter.py:438 — When renew_lease returns None, the renew loop silently returns but _run_gpu_task keeps executing GPU work with no active lease. Fix: _renew_lease_loop now accepts an optional task_to_cancel parameter and cancels the parent asyncio Task on renewal failure or expiry. 2. gpu_arbiter.py:617 — Bare except Exception around _drain_queue() swallows programming bugs (KeyError, TypeError, etc.) and retries forever in a tight 2s loop. Fix: narrow to transient errors (NoResourceAvailableError, TimeoutError, OSError), add exponential backoff (2s→60s), and stop queue processor after 10 consecutive failures.
| "times — stopping queue processor", | ||
| consecutive_failures, | ||
| ) | ||
| return |
There was a problem hiding this comment.
WARNING: Permanent return on MAX_CONSECUTIVE_FAILURES permanently kills the queue processor.
_process_queue is started exactly once in start() (line 145) and there is no supervisor/restart logic for _queue_processor_task. Previously the loop swallowed errors and retried forever; now, after 10 consecutive (transient) NoResourceAvailableError/TimeoutError/OSError — e.g. a temporarily unavailable cluster manager, a brief network partition, or a maintenance window — the processor returns and the GPU admission queue is silently dead for the rest of the process lifetime. No queued tasks are ever admitted again, with no log line saying the processor exited for a recoverable reason beyond one logger.critical.
Consider either (a) not permanently exiting — reset consecutive_failures and back off but keep the loop alive, or (b) if a hard stop is intended, surface a recoverable restart path (supervisor task or a health check that re-calls start()) so transient failures don't cause permanent admission starvation.
…RES instead of dying permanently Previously _process_queue returned after 10 consecutive transient failures (NoResourceAvailableError, TimeoutError, OSError), permanently killing the queue processor with no restart logic. A temporary cluster outage or GPU restart would kill the queue forever until manual intervention. Wrap the inner drain loop in an outer restart loop that waits a 300s cooldown after hitting the failure limit, then re-enters the inner loop with a fresh failure counter. CancelledError propagation is preserved. Kilo WARNING jaylfc#1984
|
Addressed Kilo WARNING: queue processor permanently dies after MAX_CONSECUTIVE_FAILURES. What changed: Wrapped the inner drain loop in an outer restart loop. After 10 consecutive transient failures (NoResourceAvailableError, TimeoutError, OSError), the processor now waits a 300s cooldown then restarts the inner loop with a fresh failure counter — instead of returning permanently. CancelledError propagation is preserved. Before: All 29 GPU arbiter tests + 192 scheduler/cluster tests pass. |
| async def _process_queue(self) -> None: | ||
| try: | ||
| MAX_CONSECUTIVE_FAILURES = 10 | ||
| COOLDOWN_BACKOFF = 300 # seconds to wait before restarting the processor |
There was a problem hiding this comment.
SUGGESTION: COOLDOWN_BACKOFF = 300 (5 min) imposes a full admission freeze after the inner loop exhausts MAX_CONSECUTIVE_FAILURES. Combined with the in-loop exponential backoff (capped at 60s), this means after ~10 failures the queue stops admitting for 5 minutes, even if the underlying failure was transient (e.g. a brief cluster blip). Consider making this backoff shorter, jittered, or exponential/limited so transient failures don't cause a 5-minute admission starvation window.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
1 participant
Summary
GPU arbiter operational improvements and pre-deploy defect fixes, all gated behind
TAOS_GPU_QUEUE(off by default):Earlier commits (A2 operations slice)
inference/load) — validated on submit, exposed in queue snapshot + position_queued_entries) — O(1) lookups for cancel_op, position, snapshotcancel_op— cancel queued tasks with VRAM reservation release; evict running tasks; atomic across physical queue + shadow dict + tombstonesqueue_position— 1-based position, global for loads, per-model for inferencequeue_snapshot— non-destructive read with position + op/model/backend fields_cancelled_idsafter await windows in drain + retry loopsNew (pre-deploy defect fixes, #1864)
_process_queuewraps_drain_queuein try/except so one bad task doesn't permanently kill the loop._check_cluster_admissionskips non-NVIDIA workers (free_vram_mb is None)._QueuedGpuTask, threaded throughsubmit_gpu→_drain_queue→_run_gpu_task, so drained tasks hold a cluster lease on the correct resource._renew_lease_loopbackground coroutine callsrenew_leaseevery 200 s; stops in finally.Testing
Task: t_68584724 (defect fixes), t_c34f79c7 (cancellation atomicity), prior A2 work
Fixes: #1864