540 lines
18 KiB
Python
540 lines
18 KiB
Python
"""US-004 integration tests: node self-configuring startup sequence."""
|
||
|
||
import json
|
||
import io
|
||
import sys
|
||
import types
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from meshnet_node.downloader import download_shard, write_shard_archive
|
||
from meshnet_node.hardware import detect_hardware
|
||
from meshnet_node.startup import _probationary_status_line, run_startup
|
||
from meshnet_node.wallet import _b58encode, load_or_create_wallet
|
||
from meshnet_contracts import LocalSolanaContracts
|
||
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_logs_huggingface_source(tmp_path, monkeypatch, capsys):
|
||
"""Shard download status tells the node operator when HuggingFace was used."""
|
||
|
||
def fake_snapshot_download(repo_id, cache_dir, local_dir):
|
||
Path(local_dir).mkdir(parents=True, exist_ok=True)
|
||
(Path(local_dir) / "weights.json").write_text(json.dumps({"repo_id": repo_id}))
|
||
return local_dir
|
||
|
||
monkeypatch.setitem(
|
||
sys.modules,
|
||
"huggingface_hub",
|
||
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
|
||
)
|
||
|
||
download_shard(
|
||
"tiny-llama",
|
||
0,
|
||
3,
|
||
cache_dir=tmp_path,
|
||
hf_repo="org/tiny-llama-shards",
|
||
)
|
||
|
||
assert "download source: HuggingFace" in capsys.readouterr().out
|
||
|
||
|
||
def test_download_shard_rejects_peer_checksum_mismatch_before_fallback(
|
||
tmp_path,
|
||
monkeypatch,
|
||
):
|
||
"""Corrupt peer chunks are not marked complete; HuggingFace remains the fallback."""
|
||
corrupt_dir = tmp_path / "corrupt"
|
||
corrupt_dir.mkdir()
|
||
(corrupt_dir / "weights.json").write_text(json.dumps({"payload": "corrupt"}))
|
||
archive = io.BytesIO()
|
||
write_shard_archive(corrupt_dir, archive)
|
||
|
||
class FakePeerResponse:
|
||
def __init__(self, payload: bytes):
|
||
self._payload = io.BytesIO(payload)
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
return False
|
||
|
||
def read(self, size: int = -1) -> bytes:
|
||
return self._payload.read(size)
|
||
|
||
monkeypatch.setattr(
|
||
urllib.request,
|
||
"urlopen",
|
||
lambda *args, **kwargs: FakePeerResponse(archive.getvalue()),
|
||
)
|
||
hf_calls = []
|
||
|
||
def fake_snapshot_download(repo_id, cache_dir, local_dir):
|
||
hf_calls.append(repo_id)
|
||
shard_dir = Path(local_dir)
|
||
shard_dir.mkdir(parents=True, exist_ok=True)
|
||
(shard_dir / "weights.json").write_text(json.dumps({"payload": "hf"}))
|
||
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 / "cache",
|
||
hf_repo="org/tiny-llama-shards",
|
||
peers=[{"endpoint": "http://peer", "checksum": "not-the-corrupt-checksum"}],
|
||
progress=False,
|
||
)
|
||
|
||
assert hf_calls == ["org/tiny-llama-shards"]
|
||
assert json.loads((shard_dir / "weights.json").read_text()) == {"payload": "hf"}
|
||
|
||
|
||
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()
|
||
|
||
|
||
def test_startup_formats_probationary_jobs_remaining():
|
||
"""Startup status tells a node how many free jobs remain before earning."""
|
||
contracts = LocalSolanaContracts(probationary_job_count=50)
|
||
for _ in range(12):
|
||
contracts.registry.record_completed_job("node-wallet-a")
|
||
|
||
line = _probationary_status_line(contracts, "node-wallet-a")
|
||
|
||
assert line == "Probationary period: 38 jobs remaining before earning"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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()
|
||
|
||
|
||
def test_tracker_assign_lists_peers_for_same_model_shard():
|
||
"""A registered node with a completed shard is returned as a same-shard peer."""
|
||
import json as _json
|
||
import urllib.request as _ur
|
||
|
||
tracker = TrackerServer(model_presets={
|
||
"tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
|
||
})
|
||
port = tracker.start()
|
||
try:
|
||
data = _json.dumps({
|
||
"endpoint": "http://127.0.0.1:9100",
|
||
"model": "tiny-llama",
|
||
"shard_start": 0,
|
||
"shard_end": 15,
|
||
"shard_checksum": "abc123",
|
||
"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()
|
||
|
||
resp = _get_json(
|
||
f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cpu&vram_mb=0"
|
||
)
|
||
|
||
assert resp["shard_start"] == 0
|
||
assert resp["shard_end"] == 15
|
||
assert resp["hf_repo"] == "org/tiny-llama-shards"
|
||
assert resp["peers"] == [{
|
||
"endpoint": "http://127.0.0.1:9100",
|
||
"checksum": "abc123",
|
||
}]
|
||
finally:
|
||
tracker.stop()
|
||
|
||
|
||
def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, capsys):
|
||
"""Real-model startup summary prints the shard range plus total model layers."""
|
||
import meshnet_node.startup as startup_mod
|
||
|
||
class FakeBackend:
|
||
total_layers = 24
|
||
|
||
class FakeTorchNodeServer:
|
||
def __init__(self, **kwargs):
|
||
self.kwargs = kwargs
|
||
self.backend = FakeBackend()
|
||
self.port = None
|
||
|
||
def start(self):
|
||
self.port = 8001
|
||
return self.port
|
||
|
||
monkeypatch.setattr(
|
||
startup_mod,
|
||
"detect_hardware",
|
||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||
)
|
||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||
|
||
node = run_startup(
|
||
tracker_url="http://127.0.0.1:8080",
|
||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||
shard_start=0,
|
||
shard_end=23,
|
||
wallet_path=tmp_path / "wallet.json",
|
||
)
|
||
|
||
assert node.backend.total_layers == 24
|
||
output = capsys.readouterr().out
|
||
assert "Shard: layers 0–23; 24 of 24" in output
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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_second_node_downloads_same_shard_from_peer_without_huggingface(
|
||
tmp_path,
|
||
monkeypatch,
|
||
capsys,
|
||
):
|
||
"""Node A downloads from HF stub; node B downloads same assignment from node A."""
|
||
import meshnet_node.startup as startup_mod
|
||
|
||
monkeypatch.setattr(
|
||
startup_mod,
|
||
"detect_hardware",
|
||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||
)
|
||
hf_calls = []
|
||
|
||
def fake_snapshot_download(repo_id, cache_dir, local_dir):
|
||
hf_calls.append({"repo_id": repo_id, "local_dir": local_dir})
|
||
shard_dir = Path(local_dir)
|
||
shard_dir.mkdir(parents=True, exist_ok=True)
|
||
(shard_dir / "weights.json").write_text(json.dumps({
|
||
"repo_id": repo_id,
|
||
"payload": "node-a-hf-stub",
|
||
}))
|
||
return local_dir
|
||
|
||
monkeypatch.setitem(
|
||
sys.modules,
|
||
"huggingface_hub",
|
||
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
|
||
)
|
||
|
||
tracker = TrackerServer(model_presets={
|
||
"tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
|
||
})
|
||
tracker_port = tracker.start()
|
||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||
nodes = []
|
||
try:
|
||
node_a = run_startup(
|
||
tracker_url=tracker_url,
|
||
model="tiny-llama",
|
||
wallet_path=tmp_path / "wallet-a.json",
|
||
cache_dir=tmp_path / "node-a-shards",
|
||
)
|
||
nodes.append(node_a)
|
||
assert len(hf_calls) == 1
|
||
|
||
node_b = run_startup(
|
||
tracker_url=tracker_url,
|
||
model="tiny-llama",
|
||
wallet_path=tmp_path / "wallet-b.json",
|
||
cache_dir=tmp_path / "node-b-shards",
|
||
)
|
||
nodes.append(node_b)
|
||
|
||
assert len(hf_calls) == 1
|
||
assert (tmp_path / "node-b-shards" / "tiny-llama" / "layers_0-15" / "weights.json").exists()
|
||
output = capsys.readouterr().out
|
||
assert "download source: HuggingFace" in output
|
||
assert "download source: peer" in output
|
||
finally:
|
||
for node in reversed(nodes):
|
||
node.stop()
|
||
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()
|