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