fix(cluster): async drain_worker/cancel_drain + _lease_lock on force-release + wire arbiter task cancellation#1718
fix(cluster): async drain_worker/cancel_drain + _lease_lock on force-release + wire arbiter task cancellation#1718hognek wants to merge 17 commits into
Conversation
…ction Add GpuArbiter module that wraps the resource Scheduler with GPU-specific admission control to prevent concurrent-load driver crashes (Xid 62). Key features: - VRAM-aware admission: checks local VRAM probes and cluster worker leases before admitting GPU tasks - Priority queue: pending GPU tasks wait when VRAM is insufficient, dequeued in priority order when resources free up - Eviction: lower-priority running tasks can be evicted to make room for higher-priority work - Lease integration: claims/releases GPU leases via ClusterManager for distributed coordination - GPU resource registration: discovery.py now registers gpu-cuda-N resources when GPU backends are healthy - GPU_POTENTIAL_CAPABILITIES constant for UI latent-capability display 19 unit tests pass. Builds on Slice 1 (VRAM endpoint, jaylfc#893 leases). Task: t_d7208884. Fixes jaylfc#894.
… futures _evict_task previously only cancelled the _arbiter_future on the Task object. Directly-admitted tasks never have one set, and for queued-then- admitted tasks the future is resolved by _drain_queue before _run_gpu_task completes. In both cases the running GPU work keeps executing and VRAM is never freed until natural completion. Now _run_gpu_task registers its asyncio.Task in _running_tasks, and _evict_task cancels that Task directly. CancelledError propagates to the payload, the finally block cleans up VRAM/leases, and eviction actually stops running GPU work. - Added _running_tasks dict to track asyncio.Tasks - _run_gpu_task registers current task and handles CancelledError - _evict_task cancels the asyncio.Task (in addition to _arbiter_future) - _drain_queue catches CancelledError from evicted tasks to keep queue processor alive - Finally block checks if task was already evicted (entry popped) before releasing lease — prevents double-release Fixes jaylfc#894.
_evict_task now uses pop(task_id, None) atomically instead of the check-then-pop pattern, eliminating the TOCTOU KeyError that could occur when _run_gpu_task's finally block pops the entry between the existence check and the pop. Ownership is now explicit: whoever pops self._running first releases the lease. _evict_task's pop is the sole lease releaser for evicted tasks; _run_gpu_task's finally only releases on normal completion (when it still finds the entry). Comments document the invariant. Refs jaylfc#894. Task: t_0c86d08e.
…ncompatible with stack The file was incorrectly brought in from the jaylfc#1689 version and references VramAllocation, total_vram_mb, headroom_mb, and evict_callback that don't exist in the stack's GpuArbiter. test_gpu_arbiter_894.py covers the same territory with the correct API.
…M reservation _addmission_check read live nvidia-smi free VRAM with no reservation, so two concurrent submit_gpu calls could both pass before either model loaded — the exact concurrent-load crash (NVIDIA Xid 62) the arbiter exists to prevent. Add _reserved_vram_mb counter and _pending_reservations dict, guarded by _reservation_lock. _reserve_and_check() atomically checks admission against (probe - _reserved_vram_mb) and reserves on success. Reservations are released in _run_gpu_task's finally block and in _evict_task. _stats() now includes reserved_vram_mb and pending_reservations for observability. Fixes: PR jaylfc#1683 review feedback #2 Tests: 168 pass (25 arbiter + 143 cluster/worker)
…ed VRAM in cluster path taOS jaylfc#1705 — reservation leak fix: - Move lease claim and _running registration inside try block so the finally always releases the VRAM reservation, even on claim failure. - Await claim_lease / release_lease (both are async on ClusterManager). - Make _evict_task and evict_lowest_priority async (they await release_lease). - Subtract _reserved_vram_mb in the cluster-mode admission path, matching the local-path behaviour so in-flight reservations are accounted for. - Update tests: FakeClusterManager methods made async, all _evict_task and evict_lowest_priority calls now awaited. Tests: 16/16 arbiter + 111/111 cluster pass.
…iter _drain_queue no longer blocks on _run_gpu_task. The admitted task is spawned as a background asyncio Task with a done-callback that propagates the result/exception to the submitter's _arbiter_future. This keeps the drain loop responsive on subsequent ticks so eviction-to-make-room can kick in when higher-priority tasks arrive while VRAM is full. Also adds eviction-to-make-room: when admission fails for a queued task the drain now calls evict_lowest_priority with the task's priority as the floor, then re-checks admission. Lower-priority running tasks are evicted to make room; higher-priority runners are left alone. Queue processing remains intentionally serial — only one task is admitted per drain cycle to avoid flooding the GPU with concurrent loads. PR jaylfc#1683 review feedback jaylfc#4.
taOS jaylfc#1706 — TOCTOU fix: - After eviction-to-make-room frees VRAM, the re-admit in _drain_queue used _check_admission (no reservation) instead of _reserve_and_check (atomic reservation). This spawned tasks with no reservation entry, re-opening the TOCTOU window and causing _reserved_vram_mb drift. - Changed to _reserve_and_check so post-eviction admission is atomic and properly tracked. - Also fixes the missing await on evict_lowest_priority (made async by the base jaylfc#1705 fix) and sync _evict_task calls in new tests. Tests: 22/22 arbiter pass.
…biter Make evict_lowest_priority, _evict_task, release_tasks_for_worker, stats, and running_tasks async and guard _running dict access with _running_lock. Previously _running was mutated under the lock in _run_gpu_task but read/popped without it in eviction paths — safe under single-thread-no-await but fragile. Lease release and task cancellation happen outside the lock to avoid deadlock. Update all callers: drain_worker → async, route handlers, and tests. 69 targeted tests pass.
- Added vram_probe=_probe_nvidia_vram so the arbiter uses actual GPU VRAM instead of the default (0,0) probe that makes all checks pass. - Wired cluster_manager._gpu_arbiter so eviction paths that access the cluster manager can reach the arbiter. - Added gpu_arbiter.stop() call in shutdown sequence. - Added _gpu_arbiter attribute to ClusterManager.__init__.
taOS jaylfc#1707 — _running_lock coverage fix: - _evict_task now holds _running_lock during the pop, preventing races with concurrent drain operations. - stats() made async (needed for _running_lock consistency). - All async call sites now properly awaited (evict_lowest_priority, _evict_task, stats, claim_lease, release_lease). Tests: 22/22 arbiter pass.
…, restart Adds drain_worker/cancel_drain to ClusterManager for graceful worker detach without dropping inflight GPU tasks (taOS jaylfc#890). ClusterManager: - drain_worker(name, graceful=True): worker enters 'draining' status, excluded from routing/catalog/lease claims. When graceful=False, all leases are force-released and worker is marked offline. - cancel_drain(name): returns draining worker to 'online'. - _monitor_loop: auto-completes drain when all leases released; force-finishes stale drains on heartbeat timeout. - _worker_for_resource, get_workers_for_capability, aggregate_catalog: all exclude draining workers. Routes: - POST /api/cluster/workers/{name}/drain — begin graceful/force drain - POST /api/cluster/workers/{name}/cancel-drain — cancel in-progress drain - POST /api/cluster/workers/{name}/update — full auto-update orchestration: drain → deploy update-worker → restart → re-register Tests: 9 new tests (25 total in test_cluster.py), 191 cluster tests pass.
Replace startswith(name + ':') with _parse_resource_id exact match to prevent draining worker 'foo' from releasing leases belonging to worker 'foo-bar' (same prefix-collision class as TOCTOU fixes).
taOS jaylfc#796 features extracted from jaylfc#1689 and rebased onto the arbiter stack: 1. pause/resume queue control: GpuArbiter.pause()/resume() halt/restart queue processing. Running tasks finish; queued tasks wait. Paused queue skips _drain_queue in _process_queue loop. 2. Hardware-aware LLM admission: submit_gpu() accepts required_gpu_arch (e.g. 'sm_86'). _check_gpu_arch_compatibility() checks cluster workers for compatible GPU hardware via model name and compute_cap field before admitting tasks.
…r_894 _evict_task was made async in commit 1789ad0 (running-lock fix). Three cleanup calls were missing await.
…on force-release, wire arbiter task cancellation jaylfc#1690 lock fix: drain_worker force-release and _monitor_loop stale-drain path now hold _lease_lock during self._leases mutation, serialized with claim/release/sweep. drain_worker and cancel_drain are now async (consistent with claim_lease/release_lease — all callers were already in async handlers). Cross-cutting wiring: added GpuArbiter.cancel_running_for_leases() and wired it into drain_worker(graceful=False) and _monitor_loop stale-drain path. When leases are force-released, any running GPU arbiter tasks for those leases are cancelled via _evict_task — eviction is no longer a prod no-op when the arbiter is running. Added TYPE_CHECKING import for GpuArbiter type annotation on ClusterManager._gpu_arbiter.
|
Warning Review limit reached
Next review available in: 25 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 (11)
✨ 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}
| self._leases.pop(lid, None) | ||
| released += 1 | ||
| logger.info("Worker '%s' drain: force-released %d leases", name, released) | ||
| worker.status = "offline" |
There was a problem hiding this comment.
WARNING: worker.status = "offline" is set after the async with self._lease_lock: block exits. The lease pops are serialized under the lock, but the status transition is not, so a concurrent claim_lease (or any reader that filters on worker.status) can observe the worker as still "draining" while the leases are already gone, producing a transient state where _worker_for_resource rejects new claims for a worker that no longer owns any leases. Move the status flip inside the lock block (or set it before acquiring the lock) so the lease-table mutation and the status transition are atomic with respect to readers.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| ] | ||
| for lid in lids: | ||
| self._leases.pop(lid, None) | ||
| worker.status = "offline" |
There was a problem hiding this comment.
CRITICAL: Race in stale-drain path — worker.status = "offline" is set after releasing self._lease_lock. Because _monitor_loop snapshots workers via list(self._workers.values()) and the lock is only held for the lease-table mutation, the next 5 s tick (or any concurrent iteration that re-reads worker.status) can see the worker still in "draining" with zero active leases and re-enter the if not active_leases: branch above (line 641), emitting a duplicate worker.leave notification and a second "all leases released" log line. Set worker.status = "offline" while still holding self._lease_lock (or before the lock block), so the lease-pop and the status flip are atomic with respect to the monitor's next read.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # released leases (taOS cross-cutting wiring). | ||
| if lids and self._gpu_arbiter is not None: | ||
| try: | ||
| await self._gpu_arbiter.cancel_running_for_leases(set(lids)) |
There was a problem hiding this comment.
WARNING: cancel_running_for_leases(set(lids)) is invoked after _lease_lock is released and after worker.status = "offline". The list lids was captured under the lock so it's accurate, but the await here runs without any lock held — combined with the status being already "offline", a concurrent register_worker for the same name could race against the cancellation. Worth holding a brief status guard, or at minimum logging when set(lids) is empty after lock release (no leases to cancel) to avoid a noisy "cancelled 0 tasks" trace. Also note: this call is reached only when the worker is stale (heartbeat timeout); the in-flight _run_gpu_task may have already crashed — the except Exception correctly swallows but the operator gets no signal that a running LLM was force-killed.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| """ | ||
| if not lease_ids: | ||
| return 0 | ||
| victim_ids: list[str] = [] |
There was a problem hiding this comment.
SUGGESTION: cancel_running_for_leases reads self._running under self._running_lock, drops the lock, then awaits _evict_task per victim. Between releasing the lock and the _evict_task await, a concurrent natural completion can remove the entry from _running, so _evict_task(task_id) silently returns 0 and the cancellation is dropped (benign for completed tasks, but the caller can't tell). Consider returning distinct counts (e.g. {cancelled, already_done, missing}) or just logging len(victim_ids) - cancelled as already_completed so operators can see force-kill vs natural-completion. Also: this method iterates self._running.items() which is a dict[str, tuple] — fine, but the unpacking _task, lease_id, _pri, _vram discards lease_id without using it after the membership test; you could simplify with if any(lease_id in lease_ids for ...) if you only need the task_ids.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 1 Issue Found | Recommendation: Approve with minor suggestion Overview
Issue Details (click to expand)SUGGESTION
Previous Issues — Resolved by b6e17e6
Files Reviewed (2 files)
Fix these issues in Kilo Cloud Previous Review Summary (commit 814e5fd)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 814e5fd)Status: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (4 files)
Reviewed by minimax-m3 · Input: 41.5K · Output: 6.5K · Cached: 738.8K |
|
|
… (cancelled, already_completed) from cancel_running_for_leases Addresses Kilo review findings on PR jaylfc#1718: CRITICAL: Race in stale-drain path — worker.status='offline' was set after releasing _lease_lock, allowing duplicate worker.leave on next monitor tick. Now set inside the lock (both drain_worker force-release and _monitor_loop stale-drain paths). WARNING: Same pattern in drain_worker(graceful=False) — fixed. SUGGESTION: cancel_running_for_leases now returns (cancelled, already_completed) so operators can distinguish force-kills from natural completions. Callers log both counts.
|
Addressed all 4 Kilo review findings in b6e17e6:
47/47 tests pass. |
| if await self._evict_task(task_id): | ||
| cancelled += 1 | ||
| else: | ||
| already_completed += 1 |
There was a problem hiding this comment.
SUGGESTION: already_completed is misleading — _evict_task returns 0 whenever task_id not in self._running, which includes not only natural completion but also prior eviction by evict_lowest_priority, error-path cleanup, or a concurrent admin action. Renaming to already_gone or splitting into {cancelled, completed, missing} would let operators distinguish a normal completion from an unexpected removal. Worth considering before this signal becomes load-bearing for alerting.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
Did a full SET review of the arbiter epic across every branch (#1683/1690/1705/1706/1707/1718/1725/1726/1727). #1718 is the one to converge on: it is the rebased superset composing #1683 + #1705 + #1706 + #1707 + #1690 plus pause/resume and the blocker-2/4 fixes, and kilo is clean on it (0 critical, 0 warning). My four original composed-state blockers are all confirmed fixed in #1718:
Thank you - the mechanical blockers landed. But the epic is not mergeable as a set yet. Blocking:
MERGE PLAN (I will land it as a coherent package, not piecemeal):
Once 1-6 are green together I will merge #1718 + #1726 + #1725 as the package. Solid work clearing the mechanical blockers; the rest is integration + the eviction-ordering safety. |
|
Agreed on converging everything on #1718 — the composed-state analysis is exactly right, thanks for the thorough SET pass. On item 1 (the #185 gate): let's do it properly. Wire a real admission path through One shape question before I build the reconciliation, so I match your intended API the first time: should the arbiter be the front door with I'll land item 2's eviction-ordering fix in the same change, since wiring Items 3 (#1725 min_ram_mb→VRAM estimate + the Target sequence, landing as the one package you described: #1718 (eviction-ordering + resp.json + submit_gpu wiring) → #1726 → #1725 reconciled into the arbiter. Say the word if you'd rather I split the |
|
Answering the shape question so you can build the reconciliation once: make the arbiter the front door. The whole point of #185 is one VRAM authority. If On sequencing: your target is good and I will take it as the package. Keeping item 2's eviction-ordering fix in the same change as the The one hard constraint from #185: nothing merges piecemeal. #1718 (mechanical fixes plus eviction-ordering plus resp.json plus submit_gpu wiring) plus #1726 (pop inside |
|
Hi @hognek — I did a full cross-PR pass over the arbiter epic (#1705/#1706/#1707/#1718/#1759 + #1690/#1726) to untangle it before merging, since they all rewrite Supersession (verified by commit ancestry, not titles): #1705 ⊂ #1706 ⊂ #1707 ⊂ #1718 is a strict git-ancestor chain — the first three are older snapshots of this same branch, so I am closing them as superseded (pointing here). #1759 branches off separately into a competing design (a second class also named The keepers: #1690 (worker auto-update) merges clean and is the right home for that feature; #1726 (release leases on unregister) is a real fix (one trivial test conflict). Both land on their own. The one architectural blocker on this PR: dev already merged the epic's first reservation layer — Correctness items to fold in the rework (all in
Full write-up is long; happy to drop the rest (D-6/7 queued-caller-on-stop, D-8 double-count, D-9 |
|
Hi @hognek, consolidating what is left to close the arbiter set as one package, since #1726 rebased cleanly and #1690 now gates the destructive routes with _require_admin (that closes the earlier auth CRITICAL, thanks). Remaining before I do the SET re-review and merge #1718 + #1726 + #1725 together:
Optional, not blocking: rate limiting on the /drain, /cancel-drain, /update routes (already admin-gated). Ping me when #1718 + #1726 + #1725 are in final shape and I will re-review the composed state end to end. Nothing merges piecemeal. Solid work, the drain and eviction design is right, it just needs to converge on the one reservation authority. |
2 participants
Summary
Fixes #1690 (lock fix) and cross-cutting arbiter wiring for the taOS arbiter epic (#1683/#1705/#1706/#1707/#1690).
Changes
#1690 lock fix —
drain_workerforce-release and_monitor_loopstale-drain path now hold_lease_lockduringself._leasesmutation, serialized with claim/release/sweep.drain_workerandcancel_drainare nowasync def(consistent withclaim_lease/release_lease— all callers were already in async handlers).Cross-cutting wiring — Added
GpuArbiter.cancel_running_for_leases()and wired it into:drain_worker(graceful=False)— after force-releasing leases_monitor_loopstale-drain branch — when a draining worker stops heartbeatingWhen leases are force-released, any running GPU arbiter tasks for those leases are cancelled via
_evict_task. Eviction is no longer a prod no-op when the arbiter is running.Also: Added
TYPE_CHECKINGimport forGpuArbitertype annotation onClusterManager._gpu_arbiter.Files changed
tinyagentos/cluster/manager.py— async drain_worker/cancel_drain, lock fix, arbiter cancellation wiringtinyagentos/routes/cluster.py— await on async drain_worker/cancel_drain callstinyagentos/scheduler/gpu_arbiter.py— newcancel_running_for_leases()methodtests/test_cluster.py— await on all 10 drain_worker/cancel_drain test callsTests
47/47 pass:
tests/test_cluster.py(26),tests/test_gpu_arbiter_894.py(14),tests/test_gpu_arbiter_toctou.py(7)