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:
Dobromir Popov
2026-06-30 09:05:21 +03:00
parent bbe57d5f07
commit 748d535c46
6 changed files with 898 additions and 10 deletions

View File

@@ -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