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

tracker rebalancing tweaks

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 18:24:09 +03:00
parent fdeb881c83
commit f220fd2210
4 changed files with 172 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
tests/test_tracker_routing.py
View File

@@ -1114,12 +1114,14 @@ def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap():
"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")
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 route_resp["nodes"]
for node in net["nodes"]
}
assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"]
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()
@@ -1147,7 +1149,8 @@ def test_tracker_registration_directive_is_not_replayed_on_heartbeat():
tracker.stop()
def test_tracker_reassignment_emits_drop_before_load():
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,
@@ -1156,21 +1159,69 @@ def test_tracker_reassignment_emits_drop_before_load():
})
tracker_port = tracker.start()
try:
slow = _post_json(
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},
)
_post_json(
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},
)
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"]
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()
@@ -1208,22 +1259,20 @@ def test_tracker_faster_node_receives_wider_range_when_capacity_ties():
_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"],
"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": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
"vram_bytes": 5_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"]
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()