421 lines
15 KiB
Python
421 lines
15 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,
|
|
"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}
|
|
"""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
import uuid
|
|
from typing import Any
|
|
|
|
|
|
DEFAULT_MODEL_PRESETS: dict[str, dict] = {
|
|
"stub-model": {"layers_start": 0, "layers_end": 31},
|
|
}
|
|
|
|
|
|
class _NodeEntry:
|
|
__slots__ = (
|
|
"node_id", "endpoint", "shard_start", "shard_end",
|
|
"hardware_profile", "wallet_address", "score", "last_heartbeat",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
node_id: str,
|
|
endpoint: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
hardware_profile: dict,
|
|
wallet_address: str | None,
|
|
score: float,
|
|
) -> None:
|
|
self.node_id = node_id
|
|
self.endpoint = endpoint
|
|
self.shard_start = shard_start
|
|
self.shard_end = shard_end
|
|
self.hardware_profile = hardware_profile
|
|
self.wallet_address = wallet_address
|
|
self.score = score
|
|
self.last_heartbeat: float = time.monotonic()
|
|
|
|
|
|
def _select_route(
|
|
nodes: list[_NodeEntry],
|
|
required_start: int,
|
|
required_end: int,
|
|
) -> tuple[list[_NodeEntry], str]:
|
|
"""Greedy interval-cover. Returns (ordered route, error_message)."""
|
|
candidates = sorted(nodes, key=lambda n: (n.shard_start, -n.shard_end))
|
|
route: list[_NodeEntry] = []
|
|
covered_up_to = required_start - 1
|
|
|
|
while covered_up_to < required_end:
|
|
best: _NodeEntry | None = None
|
|
for node in candidates:
|
|
if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
|
|
if best is None or node.shard_end > best.shard_end:
|
|
best = node
|
|
if best is None:
|
|
missing = covered_up_to + 1
|
|
return [], f"no route available: no registered node covers layer {missing}"
|
|
route.append(best)
|
|
covered_up_to = best.shard_end
|
|
candidates = [n for n in candidates if n is not best]
|
|
|
|
return route, ""
|
|
|
|
|
|
class _TrackerHTTPServer(http.server.HTTPServer):
|
|
def __init__(
|
|
self,
|
|
addr: tuple,
|
|
handler,
|
|
registry: dict,
|
|
lock: threading.Lock,
|
|
heartbeat_timeout: float,
|
|
model_presets: dict,
|
|
contracts: Any | None,
|
|
minimum_stake: int,
|
|
) -> 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
|
|
|
|
|
|
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()
|
|
self.wfile.write(body)
|
|
|
|
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]
|
|
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:
|
|
del server.registry[node_id]
|
|
|
|
def do_POST(self):
|
|
if self.path == "/v1/nodes/register":
|
|
self._handle_register()
|
|
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/nodes/assign":
|
|
self._handle_assign(parsed)
|
|
elif parsed.path == "/v1/models":
|
|
self._handle_models()
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def _handle_models(self):
|
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
|
created = int(time.time())
|
|
data = [
|
|
{"id": name, "object": "model", "created": created, "owned_by": "meshnet"}
|
|
for name in server.model_presets
|
|
]
|
|
self._send_json(200, {"object": "list", "data": data})
|
|
|
|
def _handle_register(self):
|
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
|
|
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
|
|
|
|
try:
|
|
shard_start = int(body["shard_start"])
|
|
shard_end = int(body["shard_end"])
|
|
score = float(body.get("score", 1.0))
|
|
except (KeyError, TypeError, ValueError):
|
|
self._send_json(400, {"error": "shard_start, shard_end, and score 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
|
|
|
|
hardware_profile = body.get("hardware_profile", {})
|
|
if not isinstance(hardware_profile, dict):
|
|
self._send_json(400, {"error": "hardware_profile must be an object"})
|
|
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
|
|
|
|
node_id = str(uuid.uuid4())
|
|
entry = _NodeEntry(
|
|
node_id=node_id,
|
|
endpoint=endpoint.rstrip("/"),
|
|
shard_start=shard_start,
|
|
shard_end=shard_end,
|
|
hardware_profile=hardware_profile,
|
|
wallet_address=wallet_address,
|
|
score=score,
|
|
)
|
|
with server.lock:
|
|
self._purge_expired_nodes()
|
|
server.registry[node_id] = entry
|
|
|
|
self._send_json(200, {"node_id": node_id})
|
|
|
|
def _handle_heartbeat(self, node_id: str):
|
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
|
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()
|
|
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: int = preset["layers_start"]
|
|
required_end: int = preset["layers_end"]
|
|
|
|
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
|
|
]
|
|
|
|
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],
|
|
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)
|
|
|
|
self._send_json(200, {
|
|
"shard_start": shard_start,
|
|
"shard_end": shard_end,
|
|
"model": model,
|
|
"model_layers_end": required_end,
|
|
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
|
})
|
|
|
|
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)
|
|
if preset is None:
|
|
self._send_json(404, {"error": f"unknown model preset: {model!r}"})
|
|
return
|
|
|
|
required_start: int = preset["layers_start"]
|
|
required_end: int = preset["layers_end"]
|
|
|
|
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
|
|
]
|
|
|
|
route, error = _select_route(alive, required_start, required_end)
|
|
if error:
|
|
self._send_json(503, {"error": error})
|
|
return
|
|
|
|
self._send_json(200, {
|
|
"route": [e.endpoint for e in route],
|
|
"nodes": [
|
|
{
|
|
"endpoint": e.endpoint,
|
|
"wallet_address": e.wallet_address,
|
|
"shard_start": e.shard_start,
|
|
"shard_end": e.shard_end,
|
|
"score": e.score,
|
|
}
|
|
for e in route
|
|
],
|
|
})
|
|
|
|
|
|
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 = 30.0,
|
|
model_presets: dict | None = None,
|
|
contracts: Any | None = None,
|
|
minimum_stake: int = 0,
|
|
) -> None:
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._heartbeat_timeout = heartbeat_timeout
|
|
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._registry: dict[str, _NodeEntry] = {}
|
|
self._lock = threading.Lock()
|
|
self._server: _TrackerHTTPServer | None = None
|
|
self._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")
|
|
self._server = _TrackerHTTPServer(
|
|
(self._host, self._requested_port),
|
|
_TrackerHandler,
|
|
self._registry,
|
|
self._lock,
|
|
self._heartbeat_timeout,
|
|
self._model_presets,
|
|
self._contracts,
|
|
self._minimum_stake,
|
|
)
|
|
self.port = self._server.server_address[1]
|
|
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
|
self._thread.start()
|
|
return self.port
|
|
|
|
def stop(self) -> None:
|
|
if self._server is None:
|
|
return
|
|
self._server.shutdown()
|
|
self._server.server_close()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=1)
|
|
self._server = None
|
|
self._thread = None
|
|
self.port = None
|