diff --git a/tests/test_cluster_worker_protocol.py b/tests/test_cluster_worker_protocol.py index 0fcf010a5..ce00238cb 100644 --- a/tests/test_cluster_worker_protocol.py +++ b/tests/test_cluster_worker_protocol.py @@ -120,3 +120,134 @@ async def test_heartbeat_updates_kv_quant(self): mgr.heartbeat("w1", kv_cache_quant_support=["fp16", "turboquant-k3v2"]) result = mgr.kv_quant_union() assert "turboquant-k3v2" in result + + +# --------------------------------------------------------------------------- +# WorkerInfo.available_models field — sidecar-manifest-derived field +# --------------------------------------------------------------------------- + +class TestWorkerInfoAvailableModels: + def test_default_is_empty_list(self): + w = WorkerInfo(name="w", url="http://localhost:9000") + assert w.available_models == [] + + def test_custom_value_stored(self): + models = [{"model_id": "qwen3.6-35b-a3b", "status": "available"}] + w = WorkerInfo(name="w", url="http://localhost:9000", available_models=models) + assert w.available_models == models + + def test_serialises_via_asdict(self): + models = [ + {"model_id": "qwen3.6-35b-a3b", "capability": "llm-chat", "status": "available"}, + {"model_id": "qwen3-embedding-8b", "capability": "embedding", "status": "loaded"}, + ] + w = WorkerInfo(name="w", url="http://localhost:9000", available_models=models) + d = asdict(w) + assert "available_models" in d + assert d["available_models"] == models + + def test_roundtrip_default(self): + w = WorkerInfo(name="w", url="http://localhost:9000") + d = asdict(w) + w2 = WorkerInfo(**{k: v for k, v in d.items()}) + assert w2.available_models == [] + + def test_roundtrip_custom(self): + models = [{"model_id": "test-model", "status": "available"}] + w = WorkerInfo(name="w", url="http://localhost:9000", available_models=models) + d = asdict(w) + w2 = WorkerInfo(**{k: v for k, v in d.items()}) + assert w2.available_models == models + + +# --------------------------------------------------------------------------- +# ClusterManager.heartbeat — available_models derivation from backends +# --------------------------------------------------------------------------- + +class TestHeartbeatAvailableModels: + async def _register(self, mgr: ClusterManager, name: str) -> None: + w = WorkerInfo(name=name, url=f"http://localhost:900{name[-1]}") + await mgr.register_worker(w) + + def test_heartbeat_derives_available_models_from_backends(self): + mgr = ClusterManager() + w = WorkerInfo(name="w1", url="http://localhost:9000") + mgr._workers["w1"] = w + + backends = [ + { + "name": "llama-cpp:9090", + "type": "llama-cpp", + "url": "http://localhost:9090", + "capabilities": ["llm-chat"], + "models": [{"name": "qwen35-9b"}], + "status": "ok", + "available_models": [ + {"model_id": "qwen3.6-35b-a3b", "capability": "llm-chat", "status": "available"}, + {"model_id": "harmonic-hermes-9b", "capability": "llm-chat", "status": "available"}, + ], + }, + ] + mgr.heartbeat("w1", backends=backends) + assert len(mgr._workers["w1"].available_models) == 2 + assert mgr._workers["w1"].available_models[0]["model_id"] == "qwen3.6-35b-a3b" + + def test_heartbeat_no_available_models_in_backends(self): + """Backends without available_models should leave the field empty.""" + mgr = ClusterManager() + w = WorkerInfo(name="w1", url="http://localhost:9000") + mgr._workers["w1"] = w + + backends = [ + { + "name": "ollama:11434", + "type": "ollama", + "url": "http://localhost:11434", + "models": [{"name": "llama3"}], + "status": "ok", + }, + ] + mgr.heartbeat("w1", backends=backends) + assert mgr._workers["w1"].available_models == [] + + def test_heartbeat_deduplicates_across_backends(self): + """Same model_id appearing in multiple backends is deduplicated.""" + mgr = ClusterManager() + w = WorkerInfo(name="w1", url="http://localhost:9000") + mgr._workers["w1"] = w + + backends = [ + { + "name": "llama-cpp:9090", + "type": "llama-cpp", + "url": "http://localhost:9090", + "available_models": [ + {"model_id": "shared-model", "status": "available"}, + ], + }, + { + "name": "llama-cpp:9091", + "type": "llama-cpp", + "url": "http://localhost:9091", + "available_models": [ + {"model_id": "shared-model", "status": "available"}, + {"model_id": "unique-model", "status": "available"}, + ], + }, + ] + mgr.heartbeat("w1", backends=backends) + # shared-model should appear only once + model_ids = [m["model_id"] for m in mgr._workers["w1"].available_models] + assert len(model_ids) == 2 + assert model_ids == ["shared-model", "unique-model"] + + def test_heartbeat_available_models_is_additive(self): + """When heartbeat omits backends, existing available_models is untouched.""" + mgr = ClusterManager() + models = [{"model_id": "existing", "status": "available"}] + w = WorkerInfo(name="w1", url="http://localhost:9000", available_models=models) + mgr._workers["w1"] = w + + # Heartbeat without backends — available_models unchanged + mgr.heartbeat("w1", load=0.5) + assert mgr._workers["w1"].available_models == models diff --git a/tests/test_skald_manifest.py b/tests/test_skald_manifest.py new file mode 100644 index 000000000..f6c26c9b7 --- /dev/null +++ b/tests/test_skald_manifest.py @@ -0,0 +1,173 @@ +"""Tests for tinyagentos.worker.skald_manifest — sidecar manifest parser.""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from unittest.mock import patch + +import pytest + +from tinyagentos.worker.skald_manifest import ( + DEFAULT_MANIFEST_PATH, + SOFTWARE_TO_BACKEND_TYPE, + load_manifest, +) + + +# --------------------------------------------------------------------------- +# SOFTWARE_TO_BACKEND_TYPE mapping +# --------------------------------------------------------------------------- + +class TestSoftwareToBackendType: + def test_llamacpp_maps_to_llama_cpp(self): + assert SOFTWARE_TO_BACKEND_TYPE["llamacpp"] == "llama-cpp" + + def test_embed_maps_to_llama_cpp(self): + assert SOFTWARE_TO_BACKEND_TYPE["embed"] == "llama-cpp" + + def test_kokoro_maps_to_kokoro(self): + assert SOFTWARE_TO_BACKEND_TYPE["kokoro"] == "kokoro" + + def test_whisper_maps_to_whisper(self): + assert SOFTWARE_TO_BACKEND_TYPE["whisper"] == "whisper" + + def test_unknown_software_raises_keyerror(self): + with pytest.raises(KeyError): + _ = SOFTWARE_TO_BACKEND_TYPE["nonexistent-software"] + + +# --------------------------------------------------------------------------- +# load_manifest — happy paths +# --------------------------------------------------------------------------- + +MANIFEST_FIXTURE = { + "resource_id": "gtx1080", + "models": [ + { + "model_id": "qwen3.6-35b-a3b", + "capability": "llm-chat", + "software": "llamacpp", + "port": 9090, + "vram_required_gb": 7.2, + "health_url": "http://127.0.0.1:9090/health", + }, + { + "model_id": "qwen3-embedding-8b", + "capability": "embedding", + "software": "embed", + "port": 8081, + "vram_required_gb": 4.5, + "health_url": "http://127.0.0.1:8081/health", + }, + { + "model_id": "whisper-large-v3", + "capability": "stt", + "software": "whisper", + "port": 0, + "vram_required_gb": 1.5, + "health_url": "", + }, + ], +} + + +class TestLoadManifestHappy: + def test_loads_valid_manifest_file(self, tmp_path: Path): + manifest_path = tmp_path / "skald-models.json" + manifest_path.write_text(json.dumps(MANIFEST_FIXTURE)) + + result = load_manifest(str(manifest_path)) + assert result["resource_id"] == "gtx1080" + assert len(result["models"]) == 3 + assert result["models"][0]["model_id"] == "qwen3.6-35b-a3b" + + def test_returns_model_fields(self, tmp_path: Path): + manifest_path = tmp_path / "skald-models.json" + manifest_path.write_text(json.dumps(MANIFEST_FIXTURE)) + + result = load_manifest(str(manifest_path)) + m0 = result["models"][0] + assert m0["capability"] == "llm-chat" + assert m0["software"] == "llamacpp" + assert m0["port"] == 9090 + assert m0["vram_required_gb"] == 7.2 + assert m0["health_url"] == "http://127.0.0.1:9090/health" + + def test_env_var_overrides_path(self, tmp_path: Path, monkeypatch): + manifest_path = tmp_path / "custom-manifest.json" + manifest_path.write_text(json.dumps(MANIFEST_FIXTURE)) + monkeypatch.setenv("TAOS_SKALD_MANIFEST", str(manifest_path)) + + result = load_manifest() + assert result["resource_id"] == "gtx1080" + + +# --------------------------------------------------------------------------- +# load_manifest — missing file +# --------------------------------------------------------------------------- + +class TestLoadManifestAbsent: + def test_missing_file_returns_empty_manifest(self, tmp_path: Path): + nonexistent = tmp_path / "does-not-exist.json" + result = load_manifest(str(nonexistent)) + assert result == {"resource_id": "", "models": []} + + def test_default_path_absent_returns_empty(self, monkeypatch): + # Ensure env var is unset and default path doesn't exist + monkeypatch.delenv("TAOS_SKALD_MANIFEST", raising=False) + # Patch Path.exists to return False for the default path + with patch.object(Path, "exists", return_value=False): + result = load_manifest() + assert result == {"resource_id": "", "models": []} + + +# --------------------------------------------------------------------------- +# load_manifest — edge cases +# --------------------------------------------------------------------------- + +class TestLoadManifestEdgeCases: + def test_empty_models_list(self, tmp_path: Path): + manifest = {"resource_id": "gtx1080", "models": []} + manifest_path = tmp_path / "empty.json" + manifest_path.write_text(json.dumps(manifest)) + + result = load_manifest(str(manifest_path)) + assert result["resource_id"] == "gtx1080" + assert result["models"] == [] + + def test_no_resource_id(self, tmp_path: Path): + manifest = {"models": [{"model_id": "test"}]} + manifest_path = tmp_path / "no-resource.json" + manifest_path.write_text(json.dumps(manifest)) + + result = load_manifest(str(manifest_path)) + # JSON parsing succeeds; missing keys aren't added + assert "resource_id" not in result # it won't be added by load_manifest + assert len(result["models"]) == 1 + + def test_invalid_json_raises(self, tmp_path: Path): + manifest_path = tmp_path / "bad.json" + manifest_path.write_text("not valid json {{{") + + with pytest.raises(json.JSONDecodeError): + load_manifest(str(manifest_path)) + + def test_none_path_falls_back_to_env(self, tmp_path: Path, monkeypatch): + manifest_path = tmp_path / "from-env.json" + manifest_path.write_text(json.dumps(MANIFEST_FIXTURE)) + monkeypatch.setenv("TAOS_SKALD_MANIFEST", str(manifest_path)) + + result = load_manifest(None) + assert result["resource_id"] == "gtx1080" + + def test_no_models_key_returns_empty_models(self, tmp_path: Path): + """A valid JSON doc without a 'models' key — treated as empty.""" + manifest = {"resource_id": "bare"} + manifest_path = tmp_path / "bare.json" + manifest_path.write_text(json.dumps(manifest)) + + result = load_manifest(str(manifest_path)) + assert result["resource_id"] == "bare" + # load_manifest doesn't add missing keys, consumer code checks .get("models") + assert "models" not in result diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index f7553fbdc..a459e26aa 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -193,6 +193,18 @@ def heartbeat( if name_m and name_m not in flat_models: flat_models.append(name_m) worker.models = flat_models + # Derive available_models from the live backend catalog. + # Each backend may carry a sidecar-manifest-enriched + # available_models list; flatten across all backends. + flat_available: list[dict] = [] + seen_ids: set[str] = set() + for b in backends: + for m in b.get("available_models") or []: + model_id = m.get("model_id") or "" + if model_id and model_id not in seen_ids: + seen_ids.add(model_id) + flat_available.append(m) + worker.available_models = flat_available if capabilities is not None: worker.capabilities = list(capabilities) if kv_cache_quant_support is not None: diff --git a/tinyagentos/cluster/worker_protocol.py b/tinyagentos/cluster/worker_protocol.py index 7563d2aae..b3994b3ab 100644 --- a/tinyagentos/cluster/worker_protocol.py +++ b/tinyagentos/cluster/worker_protocol.py @@ -14,6 +14,7 @@ class WorkerInfo: hardware: dict = field(default_factory=dict) # From hardware detection backends: list[dict] = field(default_factory=list) # Available inference backends models: list[str] = field(default_factory=list) # Currently loaded models + available_models: list[dict] = field(default_factory=list) # Models this worker CAN load (from Skald sidecar manifest) capabilities: list[str] = field(default_factory=list) # embed, chat, rerank, image-gen, tts, etc status: str = "online" # online | offline | busy last_heartbeat: float = 0 diff --git a/tinyagentos/worker/agent.py b/tinyagentos/worker/agent.py index 764ad483d..4ed9430c8 100644 --- a/tinyagentos/worker/agent.py +++ b/tinyagentos/worker/agent.py @@ -124,6 +124,42 @@ async def detect_backends(self) -> list[dict]: # its cluster-level kv_cache_quant_support advertisement. "kv_quant_support": kv_quant, }) + + # Enrich each backend with available models from the Skald sidecar + # manifest. The manifest declares which models this machine *can* + # run (per the GPU catalog), independently of what is currently + # loaded. The controller then sees both "loaded" and "available" + # states so the cluster-wide view reflects total capacity. + from tinyagentos.worker.skald_manifest import ( + load_manifest, + 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 + return backends async def _probe_kv_quant( diff --git a/tinyagentos/worker/skald_manifest.py b/tinyagentos/worker/skald_manifest.py new file mode 100644 index 000000000..12025acfd --- /dev/null +++ b/tinyagentos/worker/skald_manifest.py @@ -0,0 +1,49 @@ +"""Read the Skald sidecar manifest that declares which models this machine +*can* run, independently of whatever is currently loaded in RAM. + +The manifest is generated by Skald's ``export_taos_manifest.py`` and written +to ``/etc/taos/skald-models.json`` (configurable via ``TAOS_SKALD_MANIFEST``). +Each entry describes a model the GPU catalog knows about, including the +backend software, port, VRAM requirement, and health URL. + +TAOS uses this information to advertise *available* (not-yet-loaded) models +alongside the *loaded* models discovered by live backend probing. The +controller's cluster view then knows what a worker *could* load without +having to consult an external catalog. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +DEFAULT_MANIFEST_PATH = "/etc/taos/skald-models.json" + +# Map Skald software names to TAOS backend types so the worker can +# attach manifest entries to the correct backend dict. +SOFTWARE_TO_BACKEND_TYPE: dict[str, str] = { + "llamacpp": "llama-cpp", + "embed": "llama-cpp", + "kokoro": "kokoro", + "whisper": "whisper", +} + + +def load_manifest(path: str | None = None) -> dict[str, Any]: + """Read the Skald sidecar manifest from disk. + + Args: + path: Override path. Falls back to ``TAOS_SKALD_MANIFEST`` env var, + then ``/etc/taos/skald-models.json``. + + 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 alongside Skald). + """ + manifest_path = Path(path or os.getenv("TAOS_SKALD_MANIFEST", DEFAULT_MANIFEST_PATH)) + if not manifest_path.exists(): + return {"resource_id": "", "models": []} + return json.loads(manifest_path.read_text())