routing improvements - dynamic (wip)

packages/tracker/meshnet_tracker/server.py

@@ -19,6 +19,8 @@ HTTP API contract:
 - GET /v1/routes?model=<preset>&redundancy=<n>
   Response 200: {"routes": [{"route": [...], "nodes": [...]}]}
   Response 400/404/503: {"error": str}
+- GET /v1/routing?model=<name>  (ADR-0021 learned route table)
+  Response 200: {"config": {...}, "models": {model: {"epoch": int, "routes": [...]}}}
 """
 
 import http.server
@@ -26,6 +28,7 @@ import hashlib
 import itertools
 import json
 import os
+import random
 import socketserver
 import sqlite3
 import tarfile
@@ -47,6 +50,14 @@ from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger
 from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
 from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price
 from .gossip import NodeGossip
+from .routing_stats import (
+    RouteCandidate,
+    RouteStatsStore,
+    RoutingConfig,
+    choose_route,
+    route_signature,
+    route_table,
+)
 from .model_files import files_for_layer_range, snapshot_dir_for_repo
 from .raft import RaftNode
 
@@ -738,6 +749,64 @@ def _select_route(
     return route, ""
 
 
+def _enumerate_routes(
+    nodes: list["_NodeEntry"],
+    required_start: int,
+    required_end: int,
+    model: str | None = None,
+    contracts: Any | None = None,
+    max_candidates: int = 8,
+) -> list["RouteCandidate"]:
+    """Enumerate viable route candidates for bandit selection (ADR-0021).
+
+    One candidate per distinct head (a node that can embed the prompt, i.e.
+    covers `required_start` from layer 0 of the range), each greedily completed
+    with the longest-advancing hops. The route's prior throughput estimate is
+    its bottleneck hop's queue-adjusted effective throughput — used only until
+    observed route samples exist.
+    """
+    sharded = [
+        n for n in nodes
+        if n.shard_start is not None and n.shard_end is not None
+    ]
+    # Heads must start the pipeline at the first required layer (they tokenize
+    # and embed the prompt — same condition as tracker_mode registration).
+    heads = [n for n in sharded if n.shard_start == required_start]
+    candidates: dict[str, RouteCandidate] = {}
+    for head in heads:
+        route = [head]
+        covered_up_to = head.shard_end
+        pool = [n for n in sharded if n is not head]
+        while covered_up_to < required_end:
+            best = None
+            for n in pool:
+                if n.shard_start <= covered_up_to + 1 and n.shard_end > covered_up_to:
+                    if best is None or n.shard_end > best.shard_end or (
+                        n.shard_end == best.shard_end
+                        and _effective_throughput(n, model) * _reputation_multiplier(n, contracts)
+                        > _effective_throughput(best, model) * _reputation_multiplier(best, contracts)
+                    ):
+                        best = n
+            if best is None:
+                route = []
+                break
+            route.append(best)
+            covered_up_to = best.shard_end
+            pool = [n for n in pool if n is not best]
+        if not route:
+            continue
+        signature = route_signature(model or "?", route)
+        if signature in candidates:
+            continue
+        prior = min(
+            _effective_throughput(n, model) * _reputation_multiplier(n, contracts)
+            for n in route
+        )
+        candidates[signature] = RouteCandidate(nodes=route, signature=signature, prior_tps=prior)
+    ranked = sorted(candidates.values(), key=lambda c: -c.prior_tps)
+    return ranked[:max_candidates]
+
+
 def _node_route_summary(nodes: list["_NodeEntry"]) -> list[dict]:
     return [
         {
@@ -1720,6 +1789,14 @@ def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantiza
     }
 
 
+def _route_stats_keys(server: "_TrackerHTTPServer", entry: "_NodeEntry") -> list[str]:
+    """All stats keys a node's routes may be recorded under (model, repo, resolved preset)."""
+    keys = {entry.model, entry.hf_repo}
+    resolved, _ = _resolve_model_preset(server.model_presets, entry.hf_repo or entry.model)
+    keys.add(resolved)
+    return [key for key in keys if key]
+
+
 def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
     now = time.monotonic()
     expired_ids = [
@@ -1732,6 +1809,11 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
         expired_entries.append((node_id, entry))
     if expired_ids:
         _rebalance_all_locked(server)
+        server.route_stats.bump_epoch(
+            key
+            for _, entry in expired_entries
+            for key in _route_stats_keys(server, entry)
+        )
     for node_id, entry in expired_entries:
         _tracker_log(
             server,
@@ -2289,6 +2371,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
         toploc_backend: Any | None = None,
         hf_pricing_log: "HfPricingLog | None" = None,
         models_dir: Path | None = None,
+        route_stats: "RouteStatsStore | None" = None,
     ) -> None:
         super().__init__(addr, handler)
         self.registry = registry
@@ -2324,6 +2407,8 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
         self.console_lock = threading.Lock()
         self.active_proxies: dict[str, _ActiveProxyContext] = {}
         self.active_proxies_lock = threading.Lock()
+        self.route_stats: RouteStatsStore = route_stats or RouteStatsStore()
+        self.route_rng = random.Random()
 
 
 class _TrackerHandler(http.server.BaseHTTPRequestHandler):
@@ -2499,6 +2584,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             self._handle_route(parsed)
         elif parsed.path == "/v1/routes":
             self._handle_routes(parsed)
+        elif parsed.path == "/v1/routing":
+            self._handle_routing(parsed)
         elif parsed.path == "/v1/nodes/assign":
             self._handle_assign(parsed)
         elif parsed.path == "/v1/network/assign":
@@ -2964,25 +3051,50 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
         if pinned_nodes is not None:
             route_nodes = pinned_nodes
+            routing_decision = {"mode": "pinned"}
         else:
-            route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts)
-            if route_error:
-                _tracker_log(
-                    server,
-                    "warn",
-                    "route unavailable",
-                    model=model,
-                    route_model=route_model,
-                    error=route_error,
-                    candidate_count=len(all_nodes),
-                    candidates=_node_route_summary(all_nodes),
-                )
-                self._send_json(503, {"error": {
-                    "message": route_error,
-                    "type": "service_unavailable",
-                    "code": "route_not_available",
-                }})
-                return
+            # ADR-0021: enumerate viable routes and pick one bandit-style —
+            # ε-scout among unproven routes, otherwise weighted ∝ observed tps^α.
+            route_candidates = _enumerate_routes(
+                all_nodes, rs, re,
+                model=route_model,
+                contracts=server.contracts,
+                max_candidates=server.route_stats.config.max_candidate_routes,
+            )
+            picked, routing_decision = choose_route(
+                route_candidates, server.route_stats, route_model, rng=server.route_rng,
+            )
+            if picked is not None:
+                route_nodes = picked.nodes
+            else:
+                # No head-anchored candidate — legacy greedy cover as fallback
+                # (also produces the layer-gap error message).
+                route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts)
+                routing_decision = {"mode": "greedy-fallback"}
+                if route_error:
+                    _tracker_log(
+                        server,
+                        "warn",
+                        "route unavailable",
+                        model=model,
+                        route_model=route_model,
+                        error=route_error,
+                        candidate_count=len(all_nodes),
+                        candidates=_node_route_summary(all_nodes),
+                    )
+                    self._send_json(503, {"error": {
+                        "message": route_error,
+                        "type": "service_unavailable",
+                        "code": "route_not_available",
+                    }})
+                    return
+            # The proxy target must be the route's own head: an independently
+            # chosen fastest node may not be part of the planned route, which
+            # previously injected downstream hops with wrong start layers
+            # (ADR-0020 mixed-topology flaw).
+            if route_nodes:
+                node = route_nodes[0]
+                target_url = f"{node.endpoint}/v1/chat/completions"
         # 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
@@ -3026,6 +3138,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             head_endpoint=node.endpoint,
             downstream=downstream_urls,
             route=route_debug or "<empty>",
+            routing=routing_decision,
             nodes=_node_route_summary(route_nodes),
         )
 
@@ -3402,6 +3515,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         can refine this later without changing the external stats shape.
         """
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        if route_model and route_nodes:
+            # Route-level bandit sample (ADR-0021); the store itself rejects
+            # near-empty completions that would poison the arm.
+            server.route_stats.record_sample(
+                route_model,
+                route_signature(route_model, route_nodes),
+                total_tokens,
+                elapsed_seconds,
+            )
         if server.stats is None or total_tokens <= 0:
             return
         elapsed_seconds = max(elapsed_seconds, 1e-6)
@@ -3825,7 +3947,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             for eid in stale_ids:
                 old = server.registry.pop(eid)
                 stale_entries.append((eid, old))
+            is_topology_change = node_id not in stale_ids or any(
+                (old.shard_start, old.shard_end) != (entry.shard_start, entry.shard_end)
+                for eid, old in stale_entries
+                if eid == node_id
+            )
             server.registry[node_id] = entry
+            if is_topology_change:
+                server.route_stats.bump_epoch(_route_stats_keys(server, entry))
             if entry.managed_assignment and not explicit_shard:
                 if entry.hf_repo:
                     _rebalance_hf_model_locked(server, entry.hf_repo)
@@ -5280,6 +5409,61 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             ],
         })
 
+    def _handle_routing(self, parsed: urllib.parse.ParseResult):
+        """Learned route table (ADR-0021): per-model candidates with observed
+        tps, coefficient vs the best proven route, and expected traffic share."""
+        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        params = urllib.parse.parse_qs(parsed.query)
+        only_model = (params.get("model") or [None])[0]
+        with server.lock:
+            self._purge_expired_nodes()
+            alive = list(server.registry.values())
+        model_keys: dict[str, list] = {}
+        for node in alive:
+            if node.shard_start is None or node.shard_end is None:
+                continue
+            key = node.hf_repo or node.model
+            if not key:
+                continue
+            model_keys.setdefault(key, []).append(node)
+        cfg = server.route_stats.config
+        out: dict[str, dict] = {}
+        for key, nodes in model_keys.items():
+            if only_model and key != only_model and not any(
+                key == alias for alias in _model_aliases(only_model)
+            ):
+                continue
+            resolved_name, preset = _resolve_model_preset(server.model_presets, key)
+            # Stats are recorded under the proxy's resolved route_model —
+            # use the same key here or lookups always miss.
+            stats_key = resolved_name or key
+            if preset is not None:
+                rs, re = _preset_layer_bounds(preset)
+            else:
+                layer_counts = [n.num_layers for n in nodes if n.num_layers is not None]
+                if not layer_counts:
+                    continue
+                rs, re = 0, max(layer_counts) - 1
+            candidates = _enumerate_routes(
+                nodes, rs, re,
+                model=stats_key,
+                contracts=server.contracts,
+                max_candidates=cfg.max_candidate_routes,
+            )
+            out[stats_key] = {
+                "epoch": server.route_stats.epoch(stats_key),
+                "routes": route_table(candidates, server.route_stats, stats_key),
+            }
+        self._send_json(200, {
+            "config": {
+                "explore_share": cfg.explore_share,
+                "weight_alpha": cfg.weight_alpha,
+                "stats_half_life_seconds": cfg.stats_half_life_seconds,
+                "min_sample_tokens": cfg.min_sample_tokens,
+            },
+            "models": out,
+        })
+
     def _handle_routes(self, parsed: urllib.parse.ParseResult):
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
         params = urllib.parse.parse_qs(parsed.query)
@@ -5398,6 +5582,7 @@ class TrackerServer:
         hf_pricing_refresh_interval: float = 86400.0,
         hf_pricing_fetch_html: Any | None = None,
         models_dir: str | Path | None = None,
+        routing_config: RoutingConfig | None = None,
     ) -> None:
         self._host = host
         self._requested_port = port
@@ -5500,6 +5685,18 @@ class TrackerServer:
         self._models_dir = Path(raw_models_dir).expanduser() if raw_models_dir else None
         self._hf_pricing_stop = threading.Event()
         self._hf_pricing_thread: threading.Thread | None = None
+        if routing_config is None:
+            routing_config = RoutingConfig(
+                explore_share=float(os.environ.get("MESHNET_ROUTE_EXPLORE_SHARE", RoutingConfig.explore_share)),
+                weight_alpha=float(os.environ.get("MESHNET_ROUTE_WEIGHT_ALPHA", RoutingConfig.weight_alpha)),
+                stats_half_life_seconds=float(
+                    os.environ.get("MESHNET_ROUTE_STATS_HALF_LIFE", RoutingConfig.stats_half_life_seconds)
+                ),
+                min_sample_tokens=int(
+                    os.environ.get("MESHNET_ROUTE_MIN_SAMPLE_TOKENS", RoutingConfig.min_sample_tokens)
+                ),
+            )
+        self._route_stats = RouteStatsStore(routing_config)
         self.port: int | None = None
 
     def start(self) -> int:
@@ -5537,6 +5734,7 @@ class TrackerServer:
             toploc_backend=self._toploc_backend,
             hf_pricing_log=self._hf_pricing_log,
             models_dir=self._models_dir,
+            route_stats=self._route_stats,
         )
         self.port = self._server.server_address[1]