routing improvements - dynamic (wip)

2026-07-07 21:25:28 +02:00
parent e2b20883ca
commit 518c259cd3
6 changed files with 964 additions and 19 deletions
--- a/tests/test_dynamic_routing.py
+++ b/tests/test_dynamic_routing.py
@@ -0,0 +1,290 @@
+"""ADR-0021: dynamic bandit-style route selection with learned statistics."""
+
+import http.server
+import json
+import random
+import threading
+import types
+import urllib.request
+
+from meshnet_tracker.routing_stats import (
+    RouteCandidate,
+    RouteStatsStore,
+    RoutingConfig,
+    choose_route,
+    route_signature,
+    route_table,
+)
+from meshnet_tracker.server import TrackerServer, _enumerate_routes
+
+
+def _post_json(url: str, payload: dict) -> dict:
+    req = urllib.request.Request(
+        url,
+        data=json.dumps(payload).encode(),
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    with urllib.request.urlopen(req, timeout=10.0) as resp:
+        return json.loads(resp.read())
+
+
+def _get_json(url: str) -> dict:
+    with urllib.request.urlopen(url, timeout=10.0) as resp:
+        return json.loads(resp.read())
+
+
+def _fake_node(node_id, shard_start, shard_end, benchmark=100.0, endpoint=None):
+    return types.SimpleNamespace(
+        node_id=node_id,
+        endpoint=endpoint or f"http://{node_id}:7000",
+        model="qwen3.6-35b-a3b",
+        hf_repo="unsloth/Qwen3.6-35B-A3B",
+        shard_start=shard_start,
+        shard_end=shard_end,
+        num_layers=40,
+        benchmark_tokens_per_sec=benchmark,
+        model_tokens_per_sec={},
+        queue_depth=0,
+        proxy_inflight=0,
+        wallet_address=None,
+        relay_addr=None,
+    )
+
+
+# ---- RouteStatsStore ----------------------------------------------------
+
+
+def test_route_stats_sample_becomes_proven_and_decays():
+    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=100.0))
+    sig = "m|a[0-39]"
+    assert store.snapshot(sig, "m", now=0.0)["status"] == "unsampled"
+    assert store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
+    snap = store.snapshot(sig, "m", now=1.0)
+    assert snap["status"] == "proven"
+    assert snap["tps"] == 10.0
+    # After many half-lives the sample mass decays below the proven threshold.
+    assert store.snapshot(sig, "m", now=1000.0)["status"] == "decayed"
+
+
+def test_route_stats_rejects_near_empty_samples():
+    store = RouteStatsStore(RoutingConfig(min_sample_tokens=8))
+    assert not store.record_sample("m", "sig", tokens=3, elapsed_seconds=1.0)
+    assert store.snapshot("sig", "m")["samples"] == 0
+
+
+def test_route_stats_epoch_bump_marks_stale():
+    store = RouteStatsStore()
+    sig = "m|a[0-39]"
+    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)
+    assert store.snapshot(sig, "m", now=1.0)["status"] == "proven"
+    store.bump_epoch(["m"])
+    snap = store.snapshot(sig, "m", now=1.0)
+    assert snap["status"] == "stale"
+    assert snap["tps"] == 10.0  # EWMA kept as a prior for display
+    # A fresh sample under the new epoch re-proves the route.
+    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=2.0)
+    assert store.snapshot(sig, "m", now=3.0)["status"] == "proven"
+
+
+def test_route_stats_ewma_averages_samples():
+    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
+    sig = "m|a"
+    store.record_sample("m", sig, tokens=100, elapsed_seconds=10.0, now=0.0)  # 10 tps
+    store.record_sample("m", sig, tokens=200, elapsed_seconds=10.0, now=1.0)  # 20 tps
+    snap = store.snapshot(sig, "m", now=2.0)
+    assert 14.9 < snap["tps"] < 15.1
+
+
+# ---- choose_route --------------------------------------------------------
+
+
+def _candidates_two_routes():
+    fast = RouteCandidate(nodes=[], signature="m|fast", prior_tps=100.0)
+    slow = RouteCandidate(nodes=[], signature="m|slow", prior_tps=50.0)
+    return fast, slow
+
+
+def test_choose_route_without_samples_is_deterministic_best_prior():
+    store = RouteStatsStore()
+    fast, slow = _candidates_two_routes()
+    for _ in range(20):
+        picked, decision = choose_route([slow, fast], store, "m", rng=random.Random(7))
+        assert picked is fast
+        assert decision["mode"] == "scout"
+
+
+def test_choose_route_traffic_proportional_to_tps():
+    store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9))
+    fast, slow = _candidates_two_routes()
+    now = 0.0
+    for _ in range(5):
+        now += 1.0
+        store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
+        store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
+    rng = random.Random(42)
+    picks = {"m|fast": 0, "m|slow": 0}
+    for _ in range(4000):
+        picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
+        assert decision["mode"] == "exploit"
+        picks[picked.signature] += 1
+    share = picks["m|fast"] / 4000
+    # 15 tps vs 10 tps at alpha=1 → expected fast share 0.6
+    assert 0.55 < share < 0.65
+
+
+def test_choose_route_scouts_unproven_routes_at_explore_share():
+    store = RouteStatsStore(RoutingConfig(explore_share=0.25, stats_half_life_seconds=1e9))
+    fast, slow = _candidates_two_routes()
+    now = 1.0
+    store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
+    rng = random.Random(11)
+    scouted = 0
+    for _ in range(4000):
+        picked, decision = choose_route([fast, slow], store, "m", rng=rng, now=now)
+        if decision["mode"] == "scout":
+            scouted += 1
+            assert picked is slow
+    assert 0.20 < scouted / 4000 < 0.30
+
+
+# ---- _enumerate_routes ---------------------------------------------------
+
+
+def test_enumerate_routes_mixed_topology_yields_both_routes():
+    gpu = _fake_node("gpu", 0, 21, benchmark=11000.0)
+    cpu = _fake_node("cpu", 0, 39, benchmark=425.0)
+    candidates = _enumerate_routes([gpu, cpu], 0, 39, model="qwen3.6-35b-a3b")
+    signatures = {c.signature for c in candidates}
+    assert signatures == {
+        route_signature("qwen3.6-35b-a3b", [gpu, cpu]),
+        route_signature("qwen3.6-35b-a3b", [cpu]),
+    }
+    hybrid = next(c for c in candidates if len(c.nodes) == 2)
+    assert [n.node_id for n in hybrid.nodes] == ["gpu", "cpu"]
+    # Hybrid route's prior is its bottleneck hop, not the fast head.
+    assert hybrid.prior_tps == 425.0
+
+
+def test_enumerate_routes_requires_head_at_first_layer():
+    tail_only = _fake_node("tail", 22, 39)
+    assert _enumerate_routes([tail_only], 0, 39, model="m") == []
+
+
+def test_route_table_reports_coefficient_and_share():
+    store = RouteStatsStore(RoutingConfig(explore_share=0.3, stats_half_life_seconds=1e9))
+    fast, slow = _candidates_two_routes()
+    now = 1.0
+    for _ in range(3):
+        store.record_sample("m", fast.signature, tokens=150, elapsed_seconds=10.0, now=now)
+        store.record_sample("m", slow.signature, tokens=100, elapsed_seconds=10.0, now=now)
+        now += 1.0
+    rows = route_table([fast, slow], store, "m", now=now)
+    by_sig = {r["signature"]: r for r in rows}
+    assert by_sig["m|fast"]["coefficient"] == 1.0
+    assert abs(by_sig["m|slow"]["coefficient"] - (10.0 / 15.0)) < 0.01
+    # No scouts → full exploit budget split 0.6 / 0.4.
+    assert abs(by_sig["m|fast"]["expected_share"] - 0.6) < 0.01
+    assert abs(by_sig["m|slow"]["expected_share"] - 0.4) < 0.01
+
+
+# ---- integration: proxy uses route head + /v1/routing --------------------
+
+
+def test_proxy_head_is_route_head_and_routing_endpoint_lists_routes():
+    """Mixed topology (partial head 0-21 + full node 0-39): the proxy target
+    must be the selected route's own head, downstream hops must continue at
+    head.shard_end + 1 (the ADR-0020 flaw), and /v1/routing must list both
+    candidate routes."""
+
+    class ChatHandler(http.server.BaseHTTPRequestHandler):
+        def log_message(self, *args):  # noqa: ARG002
+            pass
+
+        def do_POST(self):
+            length = int(self.headers.get("Content-Length", 0))
+            self.rfile.read(length)
+            route_header = self.headers.get("X-Meshnet-Route") or "[]"
+            body = json.dumps({
+                "choices": [{"message": {"role": "assistant", "content": route_header}}],
+                "usage": {"prompt_tokens": 10, "completion_tokens": 40},
+            }).encode()
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+    stubs = []
+    threads = []
+    for _ in range(2):
+        stub = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
+        thread = threading.Thread(target=stub.serve_forever, daemon=True)
+        thread.start()
+        stubs.append(stub)
+        threads.append(thread)
+    gpu_stub, cpu_stub = stubs
+
+    tracker = TrackerServer(model_presets={
+        "qwen3.6-35b-a3b": {
+            "layers_start": 0,
+            "layers_end": 39,
+            "hf_repo": "unsloth/Qwen3.6-35B-A3B",
+            "aliases": ["Qwen3.6-35B-A3B"],
+        }
+    })
+    tracker_port = tracker.start()
+    try:
+        tracker._server.route_rng = random.Random(3)
+        for stub, shard_end, bench in ((gpu_stub, 21, 11000.0), (cpu_stub, 39, 425.0)):
+            _post_json(
+                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+                {"endpoint": f"http://127.0.0.1:{stub.server_address[1]}",
+                 "model": "qwen3.6-35b-a3b",
+                 "hf_repo": "unsloth/Qwen3.6-35B-A3B",
+                 "num_layers": 40,
+                 "shard_start": 0,
+                 "shard_end": shard_end,
+                 "tracker_mode": True,
+                 "benchmark_tokens_per_sec": bench,
+                 "hardware_profile": {},
+                 "score": 1.0},
+            )
+
+        for _ in range(8):
+            _post_json(
+                f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
+                {"model": "Qwen3.6-35B-A3B",
+                 "messages": [{"role": "user", "content": "hi"}]},
+            )
+
+        console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
+        routing = _get_json(f"http://127.0.0.1:{tracker_port}/v1/routing")
+    finally:
+        tracker.stop()
+        for stub, thread in zip(stubs, threads):
+            stub.shutdown()
+            stub.server_close()
+            thread.join(timeout=1.0)
+
+    gpu_endpoint = f"http://127.0.0.1:{gpu_stub.server_address[1]}"
+    cpu_endpoint = f"http://127.0.0.1:{cpu_stub.server_address[1]}"
+    selected = [e for e in console["events"] if e["message"] == "proxy route selected"]
+    assert selected
+    for event in selected:
+        fields = event["fields"]
+        nodes = fields["nodes"]
+        # The proxy head must be the route's first hop (ADR-0020 regression).
+        assert fields["head_endpoint"] == nodes[0]["endpoint"]
+        downstream = json.loads(fields["downstream"])
+        if fields["head_endpoint"] == gpu_endpoint:
+            # Partial head: downstream continues at layer 22, never 0.
+            assert downstream == [{"endpoint": cpu_endpoint, "start_layer": 22}]
+        else:
+            assert fields["head_endpoint"] == cpu_endpoint
+            assert downstream == []
+
+    table = routing["models"]["qwen3.6-35b-a3b"]
+    assert len(table["routes"]) == 2
+    sampled = [r for r in table["routes"] if r["samples"] > 0]
+    assert sampled, "completed requests must produce route samples"