feat: add node startup flow
This commit is contained in:
@@ -93,11 +93,12 @@
|
|||||||
"Commit only this story's code changes with a focused conventional commit message"
|
"Commit only this story's code changes with a focused conventional commit message"
|
||||||
],
|
],
|
||||||
"priority": 4,
|
"priority": 4,
|
||||||
"passes": false,
|
"passes": true,
|
||||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/04-node-client-startup.md",
|
"notes": "Source issue: .scratch/distributed-inference-network/issues/04-node-client-startup.md",
|
||||||
"dependsOn": [
|
"dependsOn": [
|
||||||
"US-003"
|
"US-003"
|
||||||
]
|
],
|
||||||
|
"completionNotes": "Completed by Ralph iteration 86510a10; verified by pytest, compileall, editable installs, CLI help, and final review."
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"id": "US-005",
|
"id": "US-005",
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import sys
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
@@ -12,20 +13,34 @@ def main() -> None:
|
|||||||
subparsers = parser.add_subparsers(dest="command")
|
subparsers = parser.add_subparsers(dest="command")
|
||||||
|
|
||||||
start_cmd = subparsers.add_parser("start", help="Start the node server")
|
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("--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()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.command == "start":
|
if args.command == "start":
|
||||||
from meshnet_node.server import StubNodeServer
|
from meshnet_node.startup import run_startup
|
||||||
|
|
||||||
node = StubNodeServer(port=args.port)
|
node = run_startup(
|
||||||
port = node.start()
|
tracker_url=args.tracker,
|
||||||
print(f"meshnet-node listening on port {port}")
|
port=args.port,
|
||||||
print(f"Tracker URL: {args.tracker}")
|
model=args.model,
|
||||||
|
host=args.host,
|
||||||
|
advertise_host=args.advertise_host,
|
||||||
|
)
|
||||||
try:
|
try:
|
||||||
import time
|
|
||||||
while True:
|
while True:
|
||||||
time.sleep(1)
|
time.sleep(1)
|
||||||
except KeyboardInterrupt:
|
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"
|
description = "Distributed Inference Network node client"
|
||||||
requires-python = ">=3.10"
|
requires-python = ">=3.10"
|
||||||
|
|
||||||
|
dependencies = [
|
||||||
|
"cryptography>=41",
|
||||||
|
"huggingface-hub>=0.20",
|
||||||
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
meshnet-node = "meshnet_node.cli:main"
|
meshnet-node = "meshnet_node.cli:main"
|
||||||
|
|
||||||
|
|||||||
@@ -4,7 +4,8 @@ HTTP API contract:
|
|||||||
- POST /v1/nodes/register
|
- POST /v1/nodes/register
|
||||||
Request JSON: {
|
Request JSON: {
|
||||||
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
|
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
|
||||||
"hardware_profile": object, "score": number optional
|
"hardware_profile": object, "wallet_address": str optional,
|
||||||
|
"score": number optional
|
||||||
}
|
}
|
||||||
Response 200: {"node_id": str}
|
Response 200: {"node_id": str}
|
||||||
Response 400: {"error": str}
|
Response 400: {"error": str}
|
||||||
@@ -32,7 +33,7 @@ DEFAULT_MODEL_PRESETS: dict[str, dict] = {
|
|||||||
class _NodeEntry:
|
class _NodeEntry:
|
||||||
__slots__ = (
|
__slots__ = (
|
||||||
"node_id", "endpoint", "shard_start", "shard_end",
|
"node_id", "endpoint", "shard_start", "shard_end",
|
||||||
"hardware_profile", "score", "last_heartbeat",
|
"hardware_profile", "wallet_address", "score", "last_heartbeat",
|
||||||
)
|
)
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -42,6 +43,7 @@ class _NodeEntry:
|
|||||||
shard_start: int,
|
shard_start: int,
|
||||||
shard_end: int,
|
shard_end: int,
|
||||||
hardware_profile: dict,
|
hardware_profile: dict,
|
||||||
|
wallet_address: str | None,
|
||||||
score: float,
|
score: float,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.node_id = node_id
|
self.node_id = node_id
|
||||||
@@ -49,6 +51,7 @@ class _NodeEntry:
|
|||||||
self.shard_start = shard_start
|
self.shard_start = shard_start
|
||||||
self.shard_end = shard_end
|
self.shard_end = shard_end
|
||||||
self.hardware_profile = hardware_profile
|
self.hardware_profile = hardware_profile
|
||||||
|
self.wallet_address = wallet_address
|
||||||
self.score = score
|
self.score = score
|
||||||
self.last_heartbeat: float = time.monotonic()
|
self.last_heartbeat: float = time.monotonic()
|
||||||
|
|
||||||
@@ -146,6 +149,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
parsed = urllib.parse.urlparse(self.path)
|
parsed = urllib.parse.urlparse(self.path)
|
||||||
if parsed.path == "/v1/route":
|
if parsed.path == "/v1/route":
|
||||||
self._handle_route(parsed)
|
self._handle_route(parsed)
|
||||||
|
elif parsed.path == "/v1/nodes/assign":
|
||||||
|
self._handle_assign(parsed)
|
||||||
else:
|
else:
|
||||||
self.send_response(404)
|
self.send_response(404)
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
@@ -180,6 +185,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if not isinstance(hardware_profile, dict):
|
if not isinstance(hardware_profile, dict):
|
||||||
self._send_json(400, {"error": "hardware_profile must be an object"})
|
self._send_json(400, {"error": "hardware_profile must be an object"})
|
||||||
return
|
return
|
||||||
|
wallet_address = body.get("wallet_address")
|
||||||
|
if wallet_address is not None and not isinstance(wallet_address, str):
|
||||||
|
self._send_json(400, {"error": "wallet_address must be a string"})
|
||||||
|
return
|
||||||
|
|
||||||
node_id = str(uuid.uuid4())
|
node_id = str(uuid.uuid4())
|
||||||
entry = _NodeEntry(
|
entry = _NodeEntry(
|
||||||
@@ -188,6 +197,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
shard_start=shard_start,
|
shard_start=shard_start,
|
||||||
shard_end=shard_end,
|
shard_end=shard_end,
|
||||||
hardware_profile=hardware_profile,
|
hardware_profile=hardware_profile,
|
||||||
|
wallet_address=wallet_address,
|
||||||
score=score,
|
score=score,
|
||||||
)
|
)
|
||||||
with server.lock:
|
with server.lock:
|
||||||
@@ -207,6 +217,81 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
entry.last_heartbeat = time.monotonic()
|
entry.last_heartbeat = time.monotonic()
|
||||||
self._send_json(200, {})
|
self._send_json(200, {})
|
||||||
|
|
||||||
|
def _handle_assign(self, parsed: urllib.parse.ParseResult):
|
||||||
|
"""Return an optimal shard assignment for a node given its hardware profile.
|
||||||
|
|
||||||
|
Query params:
|
||||||
|
model — model preset name (default: first preset)
|
||||||
|
device — "cuda" | "cpu"
|
||||||
|
vram_mb — integer VRAM in MB (0 for CPU)
|
||||||
|
|
||||||
|
The greedy strategy: find the first gap in current layer coverage
|
||||||
|
and assign it. If no gap exists, assign the full model range so the
|
||||||
|
node provides redundant coverage.
|
||||||
|
"""
|
||||||
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||||
|
params = urllib.parse.parse_qs(parsed.query)
|
||||||
|
|
||||||
|
model_list = params.get("model")
|
||||||
|
if not model_list:
|
||||||
|
model = next(iter(server.model_presets), None)
|
||||||
|
if model is None:
|
||||||
|
self._send_json(503, {"error": "no model presets configured"})
|
||||||
|
return
|
||||||
|
else:
|
||||||
|
model = model_list[0]
|
||||||
|
|
||||||
|
preset = server.model_presets.get(model)
|
||||||
|
if preset is None:
|
||||||
|
self._send_json(404, {"error": f"unknown model preset: {model!r}"})
|
||||||
|
return
|
||||||
|
|
||||||
|
required_start: int = preset["layers_start"]
|
||||||
|
required_end: int = preset["layers_end"]
|
||||||
|
|
||||||
|
with server.lock:
|
||||||
|
self._purge_expired_nodes()
|
||||||
|
alive = list(server.registry.values())
|
||||||
|
|
||||||
|
device = params.get("device", ["cpu"])[0]
|
||||||
|
try:
|
||||||
|
vram_mb = int(params.get("vram_mb", ["0"])[0])
|
||||||
|
except ValueError:
|
||||||
|
vram_mb = 0
|
||||||
|
max_layers = required_end - required_start + 1
|
||||||
|
if device != "cuda" or vram_mb < 8192:
|
||||||
|
max_layers = min(max_layers, 16)
|
||||||
|
|
||||||
|
# Collect covered intervals sorted by start layer.
|
||||||
|
covered = sorted(
|
||||||
|
[(n.shard_start, n.shard_end) for n in alive],
|
||||||
|
key=lambda t: t[0],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Walk from required_start to find the first uncovered layer.
|
||||||
|
gap_start = required_start
|
||||||
|
for s, e in covered:
|
||||||
|
if s <= gap_start:
|
||||||
|
gap_start = max(gap_start, e + 1)
|
||||||
|
else:
|
||||||
|
break # gap found before this node
|
||||||
|
|
||||||
|
if gap_start > required_end:
|
||||||
|
# Full coverage exists — assign the full range for redundancy.
|
||||||
|
shard_start = required_start
|
||||||
|
shard_end = min(required_end, shard_start + max_layers - 1)
|
||||||
|
else:
|
||||||
|
shard_start = gap_start
|
||||||
|
shard_end = min(required_end, shard_start + max_layers - 1)
|
||||||
|
|
||||||
|
self._send_json(200, {
|
||||||
|
"shard_start": shard_start,
|
||||||
|
"shard_end": shard_end,
|
||||||
|
"model": model,
|
||||||
|
"model_layers_end": required_end,
|
||||||
|
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
||||||
|
})
|
||||||
|
|
||||||
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
||||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||||
params = urllib.parse.parse_qs(parsed.query)
|
params = urllib.parse.parse_qs(parsed.query)
|
||||||
|
|||||||
297
tests/test_node_startup.py
Normal file
297
tests/test_node_startup.py
Normal file
@@ -0,0 +1,297 @@
|
|||||||
|
"""US-004 integration tests: node self-configuring startup sequence."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
import urllib.request
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from meshnet_node.downloader import download_shard
|
||||||
|
from meshnet_node.hardware import detect_hardware
|
||||||
|
from meshnet_node.startup import run_startup
|
||||||
|
from meshnet_node.wallet import _b58encode, load_or_create_wallet
|
||||||
|
from meshnet_tracker.server import TrackerServer
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Unit tests — hardware, wallet, downloader
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def test_detect_hardware_returns_valid_profile():
|
||||||
|
"""Hardware detection always returns a dict with required keys."""
|
||||||
|
hw = detect_hardware()
|
||||||
|
assert hw["device"] in {"cuda", "cpu"}
|
||||||
|
assert isinstance(hw.get("vram_mb"), int)
|
||||||
|
if hw["device"] == "cpu":
|
||||||
|
assert hw["gpu_name"] is None
|
||||||
|
assert hw["vram_mb"] == 0
|
||||||
|
else:
|
||||||
|
assert isinstance(hw["gpu_name"], str) and hw["gpu_name"]
|
||||||
|
assert hw["vram_mb"] > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_wallet_generates_new_keypair(tmp_path):
|
||||||
|
"""A new wallet is created when none exists, saved to disk."""
|
||||||
|
wallet_file = tmp_path / "wallet.json"
|
||||||
|
assert not wallet_file.exists()
|
||||||
|
|
||||||
|
secret, public, address = load_or_create_wallet(path=wallet_file)
|
||||||
|
|
||||||
|
assert wallet_file.exists()
|
||||||
|
assert wallet_file.stat().st_mode & 0o777 == 0o600
|
||||||
|
assert len(secret) == 32
|
||||||
|
assert len(public) == 32
|
||||||
|
assert len(address) > 20 # base58-encoded 32 bytes is 43-44 chars typically
|
||||||
|
# Solana addresses are base58 — no '0', 'O', 'I', 'l'
|
||||||
|
for ch in "0OIl":
|
||||||
|
assert ch not in address, f"Invalid base58 char {ch!r} in address {address!r}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_wallet_loads_existing_keypair(tmp_path):
|
||||||
|
"""Loading the same wallet file twice returns identical keys and address."""
|
||||||
|
wallet_file = tmp_path / "wallet.json"
|
||||||
|
|
||||||
|
secret1, public1, address1 = load_or_create_wallet(path=wallet_file)
|
||||||
|
secret2, public2, address2 = load_or_create_wallet(path=wallet_file)
|
||||||
|
|
||||||
|
assert secret1 == secret2
|
||||||
|
assert public1 == public2
|
||||||
|
assert address1 == address2
|
||||||
|
|
||||||
|
|
||||||
|
def test_wallet_load_repairs_insecure_permissions(tmp_path):
|
||||||
|
"""Existing private key files are tightened to owner-only permissions."""
|
||||||
|
wallet_file = tmp_path / "wallet.json"
|
||||||
|
load_or_create_wallet(path=wallet_file)
|
||||||
|
wallet_file.chmod(0o644)
|
||||||
|
|
||||||
|
load_or_create_wallet(path=wallet_file)
|
||||||
|
|
||||||
|
assert wallet_file.stat().st_mode & 0o777 == 0o600
|
||||||
|
|
||||||
|
|
||||||
|
def test_base58_counts_only_leading_zero_bytes():
|
||||||
|
"""Zero bytes inside the public key do not become extra base58 leading ones."""
|
||||||
|
assert _b58encode(bytes([0, 1, 0])) == "15R"
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_shard_stub_creates_cache(tmp_path):
|
||||||
|
"""Stub-model shard creates a local cache file without network access."""
|
||||||
|
shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path)
|
||||||
|
assert shard_dir.exists()
|
||||||
|
weights = shard_dir / "weights.json"
|
||||||
|
assert weights.exists()
|
||||||
|
data = json.loads(weights.read_text())
|
||||||
|
assert data["stub"] is True
|
||||||
|
assert data["shard_start"] == 0
|
||||||
|
assert data["shard_end"] == 31
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeypatch):
|
||||||
|
"""Non-stub shards use the HuggingFace snapshot_download path."""
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def fake_snapshot_download(repo_id, cache_dir, local_dir):
|
||||||
|
calls.append({"repo_id": repo_id, "cache_dir": cache_dir, "local_dir": local_dir})
|
||||||
|
Path(local_dir).mkdir(parents=True, exist_ok=True)
|
||||||
|
return local_dir
|
||||||
|
|
||||||
|
monkeypatch.setitem(
|
||||||
|
sys.modules,
|
||||||
|
"huggingface_hub",
|
||||||
|
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
|
||||||
|
)
|
||||||
|
|
||||||
|
shard_dir = download_shard(
|
||||||
|
"tiny-llama",
|
||||||
|
0,
|
||||||
|
3,
|
||||||
|
cache_dir=tmp_path,
|
||||||
|
hf_repo="org/tiny-llama-shards",
|
||||||
|
progress=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert shard_dir == tmp_path / "tiny-llama" / "layers_0-3"
|
||||||
|
assert calls == [{
|
||||||
|
"repo_id": "org/tiny-llama-shards",
|
||||||
|
"cache_dir": str(tmp_path),
|
||||||
|
"local_dir": str(shard_dir),
|
||||||
|
}]
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_shard_stub_idempotent(tmp_path):
|
||||||
|
"""Calling download_shard twice does not error — file already exists."""
|
||||||
|
download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False)
|
||||||
|
shard_dir = download_shard("stub-model", 0, 31, cache_dir=tmp_path, progress=False)
|
||||||
|
assert shard_dir.exists()
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Tracker assign endpoint
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _get_json(url: str) -> dict:
|
||||||
|
with urllib.request.urlopen(url, timeout=5) as r:
|
||||||
|
return json.loads(r.read())
|
||||||
|
|
||||||
|
|
||||||
|
def test_tracker_assign_returns_shard_for_empty_registry():
|
||||||
|
"""Tracker assigns the full layer range when no nodes are registered."""
|
||||||
|
tracker = TrackerServer()
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
resp = _get_json(
|
||||||
|
f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0"
|
||||||
|
)
|
||||||
|
assert resp["shard_start"] == 0
|
||||||
|
assert resp["shard_end"] == 15
|
||||||
|
assert resp["model"] == "stub-model"
|
||||||
|
assert resp["model_layers_end"] == 31
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_tracker_assign_fills_gap():
|
||||||
|
"""Tracker assigns the first uncovered layer range when a node is already registered."""
|
||||||
|
import json as _json
|
||||||
|
import urllib.request as _ur
|
||||||
|
|
||||||
|
tracker = TrackerServer()
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
# Register a node covering layers 0-15
|
||||||
|
data = _json.dumps({
|
||||||
|
"endpoint": "http://127.0.0.1:9100",
|
||||||
|
"shard_start": 0,
|
||||||
|
"shard_end": 15,
|
||||||
|
"hardware_profile": {},
|
||||||
|
"score": 1.0,
|
||||||
|
}).encode()
|
||||||
|
req = _ur.Request(
|
||||||
|
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||||
|
data=data,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
method="POST",
|
||||||
|
)
|
||||||
|
with _ur.urlopen(req) as r:
|
||||||
|
r.read()
|
||||||
|
|
||||||
|
# Assignment should fill the gap: layers 16-31
|
||||||
|
resp = _get_json(
|
||||||
|
f"http://127.0.0.1:{port}/v1/nodes/assign?model=stub-model&device=cpu&vram_mb=0"
|
||||||
|
)
|
||||||
|
assert resp["shard_start"] == 16
|
||||||
|
assert resp["shard_end"] == 31
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_tracker_assign_returns_huggingface_repo_when_configured():
|
||||||
|
"""Tracker includes the HuggingFace repo identifier in shard assignments."""
|
||||||
|
tracker = TrackerServer(model_presets={
|
||||||
|
"tiny-llama": {"layers_start": 0, "layers_end": 7, "hf_repo": "org/tiny-llama-shards"}
|
||||||
|
})
|
||||||
|
port = tracker.start()
|
||||||
|
try:
|
||||||
|
resp = _get_json(
|
||||||
|
f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cuda&vram_mb=24576"
|
||||||
|
)
|
||||||
|
assert resp["model"] == "tiny-llama"
|
||||||
|
assert resp["hf_repo"] == "org/tiny-llama-shards"
|
||||||
|
assert resp["shard_start"] == 0
|
||||||
|
assert resp["shard_end"] == 7
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Full startup integration test
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def test_full_startup_sequence(tmp_path):
|
||||||
|
"""Full startup: hardware → wallet → assign → download → start → register."""
|
||||||
|
tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||||
|
|
||||||
|
wallet_path = tmp_path / "wallet.json"
|
||||||
|
cache_dir = tmp_path / "shards"
|
||||||
|
|
||||||
|
try:
|
||||||
|
node = run_startup(
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
model="stub-model",
|
||||||
|
wallet_path=wallet_path,
|
||||||
|
cache_dir=cache_dir,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# Wallet was created on disk
|
||||||
|
assert wallet_path.exists()
|
||||||
|
wallet_data = json.loads(wallet_path.read_text())
|
||||||
|
assert len(wallet_data) == 64
|
||||||
|
|
||||||
|
# Shard was cached
|
||||||
|
assert any(cache_dir.rglob("weights.json"))
|
||||||
|
|
||||||
|
# Node appears in tracker registry (route resolves successfully)
|
||||||
|
route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model")
|
||||||
|
assert len(route_resp["route"]) >= 1
|
||||||
|
assert node.port is not None
|
||||||
|
assert any(str(node.port) in ep for ep in route_resp["route"])
|
||||||
|
|
||||||
|
# Node accepts an inference request
|
||||||
|
payload = json.dumps({
|
||||||
|
"messages": [{"role": "user", "content": "hello"}],
|
||||||
|
}).encode()
|
||||||
|
req = urllib.request.Request(
|
||||||
|
f"http://127.0.0.1:{node.port}/v1/infer",
|
||||||
|
data=payload,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
method="POST",
|
||||||
|
)
|
||||||
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||||
|
assert resp.status == 200
|
||||||
|
body = json.loads(resp.read())
|
||||||
|
# Last-shard node returns text; first-shard returns activations
|
||||||
|
assert "text" in body or "activations" in body
|
||||||
|
finally:
|
||||||
|
node.stop()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_startup_cpu_fallback(tmp_path, monkeypatch):
|
||||||
|
"""Node starts with CPU warning when no GPU is detected."""
|
||||||
|
import meshnet_node.startup as startup_mod
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"detect_hardware",
|
||||||
|
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||||
|
)
|
||||||
|
|
||||||
|
tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||||
|
|
||||||
|
try:
|
||||||
|
node = run_startup(
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
model="stub-model",
|
||||||
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
cache_dir=tmp_path / "shards",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# Node is running even on CPU
|
||||||
|
assert node.port is not None
|
||||||
|
route_resp = _get_json(f"{tracker_url}/v1/route?model=stub-model")
|
||||||
|
assert len(route_resp["route"]) >= 1
|
||||||
|
finally:
|
||||||
|
node.stop()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
Reference in New Issue
Block a user