From 747e24becb3e578cc509121c0aef23ab570aeed3 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Sun, 12 Jul 2026 17:08:17 +0100 Subject: [PATCH] fix(models): VRAM reservation TTL sweep + #1766 acceptance coverage Completes the remaining #1766 work after the #1767 hotfix (fail-open on no-probe hardware, backend-level min_ram_mb gate, asyncio.to_thread probe). - Reclaim VramReservation entries older than a configurable TTL (default 1h) so a hung installer cannot hold capacity until controller restart. Sweep runs from reserve(), available_vram(), stats(), and public sweep_stale(). - Extract _estimated_vram_mb() so the rkllama pull gate clearly reads requires.backends[].min_ram_mb (max across backends, variant fallback). - Tests: no-probe + real backend min_ram proceeds (no 503), concurrent large reserves on NVIDIA still atomic, TTL reclaim, and route-level 503 when free VRAM is measurable and insufficient. Closes #1766 --- tests/test_routes_models.py | 101 ++++++++++++++++++++- tests/test_vram_reservation.py | 150 +++++++++++++++++++++++++++++--- tinyagentos/routes/models.py | 30 ++++--- tinyagentos/vram_reservation.py | 74 +++++++++++++++- 4 files changed, 329 insertions(+), 26 deletions(-) diff --git a/tests/test_routes_models.py b/tests/test_routes_models.py index a4918fb27..4e1b3ba4f 100644 --- a/tests/test_routes_models.py +++ b/tests/test_routes_models.py @@ -6,7 +6,12 @@ from unittest.mock import AsyncMock, patch, MagicMock from tinyagentos.app import create_app from tinyagentos.installers.model_paths import models_root -from tinyagentos.routes.models import DEFAULT_MODELS_DIR, get_downloaded_models +from tinyagentos.routes.models import ( + DEFAULT_MODELS_DIR, + _estimated_vram_mb, + get_downloaded_models, +) +from tinyagentos.vram_reservation import VramReservationManager class TestDefaultModelsDir: @@ -41,10 +46,14 @@ def catalog_with_models(tmp_path): {"id": "small", "name": "Small GGUF", "format": "gguf", "size_mb": 100, "min_ram_mb": 512, "download_url": "https://example.com/small.gguf", "backend": ["ollama", "llama-cpp"]}, + # Catalog-faithful rkllm shape: variant min_ram_mb is 0; the real + # floor lives on requires.backends[].min_ram_mb (#1766 MEDIUM). {"id": "npu", "name": "NPU RKLLM", "format": "rkllm", "size_mb": 200, "min_ram_mb": 0, "download_url": "https://example.com/npu.rkllm", "backend": ["rkllama"], "requires_npu": ["rk3588"], - "requires": {"backends": [{"id": "rkllama"}]}}, + "requires": {"backends": [ + {"id": "rkllama", "targets": ["rockchip"], "min_ram_mb": 2048}, + ]}}, ], "hardware_tiers": {"arm-npu-16gb": {"recommended": "npu", "fallback": "small"}}, "install": {"method": "download"}, @@ -387,6 +396,94 @@ async def test_rkllama_variant_failure_skips_registry_and_catalog_refresh( installed = models_app.state.registry.list_installed() assert not any(row["id"] == "test-model" for row in installed) + async def test_rkllama_no_probe_backend_min_ram_no_503( + self, models_app, models_client + ): + """Acceptance (#1766): no nvidia-smi + real backend-level min_ram_mb + proceeds with bookkeeping and does not 503. + """ + mgr = VramReservationManager() + mgr._probe_vram = staticmethod(lambda: None) # type: ignore[method-assign] + models_app.state.vram_reservation = mgr + + with patch( + "tinyagentos.installers.rkllama_installer.RkllamaInstaller.install", + new=AsyncMock(return_value={"success": True, "app_id": "test-model"}), + ): + resp = await models_client.post( + "/api/models/download", + json={"app_id": "test-model", "variant_id": "npu"}, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["status"] == "started" + # While the installer is in flight the reservation is held. + # Drain it so release runs in the finally block. + download_id = data["download_id"] + await models_app.state.download_manager._running[download_id] + # After success the reservation is released. + assert mgr.pending_count == 0 + + async def test_rkllama_insufficient_vram_503_on_nvidia( + self, models_app, models_client + ): + """On measurable NVIDIA VRAM the atomic check still denies when + free capacity cannot cover the backend min_ram_mb floor. + """ + mgr = VramReservationManager() + mgr._probe_vram = staticmethod(lambda: (1024, 8192)) # type: ignore[method-assign] + models_app.state.vram_reservation = mgr + + with patch( + "tinyagentos.installers.rkllama_installer.RkllamaInstaller.install", + new=AsyncMock(return_value={"success": True, "app_id": "test-model"}), + ) as mock_install: + resp = await models_client.post( + "/api/models/download", + json={"app_id": "test-model", "variant_id": "npu"}, + ) + assert resp.status_code == 503 + body = resp.json() + assert "Insufficient VRAM" in body["error"] + assert "2048" in body["error"] + mock_install.assert_not_awaited() + assert mgr.pending_count == 0 + + +class TestEstimatedVramMb: + """Unit coverage for the #1766 backend-level min_ram_mb gate.""" + + def test_reads_backend_min_ram_when_variant_is_zero(self): + variant = { + "min_ram_mb": 0, + "requires": { + "backends": [ + {"id": "rkllama", "min_ram_mb": 2048}, + ], + }, + } + assert _estimated_vram_mb(variant) == 2048 + + def test_max_across_backends(self): + variant = { + "min_ram_mb": 0, + "requires": { + "backends": [ + {"id": "a", "min_ram_mb": 1024}, + {"id": "b", "min_ram_mb": 4096}, + ], + }, + } + assert _estimated_vram_mb(variant) == 4096 + + def test_falls_back_to_variant_key(self): + variant = {"min_ram_mb": 512, "requires": {"backends": [{"id": "ollama"}]}} + assert _estimated_vram_mb(variant) == 512 + + def test_missing_requires_uses_variant_only(self): + assert _estimated_vram_mb({"min_ram_mb": 256}) == 256 + assert _estimated_vram_mb({}) == 0 + @pytest.mark.asyncio class TestModelRecommendations: diff --git a/tests/test_vram_reservation.py b/tests/test_vram_reservation.py index 322e7b0f1..999ebe772 100644 --- a/tests/test_vram_reservation.py +++ b/tests/test_vram_reservation.py @@ -1,4 +1,4 @@ -"""Tests for atomic VRAM reservation (taOS #1706 TOCTOU fix). +"""Tests for atomic VRAM reservation (taOS #1706 TOCTOU fix / #1766 audit). Verifies that: - Concurrent reserve calls cannot over-commit VRAM. @@ -7,21 +7,31 @@ - release() is idempotent. - Zero-VRAM reservations always succeed. - stats() reflects the current state. +- No-probe hardware fails OPEN with bookkeeping (#1766 HIGH). +- Stale reservations past TTL are reclaimed (#1766 LOW). """ import asyncio +import time import pytest -from tinyagentos.vram_reservation import VramReservationManager +from tinyagentos.vram_reservation import ( + DEFAULT_RESERVATION_TTL_SECONDS, + VramReservationManager, +) # ── test helpers ──────────────────────────────────────────────────── -def _manager_with_probe(free_mb: int, total_mb: int) -> VramReservationManager: +def _manager_with_probe( + free_mb: int, + total_mb: int, + ttl_seconds: float = DEFAULT_RESERVATION_TTL_SECONDS, +) -> VramReservationManager: """Return a manager whose _probe_vram returns fixed values.""" - mgr = VramReservationManager() + mgr = VramReservationManager(ttl_seconds=ttl_seconds) def _fake_probe() -> tuple[int, int]: return free_mb, total_mb @@ -228,22 +238,24 @@ class TestRealProbe: @pytest.mark.asyncio async def test_reserve_with_real_probe(self): - """Reserving 1 MiB should always succeed on a real GPU.""" + """Reserving 1 MiB should always succeed on a real GPU, and on + no-probe hardware fail-open also admits (bookkeeping only).""" 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) + # With a real GPU and free VRAM this succeeds. Without nvidia-smi the + # probe returns None and reserve() fails open, so this also succeeds. + assert res is not None + mgr.release(res.reservation_id) # ── no-probe hardware (AMD / Apple / Rockchip): fail open ──────────── -def _manager_without_probe() -> VramReservationManager: +def _manager_without_probe( + ttl_seconds: float = 3600.0, +) -> VramReservationManager: """Return a manager whose probe reports no VRAM visibility (None).""" - mgr = VramReservationManager() + mgr = VramReservationManager(ttl_seconds=ttl_seconds) mgr._probe_vram = staticmethod(lambda: None) # type: ignore[method-assign] return mgr @@ -265,12 +277,26 @@ async def test_reserve_admits_without_probe(self): mgr.release(res.reservation_id) assert mgr.reserved_vram_mb == 0 + @pytest.mark.asyncio + async def test_no_probe_real_backend_min_ram_proceeds(self): + """Acceptance (#1766): on a box with no nvidia-smi, a pull with a + real backend-level min_ram_mb proceeds with bookkeeping and no deny. + """ + mgr = _manager_without_probe() + # 2048 is the catalog figure for qwen2.5-1.5b-rkllm backends.min_ram_mb + res = await mgr.reserve(2048, caller="rkllama-pull:qwen2.5-1.5b-rkllm") + assert res is not None + assert res.vram_mb == 2048 + assert mgr.reserved_vram_mb == 2048 + assert mgr.pending_count == 1 + @pytest.mark.asyncio async def test_stats_reports_probe_unavailable(self): mgr = _manager_without_probe() s = mgr.stats() assert s["probe_available"] is False assert s["free_vram_mb"] == 0 and s["total_vram_mb"] == 0 + assert "ttl_seconds" in s @pytest.mark.asyncio async def test_probe_present_still_denies_real_insufficiency(self): @@ -279,3 +305,103 @@ async def test_probe_present_still_denies_real_insufficiency(self): mgr = _manager_with_probe(free_mb=1024, total_mb=8192) res = await mgr.reserve(4096, caller="too-big") assert res is None + + +# ── concurrent rkllm-style atomic check on NVIDIA ──────────────────── + + +class TestConcurrentRkllmStyle: + """Acceptance (#1766): two concurrent large-model pulls on NVIDIA still + hit the atomic check (only one admitted when free VRAM cannot fit both). + """ + + @pytest.mark.asyncio + async def test_two_concurrent_large_reserves_only_one_admitted(self): + # 8192 free, each "rkllm" pull wants 5000 (like a large model). + mgr = _manager_with_probe(8192, 8192) + results: list[bool] = [] + + async def try_reserve(caller: str): + res = await mgr.reserve(5000, caller=caller) + results.append(res is not None) + + await asyncio.gather( + try_reserve("rkllama-pull:model-a"), + try_reserve("rkllama-pull:model-a"), + ) + + assert sum(results) == 1 + assert mgr.reserved_vram_mb == 5000 + assert mgr.pending_count == 1 + + +# ── TTL sweep for hung installers ──────────────────────────────────── + + +class TestTtlSweep: + """Acceptance (#1766): a reservation older than its TTL is reclaimed.""" + + @pytest.mark.asyncio + async def test_sweep_reclaims_stale_reservation(self): + mgr = _manager_with_probe(8192, 8192) + mgr._ttl_seconds = 60.0 + res = await mgr.reserve(4000, caller="hung-installer") + assert res is not None + assert mgr.reserved_vram_mb == 4000 + + # Age the reservation past the TTL. + res.created_at = time.time() - 120.0 + reclaimed = mgr.sweep_stale() + assert reclaimed == 1 + assert mgr.reserved_vram_mb == 0 + assert mgr.pending_count == 0 + # Idempotent: already released id is unknown. + assert mgr.release(res.reservation_id) is False + + @pytest.mark.asyncio + async def test_fresh_reservation_not_reclaimed(self): + mgr = _manager_with_probe(8192, 8192) + mgr._ttl_seconds = 3600.0 + res = await mgr.reserve(2000, caller="active") + assert res is not None + assert mgr.sweep_stale() == 0 + assert mgr.pending_count == 1 + assert mgr.reserved_vram_mb == 2000 + + @pytest.mark.asyncio + async def test_reserve_sweeps_before_admission(self): + """Stale holds must not permanently starve a later pull.""" + mgr = _manager_with_probe(5000, 8192) + mgr._ttl_seconds = 30.0 + stuck = await mgr.reserve(4000, caller="hung") + assert stuck is not None + stuck.created_at = time.time() - 120.0 + + # Without sweep this would deny (5000 free - 4000 reserved < 4000). + # With sweep the stale hold is reclaimed first. + next_res = await mgr.reserve(4000, caller="fresh") + assert next_res is not None + assert mgr.reserved_vram_mb == 4000 + assert mgr.pending_count == 1 + + @pytest.mark.asyncio + async def test_ttl_zero_disables_expiry(self): + mgr = VramReservationManager(ttl_seconds=0) + mgr._probe_vram = staticmethod(lambda: (8192, 8192)) # type: ignore[method-assign] + res = await mgr.reserve(1000, caller="permanent") + assert res is not None + res.created_at = time.time() - 10_000.0 + assert mgr.sweep_stale() == 0 + assert mgr.pending_count == 1 + + @pytest.mark.asyncio + async def test_available_vram_sweeps_stale(self): + mgr = _manager_with_probe(8192, 8192) + mgr._ttl_seconds = 10.0 + res = await mgr.reserve(3000, caller="old") + assert res is not None + res.created_at = time.time() - 60.0 + free, total = mgr.available_vram() + assert free == 8192 # stale hold reclaimed + assert total == 8192 + assert mgr.pending_count == 0 \ No newline at end of file diff --git a/tinyagentos/routes/models.py b/tinyagentos/routes/models.py index 984df17ed..5452c64f8 100644 --- a/tinyagentos/routes/models.py +++ b/tinyagentos/routes/models.py @@ -149,6 +149,24 @@ def _is_service_installed(variant: dict) -> bool: return False +def _estimated_vram_mb(variant: dict) -> int: + """VRAM floor for the pull TOCTOU gate (#1706 / #1766). + + Every rkllm variant in the catalog has variant-level ``min_ram_mb: 0``; + the real requirement (e.g. 2048) lives under + ``requires.backends[].min_ram_mb``. Take the max across required + backends, falling back to the variant-level key. + """ + 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) + ] + candidates.append(int(variant.get("min_ram_mb", 0) or 0)) + return max(candidates) if candidates else 0 + + def _variant_compatibility(variant: dict, hardware_profile) -> str: """Return green/yellow/red based on hardware compatibility.""" ram_mb = hardware_profile.ram_mb @@ -386,17 +404,7 @@ async def download_model(request: Request, body: DownloadRequest): # TODO(#1725): replace min_ram_mb heuristic with a real per-model VRAM # estimate. On CUDA GGUF the host-RAM floor over-reserves and causes # occasional false 503s in multi-model setups. - # The real requirement lives on the backend requirements, not the - # variant top level: every rkllm variant in the catalog has - # variant-level min_ram_mb 0 with the true value (e.g. 2048) under - # requires.backends[].min_ram_mb, so reading only the variant key made - # this gate dead code (audit follow-up, #1766). Take the max across - # required backends, falling back to the variant-level key. - _backend_reqs = (variant.get("requires") or {}).get("backends") or [] - estimated_vram = max( - [int(b.get("min_ram_mb", 0) or 0) for b in _backend_reqs if isinstance(b, dict)] - + [int(variant.get("min_ram_mb", 0) or 0)] - ) + estimated_vram = _estimated_vram_mb(variant) if vram_mgr is not None and estimated_vram > 0: reservation = await vram_mgr.reserve( estimated_vram, caller=f"rkllama-pull:{body.app_id}", diff --git a/tinyagentos/vram_reservation.py b/tinyagentos/vram_reservation.py index af300a4f1..e248aafd7 100644 --- a/tinyagentos/vram_reservation.py +++ b/tinyagentos/vram_reservation.py @@ -22,6 +22,12 @@ The reservation is *not* a context manager — callers must release explicitly (typically in a ``finally`` block) so the reservation can outlive a single ``async with`` block while the model loads. + +When the host has no nvidia-smi (AMD/Apple/Rockchip), the probe returns +``None`` and :meth:`reserve` fails OPEN with bookkeeping only, matching +cluster ``claim_lease`` when ``free_vram_mb is None``. Stale reservations +(hung installers that never call :meth:`release`) are reclaimed by a TTL +sweep so later pulls are not stuck at 503 until controller restart. """ from __future__ import annotations @@ -34,6 +40,11 @@ logger = logging.getLogger(__name__) +# Default TTL for a reservation held by a hung installer. Model pulls can +# take a while on slow links; 1h is generous for normal installs but still +# reclaims stuck capacity so later pulls are not permanently denied. +DEFAULT_RESERVATION_TTL_SECONDS = 3600.0 + @dataclass class VramReservation: @@ -54,12 +65,20 @@ class VramReservationManager: Zero-VRAM reservations (``vram_mb <= 0``) always succeed as a lightweight marker — no lock acquisition, no probe, no accounting. + + Reservations older than ``ttl_seconds`` are reclaimed on the next + :meth:`reserve` / :meth:`sweep_stale` / accessor call so a crashed + installer cannot hold capacity forever. """ - def __init__(self) -> None: + def __init__( + self, + ttl_seconds: float = DEFAULT_RESERVATION_TTL_SECONDS, + ) -> None: self._lock = asyncio.Lock() self._reserved_vram_mb: int = 0 self._pending: dict[str, VramReservation] = {} + self._ttl_seconds = float(ttl_seconds) # ── public API ────────────────────────────────────────────────── @@ -77,6 +96,10 @@ async def reserve( *vram_mb* ≤ 0 always succeeds without probing VRAM or acquiring the lock (a lightweight no-op marker). + + When no hardware probe is available the call fails OPEN: the + reservation is recorded for bookkeeping but never refused, matching + cluster ``claim_lease`` when free VRAM is unknown. """ if vram_mb <= 0: return VramReservation( @@ -87,6 +110,10 @@ async def reserve( ) async with self._lock: + # Drop hung installers before admission so stale holds cannot + # permanently starve later pulls. + self._sweep_stale_unlocked() + # Probe in a thread: nvidia-smi is a blocking subprocess and this # runs on a request hot path while holding the lock. probe = await asyncio.to_thread(self._probe_vram) @@ -150,6 +177,15 @@ def release(self, reservation_id: str) -> bool: return True return False + def sweep_stale(self, now: float | None = None) -> int: + """Reclaim reservations older than the configured TTL. + + Returns the number of reservations reclaimed. Called automatically + from :meth:`reserve` and the sync accessors; exposed for tests and + optional background sweeps. + """ + return self._sweep_stale_unlocked(now=now) + # ── read-only accessors ───────────────────────────────────────── @property @@ -162,6 +198,11 @@ def pending_count(self) -> int: """Number of active reservations.""" return len(self._pending) + @property + def ttl_seconds(self) -> float: + """Configured reservation TTL in seconds (0 disables expiry).""" + return self._ttl_seconds + def available_vram(self) -> tuple[int, int]: """Return ``(effective_free_mb, total_mb)`` after subtracting in-flight reservations from the hardware probe. @@ -170,6 +211,7 @@ def available_vram(self) -> tuple[int, int]: path should already have been admitted fail-open by :meth:`reserve` (a deny message never reports these zeros as real capacity). """ + self._sweep_stale_unlocked() probe = self._probe_vram() if probe is None: return 0, 0 @@ -179,6 +221,7 @@ def available_vram(self) -> tuple[int, int]: def stats(self) -> dict: """Return a snapshot for monitoring / debug endpoints.""" + self._sweep_stale_unlocked() probe = self._probe_vram() free_mb, total_mb = probe if probe is not None else (0, 0) return { @@ -188,10 +231,39 @@ def stats(self) -> dict: "reserved_vram_mb": self._reserved_vram_mb, "effective_free_mb": max(0, free_mb - self._reserved_vram_mb), "pending_reservations": self.pending_count, + "ttl_seconds": self._ttl_seconds, } # ── internal ──────────────────────────────────────────────────── + def _sweep_stale_unlocked(self, now: float | None = None) -> int: + """Drop reservations whose age exceeds the TTL. Not lock-guarded + on its own; :meth:`reserve` calls it under ``_lock``. + """ + if self._ttl_seconds <= 0: + return 0 + now = time.time() if now is None else now + stale_ids = [ + rid + for rid, res in self._pending.items() + if (now - res.created_at) > self._ttl_seconds + ] + for rid in stale_ids: + res = self._pending.pop(rid, None) + if res is None: + continue + self._reserved_vram_mb -= res.vram_mb + logger.info( + "vram-reservation: reclaimed stale %d MiB for %r (id=%s, " + "age=%.0fs > ttl=%.0fs)", + res.vram_mb, res.caller, rid, + now - res.created_at, self._ttl_seconds, + ) + if stale_ids and self._reserved_vram_mb < 0: + # Defensive: never let accounting go negative after reclaim. + self._reserved_vram_mb = 0 + return len(stale_ids) + @staticmethod def _probe_vram() -> tuple[int, int] | None: """Probe real-time free/total VRAM via nvidia-smi.