neuron-tai/packages/tracker/meshnet_tracker/server.py

"""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 json
import threading
import time
import urllib.parse
import uuid
from typing import Any

from .gossip import NodeGossip
from .raft import RaftNode


DEFAULT_MODEL_PRESETS: dict[str, dict] = {
    "stub-model": {
        "layers_start": 0,
        "layers_end": 31,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
    },
    "openai-community/gpt2": {
        "layers_start": 0,
        "layers_end": 11,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
    },
}

DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024
DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024
DEFAULT_QUANTIZATIONS = ["bfloat16"]
DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0


class _NodeEntry:
    __slots__ = (
        "node_id", "endpoint", "shard_start", "shard_end",
        "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
        "score", "vram_bytes", "ram_bytes", "quantizations",
        "benchmark_tokens_per_sec", "quantization", "managed_assignment",
        "pending_directives", "last_heartbeat", "tracker_mode",
    )

    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,
    ) -> 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.pending_directives: list[dict] = []
        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(
        [node for node in nodes if node.shard_start is not None and node.shard_end is not None],
        key=lambda n: (n.shard_start, -n.shard_end),  # type: ignore[operator]
    )
    route: list[_NodeEntry] = []
    covered_up_to = required_start - 1

    while covered_up_to < required_end:
        best: _NodeEntry | None = None
        for node in candidates:
            if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
                if best is None or node.shard_end > best.shard_end:
                    best = node
        if best is None:
            missing = covered_up_to + 1
            return [], f"no route available: no registered node covers layer {missing}"
        route.append(best)
        covered_up_to = best.shard_end
        candidates = [n for n in candidates if n is not best]

    return route, ""


def _coverage_percentage(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> float:
    required_layers = required_end - required_start + 1
    if required_layers <= 0:
        return 0.0

    intervals = sorted(
        (
            (max(required_start, node.shard_start), min(required_end, node.shard_end))
            for node in nodes
            if node.shard_start is not None
            and node.shard_end is not None
            if node.shard_end >= required_start and node.shard_start <= required_end
        ),
        key=lambda interval: interval[0],
    )
    covered = 0
    covered_until = required_start - 1
    for start, end in intervals:
        if end <= covered_until:
            continue
        next_start = max(start, covered_until + 1)
        if next_start > end:
            continue
        covered += end - next_start + 1
        covered_until = end
    return round((covered / required_layers) * 100, 2)


def _preset_layer_bounds(preset: dict) -> tuple[int, int]:
    start = int(preset.get("layers_start", 0))
    if "layers_end" in preset:
        return start, int(preset["layers_end"])
    return start, start + int(preset["total_layers"]) - 1


def _preset_bytes_per_layer(preset: dict) -> dict[str, int]:
    raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {}))
    if isinstance(raw, dict) and raw:
        return {str(quant): int(value) for quant, value in raw.items()}
    return {"bfloat16": 30 * 1024 * 1024}


def _node_quantization(node: _NodeEntry, preset: dict) -> str:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    if node.quantization in bytes_per_layer:
        return node.quantization
    for quantization in node.quantizations:
        if quantization in bytes_per_layer:
            return quantization
    return next(iter(bytes_per_layer))


def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    quantization = _node_quantization(node, preset)
    layer_bytes = bytes_per_layer[quantization]
    if layer_bytes <= 0:
        return 0
    return int((node.vram_bytes * 0.8) // layer_bytes)


def _coverage_map(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> list[dict]:
    layer_counts = []
    for layer in range(required_start, required_end + 1):
        count = 0
        for node in nodes:
            if node.shard_start is None or node.shard_end is None:
                continue
            if node.shard_start <= layer <= node.shard_end:
                count += 1
        layer_counts.append((layer, count))

    coverage: list[dict] = []
    for layer, count in layer_counts:
        if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1:
            coverage[-1]["end_layer"] = layer
        else:
            coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count})
    return coverage


def _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 _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[:8]}  {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


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,
        raft: "RaftNode | None" = None,
        gossip: "NodeGossip | None" = None,
    ) -> None:
        super().__init__(addr, handler)
        self.registry = registry
        self.lock = lock
        self.heartbeat_timeout = heartbeat_timeout
        self.model_presets = model_presets
        self.contracts = contracts
        self.minimum_stake = minimum_stake
        self.raft = raft
        self.gossip = gossip


class _TrackerHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def _send_json(self, status: int, data: dict) -> None:
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        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]
        _purge_expired_nodes_locked(server)

    def do_POST(self):
        if self.path == "/v1/chat/completions":
            self._handle_proxy_chat()
            return
        if self.path == "/v1/nodes/register":
            self._handle_register()
            return
        if self.path == "/v1/raft/vote":
            self._handle_raft_vote()
            return
        if self.path == "/v1/raft/append":
            self._handle_raft_append()
            return
        if self.path == "/v1/gossip":
            self._handle_gossip()
            return
        parts = self.path.split("/")
        # /v1/nodes/<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/models":
            self._handle_models()
        elif parsed.path.startswith("/v1/coverage/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/"))
            self._handle_coverage(model)
        elif parsed.path.startswith("/v1/tracker-nodes/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/"))
            self._handle_tracker_nodes(model)
        elif parsed.path == "/v1/raft/status":
            self._handle_raft_status()
        elif parsed.path == "/v1/health":
            self._send_json(200, {"status": "ok"})
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_models(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        with server.lock:
            self._purge_expired_nodes()
            alive = list(server.registry.values())
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        data = []
        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,
            )
            data.append({
                "id": name,
                "object": "model",
                "created": created,
                "owned_by": "meshnet",
                "shard_coverage_percentage": coverage,
            })
        self._send_json(200, {"object": "list", "data": data})

    def _handle_coverage(self, model: str):
        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, 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
            ]
        coverage = _coverage_map(alive, required_start, required_end)
        self._send_json(200, {"model": model, "coverage": coverage})

    def _handle_tracker_nodes(self, model: str):
        """Return nodes registered with tracker_mode=True whose shard starts at layer 0."""
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return
        required_start, _ = _preset_layer_bounds(preset)
        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        tracker_nodes = [
            node for node in alive
            if node.shard_start is not None
            and node.shard_start == required_start
            and node.tracker_mode
        ]
        self._send_json(200, {
            "model": model,
            "tracker_nodes": [
                {
                    "node_id": node.node_id,
                    "endpoint": node.endpoint,
                    "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
                }
                for node in tracker_nodes
            ],
        })

    # ---------------------------------------------------------------- OpenAI proxy

    def _handle_proxy_chat(self) -> None:
        """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node.

        Picks a live tracker-mode node for the requested model using round-robin,
        then forwards the request verbatim and relays the response (including
        streaming SSE chunks) back to the caller.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]

        length = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(length) if length else b"{}"

        try:
            body = json.loads(raw_body)
        except json.JSONDecodeError:
            self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}})
            return

        model: str = body.get("model", "")
        is_stream: bool = bool(body.get("stream", False))

        # Find a live tracker-mode node for this model
        with server.lock:
            self._purge_expired_nodes()
            candidates = [
                n for n in server.registry.values()
                if n.tracker_mode and (n.model == model or n.hf_repo == model)
            ]

        if not candidates:
            # Fall back: any node serving shard_start=0 for this model
            with server.lock:
                candidates = [
                    n for n in server.registry.values()
                    if n.shard_start == 0 and (n.model == model or n.hf_repo == model)
                ]

        if not candidates:
            self._send_json(503, {"error": {
                "message": f"no nodes available for model {model!r}",
                "type": "service_unavailable",
                "code": "model_not_available",
            }})
            return

        # Simple round-robin via list length modulo (stateless, good enough)
        node = candidates[int(time.time() * 1000) % len(candidates)]
        target_url = f"{node.endpoint}/v1/chat/completions"

        req = urllib.request.Request(
            target_url,
            data=raw_body,
            headers={"Content-Type": "application/json"},
            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)
        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:
            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()
                self.send_response(200)
                self.send_header("Content-Type", content_type)
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                try:
                    self.wfile.write(resp_body)
                except BrokenPipeError:
                    pass

    def _handle_register(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return

        # --- Raft cluster mode: forward to leader or propose via Raft ---
        if server.raft is not None:
            if not server.raft.is_leader:
                leader = server.raft.leader()
                if leader is None:
                    self._send_json(503, {"error": "no leader elected — retry in a moment"})
                    return
                # Proxy to leader
                try:
                    data = json.dumps(body).encode()
                    req = urllib.request.Request(
                        f"{leader}/v1/nodes/register",
                        data=data,
                        headers={"Content-Type": "application/json"},
                        method="POST",
                    )
                    with urllib.request.urlopen(req, timeout=5.0) as r:
                        resp_body = r.read()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(resp_body)))
                    self.end_headers()
                    self.wfile.write(resp_body)
                except Exception as exc:
                    self._send_json(503, {"error": f"leader proxy failed: {exc}"})
                return
            # Leader path: fall through to normal registration, then replicate via Raft.
            # We let the registration run first (to generate node_id), then asynchronously
            # propose to the Raft log so followers can replicate the entry.
            # Raft proposal happens after the response is sent (fire-and-forget replication).

        endpoint = body.get("endpoint")
        if not isinstance(endpoint, str) or not endpoint:
            self._send_json(400, {"error": "endpoint is required"})
            return
        parsed_endpoint = urllib.parse.urlparse(endpoint)
        if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc:
            self._send_json(400, {"error": "endpoint must be an http(s) URL"})
            return

        shard_start: int | None
        shard_end: int | None
        explicit_shard = "shard_start" in body or "shard_end" in body
        if explicit_shard:
            try:
                shard_start = int(body["shard_start"])
                shard_end = int(body["shard_end"])
            except (KeyError, TypeError, ValueError):
                self._send_json(400, {"error": "shard_start and shard_end must be numeric"})
                return
            if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
                self._send_json(400, {"error": "shard range must be non-negative and ordered"})
                return
        else:
            shard_start = None
            shard_end = None
        try:
            score = float(body.get("score", 1.0))
        except (TypeError, ValueError):
            self._send_json(400, {"error": "score must be numeric"})
            return

        hardware_profile = body.get("hardware_profile", {})
        if not isinstance(hardware_profile, dict):
            self._send_json(400, {"error": "hardware_profile must be an object"})
            return
        model = body.get("model")
        if model is None:
            model = "stub-model"
        if not isinstance(model, str):
            self._send_json(400, {"error": "model must be a string"})
            return
        shard_checksum = body.get("shard_checksum")
        if shard_checksum is not None and not isinstance(shard_checksum, str):
            self._send_json(400, {"error": "shard_checksum must be a string"})
            return
        try:
            vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES))
            ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES))
            benchmark_tokens_per_sec = float(
                body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC)
            )
        except (TypeError, ValueError):
            self._send_json(400, {"error": "vram_bytes, ram_bytes, and benchmark_tokens_per_sec must be numeric"})
            return
        if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0:
            self._send_json(400, {"error": "capability values must be positive"})
            return
        quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS)
        if not (
            isinstance(quantizations_body, list)
            and quantizations_body
            and all(isinstance(item, str) and item for item in quantizations_body)
        ):
            self._send_json(400, {"error": "quantizations must be a non-empty string array"})
            return
        quantizations = list(quantizations_body)
        quantization = body.get("quantization")
        if quantization is not None and not isinstance(quantization, str):
            self._send_json(400, {"error": "quantization must be a string"})
            return
        wallet_address = body.get("wallet_address")
        if wallet_address is not None and not isinstance(wallet_address, str):
            self._send_json(400, {"error": "wallet_address must be a string"})
            return
        ban_error = _registration_ban_error(server.contracts, wallet_address)
        if ban_error:
            self._send_json(403, {"error": ban_error})
            return

        tracker_mode = bool(body.get("tracker_mode", False))
        hf_repo = body.get("hf_repo")
        if hf_repo is not None and not isinstance(hf_repo, str):
            self._send_json(400, {"error": "hf_repo must be a string"})
            return
        num_layers_body = body.get("num_layers")
        num_layers: int | None = None
        if num_layers_body is not None:
            try:
                num_layers = int(num_layers_body)
            except (TypeError, ValueError):
                self._send_json(400, {"error": "num_layers must be an integer"})
                return

        node_id = str(uuid.uuid4())
        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,
        )
        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[:8]} with {node_id[:8]}"
                    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[:8]}  {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]
        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()
            _rebalance_model_locked(server, entry.model or "stub-model")
            directives = list(entry.pending_directives)
            entry.pending_directives.clear()
        if server.gossip is not None:
            server.gossip.record(node_id)
        if directives:
            self._send_json(200, {"directives": directives})
        else:
            self._send_json(200, {})

    # ---------------------------------------------------------------- Raft handlers

    def _handle_raft_vote(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_request_vote(body)
        self._send_json(200, result)

    def _handle_raft_append(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_append_entries(body)
        self._send_json(200, result)

    def _handle_raft_status(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        if server.raft is None:
            self._send_json(200, {"role": "standalone", "term": 0, "leader": None})
            return
        self._send_json(200, server.raft.status())

    def _handle_gossip(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.gossip is not None and isinstance(body, dict):
            server.gossip.merge({k: float(v) for k, v in body.items()})
        self._send_json(200, {})

    def _handle_assign(self, parsed: urllib.parse.ParseResult):
        """Return an optimal shard assignment for a node given its hardware profile.

        Query params:
          model   — model preset name (default: first preset)
          device  — "cuda" | "cpu"
          vram_mb — integer VRAM in MB (0 for CPU)

        The greedy strategy: find the first gap in current layer coverage
        and assign it.  If no gap exists, assign the full model range so the
        node provides redundant coverage.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)

        model_list = params.get("model")
        if not model_list:
            model = next(iter(server.model_presets), None)
            if model is None:
                self._send_json(503, {"error": "no model presets configured"})
                return
        else:
            model = model_list[0]

        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start, required_end = _preset_layer_bounds(preset)

        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        device = params.get("device", ["cpu"])[0]
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        max_layers = required_end - required_start + 1
        if device != "cuda" or vram_mb < 8192:
            max_layers = min(max_layers, 16)

        # Collect covered intervals sorted by start layer.
        covered = sorted(
            [
                (n.shard_start, n.shard_end)
                for n in alive
                if n.shard_start is not None and n.shard_end is not None
            ],
            key=lambda t: t[0],
        )

        # Walk from required_start to find the first uncovered layer.
        gap_start = required_start
        for s, e in covered:
            if s <= gap_start:
                gap_start = max(gap_start, e + 1)
            else:
                break  # gap found before this node

        if gap_start > required_end:
            # Full coverage exists — assign the full range for redundancy.
            shard_start = required_start
            shard_end = min(required_end, shard_start + max_layers - 1)
        else:
            shard_start = gap_start
            shard_end = min(required_end, shard_start + max_layers - 1)

        peers = [
            {"endpoint": node.endpoint, "checksum": node.shard_checksum}
            for node in alive
            if node.model == model
            and node.shard_start == shard_start
            and node.shard_end == shard_end
            and node.shard_checksum
        ]

        self._send_json(200, {
            "shard_start": shard_start,
            "shard_end": shard_end,
            "model": model,
            "model_layers_end": required_end,
            "peers": peers,
            **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
        })

    def _handle_network_assign(self, parsed: urllib.parse.ParseResult):
        """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes.

        Query params:
          vram_mb  — integer VRAM in MB (0 = CPU-only node)
          device   — "cuda" | "cpu"
          hf_repo  — optional; if set, restrict search to this repo only

        Returns:
          {hf_repo, shard_start, shard_end, num_layers, gap_found}
          gap_found=true means a real uncovered gap was assigned; false means redundancy.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        device = params.get("device", ["cpu"])[0]
        filter_repo = params.get("hf_repo", [None])[0]  # optional repo filter

        with server.lock:
            self._purge_expired_nodes()
            all_nodes = list(server.registry.values())

        # Collect only nodes that registered a real HF model (have hf_repo + shard bounds).
        hf_nodes = [
            n for n in all_nodes
            if n.hf_repo
            and n.shard_start is not None
            and n.shard_end is not None
            and n.num_layers is not None
            and (filter_repo is None or n.hf_repo == filter_repo)
        ]

        if not hf_nodes:
            msg = (
                f"no HF-model nodes registered for {filter_repo!r}"
                if filter_repo
                else "no HF-model nodes registered; cannot assign shards"
            )
            self._send_json(503, {"error": msg})
            return

        # Group by hf_repo; pick the one with the largest total_layers and biggest gap.
        from collections import defaultdict
        repo_groups: dict = defaultdict(list)
        repo_layers: dict = {}
        for n in hf_nodes:
            repo_groups[n.hf_repo].append(n)
            # Use the largest num_layers seen for this repo.
            if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]:
                repo_layers[n.hf_repo] = n.num_layers

        # Pick the repo where the gap is largest (most work to do).
        best_repo = None
        best_gap_size = -1
        best_gap_start = 0
        best_num_layers = 0

        for repo, nodes in repo_groups.items():
            total = repo_layers[repo]
            covered = sorted(
                [(n.shard_start, n.shard_end) for n in nodes],
                key=lambda t: t[0],
            )
            # Walk from 0 to find first uncovered layer.
            gap_start = 0
            for s, e in covered:
                if s <= gap_start:
                    gap_start = max(gap_start, e + 1)
                else:
                    break
            gap_size = max(0, (total - 1) - gap_start + 1)  # layers remaining uncovered
            if gap_size > best_gap_size:
                best_gap_size = gap_size
                best_gap_start = gap_start
                best_repo = repo
                best_num_layers = total

        gap_found = best_gap_size > 0
        if not gap_found:
            # All shards are covered — still assign to the model with most nodes for redundancy.
            best_repo = max(repo_groups, key=lambda r: len(repo_groups[r]))
            best_gap_start = 0
            best_num_layers = repo_layers[best_repo]

        # Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM.
        total_l = best_num_layers
        if device == "cuda" and vram_mb >= 8192:
            max_layers = total_l
        else:
            max_layers = max(1, total_l // 2)

        shard_start = best_gap_start
        shard_end = min(total_l - 1, shard_start + max_layers - 1)

        self._send_json(200, {
            "hf_repo": best_repo,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "num_layers": total_l,
            "gap_found": gap_found,
        })

    def _handle_route(self, parsed: urllib.parse.ParseResult):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model_list = params.get("model")
        if not model_list:
            self._send_json(400, {"error": "missing 'model' query parameter"})
            return

        model = model_list[0]
        preset = server.model_presets.get(model)

        with server.lock:
            self._purge_expired_nodes()
            if preset is not None:
                # Preset-based routing (stub-model system).
                alive = [node for node in server.registry.values() if node.model == model]
                required_start, required_end = _preset_layer_bounds(preset)
            else:
                # HF model routing: match by hf_repo (full) or model short name.
                alive = [
                    node for node in server.registry.values()
                    if (node.hf_repo == model or node.model == model)
                    and node.shard_start is not None
                    and node.shard_end is not None
                    and node.num_layers is not None
                ]
                if not alive:
                    self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
                    return
                required_start = 0
                required_end = max(n.num_layers for n in alive) - 1  # type: ignore[type-var]

        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        route, error = _select_route(alive, required_start, required_end)
        if error:
            self._send_json(503, {"error": error})
            return

        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,
                    "model": e.model,
                    "shard_checksum": e.shard_checksum,
                    "score": e.score,
                }
                for e in route
            ],
        })

    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": [
                    {
                        "endpoint": e.endpoint,
                        "wallet_address": e.wallet_address,
                        "shard_start": e.shard_start,
                        "shard_end": e.shard_end,
                        "model": e.model,
                        "shard_checksum": e.shard_checksum,
                        "score": e.score,
                    }
                    for e in route
                ],
            })

        self._send_json(200, {"routes": routes})


class TrackerServer:
    """HTTP tracker that manages node registration and resolves inference routes.

    Nodes register via POST /v1/nodes/register and keep themselves alive with
    POST /v1/nodes/<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,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._heartbeat_timeout = heartbeat_timeout
        self._rebalance_interval = rebalance_interval
        self._model_presets: dict = (
            model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS)
        )
        self._contracts = contracts
        self._minimum_stake = minimum_stake
        self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else []
        self._cluster_self_url: str | None = cluster_self_url
        self._registry: dict[str, _NodeEntry] = {}
        self._lock = threading.Lock()
        self._server: _TrackerHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self._rebalance_stop = threading.Event()
        self._rebalance_thread: threading.Thread | None = None
        self._raft: RaftNode | None = None
        self._gossip: NodeGossip | None = None
        self.port: int | None = None

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("TrackerServer is already running")

        # Start HTTP server first so we know our port
        self._server = _TrackerHTTPServer(
            (self._host, self._requested_port),
            _TrackerHandler,
            self._registry,
            self._lock,
            self._heartbeat_timeout,
            self._model_presets,
            self._contracts,
            self._minimum_stake,
        )
        self.port = self._server.server_address[1]

        # Start Raft + gossip if cluster peers are configured
        if self._cluster_peers:
            self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
            self._raft = RaftNode(
                self_url=self_url,
                peers=self._cluster_peers,
                apply_fn=self._raft_apply,
            )
            self._gossip = NodeGossip(peers=self._cluster_peers)
            self._server.raft = self._raft
            self._server.gossip = self._gossip
            self._raft.start()
            self._gossip.start()

        self._rebalance_stop.clear()
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True)
        self._rebalance_thread.start()
        return self.port

    def _raft_apply(self, command: str, payload: dict) -> None:
        """Called by RaftNode when a log entry is committed — replicate to local registry."""
        if command != "register":
            return
        # Re-apply the registration to our local registry (follower path).
        # On the leader this is a no-op since it already registered locally.
        node_id = payload.get("node_id")
        if not node_id or node_id in self._registry:
            return  # already present (leader case) or malformed
        endpoint = payload.get("endpoint", "")
        try:
            shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None
            shard_end   = int(payload["shard_end"])   if payload.get("shard_end")   is not None else None
        except (TypeError, ValueError):
            return
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            model=payload.get("model", "stub-model"),
            shard_checksum=payload.get("shard_checksum"),
            hardware_profile=payload.get("hardware_profile", {}),
            wallet_address=payload.get("wallet_address"),
            score=float(payload.get("score", 1.0)),
            tracker_mode=bool(payload.get("tracker_mode", False)),
            hf_repo=payload.get("hf_repo"),
            num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None,
        )
        with self._lock:
            self._registry[node_id] = entry

    def _rebalance_loop(self) -> None:
        while not self._rebalance_stop.wait(self._rebalance_interval):
            server = self._server
            if server is None:
                return
            with self._lock:
                _purge_expired_nodes_locked(server)
                _rebalance_all_locked(server)

    def stop(self) -> None:
        if self._raft is not None:
            self._raft.stop()
        if self._gossip is not None:
            self._gossip.stop()
        if self._server is None:
            return
        self._rebalance_stop.set()
        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        if self._rebalance_thread is not None:
            self._rebalance_thread.join(timeout=1)
        self._server = None
        self._thread = None
        self._rebalance_thread = None
        self._raft = None
        self._gossip = None
        self.port = None