dual billing; tracker to node model sharing

packages/node/meshnet_node/downloader.py

@@ -13,6 +13,7 @@ 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
 
@@ -105,6 +106,113 @@ def _download_shard_from_peer(
             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 download_shard(
     model: str,
     shard_start: int,
@@ -113,6 +221,7 @@ def download_shard(
     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.
@@ -157,18 +266,26 @@ def download_shard(
                 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,
         )
+    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]
+
+    if progress:
         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)
+    return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, None)