feat: add p2p shard swarm

packages/node/meshnet_node/downloader.py

@@ -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,