feat(models): atomic VRAM check-and-reserve before model load (TOCTOU fix, #1706)#1725
Conversation
… fix) Adds VramReservationManager — an asyncio.Lock-protected VRAM reservation system that atomically checks free VRAM (via nvidia-smi) and reserves the required amount before a model load begins. Concurrent callers see the in-flight reservation and back off, preventing the TOCTOU race where two model pulls both see enough free VRAM and both try to load, causing OOM. Integration points: - POST /api/models/pull — reserve before ollama pull, release after - POST /api/models/download (rkllama path) — reserve before installer, release in the background task's finally block - GET /api/models/vram-reservations — stats endpoint for monitoring New files: - tinyagentos/vram_reservation.py — atomic reserve/release manager - tests/test_vram_reservation.py — 15 tests including concurrent safety Wired as app.state.vram_reservation in the FastAPI lifespan alongside the cluster manager and resource scheduler. Issue: jaylfc#1706
📝 WalkthroughWalkthroughAdds a new ChangesAtomic VRAM Reservation
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ModelsRoute
participant VramReservationManager
participant NvidiaSmiProbe
Client->>ModelsRoute: POST /api/models/pull (required_vram_mb)
ModelsRoute->>VramReservationManager: reserve(required_vram_mb)
VramReservationManager->>NvidiaSmiProbe: probe free/total VRAM
NvidiaSmiProbe-->>VramReservationManager: (free_mb, total_mb)
alt insufficient VRAM
VramReservationManager-->>ModelsRoute: None
ModelsRoute-->>Client: 503 Insufficient VRAM
else reserved
VramReservationManager-->>ModelsRoute: VramReservation
ModelsRoute->>ModelsRoute: run pull task
ModelsRoute->>VramReservationManager: release(reservation_id)
ModelsRoute-->>Client: 200 success
end
🚥 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}
| # check and OOM. | ||
| vram_mgr = getattr(request.app.state, "vram_reservation", None) | ||
| reservation = None | ||
| estimated_vram = int(variant.get("min_ram_mb", 0) or 0) # heuristic |
There was a problem hiding this comment.
CRITICAL: Heuristic uses min_ram_mb for VRAM reservation on rkllama pulls.
min_ram_mb declares minimum host RAM (MiB), not VRAM, and on many systems host RAM is much larger than GPU VRAM. Reserving min_ram_mb MiB of VRAM will frequently over-estimate the actual VRAM requirement and cause false 503s, blocking legitimate downloads on systems where the GPU has plenty of free VRAM.
Use a VRAM-specific field (e.g. variant.get("vram_mb") or variant.get("min_vram_mb")) and fall back to a sensible fraction of min_ram_mb (e.g. min_ram_mb // 2) only when that field is absent. At minimum, document that this is a deliberate heuristic known to over-reserve.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| if vram_mgr is not None and reservation is not None: | ||
| vram_mgr.release(reservation.reservation_id) | ||
|
|
||
| dm.start_installer_task(download_id, _install_and_record) |
There was a problem hiding this comment.
WARNING: VRAM reservation leaks if start_installer_task raises.
The reservation is acquired at line 379 and released only inside _install_and_record's finally block (line 426). If dm.start_installer_task(...) itself raises — between the reserve on line 379 and the task actually running its finally — the reservation never gets released and pending_count permanently grows by one for the lifetime of the process. (Same risk applies if the FastAPI request handler returns before the task is scheduled.)
Wrap the call in a try so the reservation is released on the synchronous failure path:
try:
dm.start_installer_task(download_id, _install_and_record)
except Exception:
if vram_mgr is not None and reservation is not None:
vram_mgr.release(reservation.reservation_id)
raiseReply with @kilocode-bot fix it to have Kilo Code address this issue.
| ) | ||
|
|
||
| async with self._lock: | ||
| free_mb, total_mb = self._probe_vram() |
There was a problem hiding this comment.
WARNING: _probe_vram() runs outside the lock in stats() and available_vram().
_probe_vram calls read_nvidia_vram(), which executes nvidia-smi as a subprocess. While that's running, a concurrent reserve() may add or another release() may subtract from _reserved_vram_mb. The dict returned by stats() (and the tuple returned by available_vram()) are then not a self-consistent snapshot: free_vram_mb may be stale relative to reserved_vram_mb, and effective_free_mb may disagree with both. Monitoring clients will see values that don't add up.
Either read _reserved_vram_mb and probe under the same lock and return an already-computed snapshot, or document the snapshot as eventually consistent. Given the lock only ever guards a fast in-memory check, taking it briefly in stats()/available_vram() is cheap and removes the inconsistency.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
|
|
||
| # ── internal ──────────────────────────────────────────────────── | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
SUGGESTION: _probe_vram is declared as @staticmethod, but the test helper at tests/test_vram_reservation.py:30 rebinds it via mgr._probe_vram = staticmethod(_fake_probe). Because it's a staticmethod, the descriptor wraps it on the class, not the instance, so the test rebinding works but it's brittle. More importantly, a staticmethod can't be easily replaced by tests that want to inject behaviour via subclassing.
Consider making it a regular instance method (def _probe_vram(self) -> ...) — the lock can still serialize calls into it, and tests can subclass or monkeypatch more cleanly.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (4 files)
Fix these issues in Kilo Cloud Previous Review Summary (commit 619e225)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit 619e225)Status: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
WARNING
SUGGESTION
Files Reviewed (4 files)
Reviewed by hy3-20260706:free · Input: 40.1K · Output: 4.8K · Cached: 99.4K |
|
Confirming the direction (full detail on #894): this is the VRAM guard we ship first, as the single authority on the model-load path. Kilo is clean and the atomic check-and-reserve is exactly what prevents the concurrent-load over-subscription. Two things before it merges:
Once it is rebased and green I will merge it. The |
jaylfc review: min_ram_mb heuristic is safe for v1 (over-reserves, fails closed), but should be replaced with a real per-model VRAM estimate on CUDA GGUF to avoid false 503s in multi-model setups.
Head branch was pushed to by a user without write access
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/test_vram_reservation.py (1)
226-238: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding an assertion or skip marker to the real-probe smoke test.
test_reserve_with_real_probehas no assertions and always passes — on GPU-less systems it passes vacuously, and on GPU systems it only verifies the code doesn't crash. Consider either assertingres is not Noneunder a GPU-availability skip marker, or usingpytest.xfailfor non-GPU environments so the test actually communicates intent.♻️ Optional improvement
class TestRealProbe: """Smoke-test with the real nvidia-smi probe (if available).""" + `@pytest.mark.asyncio` async def test_reserve_with_real_probe(self): """Reserving 1 MiB should always succeed on a real GPU.""" mgr = VramReservationManager() res = await mgr.reserve(1, caller="smoke-test") - # On a system with nvidia-smi and some free VRAM, this should succeed. - # On a system without nvidia-smi, the probe returns (0, 0) so - # this will fail — that's fine, the test is a best-effort smoke. - if res is not None: - mgr.release(res.reservation_id) + if res is None: + pytest.skip("No GPU or nvidia-smi available — probe returned 0 VRAM") + try: + assert res.reservation_id.startswith("vr_") + finally: + mgr.release(res.reservation_id)🤖 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_vram_reservation.py` around lines 226 - 238, The real-probe smoke test in TestRealProbe currently has no meaningful assertion, so it passes even when nothing is actually verified. Update test_reserve_with_real_probe to either assert that VramReservationManager.reserve returns a reservation when a real GPU is available, guarded by an appropriate skip condition for GPU-less environments, or mark the non-GPU path as xfail so the intent is explicit. Use the existing reserve and release flow in TestRealProbe to keep the cleanup behavior intact.
🤖 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/vram_reservation.py`:
- Line 51: The docstring in VRAMReservationManager overstates the concurrency
guarantee by calling it “thread-safe” when it only uses asyncio.Lock for
coroutine-level serialization. Update the class docstring near
VRAMReservationManager to describe it as coroutine-safe (or equivalent), and
make sure the wording matches the actual behavior of acquire()/release(),
especially since release() mutates shared state without thread protection.
- Around line 169-186: The VRAM probe is currently synchronous in
`_probe_vram()` because it calls `read_nvidia_vram()` directly, which can block
the event loop via `subprocess.run` and stall `reserve()`, `available_vram()`,
and `stats()`. Move the probe work off the loop by wrapping the
`read_nvidia_vram` call with `asyncio.to_thread(...)` or replacing it with an
async subprocess/cache approach, and then update the read-only accessors and any
call sites in `VRAMReservation` to await the async probe instead of calling it
inline.
---
Nitpick comments:
In `@tests/test_vram_reservation.py`:
- Around line 226-238: The real-probe smoke test in TestRealProbe currently has
no meaningful assertion, so it passes even when nothing is actually verified.
Update test_reserve_with_real_probe to either assert that
VramReservationManager.reserve returns a reservation when a real GPU is
available, guarded by an appropriate skip condition for GPU-less environments,
or mark the non-GPU path as xfail so the intent is explicit. Use the existing
reserve and release flow in TestRealProbe to keep the cleanup behavior intact.
🪄 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: 0fd38523-7b24-4a10-a034-dcdb9af1226b
📒 Files selected for processing (4)
tests/test_vram_reservation.pytinyagentos/app.pytinyagentos/routes/models.pytinyagentos/vram_reservation.py
| class VramReservationManager: | ||
| """Atomic VRAM check-and-reserve for model loads. | ||
|
|
||
| Thread-safe via ``asyncio.Lock``. Reads real-time VRAM from |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring overstates concurrency guarantee.
asyncio.Lock serializes coroutines on a single event loop; it is not thread-safe. Since release() also mutates _reserved_vram_mb/_pending without the lock, calling this manager from multiple threads (e.g. via run_in_executor) would race. Consider rewording to "coroutine-safe" to avoid a false sense of thread safety.
🤖 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/vram_reservation.py` at line 51, The docstring in
VRAMReservationManager overstates the concurrency guarantee by calling it
“thread-safe” when it only uses asyncio.Lock for coroutine-level serialization.
Update the class docstring near VRAMReservationManager to describe it as
coroutine-safe (or equivalent), and make sure the wording matches the actual
behavior of acquire()/release(), especially since release() mutates shared state
without thread protection.
| @staticmethod | ||
| def _probe_vram() -> tuple[int, int]: | ||
| """Probe real-time free/total VRAM via nvidia-smi. | ||
|
|
||
| Returns ``(free_mb, total_mb)``, or ``(0, 0)`` when no | ||
| nvidia-smi is available or the probe fails. | ||
| """ | ||
| try: | ||
| from tinyagentos.system_stats import read_nvidia_vram | ||
|
|
||
| pair = read_nvidia_vram() | ||
| if pair is not None: | ||
| used_mb, total_mb = pair | ||
| free_mb = max(0, total_mb - used_mb) | ||
| return free_mb, total_mb | ||
| except Exception: | ||
| logger.debug("vram-reservation: probe failed", exc_info=True) | ||
| return 0, 0 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f 'system_stats.py' -x sed -n '1,80p' {}
rg -nP -C3 'def read_nvidia_vram'
rg -nP -C3 'subprocess|check_output|Popen|run\(' --glob '**/system_stats.py'Repository: jaylfc/taOS
Length of output: 5276
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant file around the probe and its call sites.
sed -n '1,240p' tinyagentos/vram_reservation.py
# Find where the reservation helpers are used.
rg -n "available_vram\(|stats\(|reserve\(|_probe_vram\(" -S tinyagentosRepository: jaylfc/taOS
Length of output: 11278
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Narrowly inspect the synchronous NVIDIA probe implementation.
cat -n tinyagentos/system_stats.py | sed -n '1,120p'Repository: jaylfc/taOS
Length of output: 4417
Move the VRAM probe off the event loop. read_nvidia_vram() shells out to nvidia-smi with a synchronous subprocess.run(..., timeout=3), so _probe_vram() blocks for up to 3s per call. Because reserve() runs this under asyncio.Lock, and available_vram()/stats() call it from the request path, a single probe can stall unrelated async work. Use asyncio.to_thread(...) or an async subprocess/cache, and make the read-only accessors async as well.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 184-184: Do not catch blind exception: Exception
(BLE001)
🤖 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/vram_reservation.py` around lines 169 - 186, The VRAM probe is
currently synchronous in `_probe_vram()` because it calls `read_nvidia_vram()`
directly, which can block the event loop via `subprocess.run` and stall
`reserve()`, `available_vram()`, and `stats()`. Move the probe work off the loop
by wrapping the `read_nvidia_vram` call with `asyncio.to_thread(...)` or
replacing it with an async subprocess/cache approach, and then update the
read-only accessors and any call sites in `VRAMReservation` to await the async
probe instead of calling it inline.
Source: Linters/SAST tools
2 participants
Summary
Implements atomic VRAM reservation to prevent TOCTOU races during concurrent model loads.
Problem: Two concurrent model pulls (ollama/rkllama) both check free VRAM via nvidia-smi, both see enough, both start loading — causing OOM or driver crashes.
Solution:
VramReservationManager— anasyncio.Lock-protected reservation system that atomically checks free VRAM and reserves the required amount. Concurrent callers see the in-flight reservation and back off with HTTP 503.Changes
New files
tinyagentos/vram_reservation.py— atomic reserve/release manager with real-time nvidia-smi probetests/test_vram_reservation.py— 15 tests covering reserve/release, concurrent safety, zero-VRAM, stats, and real-probe smoke testModified files
tinyagentos/app.py— wiresVramReservationManagerasapp.state.vram_reservationin the FastAPI lifespantinyagentos/routes/models.py— integrates reservation into:POST /api/models/pull— reserve before ollama pull, release in finallyPOST /api/models/download(rkllama path) — reserve before installer, release in background task finallyGET /api/models/vram-reservations— new monitoring endpointDesign
The manager follows the same pattern as the GPU arbiter's TOCTOU fix (#894 review feedback #2): an
asyncio.Lockserializes the check-and-reserve sequence so two concurrentreserve()calls never both see the same free VRAM. In-flight reservations are subtracted from the probe result so the effective free VRAM shrinks for subsequent callers.Zero-VRAM reservations (
vram_mb <= 0) always succeed as lightweight markers (no lock, no probe).Testing
No regressions in existing model, config, cluster, installer, or catalog tests.
Issue: #1706
Summary by CodeRabbit
New Features
Bug Fixes
Tests