080d49b2c2
Layer count is now fetched from the curated catalog (zero network calls for known models) or via AutoConfig.from_pretrained() (~1 KB config.json only) when model_id is given without --shard-start/--shard-end. - model_catalog: add detect_num_layers(), two small Qwen models at top - startup: _detect_num_layers() helper; shard range auto-derived - wizard: show detected layer count for custom HF repos - tests: 3 new tests for auto-shard; fix catalog-order assumptions Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
"""Full node startup sequence — self-configuring, non-interactive."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import socket
|
||
import sys
|
||
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 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 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."
|
||
)
|
||
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,
|
||
)
|
||
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}"
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready\n"
|
||
f" Wallet: {address}\n"
|
||
f" Model ID: {model_id}\n"
|
||
f" Shard: layers {shard_start}–{shard_end}\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")
|
||
|
||
# 3. Shard assignment from tracker
|
||
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"
|