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

fix model load/unload

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-15 12:35:32 +02:00
parent 97e2784b37
commit 22f28bd69a
4 changed files with 111 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -431,6 +431,7 @@ def _start_heartbeat(
which is logged for now (hot-reload implemented in US-026).
"""
_start_time = start_time or time.monotonic()
completed_directives: list[dict] = []
def _current_requests_snapshot() -> list[dict]:
if node_ref is None:
@@ -455,6 +456,8 @@ def _start_heartbeat(
current_requests = _current_requests_snapshot()
if current_requests:
stats["current_requests"] = current_requests
if completed_directives:
stats["completed_directives"] = list(completed_directives)
return stats
def _sleep_interval() -> float:
@@ -512,21 +515,26 @@ def _start_heartbeat(
except Exception as exc:
print(f" [node] WARNING: additional model registration failed: {exc}", flush=True)
def _apply_directives(directives: list[dict]) -> None:
def _apply_directives(directives: list[dict]) -> dict | None:
if not directives:
return
return None
if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
print(f" [node] tracker directives received: {directives}", flush=True)
return
return None
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
return None
if applied:
completed_directives.append(dict(applied))
if applied.get("action") == "ADD_SHARD":
_register_additional_assignment(applied)
return
return applied
if applied.get("action") == "DROP_SHARD":
# A release has no replacement range. It is not a failed
# heartbeat and must not re-register the released assignment.
return 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
@@ -534,6 +542,7 @@ def _start_heartbeat(
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))
return applied
def _loop() -> None:
nonlocal node_id
@@ -561,7 +570,10 @@ def _start_heartbeat(
continue
try:
resp = _post_json(hb_url, _get_stats())
heartbeat = _get_stats()
resp = _post_json(hb_url, heartbeat)
if heartbeat.get("completed_directives"):
completed_directives.clear()
_apply_directives(resp.get("directives", []))
new_asgn = resp.get("new_assignment")
if new_asgn: