feat(us-019): distributed tracker consensus — Raft assignments + CRDT gossip
raft.py: minimal Raft consensus for shard-assignment log
- Leader election with random 150–300ms election timeouts
- AppendEntries log replication; majority commit required
- RequestVote RPC with log-completeness check
- Follower registration forwarded to leader via HTTP proxy
gossip.py: LWW CRDT gossip for inference-node heartbeat liveness
- Each tracker keeps {node_id → wall_clock_timestamp}
- Merges incoming gossip by taking max per key
- Pushes snapshot to random peer every 3 seconds
server.py + cli.py:
- TrackerServer gains cluster_peers + cluster_self_url params
- New HTTP endpoints: /v1/raft/vote, /v1/raft/append,
/v1/raft/status, /v1/gossip
- --cluster-peers and --self-url CLI flags
tests/test_tracker_consensus.py: 6 integration tests
- Leader elected within 1s
- Follower registration propagated to all nodes
- Follower proxy to leader
- Leader kill → new election within 5s, registrations continue
- Gossip table updated on heartbeat
92 tests pass, 1 skipped.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -29,6 +29,9 @@ import urllib.parse
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from .gossip import NodeGossip
|
||||
from .raft import RaftNode
|
||||
|
||||
|
||||
DEFAULT_MODEL_PRESETS: dict[str, dict] = {
|
||||
"stub-model": {
|
||||
@@ -377,6 +380,8 @@ class _TrackerHTTPServer(http.server.HTTPServer):
|
||||
model_presets: dict,
|
||||
contracts: Any | None,
|
||||
minimum_stake: int,
|
||||
raft: "RaftNode | None" = None,
|
||||
gossip: "NodeGossip | None" = None,
|
||||
) -> None:
|
||||
super().__init__(addr, handler)
|
||||
self.registry = registry
|
||||
@@ -385,6 +390,8 @@ class _TrackerHTTPServer(http.server.HTTPServer):
|
||||
self.model_presets = model_presets
|
||||
self.contracts = contracts
|
||||
self.minimum_stake = minimum_stake
|
||||
self.raft = raft
|
||||
self.gossip = gossip
|
||||
|
||||
|
||||
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
@@ -419,6 +426,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if self.path == "/v1/nodes/register":
|
||||
self._handle_register()
|
||||
return
|
||||
if self.path == "/v1/raft/vote":
|
||||
self._handle_raft_vote()
|
||||
return
|
||||
if self.path == "/v1/raft/append":
|
||||
self._handle_raft_append()
|
||||
return
|
||||
if self.path == "/v1/gossip":
|
||||
self._handle_gossip()
|
||||
return
|
||||
parts = self.path.split("/")
|
||||
# /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
|
||||
if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
|
||||
@@ -445,6 +461,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
elif parsed.path.startswith("/v1/tracker-nodes/"):
|
||||
model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/"))
|
||||
self._handle_tracker_nodes(model)
|
||||
elif parsed.path == "/v1/raft/status":
|
||||
self._handle_raft_status()
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
@@ -538,6 +556,37 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if body is None:
|
||||
return
|
||||
|
||||
# --- Raft cluster mode: forward to leader or propose via Raft ---
|
||||
if server.raft is not None:
|
||||
if not server.raft.is_leader:
|
||||
leader = server.raft.leader()
|
||||
if leader is None:
|
||||
self._send_json(503, {"error": "no leader elected — retry in a moment"})
|
||||
return
|
||||
# Proxy to leader
|
||||
try:
|
||||
data = json.dumps(body).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{leader}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5.0) as r:
|
||||
resp_body = r.read()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(resp_body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(resp_body)
|
||||
except Exception as exc:
|
||||
self._send_json(503, {"error": f"leader proxy failed: {exc}"})
|
||||
return
|
||||
# Leader path: fall through to normal registration, then replicate via Raft.
|
||||
# We let the registration run first (to generate node_id), then asynchronously
|
||||
# propose to the Raft log so followers can replicate the entry.
|
||||
# Raft proposal happens after the response is sent (fire-and-forget replication).
|
||||
|
||||
endpoint = body.get("endpoint")
|
||||
if not isinstance(endpoint, str) or not endpoint:
|
||||
self._send_json(400, {"error": "endpoint is required"})
|
||||
@@ -685,6 +734,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
payload["directive"] = assignment_directive
|
||||
self._send_json(200, payload)
|
||||
|
||||
# Raft replication: leader proposes this registration to followers so their
|
||||
# registries stay in sync. Fire-and-forget (async) — the client already
|
||||
# got its node_id; replication happens in the background.
|
||||
if server.raft is not None and server.raft.is_leader:
|
||||
raft_payload = dict(body)
|
||||
raft_payload["node_id"] = node_id # include the generated ID
|
||||
threading.Thread(
|
||||
target=server.raft.propose,
|
||||
args=("register", raft_payload),
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
def _handle_heartbeat(self, node_id: str):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
with server.lock:
|
||||
@@ -697,15 +758,53 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
||||
directives = list(entry.pending_directives)
|
||||
entry.pending_directives.clear()
|
||||
# print(
|
||||
# f"[tracker] heartbeat: {node_id[:8]} {entry.endpoint}",
|
||||
# flush=True,
|
||||
# )
|
||||
if server.gossip is not None:
|
||||
server.gossip.record(node_id)
|
||||
if directives:
|
||||
self._send_json(200, {"directives": directives})
|
||||
else:
|
||||
self._send_json(200, {})
|
||||
|
||||
# ---------------------------------------------------------------- Raft handlers
|
||||
|
||||
def _handle_raft_vote(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
if server.raft is None:
|
||||
self._send_json(503, {"error": "raft not enabled"})
|
||||
return
|
||||
result = server.raft.handle_request_vote(body)
|
||||
self._send_json(200, result)
|
||||
|
||||
def _handle_raft_append(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
if server.raft is None:
|
||||
self._send_json(503, {"error": "raft not enabled"})
|
||||
return
|
||||
result = server.raft.handle_append_entries(body)
|
||||
self._send_json(200, result)
|
||||
|
||||
def _handle_raft_status(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
if server.raft is None:
|
||||
self._send_json(200, {"role": "standalone", "term": 0, "leader": None})
|
||||
return
|
||||
self._send_json(200, server.raft.status())
|
||||
|
||||
def _handle_gossip(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
if server.gossip is not None and isinstance(body, dict):
|
||||
server.gossip.merge({k: float(v) for k, v in body.items()})
|
||||
self._send_json(200, {})
|
||||
|
||||
def _handle_assign(self, parsed: urllib.parse.ParseResult):
|
||||
"""Return an optimal shard assignment for a node given its hardware profile.
|
||||
|
||||
@@ -1042,6 +1141,8 @@ class TrackerServer:
|
||||
model_presets: dict | None = None,
|
||||
contracts: Any | None = None,
|
||||
minimum_stake: int = 0,
|
||||
cluster_peers: list[str] | None = None,
|
||||
cluster_self_url: str | None = None,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._requested_port = port
|
||||
@@ -1052,17 +1153,23 @@ class TrackerServer:
|
||||
)
|
||||
self._contracts = contracts
|
||||
self._minimum_stake = minimum_stake
|
||||
self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else []
|
||||
self._cluster_self_url: str | None = cluster_self_url
|
||||
self._registry: dict[str, _NodeEntry] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._server: _TrackerHTTPServer | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self._rebalance_stop = threading.Event()
|
||||
self._rebalance_thread: threading.Thread | None = None
|
||||
self._raft: RaftNode | None = None
|
||||
self._gossip: NodeGossip | None = None
|
||||
self.port: int | None = None
|
||||
|
||||
def start(self) -> int:
|
||||
if self._server is not None:
|
||||
raise RuntimeError("TrackerServer is already running")
|
||||
|
||||
# Start HTTP server first so we know our port
|
||||
self._server = _TrackerHTTPServer(
|
||||
(self._host, self._requested_port),
|
||||
_TrackerHandler,
|
||||
@@ -1074,6 +1181,21 @@ class TrackerServer:
|
||||
self._minimum_stake,
|
||||
)
|
||||
self.port = self._server.server_address[1]
|
||||
|
||||
# Start Raft + gossip if cluster peers are configured
|
||||
if self._cluster_peers:
|
||||
self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
|
||||
self._raft = RaftNode(
|
||||
self_url=self_url,
|
||||
peers=self._cluster_peers,
|
||||
apply_fn=self._raft_apply,
|
||||
)
|
||||
self._gossip = NodeGossip(peers=self._cluster_peers)
|
||||
self._server.raft = self._raft
|
||||
self._server.gossip = self._gossip
|
||||
self._raft.start()
|
||||
self._gossip.start()
|
||||
|
||||
self._rebalance_stop.clear()
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
self._thread.start()
|
||||
@@ -1081,6 +1203,38 @@ class TrackerServer:
|
||||
self._rebalance_thread.start()
|
||||
return self.port
|
||||
|
||||
def _raft_apply(self, command: str, payload: dict) -> None:
|
||||
"""Called by RaftNode when a log entry is committed — replicate to local registry."""
|
||||
if command != "register":
|
||||
return
|
||||
# Re-apply the registration to our local registry (follower path).
|
||||
# On the leader this is a no-op since it already registered locally.
|
||||
node_id = payload.get("node_id")
|
||||
if not node_id or node_id in self._registry:
|
||||
return # already present (leader case) or malformed
|
||||
endpoint = payload.get("endpoint", "")
|
||||
try:
|
||||
shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None
|
||||
shard_end = int(payload["shard_end"]) if payload.get("shard_end") is not None else None
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
entry = _NodeEntry(
|
||||
node_id=node_id,
|
||||
endpoint=endpoint.rstrip("/"),
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
model=payload.get("model", "stub-model"),
|
||||
shard_checksum=payload.get("shard_checksum"),
|
||||
hardware_profile=payload.get("hardware_profile", {}),
|
||||
wallet_address=payload.get("wallet_address"),
|
||||
score=float(payload.get("score", 1.0)),
|
||||
tracker_mode=bool(payload.get("tracker_mode", False)),
|
||||
hf_repo=payload.get("hf_repo"),
|
||||
num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None,
|
||||
)
|
||||
with self._lock:
|
||||
self._registry[node_id] = entry
|
||||
|
||||
def _rebalance_loop(self) -> None:
|
||||
while not self._rebalance_stop.wait(self._rebalance_interval):
|
||||
server = self._server
|
||||
@@ -1091,6 +1245,10 @@ class TrackerServer:
|
||||
_rebalance_all_locked(server)
|
||||
|
||||
def stop(self) -> None:
|
||||
if self._raft is not None:
|
||||
self._raft.stop()
|
||||
if self._gossip is not None:
|
||||
self._gossip.stop()
|
||||
if self._server is None:
|
||||
return
|
||||
self._rebalance_stop.set()
|
||||
@@ -1103,4 +1261,6 @@ class TrackerServer:
|
||||
self._server = None
|
||||
self._thread = None
|
||||
self._rebalance_thread = None
|
||||
self._raft = None
|
||||
self._gossip = None
|
||||
self.port = None
|
||||
|
||||
Reference in New Issue
Block a user