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

node registration fixes

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-15 10:34:41 +02:00
parent ba7c656364
commit 97e2784b37
5 changed files with 214 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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

11
packages/node/meshnet_node/torch_server.py
View File

@@ -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(

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

@@ -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

67
tests/test_node_admission.py
View File

@@ -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

7
tests/test_tracker_routing.py
View File

@@ -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)