0be35d257b
Previous logic caught 404 on heartbeat and re-registered, but the re- registration failed silently if the tracker wasn't fully ready yet. New approach: after any connection-refused streak, the heartbeat loop switches to re-registration mode and keeps retrying until the tracker accepts the registration (instead of sending heartbeats into a fresh registry that doesn't know this node). Falls back to the 404 path for the case where a node is purged without a full tracker restart. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
418 lines
16 KiB
Python
418 lines
16 KiB
Python
"""Full node startup sequence — self-configuring, non-interactive."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import socket
|
||
import sys
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from .downloader import compute_shard_checksum, download_shard
|
||
from .hardware import detect_hardware
|
||
from .server import StubNodeServer
|
||
from .torch_server import TorchNodeServer
|
||
from .wallet import load_or_create_wallet
|
||
|
||
|
||
def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
|
||
data = json.dumps(payload).encode()
|
||
req = urllib.request.Request(
|
||
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
||
)
|
||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def _get_json(url: str, timeout: float = 10.0) -> dict:
|
||
with urllib.request.urlopen(url, timeout=timeout) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def _start_heartbeat(
|
||
tracker_url: str,
|
||
node_id: str,
|
||
register_payload: dict,
|
||
interval: float = 20.0,
|
||
) -> threading.Thread:
|
||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts."""
|
||
def _reregister() -> bool:
|
||
nonlocal node_id
|
||
try:
|
||
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
|
||
node_id = resp.get("node_id", node_id)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def _loop() -> None:
|
||
nonlocal node_id
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
outage_streak = 0 # consecutive intervals where tracker was unreachable
|
||
|
||
while True:
|
||
time.sleep(interval)
|
||
|
||
if outage_streak > 0:
|
||
# Tracker was down — attempt re-registration first (it may have restarted
|
||
# with a clean slate and won't know this node).
|
||
if _reregister():
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
print(f" [node] re-registered after outage — node ID: {node_id}", flush=True)
|
||
outage_streak = 0
|
||
else:
|
||
outage_streak += 1
|
||
if outage_streak <= 3 or outage_streak % 10 == 0:
|
||
print(
|
||
f" [node] WARNING: tracker still unreachable "
|
||
f"({outage_streak * interval:.0f}s)",
|
||
flush=True,
|
||
)
|
||
continue
|
||
|
||
try:
|
||
_post_json(hb_url, {})
|
||
except urllib.error.HTTPError as exc:
|
||
if exc.code == 404:
|
||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||
print(" [node] tracker lost registration — re-registering...", flush=True)
|
||
if _reregister():
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
print(f" [node] re-registered — node ID: {node_id}", flush=True)
|
||
else:
|
||
print(" [node] WARNING: re-registration failed", flush=True)
|
||
outage_streak = 1
|
||
else:
|
||
print(f" [node] WARNING: heartbeat failed ({exc.code}): {exc}", flush=True)
|
||
except Exception as exc:
|
||
outage_streak = 1
|
||
print(f" [node] WARNING: tracker unreachable: {exc}", flush=True)
|
||
|
||
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
|
||
t.start()
|
||
return t
|
||
|
||
|
||
def run_startup(
|
||
tracker_url: str,
|
||
port: int = 0,
|
||
model: str = "stub-model",
|
||
model_id: str | None = None,
|
||
shard_start: int | None = None,
|
||
shard_end: int | None = None,
|
||
quantization: str = "bfloat16",
|
||
wallet_path: Path | None = None,
|
||
cache_dir: Path | None = None,
|
||
host: str = "127.0.0.1",
|
||
advertise_host: str | None = None,
|
||
contracts: Any | None = None,
|
||
) -> StubNodeServer | TorchNodeServer:
|
||
"""Execute the full startup sequence and return a running node server.
|
||
|
||
Steps (all non-interactive):
|
||
1. Detect GPU / hardware profile
|
||
2. Load or generate Solana wallet keypair
|
||
3. Query tracker for optimal shard assignment
|
||
4. Download (or stub) the assigned shard from peers, then HuggingFace
|
||
5. Start local HTTP server
|
||
6. Register with tracker
|
||
|
||
Prints a compact status summary on completion.
|
||
"""
|
||
|
||
tracker_url = tracker_url.rstrip("/")
|
||
|
||
# 1. Hardware detection
|
||
print("Detecting hardware...", flush=True)
|
||
hw = detect_hardware()
|
||
device: str = hw["device"]
|
||
gpu_name: str | None = hw.get("gpu_name")
|
||
vram_mb: int = hw.get("vram_mb", 0)
|
||
|
||
if device == "cpu":
|
||
print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
|
||
else:
|
||
print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
|
||
|
||
# 2. Wallet
|
||
print("Loading wallet...", flush=True)
|
||
wallet_kwargs: dict = {}
|
||
if wallet_path is not None:
|
||
wallet_kwargs["path"] = wallet_path
|
||
_, _, address = load_or_create_wallet(**wallet_kwargs)
|
||
print(f" Wallet: {address}", flush=True)
|
||
probationary_line = _probationary_status_line(contracts, address)
|
||
if probationary_line is not None:
|
||
print(f" {probationary_line}", flush=True)
|
||
|
||
if model_id: # treat "" the same as None — no explicit model given
|
||
# Auto-detect shard range from model config if not explicitly provided
|
||
if shard_start is None or shard_end is None:
|
||
detected = _detect_num_layers(model_id)
|
||
if detected is None:
|
||
raise ValueError(
|
||
f"Could not read num_hidden_layers from {model_id} config. "
|
||
"Pass --shard-start and --shard-end explicitly."
|
||
)
|
||
# When no explicit shard range given, ask the tracker if there's a gap for this model.
|
||
if shard_start is None and shard_end is None:
|
||
try:
|
||
qs = urllib.parse.urlencode({
|
||
"device": device, "vram_mb": vram_mb, "hf_repo": model_id,
|
||
})
|
||
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
|
||
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
|
||
shard_start = net_asgn["shard_start"]
|
||
shard_end = net_asgn["shard_end"]
|
||
print(
|
||
f" Tracker found uncovered shard: "
|
||
f"layers {shard_start}–{shard_end} (of {detected})",
|
||
flush=True,
|
||
)
|
||
except Exception:
|
||
pass # No other nodes registered yet — default to full model below
|
||
shard_start = shard_start if shard_start is not None else 0
|
||
shard_end = shard_end if shard_end is not None else detected - 1
|
||
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
|
||
|
||
print("Loading real PyTorch model shard...", flush=True)
|
||
node = TorchNodeServer(
|
||
host=host,
|
||
port=port,
|
||
model_id=model_id,
|
||
shard_start=shard_start,
|
||
shard_end=shard_end,
|
||
quantization=quantization,
|
||
tracker_url=tracker_url,
|
||
)
|
||
actual_port = node.start()
|
||
total_layers = getattr(node.backend, "total_layers", None)
|
||
if isinstance(total_layers, int) and total_layers > 0:
|
||
layer_count = shard_end - shard_start + 1
|
||
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
|
||
else:
|
||
shard_label = f"layers {shard_start}–{shard_end}"
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
# Register with tracker so other nodes can auto-join this model.
|
||
total_layers = getattr(node.backend, "total_layers", None)
|
||
reg_payload = {
|
||
"endpoint": endpoint,
|
||
"model": model_id.split("/")[-1],
|
||
"hf_repo": model_id,
|
||
"num_layers": total_layers,
|
||
"shard_start": shard_start,
|
||
"shard_end": shard_end,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"quantization": quantization,
|
||
"score": 1.0,
|
||
"tracker_mode": (shard_start == 0),
|
||
}
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||
node_id = reg_resp.get("node_id", "?")
|
||
print(f" Registered with tracker — node ID: {node_id}", flush=True)
|
||
_start_heartbeat(tracker_url, node_id, reg_payload)
|
||
except Exception as exc:
|
||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready\n"
|
||
f" Wallet: {address}\n"
|
||
f" Model ID: {model_id}\n"
|
||
f" Shard: {shard_label}\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Hardware: {device.upper()}\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
return node
|
||
if shard_start is not None or shard_end is not None:
|
||
raise ValueError("--shard-start / --shard-end require --model-id")
|
||
|
||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||
print("Querying tracker for network assignment...", flush=True)
|
||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb})
|
||
net_assignment: dict = {}
|
||
try:
|
||
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
|
||
except Exception as exc:
|
||
print(f" (auto-join unavailable: {exc})", flush=True)
|
||
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
||
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
||
|
||
if assigned_hf_repo and _gap_found:
|
||
assigned_shard_start: int = net_assignment["shard_start"]
|
||
assigned_shard_end: int = net_assignment["shard_end"]
|
||
assigned_num_layers: int = net_assignment["num_layers"]
|
||
print(
|
||
f" Assigned: {assigned_hf_repo} "
|
||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"(of {assigned_num_layers})",
|
||
flush=True,
|
||
)
|
||
print("Loading real PyTorch model shard...", flush=True)
|
||
node = TorchNodeServer(
|
||
host=host,
|
||
port=port,
|
||
model_id=assigned_hf_repo,
|
||
shard_start=assigned_shard_start,
|
||
shard_end=assigned_shard_end,
|
||
quantization=quantization,
|
||
tracker_url=tracker_url,
|
||
)
|
||
actual_port = node.start()
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
auto_reg_payload = {
|
||
"endpoint": endpoint,
|
||
"model": assigned_hf_repo.split("/")[-1],
|
||
"hf_repo": assigned_hf_repo,
|
||
"num_layers": assigned_num_layers,
|
||
"shard_start": assigned_shard_start,
|
||
"shard_end": assigned_shard_end,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"quantization": quantization,
|
||
"score": 1.0,
|
||
"tracker_mode": (assigned_shard_start == 0),
|
||
}
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
|
||
node_id = reg_resp.get("node_id", "?")
|
||
print(f" Registered with tracker — node ID: {node_id}", flush=True)
|
||
_start_heartbeat(tracker_url, node_id, auto_reg_payload)
|
||
except Exception as exc:
|
||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready (auto-joined)\n"
|
||
f" Wallet: {address}\n"
|
||
f" Model ID: {assigned_hf_repo}\n"
|
||
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"({shard_count} of {assigned_num_layers})\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Hardware: {device.upper()}\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
return node
|
||
|
||
# 3b. Shard assignment from tracker (stub-model / preset-based path)
|
||
print("Querying tracker for shard assignment...", flush=True)
|
||
assign_qs = urllib.parse.urlencode({
|
||
"model": model,
|
||
"device": device,
|
||
"vram_mb": vram_mb,
|
||
})
|
||
try:
|
||
assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
|
||
except urllib.error.URLError as exc:
|
||
print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
|
||
raise
|
||
|
||
shard_start: int = assignment["shard_start"]
|
||
shard_end: int = assignment["shard_end"]
|
||
assigned_model: str = assignment.get("model", model)
|
||
hf_repo: str | None = assignment.get("hf_repo")
|
||
peers: list[dict] = assignment.get("peers", [])
|
||
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
||
|
||
# 4. Download shard
|
||
print("Downloading shard...", flush=True)
|
||
dl_kwargs: dict = {}
|
||
if cache_dir is not None:
|
||
dl_kwargs["cache_dir"] = cache_dir
|
||
if hf_repo is not None:
|
||
dl_kwargs["hf_repo"] = hf_repo
|
||
if peers:
|
||
dl_kwargs["peers"] = peers
|
||
shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
|
||
shard_checksum = compute_shard_checksum(shard_path)
|
||
print(f" Cached at: {shard_path}", flush=True)
|
||
|
||
# 5. Start HTTP server
|
||
is_last = shard_end >= assignment.get("model_layers_end", shard_end)
|
||
node = StubNodeServer(
|
||
host=host,
|
||
port=port,
|
||
shard_start=shard_start,
|
||
shard_end=shard_end,
|
||
is_last_shard=is_last,
|
||
model=assigned_model,
|
||
shard_path=shard_path,
|
||
)
|
||
actual_port = node.start()
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
|
||
# 6. Register with tracker
|
||
print("Registering with tracker...", flush=True)
|
||
try:
|
||
reg_resp = _post_json(
|
||
f"{tracker_url}/v1/nodes/register",
|
||
{
|
||
"endpoint": endpoint,
|
||
"model": assigned_model,
|
||
"shard_start": shard_start,
|
||
"shard_end": shard_end,
|
||
"shard_checksum": shard_checksum,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"score": 1.0,
|
||
},
|
||
)
|
||
node_id: str = reg_resp["node_id"]
|
||
except Exception:
|
||
node.stop()
|
||
raise
|
||
|
||
# Status summary
|
||
hw_str = device.upper()
|
||
if gpu_name:
|
||
hw_str += f" ({gpu_name}, {vram_mb} MB)"
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready\n"
|
||
f" Wallet: {address}\n"
|
||
f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {node_id}\n"
|
||
f" Hardware: {hw_str}\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
|
||
return node
|
||
|
||
|
||
def _detect_num_layers(model_id: str) -> int | None:
|
||
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
|
||
try:
|
||
from transformers import AutoConfig # type: ignore[import]
|
||
cfg = AutoConfig.from_pretrained(model_id)
|
||
return int(cfg.num_hidden_layers)
|
||
except Exception as exc:
|
||
print(f" Warning: could not read model config from HF: {exc}", flush=True)
|
||
return None
|
||
|
||
|
||
def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
|
||
if contracts is None:
|
||
return None
|
||
remaining = contracts.registry.probationary_jobs_remaining(wallet_address)
|
||
if remaining <= 0:
|
||
return "Probationary period complete: earning enabled"
|
||
suffix = "job" if remaining == 1 else "jobs"
|
||
return f"Probationary period: {remaining} {suffix} remaining before earning"
|