342 lines
12 KiB
Python
342 lines
12 KiB
Python
"""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 concurrent.futures import ThreadPoolExecutor, as_completed
|
|
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_model_source(
|
|
source: dict,
|
|
shard_dir: Path,
|
|
timeout: float,
|
|
) -> Path | None:
|
|
url = source.get("url")
|
|
if not isinstance(url, str) or not url:
|
|
endpoint = source.get("endpoint")
|
|
if not isinstance(endpoint, str):
|
|
return None
|
|
url = f"{endpoint.rstrip('/')}/v1/model-files/download"
|
|
shard_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="meshnet-model-source-", dir=shard_dir.parent) as tmp:
|
|
tmp_root = Path(tmp)
|
|
archive_path = tmp_root / "model-files.tar"
|
|
extract_dir = tmp_root / "extract"
|
|
extract_dir.mkdir()
|
|
try:
|
|
with urllib.request.urlopen(url, 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)
|
|
shutil.move(str(extract_dir), str(shard_dir))
|
|
return shard_dir
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _download_huggingface_subset(
|
|
hf_repo: str,
|
|
cache_dir: Path,
|
|
shard_dir: Path,
|
|
allow_patterns: list[str] | None,
|
|
) -> Path:
|
|
from huggingface_hub import snapshot_download # type: ignore[import]
|
|
|
|
kwargs = {
|
|
"repo_id": hf_repo,
|
|
"cache_dir": str(cache_dir),
|
|
"local_dir": str(shard_dir),
|
|
}
|
|
if allow_patterns:
|
|
kwargs["allow_patterns"] = allow_patterns
|
|
try:
|
|
return Path(snapshot_download(**kwargs))
|
|
except TypeError:
|
|
kwargs.pop("allow_patterns", None)
|
|
return Path(snapshot_download(**kwargs))
|
|
|
|
|
|
def _download_from_fastest_source(
|
|
*,
|
|
model_sources: list[dict],
|
|
hf_repo: str,
|
|
cache_dir: Path,
|
|
shard_dir: Path,
|
|
progress: bool,
|
|
timeout: float,
|
|
) -> tuple[str, Path] | None:
|
|
shard_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="meshnet-race-", dir=shard_dir.parent) as tmp:
|
|
tmp_root = Path(tmp)
|
|
jobs: dict[Any, tuple[str, Path]] = {}
|
|
pool = ThreadPoolExecutor(max_workers=min(4, len(model_sources) + 1))
|
|
try:
|
|
for index, source in enumerate(model_sources):
|
|
label = str(source.get("type") or "model-source")
|
|
candidate = tmp_root / f"source-{index}"
|
|
jobs[pool.submit(_download_model_source, source, candidate, timeout)] = (label, candidate)
|
|
allow_patterns = _allow_patterns_from_sources(model_sources)
|
|
hf_candidate = tmp_root / "huggingface"
|
|
jobs[pool.submit(_download_huggingface_subset, hf_repo, cache_dir, hf_candidate, allow_patterns)] = (
|
|
"HuggingFace",
|
|
hf_candidate,
|
|
)
|
|
for future in as_completed(jobs):
|
|
label, candidate = jobs[future]
|
|
try:
|
|
result = future.result()
|
|
except Exception:
|
|
continue
|
|
if result is None:
|
|
continue
|
|
if shard_dir.exists():
|
|
shutil.rmtree(shard_dir)
|
|
shutil.move(str(candidate), str(shard_dir))
|
|
if progress:
|
|
print(f" download source: {label}", flush=True)
|
|
pool.shutdown(wait=False, cancel_futures=True)
|
|
return label, shard_dir
|
|
finally:
|
|
pool.shutdown(wait=False, cancel_futures=True)
|
|
return None
|
|
|
|
|
|
def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
|
|
patterns: set[str] = set()
|
|
for source in model_sources:
|
|
for rel in source.get("files") or []:
|
|
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
|
|
patterns.add(rel)
|
|
return sorted(patterns) if patterns else None
|
|
|
|
|
|
def _allow_patterns_from_remote_index(
|
|
hf_repo: str,
|
|
cache_dir: Path,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
) -> list[str] | None:
|
|
"""Fetch just the SafeTensors index + config (a few KB) from HF and compute
|
|
which weight files the assigned layer range needs, so a HuggingFace fallback
|
|
download stays layer-scoped even when the tracker has no model_sources
|
|
(e.g. it has no local snapshot for this repo cached yet)."""
|
|
try:
|
|
from huggingface_hub import hf_hub_download # type: ignore[import]
|
|
|
|
from .safetensors_selection import (
|
|
INDEX_FILENAME,
|
|
METADATA_FILENAMES,
|
|
layers_from_config_dict,
|
|
select_files_for_layers_from_index,
|
|
)
|
|
|
|
index_path = hf_hub_download(repo_id=hf_repo, filename=INDEX_FILENAME, cache_dir=str(cache_dir))
|
|
weight_map = json.loads(Path(index_path).read_text(encoding="utf-8")).get("weight_map")
|
|
except Exception:
|
|
return None
|
|
if not isinstance(weight_map, dict):
|
|
return None
|
|
|
|
total_layers: int | None = None
|
|
try:
|
|
config_path = hf_hub_download(repo_id=hf_repo, filename="config.json", cache_dir=str(cache_dir))
|
|
config = json.loads(Path(config_path).read_text(encoding="utf-8"))
|
|
total_layers = layers_from_config_dict(config)
|
|
except Exception:
|
|
pass
|
|
|
|
selected = select_files_for_layers_from_index(
|
|
weight_map, shard_start, shard_end, total_layers=total_layers
|
|
)
|
|
return sorted(selected | METADATA_FILENAMES)
|
|
|
|
|
|
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,
|
|
model_sources: 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}"
|
|
if progress:
|
|
print(f" Target location: {shard_dir}", flush=True)
|
|
|
|
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
|
|
|
|
if progress:
|
|
print(
|
|
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
|
|
flush=True,
|
|
)
|
|
if model_sources:
|
|
if progress:
|
|
print(" Racing tracker model source against HuggingFace ...", flush=True)
|
|
raced = _download_from_fastest_source(
|
|
model_sources=model_sources,
|
|
hf_repo=hf_repo,
|
|
cache_dir=cache_dir,
|
|
shard_dir=shard_dir,
|
|
progress=progress,
|
|
timeout=peer_timeout,
|
|
)
|
|
if raced is not None:
|
|
return raced[1]
|
|
|
|
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
|
|
if progress:
|
|
if allow_patterns:
|
|
print(" download source: HuggingFace (layer-filtered)", flush=True)
|
|
else:
|
|
print(
|
|
" download source: HuggingFace (full snapshot — no SafeTensors index found)",
|
|
flush=True,
|
|
)
|
|
|
|
return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, allow_patterns)
|