1398 lines
53 KiB
Python
1398 lines
53 KiB
Python
"""US-003 tests: tracker node registration and route selection."""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
from meshnet_gateway.server import GatewayServer, _banned_route_wallet
|
|
from meshnet_node.server import StubNodeServer
|
|
from meshnet_contracts import LocalSolanaContracts
|
|
from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error
|
|
|
|
|
|
def _post_json(url: str, payload: dict) -> dict:
|
|
data = json.dumps(payload).encode()
|
|
req = urllib.request.Request(
|
|
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
|
)
|
|
with urllib.request.urlopen(req) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def _get_json(url: str) -> dict:
|
|
with urllib.request.urlopen(url) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect():
|
|
"""A disconnected client must not dump a BrokenPipe traceback from the tracker."""
|
|
from meshnet_tracker.server import _TrackerHandler
|
|
|
|
class BrokenPipeWriter:
|
|
def write(self, _body):
|
|
raise BrokenPipeError
|
|
|
|
class DummyHandler:
|
|
wfile = BrokenPipeWriter()
|
|
|
|
def send_response(self, _status):
|
|
pass
|
|
|
|
def send_header(self, _name, _value):
|
|
pass
|
|
|
|
def end_headers(self):
|
|
pass
|
|
|
|
_TrackerHandler._send_json(DummyHandler(), 200, {"ok": True})
|
|
|
|
|
|
def test_tracker_serves_health_while_proxy_request_is_in_flight():
|
|
"""Long inference proxy requests must not block heartbeats/health checks."""
|
|
|
|
class SlowChatHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args):
|
|
pass
|
|
|
|
def do_POST(self):
|
|
if self.path != "/v1/chat/completions":
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
return
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
self.rfile.read(length)
|
|
time.sleep(2.0)
|
|
body = json.dumps({"choices": [{"message": {"content": "ok"}}]}).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)
|
|
|
|
slow_node = http.server.HTTPServer(("127.0.0.1", 0), SlowChatHandler)
|
|
slow_thread = threading.Thread(target=slow_node.serve_forever, daemon=True)
|
|
slow_thread.start()
|
|
tracker = TrackerServer(heartbeat_timeout=60.0)
|
|
tracker_port = tracker.start()
|
|
proxy_error = []
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{slow_node.server_address[1]}",
|
|
"model": "slow-model", "hf_repo": "org/slow-model", "num_layers": 1,
|
|
"shard_start": 0, "shard_end": 0, "tracker_mode": True,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
def call_proxy():
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
|
{"model": "slow-model", "messages": [{"role": "user", "content": "hi"}]},
|
|
)
|
|
except Exception as exc:
|
|
proxy_error.append(exc)
|
|
|
|
proxy_thread = threading.Thread(target=call_proxy)
|
|
proxy_thread.start()
|
|
time.sleep(0.2)
|
|
|
|
with urllib.request.urlopen(f"http://127.0.0.1:{tracker_port}/v1/health", timeout=1.0) as resp:
|
|
assert resp.status == 200
|
|
|
|
proxy_thread.join(timeout=3.0)
|
|
assert not proxy_thread.is_alive()
|
|
assert not proxy_error
|
|
finally:
|
|
tracker.stop()
|
|
slow_node.shutdown()
|
|
slow_node.server_close()
|
|
slow_thread.join(timeout=1.0)
|
|
|
|
|
|
def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix():
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
wallet = "7j77FsPY1evV8tuf4Z73AVrWwxBEW1pvKwi4EvcRD3g"
|
|
try:
|
|
first = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 0, "shard_end": 21, "wallet_address": wallet,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
second = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 0, "shard_end": 21, "wallet_address": wallet,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
different_endpoint = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 20, "shard_end": 23, "wallet_address": wallet,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
assert first["node_id"].startswith("7j77FsPY-")
|
|
assert second["node_id"] == first["node_id"]
|
|
assert different_endpoint["node_id"].startswith("7j77FsPY-")
|
|
assert different_endpoint["node_id"] != first["node_id"]
|
|
|
|
route_resp = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
|
|
)
|
|
assert route_resp["nodes"][0]["node_id"] == first["node_id"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_node_registration():
|
|
"""A node can register with the tracker and receives a node_id."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
resp = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{
|
|
"endpoint": "http://127.0.0.1:9000",
|
|
"shard_start": 0,
|
|
"shard_end": 15,
|
|
"hardware_profile": {"gpu": "rtx3090"},
|
|
"score": 1.0,
|
|
},
|
|
)
|
|
assert "node_id" in resp
|
|
assert isinstance(resp["node_id"], str)
|
|
assert len(resp["node_id"]) > 0
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_selection():
|
|
"""Tracker returns ordered route when nodes collectively cover all layers."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
route_resp = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
|
)
|
|
assert "route" in route_resp
|
|
assert route_resp["route"] == [
|
|
"http://127.0.0.1:9001",
|
|
"http://127.0.0.1:9002",
|
|
]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_error_no_coverage():
|
|
"""Tracker returns 503 when registered nodes do not cover all required layers."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
# Only half the layers covered — route must fail.
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 but request succeeded")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_coverage_endpoint_reports_uncovered_ranges():
|
|
"""Coverage endpoint returns compressed layer ranges with node counts."""
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 6,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
|
"shard_start": 0, "shard_end": 2, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
|
|
bands = coverage_resp["coverage"]
|
|
assert len(bands) == 2
|
|
# Covered band: layers 0-2, one node
|
|
assert bands[0]["start_layer"] == 0
|
|
assert bands[0]["end_layer"] == 2
|
|
assert bands[0]["node_count"] == 1
|
|
assert len(bands[0]["nodes"]) == 1
|
|
node_info = bands[0]["nodes"][0]
|
|
assert node_info["endpoint"] == "http://127.0.0.1:9001"
|
|
assert node_info["alive"] is True
|
|
assert "last_seen_seconds_ago" in node_info
|
|
assert "heartbeat_success_rate" in node_info
|
|
assert "inference_success_rate" in node_info
|
|
# Gap band: layers 3-5, no nodes
|
|
assert bands[1]["start_layer"] == 3
|
|
assert bands[1]["end_layer"] == 5
|
|
assert bands[1]["node_count"] == 0
|
|
assert bands[1]["nodes"] == []
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
|
raise AssertionError("Expected 503 for zero-coverage range")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name():
|
|
"""Coverage endpoint supports real HF models registered outside preset catalog."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 0, "shard_end": 21, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 20, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
by_repo = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen/Qwen2.5-0.5B-Instruct"
|
|
)
|
|
by_short_name = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen2.5-0.5B-Instruct"
|
|
)
|
|
|
|
assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
|
bands = by_repo["coverage"]
|
|
assert len(bands) == 3
|
|
assert bands[0]["start_layer"] == 0 and bands[0]["end_layer"] == 19 and bands[0]["node_count"] == 1
|
|
assert bands[1]["start_layer"] == 20 and bands[1]["end_layer"] == 21 and bands[1]["node_count"] == 2
|
|
assert bands[2]["start_layer"] == 22 and bands[2]["end_layer"] == 23 and bands[2]["node_count"] == 1
|
|
assert all("nodes" in b for b in bands)
|
|
assert by_short_name["coverage"] == by_repo["coverage"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_models_endpoint_lists_registered_hf_repo_and_short_name_alias():
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9111", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 0, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
models_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/models")
|
|
|
|
model = next(item for item in models_resp["data"] if item["id"] == "Qwen/Qwen2.5-0.5B-Instruct")
|
|
assert model["name"] == "Qwen2.5-0.5B-Instruct"
|
|
assert model["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
|
assert model["aliases"] == ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen2.5-0.5B-Instruct"]
|
|
assert model["shard_coverage_percentage"] == 100.0
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_auto_assigns_new_node_to_uncovered_range_first():
|
|
"""Capability-driven registration fills the first uncovered layer gap."""
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 8,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
|
"shard_start": 0, "shard_end": 3, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
|
|
"vram_bytes": 10_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
assert reg["directive"]["action"] == "LOAD_SHARD"
|
|
assert reg["directive"]["start_layer"] == 4
|
|
assert reg["directive"]["end_layer"] == 7
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_speed_weighted_vram_assignment_covers_model():
|
|
"""Three auto-assigned nodes reach 100% coverage with widest range on largest VRAM."""
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 12,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
|
"vram_bytes": 4_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
|
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9003", "model": "tiny-model",
|
|
"vram_bytes": 7_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
|
widths = {
|
|
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
|
for node in route_resp["nodes"]
|
|
}
|
|
assert widths["http://127.0.0.1:9003"] == 5
|
|
assert widths["http://127.0.0.1:9002"] == 4
|
|
assert widths["http://127.0.0.1:9001"] == 3
|
|
|
|
coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
|
|
assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"])
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap():
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 12,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9011", "model": "tiny-model",
|
|
"vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9012", "model": "tiny-model",
|
|
"vram_bytes": 15_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
|
widths = {
|
|
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
|
for node in route_resp["nodes"]
|
|
}
|
|
assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_registration_directive_is_not_replayed_on_heartbeat():
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 2,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9013", "model": "tiny-model",
|
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
assert reg["directive"]["action"] == "LOAD_SHARD"
|
|
|
|
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{reg['node_id']}/heartbeat", {})
|
|
assert hb == {}
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_reassignment_emits_drop_before_load():
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 4,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
slow = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9015", "model": "tiny-model",
|
|
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9016", "model": "tiny-model",
|
|
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{slow['node_id']}/heartbeat", {})
|
|
assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_periodic_rebalance_purges_expired_nodes_without_requests():
|
|
tracker = TrackerServer(
|
|
heartbeat_timeout=0.05,
|
|
rebalance_interval=0.02,
|
|
model_presets={"tiny-model": {"total_layers": 1, "bytes_per_layer": {"bfloat16": 1_000}}},
|
|
)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9014", "model": "tiny-model",
|
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
assert reg["node_id"] in tracker._registry
|
|
time.sleep(0.12)
|
|
assert reg["node_id"] not in tracker._registry
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_faster_node_receives_wider_range_when_capacity_ties():
|
|
tracker = TrackerServer(model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 12,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
|
"vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
|
|
"vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
|
widths = {
|
|
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
|
for node in route_resp["nodes"]
|
|
}
|
|
assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_rebalances_after_middle_range_node_timeout():
|
|
"""Killing a middle shard queues LOAD_SHARD and restores coverage."""
|
|
tracker = TrackerServer(
|
|
heartbeat_timeout=0.15,
|
|
model_presets={
|
|
"tiny-model": {
|
|
"total_layers": 12,
|
|
"bytes_per_layer": {"bfloat16": 1_000},
|
|
},
|
|
},
|
|
)
|
|
tracker_port = tracker.start()
|
|
ids: dict[str, str] = {}
|
|
try:
|
|
for port in (9001, 9002, 9003, 9004):
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port}", "model": "tiny-model",
|
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
ids[str(port)] = reg["node_id"]
|
|
|
|
for node_id in ids.values():
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})
|
|
|
|
time.sleep(0.10)
|
|
for port in ("9001", "9003", "9004"):
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {})
|
|
time.sleep(0.10)
|
|
|
|
coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
|
|
assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"])
|
|
|
|
heartbeat_responses = [
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {})
|
|
for port in ("9001", "9003", "9004")
|
|
]
|
|
directives = [
|
|
directive
|
|
for response in heartbeat_responses
|
|
for directive in response.get("directives", [])
|
|
]
|
|
assert any(directive["action"] == "LOAD_SHARD" for directive in directives)
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_rebalances_managed_hf_node_after_peer_timeout():
|
|
"""HF nodes auto-assigned by the tracker receive LOAD_SHARD after a peer dies."""
|
|
tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
managed = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 0, "shard_end": 21, "managed_assignment": True,
|
|
"vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
expired = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
|
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
|
"shard_start": 22, "shard_end": 23,
|
|
"vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
time.sleep(0.10)
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
|
|
time.sleep(0.10)
|
|
|
|
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
|
|
|
|
assert expired["node_id"] not in tracker._registry
|
|
load_directives = [d for d in hb.get("directives", []) if d["action"] == "LOAD_SHARD"]
|
|
assert load_directives
|
|
assert load_directives[-1]["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
|
assert load_directives[-1]["start_layer"] == 0
|
|
assert load_directives[-1]["end_layer"] == 23
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_error_no_nodes():
|
|
"""Tracker returns 503 with clear error when the registry is empty."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 but request succeeded")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_updates_node():
|
|
"""Sending a heartbeat for a registered node succeeds."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
node_id = reg["node_id"]
|
|
|
|
hb_resp = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}
|
|
)
|
|
assert hb_resp == {}
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_expiry():
|
|
"""Nodes that miss their heartbeat window are excluded from routes."""
|
|
tracker = TrackerServer(heartbeat_timeout=0.05) # 50 ms
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
# Immediately after registration the node is alive.
|
|
route_resp = _get_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
|
)
|
|
assert route_resp["route"] == ["http://127.0.0.1:9001"]
|
|
|
|
# Let the heartbeat window expire.
|
|
time.sleep(0.15)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 after heartbeat expiry")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_heartbeat_expiry_removes_node_from_registry():
|
|
"""Expired nodes are removed, not merely hidden from route responses."""
|
|
tracker = TrackerServer(heartbeat_timeout=0.05)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
reg = _post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
node_id = reg["node_id"]
|
|
time.sleep(0.15)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 after heartbeat expiry")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
|
|
try:
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})
|
|
raise AssertionError("Expected 404 for removed node")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 404
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_route_rejects_non_extending_overlap():
|
|
"""Overlapping shards that do not extend coverage cannot form a complete route."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 10,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 5,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 for incomplete overlapping route")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_registration_rejects_invalid_payload():
|
|
"""Registration errors return a defined JSON 400 response."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "not-a-url", "shard_start": 10, "shard_end": 0,
|
|
"hardware_profile": []},
|
|
)
|
|
raise AssertionError("Expected 400 for invalid registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 400
|
|
body = json.loads(exc.read())
|
|
assert "error" in body
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_routes_only_nodes_for_requested_model():
|
|
"""A node registered for one model cannot satisfy another model's route."""
|
|
tracker = TrackerServer(model_presets={
|
|
"model-a": {"layers_start": 0, "layers_end": 31},
|
|
"model-b": {"layers_start": 0, "layers_end": 31},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
|
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=model-b")
|
|
raise AssertionError("Expected model-b route to ignore model-a nodes")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_assignment_coverage_is_model_scoped():
|
|
"""Shard assignment gaps are computed only from nodes serving the same model."""
|
|
tracker = TrackerServer(model_presets={
|
|
"model-a": {"layers_start": 0, "layers_end": 31},
|
|
"model-b": {"layers_start": 0, "layers_end": 31},
|
|
})
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
|
|
assignment = _get_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/assign?model=model-b&device=cpu")
|
|
assert assignment["shard_start"] == 0
|
|
assert assignment["shard_end"] == 15
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_excludes_banned_wallets_from_routes():
|
|
"""Tracker refuses route candidates whose wallet is banned on-chain."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
raise AssertionError("Expected 403 for banned wallet registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 403
|
|
|
|
try:
|
|
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
|
raise AssertionError("Expected 503 when only banned wallet covers a layer")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
body = json.loads(exc.read())
|
|
assert "layer 16" in body["error"]
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_rejects_banned_wallet_registration():
|
|
"""A banned wallet cannot register with the tracker."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
try:
|
|
try:
|
|
_post_json(
|
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
|
{"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
raise AssertionError("Expected 403 for banned wallet registration")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 403
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "wallet is banned"}
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_tracker_ban_registration_guard_reads_contract_state():
|
|
"""The tracker registration guard reads ban status from the contract facade."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
|
|
assert _registration_ban_error(contracts, "wallet-b") == "wallet is banned"
|
|
assert _registration_ban_error(contracts, "wallet-a") is None
|
|
assert _registration_ban_error(None, "wallet-b") is None
|
|
|
|
|
|
def test_gateway_route_guard_rejects_banned_wallet_metadata():
|
|
"""Gateway rejects tracker routes that include a wallet now banned on-chain."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.ban_wallet("wallet-b")
|
|
|
|
banned_wallet = _banned_route_wallet(
|
|
contracts,
|
|
[{"endpoint": "http://node", "wallet_address": "wallet-b"}],
|
|
)
|
|
|
|
assert banned_wallet == "wallet-b"
|
|
|
|
|
|
def test_gateway_routes_even_when_compute_node_has_low_stake():
|
|
"""Routing is not based on a compute node's token balance."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 5)
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
|
port_a = node_a.start()
|
|
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
|
port_b = node_b.start()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(
|
|
tracker_url=tracker_url,
|
|
contracts=contracts,
|
|
minimum_stake=100,
|
|
)
|
|
gateway_port = gateway.start()
|
|
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json", "Authorization": "Bearer api-key-a"},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
assert node_b.received_activations
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def test_gateway_records_compute_attribution_after_inference_session():
|
|
"""Gateway records layer and token attribution for every routed node."""
|
|
contracts = LocalSolanaContracts()
|
|
contracts.registry.submit_stake("wallet-a", 1_000)
|
|
contracts.registry.submit_stake("wallet-b", 1_000)
|
|
contracts.payment.fund_api_key("api-key-a", lamports=10_000)
|
|
tracker = TrackerServer(contracts=contracts, minimum_stake=100)
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
|
port_a = node_a.start()
|
|
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
|
port_b = node_b.start()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(
|
|
tracker_url=tracker_url,
|
|
contracts=contracts,
|
|
minimum_stake=100,
|
|
)
|
|
gateway_port = gateway.start()
|
|
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json", "Authorization": "Bearer api-key-a"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
|
|
records = contracts.payment.list_attributions()
|
|
assert [(r.api_key, r.node_wallet, r.layer_start, r.layer_end, r.tokens) for r in records] == [
|
|
("api-key-a", "wallet-a", 0, 15, 1),
|
|
("api-key-a", "wallet-b", 16, 31, 1),
|
|
]
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def test_gateway_returns_json_503_when_tracker_unavailable():
|
|
"""Tracker-backed gateway failures are reported as JSON HTTP errors."""
|
|
gateway = GatewayServer(tracker_url="http://127.0.0.1:9")
|
|
gateway_port = gateway.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
urllib.request.urlopen(req, timeout=5)
|
|
raise AssertionError("Expected 503 for unavailable tracker")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
assert exc.headers["Content-Type"] == "application/json"
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "tracker unavailable"}
|
|
finally:
|
|
gateway.stop()
|
|
|
|
|
|
def test_gateway_returns_json_503_for_malformed_tracker_route():
|
|
"""Malformed tracker route payloads do not drop the gateway connection."""
|
|
|
|
class MalformedRouteHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args):
|
|
pass
|
|
|
|
def do_GET(self):
|
|
body = json.dumps({"route": []}).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)
|
|
|
|
fake_tracker = http.server.HTTPServer(("127.0.0.1", 0), MalformedRouteHandler)
|
|
fake_thread = threading.Thread(target=fake_tracker.serve_forever, daemon=True)
|
|
fake_thread.start()
|
|
gateway = GatewayServer(tracker_url=f"http://127.0.0.1:{fake_tracker.server_address[1]}")
|
|
gateway_port = gateway.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
urllib.request.urlopen(req, timeout=5)
|
|
raise AssertionError("Expected 503 for malformed tracker route")
|
|
except urllib.error.HTTPError as exc:
|
|
assert exc.code == 503
|
|
assert exc.headers["Content-Type"] == "application/json"
|
|
body = json.loads(exc.read())
|
|
assert body == {"error": "tracker unavailable"}
|
|
finally:
|
|
gateway.stop()
|
|
fake_tracker.shutdown()
|
|
fake_tracker.server_close()
|
|
fake_thread.join(timeout=1)
|
|
|
|
|
|
def test_two_node_pipeline_via_tracker():
|
|
"""End-to-end: nodes register with tracker; gateway discovers route dynamically."""
|
|
tracker = TrackerServer()
|
|
tracker_port = tracker.start()
|
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
|
|
|
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
|
port_a = node_a.start()
|
|
|
|
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
|
port_b = node_b.start()
|
|
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
_post_json(
|
|
f"{tracker_url}/v1/nodes/register",
|
|
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
|
"hardware_profile": {}, "score": 1.0},
|
|
)
|
|
|
|
gateway = GatewayServer(tracker_url=tracker_url)
|
|
gateway_port = gateway.start()
|
|
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "stub-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}).encode()
|
|
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req) as resp:
|
|
assert resp.status == 200
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["id"] == "chatcmpl-stub"
|
|
assert body["object"] == "chat.completion"
|
|
assert body["choices"][0]["message"]["content"] == "stub response to: hello"
|
|
assert node_b.received_activations, "node B did not receive activations from node A"
|
|
assert not node_a.received_activations, "node A should not receive activations"
|
|
finally:
|
|
gateway.stop()
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
# ---------------------------------------------------------------- US-025 stats
|
|
|
|
def test_stats_endpoint_returns_empty_models_initially():
|
|
tracker = TrackerServer()
|
|
port = tracker.start()
|
|
try:
|
|
resp = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
|
assert "models" in resp
|
|
assert isinstance(resp["models"], dict)
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
def test_stats_rolling_counter_rpm_reflects_recorded_requests():
|
|
from meshnet_tracker.server import _RollingCounter
|
|
|
|
counter = _RollingCounter(num_buckets=60, bucket_seconds=60)
|
|
now = 1_000_000.0
|
|
for _ in range(60):
|
|
counter.record(now)
|
|
|
|
rpm = counter.rpm(now)
|
|
assert rpm == 1.0 # 60 requests over 60 minutes = 1 rpm
|
|
|
|
|
|
def test_stats_collector_records_and_returns_rpms():
|
|
from meshnet_tracker.server import _StatsCollector
|
|
|
|
collector = _StatsCollector()
|
|
now = 1_000_000.0
|
|
for _ in range(30):
|
|
collector.record_request("my-model", now=now)
|
|
|
|
combined = collector.get_combined_stats(now=now)
|
|
assert "my-model" in combined
|
|
assert combined["my-model"]["rpm_last_hour"] > 0
|
|
assert combined["my-model"]["rpm_last_day"] > 0
|
|
assert combined["my-model"]["rpm_last_month"] > 0
|
|
|
|
|
|
def test_stats_collector_merges_peer_rpms_additively():
|
|
from meshnet_tracker.server import _StatsCollector
|
|
|
|
collector = _StatsCollector()
|
|
now = 1_000_000.0
|
|
collector.record_request("modelX", now=now) # 1 local request in the current minute bucket
|
|
|
|
peer_rpms = {"modelX": {"rpm_last_hour": 5.0, "rpm_last_day": 3.0, "rpm_last_month": 1.0}}
|
|
collector.merge_peer_rpms("http://peer-tracker:8080", peer_rpms)
|
|
|
|
combined = collector.get_combined_stats(now=now)
|
|
assert combined["modelX"]["rpm_last_hour"] > 5.0 # local + peer
|
|
assert combined["modelX"]["rpm_last_day"] > 3.0
|
|
assert combined["modelX"]["rpm_last_month"] >= 1.0 # local rounds to 0 at 4dp, peer is 1.0
|
|
|
|
|
|
def test_stats_sqlite_persistence_survives_restart(tmp_path):
|
|
from meshnet_tracker.server import _StatsCollector
|
|
|
|
db = str(tmp_path / "stats.db")
|
|
now = 1_000_000.0
|
|
|
|
collector = _StatsCollector(db_path=db)
|
|
for _ in range(10):
|
|
collector.record_request("persistent-model", now=now)
|
|
collector.save_to_db()
|
|
|
|
collector2 = _StatsCollector(db_path=db)
|
|
combined = collector2.get_combined_stats(now=now)
|
|
assert "persistent-model" in combined
|
|
assert combined["persistent-model"]["rpm_last_hour"] > 0
|
|
|
|
|
|
def test_stats_gossip_endpoint_merges_peer_slice():
|
|
tracker = TrackerServer()
|
|
port = tracker.start()
|
|
try:
|
|
peer_payload = {
|
|
"tracker_url": "http://peer:9000",
|
|
"stats": {
|
|
"remote-model": {
|
|
"rpm_last_hour": 2.5,
|
|
"rpm_last_day": 1.0,
|
|
"rpm_last_month": 0.5,
|
|
}
|
|
},
|
|
}
|
|
_post_json(f"http://127.0.0.1:{port}/v1/stats/gossip", peer_payload)
|
|
|
|
combined = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
|
assert "remote-model" in combined["models"]
|
|
assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5
|
|
finally:
|
|
tracker.stop()
|
|
|
|
|
|
# ---------------------------------------------------------------- US-027 / US-028 routing
|
|
|
|
def _make_node(node_id, shard_start, shard_end, bench=1.0, queue_depth=0):
|
|
"""Helper: build a _NodeEntry with minimal fields."""
|
|
from meshnet_tracker.server import _NodeEntry
|
|
n = _NodeEntry(
|
|
node_id=node_id, endpoint=f"http://fake:{9000 + shard_start}",
|
|
shard_start=shard_start, shard_end=shard_end,
|
|
model="m", shard_checksum=None, hardware_profile={},
|
|
wallet_address=None, score=1.0,
|
|
benchmark_tokens_per_sec=bench,
|
|
)
|
|
n.queue_depth = queue_depth
|
|
return n
|
|
|
|
|
|
def test_select_route_overlapping_shards_uses_start_layer_protocol():
|
|
"""_select_route with overlapping shards: A(0-22) and B(20-24).
|
|
Route must be [A, B]; B's effective start_layer = 23 (after A ends at 22).
|
|
"""
|
|
from meshnet_tracker.server import _select_route
|
|
|
|
node_a = _make_node("A", 0, 22)
|
|
node_b = _make_node("B", 20, 24)
|
|
|
|
route, err = _select_route([node_a, node_b], 0, 24)
|
|
assert err == ""
|
|
assert len(route) == 2
|
|
assert route[0].node_id == "A"
|
|
assert route[1].node_id == "B"
|
|
|
|
|
|
def test_select_route_no_overlap_three_nodes():
|
|
"""Non-overlapping shards are covered in the right order."""
|
|
from meshnet_tracker.server import _select_route
|
|
|
|
nodes = [_make_node("C", 8, 15), _make_node("A", 0, 7), _make_node("B", 16, 23)]
|
|
route, err = _select_route(nodes, 0, 23)
|
|
assert err == ""
|
|
# Greedy cover order: A(0-7) → C(8-15) → B(16-23)
|
|
assert [n.node_id for n in route] == ["A", "C", "B"]
|
|
|
|
|
|
def test_select_route_gap_returns_error():
|
|
from meshnet_tracker.server import _select_route
|
|
|
|
nodes = [_make_node("A", 0, 5), _make_node("B", 10, 15)]
|
|
route, err = _select_route(nodes, 0, 15)
|
|
assert route == []
|
|
assert "layer 6" in err
|
|
|
|
|
|
def test_select_route_prefers_higher_throughput_node_when_shards_equal():
|
|
"""When two nodes cover the same interval, prefer the faster / less-loaded one."""
|
|
from meshnet_tracker.server import _select_route
|
|
|
|
fast = _make_node("fast", 0, 11, bench=10.0, queue_depth=0) # effective = 10.0
|
|
slow = _make_node("slow", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
|
|
|
route, err = _select_route([slow, fast], 0, 11)
|
|
assert err == ""
|
|
assert len(route) == 1
|
|
assert route[0].node_id == "fast"
|
|
|
|
|
|
def test_select_route_throughput_accounts_for_queue_depth():
|
|
"""A nominally faster node with high queue depth loses to a slower idle node."""
|
|
from meshnet_tracker.server import _select_route
|
|
|
|
busy_fast = _make_node("busy", 0, 11, bench=20.0, queue_depth=9) # effective = 2.0
|
|
idle_slow = _make_node("idle", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
|
|
|
route, err = _select_route([busy_fast, idle_slow], 0, 11)
|
|
assert err == ""
|
|
# tie → first in list wins; exact tie does not matter as long as no error
|
|
assert len(route) == 1
|
|
assert route[0].node_id in ("busy", "idle")
|
|
|
|
|
|
def test_two_stub_nodes_complete_pipeline_via_tracker():
|
|
"""Integration: two StubNodeServer instances serving complementary shards
|
|
produce a full inference response through the tracker route."""
|
|
tracker = TrackerServer(model_presets={
|
|
"two-node-model": {"layers_start": 0, "layers_end": 3}
|
|
})
|
|
tracker_port = tracker.start()
|
|
|
|
from meshnet_node.server import StubNodeServer
|
|
node_a = StubNodeServer()
|
|
node_b = StubNodeServer()
|
|
port_a = node_a.start()
|
|
port_b = node_b.start()
|
|
|
|
try:
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
|
"endpoint": f"http://127.0.0.1:{port_a}",
|
|
"model": "two-node-model", "shard_start": 0, "shard_end": 1,
|
|
"hardware_profile": {}, "score": 1.0,
|
|
})
|
|
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
|
"endpoint": f"http://127.0.0.1:{port_b}",
|
|
"model": "two-node-model", "shard_start": 2, "shard_end": 3,
|
|
"hardware_profile": {}, "score": 1.0,
|
|
})
|
|
|
|
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=two-node-model")
|
|
assert len(route_resp["route"]) == 2
|
|
finally:
|
|
node_a.stop()
|
|
node_b.stop()
|
|
tracker.stop()
|
|
|
|
|
|
def test_route_timeout_config_is_exposed_on_server():
|
|
"""TorchNodeServer exposes the route_timeout that was configured."""
|
|
from meshnet_node.torch_server import TorchNodeServer
|
|
|
|
class _MinimalBackend:
|
|
model_id = "m"
|
|
total_layers = 4
|
|
is_head = True
|
|
is_tail = True
|
|
def generate_text(self, *a, **kw): return ""
|
|
def count_prompt_tokens(self, *a): return 0
|
|
def count_text_tokens(self, *a): return 0
|
|
|
|
node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0)
|
|
assert node.route_timeout == 45.0
|
|
|
|
|
|
def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
|
|
from meshnet_node import torch_server
|
|
from meshnet_node.torch_server import TorchNodeServer
|
|
|
|
class _MinimalBackend:
|
|
def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=21, quantization="bfloat16"):
|
|
self.model_id = model_id
|
|
self.shard_start = shard_start
|
|
self.shard_end = shard_end
|
|
self.quantization = quantization
|
|
self.total_layers = 24
|
|
self.is_head = shard_start == 0
|
|
self.is_tail = shard_end == 23
|
|
|
|
def generate_text(self, *a, **kw): return ""
|
|
def count_prompt_tokens(self, *a): return 0
|
|
def count_text_tokens(self, *a): return 0
|
|
|
|
loaded = []
|
|
|
|
def fake_load_backend(model_id, shard_start, shard_end, quantization):
|
|
loaded.append((model_id, shard_start, shard_end, quantization))
|
|
return _MinimalBackend(model_id, shard_start, shard_end, quantization)
|
|
|
|
monkeypatch.setattr(torch_server, "_load_backend", fake_load_backend)
|
|
node = TorchNodeServer(backend=_MinimalBackend())
|
|
|
|
applied = node.apply_tracker_directives([
|
|
{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 21},
|
|
{"action": "LOAD_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 23,
|
|
"quantization": "bfloat16"},
|
|
])
|
|
|
|
assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
|
|
assert applied == {
|
|
"model": "Qwen/Qwen2.5-0.5B-Instruct",
|
|
"shard_start": 0,
|
|
"shard_end": 23,
|
|
"quantization": "bfloat16",
|
|
"tracker_mode": True,
|
|
}
|
|
assert node.backend.shard_end == 23
|