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

BIN
billing.sqlite
View File

Binary file not shown.

BIN
meshnet_registry.sqlite3
View File

Binary file not shown.

164
packages/tracker/meshnet_tracker/server.py
View File

@@ -1099,6 +1099,76 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
]
def _unassigned_managed_nodes(nodes: list["_NodeEntry"]) -> list["_NodeEntry"]:
return [
node for node in nodes
if node.managed_assignment
and (node.shard_start is None or node.shard_end is None)
]
def _emit_shard_change_directives(
node: "_NodeEntry",
model: str,
previous_range: tuple[int | None, int | None, str | None],
preset: dict,
) -> None:
"""Queue DROP/LOAD directives when a managed node's shard assignment changes."""
previous_start, previous_end, previous_quantization = previous_range
current_range = (node.shard_start, node.shard_end, node.quantization)
if node.shard_start is None or node.shard_end is None or current_range == previous_range:
return
if previous_start is not None and previous_end is not None:
node.pending_directives.append(
_drop_directive(
node,
model,
previous_start,
previous_end,
previous_quantization or _node_quantization(node, preset),
)
)
node.pending_directives.append(
_load_directive(
node,
model,
node.shard_start,
node.shard_end,
node.quantization or _node_quantization(node, preset),
)
)
def _assign_redundant_managed_nodes(
managed_nodes: list["_NodeEntry"],
model: str,
preset: dict,
required_start: int,
required_end: int,
) -> None:
"""Give newly joined managed nodes their own copy without reshuffling incumbents."""
total_layers = required_end - required_start + 1
for node in sorted(
managed_nodes,
key=lambda entry: (
-entry.benchmark_tokens_per_sec,
-_node_layer_capacity(entry, preset),
entry.node_id,
),
):
if node.shard_start is not None and node.shard_end is not None:
continue
capacity = min(_node_layer_capacity(node, preset), total_layers)
if capacity <= 0:
continue
previous_range = (node.shard_start, node.shard_end, node.quantization)
quantization = _node_quantization(node, preset)
node.quantization = quantization
node.shard_start = required_start
node.shard_end = min(required_end, required_start + capacity - 1)
_emit_shard_change_directives(node, model, previous_range, preset)
def _relay_http_request_frames(
relay_addr: str,
path: str,
@@ -1439,6 +1509,17 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
if not managed_nodes:
return
coverage = _coverage_map(model_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
unassigned = _unassigned_managed_nodes(managed_nodes)
if not gaps and not unassigned:
return
if not gaps and unassigned:
_assign_redundant_managed_nodes(
unassigned, model, preset, required_start, required_end,
)
return
previous_ranges = {
node.node_id: (node.shard_start, node.shard_end, node.quantization)
for node in managed_nodes
@@ -1455,10 +1536,8 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
)
)
base_nodes = [node for node in model_nodes if not node.managed_assignment]
coverage = _coverage_map(base_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
if not gaps:
gaps = [(required_start, required_end)]
base_coverage = _coverage_map(base_nodes, required_start, required_end)
gaps = _coverage_gaps(base_coverage)
eligible_nodes = [
node for node in managed_nodes
@@ -1487,28 +1566,11 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
node_index += 1
for node in managed_nodes:
previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
current_range = (node.shard_start, node.shard_end, node.quantization)
if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
continue
if previous_start is not None and previous_end is not None:
node.pending_directives.append(
_drop_directive(
node,
model,
previous_start,
previous_end,
previous_quantization or _node_quantization(node, preset),
)
)
node.pending_directives.append(
_load_directive(
node,
model,
node.shard_start,
node.shard_end,
node.quantization or _node_quantization(node, preset),
)
_emit_shard_change_directives(
node,
model,
previous_ranges[node.node_id],
preset,
)
@@ -1537,6 +1599,17 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
if total_layers <= 0:
return
coverage = _coverage_map(model_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
unassigned = _unassigned_managed_nodes(managed_nodes)
if not gaps and not unassigned:
return
if not gaps and unassigned:
_assign_redundant_managed_nodes(
unassigned, hf_repo, preset, required_start, required_end,
)
return
previous_ranges = {
node.node_id: (node.shard_start, node.shard_end, node.quantization)
for node in managed_nodes
@@ -1553,10 +1626,8 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
)
)
base_nodes = [node for node in model_nodes if not node.managed_assignment]
coverage = _coverage_map(base_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
if not gaps:
gaps = [(required_start, required_end)]
base_coverage = _coverage_map(base_nodes, required_start, required_end)
gaps = _coverage_gaps(base_coverage)
eligible_nodes = [
node for node in managed_nodes
@@ -1585,28 +1656,11 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
node_index += 1
for node in managed_nodes:
previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
current_range = (node.shard_start, node.shard_end, node.quantization)
if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
continue
if previous_start is not None and previous_end is not None:
node.pending_directives.append(
_drop_directive(
node,
hf_repo,
previous_start,
previous_end,
previous_quantization or _node_quantization(node, preset),
)
)
node.pending_directives.append(
_load_directive(
node,
hf_repo,
node.shard_start,
node.shard_end,
node.quantization or _node_quantization(node, preset),
)
_emit_shard_change_directives(
node,
hf_repo,
previous_ranges[node.node_id],
preset,
)
@@ -3257,7 +3311,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
old = server.registry.pop(eid)
stale_entries.append((eid, old))
server.registry[node_id] = entry
if entry.managed_assignment:
if entry.managed_assignment and not explicit_shard:
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
@@ -3358,10 +3412,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
entry.uptime_seconds = float(body["uptime_seconds"])
if "status" in body and body["status"] in ("ready", "loading"):
entry.status = body["status"]
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
_rebalance_model_locked(server, entry.model or "stub-model")
directives = list(entry.pending_directives)
entry.pending_directives.clear()
new_assignment = entry.pending_new_assignment

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()