diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index b9fa8b3..c0f6c73 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -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: diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 205b2bf..4e7875d 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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")) diff --git a/tests/test_dynamic_routing.py b/tests/test_dynamic_routing.py index 85c7ad6..cc82284 100644 --- a/tests/test_dynamic_routing.py +++ b/tests/test_dynamic_routing.py @@ -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(): diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 1c302a7..264c530 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -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