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

tracker rebalance events

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-30 20:08:18 +03:00
parent 4f00a37d72
commit 7866723c82
4 changed files with 248 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
packages/node/meshnet_node/startup.py
View File

@@ -123,6 +123,26 @@ def _start_heartbeat(
except Exception:
return False
def _apply_directives(directives: list[dict]) -> None:
if not directives:
return
if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
print(f" [node] tracker directives received: {directives}", flush=True)
return
try:
applied = node_ref.apply_tracker_directives(directives)
except Exception as exc:
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
return
if applied:
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
register_payload["model"] = str(model_id).split("/")[-1]
register_payload["hf_repo"] = model_id
register_payload["shard_start"] = applied["shard_start"]
register_payload["shard_end"] = applied["shard_end"]
register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization"))
register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False))
def _loop() -> None:
nonlocal node_id
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
@@ -150,6 +170,7 @@ def _start_heartbeat(
try:
resp = _post_json(hb_url, _get_stats())
_apply_directives(resp.get("directives", []))
new_asgn = resp.get("new_assignment")
if new_asgn:
print(
@@ -282,6 +303,7 @@ def run_startup(
print(f" {probationary_line}", flush=True)
if model_id: # treat "" the same as None — no explicit model given
user_pinned_shard = shard_start is not None or shard_end is not None
# Auto-detect shard range from model config if not explicitly provided
if shard_start is None or shard_end is None:
detected = _detect_num_layers(model_id)
@@ -355,6 +377,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
**relay_fields,
}
tracker_node_id: str | None = None
@@ -442,6 +465,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
"managed_assignment": True,
**relay_fields,
}
tracker_node_id = None

34
packages/node/meshnet_node/torch_server.py
View File

@@ -729,6 +729,40 @@ class TorchNodeServer:
def queue_depth(self) -> int:
return self._server.queue_depth if self._server is not None else 0
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
"""Apply tracker LOAD_SHARD directives by hot-swapping the loaded backend."""
load_directive = next(
(directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
None,
)
if load_directive is None:
return None
shard_start = int(load_directive["shard_start"])
shard_end = int(load_directive["shard_end"])
quantization = str(load_directive.get("quantization") or self._backend.quantization)
model_id = str(load_directive.get("model") or self._backend.model_id)
print(
f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
flush=True,
)
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
self._backend = new_backend
self._tracker_mode = shard_start == 0
if self._server is not None:
self._server.backend = new_backend
self._server.tracker_mode = self._tracker_mode
print(
f" [node] loaded reassigned shard: {model_id} layers {shard_start}-{shard_end}",
flush=True,
)
return {
"model": model_id,
"shard_start": shard_start,
"shard_end": shard_end,
"quantization": quantization,
"tracker_mode": self._tracker_mode,
}
def start(self) -> int:
if self._server is not None:
raise RuntimeError("TorchNodeServer is already running")

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

@@ -685,9 +685,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:
@@ -1314,6 +1414,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"})
@@ -1353,7 +1454,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,
@@ -1377,7 +1478,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()
@@ -1441,7 +1545,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