155 lines
5.1 KiB
Python
155 lines
5.1 KiB
Python
"""US-030: manual route selection + hop-penalty benchmarking.
|
|
|
|
Pinned "route" in the chat body overrides auto-selection; the privileged
|
|
benchmark endpoint fans the same prompt through 1/2/3-node routes and appends
|
|
per-hop latency records to benchmark_results.json.
|
|
"""
|
|
|
|
import json
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.server import StubNodeServer
|
|
from meshnet_tracker.server import TrackerServer
|
|
|
|
MODEL = "openai-community/gpt2"
|
|
|
|
|
|
def _post_json(url: str, payload: dict, headers: dict | None = None) -> dict:
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(payload).encode(),
|
|
headers={"Content-Type": "application/json", **(headers or {})},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
@pytest.fixture
|
|
def benchmark_setup(tmp_path):
|
|
"""Tracker + one full-coverage node and one 0-5/6-11 pair, with node ids."""
|
|
results_path = str(tmp_path / "benchmark_results.json")
|
|
tracker = TrackerServer(
|
|
model_presets={
|
|
MODEL: {
|
|
"layers_start": 0,
|
|
"layers_end": 11,
|
|
"bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
|
|
}
|
|
},
|
|
benchmark_results_path=results_path,
|
|
)
|
|
port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{port}"
|
|
|
|
stubs = []
|
|
node_ids = {}
|
|
for name, (start, end, prefix, last) in {
|
|
"full": (0, 11, "full:", True),
|
|
"head": (0, 5, "head:", False),
|
|
"tail": (6, 11, "tail:", True),
|
|
}.items():
|
|
stub = StubNodeServer(
|
|
shard_start=start, shard_end=end, is_last_shard=last,
|
|
response_prefix=prefix, model=MODEL, tracker_mode=(start == 0),
|
|
)
|
|
sport = stub.start()
|
|
stubs.append(stub)
|
|
reply = _post_json(f"{tracker_url}/v1/nodes/register", {
|
|
"endpoint": f"http://127.0.0.1:{sport}",
|
|
"shard_start": start,
|
|
"shard_end": end,
|
|
"model": MODEL,
|
|
"hardware_profile": {},
|
|
"score": 1.0,
|
|
"tracker_mode": start == 0,
|
|
"node_id": f"node-{name}",
|
|
})
|
|
node_ids[name] = reply.get("node_id", f"node-{name}")
|
|
|
|
yield tracker_url, node_ids, results_path
|
|
|
|
for stub in stubs:
|
|
stub.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def _chat(tracker_url: str, route: list[str] | None = None) -> dict:
|
|
body: dict = {
|
|
"model": MODEL,
|
|
"messages": [{"role": "user", "content": "ping"}],
|
|
}
|
|
if route is not None:
|
|
body["route"] = route
|
|
return _post_json(f"{tracker_url}/v1/chat/completions", body)
|
|
|
|
|
|
def test_pinned_route_uses_named_node(benchmark_setup):
|
|
tracker_url, node_ids, _ = benchmark_setup
|
|
reply = _chat(tracker_url, route=[node_ids["full"]])
|
|
content = reply["choices"][0]["message"]["content"]
|
|
assert content.startswith("full:")
|
|
|
|
|
|
def test_unknown_route_node_is_400(benchmark_setup):
|
|
tracker_url, _, _ = benchmark_setup
|
|
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
|
_chat(tracker_url, route=["no-such-node"])
|
|
assert exc_info.value.code == 400
|
|
detail = json.loads(exc_info.value.read())
|
|
assert "no-such-node" in detail["error"]["message"]
|
|
|
|
|
|
def test_invalid_route_shape_is_400(benchmark_setup):
|
|
tracker_url, _, _ = benchmark_setup
|
|
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
|
_chat(tracker_url, route=[])
|
|
assert exc_info.value.code == 400
|
|
|
|
|
|
def test_clients_without_route_are_unaffected(benchmark_setup):
|
|
tracker_url, _, _ = benchmark_setup
|
|
reply = _chat(tracker_url)
|
|
assert reply["choices"][0]["message"]["content"]
|
|
|
|
|
|
def test_benchmark_requires_auth(benchmark_setup):
|
|
tracker_url, _, _ = benchmark_setup
|
|
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
|
_post_json(f"{tracker_url}/v1/benchmark/hop-penalty", {"model": MODEL})
|
|
assert exc_info.value.code == 401
|
|
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
|
urllib.request.urlopen(f"{tracker_url}/v1/benchmark/results")
|
|
assert exc_info.value.code == 401
|
|
|
|
|
|
def test_benchmark_records_one_and_two_node_routes(benchmark_setup):
|
|
tracker_url, _, results_path = benchmark_setup
|
|
record = _post_json(
|
|
f"{tracker_url}/v1/benchmark/hop-penalty",
|
|
{"model": MODEL, "prompt": "2+2", "max_new_tokens": 8},
|
|
headers={"Authorization": "Bearer bench"},
|
|
)
|
|
by_hops = {len(r["route"]): r for r in record["routes"] if "total_ms" in r}
|
|
assert 1 in by_hops and 2 in by_hops # 3-node coverage impossible → skipped
|
|
assert 3 not in by_hops
|
|
assert len(by_hops[1]["per_hop_ms"]) == 1
|
|
assert len(by_hops[2]["per_hop_ms"]) == 2
|
|
assert by_hops[2]["total_ms"] > 0
|
|
|
|
stored = json.loads(open(results_path, encoding="utf-8").read())
|
|
assert isinstance(stored, list) and len(stored) == 1
|
|
assert stored[0]["model"] == MODEL
|
|
assert stored[0]["prompt_hash"]
|
|
|
|
fetched = json.loads(
|
|
urllib.request.urlopen(urllib.request.Request(
|
|
f"{tracker_url}/v1/benchmark/results",
|
|
headers={"Authorization": "Bearer bench"},
|
|
)).read()
|
|
)
|
|
assert len(fetched["results"]) == 1
|