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:
@@ -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")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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")
|
||||||
@@ -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",
|
||||||
|
}}
|
||||||
@@ -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",
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user