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

Merge commit '47b243cd98fd94da7918cacf5725373b099208e5' into ralph/distributed-gguf-runtime

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-15 23:04:52 +02:00
parent 1fe31ef38d 47b243cd98
commit 254627629b
12 changed files with 620 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
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,
@@ -420,6 +420,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.
@@ -431,6 +432,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 +457,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:
@@ -462,9 +466,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:
@@ -486,6 +507,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')}",
@@ -494,21 +516,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") in {"DROP_SHARD", "DROP_ALL_SHARDS"}:
# 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
@@ -516,6 +543,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
@@ -543,7 +571,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:
@@ -580,6 +611,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:
@@ -587,7 +619,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)
@@ -599,6 +638,7 @@ def _register_with_tracker(
reg_payload,
node_ref=node,
start_time=start_time,
refresh_capability=refresh_capability,
)
return None
@@ -748,6 +788,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,
@@ -1079,6 +1167,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(
@@ -1268,6 +1365,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(
proof_shard.start,
@@ -1477,6 +1583,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"
@@ -1564,7 +1679,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