popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: coverage-first tracker shard assignment (US-013)

Browse Source 
Recovered from interrupted Codex session (ended 13:02, changes uncommitted).
All 74 tests pass.

- Tracker accepts vram_bytes, ram_bytes, quantizations[], benchmark_tokens_per_sec
  in node registration payload
- GET /v1/coverage/<model_preset> returns [{start_layer, end_layer, node_count}]
- Coverage-first bin-packing: fills gaps before adding redundancy
- Speed-weighted assignment: faster nodes get wider shard ranges
- LOAD_SHARD/DROP_SHARD rebalance directives delivered via heartbeat responses
- Model is unroutable when any layer range has node_count=0
- model_presets.json config for bytes_per_layer at each quantization level

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-29 16:37:48 +03:00
parent 85c13e4e82
commit 753f553766
2 changed files with 639 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
tests/test_tracker_routing.py
View File

@@ -10,7 +10,7 @@ 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, _registration_ban_error
from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error
def _post_json(url: str, payload: dict) -> dict:
@@ -100,6 +100,294 @@ def test_tracker_route_error_no_coverage():
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")
assert coverage_resp["coverage"] == [
{"start_layer": 0, "end_layer": 2, "node_count": 1},
{"start_layer": 3, "end_layer": 5, "node_count": 0},
]
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_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_route_error_no_nodes():
"""Tracker returns 503 with clear error when the registry is empty."""
tracker = TrackerServer()