neuron-tai/packages/node/meshnet_node/downloader.py

"""Shard downloader — fetches model shards from peers or HuggingFace Hub.

Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/

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(
    model: str,
    shard_start: int,
    shard_end: int,
    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.

    When *hf_repo* is None (or *model* is ``"stub-model"``), a placeholder
    weights file is created locally — no network access required.  This keeps
    the test suite hermetic while the real download path is exercised by
    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":
        stub_file = shard_dir / "weights.json"
        if not stub_file.exists():
            stub_file.write_text(json.dumps({
                "model": model,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "stub": True,
            }))
            if progress:
                print(f"  [stub] shard placeholder written to {stub_file}", flush=True)
        else:
            if progress:
                print(f"  [stub] shard already cached at {shard_dir}", flush=True)
        return shard_dir

    from huggingface_hub import snapshot_download  # type: ignore[import]

    if progress:
        print(
            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,
        cache_dir=str(cache_dir),
        local_dir=str(shard_dir),
    )
    return Path(local_dir)