popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
7cf8d9bcf31ed66308aa4a261a61d62770faca1d
neuron-tai/tests/test_dynamic_routing.py
Dobromir Popov 7cf8d9bcf3 test descriptions
2026-07-11 22:25:30 +03:00

368 lines
15 KiB
Python
Raw Blame History

"""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():
"Route stats sample becomes proven and decays\n\nTags: http, performance, routing, tracker"
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():
"Route stats rejects near empty samples\n\nTags: http, performance, routing, tracker"
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():
"Route stats epoch bump marks stale\n\nTags: http, performance, routing, tracker"
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():
"Route stats ewma averages samples\n\nTags: http, performance, routing, tracker"
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):
"Route stats persist historical hop latency across restart\n\nTags: http, performance, persistence, routing, tracker"
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():
"Choose route without samples is deterministic best prior\n\nTags: http, performance, routing, tracker"
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():
"Choose route traffic proportional to tps\n\nTags: http, performance, routing, tracker"
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():
"Choose route scouts unproven routes at explore share\n\nTags: http, performance, routing, tracker"
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():
"Enumerate routes mixed topology yields both routes\n\nTags: http, performance, routing, tracker"
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():
"Enumerate routes requires head at first layer\n\nTags: http, performance, routing, tracker"
tail_only = _fake_node("tail", 22, 39)
assert _enumerate_routes([tail_only], 0, 39, model="m") == []
def test_route_table_reports_coefficient_and_share():
"Route table reports coefficient and share\n\nTags: http, performance, routing, tracker"
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.\n\nTags: http, performance, routing, tracker"
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():
"Admin model load request queues directive on joined node\n\nTags: http, performance, routing, tracker"
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():
"Endpoint key distinguishes same port different hosts\n\nTags: http, performance, routing, tracker"
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},
]
Reference in New Issue View Git Blame Copy Permalink