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

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,