popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add node startup flow

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

297
tests/test_node_startup.py Normal file
View 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()