Task 11: cluster protocol — subscribe(), ping/pong health, 9 message types

- amqp.py: subscribe() creates exclusive queues bound to routing keys,
  _rebind_subscriptions() recreates them after reconnect
- cluster.py: CLUSTER_NODES, CLUSTER_EVENTS (bounded 1000),
  CLUSTER_COORDINATOR with auto-promotion, all 6 handlers
  (register, deregister, pong, event, coordinator_query),
  ping_node() with 5s timeout + auto-deregister on failure
- routers/cluster.py: GET /api/cluster
- Wired into app.py lifespan after amqp_connect()
- 13 tests covering all handlers, boundaries, and API shape
- No passive heartbeats — workers assumed present until ping
  timeout at work-routing time
This commit is contained in:
gramps
2026-07-06 19:22:31 -07:00
parent 454fb3a380
commit 78f7b79494
5 changed files with 575 additions and 1 deletions
+46
View File
@@ -26,6 +26,51 @@ log = logging.getLogger("jarvischat")
_connection = None _connection = None
_channel = None _channel = None
_lock = asyncio.Lock() _lock = asyncio.Lock()
_subscriptions = [] # list of (exchange, routing_keys, handler_fn)
async def subscribe(exchange: str, routing_keys: list[str], handler) -> None:
"""Subscribe to routing keys on an exchange.
Creates an exclusive anonymous queue bound to the given routing keys.
Handler receives (exchange, routing_key, payload_dict).
On reconnect: _rebind_subscriptions recreates all subscriptions.
"""
if not HAS_AIO_PIKA:
log.warning("aio-pika not installed — cannot subscribe")
return
ch = await get_channel()
if ch is None:
log.error("cannot subscribe — no AMQP channel")
return
try:
queue = await ch.declare_queue("", exclusive=True)
ex = await ch.get_exchange(exchange)
for rk in routing_keys:
await queue.bind(ex, rk)
async def _dispatch(msg: aio_pika.IncomingMessage):
async with msg.process():
try:
payload = json.loads(msg.body.decode())
await handler(exchange, msg.routing_key, payload)
except Exception:
log.exception("AMQP handler error for %s %s", exchange, msg.routing_key)
await queue.consume(_dispatch)
_subscriptions.append((exchange, routing_keys, handler))
except Exception:
log.exception("AMQP subscribe failed for %s %s", exchange, routing_keys)
async def _rebind_subscriptions() -> None:
"""Re-create all subscriptions after reconnect."""
subs = list(_subscriptions)
_subscriptions.clear()
for exchange, routing_keys, handler in subs:
await subscribe(exchange, routing_keys, handler)
if subs:
log.info("AMQP subscriptions rebound")
async def connect() -> None: async def connect() -> None:
@@ -48,6 +93,7 @@ async def connect() -> None:
await ch.declare_exchange(ex, ExchangeType.TOPIC, durable=True) await ch.declare_exchange(ex, ExchangeType.TOPIC, durable=True)
_connection = conn _connection = conn
_channel = ch _channel = ch
await _rebind_subscriptions()
log.info("AMQP connected, exchanges declared") log.info("AMQP connected, exchanges declared")
+4 -1
View File
@@ -15,6 +15,7 @@ from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates from fastapi.templating import Jinja2Templates
from amqp import connect as amqp_connect, disconnect as amqp_disconnect from amqp import connect as amqp_connect, disconnect as amqp_disconnect
from cluster import start_cluster_subscriptions
from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR, RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, RAG_EVICTION_BATCH from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR, RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, RAG_EVICTION_BATCH
from db import init_db from db import init_db
from hardware import assess_hardware from hardware import assess_hardware
@@ -39,6 +40,7 @@ import routers.upload as upload
import routers.ingest as ingest import routers.ingest as ingest
import routers.hardware as hardware import routers.hardware as hardware
import routers.rag_admin as rag_admin import routers.rag_admin as rag_admin
import routers.cluster as cluster_router
# --- Logging --- # --- Logging ---
log = logging.getLogger("jarvischat") log = logging.getLogger("jarvischat")
@@ -59,6 +61,7 @@ async def lifespan(app: FastAPI):
log.info(f"Memory system: {get_memory_count()} memories loaded") log.info(f"Memory system: {get_memory_count()} memories loaded")
await assess_hardware() await assess_hardware()
await amqp_connect() await amqp_connect()
await start_cluster_subscriptions()
if RAG_MAX_VECTORS > 0: if RAG_MAX_VECTORS > 0:
if RAG_EVICTION_HIGH_WATER <= RAG_EVICTION_LOW_WATER: if RAG_EVICTION_HIGH_WATER <= RAG_EVICTION_LOW_WATER:
@@ -162,7 +165,7 @@ for router_module in [
auth_router, conversations.router, memories.router, models.router, auth_router, conversations.router, memories.router, models.router,
presets.router, profile.router, settings.router, skills.router, presets.router, profile.router, settings.router, skills.router,
chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router, chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router,
rag_admin.router, rag_admin.router, cluster_router.router,
]: ]:
app.include_router(router_module) app.include_router(router_module)
+189
View File
@@ -0,0 +1,189 @@
"""
JarvisChat — Cluster protocol implementation.
Maintains node registry, event log, coordinator state, and ping-based health checks.
"""
import asyncio
import logging
import uuid
from collections import deque
from datetime import datetime, timezone
from amqp import publish, subscribe
from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
log = logging.getLogger("jarvischat")
CLUSTER_NODES: dict[str, dict] = {}
CLUSTER_EVENTS: deque = deque(maxlen=1000)
CLUSTER_COORDINATOR: str | None = None
_pending_pings: dict[str, tuple[str, asyncio.Event]] = {}
NODE_NAME: str = "ultron"
PING_TIMEOUT: float = 5.0
def _push_event(category: str, severity: str, node_name: str | None, message: str, details: dict | None = None) -> dict:
record = {
"category": category,
"severity": severity,
"node": node_name,
"message": message,
"details": details or {},
"timestamp": datetime.now(timezone.utc).isoformat() + "Z",
}
CLUSTER_EVENTS.append(record)
level = {"info": logging.INFO, "warn": logging.WARNING, "error": logging.ERROR, "critical": logging.CRITICAL}.get(severity, logging.INFO)
log.log(level, "[cluster] %s: %s", node_name or "system", message)
return record
async def handle_registration(exchange: str, routing_key: str, payload: dict) -> None:
global CLUSTER_COORDINATOR
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
node_type = payload.get("node_type", "worker")
if node_name in CLUSTER_NODES:
_push_event("cluster", "warn", node_name, "Duplicate registration rejected")
await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.rejected", {
"from": NODE_NAME, "node_name": node_name, "type": "rejected",
"reason": "duplicate_node_name",
"timestamp": datetime.now(timezone.utc).isoformat() + "Z",
})
return
if "node_type" not in payload or "capabilities" not in payload:
_push_event("cluster", "warn", node_name, "Malformed registration rejected")
await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.rejected", {
"from": NODE_NAME, "node_name": node_name, "type": "rejected",
"reason": "malformed_payload",
"timestamp": datetime.now(timezone.utc).isoformat() + "Z",
})
return
now = datetime.now(timezone.utc).isoformat() + "Z"
node = {
"name": node_name,
"type": node_type,
"status": "active",
"capabilities": payload.get("capabilities", []),
"active_model": payload.get("active_model"),
"load": payload.get("load"),
"registered_at": now,
"last_seen": now,
}
CLUSTER_NODES[node_name] = node
_push_event("cluster", "info", node_name, f"Node registered (type={node_type})")
await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.admitted", {
"from": NODE_NAME, "node_name": node_name, "type": "admitted",
"coordinator": CLUSTER_COORDINATOR,
"timestamp": now,
})
if node_type == "coordinator" and CLUSTER_COORDINATOR is None:
CLUSTER_COORDINATOR = node_name
_push_event("cluster", "info", node_name, "Elected as coordinator")
await publish(AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.response", {
"from": NODE_NAME, "type": "coord_response",
"coordinator": node_name, "nodes": list(CLUSTER_NODES.keys()),
"timestamp": now,
})
async def handle_deregistration(exchange: str, routing_key: str, payload: dict) -> None:
global CLUSTER_COORDINATOR
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
removed = CLUSTER_NODES.pop(node_name, None)
if not removed:
log.warning("[cluster] deregistration for unknown node %s", node_name)
return
_push_event("cluster", "info", node_name, "Node deregistered")
if CLUSTER_COORDINATOR == node_name:
CLUSTER_COORDINATOR = None
_push_event("cluster", "warn", node_name, "Coordinator deregistered — no coordinator active")
async def handle_pong(exchange: str, routing_key: str, payload: dict) -> None:
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
correlation_id = payload.get("correlation_id")
if correlation_id and correlation_id in _pending_pings:
_, event = _pending_pings.pop(correlation_id)
event.set()
if node_name in CLUSTER_NODES:
now = datetime.now(timezone.utc).isoformat() + "Z"
CLUSTER_NODES[node_name]["last_seen"] = now
if "status" in payload:
CLUSTER_NODES[node_name]["status"] = payload["status"]
if "active_model" in payload:
CLUSTER_NODES[node_name]["active_model"] = payload["active_model"]
if "load" in payload:
CLUSTER_NODES[node_name]["load"] = payload["load"]
else:
log.warning("[cluster] pong from unknown node %s", node_name)
async def handle_event(exchange: str, routing_key: str, payload: dict) -> None:
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
severity = payload.get("severity", "info")
message = payload.get("message", "")
details = payload.get("details")
_push_event("application", severity, node_name, message, details)
async def handle_coordinator_query(exchange: str, routing_key: str, payload: dict) -> None:
if CLUSTER_COORDINATOR is None:
return
now = datetime.now(timezone.utc).isoformat() + "Z"
await publish(AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.response", {
"from": NODE_NAME, "type": "coord_response",
"coordinator": CLUSTER_COORDINATOR,
"nodes": list(CLUSTER_NODES.keys()),
"timestamp": now,
})
async def ping_node(node_name: str) -> bool:
global CLUSTER_COORDINATOR
if node_name not in CLUSTER_NODES:
return False
correlation_id = str(uuid.uuid4())
event = asyncio.Event()
_pending_pings[correlation_id] = (node_name, event)
now = datetime.now(timezone.utc).isoformat() + "Z"
await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.ping", {
"from": NODE_NAME, "node_name": node_name, "type": "ping",
"correlation_id": correlation_id, "timestamp": now,
})
try:
await asyncio.wait_for(event.wait(), timeout=PING_TIMEOUT)
return True
except asyncio.TimeoutError:
_pending_pings.pop(correlation_id, None)
CLUSTER_NODES.pop(node_name, None)
_push_event("cluster", "warn", node_name, "Node unresponsive — deregistered after ping timeout")
if CLUSTER_COORDINATOR == node_name:
CLUSTER_COORDINATOR = None
_push_event("cluster", "warn", node_name, "Coordinator unresponsive — no coordinator active")
return False
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),
]
async def start_cluster_subscriptions() -> None:
for exchange, routing_keys, handler in SUBSCRIBE_TABLE:
await subscribe(exchange, routing_keys, handler)
log.info("cluster subscriptions started")
+23
View File
@@ -0,0 +1,23 @@
"""JarvisChat routers - Cluster status API."""
from fastapi import APIRouter
import cluster
router = APIRouter()
@router.get("/api/cluster")
async def cluster_status():
return {
"nodes": {name: _strip_internal(node) for name, node in cluster.CLUSTER_NODES.items()},
"node_count": len(cluster.CLUSTER_NODES),
"coordinator": cluster.CLUSTER_COORDINATOR,
"events": list(cluster.CLUSTER_EVENTS),
}
def _strip_internal(node: dict) -> dict:
return {k: v for k, v in node.items() if k in {
"name", "type", "status", "capabilities", "active_model", "load",
"registered_at", "last_seen",
}}
+313
View File
@@ -0,0 +1,313 @@
"""Tests for cluster.py — no live AMQP, all handlers called directly."""
import asyncio
from collections import deque
import cluster
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. Valid worker registration ----------
def test_valid_worker_registration(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
))
assert "jarvis" in cluster.CLUSTER_NODES
assert cluster.CLUSTER_NODES["jarvis"]["type"] == "worker"
assert cluster.CLUSTER_NODES["jarvis"]["status"] == "active"
assert len(cluster.CLUSTER_EVENTS) == 1
assert cluster.CLUSTER_EVENTS[0]["category"] == "cluster"
assert cluster.CLUSTER_EVENTS[0]["message"] == "Node registered (type=worker)"
assert len(_published) == 1
exchange, rk, payload = _published[0]
assert exchange == AMQP_EXCHANGE_ADMIN
assert rk == "node.jarvis.admitted"
assert payload["type"] == "admitted"
assert payload["node_name"] == "jarvis"
# ---------- 2. First coordinator auto-promotion ----------
def test_first_coordinator_auto_promotion(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.ultron.register",
{"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm", "rag"]},
))
assert cluster.CLUSTER_COORDINATOR == "ultron"
assert len(cluster.CLUSTER_EVENTS) == 2
assert cluster.CLUSTER_EVENTS[1]["message"] == "Elected as coordinator"
assert len(_published) == 2
# First publish: admitted
assert _published[0][1] == "node.ultron.admitted"
# Second publish: coord_response
exchange, rk, payload = _published[1]
assert exchange == AMQP_EXCHANGE_SYSTEM
assert rk == "cluster.coordinator.response"
assert payload["type"] == "coord_response"
assert payload["coordinator"] == "ultron"
# ---------- 3. Duplicate node name rejected ----------
def test_duplicate_node_rejected(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
))
_published.clear()
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
))
assert len(cluster.CLUSTER_NODES) == 1
assert len(_published) == 1
exchange, rk, payload = _published[0]
assert exchange == AMQP_EXCHANGE_ADMIN
assert rk == "node.jarvis.rejected"
assert payload["reason"] == "duplicate_node_name"
# ---------- 4. Malformed payload rejected ----------
def test_malformed_payload_rejected(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis"}, # missing node_type and capabilities
))
assert len(cluster.CLUSTER_NODES) == 0
assert len(_published) == 1
exchange, rk, payload = _published[0]
assert rk == "node.jarvis.rejected"
assert payload["reason"] == "malformed_payload"
# ---------- 5. Graceful deregistration ----------
def test_graceful_deregistration(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
))
_published.clear()
asyncio.run(cluster.handle_deregistration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.deregister",
{"node_name": "jarvis"},
))
assert "jarvis" not in cluster.CLUSTER_NODES
assert len(cluster.CLUSTER_EVENTS) == 2
assert cluster.CLUSTER_EVENTS[1]["message"] == "Node deregistered"
def test_deregister_coordinator_clears(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.ultron.register",
{"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]},
))
assert cluster.CLUSTER_COORDINATOR == "ultron"
_published.clear()
asyncio.run(cluster.handle_deregistration(
AMQP_EXCHANGE_ADMIN, "node.ultron.deregister",
{"node_name": "ultron"},
))
assert "ultron" not in cluster.CLUSTER_NODES
assert cluster.CLUSTER_COORDINATOR is None
# ---------- 6. Pong from known node ----------
def test_pong_updates_known_node(monkeypatch):
_reset()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
{"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
))
original_seen = cluster.CLUSTER_NODES["jarvis"]["last_seen"]
asyncio.run(cluster.handle_pong(
AMQP_EXCHANGE_ADMIN, "node.jarvis.pong",
{"node_name": "jarvis", "status": "busy", "load": {"cpu_pct": 80}},
))
assert cluster.CLUSTER_NODES["jarvis"]["status"] == "busy"
assert cluster.CLUSTER_NODES["jarvis"]["load"] == {"cpu_pct": 80}
assert cluster.CLUSTER_NODES["jarvis"]["last_seen"] != original_seen
# ---------- 7. Pong from unknown node ----------
def test_pong_from_unknown_node(caplog, monkeypatch):
_reset()
caplog.set_level("WARNING")
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_pong(
AMQP_EXCHANGE_ADMIN, "node.stranger.pong",
{"node_name": "stranger", "status": "active"},
))
assert "stranger" not in cluster.CLUSTER_NODES
assert any("pong from unknown node" in rec.message for rec in caplog.records)
# ---------- 8. Event stored in log ----------
def test_event_appended_to_log(monkeypatch):
_reset()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_event(
AMQP_EXCHANGE_SYSTEM, "node.jarvis.event",
{"node_name": "jarvis", "severity": "error", "message": "OOM on GPU"},
))
assert len(cluster.CLUSTER_EVENTS) == 1
ev = cluster.CLUSTER_EVENTS[0]
assert ev["category"] == "application"
assert ev["severity"] == "error"
assert ev["node"] == "jarvis"
def test_event_log_bounded(monkeypatch):
_reset()
monkeypatch.setattr(cluster, "publish", _fake_publish)
for i in range(1001):
cluster._push_event("application", "info", "jarvis", f"event {i}")
assert len(cluster.CLUSTER_EVENTS) == 1000
assert cluster.CLUSTER_EVENTS[0]["message"] == "event 1"
assert cluster.CLUSTER_EVENTS[-1]["message"] == "event 1000"
# ---------- 9. Coordinator query produces response ----------
def test_coordinator_query_response(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_registration(
AMQP_EXCHANGE_ADMIN, "node.ultron.register",
{"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]},
))
_published.clear()
asyncio.run(cluster.handle_coordinator_query(
AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query",
{"from": "jarvis", "type": "coord_query"},
))
assert len(_published) == 1
exchange, rk, payload = _published[0]
assert exchange == AMQP_EXCHANGE_SYSTEM
assert rk == "cluster.coordinator.response"
assert payload["coordinator"] == "ultron"
assert "nodes" in payload
def test_coordinator_query_no_coordinator(monkeypatch):
_reset()
_published.clear()
monkeypatch.setattr(cluster, "publish", _fake_publish)
asyncio.run(cluster.handle_coordinator_query(
AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query",
{"from": "jarvis", "type": "coord_query"},
))
assert len(_published) == 0
# ---------- 10. GET /api/cluster shape ----------
def test_cluster_api_shape():
_reset()
cluster.CLUSTER_NODES["jarvis"] = {
"name": "jarvis", "type": "worker", "status": "active",
"capabilities": ["llm"], "active_model": None, "load": None,
"registered_at": "2026-01-01T00:00:00Z", "last_seen": "2026-01-01T00:00:00Z",
}
cluster.CLUSTER_COORDINATOR = "ultron"
cluster._push_event("cluster", "info", "ultron", "Elected")
from routers.cluster import cluster_status
resp = asyncio.run(cluster_status())
assert "nodes" in resp
assert "node_count" in resp
assert "coordinator" in resp
assert "events" in resp
assert resp["node_count"] == 1
assert resp["coordinator"] == "ultron"
assert "jarvis" in resp["nodes"]
assert len(resp["events"]) == 1
# No internal keys leaked
for node in resp["nodes"].values():
for key in node:
assert key in {
"name", "type", "status", "capabilities", "active_model",
"load", "registered_at", "last_seen",
}