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

feat: add p2p shard swarm

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 10:15:01 +03:00
parent 792a9fd97f
commit e9c9bf63bc
7 changed files with 448 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

195
tests/test_node_startup.py
View File

@@ -1,6 +1,7 @@
"""US-004 integration tests: node self-configuring startup sequence."""
import json
import io
import sys
import types
import urllib.request
@@ -8,7 +9,7 @@ from pathlib import Path
import pytest
from meshnet_node.downloader import download_shard
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
@@ -123,6 +124,89 @@ def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeyp
}]
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)
@@ -220,6 +304,49 @@ def test_tracker_assign_returns_huggingface_repo_when_configured():
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()
# ---------------------------------------------------------------------------
# Full startup integration test
# ---------------------------------------------------------------------------
@@ -277,6 +404,72 @@ def test_full_startup_sequence(tmp_path):
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

44
tests/test_tracker_routing.py
View File

@@ -238,6 +238,50 @@ def test_tracker_registration_rejects_invalid_payload():
tracker.stop()
def test_tracker_routes_only_nodes_for_requested_model():
"""A node registered for one model cannot satisfy another model's route."""
tracker = TrackerServer(model_presets={
"model-a": {"layers_start": 0, "layers_end": 31},
"model-b": {"layers_start": 0, "layers_end": 31},
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
)
try:
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=model-b")
raise AssertionError("Expected model-b route to ignore model-a nodes")
except urllib.error.HTTPError as exc:
assert exc.code == 503
finally:
tracker.stop()
def test_tracker_assignment_coverage_is_model_scoped():
"""Shard assignment gaps are computed only from nodes serving the same model."""
tracker = TrackerServer(model_presets={
"model-a": {"layers_start": 0, "layers_end": 31},
"model-b": {"layers_start": 0, "layers_end": 31},
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
"model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
)
assignment = _get_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/assign?model=model-b&device=cpu")
assert assignment["shard_start"] == 0
assert assignment["shard_end"] == 15
finally:
tracker.stop()
def test_tracker_excludes_banned_wallets_from_routes():
"""Tracker refuses route candidates whose wallet is banned on-chain."""
contracts = LocalSolanaContracts()