2040 lines
79 KiB
Python
2040 lines
79 KiB
Python
"""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=<preset>
|
||
Response 200: {"route": ["http://node-a", "http://node-b"]}
|
||
Response 400/404/503: {"error": str}
|
||
- GET /v1/routes?model=<preset>&redundancy=<n>
|
||
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
|
||
|
||
|
||
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 _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", "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 _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_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 _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.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,
|
||
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/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', '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,
|
||
"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,
|
||
"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())
|
||
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,
|
||
"shard_start": node.shard_start,
|
||
"shard_end": node.shard_end,
|
||
"tracker_mode": node.tracker_mode,
|
||
"last_heartbeat": node.last_heartbeat,
|
||
}
|
||
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 (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 '<empty>'}",
|
||
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:
|
||
if node.relay_addr:
|
||
print(
|
||
f"[tracker] direct proxy failed {request_id}: {target_url}: {exc}; "
|
||
f"trying relay {node.relay_addr}",
|
||
flush=True,
|
||
)
|
||
relayed = _relay_http_request(
|
||
node.relay_addr,
|
||
path="/v1/chat/completions",
|
||
body=raw_body,
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"X-Meshnet-Route": downstream_urls,
|
||
**({"Authorization": auth} if auth else {}),
|
||
},
|
||
)
|
||
if relayed is not None:
|
||
self._send_relayed_response(relayed)
|
||
return
|
||
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))
|
||
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_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)
|
||
|
||
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
|
||
|
||
# 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: 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,
|
||
"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.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,
|
||
"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, 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/<id>/heartbeat. The gateway queries GET /v1/route?model=<name>
|
||
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")
|
||
|
||
# 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=self._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,
|
||
)
|
||
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
|