291 lines
11 KiB
Python
291 lines
11 KiB
Python
"""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"
|