654 lines
26 KiB
Python
654 lines
26 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 .relay_bridge import RelayHttpBridge, peer_id_from_wallet
|
||
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 _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", 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 = "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,
|
||
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)
|
||
|
||
if vram_mb_override is not None:
|
||
vram_mb = vram_mb_override
|
||
print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True)
|
||
elif 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)
|
||
registration_capabilities = {
|
||
"max_loaded_shards": max_loaded_shards,
|
||
}
|
||
if vram_mb_override is not None or vram_mb > 0:
|
||
registration_capabilities["vram_bytes"] = max(0, int(vram_mb)) * 1024 * 1024
|
||
|
||
# 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:
|
||
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,
|
||
route_timeout=route_timeout,
|
||
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,
|
||
**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" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\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,
|
||
route_timeout=route_timeout,
|
||
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,
|
||
**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" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\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}"
|
||
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} 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"
|