fix(scheduler): add _running_lock coverage to eviction paths in GpuArbiter#1707
fix(scheduler): add _running_lock coverage to eviction paths in GpuArbiter#1707hognek wants to merge 11 commits into
Conversation
|
Warning Review limit reached
Next review available in: 52 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 (9)
✨ 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}
| Running tasks whose lease belongs to this worker are evicted | ||
| (cancelled) so they can be resubmitted to another worker. | ||
|
|
||
| Returns the number of tasks evicted. |
There was a problem hiding this comment.
CRITICAL: _running_lock held while iterating cluster_manager._leases introduces a race with concurrent lease mutation.
The new code reads self._cluster_manager._leases (a separate dict, owned by the cluster manager, not guarded by _running_lock) inside the async with self._running_lock: block:
async with self._running_lock:
victim_ids = []
for tid in list(self._running.keys()):
_task, lid, _pri, _vram = self._running.get(tid, (None, None, 0, 0))
if lid is None:
continue
if self._cluster_manager is not None:
lease = getattr(self._cluster_manager, "_leases", {}).get(lid)
..._leases is mutated by other coroutines: _sweep_expired_leases (runs every 5s in the monitor loop), release_lease() from any caller, _monitor_loop's offline-worker branch. Iterating it indirectly is safe here (it's a .get()), but pairing that with the lock-widening PR is misleading: the lock doesn't actually protect the lease lookup, and the lock scope now spans dict accesses the lock does not guard against.
More importantly, this widens the lock to cover I/O-shaped work that doesn't need it. The intent stated in the PR description ("Lease release and task cancellation happen outside the lock to avoid deadlock") is correct for lease release/cancel, but here we still iterate the external _leases dict while holding _running_lock. Build the victim_ids snapshot under the lock based only on data the lock protects (_running and whatever cached lease_id is stored with the running task), and resolve lease ownership outside the lock.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| Returns the number of tasks evicted. | ||
| """ | ||
| async with self._running_lock: |
There was a problem hiding this comment.
SUGGESTION: Redundant _running.get() after list(self._running.keys()) under the lock.
Since the loop holds _running_lock, _running cannot change between list(self._running.keys()) and the subsequent self._running.get(tid, (None, None, 0, 0)) lookups. The get() always succeeds with the value already known from the key iteration. Capture the value once:
| async with self._running_lock: | |
| for tid, (_task, lid, _pri, _vram) in self._running.items(): |
This also avoids relying on _running.keys() and _running.get() staying consistent across statements — a reasonable pattern when not holding a lock, but pure overhead here.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| async def evict_lowest_priority(self, min_priority: int | None = None) -> int: | ||
| if not self._eviction_enabled: | ||
| return 0 |
There was a problem hiding this comment.
WARNING: asyncio_task is popped inside the lock but task.cancel() happens after exiting the lock. The popped asyncio_task is held in a local variable, so there is no race on the dict entry itself — this part is safe. However, the task body (task.payload) may resume and re-enter submit_gpu / _run_gpu_task before our cancellation propagates, and the PR introduces this ordering by design. Add a brief comment near the cancel() call to document the intentional ordering so a future reviewer doesn't try to "fix" it by moving the cancel back under the lock (which would deadlock). The PR description already covers the intent; mirroring it as a code comment at line 332-333 helps future maintenance.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 0 NEW Issues | Recommendation: All 14 prior findings resolved Overview
All 14 inline comments from the previous review at
No new issues introduced by the current diff. Files Reviewed (9 files in diff)
Previous Review Summaries (5 snapshots, latest commit 4fc57cb)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 4fc57cb)Status: 14 Issues Found (no change since previous review) | Recommendation: Address before merge Overview
Incremental diff vs prior review SHA ( Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (incremental, 7 files in PR since `4f94c83d`)
Fix these issues in Kilo Cloud Previous review (commit 4f94c83)Status: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (9 files)
Fix these issues in Kilo Cloud Previous review (commit 6d77b76)Status: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
Files Reviewed (incremental, 8 files in PR)
Fix these issues in Kilo Cloud Previous review (commit 854942c)Status: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (13 files in PR)
Fix these issues in Kilo Cloud Previous review (commit 69a34e1)Status: 3 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (19 files in PR; focus was commits 7–8, the lock-coverage fix)
Reviewed by minimax-m3 · Input: 58.7K · Output: 2.2K · Cached: 196.9K |
|
|
hognek
force-pushed
the
fix/894-lock-running-eviction
branch
2 times, most recently
from
0ccf71b to
854942c
Compare
| async with self._running_lock: | ||
| self._running[task.id] = (task, lease_id, int(task.priority), required_vram_mb) | ||
| if current is not None: | ||
| self._running_tasks[task.id] = current |
There was a problem hiding this comment.
CRITICAL: asyncio.current_task() inside _run_gpu_task records whatever task is currently executing — when this function is called from _drain_queue (line 376), that's the queue processor task, not the submitter. Cancelling it (via _evict_task → asyncio_task.cancel() at line 346) terminates the entire background drain loop for the lifetime of the process, halting all queued GPU work after a single eviction.
Distinguish the two callers — e.g., only register in _running_tasks when invoked from submit_gpu, or store the submitter task explicitly via a parameter — so eviction always targets the user-visible task and never the internal processor.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| ] | ||
| for lid in lids: | ||
| self._leases.pop(lid, None) | ||
| released += 1 |
There was a problem hiding this comment.
WARNING: drain_worker (force branch) mutates self._leases (line 437-439 loop) and reads self._leases.items() (line 433-435) without holding self._lease_lock. Other coroutines — claim_lease, release_lease, _sweep_expired_leases, the monitor-loop drain-completion branch — all serialize on that lock. Concurrent mutation can raise RuntimeError: dictionary changed size during iteration or silently miss/stale leases. Wrap the iteration+pop in async with self._lease_lock:.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if isinstance(body, dict): | ||
| graceful = body.get("graceful", True) | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
WARNING: Bare except Exception: pass swallows every JSON-parse failure as if no body was sent. A caller POSTing Content-Type: application/json with body {"graceful": false} that fails to parse (truncated body, encoding error) silently falls back to the default graceful=True, producing a graceful drain instead of the requested force-drain — a correctness footgun on a destructive admin endpoint. Distinguish "no body / empty" from "parse error" and surface the latter as a 400.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
hognek
force-pushed
the
fix/894-lock-running-eviction
branch
from
854942c to
6d77b76
Compare
| entry = self._queue.get_nowait() | ||
| admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) | ||
| if not admission.admitted: | ||
| # Try eviction-to-make-room for higher-priority queued tasks. |
There was a problem hiding this comment.
CRITICAL: Missing await on evict_lowest_priority — will raise TypeError at runtime.
Patch 5 added this call to _drain_queue:
evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority))
Patch 6 then changed evict_lowest_priority to async def, but did not add await here. As a result evicted is a coroutine object, not an int, so if evicted > 0 raises TypeError: '>' not supported between instances of 'coroutine' and 'int' the first time a higher-priority queued task needs eviction-to-make-room. The whole eviction-to-make-room code path (the headline feature of patch 5) is unreachable.
| # Try eviction-to-make-room for higher-priority queued tasks. | |
| evicted = await self.evict_lowest_priority(min_priority=int(entry.task.priority)) |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if not admission.admitted: | ||
| # Try eviction-to-make-room for higher-priority queued tasks. | ||
| # evict_lowest_priority(min_priority=N) skips running tasks | ||
| # whose priority value is *lower* than N (i.e. tasks that are |
There was a problem hiding this comment.
WARNING: Non-atomic re-check after eviction defeats the TOCTOU fix.
After evict_lowest_priority frees some VRAM, this line re-checks admission with a plain _check_admission call — which does not go through _reservation_lock and does not reserve the freed VRAM. A concurrent submit_gpu can grab that same VRAM between the eviction and this re-check, and now we admit a task whose VRAM is no longer actually free. This re-opens the TOCTOU window that patch 4 explicitly closed.
It should use the same atomic helper:
| # whose priority value is *lower* than N (i.e. tasks that are | |
| admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
hognek
force-pushed
the
fix/894-lock-running-eviction
branch
from
c9e2705 to
4f94c83
Compare
…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.
| # whose priority value is *lower* than N (i.e. tasks that are | ||
| # actually higher priority), so only lower-or-equal priority | ||
| # running tasks are candidates for eviction. | ||
| evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority)) |
There was a problem hiding this comment.
CRITICAL: Missing await on evict_lowest_priority — will raise TypeError at runtime.
evict_lowest_priority is declared async (line 314) and returns a coroutine. This call drops the coroutine on the floor and assigns it to evicted, which is always a non-zero truthy coroutine object — so the subsequent if evicted > 0 branch is unreachable in the intended way, and the eviction that was supposed to free VRAM never actually runs (its task is never cancelled, its lease is never released, its reservation is never freed). The _drain_queue path that exists specifically to handle eviction-to-make-room is dead in practice.
| evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority)) | |
| evicted = await self.evict_lowest_priority(min_priority=int(entry.task.priority)) |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # running tasks are candidates for eviction. | ||
| evicted = self.evict_lowest_priority(min_priority=int(entry.task.priority)) | ||
| if evicted > 0: | ||
| admission = self._check_admission(entry.task, entry.required_vram_mb) |
There was a problem hiding this comment.
WARNING: Non-atomic re-check after eviction re-opens the TOCTOU window.
After evict_lowest_priority frees some VRAM (line 383), this line re-checks admission with a plain _check_admission call — which does not reserve the freed VRAM. A concurrent submit_gpu (or another drain tick) can race in between, grab the VRAM that was just freed, and put us right back where we started. The TOCTOU reservation introduced by _reserve_and_check should be used here too.
| admission = self._check_admission(entry.task, entry.required_vram_mb) | |
| admission = await self._reserve_and_check(entry.task.id, entry.required_vram_mb) |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if lease_id is not None and self._cluster_manager is not None: | ||
| self._cluster_manager.release_lease(lease_id) | ||
| # Cancel the running asyncio Task — stops _run_gpu_task and frees VRAM | ||
| if asyncio_task is not None and not asyncio_task.done(): |
There was a problem hiding this comment.
WARNING: asyncio_task.cancel() is invoked after exiting _running_lock.
Lines 329–333 pop both _running[task_id] and _running_tasks[task_id] under the lock. After the lock is released, asyncio_task.cancel() is called (line 340). Between release and cancel, a concurrent coroutine (e.g. _run_gpu_task's finally block at 303–305) can observe the missing entry, treat the task as already-evicted, and the cancel races with that path. Acquire the lock for the full eviction sequence or cancel under the lock.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| # cancel_and_wait budget alongside the supervised background tasks. | ||
| await cluster_manager.stop() | ||
| if app.state.gpu_arbiter is not None: | ||
| await app.state.gpu_arbiter.stop() |
There was a problem hiding this comment.
WARNING: await app.state.gpu_arbiter.stop() is not wrapped in try/except.
Every other teardown step in this block (_mdns.stop(), stop_taos_opencode_server, auto_updater.stop(), etc.) is guarded by try/except Exception + logger.exception(...). If gpu_arbiter.stop() raises (e.g. the _process_queue cancel/await path trips on a hung task, or the _running_lock is contended), shutdown of llm_proxy.stop() / stop_taos_opencode_server / monitor.stop() / auto_updater.stop() is skipped entirely and the FastAPI lifespan task crashes — leaving background tasks orphaned.
| await app.state.gpu_arbiter.stop() | |
| if app.state.gpu_arbiter is not None: | |
| try: | |
| await app.state.gpu_arbiter.stop() | |
| except Exception: | |
| logger.exception("gpu arbiter stop failed") |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| def _gpu_vram_probe() -> int: | ||
| free, _total = _probe_nvidia_vram() | ||
| return free if free > 0 else 999_999 # optimistic |
There was a problem hiding this comment.
SUGGESTION: Optimistic 999_999 MiB fallback masks broken nvidia-smi probes.
_probe_nvidia_vram() swallows every exception and returns (0, 0) when nvidia-smi is missing, broken, or returns garbage. The scheduler then reports ~1 TB of "free" VRAM and the resource looks healthier than it is — the scheduler will happily admit a workload the GPU arbiter (gpu_arbiter._check_admission) will then immediately reject. Either propagate the failure (return -1 to signal "unknown" and let the Resource skip the memory_probe) or log a warning when the probe returns 0.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| submitter: str = "unknown" | ||
| estimated_seconds: float = 1.0 | ||
| estimated_memory_mb: int = 0 | ||
| estimated_vram_mb: int = 0 |
There was a problem hiding this comment.
SUGGESTION: estimated_vram_mb field is added but never read by production code.
A repo-wide search shows only three test files (tests/test_gpu_arbiter*.py) construct Tasks with this field. The GPU arbiter's admission path uses required_vram_mb from the explicit submit_gpu(required_vram_mb=...) parameter, never task.estimated_vram_mb. Either wire it into submit_gpu (e.g. default required_vram_mb from task.estimated_vram_mb when caller doesn't override) or remove the field — leaving it dead is a maintenance trap.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…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)
hognek
force-pushed
the
fix/894-lock-running-eviction
branch
from
4f94c83 to
4fc57cb
Compare
…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.
hognek
force-pushed
the
fix/894-lock-running-eviction
branch
from
4fc57cb to
82e7611
Compare
2 participants
Summary
PR #1683 review feedback #5. Closes #894.
_runningwas mutated under_running_lockin_run_gpu_taskbut read/popped without the lock inevict_lowest_priority,_evict_task,release_tasks_for_worker,stats, andrunning_tasks. Safe under single-thread-no-await but fragile.Changes
gpu_arbiter.py: Makeevict_lowest_priority,_evict_task,release_tasks_for_worker,stats,running_tasksasync; guard_runningdict access withasync with self._running_lock:. Lease release and task cancellation happen outside the lock to avoid deadlock.cluster/manager.py:drain_worker→ async, awaitrelease_tasks_for_workerroutes/cluster.py: awaitstats()andrunning_tasks(); awaitdrain_workerawaitand@pytest.mark.asynciowhere neededTests
69/69 targeted tests pass:
test_gpu_arbiter_894.py,test_gpu_arbiter_796.py,test_gpu_arbiter.py,test_cluster_drain.py