feat: add node startup flow

This commit is contained in:
Dobromir Popov
2026-06-29 01:30:12 +03:00
parent 15897d5c95
commit 0199c99f6b
9 changed files with 727 additions and 11 deletions

View File

@@ -2,6 +2,7 @@
import argparse
import sys
import time
def main() -> None:
@@ -12,20 +13,34 @@ def main() -> None:
subparsers = parser.add_subparsers(dest="command")
start_cmd = subparsers.add_parser("start", help="Start the node server")
start_cmd.add_argument("--tracker", default="http://localhost:8080", help="Tracker URL")
start_cmd.add_argument(
"--tracker", default="http://localhost:8080", help="Tracker URL"
)
start_cmd.add_argument("--port", type=int, default=7000, help="Port to listen on")
start_cmd.add_argument(
"--model", default="stub-model", help="Model preset to request from tracker"
)
start_cmd.add_argument(
"--host", default="0.0.0.0", help="Interface to bind to"
)
start_cmd.add_argument(
"--advertise-host",
help="Reachable host/IP to advertise to the tracker (defaults to FQDN when binding 0.0.0.0)",
)
args = parser.parse_args()
if args.command == "start":
from meshnet_node.server import StubNodeServer
from meshnet_node.startup import run_startup
node = StubNodeServer(port=args.port)
port = node.start()
print(f"meshnet-node listening on port {port}")
print(f"Tracker URL: {args.tracker}")
node = run_startup(
tracker_url=args.tracker,
port=args.port,
model=args.model,
host=args.host,
advertise_host=args.advertise_host,
)
try:
import time
while True:
time.sleep(1)
except KeyboardInterrupt:

View File

@@ -0,0 +1,62 @@
"""Shard downloader — fetches or stubs a model shard from HuggingFace Hub.
Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/
For "stub-model" (no HF repo), a placeholder JSON file is written so the
test suite never touches the network.
"""
import json
from pathlib import Path
_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"
def download_shard(
model: str,
shard_start: int,
shard_end: int,
cache_dir: Path = _DEFAULT_CACHE,
hf_repo: str | None = None,
progress: bool = True,
) -> Path:
"""Ensure the shard is present in *cache_dir* and return its local path.
When *hf_repo* is None (or *model* is ``"stub-model"``), a placeholder
weights file is created locally — no network access required. This keeps
the test suite hermetic while the real download path is exercised by
passing a non-stub *hf_repo*.
"""
shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
shard_dir.mkdir(parents=True, exist_ok=True)
if hf_repo is None or model == "stub-model":
stub_file = shard_dir / "weights.json"
if not stub_file.exists():
stub_file.write_text(json.dumps({
"model": model,
"shard_start": shard_start,
"shard_end": shard_end,
"stub": True,
}))
if progress:
print(f" [stub] shard placeholder written to {stub_file}", flush=True)
else:
if progress:
print(f" [stub] shard already cached at {shard_dir}", flush=True)
return shard_dir
from huggingface_hub import snapshot_download # type: ignore[import]
if progress:
print(
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
flush=True,
)
local_dir = snapshot_download(
repo_id=hf_repo,
cache_dir=str(cache_dir),
local_dir=str(shard_dir),
)
return Path(local_dir)

View File

@@ -0,0 +1,33 @@
"""GPU hardware detection with graceful CPU fallback."""
import subprocess
def detect_hardware() -> dict:
"""Detect GPU model and available VRAM. Returns hardware profile dict."""
try:
import torch # type: ignore[import]
if torch.cuda.is_available():
idx = torch.cuda.current_device()
name = torch.cuda.get_device_name(idx)
props = torch.cuda.get_device_properties(idx)
vram_mb = props.total_memory // (1024 * 1024)
return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb}
except ImportError:
pass
try:
result = subprocess.run(
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
capture_output=True, text=True, timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
line = result.stdout.strip().splitlines()[0]
parts = line.split(",", 1)
gpu_name = parts[0].strip()
vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0
return {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb}
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
pass
return {"device": "cpu", "gpu_name": None, "vram_mb": 0}

View File

@@ -0,0 +1,154 @@
"""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 .downloader import download_shard
from .hardware import detect_hardware
from .server import StubNodeServer
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",
wallet_path: Path | None = None,
cache_dir: Path | None = None,
host: str = "127.0.0.1",
advertise_host: str | None = None,
) -> StubNodeServer:
"""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
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)
# 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")
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
shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
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,
)
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,
"shard_start": shard_start,
"shard_end": shard_end,
"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

View File

@@ -0,0 +1,64 @@
"""Solana wallet management — load or generate an Ed25519 keypair.
Solana keypair format: 64-byte array stored as JSON list of ints.
bytes 0-31: private scalar (seed)
bytes 32-63: public key (on-curve point)
Wallet address: base58-encoded 32-byte public key.
"""
import json
from pathlib import Path
_DEFAULT_WALLET_PATH = Path.home() / ".config" / "meshnet" / "wallet.json"
_B58_ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
def _b58encode(data: bytes) -> str:
n = int.from_bytes(data, "big")
chars: list[str] = []
while n:
n, rem = divmod(n, 58)
chars.append(_B58_ALPHABET[rem])
leading = 0
for b in data:
if b == 0:
leading += 1
else:
break
return "1" * leading + "".join(reversed(chars))
def load_or_create_wallet(
path: Path = _DEFAULT_WALLET_PATH,
) -> tuple[bytes, bytes, str]:
"""Return (secret_32, public_32, address_base58).
Loads from *path* if it exists; otherwise generates a new keypair and
saves it. The parent directory is created if needed.
"""
if path.exists():
mode = path.stat().st_mode & 0o777
if mode & 0o077:
path.chmod(0o600)
raw = bytes(json.loads(path.read_text()))
if len(raw) != 64:
raise ValueError(f"Wallet at {path} has unexpected length {len(raw)}, expected 64")
secret, public = raw[:32], raw[32:]
return secret, public, _b58encode(public)
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.serialization import (
Encoding, NoEncryption, PrivateFormat, PublicFormat,
)
priv = Ed25519PrivateKey.generate()
secret = priv.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
public = priv.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
keypair = list(secret) + list(public)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(keypair))
path.chmod(0o600)
return secret, public, _b58encode(public)

View File

@@ -8,6 +8,11 @@ version = "0.1.0"
description = "Distributed Inference Network node client"
requires-python = ">=3.10"
dependencies = [
"cryptography>=41",
"huggingface-hub>=0.20",
]
[project.scripts]
meshnet-node = "meshnet_node.cli:main"