diff --git a/cluster.py b/cluster.py index 1ef3df5..70cc438 100644 --- a/cluster.py +++ b/cluster.py @@ -67,6 +67,7 @@ async def handle_registration(exchange: str, routing_key: str, payload: dict) -> "ip": payload.get("ip"), "capabilities": payload.get("capabilities", []), "active_model": payload.get("active_model"), + "inventory": payload.get("inventory", []), "load": payload.get("load"), "registered_at": now, "last_seen": now, @@ -175,12 +176,68 @@ async def ping_node(node_name: str) -> bool: return False +async def request_model_swap(node_name: str, model_filename: str) -> bool: + """Request a worker node to swap its active model.""" + if node_name not in CLUSTER_NODES: + log.warning("request_model_swap: unknown node %s", node_name) + return False + + CLUSTER_NODES[node_name]["status"] = "swapping" + _push_event("cluster", "info", node_name, f"Model swap requested: {model_filename}") + + now = datetime.now(timezone.utc).isoformat() + "Z" + await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.cmd.swap_model", { + "model_filename": model_filename, + "requested_at": now, + }) + return True + + +async def handle_model_ready(exchange: str, routing_key: str, payload: dict) -> None: + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + model_filename = payload.get("active_model") + port = payload.get("port", 8081) + + if node_name not in CLUSTER_NODES: + log.warning("handle_model_ready: unknown node %s", node_name) + return + + inventory = CLUSTER_NODES[node_name].get("inventory") or [] + model_info = None + for inv in inventory: + if inv.get("filename") == model_filename: + model_info = {**inv, "port": port} + break + if model_info is None: + model_info = {"filename": model_filename, "port": port} + + now = datetime.now(timezone.utc).isoformat() + "Z" + CLUSTER_NODES[node_name]["active_model"] = model_info + CLUSTER_NODES[node_name]["status"] = "active" + CLUSTER_NODES[node_name]["last_seen"] = now + _push_event("cluster", "info", node_name, f"Model swap complete: {model_filename}") + + +async def handle_model_failed(exchange: str, routing_key: str, payload: dict) -> None: + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + error = payload.get("error", "unknown error") + + if node_name not in CLUSTER_NODES: + log.warning("handle_model_failed: unknown node %s", node_name) + return + + CLUSTER_NODES[node_name]["status"] = "error" + _push_event("cluster", "error", node_name, f"Model swap failed: {error}") + + SUBSCRIBE_TABLE = [ (AMQP_EXCHANGE_ADMIN, ["node.*.register"], handle_registration), (AMQP_EXCHANGE_ADMIN, ["node.*.deregister"], handle_deregistration), (AMQP_EXCHANGE_ADMIN, ["node.*.pong"], handle_pong), (AMQP_EXCHANGE_SYSTEM, ["node.*.event"], handle_event), (AMQP_EXCHANGE_SYSTEM, ["cluster.coordinator.query"], handle_coordinator_query), + (AMQP_EXCHANGE_SYSTEM, ["node.*.model_ready"], handle_model_ready), + (AMQP_EXCHANGE_SYSTEM, ["node.*.model_failed"], handle_model_failed), ] diff --git a/config.py b/config.py index 4e34d2d..7b1f6d7 100644 --- a/config.py +++ b/config.py @@ -9,7 +9,7 @@ import logging log = logging.getLogger("caic") -VERSION = "v0.15.0" +VERSION = "v0.16.0" OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434") LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081") SEARXNG_BASE = "http://localhost:8888" diff --git a/tests/test_model_swap.py b/tests/test_model_swap.py new file mode 100644 index 0000000..5a69689 --- /dev/null +++ b/tests/test_model_swap.py @@ -0,0 +1,202 @@ +"""Tests for cluster.py model swap flow.""" +import asyncio + +import cluster +import triage +from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM + + +def _reset(): + cluster.CLUSTER_NODES.clear() + cluster.CLUSTER_EVENTS.clear() + cluster.CLUSTER_COORDINATOR = None + cluster._pending_pings.clear() + + +_published = [] + + +async def _fake_publish(exchange, routing_key, payload): + _published.append((exchange, routing_key, payload)) + + +# ---------- 1. request_model_swap() publishes swap command ---------- + + +def test_request_model_swap_publishes_command(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "active", + } + + asyncio.run(cluster.request_model_swap("jarvis", "qwen2.5-coder-Q4_K_M.gguf")) + + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "swapping" + assert len(_published) == 1 + exchange, rk, payload = _published[0] + assert exchange == AMQP_EXCHANGE_ADMIN + assert rk == "node.jarvis.cmd.swap_model" + assert payload["model_filename"] == "qwen2.5-coder-Q4_K_M.gguf" + assert "requested_at" in payload + + +def test_request_model_swap_unknown_node(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + result = asyncio.run(cluster.request_model_swap("ghost", "any.gguf")) + assert result is False + assert len(_published) == 0 + + +# ---------- 2. handle_model_ready() updates node ---------- + + +def test_handle_model_ready_updates_active_model(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "swapping", + "active_model": {"name": "llama3.1", "port": 8081}, + "inventory": [ + {"filename": "qwen2.5-coder-Q4_K_M.gguf", "name": "qwen2.5-coder", "version": "14b", "quant": "Q4_K_M"}, + ], + } + + asyncio.run(cluster.handle_model_ready( + AMQP_EXCHANGE_SYSTEM, "node.jarvis.model_ready", + {"node_name": "jarvis", "active_model": "qwen2.5-coder-Q4_K_M.gguf", "port": 8082}, + )) + + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "active" + am = cluster.CLUSTER_NODES["jarvis"]["active_model"] + assert am["name"] == "qwen2.5-coder" + assert am["port"] == 8082 + assert am["filename"] == "qwen2.5-coder-Q4_K_M.gguf" + + +def test_handle_model_ready_inventory_lookup_fallback(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "swapping", + "active_model": {"name": "llama3.1", "port": 8081}, + "inventory": [], + } + + asyncio.run(cluster.handle_model_ready( + AMQP_EXCHANGE_SYSTEM, "node.jarvis.model_ready", + {"node_name": "jarvis", "active_model": "unknown.gguf", "port": 9999}, + )) + + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "active" + am = cluster.CLUSTER_NODES["jarvis"]["active_model"] + assert am["filename"] == "unknown.gguf" + assert am["port"] == 9999 + + +def test_handle_model_ready_unknown_node(caplog, monkeypatch): + _reset() + caplog.set_level("WARNING") + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_model_ready( + AMQP_EXCHANGE_SYSTEM, "node.ghost.model_ready", + {"node_name": "ghost", "active_model": "any.gguf"}, + )) + + assert any("unknown node" in rec.message for rec in caplog.records) + + +# ---------- 3. handle_model_failed() sets error status ---------- + + +def test_handle_model_failed_sets_error(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "swapping", + } + + asyncio.run(cluster.handle_model_failed( + AMQP_EXCHANGE_SYSTEM, "node.jarvis.model_failed", + {"node_name": "jarvis", "error": "llama-server unhealthy after 120s"}, + )) + + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "error" + assert len(cluster.CLUSTER_EVENTS) == 1 + assert "swap failed" in cluster.CLUSTER_EVENTS[0]["message"].lower() + + +def test_handle_model_failed_unknown_node(caplog, monkeypatch): + _reset() + caplog.set_level("WARNING") + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_model_failed( + AMQP_EXCHANGE_SYSTEM, "node.ghost.model_failed", + {"node_name": "ghost", "error": "OOM"}, + )) + + assert any("unknown node" in rec.message for rec in caplog.records) + + +# ---------- 4. select_node() triggers swap when model mismatched ---------- + + +def test_select_node_code_triggers_swap(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "active", + "ip": "192.168.50.210", + "active_model": {"name": "llama3.1", "port": 8081}, + "inventory": [ + {"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder", "version": "14b", "quant": "Q4_K_M"}, + ], + } + + result = asyncio.run(triage.select_node("code")) + + assert result is None + # Swap should have been published + assert any("cmd.swap_model" in rk for _, rk, _ in _published) + # Node should now be swapping + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "swapping" + + +# ---------- 5. select_node() returns None when node is already swapping ---------- + + +def test_select_node_swapping_returns_none(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "swapping", + "ip": "192.168.50.210", + "active_model": {"name": "llama3.1", "port": 8081}, + "inventory": [ + {"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder"}, + ], + } + + result = asyncio.run(triage.select_node("code")) + + assert result is None + # No swap command should be published while already swapping + swap_published = any("cmd.swap_model" in rk for _, rk, _ in _published) + assert not swap_published diff --git a/tests/test_triage.py b/tests/test_triage.py index 86ae99c..d5311bd 100644 --- a/tests/test_triage.py +++ b/tests/test_triage.py @@ -1,4 +1,5 @@ """Tests for triage.py — query classification and node selection.""" +import asyncio import json import httpx @@ -13,6 +14,13 @@ def _reset(): cluster.CLUSTER_COORDINATOR = None +_published = [] + + +async def _fake_publish(exchange, routing_key, payload): + _published.append((exchange, routing_key, payload)) + + class _MockPostResponse: def __init__(self, json_data: dict, status_code: int = 200): self._json_data = json_data @@ -83,7 +91,7 @@ def test_select_node_code_returns_coder(): "active_model": {"name": "llama3.1", "port": 8081}, } - node = triage.select_node("code") + node = asyncio.run(triage.select_node("code")) assert node is not None assert node["name"] == "coder01" @@ -97,7 +105,7 @@ def test_select_node_general_no_match_returns_none(): "name": "coder01", "type": "worker", "status": "active", "active_model": {"name": "qwen2.5-coder-14b", "port": 8082}, } - node = triage.select_node("general") + node = asyncio.run(triage.select_node("general")) assert node is None diff --git a/triage.py b/triage.py index 8dd2dcc..8552f97 100644 --- a/triage.py +++ b/triage.py @@ -7,6 +7,11 @@ from config import TRIAGE_BASE, TRIAGE_TIMEOUT, LLAMA_SERVER_BASE log = logging.getLogger("caic") +_IDEAL_MODEL_MAP = { + "code": {"name_contains": ["coder", "qwen"]}, + "general": {"name_contains": ["mistral", "llama"]}, +} + _CLASSIFICATION_PROMPT = """Classify the following user query into exactly one category. Respond with only the category name. Categories: @@ -45,22 +50,36 @@ async def classify_query(query: str) -> str: return "general" -def select_node(classification: str) -> dict | None: +async def select_node(classification: str) -> dict | None: from cluster import CLUSTER_NODES + if classification in ("search", "rag"): + return None + + ideal = _IDEAL_MODEL_MAP.get(classification, {}) + ideal_contains = ideal.get("name_contains", []) + + # First pass: find an active node with the right model already loaded for node in CLUSTER_NODES.values(): if node.get("status") != "active": continue am = node.get("active_model") or {} name = (am.get("name") or "").lower() - if classification == "code": - if "coder" in name or "qwen" in name: - return node - elif classification == "general": - if "mistral" in name or "llama" in name: - return node - else: - return None + if any(ideal in name for ideal in ideal_contains): + return node + + # Second pass: find an active node that can swap to the right model + for node in CLUSTER_NODES.values(): + if node.get("status") != "active": + continue + inventory = node.get("inventory") or [] + for inv in inventory: + inv_name = (inv.get("name") or "").lower() + if any(ideal in inv_name for ideal in ideal_contains): + from cluster import request_model_swap + await request_model_swap(node["name"], inv["filename"]) + return None + return None @@ -70,7 +89,7 @@ async def get_inference_url(query: str) -> str: classification = await classify_query(query) if classification in ("search", "rag"): return LLAMA_SERVER_BASE - node = select_node(classification) + node = await select_node(classification) if node: am = node.get("active_model") or {} port = am.get("port", 8081)