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

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:

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

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}",