811 lines
33 KiB
Python
811 lines
33 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, benchmark_throughput_checked
|
||
from .model_catalog import model_metadata_for
|
||
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
|
||
from .server import StubNodeServer
|
||
from .torch_server import TorchNodeServer
|
||
from .wallet import load_or_create_wallet
|
||
|
||
|
||
_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
|
||
|
||
|
||
def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
|
||
"""Return the capacity budget in MB and whether it came from VRAM or RAM."""
|
||
if device == "cuda" and vram_mb > 0:
|
||
if shared_vram_mb > 0:
|
||
return vram_mb + shared_vram_mb, "VRAM + shared RAM"
|
||
return vram_mb, "VRAM"
|
||
return max(0, ram_mb), "RAM"
|
||
|
||
|
||
def _full_model_sources(model_sources: list[dict]) -> list[dict]:
|
||
"""Use tracker full-snapshot URLs for real HF model loading."""
|
||
full_sources: list[dict] = []
|
||
for source in model_sources:
|
||
full_url = source.get("full_url")
|
||
if isinstance(full_url, str) and full_url:
|
||
full_sources.append({
|
||
**source,
|
||
"url": full_url,
|
||
"files": [],
|
||
"type": f"{source.get('type') or 'model-source'}-full",
|
||
})
|
||
return full_sources
|
||
|
||
|
||
def _hardware_label(device: str, gpu_name: str | None = None) -> str:
|
||
if device == "cuda":
|
||
return "CUDA"
|
||
if gpu_name:
|
||
return "CPU (CUDA inactive)"
|
||
return "CPU"
|
||
|
||
|
||
def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int:
|
||
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
|
||
return 0
|
||
budget_bytes = memory_mb * 1024 * 1024
|
||
return min(total_layers, int((budget_bytes * 0.8) // _DEFAULT_BYTES_PER_LAYER))
|
||
|
||
|
||
def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str:
|
||
memory_gb = memory_mb / 1024
|
||
gb_str = f"{memory_gb:.1f} GB"
|
||
budget_quantization = "bfloat16" if quantization == "auto" else quantization
|
||
if total_layers is None or total_layers <= 0:
|
||
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
|
||
max_layers = _max_assignable_layers(memory_mb, total_layers)
|
||
# Remaining capacity after one full model load (rough estimate)
|
||
shard_bytes = max_layers * _DEFAULT_BYTES_PER_LAYER
|
||
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
|
||
remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
|
||
return (
|
||
f"Memory budget: {gb_str} {memory_source}; "
|
||
f"Shard budget: up to {max_layers}/{total_layers} layers at {budget_quantization}"
|
||
f"{remaining_str}"
|
||
)
|
||
|
||
|
||
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 _infer_relay_url_from_tracker(tracker_url: str) -> str | None:
|
||
"""Infer relay WebSocket URL from a public HTTPS tracker origin.
|
||
|
||
Public deployments colocate relay at /ws on the same host as the tracker API
|
||
(see QUICKSTART nginx layout). Local LAN trackers use a separate relay port
|
||
and must advertise relay_url explicitly via /v1/network/map.
|
||
"""
|
||
parsed = urllib.parse.urlparse(tracker_url)
|
||
if parsed.scheme != "https":
|
||
return None
|
||
host = parsed.hostname
|
||
if not host or host in ("127.0.0.1", "localhost"):
|
||
return None
|
||
return f"wss://{parsed.netloc}/ws"
|
||
|
||
|
||
def _discover_relay_url(tracker_url: str) -> str | None:
|
||
relay_url: str | None = None
|
||
try:
|
||
network_map = _get_json(f"{tracker_url}/v1/network/map", timeout=5.0)
|
||
raw = network_map.get("relay_url")
|
||
if isinstance(raw, str) and raw:
|
||
relay_url = raw
|
||
except Exception:
|
||
pass
|
||
return relay_url or _infer_relay_url_from_tracker(tracker_url)
|
||
|
||
|
||
def _start_relay_bridge_if_available(
|
||
tracker_url: str,
|
||
wallet_address: str,
|
||
local_base_url: str,
|
||
advertised_endpoint: str,
|
||
relay_url: str | None = None,
|
||
) -> tuple[RelayHttpBridge | None, dict]:
|
||
relay_url = relay_url or _discover_relay_url(tracker_url)
|
||
if not relay_url:
|
||
return None, {}
|
||
peer_id = peer_id_from_wallet(wallet_address)
|
||
bridge = RelayHttpBridge(
|
||
relay_url=relay_url,
|
||
peer_id=peer_id,
|
||
local_base_url=local_base_url,
|
||
advertised_addr=advertised_endpoint,
|
||
)
|
||
info = bridge.start()
|
||
if bridge.wait_connected(timeout=5.0):
|
||
print(f" Relay connected — {info.relay_addr}", flush=True)
|
||
else:
|
||
print(f" Relay configured but not connected yet — {info.relay_addr}", flush=True)
|
||
return bridge, {
|
||
"relay_addr": info.relay_addr,
|
||
"peer_id": info.peer_id,
|
||
}
|
||
|
||
|
||
def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHttpBridge | None) -> None:
|
||
setattr(node, "relay_bridge", bridge)
|
||
if bridge is None:
|
||
return
|
||
original_stop = node.stop
|
||
|
||
def _stop_with_bridge() -> None:
|
||
try:
|
||
bridge.stop()
|
||
finally:
|
||
original_stop()
|
||
|
||
node.stop = _stop_with_bridge # type: ignore[method-assign]
|
||
|
||
|
||
def _start_heartbeat(
|
||
tracker_url: str,
|
||
node_id: str,
|
||
register_payload: dict,
|
||
interval: float = 20.0,
|
||
node_ref: Any | None = None,
|
||
start_time: float | None = None,
|
||
) -> threading.Thread:
|
||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||
|
||
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||
outage and flushed on next successful heartbeat.
|
||
|
||
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||
which is logged for now (hot-reload implemented in US-026).
|
||
"""
|
||
_start_time = start_time or time.monotonic()
|
||
|
||
def _get_stats() -> dict:
|
||
uptime = time.monotonic() - _start_time
|
||
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||
if node_ref is not None:
|
||
stats["total_requests"] = getattr(
|
||
node_ref,
|
||
"total_requests",
|
||
getattr(node_ref, "chat_completion_count", 0),
|
||
)
|
||
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||
return stats
|
||
|
||
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 _apply_directives(directives: list[dict]) -> None:
|
||
if not directives:
|
||
return
|
||
if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
|
||
print(f" [node] tracker directives received: {directives}", flush=True)
|
||
return
|
||
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
|
||
if 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
|
||
register_payload["shard_start"] = applied["shard_start"]
|
||
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))
|
||
|
||
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:
|
||
resp = _post_json(hb_url, _get_stats())
|
||
_apply_directives(resp.get("directives", []))
|
||
new_asgn = resp.get("new_assignment")
|
||
if new_asgn:
|
||
print(
|
||
f" [node] tracker reassignment received: "
|
||
f"model={new_asgn.get('model')!r} "
|
||
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||
flush=True,
|
||
)
|
||
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 _warn_virtual_network_ip(ip: str | None) -> None:
|
||
"""Print a warning when the auto-detected advertise IP is in a known virtual-network range.
|
||
|
||
172.16.0.0/12 is used by Docker, WSL2, and most hypervisors. Nodes behind these
|
||
adapters are NOT directly reachable from other physical machines on the LAN, so
|
||
cross-host pipeline hops will time out. The user must pass --advertise-host with
|
||
their actual LAN IP (e.g. 192.168.x.x) to fix this.
|
||
"""
|
||
if ip is None:
|
||
return
|
||
try:
|
||
parts = [int(p) for p in ip.split(".")]
|
||
if len(parts) != 4:
|
||
return
|
||
a, b = parts[0], parts[1]
|
||
# 172.16.0.0/12 → 172.16–31.x.x
|
||
if a == 172 and 16 <= b <= 31:
|
||
print(
|
||
f"\n WARNING: auto-detected endpoint IP {ip} is in 172.16.0.0/12.\n"
|
||
f" This range is used by Docker, WSL2, and virtual machines and is\n"
|
||
f" NOT reachable from other physical machines on your LAN.\n"
|
||
f" Cross-host pipeline hops WILL time out.\n"
|
||
f" Fix: use a public tracker with relay (wss://…/ws), or pass\n"
|
||
f" --advertise-host <your-LAN-ip> (e.g. 192.168.x.x).\n",
|
||
flush=True,
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
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 = "auto",
|
||
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,
|
||
route_timeout: float = 30.0,
|
||
vram_mb_override: int | None = None,
|
||
max_loaded_shards: int = 1,
|
||
debug: bool = False,
|
||
) -> 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("/")
|
||
relay_url = _discover_relay_url(tracker_url)
|
||
if max_loaded_shards < 1:
|
||
raise ValueError("--max-shards must be at least 1")
|
||
|
||
# 1. Hardware detection
|
||
if advertise_host is None and host == "0.0.0.0":
|
||
# socket.getfqdn() returns an mDNS name (.local / .localdomain) that remote
|
||
# machines on a different OS or subnet often can't resolve. Instead, probe the
|
||
# outbound IP by opening a UDP socket toward the tracker — no data is sent.
|
||
try:
|
||
_tracker_host = urllib.parse.urlparse(tracker_url).hostname or "8.8.8.8"
|
||
_s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
_s.connect((_tracker_host, 80))
|
||
advertise_host = _s.getsockname()[0]
|
||
_s.close()
|
||
except Exception:
|
||
advertise_host = socket.getfqdn()
|
||
|
||
if relay_url:
|
||
print(f"Relay advertised by tracker — using outbound tunnel {relay_url}", flush=True)
|
||
else:
|
||
_warn_virtual_network_ip(advertise_host)
|
||
|
||
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)
|
||
shared_vram_mb: int = hw.get("shared_vram_mb", 0)
|
||
ram_mb: int = hw.get("ram_mb", 16 * 1024)
|
||
|
||
if vram_mb_override is not None:
|
||
vram_mb = vram_mb_override
|
||
shared_vram_mb = 0
|
||
print(f" Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True)
|
||
elif device == "cpu":
|
||
gpu_suffix = ""
|
||
if gpu_name and vram_mb > 0:
|
||
gpu_suffix = (
|
||
f"; CUDA inactive; detected {gpu_name} "
|
||
f"({vram_mb / 1024:.1f} GB dedicated VRAM, {shared_vram_mb / 1024:.1f} GB shared)"
|
||
)
|
||
print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM{gpu_suffix})", flush=True)
|
||
else:
|
||
shared_suffix = f", {shared_vram_mb / 1024:.1f} GB shared" if shared_vram_mb > 0 else ""
|
||
print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB dedicated VRAM{shared_suffix}, {ram_mb / 1024:.1f} GB RAM)", flush=True)
|
||
|
||
if vram_mb_override is not None:
|
||
memory_budget_mb = vram_mb
|
||
memory_budget_source = "memory override"
|
||
else:
|
||
memory_budget_mb, memory_budget_source = _memory_budget(device, vram_mb, ram_mb, shared_vram_mb)
|
||
assignment_vram_mb = memory_budget_mb if device == "cuda" or vram_mb_override is not None else 0
|
||
print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)
|
||
|
||
print("Benchmarking compute...", flush=True)
|
||
if device != "cuda" and gpu_name:
|
||
_cuda_score, cuda_ok, cuda_error = benchmark_throughput_checked("cuda")
|
||
hw["cuda_benchmark_ok"] = cuda_ok
|
||
if cuda_error:
|
||
hw["cuda_benchmark_error"] = cuda_error
|
||
if not cuda_ok:
|
||
print(f" CUDA benchmark unavailable: {cuda_error}; using CPU benchmark", flush=True)
|
||
bench_tps, bench_ok, bench_error = benchmark_throughput_checked(device)
|
||
hw["benchmark_device"] = device
|
||
hw["benchmark_ok"] = bench_ok
|
||
if bench_error:
|
||
hw["benchmark_error"] = bench_error
|
||
device_label = "GPU" if device == "cuda" else "CPU"
|
||
print(f" {device_label} throughput index: {bench_tps:,.0f}", flush=True)
|
||
|
||
registration_capabilities = {
|
||
"vram_bytes": max(0, int(assignment_vram_mb)) * 1024 * 1024,
|
||
"ram_bytes": max(0, int(ram_mb)) * 1024 * 1024,
|
||
"max_loaded_shards": max_loaded_shards,
|
||
"benchmark_tokens_per_sec": bench_tps,
|
||
}
|
||
# 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
|
||
user_pinned_shard = shard_start is not None or shard_end is not None
|
||
# Auto-detect shard range from model config if not explicitly provided
|
||
if shard_start is None or shard_end is None:
|
||
try:
|
||
detected = _detect_num_layers(model_id, cache_dir=cache_dir)
|
||
except TypeError:
|
||
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": assignment_vram_mb, "ram_mb": ram_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"]
|
||
full_sources = _full_model_sources(net_asgn.get("model_sources", []))
|
||
if full_sources:
|
||
cache_dir = download_shard(
|
||
model_id.split("/")[-1],
|
||
shard_start,
|
||
shard_end,
|
||
cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
|
||
hf_repo=model_id,
|
||
model_sources=full_sources,
|
||
)
|
||
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,
|
||
route_timeout=route_timeout,
|
||
cache_dir=cache_dir,
|
||
debug=debug,
|
||
)
|
||
_node_start_time = time.monotonic()
|
||
actual_port = node.start()
|
||
total_layers = getattr(getattr(node, "backend", None), "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}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
# Register with tracker so other nodes can auto-join this model.
|
||
total_layers = getattr(getattr(node, "backend", None), "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),
|
||
"managed_assignment": not user_pinned_shard,
|
||
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
}
|
||
tracker_node_id: str | None = None
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||
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=_node_start_time)
|
||
except Exception as exc:
|
||
setattr(node, "tracker_node_id", None)
|
||
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" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization)}\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\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.
|
||
# Skipped when the user explicitly requested a model — the shard-assignment
|
||
# query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
|
||
# tracker would otherwise print a scary 503 for the model-less auto-join.
|
||
net_assignment: dict = {}
|
||
if model and model != "stub-model":
|
||
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
|
||
else:
|
||
print("Querying tracker for network assignment...", flush=True)
|
||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
||
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"]
|
||
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
|
||
print(
|
||
f" Assigned: {assigned_hf_repo} "
|
||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"(of {assigned_num_layers})",
|
||
flush=True,
|
||
)
|
||
full_sources = _full_model_sources(assigned_model_sources)
|
||
if full_sources:
|
||
print("Downloading assigned model snapshot...", flush=True)
|
||
cache_dir = download_shard(
|
||
assigned_hf_repo.split("/")[-1],
|
||
assigned_shard_start,
|
||
assigned_shard_end,
|
||
cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
|
||
hf_repo=assigned_hf_repo,
|
||
model_sources=full_sources,
|
||
)
|
||
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,
|
||
route_timeout=route_timeout,
|
||
cache_dir=cache_dir,
|
||
debug=debug,
|
||
)
|
||
_node_start_time = time.monotonic()
|
||
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}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
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),
|
||
"managed_assignment": True,
|
||
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
}
|
||
tracker_node_id = None
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
|
||
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, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
||
except Exception as exc:
|
||
setattr(node, "tracker_node_id", None)
|
||
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" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\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,
|
||
"ram_mb": ram_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", [])
|
||
model_sources: list[dict] = assignment.get("model_sources", [])
|
||
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
|
||
if model_sources:
|
||
dl_kwargs["model_sources"] = model_sources
|
||
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}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
|
||
# 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,
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
},
|
||
)
|
||
node_id = str(reg_resp["node_id"])
|
||
setattr(node, "tracker_node_id", node_id)
|
||
except Exception:
|
||
node.stop()
|
||
raise
|
||
|
||
# Status summary
|
||
hw_str = device.upper()
|
||
if gpu_name:
|
||
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
|
||
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" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization)}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {node_id}\n"
|
||
f" Hardware: {hw_str}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
|
||
return node
|
||
|
||
|
||
def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | None:
|
||
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
|
||
try:
|
||
from transformers import AutoConfig # type: ignore[import]
|
||
|
||
from .model_catalog import layers_from_config
|
||
|
||
cfg = AutoConfig.from_pretrained(
|
||
model_id,
|
||
cache_dir=str(cache_dir) if cache_dir is not None else None,
|
||
)
|
||
layers = layers_from_config(cfg)
|
||
if layers is None:
|
||
print(
|
||
f" Warning: no layer count in {model_id} config "
|
||
"(checked top level and text_config)", flush=True,
|
||
)
|
||
return 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"
|