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

Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-30 19:28:43 +02:00
parent f1e4870124 7866723c82
commit d0307fcc84
8 changed files with 590 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -706,9 +706,109 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
)
def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict:
total_layers = max(node.num_layers or 0 for node in nodes)
return {
"layers_start": 0,
"layers_end": total_layers - 1,
"bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
}
def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None:
model_nodes = [
node for node in server.registry.values()
if node.hf_repo == hf_repo
and node.num_layers is not None
]
managed_nodes = [node for node in model_nodes if node.managed_assignment]
if not model_nodes or not managed_nodes:
return
preset = _hf_rebalance_preset(model_nodes)
required_start, required_end = _preset_layer_bounds(preset)
total_layers = required_end - required_start + 1
if total_layers <= 0:
return
previous_ranges = {
node.node_id: (node.shard_start, node.shard_end, node.quantization)
for node in managed_nodes
}
for node in managed_nodes:
node.shard_start = None
node.shard_end = None
managed_nodes.sort(
key=lambda node: (
-node.benchmark_tokens_per_sec,
-_node_layer_capacity(node, preset),
node.node_id,
)
)
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)]
eligible_nodes = [
node for node in managed_nodes
if _node_layer_capacity(node, preset) > 0
]
node_index = 0
for gap_start, gap_end in gaps:
cursor = gap_start
while cursor <= gap_end and node_index < len(eligible_nodes):
node = eligible_nodes[node_index]
remaining_layers = gap_end - cursor + 1
remaining_nodes_after = len(eligible_nodes) - node_index - 1
capacity = min(
_node_layer_capacity(node, preset),
total_layers,
max(1, remaining_layers - remaining_nodes_after),
)
if capacity <= 0:
node_index += 1
continue
quantization = _node_quantization(node, preset)
node.quantization = quantization
node.shard_start = cursor
node.shard_end = min(gap_end, cursor + capacity - 1)
cursor = node.shard_end + 1
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),
)
)
def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None:
for model in list(server.model_presets):
_rebalance_model_locked(server, model)
for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}):
_rebalance_hf_model_locked(server, hf_repo)
def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
@@ -1095,7 +1195,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
covered_up_to = rs - 1
route_hops: list[dict] = []
for rn in route_nodes:
route_hops.append({"endpoint": rn.endpoint, "start_layer": covered_up_to + 1})
hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1}
if rn.relay_addr:
hop["relay_addr"] = rn.relay_addr
route_hops.append(hop)
covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to
# Strip the first-shard node we're about to proxy to — it's already handling the request.
downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")]
@@ -1332,6 +1435,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
return
tracker_mode = bool(body.get("tracker_mode", False))
managed_assignment = bool(body.get("managed_assignment", False))
hf_repo = body.get("hf_repo")
if hf_repo is not None and not isinstance(hf_repo, str):
self._send_json(400, {"error": "hf_repo must be a string"})
@@ -1371,7 +1475,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
quantizations=quantizations,
benchmark_tokens_per_sec=benchmark_tokens_per_sec,
quantization=quantization,
managed_assignment=not explicit_shard,
managed_assignment=managed_assignment or not explicit_shard,
tracker_mode=tracker_mode,
hf_repo=hf_repo,
num_layers=num_layers,
@@ -1395,7 +1499,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
server.registry[node_id] = entry
if entry.managed_assignment:
_rebalance_model_locked(server, model)
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
_rebalance_model_locked(server, model)
assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None
if assignment_directive is not None:
entry.pending_directives.clear()
@@ -1459,7 +1566,10 @@ 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"]
_rebalance_model_locked(server, entry.model or "stub-model")
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