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:

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

@@ -3276,7 +3276,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
node.hf_repo or node.model
for node in alive
if node.model is not None
and node.model not in server.model_presets
# The same model can be registered under its HF repository while
# the catalogue exposes its short preset id. Do not emit a second
# repo-keyed entry when either node identifier resolves to a preset.
and _resolve_model_preset(
server.model_presets, node.hf_repo or node.model,
)[1] is None
and node.shard_start is not None
and node.shard_end is not None
and node.num_layers is not None
@@ -4820,6 +4825,20 @@ 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"]
completed_directives = body.get("completed_directives", [])
if isinstance(completed_directives, list):
for directive in completed_directives:
if not isinstance(directive, dict) or directive.get("action") != "DROP_SHARD":
continue
# A node has confirmed the release. Stop advertising its
# old route immediately so the dashboard and routing state
# agree with the runtime.
entry.model = "stub-model"
entry.hf_repo = None
entry.shard_start = None
entry.shard_end = None
entry.tracker_mode = False
entry.status = "ready"
if "friendly_name" in body:
try:
entry.friendly_name = _normalize_friendly_name(body.get("friendly_name"))

32
tests/test_dynamic_routing.py
View File

@@ -375,6 +375,11 @@ def test_admin_can_replace_a_served_model_and_release_it():
with urllib.request.urlopen(load) as response:
assert json.loads(response.read())["assignment"]["node_id"] == node["node_id"]
heartbeat = _post_json(f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat", {})
_post_json(
f"http://127.0.0.1:{port}/v1/nodes/{node['node_id']}/heartbeat",
{"completed_directives": [{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct"}]},
)
network = _get_json(f"http://127.0.0.1:{port}/v1/network/map")
assert heartbeat["directives"][0]["action"] == "LOAD_SHARD"
release = urllib.request.Request(
f"http://127.0.0.1:{port}/v1/models/release",
@@ -386,6 +391,33 @@ def test_admin_can_replace_a_served_model_and_release_it():
tracker.stop()
assert heartbeat["directives"][0]["action"] == "DROP_SHARD"
released_node = next(item for item in network["nodes"] if item["node_id"] == node["node_id"])
assert released_node["shard_start"] is None
assert released_node["shard_end"] is None
def test_models_list_does_not_duplicate_a_preset_registered_by_hf_repo():
"""A preset and its canonical repository are one selectable model."""
tracker = TrackerServer(enable_billing=False)
port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{port}/v1/nodes/register",
{
"endpoint": "http://127.0.0.1:9913",
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"shard_start": 0,
"shard_end": 23,
},
)
models = _get_json(f"http://127.0.0.1:{port}/v1/models")["data"]
finally:
tracker.stop()
assert [model["id"] for model in models].count("qwen2.5-0.5b-instruct") == 1
assert not any(model["id"] == "Qwen/Qwen2.5-0.5B-Instruct" for model in models)
def test_endpoint_key_distinguishes_same_port_different_hosts():

41
tests/test_node_startup.py
View File

@@ -287,6 +287,47 @@ def test_configure_torch_threads_applies_explicit_settings(monkeypatch):
assert active == {"torch_threads": 12, "torch_interop_threads": 2}
def test_heartbeat_applies_release_without_reregistering(monkeypatch):
"""DROP_SHARD has no replacement range and must not look like an outage."""
import meshnet_node.startup as startup_mod
released = threading.Event()
requests: list[tuple[str, dict]] = []
class FakeNode:
def apply_tracker_directives(self, directives):
assert directives == [{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct"}]
return {"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct"}
def fake_post(url, payload, timeout=10.0):
requests.append((url, dict(payload)))
released.set()
return {"directives": [{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct"}]}
sleep_calls = 0
def one_heartbeat(_seconds):
nonlocal sleep_calls
sleep_calls += 1
if sleep_calls > 1:
raise SystemExit
monkeypatch.setattr(startup_mod, "_post_json", fake_post)
monkeypatch.setattr(startup_mod.time, "sleep", one_heartbeat)
payload = {
"model": "Qwen2.5-0.5B-Instruct",
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"shard_start": 0,
"shard_end": 23,
}
startup_mod._start_heartbeat("http://tracker", "node-1", payload, interval=0, node_ref=FakeNode())
assert released.wait(1), "heartbeat did not receive the queued release"
assert len(requests) == 1, "release must not trigger a re-registration"
assert payload["shard_start"] == 0
assert payload["shard_end"] == 23
def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
"benchmark_tokens_per_sec from the benchmark is included in the tracker registration.\n\nTags: node, performance, startup"
import meshnet_node.startup as startup_mod