fix(models): VRAM reservation TTL sweep (#1766)

[jaylfc]

Summary

Completes the remaining work for #1766 after the urgent parts landed in #1767.

• HIGH (already in fix(worker,models): audit hotfixes - manifest crash, VRAM fail-open, real min_ram_mb key #1767, regression-tested here): _probe_vram() returns None when nvidia-smi is absent; reserve() fails OPEN with bookkeeping only (matches claim_lease when free_vram_mb is None). Atomic deny still applies when VRAM is measurable.
• MEDIUM (already in fix(worker,models): audit hotfixes - manifest crash, VRAM fail-open, real min_ram_mb key #1767, helper extracted here): rkllama pull gate reads requires.backends[].min_ram_mb (max across backends, variant fallback) via _estimated_vram_mb().
• LOW (this PR): TTL sweep reclaims stale VramReservation entries (default 1h) so a hung installer cannot hold capacity until controller restart. Sweep runs from reserve(), available_vram(), stats(), and public sweep_stale(). Probe under the async lock already uses asyncio.to_thread (from fix(worker,models): audit hotfixes - manifest crash, VRAM fail-open, real min_ram_mb key #1767).

Acceptance

• No nvidia-smi: pull with real backend-level min_ram_mb proceeds with bookkeeping, no 503
• Two concurrent large reserves on NVIDIA still hit the atomic check (only one admitted)
• Reservation older than its TTL is reclaimed

Test plan

uv run pytest tests/test_vram_reservation.py tests/test_routes_models.py -q

Local result: 61 passed.

Closes #1766

Summary by CodeRabbit

• Bug Fixes
  • Improved model download handling for RKLLM backends by accurately considering backend memory requirements.
  • Downloads can proceed when VRAM cannot be measured, while still maintaining reservation tracking.
  • Insufficient measurable VRAM now prevents downloads before installation begins.
  • Stale VRAM reservations are automatically reclaimed, preventing capacity from remaining blocked after interrupted installations.
• Reliability
  • Improved handling of concurrent large-model downloads and reservation expiration.
  • Added clearer reservation status information, including the configured expiration window.

[qodo-code-review]

Qodo reviews are paused for this user.

Troubleshooting steps vary by plan Learn more →

On a Teams plan?
Reviews resume once this user has a paid seat and their Git account is linked in Qodo.
Link Git account →

Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center?
These require an Enterprise plan - Contact us
Contact us →

[coderabbitai]

📝 Walkthrough

Walkthrough

VRAM reservation now supports fail-open behavior without a hardware probe and TTL-based stale-entry reclamation. Model downloads estimate requirements from backend manifests, and tests cover admission, rejection, concurrency, estimation, and expiration.

Changes

VRAM admission lifecycle

Layer / File(s)	Summary
Reservation fail-open and TTL lifecycle tinyagentos/vram_reservation.py, tests/test_vram_reservation.py	Reservations record bookkeeping when VRAM probing is unavailable, stale reservations are reclaimed using configurable TTLs, and admission/accessor operations perform stale sweeps.
Backend-level model VRAM gating tinyagentos/routes/models.py, tests/test_routes_models.py	rkllama downloads reserve the maximum backend-level min_ram_mb, falling back to the variant-level requirement when needed.
Reservation and download acceptance coverage tests/test_vram_reservation.py, tests/test_routes_models.py	Tests cover no-probe admission, insufficient VRAM rejection, concurrent reservations, backend-level estimation, and TTL behavior.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• jaylfc/taOS#1658: Modifies the RKLLama download installer flow in the same route.
• jaylfc/taOS#1725: Introduces the related atomic VRAM check-and-reserve behavior.
• jaylfc/taOS#1767: Overlaps with backend-level VRAM gating and probe behavior.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.05% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: VRAM reservation TTL sweeping for models.
Linked Issues check	✅ Passed	The summaries show #1766 requirements were addressed: fail-open no-probe handling, backend min_ram_mb estimation, stale TTL sweeping, and atomic reservation tests.
Out of Scope Changes check	✅ Passed	The changes stay focused on the linked VRAM reservation and model-pull behavior, with no clear unrelated code introduced.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/vram-reservation-1766

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gitar-bot]

Important

You are using the Gitar free plan. Upgrade to unlock code review, CI analysis, auto-apply, custom automations, and more.

_Gitar

[kilo-code-bot]

WARNING: sweep_stale() mutates shared state without holding the lock

This public method calls _sweep_stale_unlocked(), which mutates self._pending and self._reserved_vram_mb, but does not acquire self._lock. reserve() mutates the same structures under the lock and yields during await asyncio.to_thread(self._probe_vram), so a sweep can run in that window. The docstring explicitly invites "optional background sweeps"; if such a sweep is driven from a background task/thread it can race with an in-flight reserve() and corrupt accounting. Acquire self._lock in this method (or clearly document single-threaded-only usage).

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

SUGGESTION: available_vram() is documented as a read-only accessor but now reclaims reservations

Calling _sweep_stale_unlocked() here mutates self._pending, decrements self._reserved_vram_mb, and emits an INFO log. A monitoring endpoint that polls available_vram() will silently reclaim capacity and generate logs, which can mask genuine hangs and is a surprising side effect for a getter. Consider making the sweep explicit (call sweep_stale() deliberately) rather than implicit in a read accessor.

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

SUGGESTION: stats() returns a read-only snapshot but now reclaims reservations

Like available_vram(), this accessor calls _sweep_stale_unlocked(), so fetching stats for dashboards/monitoring mutates manager state and logs INFO messages. This couples monitoring with capacity reclamation and can hide a real hung-installer signal (the stale reservation vanishes from the snapshot before anyone inspects it). Document the side effect or move the sweep off the read path.

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

SUGGESTION: Clamping self._reserved_vram_mb to 0 can drop accounting for active reservations

The defensive if stale_ids and self._reserved_vram_mb < 0: self._reserved_vram_mb = 0 resets the entire counter when it goes negative. If an underflow is detected, accounting already drifted, but zeroing the whole counter also discards the (positive) VRAM still legitimately held by active, non-stale reservations. That can over-report free VRAM and permit over-commit. Prefer logging/asserting the inconsistency rather than silently zeroing, or only clamp the per-entry delta.

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

SUGGESTION: Dead conditional — if candidates else 0 is unreachable

candidates always has at least one element because the variant-level min_ram_mb is appended unconditionally, so max(candidates) can never raise. The else 0 branch is dead code.

Suggested change

	return max(candidates) if candidates else 0
	return max(candidates)

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

SUGGESTION: await ..._running[download_id] has no timeout

This awaits the installer task directly with no timeout. If the mocked install path ever stalls or _running[download_id] is missing, the whole test suite hangs instead of failing fast. Wrap it in asyncio.wait_for(..., timeout=...).

---

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

[kilo-code-bot]

Code Review Summary

Status: 6 Issues Found | Recommendation: Address before merge

Overview

Severity	Count
CRITICAL	0
WARNING	1
SUGGESTION	5

Issue Details (click to expand)

WARNING

File	Line	Issue
tinyagentos/vram_reservation.py	187	sweep_stale() mutates shared state without the lock; can race with reserve() if driven from a background task/thread as the docstring invites

SUGGESTION

File	Line	Issue
tinyagentos/vram_reservation.py	214	available_vram() is a read-only accessor but now reclaims reservations and logs INFO
tinyagentos/vram_reservation.py	224	stats() snapshot now reclaims reservations, coupling monitoring with reclamation
tinyagentos/vram_reservation.py	262	Clamping self._reserved_vram_mb to 0 drops accounting for active reservations and can permit over-commit
tinyagentos/routes/models.py	167	Dead conditional if candidates else 0 is unreachable (candidates always non-empty)
tests/test_routes_models.py	423	await _running[download_id] has no timeout; a stall hangs the suite

Files Reviewed (4 files)

• tinyagentos/vram_reservation.py - 4 issues
• tinyagentos/routes/models.py - 1 issue
• tests/test_routes_models.py - 1 issue
• tests/test_vram_reservation.py - 0 issues

Fix these issues in Kilo Cloud

---

_{Reviewed by hy3-20260706:free · Input: 87.2K · Output: 15.9K · Cached: 427K}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tinyagentos/vram_reservation.py (1)

206-235: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Move the VRAM probe off the event loop tinyagentos/vram_reservation.py:206-266 still probes synchronously in available_vram()/stats(), and the only call sites are the async deny path and /api/models/vram-reservations in tinyagentos/routes/models.py. Make the accessors async (or offload the probe at those call sites) so nvidia-smi doesn’t stall the request loop; if background sweeps are planned, wrap sweep_stale() in the manager lock.

🤖 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 206 - 235, Make
available_vram() and stats() asynchronous and offload their synchronous
_probe_vram() calls from the event loop, then update the async deny path and
/api/models/vram-reservations handler to await them. Preserve the existing
reservation sweep and returned values; if sweep_stale() is used from background
tasks, protect it with the manager lock.

🧹 Nitpick comments (1)

tinyagentos/routes/models.py (1)

152-168: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

_estimated_vram_mb always maxes with variant-level min_ram_mb, not a strict fallback.

Docstring states the variant-level key is a fallback, but candidates.append(int(variant.get("min_ram_mb", 0) or 0)) unconditionally folds it into max() even when backend entries already supply a real figure. Today this is inert only because the rkllm catalog convention guarantees variant-level min_ram_mb == 0 (per this same docstring); if that invariant is ever violated (e.g. a backend legitimately requiring less VRAM than the variant's unrelated host-RAM figure), this silently over-reserves — the exact failure mode the adjacent TODO (lines 404-406) already calls out for host-RAM floors.

♻️ Proposed fix — true fallback semantics

 def _estimated_vram_mb(variant: dict) -> int:
     backend_reqs = (variant.get("requires") or {}).get("backends") or []
     candidates = [
         int(b.get("min_ram_mb", 0) or 0)
         for b in backend_reqs
-        if isinstance(b, dict)
+        if isinstance(b, dict) and b.get("min_ram_mb")
     ]
-    candidates.append(int(variant.get("min_ram_mb", 0) or 0))
-    return max(candidates) if candidates else 0
+    if candidates:
+        return max(candidates)
+    return int(variant.get("min_ram_mb", 0) or 0)

🤖 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/routes/models.py` around lines 152 - 168, Update
_estimated_vram_mb so variant-level min_ram_mb is used only when no valid
backend min_ram_mb values are present; return the maximum backend requirement
otherwise. Preserve the existing integer/default handling for backend entries
and the zero result when neither source provides a value.

🤖 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.

Outside diff comments:
In `@tinyagentos/vram_reservation.py`:
- Around line 206-235: Make available_vram() and stats() asynchronous and
offload their synchronous _probe_vram() calls from the event loop, then update
the async deny path and /api/models/vram-reservations handler to await them.
Preserve the existing reservation sweep and returned values; if sweep_stale() is
used from background tasks, protect it with the manager lock.

---

Nitpick comments:
In `@tinyagentos/routes/models.py`:
- Around line 152-168: Update _estimated_vram_mb so variant-level min_ram_mb is
used only when no valid backend min_ram_mb values are present; return the
maximum backend requirement otherwise. Preserve the existing integer/default
handling for backend entries and the zero result when neither source provides a
value.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 6df80eef-0fa1-4afb-84f4-9524263d208d

📥 Commits

Reviewing files that changed from the base of the PR and between 5c4afbb and 747e24b.

📒 Files selected for processing (4)

• tests/test_routes_models.py
• tests/test_vram_reservation.py
• tinyagentos/routes/models.py
• tinyagentos/vram_reservation.py