"""Tracker HTTP server — node registry and route selection for the inference mesh. HTTP API contract: - POST /v1/nodes/register Request JSON: { "endpoint": "http://host:port", "shard_start": int, "shard_end": int, "model": str optional, "shard_checksum": str optional, "hardware_profile": object, "wallet_address": str optional, "score": number optional } Response 200: {"node_id": str} Response 400: {"error": str} - POST /v1/nodes/{node_id}/heartbeat Response 200: {} Response 404: {"error": "node not found"} - GET /v1/route?model= Response 200: {"route": ["http://node-a", "http://node-b"]} Response 400/404/503: {"error": str} - GET /v1/routes?model=&redundancy= Response 200: {"routes": [{"route": [...], "nodes": [...]}]} Response 400/404/503: {"error": str} """ import http.server import hashlib import json import socketserver import threading import time import urllib.parse from typing import Any from .gossip import NodeGossip from .raft import RaftNode DEFAULT_MODEL_PRESETS: dict[str, dict] = { "stub-model": { "layers_start": 0, "layers_end": 31, "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, }, "openai-community/gpt2": { "layers_start": 0, "layers_end": 11, "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, }, } DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024 DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024 DEFAULT_QUANTIZATIONS = ["bfloat16"] DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0 class _NodeEntry: __slots__ = ( "node_id", "endpoint", "shard_start", "shard_end", "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address", "score", "vram_bytes", "ram_bytes", "quantizations", "benchmark_tokens_per_sec", "quantization", "managed_assignment", "pending_directives", "last_heartbeat", "tracker_mode", "relay_addr", "cert_fingerprint", "peer_id", # heartbeat stats (reported by node, cumulative) "total_requests", "failed_requests", "queue_depth", "uptime_seconds", "status", # "ready" | "loading" "heartbeats_expected", "heartbeats_received", # dynamic reassignment queued by the tracker "pending_new_assignment", ) def __init__( self, node_id: str, endpoint: str, shard_start: int | None, shard_end: int | None, model: str | None, shard_checksum: str | None, hardware_profile: dict, wallet_address: str | None, score: float, vram_bytes: int = DEFAULT_VRAM_BYTES, ram_bytes: int = DEFAULT_RAM_BYTES, quantizations: list[str] | None = None, benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC, quantization: str | None = None, managed_assignment: bool = False, tracker_mode: bool = False, hf_repo: str | None = None, num_layers: int | None = None, relay_addr: str | None = None, cert_fingerprint: str | None = None, peer_id: str | None = None, ) -> None: self.node_id = node_id self.endpoint = endpoint self.shard_start = shard_start self.shard_end = shard_end self.model = model self.shard_checksum = shard_checksum self.hardware_profile = hardware_profile self.wallet_address = wallet_address self.score = score self.vram_bytes = vram_bytes self.ram_bytes = ram_bytes self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS) self.benchmark_tokens_per_sec = benchmark_tokens_per_sec self.quantization = quantization self.managed_assignment = managed_assignment self.tracker_mode = tracker_mode self.hf_repo = hf_repo self.num_layers = num_layers self.relay_addr = relay_addr self.cert_fingerprint = cert_fingerprint self.peer_id = peer_id self.pending_directives: list[dict] = [] self.last_heartbeat: float = time.monotonic() self.total_requests: int = 0 self.failed_requests: int = 0 self.queue_depth: int = 0 self.uptime_seconds: float = 0.0 self.status: str = "ready" self.heartbeats_expected: int = 0 self.heartbeats_received: int = 0 self.pending_new_assignment: dict | None = None def _select_route( nodes: list[_NodeEntry], required_start: int, required_end: int, ) -> tuple[list[_NodeEntry], str]: """Greedy interval-cover. Returns (ordered route, error_message).""" candidates = sorted( [node for node in nodes if node.shard_start is not None and node.shard_end is not None], key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] ) route: list[_NodeEntry] = [] covered_up_to = required_start - 1 while covered_up_to < required_end: best: _NodeEntry | None = None for node in candidates: if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to: if best is None or node.shard_end > best.shard_end: best = node if best is None: missing = covered_up_to + 1 return [], f"no route available: no registered node covers layer {missing}" route.append(best) covered_up_to = best.shard_end candidates = [n for n in candidates if n is not best] return route, "" def _coverage_percentage( nodes: list[_NodeEntry], required_start: int, required_end: int, ) -> float: required_layers = required_end - required_start + 1 if required_layers <= 0: return 0.0 intervals = sorted( ( (max(required_start, node.shard_start), min(required_end, node.shard_end)) for node in nodes if node.shard_start is not None and node.shard_end is not None if node.shard_end >= required_start and node.shard_start <= required_end ), key=lambda interval: interval[0], ) covered = 0 covered_until = required_start - 1 for start, end in intervals: if end <= covered_until: continue next_start = max(start, covered_until + 1) if next_start > end: continue covered += end - next_start + 1 covered_until = end return round((covered / required_layers) * 100, 2) def _preset_layer_bounds(preset: dict) -> tuple[int, int]: start = int(preset.get("layers_start", 0)) if "layers_end" in preset: return start, int(preset["layers_end"]) return start, start + int(preset["total_layers"]) - 1 def _preset_bytes_per_layer(preset: dict) -> dict[str, int]: raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {})) if isinstance(raw, dict) and raw: return {str(quant): int(value) for quant, value in raw.items()} return {"bfloat16": 30 * 1024 * 1024} def _node_quantization(node: _NodeEntry, preset: dict) -> str: bytes_per_layer = _preset_bytes_per_layer(preset) if node.quantization in bytes_per_layer: return node.quantization for quantization in node.quantizations: if quantization in bytes_per_layer: return quantization return next(iter(bytes_per_layer)) def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: bytes_per_layer = _preset_bytes_per_layer(preset) quantization = _node_quantization(node, preset) layer_bytes = bytes_per_layer[quantization] if layer_bytes <= 0: return 0 return int((node.vram_bytes * 0.8) // layer_bytes) def _coverage_map( nodes: list[_NodeEntry], required_start: int, required_end: int, ) -> list[dict]: layer_counts = [] for layer in range(required_start, required_end + 1): count = 0 for node in nodes: if node.shard_start is None or node.shard_end is None: continue if node.shard_start <= layer <= node.shard_end: count += 1 layer_counts.append((layer, count)) coverage: list[dict] = [] for layer, count in layer_counts: if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1: coverage[-1]["end_layer"] = layer else: coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count}) return coverage def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: """Per-node health detail for the availability map.""" age = time.monotonic() - node.last_heartbeat alive = age <= heartbeat_timeout hb_expected = max(1, round(node.uptime_seconds / 20.0)) # assume ~20s interval hb_rate = round(min(1.0, node.heartbeats_received / hb_expected), 4) if node.heartbeats_received else 0.0 total = node.total_requests inf_rate = round((total - node.failed_requests) / total, 4) if total > 0 else 1.0 return { "node_id": node.node_id, "endpoint": node.endpoint, "alive": alive, "last_seen_seconds_ago": round(age, 1), "status": node.status, "queue_depth": node.queue_depth, "total_requests": node.total_requests, "heartbeat_success_rate": hb_rate, "inference_success_rate": inf_rate, } def _coverage_map_detailed( nodes: list["_NodeEntry"], required_start: int, required_end: int, heartbeat_timeout: float, ) -> list[dict]: """Like _coverage_map but with per-node identity and health in each band. Includes all nodes (alive and stale) so operators can see both coverage holes and which dead nodes used to fill them. """ now = time.monotonic() def covers(node: "_NodeEntry", layer: int) -> bool: return ( node.shard_start is not None and node.shard_end is not None and node.shard_start <= layer <= node.shard_end ) # Build per-layer list of nodes (with alive flag) layer_nodes: list[list[tuple["_NodeEntry", bool]]] = [] for layer in range(required_start, required_end + 1): ns = [ (n, (now - n.last_heartbeat) <= heartbeat_timeout) for n in nodes if covers(n, layer) ] layer_nodes.append(ns) # Merge consecutive layers with identical node-set (same node_ids, same alive) coverage: list[dict] = [] for i, layer_idx in enumerate(range(required_start, required_end + 1)): ns = layer_nodes[i] node_ids = [n.node_id for n, _ in ns] alive_flags = [a for _, a in ns] # Compare with last band if ( coverage and coverage[-1]["end_layer"] == layer_idx - 1 and [nd["node_id"] for nd in coverage[-1]["nodes"]] == node_ids and [nd["alive"] for nd in coverage[-1]["nodes"]] == alive_flags ): coverage[-1]["end_layer"] = layer_idx else: coverage.append({ "start_layer": layer_idx, "end_layer": layer_idx, "node_count": len(ns), "nodes": [_node_health(n, heartbeat_timeout) for n, _ in ns], }) return coverage def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: return [ (segment["start_layer"], segment["end_layer"]) for segment in coverage if segment["node_count"] == 0 ] def _nodes_and_bounds_for_model( server: "_TrackerHTTPServer", model: str, ) -> tuple[list[_NodeEntry], int, int] | None: preset = server.model_presets.get(model) if preset is not None: required_start, required_end = _preset_layer_bounds(preset) return [node for node in server.registry.values() if node.model == model], required_start, required_end nodes = [ node for node in server.registry.values() if (node.hf_repo == model or node.model == model) and node.shard_start is not None and node.shard_end is not None and node.num_layers is not None ] if not nodes: return None return nodes, 0, max(node.num_layers for node in nodes) - 1 def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: return { "action": "LOAD_SHARD", "model": model, "start_layer": start, "end_layer": end, "shard_start": start, "shard_end": end, "quantization": quantization, } def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: return { "action": "DROP_SHARD", "model": model, "start_layer": start, "end_layer": end, "shard_start": start, "shard_end": end, "quantization": quantization, } def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: now = time.monotonic() expired_ids = [ node_id for node_id, entry in server.registry.items() if (now - entry.last_heartbeat) > server.heartbeat_timeout ] for node_id in expired_ids: entry = server.registry.pop(node_id) print( f"[tracker] node expired: {node_id} {entry.endpoint} " f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)", flush=True, ) if expired_ids: _rebalance_all_locked(server) return expired_ids def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: preset = server.model_presets.get(model) if preset is None: return required_start, required_end = _preset_layer_bounds(preset) total_layers = required_end - required_start + 1 model_nodes = [node for node in server.registry.values() if node.model == model] managed_nodes = [node for node in model_nodes if node.managed_assignment] if not managed_nodes: return previous_ranges = { node.node_id: (node.shard_start, node.shard_end, node.quantization) for node in managed_nodes } for node in managed_nodes: node.shard_start = None node.shard_end = None managed_nodes.sort( key=lambda node: ( -node.benchmark_tokens_per_sec, -_node_layer_capacity(node, preset), node.node_id, ) ) base_nodes = [node for node in model_nodes if not node.managed_assignment] coverage = _coverage_map(base_nodes, required_start, required_end) gaps = _coverage_gaps(coverage) if not gaps: gaps = [(required_start, required_end)] eligible_nodes = [ node for node in managed_nodes if _node_layer_capacity(node, preset) > 0 ] node_index = 0 for gap_start, gap_end in gaps: cursor = gap_start while cursor <= gap_end and node_index < len(eligible_nodes): node = eligible_nodes[node_index] remaining_layers = gap_end - cursor + 1 remaining_nodes_after = len(eligible_nodes) - node_index - 1 capacity = min( _node_layer_capacity(node, preset), total_layers, max(1, remaining_layers - remaining_nodes_after), ) if capacity <= 0: node_index += 1 continue quantization = _node_quantization(node, preset) node.quantization = quantization node.shard_start = cursor node.shard_end = min(gap_end, cursor + capacity - 1) cursor = node.shard_end + 1 node_index += 1 for node in managed_nodes: previous_start, previous_end, previous_quantization = previous_ranges[node.node_id] current_range = (node.shard_start, node.shard_end, node.quantization) if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]: continue if previous_start is not None and previous_end is not None: node.pending_directives.append( _drop_directive( node, model, previous_start, previous_end, previous_quantization or _node_quantization(node, preset), ) ) node.pending_directives.append( _load_directive( node, model, node.shard_start, node.shard_end, node.quantization or _node_quantization(node, preset), ) ) def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: for model in list(server.model_presets): _rebalance_model_locked(server, model) def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: if contracts is None or not wallet_address: return None if contracts.registry.get_wallet(wallet_address).banned: return "wallet is banned" return None def _node_id_for_registration( endpoint: str, model: str, wallet_address: str | None, shard_start: int | None, shard_end: int | None, hf_repo: str | None, ) -> str: wallet_prefix = wallet_address[:8] if wallet_address else "anon" stable_key = "|".join([ wallet_address or "", endpoint.rstrip("/"), model, hf_repo or "", "" if shard_start is None else str(shard_start), "" if shard_end is None else str(shard_end), ]) digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12] return f"{wallet_prefix}-{digest}" class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): daemon_threads = True def __init__( self, addr: tuple, handler, registry: dict, lock: threading.Lock, heartbeat_timeout: float, 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 self.lock = lock self.heartbeat_timeout = heartbeat_timeout self.model_presets = model_presets self.contracts = contracts self.minimum_stake = minimum_stake self.raft = raft self.gossip = gossip class _TrackerHandler(http.server.BaseHTTPRequestHandler): def log_message(self, fmt, *args): # noqa: suppress request logs in tests pass def _send_json(self, status: int, data: dict) -> None: body = json.dumps(data).encode() self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) self.end_headers() try: self.wfile.write(body) except BrokenPipeError: pass def _read_json_body(self) -> dict | None: length = int(self.headers.get("Content-Length", 0)) try: body = json.loads(self.rfile.read(length) or b"{}") except json.JSONDecodeError: self._send_json(400, {"error": "invalid JSON body"}) return None if not isinstance(body, dict): self._send_json(400, {"error": "JSON body must be an object"}) return None return body def _purge_expired_nodes(self) -> None: server: _TrackerHTTPServer = self.server # type: ignore[assignment] _purge_expired_nodes_locked(server) def do_POST(self): if self.path == "/v1/chat/completions": self._handle_proxy_chat() return 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//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": self._handle_heartbeat(parts[3]) return self.send_response(404) self.end_headers() def do_GET(self): parsed = urllib.parse.urlparse(self.path) if parsed.path == "/v1/route": self._handle_route(parsed) elif parsed.path == "/v1/routes": self._handle_routes(parsed) elif parsed.path == "/v1/nodes/assign": self._handle_assign(parsed) elif parsed.path == "/v1/network/assign": self._handle_network_assign(parsed) elif parsed.path == "/v1/models": self._handle_models() elif parsed.path.startswith("/v1/coverage/"): model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/")) self._handle_coverage(model) 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() elif parsed.path == "/v1/health": self._send_json(200, {"status": "ok"}) else: self.send_response(404) self.end_headers() def _handle_models(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] created = int(time.time()) with server.lock: self._purge_expired_nodes() alive = list(server.registry.values()) if server.contracts is not None: alive = [ node for node in alive if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] data = [] seen_ids: set[str] = set() for name, preset in server.model_presets.items(): model_nodes = [node for node in alive if node.model == name] if not model_nodes: continue required_start, required_end = _preset_layer_bounds(preset) coverage = _coverage_percentage( model_nodes, required_start, required_end, ) aliases = [name] hf_repo = preset.get("hf_repo") if hf_repo and hf_repo not in aliases: aliases.append(hf_repo) data.append({ "id": name, "object": "model", "created": created, "owned_by": "meshnet", "name": name, "hf_repo": hf_repo, "aliases": aliases, "shard_coverage_percentage": coverage, }) seen_ids.add(name) hf_model_ids = sorted({ node.hf_repo or node.model for node in alive if node.model is not None and node.model not in server.model_presets and node.shard_start is not None and node.shard_end is not None and node.num_layers is not None }) for model_id in hf_model_ids: if model_id is None or model_id in seen_ids: continue model_nodes = [ node for node in alive if node.shard_start is not None and node.shard_end is not None and node.num_layers is not None and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id)) ] if not model_nodes: continue short_names = sorted({node.model for node in model_nodes if node.model}) aliases = [model_id, *[name for name in short_names if name != model_id]] required_start = 0 required_end = max(node.num_layers for node in model_nodes) - 1 data.append({ "id": model_id, "object": "model", "created": created, "owned_by": "meshnet", "name": short_names[0] if short_names else model_id, "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None, "aliases": aliases, "shard_coverage_percentage": _coverage_percentage( model_nodes, required_start, required_end, ), }) seen_ids.add(model_id) self._send_json(200, {"object": "list", "data": data}) def _handle_coverage(self, model: str): server: _TrackerHTTPServer = self.server # type: ignore[assignment] # Do NOT purge before coverage — dead nodes are included with alive=false # so operators can see what was covering each layer band before failure. with server.lock: resolved = _nodes_and_bounds_for_model(server, model) if resolved is None: self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) return all_nodes, required_start, required_end = resolved if server.contracts is not None: all_nodes = [ node for node in all_nodes if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] coverage = _coverage_map_detailed(all_nodes, required_start, required_end, server.heartbeat_timeout) self._send_json(200, {"model": model, "coverage": coverage}) def _handle_tracker_nodes(self, model: str): """Return nodes registered with tracker_mode=True whose shard starts at layer 0.""" server: _TrackerHTTPServer = self.server # type: ignore[assignment] preset = server.model_presets.get(model) if preset is None: self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return required_start, _ = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() alive = [node for node in server.registry.values() if node.model == model] if server.contracts is not None: alive = [ node for node in alive if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] tracker_nodes = [ node for node in alive if node.shard_start is not None and node.shard_start == required_start and node.tracker_mode ] self._send_json(200, { "model": model, "tracker_nodes": [ { "node_id": node.node_id, "endpoint": node.endpoint, "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, } for node in tracker_nodes ], }) # ---------------------------------------------------------------- OpenAI proxy def _handle_proxy_chat(self) -> None: """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node. Picks a live tracker-mode node for the requested model using round-robin, then forwards the request verbatim and relays the response (including streaming SSE chunks) back to the caller. """ server: _TrackerHTTPServer = self.server # type: ignore[assignment] length = int(self.headers.get("Content-Length", 0)) raw_body = self.rfile.read(length) if length else b"{}" try: body = json.loads(raw_body) except json.JSONDecodeError: self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}}) return model: str = body.get("model", "") is_stream: bool = bool(body.get("stream", False)) # Find a live tracker-mode node for this model with server.lock: self._purge_expired_nodes() candidates = [ n for n in server.registry.values() if n.tracker_mode and (n.model == model or n.hf_repo == model) ] if not candidates: # Fall back: any node serving shard_start=0 for this model with server.lock: candidates = [ n for n in server.registry.values() if n.shard_start == 0 and (n.model == model or n.hf_repo == model) ] if not candidates: self._send_json(503, {"error": { "message": f"no nodes available for model {model!r}", "type": "service_unavailable", "code": "model_not_available", }}) return # Simple round-robin via list length modulo (stateless, good enough) node = candidates[int(time.time() * 1000) % len(candidates)] target_url = f"{node.endpoint}/v1/chat/completions" request_id = str(body.get("id") or f"req-{time.time_ns():x}") # Pre-resolve the downstream route so the first-shard node skips its own # tracker query. We already hold the full registry picture — no need for # a second round-trip. route_model = node.hf_repo or node.model or model with server.lock: if route_model in server.model_presets: preset = server.model_presets[route_model] rs, re = _preset_layer_bounds(preset) all_nodes: list = [n for n in server.registry.values() if n.model == route_model] else: all_nodes = [ n for n in server.registry.values() if (n.hf_repo == route_model or n.model == route_model) and n.shard_start is not None and n.num_layers is not None ] rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) route_nodes, _ = _select_route(all_nodes, rs, re) # Compute start_layer for each hop: each node begins where the previous ended + 1. # This allows overlapping shard registrations without double-computation. covered_up_to = rs - 1 route_hops: list[dict] = [] for rn in route_nodes: route_hops.append({"endpoint": rn.endpoint, "start_layer": covered_up_to + 1}) covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to # Strip the first-shard node we're about to proxy to — it's already handling the request. downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] downstream_urls = json.dumps(downstream_hops) route_debug = " -> ".join( f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" for n in route_nodes ) print( f"[tracker] proxy route {request_id}: model={model!r} " f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} " f"route={route_debug or ''}", flush=True, ) req = urllib.request.Request( target_url, data=raw_body, headers={ "Content-Type": "application/json", "X-Meshnet-Route": downstream_urls, }, method="POST", ) # Copy Authorization header from client if present auth = self.headers.get("Authorization") if auth: req.add_header("Authorization", auth) try: upstream = urllib.request.urlopen(req, timeout=300.0) print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True) except urllib.error.HTTPError as exc: # Relay error status + body from node err_body = exc.read() self.send_response(exc.code) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(err_body))) self.end_headers() try: self.wfile.write(err_body) except BrokenPipeError: pass return except Exception as exc: print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True) self._send_json(503, {"error": { "message": f"upstream node unreachable: {exc}", "type": "service_unavailable", "code": "upstream_error", }}) return with upstream: content_type = upstream.headers.get("Content-Type", "application/json") if is_stream or "text/event-stream" in content_type: # Relay SSE stream chunk-by-chunk self.send_response(200) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") self.end_headers() try: while True: line = upstream.readline() if not line: break self.wfile.write(line) self.wfile.flush() except BrokenPipeError: pass else: # Non-streaming: buffer and relay resp_body = upstream.read() print( f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}", flush=True, ) self.send_response(200) self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(resp_body))) self.end_headers() try: self.wfile.write(resp_body) except BrokenPipeError: pass def _handle_register(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] body = self._read_json_body() 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"}) return parsed_endpoint = urllib.parse.urlparse(endpoint) if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc: self._send_json(400, {"error": "endpoint must be an http(s) URL"}) return shard_start: int | None shard_end: int | None explicit_shard = "shard_start" in body or "shard_end" in body if explicit_shard: try: shard_start = int(body["shard_start"]) shard_end = int(body["shard_end"]) except (KeyError, TypeError, ValueError): self._send_json(400, {"error": "shard_start and shard_end must be numeric"}) return if shard_start < 0 or shard_end < 0 or shard_start > shard_end: self._send_json(400, {"error": "shard range must be non-negative and ordered"}) return else: shard_start = None shard_end = None try: score = float(body.get("score", 1.0)) except (TypeError, ValueError): self._send_json(400, {"error": "score must be numeric"}) return hardware_profile = body.get("hardware_profile", {}) if not isinstance(hardware_profile, dict): self._send_json(400, {"error": "hardware_profile must be an object"}) return model = body.get("model") if model is None: model = "stub-model" if not isinstance(model, str): self._send_json(400, {"error": "model must be a string"}) return shard_checksum = body.get("shard_checksum") if shard_checksum is not None and not isinstance(shard_checksum, str): self._send_json(400, {"error": "shard_checksum must be a string"}) return try: vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES)) ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES)) benchmark_tokens_per_sec = float( body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC) ) except (TypeError, ValueError): self._send_json(400, {"error": "vram_bytes, ram_bytes, and benchmark_tokens_per_sec must be numeric"}) return if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0: self._send_json(400, {"error": "capability values must be positive"}) return quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS) if not ( isinstance(quantizations_body, list) and quantizations_body and all(isinstance(item, str) and item for item in quantizations_body) ): self._send_json(400, {"error": "quantizations must be a non-empty string array"}) return quantizations = list(quantizations_body) quantization = body.get("quantization") if quantization is not None and not isinstance(quantization, str): self._send_json(400, {"error": "quantization must be a string"}) return wallet_address = body.get("wallet_address") if wallet_address is not None and not isinstance(wallet_address, str): self._send_json(400, {"error": "wallet_address must be a string"}) return ban_error = _registration_ban_error(server.contracts, wallet_address) if ban_error: self._send_json(403, {"error": ban_error}) return tracker_mode = bool(body.get("tracker_mode", False)) hf_repo = body.get("hf_repo") if hf_repo is not None and not isinstance(hf_repo, str): self._send_json(400, {"error": "hf_repo must be a string"}) return num_layers_body = body.get("num_layers") num_layers: int | None = None if num_layers_body is not None: try: num_layers = int(num_layers_body) except (TypeError, ValueError): self._send_json(400, {"error": "num_layers must be an integer"}) return relay_addr = body.get("relay_addr") or None cert_fingerprint = body.get("cert_fingerprint") or None peer_id = body.get("peer_id") or None node_id = _node_id_for_registration( endpoint, model, wallet_address, shard_start, shard_end, hf_repo, ) entry = _NodeEntry( node_id=node_id, endpoint=endpoint.rstrip("/"), shard_start=shard_start, shard_end=shard_end, model=model, shard_checksum=shard_checksum, hardware_profile=hardware_profile, wallet_address=wallet_address, score=score, vram_bytes=vram_bytes, ram_bytes=ram_bytes, quantizations=quantizations, benchmark_tokens_per_sec=benchmark_tokens_per_sec, quantization=quantization, managed_assignment=not explicit_shard, tracker_mode=tracker_mode, hf_repo=hf_repo, num_layers=num_layers, relay_addr=relay_addr, cert_fingerprint=cert_fingerprint, peer_id=peer_id, ) with server.lock: self._purge_expired_nodes() # Dedup: if this endpoint is already registered, remove the old entry first. stale_ids = [ eid for eid, e in server.registry.items() if e.endpoint == entry.endpoint.rstrip("/") ] for eid in stale_ids: old = server.registry.pop(eid) print( f"[tracker] node re-registered: replaced {eid} with {node_id}" f" {old.endpoint}", flush=True, ) server.registry[node_id] = entry if entry.managed_assignment: _rebalance_model_locked(server, model) assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None if assignment_directive is not None: entry.pending_directives.clear() shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" repo_info = f" [{hf_repo}]" if hf_repo else "" print( f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info}", flush=True, ) payload = {"node_id": node_id} if assignment_directive is not None: 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] body: dict = {} content_length = int(self.headers.get("Content-Length", 0)) if content_length > 0: try: body = json.loads(self.rfile.read(content_length)) except Exception: pass with server.lock: self._purge_expired_nodes() entry = server.registry.get(node_id) if entry is None: self._send_json(404, {"error": "node not found"}) return entry.last_heartbeat = time.monotonic() entry.heartbeats_received += 1 # P2P metadata if body.get("relay_addr"): entry.relay_addr = body["relay_addr"] if body.get("cert_fingerprint"): entry.cert_fingerprint = body["cert_fingerprint"] if body.get("peer_id"): entry.peer_id = body["peer_id"] # Node stats (cumulative — node always sends totals, tracker replaces) if "total_requests" in body: entry.total_requests = int(body["total_requests"]) if "failed_requests" in body: entry.failed_requests = int(body["failed_requests"]) if "queue_depth" in body: entry.queue_depth = int(body["queue_depth"]) if "uptime_seconds" in body: entry.uptime_seconds = float(body["uptime_seconds"]) if "status" in body and body["status"] in ("ready", "loading"): entry.status = body["status"] _rebalance_model_locked(server, entry.model or "stub-model") directives = list(entry.pending_directives) entry.pending_directives.clear() new_assignment = entry.pending_new_assignment if new_assignment is not None: entry.pending_new_assignment = None if server.gossip is not None: server.gossip.record(node_id) resp: dict = {} if directives: resp["directives"] = directives if new_assignment is not None: resp["new_assignment"] = new_assignment self._send_json(200, resp) # ---------------------------------------------------------------- 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. Query params: model — model preset name (default: first preset) device — "cuda" | "cpu" vram_mb — integer VRAM in MB (0 for CPU) The greedy strategy: find the first gap in current layer coverage and assign it. If no gap exists, assign the full model range so the node provides redundant coverage. """ server: _TrackerHTTPServer = self.server # type: ignore[assignment] params = urllib.parse.parse_qs(parsed.query) model_list = params.get("model") if not model_list: model = next(iter(server.model_presets), None) if model is None: self._send_json(503, {"error": "no model presets configured"}) return else: model = model_list[0] preset = server.model_presets.get(model) if preset is None: self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() alive = [node for node in server.registry.values() if node.model == model] if server.contracts is not None: alive = [ node for node in alive if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] device = params.get("device", ["cpu"])[0] try: vram_mb = int(params.get("vram_mb", ["0"])[0]) except ValueError: vram_mb = 0 max_layers = required_end - required_start + 1 if device != "cuda" or vram_mb < 8192: max_layers = min(max_layers, 16) # Collect covered intervals sorted by start layer. covered = sorted( [ (n.shard_start, n.shard_end) for n in alive if n.shard_start is not None and n.shard_end is not None ], key=lambda t: t[0], ) # Walk from required_start to find the first uncovered layer. gap_start = required_start for s, e in covered: if s <= gap_start: gap_start = max(gap_start, e + 1) else: break # gap found before this node if gap_start > required_end: # Full coverage exists — assign the full range for redundancy. shard_start = required_start shard_end = min(required_end, shard_start + max_layers - 1) else: shard_start = gap_start shard_end = min(required_end, shard_start + max_layers - 1) peers = [ {"endpoint": node.endpoint, "checksum": node.shard_checksum} for node in alive if node.model == model and node.shard_start == shard_start and node.shard_end == shard_end and node.shard_checksum ] self._send_json(200, { "shard_start": shard_start, "shard_end": shard_end, "model": model, "model_layers_end": required_end, "peers": peers, **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}), }) def _handle_network_assign(self, parsed: urllib.parse.ParseResult): """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes. Query params: vram_mb — integer VRAM in MB (0 = CPU-only node) device — "cuda" | "cpu" hf_repo — optional; if set, restrict search to this repo only Returns: {hf_repo, shard_start, shard_end, num_layers, gap_found} gap_found=true means a real uncovered gap was assigned; false means redundancy. """ server: _TrackerHTTPServer = self.server # type: ignore[assignment] params = urllib.parse.parse_qs(parsed.query) try: vram_mb = int(params.get("vram_mb", ["0"])[0]) except ValueError: vram_mb = 0 device = params.get("device", ["cpu"])[0] filter_repo = params.get("hf_repo", [None])[0] # optional repo filter with server.lock: self._purge_expired_nodes() all_nodes = list(server.registry.values()) # Collect only nodes that registered a real HF model (have hf_repo + shard bounds). hf_nodes = [ n for n in all_nodes if n.hf_repo and n.shard_start is not None and n.shard_end is not None and n.num_layers is not None and (filter_repo is None or n.hf_repo == filter_repo) ] if not hf_nodes: msg = ( f"no HF-model nodes registered for {filter_repo!r}" if filter_repo else "no HF-model nodes registered; cannot assign shards" ) self._send_json(503, {"error": msg}) return # Group by hf_repo; pick the one with the largest total_layers and biggest gap. from collections import defaultdict repo_groups: dict = defaultdict(list) repo_layers: dict = {} for n in hf_nodes: repo_groups[n.hf_repo].append(n) # Use the largest num_layers seen for this repo. if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]: repo_layers[n.hf_repo] = n.num_layers # Pick the repo where the gap is largest (most work to do). best_repo = None best_gap_size = -1 best_gap_start = 0 best_num_layers = 0 for repo, nodes in repo_groups.items(): total = repo_layers[repo] covered = sorted( [(n.shard_start, n.shard_end) for n in nodes], key=lambda t: t[0], ) # Walk from 0 to find first uncovered layer. gap_start = 0 for s, e in covered: if s <= gap_start: gap_start = max(gap_start, e + 1) else: break gap_size = max(0, (total - 1) - gap_start + 1) # layers remaining uncovered if gap_size > best_gap_size: best_gap_size = gap_size best_gap_start = gap_start best_repo = repo best_num_layers = total gap_found = best_gap_size > 0 if not gap_found: # All shards are covered — still assign to the model with most nodes for redundancy. best_repo = max(repo_groups, key=lambda r: len(repo_groups[r])) best_gap_start = 0 best_num_layers = repo_layers[best_repo] # Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM. total_l = best_num_layers if device == "cuda" and vram_mb >= 8192: max_layers = total_l else: max_layers = max(1, total_l // 2) shard_start = best_gap_start shard_end = min(total_l - 1, shard_start + max_layers - 1) self._send_json(200, { "hf_repo": best_repo, "shard_start": shard_start, "shard_end": shard_end, "num_layers": total_l, "gap_found": gap_found, }) def _handle_route(self, parsed: urllib.parse.ParseResult): server: _TrackerHTTPServer = self.server # type: ignore[assignment] params = urllib.parse.parse_qs(parsed.query) model_list = params.get("model") if not model_list: self._send_json(400, {"error": "missing 'model' query parameter"}) return model = model_list[0] preset = server.model_presets.get(model) with server.lock: self._purge_expired_nodes() if preset is not None: # Preset-based routing (stub-model system). alive = [node for node in server.registry.values() if node.model == model] required_start, required_end = _preset_layer_bounds(preset) else: # HF model routing: match by hf_repo (full) or model short name. alive = [ node for node in server.registry.values() if (node.hf_repo == model or node.model == model) and node.shard_start is not None and node.shard_end is not None and node.num_layers is not None ] if not alive: self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) return required_start = 0 required_end = max(n.num_layers for n in alive) - 1 # type: ignore[type-var] if server.contracts is not None: alive = [ node for node in alive if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] route, error = _select_route(alive, required_start, required_end) if error: self._send_json(503, {"error": error}) return covered_up_to = required_start - 1 route_with_start: list[tuple] = [] for rn in route: route_with_start.append((rn, covered_up_to + 1)) covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to self._send_json(200, { "route": [e.endpoint for e, _ in route_with_start], "nodes": [ { "node_id": e.node_id, "endpoint": e.endpoint, "start_layer": start, "wallet_address": e.wallet_address, "shard_start": e.shard_start, "shard_end": e.shard_end, "model": e.model, "shard_checksum": e.shard_checksum, "score": e.score, } for e, start in route_with_start ], }) def _handle_routes(self, parsed: urllib.parse.ParseResult): server: _TrackerHTTPServer = self.server # type: ignore[assignment] params = urllib.parse.parse_qs(parsed.query) model_list = params.get("model") if not model_list: self._send_json(400, {"error": "missing 'model' query parameter"}) return try: redundancy = int(params.get("redundancy", ["1"])[0]) except ValueError: self._send_json(400, {"error": "redundancy must be an integer"}) return if redundancy < 1: self._send_json(400, {"error": "redundancy must be at least 1"}) return model = model_list[0] preset = server.model_presets.get(model) if preset is None: self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() candidates = [node for node in server.registry.values() if node.model == model] if server.contracts is not None: candidates = [ node for node in candidates if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] routes = [] remaining = list(candidates) for _ in range(redundancy): route, error = _select_route(remaining, required_start, required_end) if error: self._send_json(503, {"error": error}) return route_endpoints = {node.endpoint for node in route} remaining = [node for node in remaining if node.endpoint not in route_endpoints] routes.append({ "route": [e.endpoint for e in route], "nodes": [ { "node_id": e.node_id, "endpoint": e.endpoint, "wallet_address": e.wallet_address, "shard_start": e.shard_start, "shard_end": e.shard_end, "model": e.model, "shard_checksum": e.shard_checksum, "score": e.score, } for e in route ], }) self._send_json(200, {"routes": routes}) class TrackerServer: """HTTP tracker that manages node registration and resolves inference routes. Nodes register via POST /v1/nodes/register and keep themselves alive with POST /v1/nodes//heartbeat. The gateway queries GET /v1/route?model= to obtain an ordered list of node endpoints whose shards cover all layers. """ def __init__( self, host: str = "127.0.0.1", port: int = 0, heartbeat_timeout: float = 90.0, rebalance_interval: float = 30.0, 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 self._heartbeat_timeout = heartbeat_timeout self._rebalance_interval = rebalance_interval self._model_presets: dict = ( model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS) ) 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, self._registry, self._lock, self._heartbeat_timeout, self._model_presets, self._contracts, 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() self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True) 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 if server is None: return with self._lock: _purge_expired_nodes_locked(server) _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() self._server.shutdown() self._server.server_close() if self._thread is not None: self._thread.join(timeout=1) if self._rebalance_thread is not None: self._rebalance_thread.join(timeout=1) self._server = None self._thread = None self._rebalance_thread = None self._raft = None self._gossip = None self.port = None