"""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 sqlite3 import threading import time import urllib.parse import urllib.request import uuid from typing import Any from .gossip import NodeGossip from .raft import RaftNode def derive_relay_url_from_public_tracker_url(url: str | None) -> str | None: """Return wss://host/ws when url is a public HTTPS tracker origin.""" if not url: return None parsed = urllib.parse.urlparse(url) if parsed.scheme != "https": return None host = parsed.hostname if not host or host in ("127.0.0.1", "localhost"): return None return f"wss://{parsed.netloc}/ws" 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 def _model_aliases(model: str | None) -> set[str]: """Return stable lookup aliases for a model repo or display name.""" if not model: return set() aliases = {model} short = model.rsplit("/", 1)[-1] aliases.add(short) lowered = short.lower() aliases.add(lowered) if lowered.endswith("-instruct"): aliases.add(lowered.removesuffix("-instruct")) return aliases def _node_matches_model(node: "_NodeEntry", model: str) -> bool: requested = _model_aliases(model) if not requested: return False return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo))) class _RollingCounter: """Circular-bucket request counter. Tracks events in `num_buckets` slots, each covering `bucket_seconds` of wall time. Old buckets are silently discarded when their epoch expires. """ def __init__(self, num_buckets: int, bucket_seconds: int) -> None: self._num = num_buckets self._bsec = bucket_seconds self._counts: list[int] = [0] * num_buckets self._epochs: list[int] = [-1] * num_buckets # which epoch each slot holds def _epoch(self, now: float) -> int: return int(now) // self._bsec def record(self, now: float | None = None) -> None: t = now if now is not None else time.time() ep = self._epoch(t) idx = ep % self._num if self._epochs[idx] != ep: self._counts[idx] = 0 self._epochs[idx] = ep self._counts[idx] += 1 def rpm(self, now: float | None = None) -> float: t = now if now is not None else time.time() cutoff = self._epoch(t) - self._num # epochs <= this are stale total = sum(self._counts[i] for i in range(self._num) if self._epochs[i] > cutoff) window_minutes = (self._num * self._bsec) / 60.0 return total / window_minutes if window_minutes > 0 else 0.0 def buckets(self) -> list[tuple[int, int]]: return [(self._epochs[i], self._counts[i]) for i in range(self._num)] def restore_buckets(self, data: list[tuple[int, int]]) -> None: for i, (ep, cnt) in enumerate(data): if i < self._num: self._epochs[i] = ep self._counts[i] = cnt class _ModelStats: """Three rolling windows for one model: last hour, last day, last month.""" def __init__(self) -> None: self.per_minute = _RollingCounter(60, 60) # 60 × 1-min buckets = 1 hour self.per_hour = _RollingCounter(24, 3600) # 24 × 1-hr buckets = 1 day self.per_day = _RollingCounter(30, 86400) # 30 × 1-day buckets = ~1 month def record(self, now: float | None = None) -> None: t = now if now is not None else time.time() self.per_minute.record(t) self.per_hour.record(t) self.per_day.record(t) def rpms(self, now: float | None = None) -> dict: t = now if now is not None else time.time() return { "rpm_last_hour": round(self.per_minute.rpm(t), 4), "rpm_last_day": round(self.per_hour.rpm(t), 4), "rpm_last_month": round(self.per_day.rpm(t), 4), } class _StatsCollector: """Thread-safe model request stats with SQLite persistence and peer slice merging.""" SAVE_INTERVAL = 60.0 # seconds between DB saves def __init__(self, db_path: str | None = None) -> None: self._lock = threading.Lock() self._local: dict[str, _ModelStats] = {} self._peer_rpms: dict[str, dict[str, dict]] = {} # tracker_url -> model -> rpms self._db_path = db_path if db_path: self._init_db() self._load_from_db() # ---------- public API ---------- def record_request(self, model: str, now: float | None = None) -> None: t = now if now is not None else time.time() with self._lock: if model not in self._local: self._local[model] = _ModelStats() self._local[model].record(t) def get_local_rpms(self, now: float | None = None) -> dict[str, dict]: t = now if now is not None else time.time() with self._lock: return {m: s.rpms(t) for m, s in self._local.items()} def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None: with self._lock: self._peer_rpms[tracker_url] = dict(rpms) def get_combined_stats(self, now: float | None = None) -> dict: t = now if now is not None else time.time() with self._lock: combined: dict[str, dict] = {} for model, ms in self._local.items(): combined[model] = ms.rpms(t) for _url, peer in self._peer_rpms.items(): for model, rpms in peer.items(): if model not in combined: combined[model] = {"rpm_last_hour": 0.0, "rpm_last_day": 0.0, "rpm_last_month": 0.0} for key in ("rpm_last_hour", "rpm_last_day", "rpm_last_month"): combined[model][key] = round(combined[model].get(key, 0.0) + rpms.get(key, 0.0), 4) return combined # ---------- persistence ---------- def _init_db(self) -> None: con = sqlite3.connect(self._db_path) # type: ignore[arg-type] con.execute( "CREATE TABLE IF NOT EXISTS model_rpm_buckets " "(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, " "PRIMARY KEY (model, window, bucket_idx))" ) con.commit() con.close() def save_to_db(self) -> None: if not self._db_path: return with self._lock: rows = [] for model, ms in self._local.items(): for window_name, counter in ( ("hour", ms.per_minute), ("day", ms.per_hour), ("month", ms.per_day), ): for idx, (ep, cnt) in enumerate(counter.buckets()): if ep >= 0: rows.append((model, window_name, idx, ep, cnt)) con = sqlite3.connect(self._db_path) # type: ignore[arg-type] con.executemany( "INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows ) con.commit() con.close() def _load_from_db(self) -> None: con = sqlite3.connect(self._db_path) # type: ignore[arg-type] rows = con.execute("SELECT model, window, bucket_idx, bucket_epoch, count FROM model_rpm_buckets").fetchall() con.close() grouped: dict[str, dict[str, list[tuple[int, int]]]] = {} for model, window, idx, ep, cnt in rows: grouped.setdefault(model, {}).setdefault(window, []).append((idx, ep, cnt)) for model, windows in grouped.items(): ms = _ModelStats() for window_name, entries in windows.items(): counter = {"hour": ms.per_minute, "day": ms.per_hour, "month": ms.per_day}.get(window_name) if counter is None: continue data = [(0, -1)] * counter._num for idx, ep, cnt in entries: if 0 <= idx < counter._num: data[idx] = (ep, cnt) counter.restore_buckets(data) self._local[model] = ms class _NodeEntry: __slots__ = ( "node_id", "endpoint", "shard_start", "shard_end", "model", "hf_repo", "num_layers", "model_metadata", "shard_checksum", "hardware_profile", "wallet_address", "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards", "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, max_loaded_shards: int = 1, 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, model_metadata: dict | 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.max_loaded_shards = max_loaded_shards 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.model_metadata = dict(model_metadata or {}) 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 _effective_throughput(node: "_NodeEntry") -> float: """Effective tokens/s accounting for current queue depth.""" return node.benchmark_tokens_per_sec / (node.queue_depth + 1) def _select_route( nodes: list[_NodeEntry], required_start: int, required_end: int, ) -> tuple[list[_NodeEntry], str]: """Greedy interval-cover biased toward fast, lightly-loaded nodes. Among nodes that equally advance coverage, prefer the one with higher effective throughput: benchmark_tokens_per_sec / (queue_depth + 1). Tiebreak: higher shard_end (fewer hops). """ 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: best = node elif node.shard_end > best.shard_end: best = node elif node.shard_end == best.shard_end and _effective_throughput(node) > _effective_throughput(best): 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_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]: """Return the memory pool used for shard-capacity planning.""" if node.vram_bytes > 0: return node.vram_bytes, "vram" if node.ram_bytes > 0: return node.ram_bytes, "ram" return DEFAULT_RAM_BYTES, "ram-default" 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 memory_budget_bytes, _ = _node_memory_budget_bytes(node) return int((memory_budget_bytes * 0.8) // layer_bytes) def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict: """Operator-facing capacity fields for inspection endpoints.""" memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node) summary = { "vram_bytes": node.vram_bytes, "ram_bytes": node.ram_bytes, "memory_budget_bytes": memory_budget_bytes, "memory_budget_source": memory_budget_source, "max_loaded_shards": node.max_loaded_shards, "quantizations": list(node.quantizations), "quantization": node.quantization, "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, "effective_throughput": round(_effective_throughput(node), 4), } if preset is not None: summary["max_assignable_layers"] = _node_layer_capacity(node, preset) return summary def _model_metadata_from_nodes(nodes: list[_NodeEntry]) -> dict: metadata: dict = {} for node in nodes: if node.model_metadata: metadata.update(node.model_metadata) if "num_layers" not in metadata: layers = [node.num_layers for node in nodes if node.num_layers is not None] if layers: metadata["num_layers"] = max(layers) return metadata 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, "capacity": _node_capacity_summary(node), } 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 _relay_http_request( relay_addr: str, path: str, body: bytes, headers: dict[str, str], timeout: float = 310.0, ) -> dict | None: """Send an HTTP-shaped request through a relay RPC WebSocket.""" try: import websockets.sync.client as wsc # type: ignore[import] request_id = str(uuid.uuid4()) with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws: ws.send(json.dumps({ "request_id": request_id, "method": "POST", "path": path, "headers": headers, "body": body.decode(errors="replace"), })) raw = ws.recv(timeout=timeout) response = json.loads(raw) if response.get("request_id") not in {None, request_id}: return None return response except Exception: return None 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_matches_model(node, 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 _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict: total_layers = max(node.num_layers or 0 for node in nodes) return { "layers_start": 0, "layers_end": total_layers - 1, "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, } def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None: model_nodes = [ node for node in server.registry.values() if node.hf_repo == hf_repo and node.num_layers is not None ] managed_nodes = [node for node in model_nodes if node.managed_assignment] if not model_nodes or not managed_nodes: return preset = _hf_rebalance_preset(model_nodes) required_start, required_end = _preset_layer_bounds(preset) total_layers = required_end - required_start + 1 if total_layers <= 0: 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, hf_repo, previous_start, previous_end, previous_quantization or _node_quantization(node, preset), ) ) node.pending_directives.append( _load_directive( node, hf_repo, 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) for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}): _rebalance_hf_model_locked(server, hf_repo) 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, relay_url: str | None = None, raft: "RaftNode | None" = None, gossip: "NodeGossip | None" = None, stats: "_StatsCollector | 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.relay_url = relay_url.rstrip("/") if relay_url else None self.raft = raft self.gossip = gossip self.stats: _StatsCollector | None = stats 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 if self.path == "/v1/stats/gossip": self._handle_stats_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/network/map": self._handle_network_map() 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/stats": self._handle_stats() 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, "metadata": dict(preset.get("metadata") or _model_metadata_from_nodes(model_nodes)), "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, "metadata": _model_metadata_from_nodes(model_nodes), "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, "relay_addr": node.relay_addr, "peer_id": node.peer_id, "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, } for node in tracker_nodes ], }) def _handle_network_map(self) -> None: server: _TrackerHTTPServer = self.server # type: ignore[assignment] with server.lock: self._purge_expired_nodes() nodes = list(server.registry.values()) def capacity_for(node: _NodeEntry) -> dict: preset = None if node.model: preset = server.model_presets.get(node.model) if preset is None and node.hf_repo and node.num_layers: preset = _hf_rebalance_preset([node]) return _node_capacity_summary(node, preset) self._send_json(200, { "relay_url": server.relay_url, "nodes": [ { "node_id": node.node_id, "endpoint": node.endpoint, "relay_addr": node.relay_addr, "peer_id": node.peer_id, "model": node.model, "hf_repo": node.hf_repo, "num_layers": node.num_layers, "model_metadata": dict(node.model_metadata), "shard_start": node.shard_start, "shard_end": node.shard_end, "tracker_mode": node.tracker_mode, "last_heartbeat": node.last_heartbeat, "capacity": capacity_for(node), "stats": _node_health(node, server.heartbeat_timeout), } for node in 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)) if model and server.stats is not None: server.stats.record_request(model) # 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 _node_matches_model(n, 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 _node_matches_model(n, 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 _node_matches_model(n, 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: hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1} if rn.relay_addr: hop["relay_addr"] = rn.relay_addr route_hops.append(hop) 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) relay_headers = { "Content-Type": "application/json", "X-Meshnet-Route": downstream_urls, **({"Authorization": auth} if auth else {}), } if node.relay_addr: print( f"[tracker] proxy via relay {request_id}: {node.relay_addr} " f"(direct endpoint {target_url})", flush=True, ) relayed = _relay_http_request( node.relay_addr, path="/v1/chat/completions", body=raw_body, headers=relay_headers, ) if relayed is not None: self._send_relayed_response(relayed) return print( f"[tracker] relay proxy failed {request_id}: {node.relay_addr}; " f"trying direct {target_url}", flush=True, ) 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: if node.relay_addr: print( f"[tracker] direct proxy failed {request_id}: {target_url}: {exc}; " f"relay already attempted for {node.relay_addr}", flush=True, ) else: 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 _send_relayed_response(self, response: dict) -> None: status = int(response.get("status", 503)) headers = response.get("headers") if isinstance(response.get("headers"), dict) else {} body_text = response.get("body") or "" body = body_text.encode() if isinstance(body_text, str) else bytes(body_text) self.send_response(status) self.send_header("Content-Type", headers.get("Content-Type", "application/json")) self.send_header("Content-Length", str(len(body))) self.end_headers() try: self.wfile.write(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)) max_loaded_shards = int(body.get("max_loaded_shards", 1)) 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, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"}) return if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 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)) managed_assignment = bool(body.get("managed_assignment", 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 model_metadata = body.get("model_metadata", {}) if model_metadata is None: model_metadata = {} if not isinstance(model_metadata, dict): self._send_json(400, {"error": "model_metadata must be an object"}) 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, max_loaded_shards=max_loaded_shards, benchmark_tokens_per_sec=benchmark_tokens_per_sec, quantization=quantization, managed_assignment=managed_assignment or not explicit_shard, tracker_mode=tracker_mode, hf_repo=hf_repo, num_layers=num_layers, model_metadata=model_metadata, 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: if entry.hf_repo: _rebalance_hf_model_locked(server, entry.hf_repo) else: _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 "" budget_bytes, budget_source = _node_memory_budget_bytes(entry) budget_gb = budget_bytes / (1024 ** 3) print( f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info} " f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}", 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"] if entry.hf_repo: _rebalance_hf_model_locked(server, entry.hf_repo) else: _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_stats(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] if server.stats is None: self._send_json(200, {"models": {}}) return self._send_json(200, {"models": server.stats.get_combined_stats()}) def _handle_stats_gossip(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] body = self._read_json_body() if body is None: return tracker_url = body.get("tracker_url", "") rpms = body.get("stats", {}) if server.stats is not None and tracker_url and isinstance(rpms, dict): server.stats.merge_peer_rpms(tracker_url, rpms) 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) ram_mb — integer system RAM in MB, used when vram_mb=0 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 try: ram_mb = int(params.get("ram_mb", ["0"])[0]) except ValueError: ram_mb = 0 max_layers = required_end - required_start + 1 memory_mb = vram_mb if vram_mb > 0 else ram_mb if memory_mb > 0: layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024) max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes))) elif 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) ram_mb — integer system RAM in MB, used when vram_mb=0 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 try: ram_mb = int(params.get("ram_mb", ["0"])[0]) except ValueError: ram_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 # Smart scoring: demand_rpm × coverage_deficit # coverage_deficit = uncovered_layers / total_layers (0 = fully covered) # demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet. demand_rpms: dict[str, float] = {} if server.stats is not None: local_rpms = server.stats.get_local_rpms() for repo in repo_groups: demand_rpms[repo] = local_rpms.get(repo, {}).get("rpm_last_hour", 0.0) best_repo = None best_score = -1.0 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) coverage_deficit = gap_size / max(total, 1) demand = demand_rpms.get(repo, 0.0) # +1.0 floor on demand so models with no recorded traffic still compete by coverage. score = (demand + 1.0) * (coverage_deficit + 0.01) if score > best_score or (score == best_score and gap_size > best_gap_size): best_score = score 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 covered — assign to highest-demand model for redundancy. best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0)) best_gap_start = 0 best_num_layers = repo_layers[best_repo] # Capacity: use the same 80%-of-memory rule as registered node planning. total_l = best_num_layers memory_mb = vram_mb if vram_mb > 0 else ram_mb if memory_mb > 0: max_layers = min( total_l, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // (30 * 1024 * 1024))), ) elif 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, "price_per_token": 0.0, }) 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_matches_model(node, 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, "relay_addr": e.relay_addr, "peer_id": e.peer_id, "wallet_address": e.wallet_address, "shard_start": e.shard_start, "shard_end": e.shard_end, "model": e.model, "hf_repo": e.hf_repo, "num_layers": e.num_layers, "model_metadata": dict(e.model_metadata), "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, "relay_addr": e.relay_addr, "peer_id": e.peer_id, "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, stats_db: str | None = None, relay_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._relay_url = relay_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._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector() self._stats_stop = threading.Event() self._stats_thread: threading.Thread | None = None self.port: int | None = None def start(self) -> int: if self._server is not None: raise RuntimeError("TrackerServer is already running") effective_relay_url = ( self._relay_url or derive_relay_url_from_public_tracker_url(self._cluster_self_url) ) # 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, relay_url=effective_relay_url, stats=self._stats, ) 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._stats_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() self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True) self._stats_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, model_metadata=payload.get("model_metadata") if isinstance(payload.get("model_metadata"), dict) 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 _stats_loop(self) -> None: """Periodically save stats to DB and push local slice to cluster peers.""" while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL): if self._stats is not None: self._stats.save_to_db() if self._stats is not None and self._cluster_peers: self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" local_rpms = self._stats.get_local_rpms() body = json.dumps({"tracker_url": self_url, "stats": local_rpms}).encode() for peer in self._cluster_peers: try: req = urllib.request.Request( f"{peer}/v1/stats/gossip", data=body, headers={"Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=2.0) as r: r.read() except Exception: pass 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._stats_stop.set() if self._stats is not None: self._stats.save_to_db() 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) if self._stats_thread is not None: self._stats_thread.join(timeout=1) self._server = None self._thread = None self._rebalance_thread = None self._stats_thread = None self._raft = None self._gossip = None self.port = None