"""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 test_route_stats_persist_historical_hop_latency_across_restart(tmp_path): db_path = str(tmp_path / "routing-stats.sqlite") signature = "m|gpu[0-11]->cpu[12-23]" store = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9), db_path=db_path) store.record_sample("m", signature, tokens=100, elapsed_seconds=2.0, now=100.0) store.save_to_db() restored = RouteStatsStore(RoutingConfig(stats_half_life_seconds=1e9), db_path=db_path) row = restored.model_rows("m", now=101.0)[0] assert row["hop_count"] == 2 assert row["samples"] == 1 assert row["tps"] == 50.0 assert row["latency_ms"] == 2000.0 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" def test_admin_model_load_request_queues_directive_on_joined_node(): tracker = TrackerServer(validator_service_token="test-admin") port = tracker.start() try: node = _post_json( f"http://127.0.0.1:{port}/v1/nodes/register", { "endpoint": "http://127.0.0.1:9911", "model": "stub-model", "shard_start": 0, "shard_end": 3, "managed_assignment": True, "memory_mb": 32768, "hardware_profile": {"host_id": "available-ram-pool"}, }, ) request = urllib.request.Request( f"http://127.0.0.1:{port}/v1/models/load", data=json.dumps({"model": "qwen2.5-0.5b-instruct"}).encode(), headers={"Content-Type": "application/json", "Authorization": "Bearer test-admin"}, method="POST", ) with urllib.request.urlopen(request) as response: result = json.loads(response.read()) heartbeat = _post_json( f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat", {}) finally: tracker.stop() assert result["status"] == "queued" assert result["assignment"]["node_id"] == node["node_id"] assert heartbeat["directives"][0]["action"] == "ADD_SHARD" assert heartbeat["directives"][0]["model"] == "Qwen/Qwen2.5-0.5B-Instruct" def test_endpoint_key_distinguishes_same_port_different_hosts(): from meshnet_node.torch_server import _clamp_downstream_hops, _endpoint_key assert _endpoint_key("http://192.168.0.20:7000") == "192.168.0.20:7000" assert _endpoint_key("http://192.168.0.179:7000") == "192.168.0.179:7000" assert _endpoint_key("http://192.168.0.20:7000") != _endpoint_key("http://192.168.0.179:7000") class Backend: shard_end = 21 hops = [{"endpoint": "http://192.168.0.179:7000", "start_layer": 0}] assert _clamp_downstream_hops(hops, Backend()) == [ {"endpoint": "http://192.168.0.179:7000", "start_layer": 22}, ]