545 lines
20 KiB
Python
545 lines
20 KiB
Python
"""Shard downloader — fetches model shards from peers or HuggingFace Hub.
|
|
|
|
Cache layout: ~/.cache/meshnet/shards/<model>/
|
|
|
|
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 os
|
|
import shutil
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
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
|
|
# Model-source tar streams are multi-GB; a short socket timeout must not kill
|
|
# them on a transient read stall. Peer probes keep the short timeout because
|
|
# they run sequentially before the race and may hit dead endpoints.
|
|
_MODEL_SOURCE_TIMEOUT_SECONDS = 30.0
|
|
_PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024
|
|
_FILE_RETRY_ATTEMPTS = 3
|
|
|
|
|
|
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."""
|
|
# dereference: HF cache snapshots are symlinks into blobs/ — ship contents.
|
|
with tarfile.open(fileobj=out_file, mode="w|", dereference=True) 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
|
|
_merge_tree(extract_dir, shard_dir)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
class _TarOnlySource(Exception):
|
|
"""The server ignored ?file= and streamed a tar — no single-file support."""
|
|
|
|
|
|
def _fetch_source_file(
|
|
file_url: str,
|
|
dest: Path,
|
|
timeout: float,
|
|
on_chunk=None,
|
|
) -> int:
|
|
"""Download one file; skip if *dest* already matches the remote size.
|
|
|
|
Returns the file's final size. Raises on any failure, including a
|
|
short read against the server's Content-Length.
|
|
"""
|
|
with urllib.request.urlopen(file_url, timeout=timeout) as resp:
|
|
ctype = resp.getheader("Content-Type") if hasattr(resp, "getheader") else None
|
|
if isinstance(ctype, str) and "x-tar" in ctype:
|
|
raise _TarOnlySource(file_url)
|
|
length = resp.getheader("Content-Length") if hasattr(resp, "getheader") else None
|
|
expected = int(length) if isinstance(length, str) and length.isdigit() else None
|
|
if expected is not None and dest.exists() and dest.stat().st_size == expected:
|
|
if on_chunk is not None:
|
|
on_chunk(expected)
|
|
return expected # complete from an earlier attempt/run
|
|
received = 0
|
|
with dest.open("wb") as out:
|
|
while True:
|
|
chunk = resp.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
received += len(chunk)
|
|
if on_chunk is not None:
|
|
on_chunk(len(chunk))
|
|
if expected is not None and received != expected:
|
|
raise OSError(f"short read: got {received} of {expected} bytes")
|
|
return received
|
|
|
|
|
|
class _SourceProgress:
|
|
"""tqdm bar over the whole per-file download (total / speed / ETA),
|
|
falling back to plain per-file prints when tqdm is unavailable."""
|
|
|
|
def __init__(self, enabled: bool, label: str, total_bytes: int | None):
|
|
self._label = label
|
|
self._bar = None
|
|
self._enabled = enabled
|
|
if enabled and total_bytes:
|
|
try:
|
|
from tqdm import tqdm # type: ignore[import]
|
|
|
|
self._bar = tqdm(
|
|
total=total_bytes,
|
|
unit="B",
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
desc=f"Downloading ({label})",
|
|
)
|
|
except Exception:
|
|
self._bar = None
|
|
|
|
def add_bytes(self, n: int) -> None:
|
|
if self._bar is not None:
|
|
self._bar.update(n)
|
|
|
|
def rewind(self, n: int) -> None:
|
|
"""Roll the bar back after a failed attempt so retries don't double-count."""
|
|
if self._bar is not None and n:
|
|
self._bar.update(-n)
|
|
|
|
def file_done(self, index: int, count: int, rel: str, size: int, reused: bool) -> None:
|
|
note = "already complete, " if reused else ""
|
|
line = f" {self._label}: [{index}/{count}] {rel} ({note}{size / 1e9:.2f} GB)"
|
|
if self._bar is not None:
|
|
self._bar.set_postfix_str(f"{index}/{count} files", refresh=False)
|
|
elif self._enabled:
|
|
print(line, flush=True)
|
|
|
|
def message(self, text: str) -> None:
|
|
if self._bar is not None:
|
|
self._bar.write(text)
|
|
elif self._enabled:
|
|
print(text, flush=True)
|
|
|
|
def close(self) -> None:
|
|
if self._bar is not None:
|
|
self._bar.close()
|
|
|
|
|
|
def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
|
|
"""Reuse an already-complete copy of a file instead of re-downloading."""
|
|
if dest.exists() and dest.stat().st_size == expected:
|
|
return True
|
|
if final.exists() and final.stat().st_size == expected:
|
|
try:
|
|
os.link(final, dest)
|
|
except OSError:
|
|
shutil.copy2(final, dest)
|
|
return True
|
|
return False
|
|
|
|
|
|
def _valid_source_rel_files(source: dict) -> list[str]:
|
|
return [
|
|
rel for rel in (source.get("files") or [])
|
|
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
|
|
]
|
|
|
|
|
|
def _source_files_cached(source: dict, shard_dir: Path) -> bool:
|
|
rel_files = _valid_source_rel_files(source)
|
|
if not rel_files:
|
|
return False
|
|
sizes = source.get("file_sizes")
|
|
if not isinstance(sizes, dict):
|
|
return False
|
|
for rel in rel_files:
|
|
expected = sizes.get(rel)
|
|
path = shard_dir / rel
|
|
if not isinstance(expected, int) or not path.exists() or path.stat().st_size != expected:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _merge_tree(src: Path, dest: Path) -> None:
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
for path in sorted(p for p in src.rglob("*") if p.is_file()):
|
|
target = dest / path.relative_to(src)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
if target.exists():
|
|
target.unlink()
|
|
shutil.move(str(path), str(target))
|
|
|
|
|
|
def _download_source_files(
|
|
source: dict,
|
|
shard_dir: Path,
|
|
timeout: float,
|
|
progress: bool,
|
|
label: str,
|
|
) -> Path | None:
|
|
"""Per-file download from a model source — retries and resumes per file.
|
|
|
|
Far more robust than one multi-GB tar stream on flaky links: a dropped
|
|
connection costs at most one file, and completed files are skipped when
|
|
the download is retried or the node is restarted.
|
|
"""
|
|
url = source.get("url")
|
|
rel_files = _valid_source_rel_files(source)
|
|
if not isinstance(url, str) or not url or not rel_files:
|
|
return None
|
|
sizes = source.get("file_sizes")
|
|
if not isinstance(sizes, dict):
|
|
sizes = {}
|
|
partial_dir = shard_dir.parent / f"{shard_dir.name}.partial"
|
|
partial_dir.mkdir(parents=True, exist_ok=True)
|
|
sep = "&" if "?" in url else "?"
|
|
known_sizes = [sizes.get(rel) for rel in rel_files]
|
|
total_bytes = sum(s for s in known_sizes if isinstance(s, int)) if all(
|
|
isinstance(s, int) for s in known_sizes
|
|
) else None
|
|
tracker_bar = _SourceProgress(progress, label, total_bytes)
|
|
try:
|
|
for index, rel in enumerate(rel_files, start=1):
|
|
dest = partial_dir / rel
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
known_size = sizes.get(rel)
|
|
if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel):
|
|
tracker_bar.add_bytes(known_size)
|
|
tracker_bar.file_done(index, len(rel_files), rel, known_size, reused=True)
|
|
continue
|
|
file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}"
|
|
for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1):
|
|
counted = 0
|
|
|
|
def _on_chunk(n: int) -> None:
|
|
nonlocal counted
|
|
counted += n
|
|
tracker_bar.add_bytes(n)
|
|
|
|
try:
|
|
size = _fetch_source_file(file_url, dest, timeout, on_chunk=_on_chunk)
|
|
tracker_bar.file_done(index, len(rel_files), rel, size, reused=False)
|
|
break
|
|
except _TarOnlySource:
|
|
tracker_bar.message(f" {label}: no single-file support — using tar stream")
|
|
return None
|
|
except Exception as exc:
|
|
tracker_bar.rewind(counted)
|
|
tracker_bar.message(
|
|
f" {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}"
|
|
)
|
|
if attempt == _FILE_RETRY_ATTEMPTS:
|
|
return None
|
|
time.sleep(1.0 * attempt)
|
|
finally:
|
|
tracker_bar.close()
|
|
_merge_tree(partial_dir, shard_dir)
|
|
try:
|
|
partial_dir.rmdir()
|
|
except OSError:
|
|
pass
|
|
return shard_dir
|
|
|
|
|
|
def _download_model_source(
|
|
source: dict,
|
|
shard_dir: Path,
|
|
timeout: float,
|
|
progress: bool = False,
|
|
label: str = "model-source",
|
|
) -> Path | None:
|
|
# Prefer per-file transfers whenever the source advertises its file list;
|
|
# fall through to the tar stream if the server lacks single-file support
|
|
# or per-file transfers keep failing.
|
|
if source.get("files"):
|
|
fetched = _download_source_files(source, shard_dir, timeout, progress, label)
|
|
if fetched is not None:
|
|
return fetched
|
|
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:
|
|
received = 0
|
|
next_report = _PROGRESS_INTERVAL_BYTES
|
|
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)
|
|
received += len(chunk)
|
|
if progress and received >= next_report:
|
|
print(f" {label}: {received / 1e9:.1f} GB received ...", flush=True)
|
|
next_report += _PROGRESS_INTERVAL_BYTES
|
|
if progress:
|
|
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
|
|
_safe_extract_shard(archive_path, extract_dir)
|
|
_merge_tree(extract_dir, shard_dir)
|
|
return shard_dir
|
|
except Exception as exc:
|
|
if progress:
|
|
print(f" {label}: download failed ({url}): {exc!r}", flush=True)
|
|
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 _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
|
|
if progress:
|
|
print(f" Target location: {shard_dir}", flush=True)
|
|
|
|
for source in model_sources or []:
|
|
label = str(source.get("type") or "model-source")
|
|
if _source_files_cached(source, shard_dir):
|
|
if progress:
|
|
print(f" [{label}] requested files already cached at {shard_dir}", flush=True)
|
|
return shard_dir
|
|
|
|
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,
|
|
)
|
|
# Tracker (or peer) model sources are preferred outright — usually LAN-fast.
|
|
# HuggingFace is only the fallback when every advertised source fails.
|
|
for source in model_sources or []:
|
|
label = str(source.get("type") or "model-source")
|
|
if progress:
|
|
print(f" Downloading from {label} model source (HuggingFace is the fallback) ...", flush=True)
|
|
fetched = _download_model_source(
|
|
source,
|
|
shard_dir,
|
|
timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
|
|
progress=progress,
|
|
label=label,
|
|
)
|
|
if fetched is not None:
|
|
if progress:
|
|
print(f" download source: {label}", flush=True)
|
|
return fetched
|
|
if model_sources and progress:
|
|
print(" All model sources failed — falling back to HuggingFace ...", flush=True)
|
|
|
|
allow_patterns = None
|
|
if model_sources:
|
|
allow_patterns = _allow_patterns_from_sources(model_sources)
|
|
if allow_patterns is None:
|
|
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)
|