From e4c33152fd4af8e2789c14fa17640cab3be49589 Mon Sep 17 00:00:00 2001 From: jaylfc Date: Thu, 9 Jul 2026 21:55:49 +0100 Subject: [PATCH] fix(worker,models): audit hotfixes - manifest crash, VRAM fail-open, real min_ram_mb key Fixes the urgent findings from the full code audit of the recent merges (follow-ups #1765/#1766 track the deeper redesign parts): - worker_manifest.load_manifest: a malformed or non-object manifest file is external input and now degrades to the empty manifest with a logged warning instead of raising. Previously one typo'd comma in /etc/taos/worker-models.json crash-looped the whole worker under systemd. - worker/agent.py enrichment: skip entries without model_id (logged) and wrap the whole enrichment in a fail-soft try/except so no manifest problem can take detect_backends down. Heartbeat send failures are now logged before being swallowed so a payload-build bug no longer masquerades as controller unreachability. - vram_reservation: the probe now returns None (cannot know) distinct from (0, 0) (known-full), and reserve() fails OPEN with bookkeeping-only admission when no probe exists, matching cluster claim_lease semantics. Previously every reservation was denied 503 on AMD/Apple/Rockchip, i.e. on exactly the rkllama hardware the guard was written for. The probe also moved to asyncio.to_thread so nvidia-smi no longer blocks the event loop while holding the reservation lock. stats() gains probe_available. - routes/models.py: the rkllama pull gate now reads the real backend-level requirement (max of variant.requires.backends[].min_ram_mb, falling back to the variant key). Every rkllm variant has variant-level min_ram_mb 0, so the previous read made the TOCTOU gate dead code against the entire catalog. - cluster/backend_services.load_managed_backends: skipped manifests are now logged so a device-side catalog edit cannot silently drop a backend out of #1743 recovery. - install-rknpu.sh: replace the stale copy-of-truth comment (old port 8080 and pre-consolidation models path) with an accurate one. Tests: the two tests that blessed the crash now assert the degrade behaviour, plus new coverage for non-object manifests, no-probe fail-open, probe-unavailable stats, and a regression guard that a real probe showing insufficiency still denies. Docs-Reviewed: behaviour hardening within existing features; no documented install step, endpoint path, or response shape changed (installer diff is a comment correction only). --- scripts/install-rknpu.sh | 9 ++-- tests/test_vram_reservation.py | 43 ++++++++++++++++ tests/test_worker_manifest.py | 32 +++++++++--- tinyagentos/cluster/backend_services.py | 12 +++++ tinyagentos/routes/models.py | 15 ++++-- tinyagentos/vram_reservation.py | 60 ++++++++++++++++------ tinyagentos/worker/agent.py | 68 ++++++++++++++++--------- tinyagentos/worker/worker_manifest.py | 27 ++++++++-- 8 files changed, 206 insertions(+), 60 deletions(-) diff --git a/scripts/install-rknpu.sh b/scripts/install-rknpu.sh index 0f8cf4a7..42df4d47 100755 --- a/scripts/install-rknpu.sh +++ b/scripts/install-rknpu.sh @@ -631,10 +631,11 @@ pull_models() { install_systemd_unit() { local unit="/etc/systemd/system/rkllama.service" local exec_start - # This ExecStart line is copy-of-truth from the live orange pi: - # rkllama_server --processor rk3588 --port 8080 \ - # --models /home/jay/rkllama/models \ - # --preload qwen3-embedding-0.6b,qwen3-reranker-0.6b,qmd-query-expansion + # Matches the unit shape running on the reference RK3588 deployment + # (port 7833, unified models root, --preload of the three shared + # models). Kept in sync with the live Pi unit as of the 1.3.0 deploy; + # the audit flagged the previous comment as stale (it referenced the + # old port 8080 and the pre-consolidation models path). exec_start="$RKLLAMA_VENV/bin/python $RKLLAMA_VENV/bin/rkllama_server --processor $SOC --port $RKLLAMA_PORT --models $RKLLAMA_MODELS --preload qwen3-embedding-0.6b,qwen3-reranker-0.6b,qmd-query-expansion" log "installing $unit" diff --git a/tests/test_vram_reservation.py b/tests/test_vram_reservation.py index 0054b383..322e7b0f 100644 --- a/tests/test_vram_reservation.py +++ b/tests/test_vram_reservation.py @@ -236,3 +236,46 @@ async def test_reserve_with_real_probe(self): # this will fail — that's fine, the test is a best-effort smoke. if res is not None: mgr.release(res.reservation_id) + + +# ── no-probe hardware (AMD / Apple / Rockchip): fail open ──────────── + + +def _manager_without_probe() -> VramReservationManager: + """Return a manager whose probe reports no VRAM visibility (None).""" + mgr = VramReservationManager() + mgr._probe_vram = staticmethod(lambda: None) # type: ignore[method-assign] + return mgr + + +class TestNoProbeFailOpen: + """Without a VRAM probe we cannot prove insufficiency, so admission + fails OPEN (bookkeeping only), matching cluster claim_lease semantics. + A closed fail here 503s every model pull on non-NVIDIA hosts.""" + + @pytest.mark.asyncio + async def test_reserve_admits_without_probe(self): + mgr = _manager_without_probe() + res = await mgr.reserve(4096, caller="rkllama-pull:test") + assert res is not None + assert res.vram_mb == 4096 + # Bookkeeping still recorded so a future probe-aware path sees it. + assert mgr.reserved_vram_mb == 4096 + assert mgr.pending_count == 1 + mgr.release(res.reservation_id) + assert mgr.reserved_vram_mb == 0 + + @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 + + @pytest.mark.asyncio + async def test_probe_present_still_denies_real_insufficiency(self): + # Regression guard: fail-open applies ONLY to the no-probe case; a + # real probe showing insufficient VRAM must still deny. + mgr = _manager_with_probe(free_mb=1024, total_mb=8192) + res = await mgr.reserve(4096, caller="too-big") + assert res is None diff --git a/tests/test_worker_manifest.py b/tests/test_worker_manifest.py index be9a74ea..60497d74 100644 --- a/tests/test_worker_manifest.py +++ b/tests/test_worker_manifest.py @@ -99,8 +99,9 @@ def test_manifest_with_empty_models(self): finally: os.unlink(tmp_path) - def test_invalid_json_raises(self): - """Invalid JSON raises json.JSONDecodeError.""" + def test_invalid_json_degrades_to_empty(self): + """Invalid JSON is external input: it must degrade to the empty + manifest (logged), never raise and take the worker down.""" with tempfile.NamedTemporaryFile( mode="w", suffix=".json", delete=False ) as f: @@ -108,8 +109,22 @@ def test_invalid_json_raises(self): tmp_path = f.name try: - with pytest.raises(json.JSONDecodeError): - load_manifest(tmp_path) + result = load_manifest(tmp_path) + assert result == {"resource_id": "", "models": []} + finally: + os.unlink(tmp_path) + + def test_non_object_top_level_degrades_to_empty(self): + """A JSON array / scalar top level degrades to the empty manifest.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f: + f.write('["not", "an", "object"]') + tmp_path = f.name + + try: + result = load_manifest(tmp_path) + assert result == {"resource_id": "", "models": []} finally: os.unlink(tmp_path) @@ -169,8 +184,9 @@ def test_default_path_fallback(self, monkeypatch): class TestLoadManifestEdgeCases: - def test_empty_file_raises(self): - """An empty file is not valid JSON.""" + def test_empty_file_degrades_to_empty(self): + """An empty file is not valid JSON; it degrades to the empty + manifest rather than raising (bad input must not brick the worker).""" with tempfile.NamedTemporaryFile( mode="w", suffix=".json", delete=False ) as f: @@ -178,8 +194,8 @@ def test_empty_file_raises(self): tmp_path = f.name try: - with pytest.raises(json.JSONDecodeError): - load_manifest(tmp_path) + result = load_manifest(tmp_path) + assert result == {"resource_id": "", "models": []} finally: os.unlink(tmp_path) diff --git a/tinyagentos/cluster/backend_services.py b/tinyagentos/cluster/backend_services.py index 8465f536..5de4d751 100644 --- a/tinyagentos/cluster/backend_services.py +++ b/tinyagentos/cluster/backend_services.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import logging import os from dataclasses import dataclass from pathlib import Path @@ -25,6 +26,8 @@ import httpx import yaml +logger = logging.getLogger(__name__) + VALID_SCOPES = {"system", "user"} @@ -54,8 +57,10 @@ def load_managed_backends(catalog_root: Path) -> list[ManagedBackend]: try: data = yaml.safe_load(manifest.read_text()) except yaml.YAMLError: + logger.warning("skipping unparseable service manifest %s", manifest) continue if not isinstance(data, dict): + logger.warning("skipping non-mapping service manifest %s", manifest) continue lifecycle = data.get("lifecycle") if not isinstance(lifecycle, dict) or lifecycle.get("auto_manage") is not True: @@ -63,6 +68,13 @@ def load_managed_backends(catalog_root: Path) -> list[ManagedBackend]: unit = lifecycle.get("unit") scope = lifecycle.get("scope") if not unit or scope not in VALID_SCOPES: + # The CI lint guarantees these fields for auto-managed services; + # a device-side catalog edit that breaks the contract would + # silently drop the backend out of recovery, so say so. + logger.warning( + "skipping managed backend %s: lifecycle.unit/scope invalid " + "(unit=%r scope=%r)", manifest, unit, scope, + ) continue health = lifecycle.get("health") if isinstance(lifecycle.get("health"), dict) else {} out.append( diff --git a/tinyagentos/routes/models.py b/tinyagentos/routes/models.py index 4bfbb4ed..984df17e 100644 --- a/tinyagentos/routes/models.py +++ b/tinyagentos/routes/models.py @@ -385,9 +385,18 @@ async def download_model(request: Request, body: DownloadRequest): reservation = None # 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. Safe for v1 (fails - # closed), but a model-card estimate would be more accurate. - estimated_vram = int(variant.get("min_ram_mb", 0) or 0) # heuristic + # 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)] + ) 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 1b2eee78..af300a4f 100644 --- a/tinyagentos/vram_reservation.py +++ b/tinyagentos/vram_reservation.py @@ -87,16 +87,32 @@ async def reserve( ) async with self._lock: - free_mb, total_mb = self._probe_vram() - effective_free = max(0, free_mb - self._reserved_vram_mb) - - if effective_free < vram_mb: - logger.debug( - "vram-reservation: denied — need %d MiB, have %d MiB " - "(%d free − %d reserved)", - vram_mb, effective_free, free_mb, self._reserved_vram_mb, + # 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) + if probe is None: + # No VRAM probe on this hardware (AMD/Apple/Rockchip or a + # failed nvidia-smi). Fail OPEN, matching the cluster lease + # ledger's convention (manager.claim_lease): we cannot prove + # insufficiency, so admit and keep bookkeeping only. A closed + # fail here would 503 every model pull on non-NVIDIA hosts. + free_mb, total_mb = 0, 0 + effective_free = vram_mb # unknown; grant-log then shows 0 + logger.info( + "vram-reservation: no VRAM probe available; admitting %d " + "MiB for %r unenforced (bookkeeping only)", vram_mb, caller, ) - return None + else: + free_mb, total_mb = probe + effective_free = max(0, free_mb - self._reserved_vram_mb) + + if effective_free < vram_mb: + logger.debug( + "vram-reservation: denied — need %d MiB, have %d MiB " + "(%d free − %d reserved)", + vram_mb, effective_free, free_mb, self._reserved_vram_mb, + ) + return None reservation_id = f"vr_{secrets.token_hex(4)}" reservation = VramReservation( @@ -148,15 +164,25 @@ def pending_count(self) -> int: def available_vram(self) -> tuple[int, int]: """Return ``(effective_free_mb, total_mb)`` after subtracting - in-flight reservations from the hardware probe.""" - free_mb, total_mb = self._probe_vram() + in-flight reservations from the hardware probe. + + When no probe is available, returns ``(0, 0)``; callers on that + path should already have been admitted fail-open by :meth:`reserve` + (a deny message never reports these zeros as real capacity). + """ + probe = self._probe_vram() + if probe is None: + return 0, 0 + free_mb, total_mb = probe effective_free = max(0, free_mb - self._reserved_vram_mb) return effective_free, total_mb def stats(self) -> dict: """Return a snapshot for monitoring / debug endpoints.""" - free_mb, total_mb = self._probe_vram() + probe = self._probe_vram() + free_mb, total_mb = probe if probe is not None else (0, 0) return { + "probe_available": probe is not None, "free_vram_mb": free_mb, "total_vram_mb": total_mb, "reserved_vram_mb": self._reserved_vram_mb, @@ -167,11 +193,13 @@ def stats(self) -> dict: # ── internal ──────────────────────────────────────────────────── @staticmethod - def _probe_vram() -> tuple[int, int]: + def _probe_vram() -> tuple[int, int] | None: """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. + Returns ``(free_mb, total_mb)``, or ``None`` when no nvidia-smi + is available or the probe fails. ``None`` (cannot know) is + deliberately distinct from ``(0, 0)`` (known-full): admission + fails open on ``None``, matching cluster claim_lease semantics. """ try: from tinyagentos.system_stats import read_nvidia_vram @@ -183,4 +211,4 @@ def _probe_vram() -> tuple[int, int]: return free_mb, total_mb except Exception: logger.debug("vram-reservation: probe failed", exc_info=True) - return 0, 0 + return None diff --git a/tinyagentos/worker/agent.py b/tinyagentos/worker/agent.py index ac71d8f3..8c5f69e1 100644 --- a/tinyagentos/worker/agent.py +++ b/tinyagentos/worker/agent.py @@ -172,30 +172,44 @@ async def detect_backends(self) -> list[dict]: SOFTWARE_TO_BACKEND_TYPE, ) - manifest = load_manifest() - if manifest.get("models"): - for backend in backends: - backend_type = backend["type"] - probed_names = { - m.get("name", "") for m in backend.get("models", []) - } - available = [] - for m in manifest["models"]: - if SOFTWARE_TO_BACKEND_TYPE.get(m.get("software", "")) != backend_type: - continue - model_id = m["model_id"] - status = "loaded" if model_id in probed_names else "available" - available.append({ - "model_id": model_id, - "capability": m.get("capability", ""), - "software": m.get("software", ""), - "port": m.get("port", 0), - "vram_required_gb": m.get("vram_required_gb", 0.0), - "health_url": m.get("health_url", ""), - "status": status, - }) - if available: - backend["available_models"] = available + # The manifest is external input: a malformed file or entry must + # degrade to "no manifest" (logged), never crash detect_backends() + # and take the whole worker down with it. + try: + manifest = load_manifest() + if manifest.get("models"): + for backend in backends: + backend_type = backend["type"] + probed_names = { + m.get("name", "") for m in backend.get("models", []) + } + available = [] + for m in manifest["models"]: + if not isinstance(m, dict): + continue + if SOFTWARE_TO_BACKEND_TYPE.get(m.get("software", "")) != backend_type: + continue + model_id = m.get("model_id") + if not model_id: + logger.warning( + "skipping worker-manifest entry without model_id: %r", m + ) + continue + status = "loaded" if model_id in probed_names else "available" + available.append({ + "model_id": model_id, + "capability": m.get("capability", ""), + "software": m.get("software", ""), + "port": m.get("port", 0), + "vram_required_gb": m.get("vram_required_gb", 0.0), + "health_url": m.get("health_url", ""), + "status": status, + }) + if available: + backend["available_models"] = available + except Exception: # noqa: BLE001 - manifest must never brick the worker + logger.warning("worker-manifest enrichment failed; continuing without it", + exc_info=True) return backends @@ -567,7 +581,11 @@ async def heartbeat(self) -> int: headers=auth_headers, ) return resp.status_code - except Exception: + except Exception as exc: # noqa: BLE001 + # Log before swallowing: a payload-build bug (not just a network + # blip) would otherwise drop the worker offline with zero + # diagnostics, masquerading as controller unreachability. + logger.warning("heartbeat send failed: %s: %s", type(exc).__name__, exc) return 0 def _log_repair_instruction(self) -> None: diff --git a/tinyagentos/worker/worker_manifest.py b/tinyagentos/worker/worker_manifest.py index bb319b20..ea735289 100644 --- a/tinyagentos/worker/worker_manifest.py +++ b/tinyagentos/worker/worker_manifest.py @@ -33,10 +33,13 @@ from __future__ import annotations import json +import logging import os from pathlib import Path from typing import Any +logger = logging.getLogger(__name__) + DEFAULT_MANIFEST_PATH = "/etc/taos/worker-models.json" # Map software names to TAOS backend types so the worker can attach manifest @@ -58,10 +61,26 @@ def load_manifest(path: str | None = None) -> dict[str, Any]: Returns: Dict with keys ``resource_id`` (str) and ``models`` (list of dicts). - Returns an empty manifest when the file is absent (the worker may - not run with a manifest). + Returns an empty manifest when the file is absent OR unreadable/ + malformed. The manifest is external input (any platform may write + it), so a typo in it must never take the worker down -- a bad file + is logged and treated as no-manifest rather than raised. """ manifest_path = Path(path or os.getenv("TAOS_WORKER_MANIFEST", DEFAULT_MANIFEST_PATH)) + empty: dict[str, Any] = {"resource_id": "", "models": []} if not manifest_path.exists(): - return {"resource_id": "", "models": []} - return json.loads(manifest_path.read_text()) + return empty + try: + data = json.loads(manifest_path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + logger.warning( + "ignoring malformed worker manifest %s: %s", manifest_path, exc + ) + return empty + if not isinstance(data, dict) or not isinstance(data.get("models", []), list): + logger.warning( + "ignoring worker manifest %s: top level must be an object with a " + "'models' list", manifest_path, + ) + return empty + return data