popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
4c6e1ed8b6ff2a37adc423ee1c48c0c223abaae1
neuron-tai/tests/test_tracker_routing.py

2848 lines
106 KiB
Python
Raw Blame History

"""US-003 tests: tracker node registration and route selection."""
import http.server
import json
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import pytest
from meshnet_gateway.server import GatewayServer, _banned_route_wallet
from meshnet_node.server import StubNodeServer
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.auth import sign_hive_request
from meshnet_tracker.server import (
TrackerServer,
_NodeEntry,
_memory_pool_map,
_registration_ban_error,
_scale_demanded_models_locked,
)
_TEST_HIVE_SECRET = "test-hive-secret"
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_exposes_registered_model_metadata():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
try:
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7100",
"model": "Kimi-K2.7-Code",
"hf_repo": "unsloth/Kimi-K2.7-Code",
"num_layers": 61,
"shard_start": 0,
"shard_end": 60,
"hardware_profile": {},
"model_metadata": {
"total_parameters": "1T",
"activated_parameters": "32B",
"context_length": 256000,
},
"downloaded_models": [{
"model": "Kimi-K2.7-Code",
"hf_repo": "unsloth/Kimi-K2.7-Code",
"shard_start": 0,
"shard_end": 60,
"local_model_percentage": 100.0,
}],
},
)
models = _get_json(f"{url}/v1/models")
network_map = _get_json(f"{url}/v1/network/map")
finally:
tracker.stop()
kimi = next(model for model in models["data"] if model["id"] == "kimi-k2.7")
assert kimi["hf_repo"] == "unsloth/Kimi-K2.7-Code"
assert kimi["metadata"]["total_parameters"] == "1T"
assert kimi["metadata"]["activated_parameters"] == "32B"
assert kimi["metadata"]["num_layers"] == 61
registered = network_map["nodes"][0]
assert registered["num_layers"] == 61
assert registered["model_metadata"]["context_length"] == 256000
assert registered["downloaded_models"][0]["local_model_percentage"] == 100.0
def test_tracker_lists_recommended_kimi_before_nodes_register():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
try:
models = _get_json(f"{url}/v1/models")
network_map = _get_json(f"{url}/v1/network/map")
finally:
tracker.stop()
kimi = next(model for model in models["data"] if model["id"] == "kimi-k2.7")
assert kimi["hf_repo"] == "unsloth/Kimi-K2.7-Code"
assert "Kimi-K2.7-Code" in kimi["aliases"]
assert kimi["metadata"]["download_size_gb"] == 595
assert kimi["deployment"]["native_quantization"] == "int4"
assert kimi["deployment"]["deployable"] is False
assert network_map["pool"]["node_count"] == 0
assert network_map["recommended_models"][0]["id"] == "kimi-k2.7"
def test_network_map_exposes_pool_size_and_speed_summary():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
try:
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7101",
"model": "inventory-a",
"vram_bytes": 10_000,
"ram_bytes": 20_000,
"benchmark_tokens_per_sec": 4.0,
"hardware_profile": {},
},
)
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7102",
"model": "inventory-b",
"vram_bytes": 0,
"ram_bytes": 30_000,
"benchmark_tokens_per_sec": 6.0,
"hardware_profile": {},
},
)
network_map = _get_json(f"{url}/v1/network/map")
finally:
tracker.stop()
pool = network_map["pool"]
assert pool["node_count"] == 2
assert pool["total_vram_bytes"] == 10_000
assert pool["total_ram_bytes"] == 50_000
assert pool["total_memory_budget_bytes"] == 40_000
assert pool["total_benchmark_tokens_per_sec"] == 10.0
assert pool["total_effective_throughput"] == 10.0
def test_network_map_exposes_served_model_copy_count():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
try:
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7201",
"model": "copy-count-test",
"hf_repo": "example/copy-count-test",
"num_layers": 37,
"shard_start": 0,
"shard_end": 21,
"hardware_profile": {},
},
)
network_map = _get_json(f"{url}/v1/network/map")
assert network_map["nodes"][0]["model_supply"]["served_model_copies"] == 0.59
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7202",
"model": "copy-count-test",
"hf_repo": "example/copy-count-test",
"num_layers": 37,
"shard_start": 22,
"shard_end": 36,
"hardware_profile": {},
},
)
network_map = _get_json(f"{url}/v1/network/map")
assert network_map["nodes"][0]["model_supply"]["served_model_copies"] == 1.0
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:7203",
"model": "copy-count-test",
"hf_repo": "example/copy-count-test",
"num_layers": 37,
"shard_start": 0,
"shard_end": 36,
"hardware_profile": {},
},
)
network_map = _get_json(f"{url}/v1/network/map")
assert network_map["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
finally:
tracker.stop()
def test_recommended_kimi_becomes_deployable_when_pool_is_large_enough():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
gib = 1024 * 1024 * 1024
try:
for idx in range(2):
_post_json(
f"{url}/v1/nodes/register",
{
"endpoint": f"http://127.0.0.1:{7200 + idx}",
"model": f"inventory-{idx}",
"vram_bytes": 0,
"ram_bytes": 400 * gib,
"benchmark_tokens_per_sec": 5.0,
"hardware_profile": {},
},
)
models = _get_json(f"{url}/v1/models")
finally:
tracker.stop()
kimi = next(model for model in models["data"] if model["id"] == "kimi-k2.7")
assert kimi["deployment"]["deployable"] is True
assert kimi["deployment"]["missing_effective_memory_bytes"] == 0
def test_network_assign_can_start_recommended_kimi_by_short_name():
tracker = TrackerServer()
port = tracker.start()
url = f"http://127.0.0.1:{port}"
try:
assignment = _get_json(
f"{url}/v1/network/assign?device=cpu&ram_mb=204800&hf_repo=kimi-k2.7"
)
finally:
tracker.stop()
assert assignment["model"] == "kimi-k2.7"
assert assignment["hf_repo"] == "unsloth/Kimi-K2.7-Code"
assert assignment["num_layers"] == 61
assert assignment["shard_start"] == 0
assert 0 <= assignment["shard_end"] < 60
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_route_log_counts_proxy_inflight_requests():
entered = threading.Event()
release = threading.Event()
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)
entered.set()
release.wait(timeout=3.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()
errors = []
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": "burst-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
def call_proxy():
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "burst-model", "messages": [{"role": "user", "content": "hi"}]},
)
except Exception as exc:
errors.append(exc)
first = threading.Thread(target=call_proxy)
second = threading.Thread(target=call_proxy)
first.start()
assert entered.wait(timeout=2.0)
second.start()
selected_events = []
deadline = time.time() + 2.0
while time.time() < deadline:
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
selected_events = [
event for event in console["events"]
if event["message"] == "proxy route selected"
]
if len(selected_events) >= 2:
break
time.sleep(0.05)
assert len(selected_events) >= 2
second_nodes = selected_events[-1]["fields"]["nodes"]
assert second_nodes[0]["queue_depth"] == 2
assert second_nodes[0]["proxy_inflight"] == 2
finally:
release.set()
first.join(timeout=3.0)
second.join(timeout=3.0)
tracker.stop()
slow_node.shutdown()
slow_node.server_close()
slow_thread.join(timeout=1.0)
assert not first.is_alive()
assert not second.is_alive()
assert not errors
def test_tracker_logs_stream_progress_before_request_completes():
chunk_sent = threading.Event()
release = threading.Event()
class StreamingChatHandler(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
self.rfile.read(int(self.headers.get("Content-Length", 0)))
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.end_headers()
payload = json.dumps({
"choices": [{"delta": {"content": "hello world"}}],
}).encode()
self.wfile.write(b"data: " + payload + b"\n\n")
self.wfile.flush()
chunk_sent.set()
release.wait(timeout=3.0)
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
response = None
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "stream-progress-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
req = urllib.request.Request(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
data=json.dumps({
"model": "stream-progress-model",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
}).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
response = urllib.request.urlopen(req, timeout=3.0)
assert response.headers["X-Accel-Buffering"] == "no"
first_line = response.readline()
assert first_line.startswith(b"data:")
assert chunk_sent.wait(timeout=1.0)
progress_events = []
deadline = time.time() + 2.0
while time.time() < deadline:
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
progress_events = [
event for event in console["events"]
if event["message"] == "proxy progress"
]
if progress_events:
break
time.sleep(0.05)
assert progress_events
fields = progress_events[-1]["fields"]
assert fields["tokens"] == 2
assert fields["tokens_per_sec"] > 0
active = [
event for event in console["events"]
if event["message"] == "proxy route selected"
]
assert active
finally:
release.set()
if response is not None:
response.close()
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
def test_tracker_stream_survives_idle_gap_between_sse_chunks():
first_chunk_sent = threading.Event()
class IdleStreamingChatHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, format, *args):
pass
def do_POST(self):
if self.path != "/v1/chat/completions":
self.send_response(404)
self.end_headers()
return
self.rfile.read(int(self.headers.get("Content-Length", 0)))
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.end_headers()
first = json.dumps({
"choices": [{"delta": {"content": "hello"}}],
}).encode()
second = json.dumps({
"choices": [{"delta": {"content": " world"}}],
}).encode()
self.wfile.write(b"data: " + first + b"\n\n")
self.wfile.flush()
first_chunk_sent.set()
time.sleep(1.0)
self.wfile.write(b"data: " + second + b"\n\n")
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
node = http.server.HTTPServer(("127.0.0.1", 0), IdleStreamingChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
response = None
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "idle-stream-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
req = urllib.request.Request(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
data=json.dumps({
"model": "idle-stream-model",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
}).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
response = urllib.request.urlopen(req, timeout=3.0)
assert response.readline().startswith(b"data:")
assert first_chunk_sent.wait(timeout=1.0)
remaining = response.read().splitlines()
assert b"data: [DONE]" in remaining
finally:
if response is not None:
response.close()
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
def test_tracker_dashboard_can_cancel_inflight_proxy():
chunk_sent = threading.Event()
release = threading.Event()
class StreamingChatHandler(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
self.rfile.read(int(self.headers.get("Content-Length", 0)))
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.end_headers()
payload = json.dumps({
"choices": [{"delta": {"content": "hello world"}}],
}).encode()
self.wfile.write(b"data: " + payload + b"\n\n")
self.wfile.flush()
chunk_sent.set()
release.wait(timeout=3.0)
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer(heartbeat_timeout=60.0)
tracker_port = tracker.start()
response = None
request_id = None
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "cancel-proxy-model", "num_layers": 1,
"shard_start": 0, "shard_end": 0,
"hardware_profile": {}, "score": 1.0},
)
req = urllib.request.Request(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
data=json.dumps({
"model": "cancel-proxy-model",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
}).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
response = urllib.request.urlopen(req, timeout=3.0)
first_line = response.readline()
assert first_line.startswith(b"data:")
assert chunk_sent.wait(timeout=1.0)
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
selected = [
event for event in console["events"]
if event["message"] == "proxy route selected"
]
assert selected
request_id = selected[-1]["fields"]["request_id"]
cancel = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/proxy/requests/{urllib.parse.quote(request_id, safe='')}/cancel",
{},
)
assert cancel["status"] == "canceled"
deadline = time.time() + 5.0
canceled_events = []
while time.time() < deadline:
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
canceled_events = [
event for event in console["events"]
if event["message"] == "proxy canceled"
and event["fields"].get("request_id") == request_id
]
if canceled_events:
break
time.sleep(0.05)
assert canceled_events
finally:
release.set()
if response is not None:
response.close()
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
def test_tracker_routes_hf_model_alias_from_quickstart():
"""The documented qwen2.5-0.5b alias resolves a full HF repo registration."""
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:9100",
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
route_resp = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/route?model=qwen2.5-0.5b"
)
finally:
tracker.stop()
assert route_resp["route"] == ["http://127.0.0.1:9100"]
def test_tracker_proxy_accepts_hf_model_alias_from_quickstart():
"""The tracker proxy accepts the same model alias used by the quickstart curl."""
class ChatHandler(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))
request_body = json.loads(self.rfile.read(length) or b"{}")
body = json.dumps({
"model": request_body["model"],
"choices": [{"message": {"content": "56"}}],
}).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)
node = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
tracker = TrackerServer()
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "qwen2.5-0.5b",
"messages": [{"role": "user", "content": "What is 7 times 8?"}]},
)
finally:
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
assert response["choices"][0]["message"]["content"] == "56"
def test_tracker_proxy_routes_split_preset_nodes_by_alias():
"""The chat proxy must build routes from alias-matched preset nodes."""
class ChatHandler(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)
route_header = self.headers.get("X-Meshnet-Route", "[]")
body = json.dumps({
"choices": [{"message": {"content": route_header}}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
}).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)
head = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
head_thread = threading.Thread(target=head.serve_forever, daemon=True)
head_thread.start()
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:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{head.server_address[1]}",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 21,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9199",
"model": "Qwen3.6-35B-A3B",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 22,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _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")
finally:
tracker.stop()
head.shutdown()
head.server_close()
head_thread.join(timeout=1.0)
route = json.loads(response["choices"][0]["message"]["content"])
assert route == [{"endpoint": "http://127.0.0.1:9199", "start_layer": 22}]
complete = [event for event in console["events"] if event["message"] == "proxy complete"][-1]
assert complete["fields"]["tokens"] == 2
assert complete["fields"]["tokens_per_sec"] >= 0
def test_tracker_route_endpoint_routes_split_preset_nodes_by_alias():
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:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9101",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 21,
"tracker_mode": True,
"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": "Qwen3.6-35B-A3B",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 22,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen3.6-35B-A3B"
)
finally:
tracker.stop()
assert response["route"] == ["http://127.0.0.1:9101", "http://127.0.0.1:9102"]
assert [node["start_layer"] for node in response["nodes"]] == [0, 22]
def test_tracker_route_endpoint_ignores_model_case_and_outer_whitespace():
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:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9101",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/route?model=%20Qwen3.6-35B-A3B%20"
)
finally:
tracker.stop()
assert response["route"] == ["http://127.0.0.1:9101"]
def test_tracker_route_prefers_distributed_over_single_full_shard():
"""When a full 0-39 node and a partial 0-21 head coexist, /v1/route
should return both hops — not the full shard alone."""
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:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://192.168.0.179:7000",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://192.168.0.20:7000",
"model": "Qwen3.6-35B-A3B",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 21,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/route?model=qwen3.6-35b-a3b"
)
finally:
tracker.stop()
assert response["route"] == [
"http://192.168.0.20:7000",
"http://192.168.0.179:7000",
]
assert [node["start_layer"] for node in response["nodes"]] == [0, 22]
def test_tracker_proxy_ignores_model_case_and_outer_whitespace():
class ChatHandler(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))
request_body = json.loads(self.rfile.read(length) or b"{}")
body = json.dumps({
"model": request_body["model"],
"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)
node = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
node_thread.start()
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:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
"model": "qwen3.6-35b-a3b",
"hf_repo": "unsloth/Qwen3.6-35B-A3B",
"num_layers": 40,
"shard_start": 0,
"shard_end": 39,
"tracker_mode": True,
"hardware_profile": {},
"score": 1.0},
)
response = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "Qwen3.6-35B-A3B ",
"messages": [{"role": "user", "content": "hi"}]},
)
finally:
tracker.stop()
node.shutdown()
node.server_close()
node_thread.join(timeout=1.0)
assert response["choices"][0]["message"]["content"] == "ok"
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},
)
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
widths = {
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
for node in net["nodes"]
}
assert widths["http://127.0.0.1:9011"] == 12
assert widths["http://127.0.0.1:9012"] == 12
assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
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_pool_join_adds_redundant_copy_without_splitting_incumbent():
"""A second managed node with capacity for the full model keeps the first copy intact."""
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 4,
"bytes_per_layer": {"bfloat16": 1_000},
},
})
tracker_port = tracker.start()
try:
first = _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},
)
second = _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},
)
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
widths = {
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
for node in net["nodes"]
}
assert widths["http://127.0.0.1:9015"] == 4
assert widths["http://127.0.0.1:9016"] == 4
assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {})
assert hb.get("directives", []) == []
finally:
tracker.stop()
def test_tracker_explicit_full_copy_join_preserves_existing_serving_node():
"""--model style joins with explicit shards add redundancy instead of reshuffling."""
tracker = TrackerServer(heartbeat_timeout=10.0)
tracker_port = tracker.start()
try:
base_reg = {
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
"managed_assignment": True,
"vram_bytes": 2_000_000_000,
"hardware_profile": {},
"score": 1.0,
}
first = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:9201"},
)
second = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:9202"},
)
coverage = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/network/map"
)
assert coverage["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
hb = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {}
)
assert hb.get("directives", []) == []
assert second["node_id"] in tracker._registry
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": 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:9002", "model": "tiny-model",
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
)
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
heads = [node for node in net["nodes"] if node["shard_start"] == 0]
assert len(heads) == 1
assert heads[0]["endpoint"] == "http://127.0.0.1:9002"
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_stores_current_requests():
"""Node-reported in-flight request snapshots appear on the network map."""
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",
"model": "progress-model",
"shard_start": 0,
"shard_end": 31,
"hardware_profile": {},
"score": 1.0,
},
)
node_id = reg["node_id"]
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat",
{
"queue_depth": 1,
"current_requests": [{
"request_id": "req-abc123",
"model": "progress-model",
"kind": "chat",
"tokens": 17,
"elapsed_seconds": 42.5,
"tokens_per_sec": 0.4,
"routing_complete": True,
}],
},
)
network = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
node = next(item for item in network["nodes"] if item["node_id"] == node_id)
assert node["stats"]["queue_depth"] == 1
assert node["stats"]["current_requests"] == [{
"request_id": "req-abc123",
"model": "progress-model",
"kind": "chat",
"tokens": 17,
"elapsed_seconds": 42.5,
"tokens_per_sec": 0.4,
"routing_complete": True,
}]
finally:
tracker.stop()
def test_normalize_current_requests_sanitizes_payload():
from meshnet_tracker.server import _normalize_current_requests
assert _normalize_current_requests(None) == []
assert _normalize_current_requests([
{"request_id": "req-1", "model": "m", "tokens": "9", "tokens_per_sec": "1.5"},
{"model": "missing-id"},
"bad",
]) == [{
"request_id": "req-1",
"model": "m",
"tokens": 9,
"tokens_per_sec": 1.5,
}]
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_console_records_model_not_available():
tracker = TrackerServer()
tracker_port = tracker.start()
try:
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
{"model": "missing-model", "messages": [{"role": "user", "content": "hi"}]},
)
raise AssertionError("Expected 503")
except urllib.error.HTTPError as exc:
assert exc.code == 503
console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console")
finally:
tracker.stop()
assert any(
event["message"] == "no nodes available for model"
and event["fields"]["model"] == "missing-model"
for event in console["events"]
)
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_registration_rejects_invalid_max_loaded_shards():
tracker = TrackerServer()
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:9001", "shard_start": 0, "shard_end": 31,
"max_loaded_shards": 0, "hardware_profile": {}, "score": 1.0},
)
raise AssertionError("Expected 400 for invalid max_loaded_shards")
except urllib.error.HTTPError as exc:
assert exc.code == 400
body = json.loads(exc.read())
assert "capability values" in body["error"]
finally:
tracker.stop()
def test_network_map_exposes_node_capacity_limits():
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:9021", "model": "tiny-model",
"shard_start": 0, "shard_end": 7,
"vram_bytes": 10_000, "ram_bytes": 20_000, "max_loaded_shards": 2,
"quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 4.0,
"hardware_profile": {}, "score": 1.0},
)
network_map = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
node = network_map["nodes"][0]
assert node["capacity"]["vram_bytes"] == 10_000
assert node["capacity"]["ram_bytes"] == 20_000
assert node["capacity"]["max_loaded_shards"] == 2
assert node["capacity"]["max_assignable_layers"] == 8
assert node["capacity"]["benchmark_tokens_per_sec"] == 4.0
finally:
tracker.stop()
def test_tracker_capacity_uses_ram_when_node_has_no_vram():
"""CPU-only nodes should expose RAM-backed shard capacity, not default GPU capacity."""
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 20,
"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:9020", "model": "tiny-model",
"vram_bytes": 0, "ram_bytes": 16_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
network_map = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
capacity = network_map["nodes"][0]["capacity"]
assert capacity["vram_bytes"] == 0
assert capacity["ram_bytes"] == 16_000
assert capacity["memory_budget_bytes"] == 16_000
assert capacity["memory_budget_source"] == "ram"
assert capacity["max_assignable_layers"] == 12
finally:
tracker.stop()
def test_rebalance_keeps_one_active_range_even_when_multiple_slots_advertised():
"""max_loaded_shards is exposed but reserved until node runtime supports multi-range serving."""
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 8,
"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:9022", "model": "tiny-model",
"vram_bytes": 20_000, "ram_bytes": 20_000, "max_loaded_shards": 2,
"quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0,
"hardware_profile": {}, "score": 1.0},
)
assert reg["directive"]["action"] == "LOAD_SHARD"
assert reg["directive"]["start_layer"] == 0
assert reg["directive"]["end_layer"] == 7
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_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_records_node_model_tokens_per_second():
from meshnet_tracker.server import _StatsCollector
collector = _StatsCollector()
now = 1_000_000.0
collector.record_node_throughput("node-a", "my-model", total_tokens=200, elapsed_seconds=2.0, now=now)
collector.record_node_throughput("node-a", "my-model", total_tokens=300, elapsed_seconds=3.0, now=now)
node_stats = collector.get_node_model_stats("node-a", "my-model", now=now)
assert node_stats["tokens_per_sec_last_hour"] == pytest.approx(100.0)
assert node_stats["tokens_per_sec_last_day"] == pytest.approx(100.0)
assert node_stats["tokens_per_sec_last_month"] == pytest.approx(100.0)
assert node_stats["sample_count_last_hour"] == 2
node_models = collector.get_node_throughput_stats(now=now)
assert node_models["node-a"]["models"]["my-model"]["tokens_per_sec_last_hour"] == pytest.approx(100.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.record_node_throughput("node-p", "persistent-model", total_tokens=120, elapsed_seconds=2.0, 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
node_models = collector2.get_node_throughput_stats(now=now)
assert node_models["node-p"]["models"]["persistent-model"]["tokens_per_sec_last_hour"] == pytest.approx(60.0)
def test_stats_gossip_endpoint_merges_peer_slice():
tracker = TrackerServer(hive_secret=_TEST_HIVE_SECRET)
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,
}
},
}
body = json.dumps(peer_payload).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{port}/v1/stats/gossip", data=body,
headers={"Content-Type": "application/json", **sign_hive_request(_TEST_HIVE_SECRET, body)},
method="POST",
)
with urllib.request.urlopen(req) as r:
r.read()
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, wallet_address=None):
"""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=wallet_address, 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_prefers_observed_model_throughput_over_startup_benchmark():
"""Observed per-model TPS is the primary routing signal once available."""
from meshnet_tracker.server import _select_route
high_benchmark = _make_node("bench-fast", 0, 11, bench=100.0, queue_depth=0)
observed_fast = _make_node("observed-fast", 0, 11, bench=1.0, queue_depth=0)
high_benchmark.model_tokens_per_sec["m"] = 10.0
observed_fast.model_tokens_per_sec["m"] = 40.0
route, err = _select_route([high_benchmark, observed_fast], 0, 11, model="m")
assert err == ""
assert route[0].node_id == "observed-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")
# ---------------------------------------------------------------- AH-009 reputation-weighted routing
def test_select_route_prefers_higher_reputation_when_throughput_equal():
"""AH-009 test-first item 3: among candidates that advance coverage
equally with the same effective throughput, the higher-reputation
wallet's node wins the tiebreak (ADR-0018 §6 -- earnings scale with
tenure/standing)."""
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.server import _select_route
contracts = LocalSolanaContracts()
contracts.registry.record_audit_outcome("wallet-veteran", passed=False) # 1.0 -> 0.7
# wallet-fresh keeps the default reputation of 1.0
veteran = _make_node("veteran", 0, 11, bench=10.0, wallet_address="wallet-veteran")
fresh = _make_node("fresh", 0, 11, bench=10.0, wallet_address="wallet-fresh")
route, err = _select_route([veteran, fresh], 0, 11, contracts=contracts)
assert err == ""
assert route[0].node_id == "fresh"
def test_select_route_reputation_never_overrides_coverage():
"""A lower-reputation node that covers more layers still wins -- the
reputation multiplier is a tiebreak among equal-coverage candidates,
never a substitute for coverage maximization."""
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.server import _select_route
contracts = LocalSolanaContracts()
contracts.registry.record_audit_outcome("wallet-shady", passed=False)
for _ in range(9):
contracts.registry.record_audit_outcome("wallet-shady", passed=False) # reputation -> 0.0
wide_shady = _make_node("wide-shady", 0, 11, bench=1.0, wallet_address="wallet-shady")
narrow_good = _make_node("narrow-good", 0, 5, bench=1.0, wallet_address="wallet-good")
route, err = _select_route([wide_shady, narrow_good], 0, 11, contracts=contracts)
assert err == ""
assert [n.node_id for n in route] == ["wide-shady"]
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 == {
"action": "LOAD_SHARD",
"model": "Qwen/Qwen2.5-0.5B-Instruct",
"shard_start": 0,
"shard_end": 23,
"quantization": "bfloat16",
"tracker_mode": True,
}
assert node.backend.shard_end == 23
def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch):
"""End-to-end heal: kill one managed node, surviving node receives LOAD_SHARD and hot-swaps.
Cycle:
1. Two managed nodes (A: 0-11, B: 12-23) register with tracker.
2. Node A stops heartbeating; tracker expires it and triggers rebalance.
3. Node B's next heartbeat response contains LOAD_SHARD(0, 23).
4. Node B (TorchNodeServer) applies the directive — backend hot-swapped.
5. Coverage endpoint confirms full model is covered by Node B alone.
"""
from meshnet_node import torch_server
from meshnet_node.torch_server import TorchNodeServer
# --- minimal fake backend (no GPU / PyTorch needed) ---
class _FakeBackend:
def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=23, quantization="int8"):
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_shards: list[tuple] = []
def fake_load(model_id, shard_start, shard_end, quantization):
loaded_shards.append((model_id, shard_start, shard_end))
return _FakeBackend(model_id, shard_start, shard_end, quantization)
monkeypatch.setattr(torch_server, "_load_backend", fake_load)
# Use a very short timeout so Node A expires quickly.
tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0)
tracker_port = tracker.start()
node_b = TorchNodeServer(backend=_FakeBackend(shard_start=12, shard_end=23))
base_reg = {
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"model": "Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"vram_bytes": 2_000_000_000,
"ram_bytes": 0,
"hardware_profile": {},
"score": 1.0,
"managed_assignment": True,
}
try:
# Step 1: register both nodes as managed.
reg_a = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19001", "shard_start": 0, "shard_end": 11},
)
reg_b = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19002", "shard_start": 12, "shard_end": 23},
)
node_a_id = reg_a["node_id"]
node_b_id = reg_b["node_id"]
# Initial heartbeat to mark both alive.
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_a_id}/heartbeat", {})
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {})
# Step 2: let Node A's heartbeat expire (only Node B keeps heartbeating).
time.sleep(0.10)
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {})
time.sleep(0.10)
# Step 3: Node B's heartbeat triggers purge of A and gets LOAD_SHARD.
hb_resp = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {}
)
directives = hb_resp.get("directives", [])
load_dirs = [d for d in directives if d["action"] == "LOAD_SHARD"]
assert load_dirs, f"Expected LOAD_SHARD directive, got: {directives}"
assert load_dirs[-1]["shard_start"] == 0
assert load_dirs[-1]["shard_end"] == 23
assert node_a_id not in tracker._registry
# Step 4: Node B applies the directive — backend hot-swapped.
applied = node_b.apply_tracker_directives(directives)
assert applied is not None
assert applied["shard_start"] == 0
assert applied["shard_end"] == 23
assert applied["tracker_mode"] is True
assert node_b.backend.shard_start == 0
assert node_b.backend.shard_end == 23
assert loaded_shards == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23)]
# Step 5: re-register Node B with its new shard so tracker reflects healed state.
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19002", "shard_start": 0, "shard_end": 23},
)
coverage_resp = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen%2FQwen2.5-0.5B-Instruct"
)
assert all(seg["node_count"] >= 1 for seg in coverage_resp["coverage"]), (
f"Coverage gap after heal: {coverage_resp['coverage']}"
)
finally:
tracker.stop()
def test_network_map_exposes_memory_pool():
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 8,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/TinyModel",
},
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9050", "model": "tiny-model",
"hf_repo": "org/TinyModel", "num_layers": 8,
"shard_start": 0, "shard_end": 7, "max_loaded_shards": 2,
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
pool = net["memory_pool"]
assert pool["total_spare_slots"] == 1
assert pool["hosts"][0]["loaded_slots"] == 1
assert pool["hosts"][0]["max_loaded_shards"] == 2
assert pool["hosts"][0]["memory_spare_bytes"] > 0
finally:
tracker.stop()
def test_same_endpoint_can_register_multiple_models():
tracker = TrackerServer()
tracker_port = tracker.start()
try:
base = {
"endpoint": "http://127.0.0.1:9055",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
"hardware_profile": {},
"score": 1.0,
"max_loaded_shards": 2,
"vram_bytes": 50_000_000,
"ram_bytes": 50_000_000,
}
first = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base, "model": "Qwen2.5-0.5B-Instruct", "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct"},
)
second = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base, "model": "OtherModel", "hf_repo": "org/OtherModel"},
)
assert first["node_id"] != second["node_id"]
assert len(tracker._registry) == 2
finally:
tracker.stop()
def test_scale_demanded_models_queues_add_shard_on_spare_host():
tracker = TrackerServer(model_presets={
"model-a": {
"total_layers": 4,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/ModelA",
},
"model-b": {
"total_layers": 4,
"bytes_per_layer": {"bfloat16": 1_000},
"hf_repo": "org/ModelB",
},
})
tracker_port = tracker.start()
try:
reg_b = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9061", "model": "model-b",
"hf_repo": "org/ModelB", "num_layers": 4,
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 2,
"vram_bytes": 20_000, "ram_bytes": 20_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:9060", "model": "model-a",
"hf_repo": "org/ModelA", "num_layers": 4,
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 1,
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
assert tracker._stats is not None
for _ in range(400):
tracker._stats.record_request("org/ModelA")
with tracker._lock:
_scale_demanded_models_locked(tracker._server) # type: ignore[arg-type]
node_b = tracker._registry[reg_b["node_id"]]
assignment = node_b.pending_new_assignment
assert assignment is not None
assert assignment["action"] == "ADD_SHARD"
assert assignment["model"] == "org/ModelA"
finally:
tracker.stop()
Reference in New Issue View Git Blame Copy Permalink