Skip to content

fix(models): VRAM reservation fails closed on non-NVIDIA + TOCTOU gate reads dead manifest key (audit follow-up to #1725) #1766

Description

@jaylfc

Follow-up to #1725 from a code audit of the merged tree. The audit found the VRAM guard is currently either inert or wrongly fail-closed depending on catalog data, so the headline TOCTOU protection is not yet protecting the path it was written for. All in tinyagentos/vram_reservation.py + tinyagentos/routes/models.py. Over to @hognek.

1. HIGH: fails closed on non-NVIDIA hardware. _probe_vram() returns (0, 0) when nvidia-smi is absent (vram_reservation.py:170-182), so any reserve(vram_mb > 0) is always denied with 503 "Insufficient VRAM ... have 0 MiB" on AMD/Apple/Rockchip. The pre-existing cluster ledger made the opposite, deliberate call: claim_lease never refuses on VRAM grounds when free_vram_mb is None because it cannot prove insufficiency (cluster/manager.py:323-330). The primary consumer here is the rkllama pull path, which runs exclusively on Rockchip boxes where nvidia-smi never exists; and /api/models/pull with a client-supplied required_vram_mb > 0 503s on any non-NVIDIA machine today. Fix shape: return None for no-probe, distinguish it from a real 0, and fail open (bookkeeping only) to match claim_lease.

2. MEDIUM: the rkllama TOCTOU gate reads the wrong manifest key and is dead against the whole catalog. routes/models.py:390 gates on variant.get("min_ram_mb", 0), but every rkllama-requiring variant in the catalog has variant-level min_ram_mb: 0; the real requirement (e.g. 2048) lives at variant.requires.backends[].min_ram_mb and is never read (verified across the rkllm manifests, e.g. app-catalog/models/qwen2.5-1.5b-rkllm/manifest.yaml lines 16 vs 27). So two concurrent rkllm pulls still both pass. Note the interaction: if someone "fixes" a manifest by setting the variant-level key while finding 1 is unfixed, every install of that model 503s on the Pi. Fix shape: read the backend-level requirement (max across required backends), after fixing finding 1.

3. LOW, same files: (a) _probe_vram runs a blocking subprocess.run(timeout=3) inside async reserve() while holding the asyncio lock (and again in available_vram()/stats()), stalling the event loop up to 3s on a request hot path; move to a thread or async subprocess. (b) VramReservation.created_at exists but nothing sweeps stale reservations, so a hung installer holds its reservation until controller restart and every later pull 503s; add a TTL sweep.

Acceptance: on a Rockchip box with no nvidia-smi, a pull with a real backend-level min_ram_mb proceeds (bookkeeping recorded, no 503); two concurrent rkllm pulls of the same large model on NVIDIA hardware still exercise the atomic check; a reservation older than its TTL is reclaimed. Usual bar: wait for Kilo/CodeRabbit and resolve every finding. No rush.

Metadata

Metadata

Assignees

No one assigned

    Labels

    fork-auditTracking issue for fork-vs-upstream drift auditskilo-auto-fixAuto-generated label by Kilokilo-triagedAuto-generated label by KilomodelsModel management and inferencenpu-rk3588RK3588 NPU work

    Projects

    Status
    Done

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions