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

5
.scratch/distributed-inference-network/prd.json
View File

@@ -223,11 +223,12 @@
"Commit only this story's code changes with a focused conventional commit message"
],
"priority": 9,
"passes": false,
"passes": true,
"notes": "Source issue: .scratch/distributed-inference-network/issues/09-p2p-shard-swarm.md",
"dependsOn": [
"US-004"
]
],
"completionNotes": "Completed by fresh Ralph/Codex session 243fae88 with controller fix for streaming tar archives; verified by pytest, compileall, and diff check."
},
{
"id": "US-010",

114
packages/node/meshnet_node/downloader.py
View File

@@ -1,4 +1,4 @@
"""Shard downloader — fetches or stubs a model shard from HuggingFace Hub.
"""Shard downloader — fetches model shards from peers or HuggingFace Hub.
Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/
@@ -6,10 +6,103 @@ For "stub-model" (no HF repo), a placeholder JSON file is written so the
test suite never touches the network.
"""
import hashlib
import json
import shutil
import tarfile
import tempfile
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any
_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"
_PEER_TIMEOUT_SECONDS = 2.0
def compute_shard_checksum(shard_dir: Path) -> str:
"""Return a stable SHA256 checksum for all regular files in a shard."""
digest = hashlib.sha256()
for path in sorted(p for p in shard_dir.rglob("*") if p.is_file()):
rel_path = path.relative_to(shard_dir).as_posix()
digest.update(rel_path.encode())
digest.update(b"\0")
with path.open("rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
digest.update(chunk)
digest.update(b"\0")
return digest.hexdigest()
def write_shard_archive(shard_dir: Path, out_file: Any) -> None:
"""Write a tar archive for *shard_dir* to a binary file-like object."""
with tarfile.open(fileobj=out_file, mode="w|") as archive:
for path in sorted(p for p in shard_dir.rglob("*") if p.is_file()):
archive.add(path, arcname=path.relative_to(shard_dir).as_posix())
def _safe_extract_shard(archive_path: Path, target_dir: Path) -> None:
target_root = target_dir.resolve()
with tarfile.open(archive_path, mode="r") as archive:
for member in archive.getmembers():
dest = (target_dir / member.name).resolve()
if target_root != dest and target_root not in dest.parents:
raise ValueError("peer shard archive contains an unsafe path")
archive.extractall(target_dir)
def _peer_download_url(
endpoint: str,
model: str,
shard_start: int,
shard_end: int,
) -> str:
query = urllib.parse.urlencode({
"model": model,
"shard_start": shard_start,
"shard_end": shard_end,
})
return f"{endpoint.rstrip('/')}/v1/shards/download?{query}"
def _download_shard_from_peer(
peer: dict,
model: str,
shard_start: int,
shard_end: int,
shard_dir: Path,
timeout: float,
) -> bool:
endpoint = peer.get("endpoint")
checksum = peer.get("checksum")
if not isinstance(endpoint, str) or not isinstance(checksum, str):
return False
shard_dir.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(prefix="meshnet-peer-", dir=shard_dir.parent) as tmp:
tmp_root = Path(tmp)
archive_path = tmp_root / "shard.tar"
extract_dir = tmp_root / "extract"
extract_dir.mkdir()
try:
with urllib.request.urlopen(
_peer_download_url(endpoint, model, shard_start, shard_end),
timeout=timeout,
) as resp, archive_path.open("wb") as out:
while True:
chunk = resp.read(1024 * 1024)
if not chunk:
break
out.write(chunk)
_safe_extract_shard(archive_path, extract_dir)
if compute_shard_checksum(extract_dir) != checksum:
return False
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(extract_dir), str(shard_dir))
return True
except Exception:
return False
def download_shard(
@@ -19,6 +112,8 @@ def download_shard(
cache_dir: Path = _DEFAULT_CACHE,
hf_repo: str | None = None,
progress: bool = True,
peers: list[dict] | None = None,
peer_timeout: float = _PEER_TIMEOUT_SECONDS,
) -> Path:
"""Ensure the shard is present in *cache_dir* and return its local path.
@@ -28,6 +123,22 @@ def download_shard(
passing a non-stub *hf_repo*.
"""
shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
for peer in peers or []:
if progress:
print(f" Trying peer shard download from {peer.get('endpoint')} ...", flush=True)
if _download_shard_from_peer(
peer,
model,
shard_start,
shard_end,
shard_dir,
timeout=peer_timeout,
):
if progress:
print(" download source: peer", flush=True)
return shard_dir
shard_dir.mkdir(parents=True, exist_ok=True)
if hf_repo is None or model == "stub-model":
@@ -53,6 +164,7 @@ def download_shard(
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
flush=True,
)
print(" download source: HuggingFace", flush=True)
local_dir = snapshot_download(
repo_id=hf_repo,

51
packages/node/meshnet_node/server.py
View File

@@ -5,6 +5,10 @@ import http.server
import json
import struct
import threading
import urllib.parse
from pathlib import Path
from .downloader import compute_shard_checksum, write_shard_archive
# Activation wire format (contract for all shard nodes):
# JSON object: {"shape": [B, S, D], "dtype": "float32",
@@ -36,12 +40,16 @@ class _StubHTTPServer(http.server.HTTPServer):
shard_end: int,
is_last_shard: bool,
response_prefix: str,
model: str,
shard_path: Path | None,
):
super().__init__(addr, handler)
self.shard_start = shard_start
self.shard_end = shard_end
self.is_last_shard = is_last_shard
self.response_prefix = response_prefix
self.model = model
self.shard_path = shard_path
self.received_activations: bool = False
@@ -56,6 +64,43 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
self.send_response(404)
self.end_headers()
def do_GET(self):
parsed = urllib.parse.urlparse(self.path)
if parsed.path == "/v1/shards/download":
self._handle_shard_download(parsed)
else:
self.send_response(404)
self.end_headers()
def _handle_shard_download(self, parsed: urllib.parse.ParseResult):
server: _StubHTTPServer = self.server # type: ignore[assignment]
params = urllib.parse.parse_qs(parsed.query)
model = params.get("model", [""])[0]
try:
shard_start = int(params.get("shard_start", [""])[0])
shard_end = int(params.get("shard_end", [""])[0])
except ValueError:
self.send_response(400)
self.end_headers()
return
if (
server.shard_path is None
or not server.shard_path.exists()
or model != server.model
or shard_start != server.shard_start
or shard_end != server.shard_end
):
self.send_response(404)
self.end_headers()
return
self.send_response(200)
self.send_header("Content-Type", "application/x-tar")
self.send_header("X-Meshnet-Shard-Checksum", compute_shard_checksum(server.shard_path))
self.end_headers()
write_shard_archive(server.shard_path, self.wfile)
def _handle_infer(self):
server: _StubHTTPServer = self.server # type: ignore[assignment]
length = int(self.headers.get("Content-Length", 0))
@@ -102,6 +147,8 @@ class StubNodeServer:
shard_end: int = 31,
is_last_shard: bool = True,
response_prefix: str = "stub response to:",
model: str = "stub-model",
shard_path: Path | None = None,
):
self._host = host
self._requested_port = port
@@ -111,6 +158,8 @@ class StubNodeServer:
self._shard_end = shard_end
self._is_last_shard = is_last_shard
self._response_prefix = response_prefix
self._model = model
self._shard_path = shard_path
self._server: _StubHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@@ -131,6 +180,8 @@ class StubNodeServer:
self._shard_end,
self._is_last_shard,
self._response_prefix,
self._model,
self._shard_path,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)

12
packages/node/meshnet_node/startup.py
View File

@@ -11,7 +11,7 @@ import urllib.request
from pathlib import Path
from typing import Any
from .downloader import download_shard
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware
from .server import StubNodeServer
from .wallet import load_or_create_wallet
@@ -47,7 +47,7 @@ def run_startup(
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
4. Download (or stub) the assigned shard from peers, then HuggingFace
5. Start local HTTP server
6. Register with tracker
@@ -96,6 +96,7 @@ def run_startup(
shard_end: int = assignment["shard_end"]
assigned_model: str = assignment.get("model", model)
hf_repo: str | None = assignment.get("hf_repo")
peers: list[dict] = assignment.get("peers", [])
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
# 4. Download shard
@@ -105,7 +106,10 @@ def run_startup(
dl_kwargs["cache_dir"] = cache_dir
if hf_repo is not None:
dl_kwargs["hf_repo"] = hf_repo
if peers:
dl_kwargs["peers"] = peers
shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
shard_checksum = compute_shard_checksum(shard_path)
print(f" Cached at: {shard_path}", flush=True)
# 5. Start HTTP server
@@ -116,6 +120,8 @@ def run_startup(
shard_start=shard_start,
shard_end=shard_end,
is_last_shard=is_last,
model=assigned_model,
shard_path=shard_path,
)
actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
@@ -128,8 +134,10 @@ def run_startup(
f"{tracker_url}/v1/nodes/register",
{
"endpoint": endpoint,
"model": assigned_model,
"shard_start": shard_start,
"shard_end": shard_end,
"shard_checksum": shard_checksum,
"hardware_profile": hw,
"wallet_address": address,
"score": 1.0,

36
packages/tracker/meshnet_tracker/server.py
View File

@@ -4,6 +4,7 @@ HTTP API contract:
- POST /v1/nodes/register
Request JSON: {
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
"model": str optional, "shard_checksum": str optional,
"hardware_profile": object, "wallet_address": str optional,
"score": number optional
}
@@ -34,7 +35,8 @@ DEFAULT_MODEL_PRESETS: dict[str, dict] = {
class _NodeEntry:
__slots__ = (
"node_id", "endpoint", "shard_start", "shard_end",
"hardware_profile", "wallet_address", "score", "last_heartbeat",
"model", "shard_checksum", "hardware_profile", "wallet_address",
"score", "last_heartbeat",
)
def __init__(
@@ -43,6 +45,8 @@ class _NodeEntry:
endpoint: str,
shard_start: int,
shard_end: int,
model: str | None,
shard_checksum: str | None,
hardware_profile: dict,
wallet_address: str | None,
score: float,
@@ -51,6 +55,8 @@ class _NodeEntry:
self.endpoint = endpoint
self.shard_start = shard_start
self.shard_end = shard_end
self.model = model
self.shard_checksum = shard_checksum
self.hardware_profile = hardware_profile
self.wallet_address = wallet_address
self.score = score
@@ -209,6 +215,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if not isinstance(hardware_profile, dict):
self._send_json(400, {"error": "hardware_profile must be an object"})
return
model = body.get("model")
if model is None:
model = "stub-model"
if not isinstance(model, str):
self._send_json(400, {"error": "model must be a string"})
return
shard_checksum = body.get("shard_checksum")
if shard_checksum is not None and not isinstance(shard_checksum, str):
self._send_json(400, {"error": "shard_checksum must be a string"})
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"})
@@ -224,6 +240,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
endpoint=endpoint.rstrip("/"),
shard_start=shard_start,
shard_end=shard_end,
model=model,
shard_checksum=shard_checksum,
hardware_profile=hardware_profile,
wallet_address=wallet_address,
score=score,
@@ -279,7 +297,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
self._purge_expired_nodes()
alive = list(server.registry.values())
alive = [node for node in server.registry.values() if node.model == model]
if server.contracts is not None:
alive = [
node for node in alive
@@ -317,11 +335,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
shard_start = gap_start
shard_end = min(required_end, shard_start + max_layers - 1)
peers = [
{"endpoint": node.endpoint, "checksum": node.shard_checksum}
for node in alive
if node.model == model
and node.shard_start == shard_start
and node.shard_end == shard_end
and node.shard_checksum
]
self._send_json(200, {
"shard_start": shard_start,
"shard_end": shard_end,
"model": model,
"model_layers_end": required_end,
"peers": peers,
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
})
@@ -344,7 +372,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
self._purge_expired_nodes()
alive = list(server.registry.values())
alive = [node for node in server.registry.values() if node.model == model]
if server.contracts is not None:
alive = [
node for node in alive
@@ -364,6 +392,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"wallet_address": e.wallet_address,
"shard_start": e.shard_start,
"shard_end": e.shard_end,
"model": e.model,
"shard_checksum": e.shard_checksum,
"score": e.score,
}
for e in route

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()