feat: add node startup flow
This commit is contained in:
@@ -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:
|
||||
|
||||
62
packages/node/meshnet_node/downloader.py
Normal file
62
packages/node/meshnet_node/downloader.py
Normal 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)
|
||||
33
packages/node/meshnet_node/hardware.py
Normal file
33
packages/node/meshnet_node/hardware.py
Normal 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}
|
||||
154
packages/node/meshnet_node/startup.py
Normal file
154
packages/node/meshnet_node/startup.py
Normal 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
|
||||
64
packages/node/meshnet_node/wallet.py
Normal file
64
packages/node/meshnet_node/wallet.py
Normal 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)
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user