fix(scheduler): close TOCTOU race in GPU admission with in-flight VRAM reservation#1705
fix(scheduler): close TOCTOU race in GPU admission with in-flight VRAM reservation#1705hognek wants to merge 6 commits into
Conversation
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? |
📝 WalkthroughWalkthroughThis PR introduces a ChangesGPU Arbiter feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant GpuArbiter
participant ClusterManager
participant Scheduler
Caller->>GpuArbiter: submit_gpu(task, required_vram_mb)
GpuArbiter->>GpuArbiter: _reserve_and_check(vram)
alt admitted
GpuArbiter->>ClusterManager: claim_lease(resource_id)
GpuArbiter->>Scheduler: submit(task)
Scheduler-->>GpuArbiter: result
GpuArbiter->>GpuArbiter: release reservation
GpuArbiter->>ClusterManager: release_lease(lease_id)
GpuArbiter-->>Caller: result
else rejected
GpuArbiter->>GpuArbiter: enqueue _QueuedGpuTask with future
GpuArbiter-->>Caller: await arbiter_future
GpuArbiter->>GpuArbiter: _drain_queue admits later
end
sequenceDiagram
participant Operator
participant GpuArbiter
participant RunningTask
participant ClusterManager
Operator->>GpuArbiter: evict_lowest_priority(min_priority)
GpuArbiter->>GpuArbiter: select lowest-priority running task
GpuArbiter->>GpuArbiter: release VRAM reservation
GpuArbiter->>ClusterManager: release_lease(lease_id)
GpuArbiter->>RunningTask: cancel()
RunningTask-->>GpuArbiter: CancelledError propagated
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 |
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| if required_vram_mb <= 0: | ||
| return GpuAdmission(admitted=True) | ||
|
|
||
| async with self._reservation_lock: |
There was a problem hiding this comment.
WARNING: _reservation_lock is held during self._vram_probe(), which can block all concurrent admissions for up to the probe's 5-second timeout (_probe_nvidia_vram). Under load this serialises every admission decision behind the slowest probe call.
Consider probing outside the lock (or caching the probe value with a short TTL) and only taking the lock to read _reserved_vram_mb and update _pending_reservations. Alternatively, snapshot _vram_probe() before the async with and pass the result into _check_admission.
| l.required_vram_mb for l in leases | ||
| if l.resource_id.startswith(worker.name + ":") and l.required_vram_mb > 0 | ||
| ) | ||
| if worker.free_vram_mb - worker_leases >= required_vram_mb: |
There was a problem hiding this comment.
WARNING: The cluster fallback path does not subtract _reserved_vram_mb from worker.free_vram_mb, so the TOCTOU fix is incomplete for cluster-mode admission. Two concurrent submit_gpu calls can both be admitted to different (or the same) cluster workers when their combined required VRAM exceeds available capacity — the same race the fix is meant to close.
Apply the same reservation accounting here: subtract _reserved_vram_mb (or a per-worker share) from the worker's effective free VRAM before admitting.
|
|
||
| # Simulate what _reserve_and_check would have done when the task | ||
| # was admitted. | ||
| arbiter._reserved_vram_mb += 4096 |
There was a problem hiding this comment.
SUGGESTION: This test manually mutates _reserved_vram_mb and _pending_reservations instead of going through the public _reserve_and_check API. That bypasses the reservation flow it's supposed to test and won't catch integration bugs (e.g. if _reserve_and_check forgets to populate _pending_reservations).
Use await arbiter._reserve_and_check("t-evict", 4096) to set up the reservation, so the test exercises the real code path that submit_gpu and _drain_queue use.
| assert "t-noevict" in arbiter._running | ||
|
|
||
| # Clean up | ||
| arbiter._running_tasks.get("t-noevict", None) |
There was a problem hiding this comment.
SUGGESTION: arbiter._running_tasks.get("t-noevict", None) is dead code — the result is discarded. The comment says "Clean up" but nothing is actually cleaned up here. Replace with arbiter._evict_task("t-noevict") (or cancel the stored asyncio task) so the test doesn't leak a running coroutine if evict_lowest_priority returns 0.
Code Review SummaryStatus: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (4 files)
Fix these issues in Kilo Cloud Previous Review Summaries (3 snapshots, latest commit 815cca8)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 815cca8)Status: 7 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (8 files)
Fix these issues in Kilo Cloud Previous review (commit d647517)Status: 5 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (8 files)
Fix these issues in Kilo Cloud Previous review (commit 04f5a69)Status: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (8 files)
Reviewed by minimax-m3 · Input: 43.1K · Output: 5.2K · Cached: 425.1K |
|
|
hognek
force-pushed
the
fix/toctou-vram-reservation
branch
from
04f5a69 to
d647517
Compare
| resource_id=resource_id, caller=task.submitter, | ||
| ttl_seconds=300, required_vram_mb=required_vram_mb, | ||
| ) | ||
| if lease is None: |
There was a problem hiding this comment.
CRITICAL: Orphaned VRAM reservation when claim_lease fails.
submit_gpu now reserves VRAM via _reserve_and_check (line 178) BEFORE calling _run_gpu_task. If claim_lease returns None here at line 277, the raise NoResourceAvailableError propagates out before the try block on line 287 — so the finally block (which calls _release_reservation(task.id)) never runs. The reservation leaks: _reserved_vram_mb and _pending_reservations[task.id] stay populated, permanently reducing the effective free VRAM seen by _check_admission for the lifetime of the arbiter. This is a regression introduced by the TOCTOU fix in patch 4.
Restructure so the reservation release is guaranteed on lease-claim failure — e.g., wrap the entire function body in try/finally, or release the reservation at the raise site.
| if lease is None: | |
| if lease_id is None and self._cluster_manager is not None and resource_id is not None: | |
| raise NoResourceAvailableError( | |
| f"GPU lease claim failed for {resource_id} (task {task.id})" | |
| ) |
Note: the suggestion above is structural; the minimal safe fix is to move the raise inside the try: block (or add an early self._release_reservation(task.id) before the raise).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…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)
hognek
force-pushed
the
fix/toctou-vram-reservation
branch
from
d647517 to
815cca8
Compare
| # Build the GPU arbiter — wraps the resource scheduler with VRAM-accounted | ||
| # admission control, queuing, and eviction for GPU-bound workloads. | ||
| try: | ||
| gpu_arbiter = GpuArbiter( |
There was a problem hiding this comment.
WARNING: GpuArbiter is constructed without a vram_probe argument, so it falls back to _default_vram_probe which returns (0, 0). With free_vram=0, _check_admission always falls through to the cluster branch (or returns admitted=True when no cluster_manager), meaning the arbiter never checks local VRAM and the entire TOCTOU reservation system (_reserved_vram_mb, _reserve_and_check, etc.) is never exercised in practice. Pass vram_probe=_probe_nvidia_vram (or a cached variant) so the arbiter actually prevents the Xid 62 crashes it was built to stop.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| max_queue_size=100, | ||
| eviction_enabled=True, | ||
| ) | ||
| await gpu_arbiter.start() |
There was a problem hiding this comment.
WARNING: gpu_arbiter.start() spawns a long-running queue-processor task (asyncio.create_task(self._process_queue())), but there is no corresponding gpu_arbiter.stop() in the FastAPI lifespan shutdown sequence (see lines ~1295-1344). On app shutdown the queue processor task is never cancelled and awaited, so the asyncio task leaks and the process may hang at interpreter shutdown while waiting for the background coroutine. Add await gpu_arbiter.stop() to the lifespan cleanup, guarded the same way as the other .stop() calls (e.g. inside a try/except or behind a getattr(app.state, "gpu_arbiter", None) check).
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…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.
| # Subtract in-flight reservations as well — the cluster-mode | ||
| # path must also account for tasks whose models are still | ||
| # loading (taOS #1705). | ||
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb |
There was a problem hiding this comment.
WARNING: Cluster path now risks TypeError on workers with no VRAM probe
WorkerInfo.free_vram_mb is typed int | None (see tinyagentos/cluster/worker_protocol.py:69) and is None on non-NVIDIA hardware (RK3588, Apple Silicon, CPU-only) and on workers that haven't reported a heartbeat yet. The expression worker.free_vram_mb - worker_leases - self._reserved_vram_mb will raise TypeError: unsupported operand type(s) for -: 'NoneType' and 'int' the moment such a worker is online, defeating the new reservation fix and crashing the admission check. Treat None as "unknown / un-leasable" and continue (mirroring the None-vs-0 contract documented on the dataclass) so this path stays consistent with the lease pre-claim check.
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb | |
| if worker.free_vram_mb is None: | |
| continue | |
| available = worker.free_vram_mb - worker_leases - self._reserved_vram_mb |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tinyagentos/app.py (1)
1149-1149: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant conditional.
resource_scheduler if resource_scheduler is not None else Nonealways evaluates toresource_scheduler.♻️ Simplify
- gpu_arbiter = GpuArbiter( - scheduler=resource_scheduler if resource_scheduler is not None else None, + gpu_arbiter = GpuArbiter( + scheduler=resource_scheduler,🤖 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 `@tinyagentos/app.py` at line 1149, The conditional expression in the scheduler assignment is redundant because it always returns the same value; simplify the argument in the relevant call site by passing resource_scheduler directly instead of using the “if not None else None” form. Update the code around the scheduler parameter in the app setup path so the expression is cleaner while preserving the existing behavior.tinyagentos/scheduler/discovery.py (1)
190-197: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winComment lists GPU backends not in
gpu_backend_types.The block comment advertises
llama-cpp(CUDA build) andsd-cpp/sd-gpu(image-generation) as GPU backends, butgpu_backend_typesonly contains{"vllm", "ollama", "exo", "mlx"}. As written, a CUDA-builtllama-cppand any image-generation backend are never treated as GPU-capable here (andllama-cpp/sd-cppare only picked up by the CPU tier below). Either drop the misleading lines from the comment or add the intended types to the set.🤖 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 `@tinyagentos/scheduler/discovery.py` around lines 190 - 197, The GPU backend classification in the discovery logic is inconsistent with the surrounding comment: `gpu_backend_types` in `discovery.py` only includes `vllm`, `ollama`, `exo`, and `mlx`, while the comment also claims `llama-cpp` and `sd-cpp`/`sd-gpu` are GPU-capable. Update the `gpu_backend_types` set in `gpu_backends` (or trim the comment) so the documented GPU backends and the actual filter match, ensuring the intended `llama-cpp` and image-generation backends are treated consistently.
🤖 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 `@tinyagentos/app.py`:
- Line 1149: The conditional expression in the scheduler assignment is redundant
because it always returns the same value; simplify the argument in the relevant
call site by passing resource_scheduler directly instead of using the “if not
None else None” form. Update the code around the scheduler parameter in the app
setup path so the expression is cleaner while preserving the existing behavior.
In `@tinyagentos/scheduler/discovery.py`:
- Around line 190-197: The GPU backend classification in the discovery logic is
inconsistent with the surrounding comment: `gpu_backend_types` in `discovery.py`
only includes `vllm`, `ollama`, `exo`, and `mlx`, while the comment also claims
`llama-cpp` and `sd-cpp`/`sd-gpu` are GPU-capable. Update the
`gpu_backend_types` set in `gpu_backends` (or trim the comment) so the
documented GPU backends and the actual filter match, ensuring the intended
`llama-cpp` and image-generation backends are treated consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ae14a23b-2a8b-4166-8a69-d943005ad17f
📒 Files selected for processing (8)
.gitignoretests/test_gpu_arbiter_894.pytests/test_gpu_arbiter_toctou.pytinyagentos/app.pytinyagentos/scheduler/__init__.pytinyagentos/scheduler/discovery.pytinyagentos/scheduler/gpu_arbiter.pytinyagentos/scheduler/types.py
2 participants
Summary
Closes the TOCTOU (time-of-check to time-of-use) race in
_check_admissionidentified in PR #1683 review feedback #2.Problem
_check_admissionreads livenvidia-smifree VRAM with no reservation mechanism. Two concurrentsubmit_gpucalls can both pass admission before either model loads — the exact concurrent-load crash (NVIDIA Xid 62) the arbiter exists to prevent.Fix
Adds in-flight VRAM reservation tracking:
_reserved_vram_mb: counter of VRAM promised to tasks whose models are still loading_pending_reservations: dict mapping task_id → reserved VRAM for idempotent release_reservation_lock:asyncio.Lockto make check+reserve atomicNew method
_reserve_and_check(task_id, required_vram_mb)atomically:_reservation_lock_check_admission(which now subtracts_reserved_vram_mbfrom the probe)_reserved_vram_mbReservations are released in:
_run_gpu_taskfinally block (normal completion, error, cancellation)_evict_task(preemption)Tests
test_gpu_arbiter_894.pytests continue to passtest_gpu_arbiter_toctou.pytests:_reserve_and_checkonly admits one of two calls exceeding capacitysubmit_gpureserves and releases correctly_check_admissionsubtracts reserved VRAM from probestats()All 168 tests pass (25 arbiter + 143 cluster/worker).
Summary by CodeRabbit