Task 14: coordinator-side model swap flow

- request_model_swap() publishes cmd.swap_model on jc.admin, sets node status to swapping
- handle_model_ready() updates active_model from inventory, restores active status
- handle_model_failed() sets node status to error with failure detail
- select_node() made async with two-pass logic: find ready node or trigger swap
- inventory field stored in node records for swap target lookup
- Both handlers subscribed on jc.system via SUBSCRIBE_TABLE

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),
 ]
 
 

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"

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

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
 
 

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:
+        if any(ideal in name for ideal in ideal_contains):
             return node
-        elif classification == "general":
-            if "mistral" in name or "llama" in name:
-                return node
-        else:
+
+    # 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)