feat: add p2p shard swarm
This commit is contained in:
Dobromir Popov
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user