diff --git a/tests/test_cluster_worker_protocol.py b/tests/test_cluster_worker_protocol.py index 0fcf010a5..90ea1349c 100644 --- a/tests/test_cluster_worker_protocol.py +++ b/tests/test_cluster_worker_protocol.py @@ -120,3 +120,82 @@ 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 VRAM fields (taOS #894 slice) +# --------------------------------------------------------------------------- + +class TestWorkerInfoVramDefaults: + def test_default_vram_is_zero(self): + w = WorkerInfo(name="w", url="http://localhost:9000") + assert w.free_vram_mb == 0 + assert w.used_vram_mb == 0 + + def test_vram_stored_and_serialised(self): + w = WorkerInfo( + name="w", + url="http://localhost:9000", + free_vram_mb=8192, + used_vram_mb=4096, + ) + assert w.free_vram_mb == 8192 + assert w.used_vram_mb == 4096 + d = asdict(w) + assert d["free_vram_mb"] == 8192 + assert d["used_vram_mb"] == 4096 + + def test_vram_roundtrip(self): + w = WorkerInfo( + name="w", + url="http://localhost:9000", + free_vram_mb=12288, + used_vram_mb=2048, + ) + d = asdict(w) + w2 = WorkerInfo(**{k: v for k, v in d.items() if not isinstance(v, bytes)}) + assert w2.free_vram_mb == 12288 + assert w2.used_vram_mb == 2048 + + +# --------------------------------------------------------------------------- +# ClusterManager heartbeat VRAM storage (taOS #894 slice) +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +class TestHeartbeatVram: + async def _register(self, mgr: ClusterManager, name: str) -> WorkerInfo: + w = WorkerInfo(name=name, url=f"http://localhost:900{name[-1]}") + await mgr.register_worker(w) + return w + + async def test_heartbeat_stores_vram(self): + mgr = ClusterManager() + await self._register(mgr, "w1") + mgr.heartbeat("w1", free_vram_mb=8192, used_vram_mb=4096) + w = mgr.get_worker("w1") + assert w.free_vram_mb == 8192 + assert w.used_vram_mb == 4096 + + async def test_heartbeat_omitting_vram_preserves_prior(self): + """Legacy heartbeat without VRAM fields leaves stored values unchanged.""" + mgr = ClusterManager() + await self._register(mgr, "w1") + mgr.heartbeat("w1", free_vram_mb=12288, used_vram_mb=4096) + # Sparse heartbeat — only load + models, no VRAM fields + mgr.heartbeat("w1", load=0.5, models=["phi3"]) + w = mgr.get_worker("w1") + assert w.free_vram_mb == 12288 + assert w.used_vram_mb == 4096 + assert w.load == 0.5 + assert w.models == ["phi3"] + + async def test_vram_zero_is_stored(self): + """Explicit zero is stored (distinct from 'not sent' which preserves prior).""" + mgr = ClusterManager() + await self._register(mgr, "w1") + mgr.heartbeat("w1", free_vram_mb=8192, used_vram_mb=4096) + mgr.heartbeat("w1", free_vram_mb=0, used_vram_mb=0) + w = mgr.get_worker("w1") + assert w.free_vram_mb == 0 + assert w.used_vram_mb == 0 diff --git a/tests/test_model_archive.py b/tests/test_model_archive.py new file mode 100644 index 000000000..2a681eca3 --- /dev/null +++ b/tests/test_model_archive.py @@ -0,0 +1,383 @@ +"""Tests for the model archive promotion engine.""" +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import pytest + +from tinyagentos.cluster.model_archive import ( + _archive_root, + _active_models_root, + _worker_can_run, + find_promotable, + list_archived_models, + promote_model, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_worker_hw( + ram_mb: int = 8192, + gpu_type: str = "nvidia", + gpu_cuda: bool = True, + vram_mb: int = 8192, + arch: str = "x86_64", +) -> dict: + return { + "ram_mb": ram_mb, + "gpu": { + "type": gpu_type, + "cuda": gpu_cuda, + "vram_mb": vram_mb, + }, + "cpu": {"arch": arch}, + "npu": {"type": "none"}, + } + + +def _write_archive_manifest( + archive_dir: Path, + model_id: str, + requirements: dict | None = None, + files: list[str] | None = None, + backend: str = "llama-cpp", + family: str = "qwen3", +) -> dict: + """Write a manifest AND create a dummy model-files dir so + :func:`promote_model` has something to move. + """ + manifest = { + "model_id": model_id, + "backend": backend, + "family": family, + "files": files or [f"{model_id}-Q4_K_M.gguf"], + "requirements": requirements or {}, + "archived_at": 1700000000.0, + } + manifest_path = archive_dir / f"{model_id}.json" + manifest_path.write_text(json.dumps(manifest)) + # Create the accompanying model files directory + files_dir = archive_dir / model_id + files_dir.mkdir(parents=True, exist_ok=True) + (files_dir / f"{model_id}-Q4_K_M.gguf").write_text("fake-model-data") + return manifest + + +# --------------------------------------------------------------------------- +# _worker_can_run +# --------------------------------------------------------------------------- + + +class TestWorkerCanRun: + def test_empty_requirements(self): + assert _worker_can_run(_make_worker_hw(), {}) is True + + def test_vram_met(self): + assert _worker_can_run( + _make_worker_hw(vram_mb=8192), + {"min_vram_mb": 4096}, + ) is True + + def test_vram_not_met(self): + assert _worker_can_run( + _make_worker_hw(vram_mb=4096), + {"min_vram_mb": 8192}, + ) is False + + def test_ram_met(self): + assert _worker_can_run( + _make_worker_hw(ram_mb=16384), + {"min_ram_mb": 8192}, + ) is True + + def test_ram_not_met(self): + assert _worker_can_run( + _make_worker_hw(ram_mb=4096), + {"min_ram_mb": 8192}, + ) is False + + def test_gpu_type_nvidia_match(self): + assert _worker_can_run( + _make_worker_hw(gpu_type="nvidia"), + {"gpu_type": "nvidia"}, + ) is True + + def test_gpu_type_nvidia_mismatch(self): + assert _worker_can_run( + _make_worker_hw(gpu_type="amd"), + {"gpu_type": "nvidia"}, + ) is False + + def test_gpu_accel_cuda_match(self): + assert _worker_can_run( + _make_worker_hw(gpu_type="nvidia", gpu_cuda=True), + {"gpu_accel": "cuda"}, + ) is True + + def test_gpu_accel_cuda_mismatch(self): + assert _worker_can_run( + _make_worker_hw(gpu_type="nvidia", gpu_cuda=False), + {"gpu_accel": "cuda"}, + ) is False + + def test_arch_match(self): + assert _worker_can_run( + _make_worker_hw(arch="x86_64"), + {"arch": "x86_64"}, + ) is True + + def test_arch_mismatch(self): + assert _worker_can_run( + _make_worker_hw(arch="aarch64"), + {"arch": "x86_64"}, + ) is False + + def test_apple_silicon_unified_memory(self): + worker_hw = { + "ram_mb": 16384, + "gpu": {"type": "apple", "vulkan": True, "vram_mb": 16384}, + "cpu": {"arch": "aarch64"}, + "npu": {"type": "none"}, + } + assert _worker_can_run(worker_hw, {"min_vram_mb": 8192}) is True + + def test_all_requirements_met(self): + worker = _make_worker_hw(vram_mb=12288, ram_mb=16384, gpu_type="nvidia", gpu_cuda=True) + reqs = { + "min_vram_mb": 8192, + "min_ram_mb": 8192, + "gpu_type": "nvidia", + "gpu_accel": "cuda", + "arch": "x86_64", + } + assert _worker_can_run(worker, reqs) is True + + def test_one_requirement_fails_whole_thing_fails(self): + worker = _make_worker_hw(vram_mb=4096, ram_mb=16384, gpu_type="nvidia", gpu_cuda=True) + reqs = { + "min_vram_mb": 8192, + "min_ram_mb": 8192, + "gpu_type": "nvidia", + "gpu_accel": "cuda", + } + assert _worker_can_run(worker, reqs) is False + + +# --------------------------------------------------------------------------- +# list_archived_models +# --------------------------------------------------------------------------- + + +class TestListArchivedModels: + def test_empty_dir(self, tmp_path: Path): + assert list_archived_models(tmp_path) == [] + + def test_nonexistent_dir(self, tmp_path: Path): + assert list_archived_models(tmp_path / "nonexistent") == [] + + def test_single_manifest(self, tmp_path: Path): + _write_archive_manifest(tmp_path, "qwen3.5-4b") + result = list_archived_models(tmp_path) + assert len(result) == 1 + assert result[0]["model_id"] == "qwen3.5-4b" + assert "manifest_path" in result[0] + + def test_skips_corrupt_manifest(self, tmp_path: Path): + (tmp_path / "bad.json").write_text("not json") + _write_archive_manifest(tmp_path, "gemma-4-e2b") + result = list_archived_models(tmp_path) + assert len(result) == 1 + assert result[0]["model_id"] == "gemma-4-e2b" + + def test_multiple_manifests_sorted(self, tmp_path: Path): + _write_archive_manifest(tmp_path, "llama3-8b") + _write_archive_manifest(tmp_path, "gemma-4-e2b") + _write_archive_manifest(tmp_path, "qwen3.5-4b") + result = list_archived_models(tmp_path) + ids = [m["model_id"] for m in result] + # Sorted alphabetically by filename + assert ids == sorted(ids) + + +# --------------------------------------------------------------------------- +# find_promotable +# --------------------------------------------------------------------------- + + +class TestFindPromotable: + def test_no_archive(self, tmp_path: Path): + worker = _make_worker_hw() + assert find_promotable(worker, "test-worker", tmp_path) == [] + + def test_no_compatible(self, tmp_path: Path): + # Archive a model requiring 16GB VRAM; worker has 8GB + _write_archive_manifest( + tmp_path, "big-model", + requirements={"min_vram_mb": 16384}, + ) + worker = _make_worker_hw(vram_mb=8192) + assert find_promotable(worker, "test-worker", tmp_path) == [] + + def test_compatible_found(self, tmp_path: Path): + _write_archive_manifest( + tmp_path, "qwen3.5-4b", + requirements={"min_vram_mb": 4096, "gpu_accel": "cuda"}, + ) + worker = _make_worker_hw(vram_mb=8192) + result = find_promotable(worker, "test-worker", tmp_path) + assert len(result) == 1 + assert result[0]["model_id"] == "qwen3.5-4b" + assert result[0]["worker_name"] == "test-worker" + + def test_mixed_compatible_and_incompatible(self, tmp_path: Path): + _write_archive_manifest( + tmp_path, "qwen3.5-4b", + requirements={"min_vram_mb": 4096}, + ) + _write_archive_manifest( + tmp_path, "big-model", + requirements={"min_vram_mb": 32768}, + ) + worker = _make_worker_hw(vram_mb=8192) + result = find_promotable(worker, "test-worker", tmp_path) + assert len(result) == 1 + assert result[0]["model_id"] == "qwen3.5-4b" + + def test_no_requirements_always_promotable(self, tmp_path: Path): + _write_archive_manifest(tmp_path, "no-reqs-model", requirements={}) + worker = _make_worker_hw(vram_mb=128) # very weak + result = find_promotable(worker, "test-worker", tmp_path) + assert len(result) == 1 + + +# --------------------------------------------------------------------------- +# promote_model +# --------------------------------------------------------------------------- + + +class TestPromoteModel: + def test_promote_moves_files_and_removes_manifest(self, tmp_path: Path, monkeypatch): + archive_dir = tmp_path / "archive" + active_dir = tmp_path / "active" + archive_dir.mkdir(parents=True, exist_ok=True) + active_dir.mkdir(parents=True, exist_ok=True) + + # Override paths for test isolation + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._archive_root", + lambda: archive_dir, + ) + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._active_models_root", + lambda: active_dir, + ) + + _write_archive_manifest( + archive_dir, "qwen3.5-4b", + requirements={"min_vram_mb": 4096}, + backend="llama-cpp", + family="qwen3.5", + ) + models = list_archived_models(archive_dir) + assert len(models) == 1 + + ok = promote_model(models[0]) + assert ok is True + + # Manifest removed + assert not (archive_dir / "qwen3.5-4b.json").exists() + # Files dir moved + assert not (archive_dir / "qwen3.5-4b").is_dir() + target = active_dir / "llama-cpp" / "qwen3.5" / "qwen3.5-4b" + assert target.is_dir() + assert (target / "qwen3.5-4b-Q4_K_M.gguf").read_text() == "fake-model-data" + + def test_promote_skips_when_target_exists(self, tmp_path: Path, monkeypatch): + archive_dir = tmp_path / "archive" + active_dir = tmp_path / "active" + archive_dir.mkdir(parents=True, exist_ok=True) + active_dir.mkdir(parents=True, exist_ok=True) + + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._archive_root", + lambda: archive_dir, + ) + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._active_models_root", + lambda: active_dir, + ) + + _write_archive_manifest( + archive_dir, "qwen3.5-4b", + backend="llama-cpp", + family="qwen3.5", + ) + # Pre-create target dir + target_dir = active_dir / "llama-cpp" / "qwen3.5" / "qwen3.5-4b" + target_dir.mkdir(parents=True, exist_ok=True) + (target_dir / "existing.gguf").write_text("preexisting") + + models = list_archived_models(archive_dir) + ok = promote_model(models[0]) + assert ok is True + # Manifest removed, but existing target untouched + assert not (archive_dir / "qwen3.5-4b.json").exists() + assert (target_dir / "existing.gguf").read_text() == "preexisting" + + def test_promote_fails_without_files_dir(self, tmp_path: Path, monkeypatch): + archive_dir = tmp_path / "archive" + active_dir = tmp_path / "active" + archive_dir.mkdir(parents=True, exist_ok=True) + active_dir.mkdir(parents=True, exist_ok=True) + + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._archive_root", + lambda: archive_dir, + ) + monkeypatch.setattr( + "tinyagentos.cluster.model_archive._active_models_root", + lambda: active_dir, + ) + + # Create manifest without files dir + manifest = { + "model_id": "orphan-model", + "backend": "llama-cpp", + "family": "orphan", + "files": ["orphan.gguf"], + "requirements": {}, + "archived_at": 1700000000.0, + } + (archive_dir / "orphan-model.json").write_text(json.dumps(manifest)) + # No files dir — don't create it + + models = list_archived_models(archive_dir) + assert len(models) == 1 + ok = promote_model(models[0]) + assert ok is False + # Manifest remains + assert (archive_dir / "orphan-model.json").exists() + + +# --------------------------------------------------------------------------- +# Archive root env var override +# --------------------------------------------------------------------------- + + +class TestArchiveRootOverride: + def test_env_var_override(self, tmp_path: Path, monkeypatch): + custom = tmp_path / "custom-archive" + monkeypatch.setenv("TAOS_ARCHIVE_ROOT", str(custom)) + assert _archive_root() == custom + + def test_default_is_home_taos(self, monkeypatch): + monkeypatch.delenv("TAOS_ARCHIVE_ROOT", raising=False) + root = _archive_root() + assert root.name == "models" + assert "taos" in str(root) + assert "archive" in str(root) diff --git a/tinyagentos/cluster/manager.py b/tinyagentos/cluster/manager.py index e4fa93fb5..aa1ac5baf 100644 --- a/tinyagentos/cluster/manager.py +++ b/tinyagentos/cluster/manager.py @@ -93,6 +93,23 @@ async def register_worker(self, info: WorkerInfo) -> None: level="info", ) + # Promote any archived models this worker can now run + try: + from tinyagentos.cluster.model_archive import ( + promote_compatible_models, + ) + + await promote_compatible_models( + worker_hardware=info.hardware, + worker_name=info.name, + notifications=self._notifications, + ) + except Exception: + logger.exception( + "model_archive: promotion scan failed for worker '%s'", + info.name, + ) + def kv_quant_union(self) -> list[str]: """Return the set-union of KV cache quant types across all online workers. @@ -148,6 +165,8 @@ def heartbeat( kv_cache_quant_k_support: list[str] | None = None, kv_cache_quant_v_support: list[str] | None = None, kv_cache_quant_boundary_layer_protect: bool | None = None, + free_vram_mb: int | None = None, + used_vram_mb: int | None = None, ) -> bool: """Accept a worker heartbeat. @@ -186,6 +205,10 @@ def heartbeat( worker.kv_cache_quant_v_support = list(kv_cache_quant_v_support) if kv_cache_quant_boundary_layer_protect is not None: worker.kv_cache_quant_boundary_layer_protect = bool(kv_cache_quant_boundary_layer_protect) + if free_vram_mb is not None: + worker.free_vram_mb = int(free_vram_mb) + if used_vram_mb is not None: + worker.used_vram_mb = int(used_vram_mb) # Fire worker.online notification when a previously-offline worker recovers. # heartbeat() is sync, so schedule the async emit as a background task. if self._notifications and prev_status in ("offline", "stale"): diff --git a/tinyagentos/cluster/model_archive.py b/tinyagentos/cluster/model_archive.py new file mode 100644 index 000000000..545424f49 --- /dev/null +++ b/tinyagentos/cluster/model_archive.py @@ -0,0 +1,257 @@ +"""Archived model store and promotion engine. + +When a model is downloaded and no current worker can run it +(PR #325 force=True "Archive anyway"), it lands in + + ~/taos/archive/models/.json + files under + ~/taos/archive/models// + +This module scans that directory on worker join and promotes +any model that the new worker can now run — moving it into the +active models tree via :func:`~tinyagentos.installers.model_paths.models_root`. + +Consumed by: +- :meth:`~tinyagentos.cluster.manager.ClusterManager.register_worker` + (automatic promotion on worker join) +- ``GET /api/cluster/promote-archived`` (manual trigger) +""" +from __future__ import annotations + +import json +import logging +import os +import shutil +import time +from dataclasses import asdict +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + + +def _archive_root() -> Path: + """Root directory for archived models. Override with TAOS_ARCHIVE_ROOT env.""" + override = os.environ.get("TAOS_ARCHIVE_ROOT") + return Path(override) if override else Path.home() / "taos" / "archive" / "models" + + +def _active_models_root() -> Path: + """The active models tree — same as all backend installers write into.""" + from tinyagentos.installers.model_paths import models_root + return models_root() + + +def list_archived_models(archive_dir: Path | None = None) -> list[dict]: + """Scan the archive directory and return every archived model's manifest. + + Each manifest is a JSON file ``.json`` in the archive root. + Returns a list of dicts with keys: ``model_id``, ``files``, + ``requirements``, ``archived_at``, ``backend``, ``manifest_path``. + """ + root = archive_dir or _archive_root() + if not root.is_dir(): + return [] + models: list[dict] = [] + for manifest_path in sorted(root.glob("*.json")): + try: + data = json.loads(manifest_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + logger.warning("model_archive: skipping unreadable manifest %s", manifest_path) + continue + data["manifest_path"] = str(manifest_path) + models.append(data) + return models + + +def _worker_can_run(worker_hardware: dict, requirements: dict) -> bool: + """True if the worker's hardware meets the model's minimum requirements. + + Checks: VRAM, RAM, GPU type (cuda/rocm/vulkan/apple/npu), architecture. + + Requirements shape (all keys optional; missing means no constraint):: + + { + "min_vram_mb": 8192, + "min_ram_mb": 4096, + "gpu_type": "nvidia", # or "amd" / "apple" / "" + "gpu_accel": "cuda", # or "rocm" / "vulkan" / "mlx" / "" + "npu_type": "rknpu", # or "" + "arch": "x86_64" # or "aarch64" / "" + } + """ + if not requirements: + return True # No requirements = runs anywhere + + hw_gpu = worker_hardware.get("gpu") or {} + hw_npu = worker_hardware.get("npu") or {} + hw_cpu_raw = worker_hardware.get("cpu") or {} + hw_cpu: dict = hw_cpu_raw if isinstance(hw_cpu_raw, dict) else {} + hw_ram = worker_hardware.get("ram_mb", 0) + + # VRAM check + min_vram = requirements.get("min_vram_mb") + if min_vram: + worker_vram = hw_gpu.get("vram_mb", 0) or 0 + # Apple Silicon unified memory counts + if hw_gpu.get("type") == "apple": + worker_vram = max(worker_vram, hw_ram) + if worker_vram < min_vram: + return False + + # RAM check + min_ram = requirements.get("min_ram_mb") + if min_ram and hw_ram < min_ram: + return False + + # GPU type / accelerator check + req_gpu_type = requirements.get("gpu_type") + req_gpu_accel = requirements.get("gpu_accel") + worker_gpu_type = hw_gpu.get("type", "none") or "none" + + if req_gpu_type: + if req_gpu_type == "nvidia" and worker_gpu_type != "nvidia": + return False + if req_gpu_type == "amd" and worker_gpu_type != "amd": + return False + if req_gpu_type == "apple" and worker_gpu_type != "apple": + return False + + if req_gpu_accel: + if req_gpu_accel == "cuda" and not hw_gpu.get("cuda"): + return False + if req_gpu_accel == "rocm" and not hw_gpu.get("rocm"): + return False + if req_gpu_accel == "vulkan" and not hw_gpu.get("vulkan"): + return False + + # NPU check + req_npu = requirements.get("npu_type") + if req_npu: + worker_npu_type = hw_npu.get("type", "none") or "none" + if worker_npu_type != req_npu and worker_npu_type not in (req_npu,): + return False + + # Architecture check + req_arch = requirements.get("arch") + if req_arch: + worker_arch = hw_cpu.get("arch", "") + if worker_arch and worker_arch != req_arch: + return False + + return True + + +def find_promotable( + worker_hardware: dict, + worker_name: str, + archive_dir: Path | None = None, +) -> list[dict]: + """Return archived models that *this* worker can now run. + + Each entry is the model manifest dict with an extra ``worker_name`` key. + Does NOT move files — call :func:`promote_model` for each. + """ + promotable: list[dict] = [] + for model in list_archived_models(archive_dir): + reqs = model.get("requirements") or {} + if _worker_can_run(worker_hardware, reqs): + model["worker_name"] = worker_name + promotable.append(model) + return promotable + + +def promote_model(model: dict) -> bool: + """Move one archived model into the active models tree. + + Args: + model: A manifest dict from :func:`list_archived_models`. + + Returns True on success, False if the move fails (model stays archived). + """ + model_id = model.get("model_id", "") + manifest_path_str = model.get("manifest_path", "") + if not model_id or not manifest_path_str: + logger.warning("model_archive: cannot promote — missing model_id or manifest_path") + return False + + manifest_path = Path(manifest_path_str) + archive_root_path = manifest_path.parent + model_files_dir = archive_root_path / model_id + + # Resolve target directory in the active models tree. + # Use the backend from the manifest if present; otherwise guess. + backend = model.get("backend", "uncategorised") + # Build a target path: ~/models//// + # The family is derived from the model_id's first token or from the manifest. + family = model.get("family", model_id.split("-", 1)[0] if "-" in model_id else model_id) + target_dir = _active_models_root() / backend / family / model_id + + if not model_files_dir.is_dir(): + logger.warning( + "model_archive: model files directory %s not found — " + "promotion requires both the manifest and the files dir", + model_files_dir, + ) + return False + + try: + target_dir.parent.mkdir(parents=True, exist_ok=True) + # If target already exists, don't overwrite — but still remove + # the archive manifest so we don't keep trying. + if target_dir.exists(): + logger.info( + "model_archive: target %s already exists; removing archive entry, skipping move", + target_dir, + ) + manifest_path.unlink(missing_ok=True) + # Clean up empty model files dir if possible + if model_files_dir.is_dir(): + try: + model_files_dir.rmdir() + except OSError: + pass + return True + + shutil.move(str(model_files_dir), str(target_dir)) + # Remove the archive manifest after successful move + manifest_path.unlink(missing_ok=True) + logger.info("model_archive: promoted %s -> %s", model_id, target_dir) + return True + except (OSError, shutil.Error) as exc: + logger.error("model_archive: failed to promote %s: %s", model_id, exc) + return False + + +async def promote_compatible_models( + worker_hardware: dict, + worker_name: str, + archive_dir: Path | None = None, + notifications=None, +) -> list[str]: + """Scan archive, promote every model compatible with this worker. + + Called by :meth:`ClusterManager.register_worker` when a new worker joins. + Sends a notification for each promoted model. + + Returns the list of promoted model IDs. + """ + promotable = find_promotable(worker_hardware, worker_name, archive_dir) + promoted: list[str] = [] + for model in promotable: + model_id = model.get("model_id", "?") + if promote_model(model): + promoted.append(model_id) + if notifications: + await notifications.emit_event( + "model.promoted", + f"Archived model '{model_id}' promoted", + f"Worker '{worker_name}' can now run '{model_id}'. " + f"Moved from archive to active models.", + level="info", + ) + if promoted: + logger.info( + "model_archive: worker '%s' promoted %d model(s): %s", + worker_name, len(promoted), ", ".join(promoted), + ) + return promoted diff --git a/tinyagentos/cluster/worker_capacity.py b/tinyagentos/cluster/worker_capacity.py index e226956eb..21ac2ee74 100644 --- a/tinyagentos/cluster/worker_capacity.py +++ b/tinyagentos/cluster/worker_capacity.py @@ -96,3 +96,24 @@ def capacity_snapshot( "storage_used_bytes": used, "bytes_deduped_total": read_bees_deduped_total(bees_status_path), } + + +def gpu_vram_snapshot() -> dict | None: + """Return real-time free/used VRAM for the first GPU, or None. + + Dispatches to the appropriate probe based on what's available on + this worker (nvidia-smi, etc). Returns a dict with ``free_vram_mb`` + and ``used_vram_mb`` integers, or ``None`` when no GPU probe is + available or the probe fails. + + Called from the worker heartbeat loop alongside capacity_snapshot(). + Best-effort: a missing nvidia-smi or a probe timeout is not an error. + """ + from tinyagentos.system_stats import read_nvidia_vram + + pair = read_nvidia_vram() + if pair is None: + return None + used_mb, total_mb = pair + free_mb = max(0, total_mb - used_mb) + return {"free_vram_mb": free_mb, "used_vram_mb": used_mb} diff --git a/tinyagentos/cluster/worker_protocol.py b/tinyagentos/cluster/worker_protocol.py index 7563d2aae..4380f6cbe 100644 --- a/tinyagentos/cluster/worker_protocol.py +++ b/tinyagentos/cluster/worker_protocol.py @@ -58,3 +58,9 @@ class WorkerInfo: worker_lxc_image_version: str | None = None # Ubuntu image version, e.g. "ubuntu/24.04/amd64" degraded: bool = False degraded_reason: str | None = None + # Real-time GPU VRAM reported on every heartbeat (free/used in MiB). + # Default 0 means unreported / no GPU; consumers (Skald's TaosDispatcher) + # use these to sort candidates by actual availability instead of total + # capacity. Worker reports them; controller stores them; API surfaces them. + free_vram_mb: int = 0 + used_vram_mb: int = 0 diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index ac40f93e6..22d6b9ea4 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -66,6 +66,12 @@ class HeartbeatBody(BaseModel): storage_cap_bytes: int | None = None storage_used_bytes: int | None = None bytes_deduped_total: int | None = None + # Real-time GPU VRAM reported by the worker on every heartbeat. + # Optional so legacy workers that don't send them leave stored values + # unchanged. Consumers (e.g. Skald's TaosDispatcher) use these to + # rank candidates by actual free memory, not total capacity. + free_vram_mb: int | None = None + used_vram_mb: int | None = None class RouteRequest(BaseModel): @@ -232,6 +238,8 @@ async def worker_heartbeat(request: Request, body: HeartbeatBody): kv_cache_quant_k_support=body.kv_cache_quant_k_support, kv_cache_quant_v_support=body.kv_cache_quant_v_support, kv_cache_quant_boundary_layer_protect=body.kv_cache_quant_boundary_layer_protect, + free_vram_mb=body.free_vram_mb, + used_vram_mb=body.used_vram_mb, ) if not ok: return JSONResponse({"error": "Worker not registered"}, status_code=404) @@ -656,3 +664,57 @@ async def worker_remote_command(request: Request, name: str, body: WorkerRemoteR return resp.json() except Exception as exc: return JSONResponse({"error": str(exc)}, status_code=502) + + +@router.get("/api/cluster/promote-archived") +async def promote_archived_models(request: Request): + """Manual trigger: scan all online workers and promote any archived + models that are now compatible with cluster hardware. + + Called by the user from the Cluster page or admin CLI. Safe to call + repeatedly — already-promoted models are skipped. + """ + cluster = request.app.state.cluster_manager + notifications = getattr(request.app.state, "notifications", None) + + workers = cluster.get_workers() + online = [w for w in workers if w.status == "online"] + + from tinyagentos.cluster.model_archive import ( + find_promotable, + promote_model, + ) + + promoted_by_worker: dict[str, list[str]] = {} + total = 0 + + for w in online: + promotable = find_promotable( + worker_hardware=w.hardware, + worker_name=w.name, + ) + for model in promotable: + model_id = model.get("model_id", "?") + if promote_model(model): + promoted_by_worker.setdefault(w.name, []).append(model_id) + total += 1 + if notifications: + try: + await notifications.emit_event( + "model.promoted", + f"Archived model '{model_id}' promoted", + f"Worker '{w.name}' can now run '{model_id}'. " + f"Moved from archive to active models.", + level="info", + ) + except Exception: + logger.exception( + "notification emit failed for model promotion %s", + model_id, + ) + + return { + "promoted": total, + "by_worker": promoted_by_worker, + "workers_scanned": len(online), + } diff --git a/tinyagentos/worker/agent.py b/tinyagentos/worker/agent.py index 764ad483d..dffdfe47f 100644 --- a/tinyagentos/worker/agent.py +++ b/tinyagentos/worker/agent.py @@ -393,7 +393,7 @@ async def heartbeat(self) -> int: the 404 case (controller restarted and forgot about us) and trigger a re-registration. """ - from tinyagentos.cluster.worker_capacity import capacity_snapshot + from tinyagentos.cluster.worker_capacity import capacity_snapshot, gpu_vram_snapshot try: load = psutil.cpu_percent() / 100.0 @@ -401,6 +401,7 @@ async def heartbeat(self) -> int: caps = self.detect_capabilities(backends) kv_quant = self.detect_kv_quant_support(backends) snap = capacity_snapshot() + vram = gpu_vram_snapshot() async with httpx.AsyncClient(timeout=5) as client: resp = await client.post( f"{self.controller_url}/api/cluster/heartbeat", @@ -416,6 +417,8 @@ async def heartbeat(self) -> int: "storage_cap_bytes": snap["storage_cap_bytes"], "storage_used_bytes": snap["storage_used_bytes"], "bytes_deduped_total": snap["bytes_deduped_total"], + "free_vram_mb": vram["free_vram_mb"] if vram else 0, + "used_vram_mb": vram["used_vram_mb"] if vram else 0, }, ) return resp.status_code