From 97e2784b372558c4bcfc38c35ff296e79726a76d Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Wed, 15 Jul 2026 10:34:41 +0200 Subject: [PATCH] node registration fixes --- packages/node/meshnet_node/startup.py | 124 ++++++++++++++++++++- packages/node/meshnet_node/torch_server.py | 11 ++ packages/tracker/meshnet_tracker/server.py | 13 ++- tests/test_node_admission.py | 67 +++++++++++ tests/test_tracker_routing.py | 7 ++ 5 files changed, 214 insertions(+), 8 deletions(-) diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 77192f0..b9fa8b3 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -12,7 +12,7 @@ import urllib.error import urllib.parse import urllib.request from pathlib import Path -from typing import Any +from typing import Any, Callable from .admission import ( AdmissionRequirement, @@ -419,6 +419,7 @@ def _start_heartbeat( interval: float = _HEARTBEAT_INTERVAL_IDLE, node_ref: Any | None = None, start_time: float | None = None, + refresh_capability: Callable[[dict], dict | None] | None = None, ) -> threading.Thread: """Daemon thread: sends heartbeats and re-registers automatically after tracker restarts. @@ -461,9 +462,26 @@ def _start_heartbeat( return _HEARTBEAT_INTERVAL_BUSY return interval + def _refresh_proof(payload: dict) -> None: + """Re-prove the current shard so a re-registration never presents an aged proof. + + The tracker refuses proofs older than its freshness budget: re-sending the + startup-time report after an outage would re-register the node unroutable. + """ + if refresh_capability is None or "capability_report" not in payload: + return + try: + fresh = refresh_capability(payload) + except Exception as exc: + print(f" [node] WARNING: capability re-validation failed: {exc}", flush=True) + return + if fresh: + payload["capability_report"] = fresh + def _reregister() -> bool: nonlocal node_id try: + _refresh_proof(register_payload) resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload) node_id = resp.get("node_id", node_id) if node_ref is not None: @@ -485,6 +503,7 @@ def _start_heartbeat( "managed_assignment": True, } try: + _refresh_proof(extra_payload) reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", extra_payload) print( f" [node] registered additional model — node ID: {reg_resp.get('node_id')}", @@ -579,6 +598,7 @@ def _register_with_tracker( reg_payload: dict, node: Any, start_time: float, + refresh_capability: Callable[[dict], dict | None] | None = None, ) -> str | None: """Register with the tracker, or start background retries when it is unreachable.""" try: @@ -586,7 +606,14 @@ def _register_with_tracker( tracker_node_id = str(reg_resp.get("node_id") or "?") setattr(node, "tracker_node_id", tracker_node_id) print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) - _start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=start_time) + _start_heartbeat( + tracker_url, + tracker_node_id, + reg_payload, + node_ref=node, + start_time=start_time, + refresh_capability=refresh_capability, + ) return tracker_node_id except Exception as exc: setattr(node, "tracker_node_id", None) @@ -598,6 +625,7 @@ def _register_with_tracker( reg_payload, node_ref=node, start_time=start_time, + refresh_capability=refresh_capability, ) return None @@ -718,6 +746,54 @@ def _admit_capability( return report +def _capability_refresher( + node: Any, + *, + manifest: RecipeManifest, + recipe: Recipe, + detected_device: str, + cache_dir: Path | None, + force_cpu: bool, + validator: CapabilityValidator | None = None, +) -> Callable[[dict], dict | None]: + """A fresh proof for what the node serves *now*, run at re-registration time. + + The startup proof ages past the tracker's freshness budget, and directives + can move the node to a shard the startup proof never covered — so every + re-registration re-proves against the currently loaded backend rather than + replaying the report captured at boot. + """ + def refresh(payload: dict) -> dict | None: + target_model = payload.get("hf_repo") or payload.get("model") + backend = None + accessor = getattr(node, "backend_for", None) + if callable(accessor) and target_model: + backend = accessor(str(target_model)) + if backend is None: + backend = getattr(node, "backend", None) + if backend is None: + return None + context = CapabilityContext( + backend=backend, + selection=DoctorSelection( + model_id=str(getattr(backend, "model_id", target_model)), + shard_start=int(getattr(backend, "shard_start", 0) or 0), + shard_end=int(getattr(backend, "shard_end", 0) or 0), + quantization=str(getattr(backend, "quantization", None) or "auto"), + cache_dir=cache_dir, + force_cpu=force_cpu, + ), + recipe=recipe, + manifest=manifest, + device=_capability_device(backend, detected_device), + ) + report = (validator or probe_capability)(context) + setattr(node, "capability_report", report) + return report.to_dict() + + return refresh + + def run_startup( tracker_url: str, port: int = 0, @@ -1026,6 +1102,15 @@ def run_startup( } tracker_node_id = _register_with_tracker( tracker_url, reg_payload, node, _node_start_time, + refresh_capability=_capability_refresher( + node, + manifest=manifest, + recipe=recipe, + detected_device=device, + cache_dir=cache_dir, + force_cpu=force_cpu, + validator=capability_validator, + ), ) print( @@ -1197,6 +1282,15 @@ def run_startup( } tracker_node_id = _register_with_tracker( tracker_url, auto_reg_payload, node, _node_start_time, + refresh_capability=_capability_refresher( + node, + manifest=manifest, + recipe=recipe, + detected_device=device, + cache_dir=cache_dir, + force_cpu=force_cpu, + validator=capability_validator, + ), ) shard_label = _format_shard_label( assigned_shard_start, @@ -1389,6 +1483,15 @@ def run_startup( } tracker_node_id = _register_with_tracker( tracker_url, reg_payload, node, _node_start_time, + refresh_capability=_capability_refresher( + node, + manifest=manifest, + recipe=recipe, + detected_device=device, + cache_dir=cache_dir, + force_cpu=force_cpu, + validator=capability_validator, + ), ) print( f"\n{'=' * 32}\n" @@ -1474,7 +1577,22 @@ def run_startup( ) node_id = str(reg_resp["node_id"]) setattr(node, "tracker_node_id", node_id) - _start_heartbeat(tracker_url, node_id, reg_payload, node_ref=node, start_time=_node_start_time) + _start_heartbeat( + tracker_url, + node_id, + reg_payload, + node_ref=node, + start_time=_node_start_time, + refresh_capability=_capability_refresher( + node, + manifest=manifest, + recipe=recipe, + detected_device=device, + cache_dir=shard_path, + force_cpu=force_cpu, + validator=capability_validator, + ), + ) except Exception: node.stop() raise diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 7bbabac..36fbeb9 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -1543,6 +1543,17 @@ class TorchNodeServer: def loaded_model_ids(self) -> list[str]: return list(self._backends.keys()) + def backend_for(self, model_id: str) -> TorchModelShard | None: + """The loaded backend serving `model_id` — full repo id or short name.""" + backend = self._backends.get(model_id) + if backend is not None: + return backend + short = model_id.split("/")[-1].lower() + for key, candidate in self._backends.items(): + if key.split("/")[-1].lower() == short: + return candidate + return None + def apply_tracker_directives(self, directives: list[dict]) -> dict | None: """Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more).""" drop_directive = next( diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index b26e975..205b2bf 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1101,12 +1101,15 @@ def _registration_quantization(body: dict, quantizations: list[str]) -> str | No An absent field predates the protocol adding it: it means "unknown", not "unsupported", so the node keeps the best precision it advertises and stays - routable. Anything the node states explicitly is taken at its word -- a null, - a non-string, or an unsupported name leaves it with no usable precision and - routing excludes it. + routable. An explicit "auto" means the same thing — the node's CLI default + delegates the choice, it does not refuse one. Anything else the node states + explicitly is taken at its word -- a null, a non-string, or an unsupported + name leaves it with no usable precision and routing excludes it. """ - if "quantization" in body: - return _normalize_quantization(body["quantization"]) + declared = body.get("quantization") + declared_auto = isinstance(declared, str) and declared.strip().lower() == "auto" + if "quantization" in body and not declared_auto: + return _normalize_quantization(declared) supported = [ normalized for value in quantizations if (normalized := _normalize_quantization(value)) is not None diff --git a/tests/test_node_admission.py b/tests/test_node_admission.py index 9b606e2..109e664 100644 --- a/tests/test_node_admission.py +++ b/tests/test_node_admission.py @@ -358,6 +358,73 @@ def test_a_stale_report_cannot_be_reused_to_register(startup_env): assert startup_env == [] +# --------------------------------------------------------------------------- +# Re-registration: the proof presented is fresh, never the one captured at boot +# --------------------------------------------------------------------------- + + +def test_run_startup_hands_the_heartbeat_a_refresher_for_the_current_shard(startup_env, monkeypatch): + "The tracker refuses aged proofs, so the heartbeat must be able to re-prove what the node serves now.\n\nTags: node, admission, startup" + import meshnet_node.startup as startup_mod + + captured: dict = {} + monkeypatch.setattr( + startup_mod, "_start_heartbeat", lambda *a, **kw: captured.update(kw) + ) + + _start(capability_validator=capability_stub()) + + refresh = captured.get("refresh_capability") + assert callable(refresh), "run_startup no longer wires a capability refresher" + fresh = refresh({"hf_repo": MODEL, "model": MODEL.split("/")[-1]}) + assert fresh is not None + assert fresh["model"]["model_id"] == MODEL + assert (fresh["shard"]["start"], fresh["shard"]["end"]) == (0, 23) + assert fresh["validated_at"] > time.time() - 60 + + +def test_a_reregistration_presents_a_refreshed_proof(monkeypatch): + "Replaying the boot-time report after an outage re-registers the node unroutable; the re-register path must present a fresh proof.\n\nTags: node, admission, startup" + import json + + import meshnet_node.startup as startup_mod + + model = "acme/refresh-model-7b" + boot_report = {"validated_at": 1.0, "marker": "boot"} + fresh_report = {"validated_at": 2.0, "marker": "fresh"} + posted: list[dict] = [] + + def _record(url, payload, timeout=10.0): + if url.endswith("/v1/nodes/register") and payload.get("hf_repo") == model: + posted.append(json.loads(json.dumps(payload))) + return {"node_id": "node-refresh"} + raise SystemExit # first heartbeat POST ends the daemon loop + + monkeypatch.setattr(startup_mod, "_post_json", _record) + + payload = { + "hf_repo": model, + "model": model.split("/")[-1], + "capability_report": dict(boot_report), + } + thread = startup_mod._start_heartbeat( + "http://tracker.invalid", + startup_mod._PENDING_NODE_ID, # forces a re-registration on the first tick + payload, + interval=0.02, + refresh_capability=lambda _payload: dict(fresh_report), + ) + + # The loop must be dead before this test returns: once monkeypatch restores + # `_post_json`, a surviving thread would re-register through whatever the + # *next* test patches in and corrupt its call counts. + thread.join(timeout=5.0) + + assert not thread.is_alive(), "the heartbeat loop outlived the test" + assert posted, "the heartbeat never re-registered" + assert posted[0]["capability_report"] == fresh_report + + def test_a_matching_passing_report_registers_and_travels_with_the_payload(startup_env): "Registration carries the proof for exactly the model/shard/recipe it advertises.\n\nTags: node, admission, startup" node = _start() # production validator against a working fake backend diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 052d386..81ad924 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -3042,6 +3042,13 @@ def test_a_node_declaring_an_unsupported_quantization_is_never_routed(): assert status == 503 +def test_a_node_declaring_auto_quantization_serves_a_default_precision_request(): + "'auto' is the CLI default that delegates the choice — it is not a refusal, so the node must resolve to its best advertised precision and route.\n\nTags: http, routing, tracker" + status, response = _proxy_chat_status(POLICY_COMPAT, quantization="auto") + assert status == 200 + assert response["choices"][0]["message"]["content"] == "ok" + + def test_a_node_declaring_a_null_quantization_is_never_routed(): "An explicit null states 'no usable precision' -- only an absent field is legacy.\n\nTags: http, routing, tracker" node = http.server.HTTPServer(("127.0.0.1", 0), _EchoChatHandler)