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)

packages/node/meshnet_node/model_backend.py

@@ -85,24 +85,25 @@ class TorchModelShard:
 
         self.torch = torch
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
         quant_config, dtype, uses_quantized_weights = _model_load_plan(
             AutoConfig,
-            model_id,
+            load_source,
             quantization,
             torch,
-            cache_dir,
+            None if load_source != model_id else cache_dir,
         )
         try:
             load_kwargs = {
                 "device_map": "auto" if uses_quantized_weights else None,
                 "dtype": dtype,
                 "low_cpu_mem_usage": True,
-                "cache_dir": str(cache_dir) if cache_dir is not None else None,
+                "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
             }
             if quant_config is not None:
                 load_kwargs["quantization_config"] = quant_config
             self.model = AutoModelForCausalLM.from_pretrained(
-                model_id,
+                load_source,
                 **load_kwargs,
             )
             if not uses_quantized_weights:
@@ -117,8 +118,8 @@ class TorchModelShard:
 
         self.model.eval()
         self.tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            cache_dir=str(cache_dir) if cache_dir is not None else None,
+            load_source,
+            cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None,
         )
         self.layers = _model_layers(self.model)
         self.total_layers = len(self.layers)

packages/node/meshnet_node/safetensors_selection.py

@@ -0,0 +1,175 @@
+"""Layer-aware SafeTensors snapshot file selection."""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+INDEX_FILENAME = "model.safetensors.index.json"
+
+_LAYER_RE = re.compile(
+    r"(?:^|\.)"
+    r"(?:model\.layers|layers|h|blocks|decoder\.layers|encoder\.layers)"
+    r"\.(\d+)(?:\.|$)"
+)
+
+_METADATA_FILENAMES = {
+    INDEX_FILENAME,
+    "config.json",
+    "generation_config.json",
+    "preprocessor_config.json",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer.model",
+    "tokenizer_config.json",
+    "vocab.json",
+    "merges.txt",
+    "added_tokens.json",
+}
+
+_METADATA_PREFIXES = ("config.", "tokenizer.", "tokenizer_", "vocab.")
+
+_HEAD_MARKERS = (
+    "embed",
+    "embedding",
+    "embed_tokens",
+    "wte",
+    "wpe",
+)
+
+_TAIL_EXACT = {
+    "lm_head.weight",
+    "lm_head.bias",
+    "model.norm.weight",
+    "model.norm.bias",
+    "transformer.ln_f.weight",
+    "transformer.ln_f.bias",
+    "decoder.final_layer_norm.weight",
+    "decoder.final_layer_norm.bias",
+}
+
+_TAIL_MARKERS = (
+    ".lm_head.",
+    ".norm.",
+    ".ln_f.",
+    ".final_layer_norm.",
+)
+
+
+def select_safetensors_files_for_layers(
+    model_dir: str | Path,
+    start_layer: int,
+    end_layer: int,
+    *,
+    total_layers: int | None = None,
+) -> list[str]:
+    """Return relative snapshot files needed for an inclusive layer range.
+
+    The returned list always includes root-level config/tokenizer metadata and
+    the SafeTensors index. Weight shard files are included only when at least one
+    tensor in the index belongs to the assigned layer range, or when the tensor
+    is needed by the head/tail shard.
+    """
+    if start_layer < 0:
+        raise ValueError("start_layer must be non-negative")
+    if end_layer < start_layer:
+        raise ValueError("end_layer must be greater than or equal to start_layer")
+
+    root = Path(model_dir)
+    index_path = root / INDEX_FILENAME
+    try:
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+    except FileNotFoundError as exc:
+        raise FileNotFoundError(f"missing SafeTensors index: {index_path}") from exc
+
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict):
+        raise ValueError(f"{INDEX_FILENAME} must contain a weight_map object")
+
+    inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root)
+    selected = _metadata_files(root)
+
+    for tensor_name, rel_file in weight_map.items():
+        if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
+            continue
+        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, inferred_total_layers):
+            selected.add(_normalise_relative_file(rel_file))
+
+    return sorted(selected)
+
+
+def _tensor_belongs_to_range(
+    tensor_name: str,
+    start_layer: int,
+    end_layer: int,
+    total_layers: int | None,
+) -> bool:
+    layer = _layer_index(tensor_name)
+    if layer is not None:
+        return start_layer <= layer <= end_layer
+
+    if start_layer == 0 and _is_head_tensor(tensor_name):
+        return True
+
+    if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(tensor_name):
+        return True
+
+    return False
+
+
+def _layer_index(tensor_name: str) -> int | None:
+    match = _LAYER_RE.search(tensor_name)
+    if match is None:
+        return None
+    return int(match.group(1))
+
+
+def _is_head_tensor(tensor_name: str) -> bool:
+    lowered = tensor_name.lower()
+    return any(marker in lowered for marker in _HEAD_MARKERS)
+
+
+def _is_tail_tensor(tensor_name: str) -> bool:
+    lowered = tensor_name.lower()
+    return lowered in _TAIL_EXACT or any(marker in lowered for marker in _TAIL_MARKERS)
+
+
+def _metadata_files(root: Path) -> set[str]:
+    files = {INDEX_FILENAME}
+    for path in root.iterdir():
+        if not path.is_file():
+            continue
+        name = path.name
+        if name in _METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
+            files.add(name)
+    return files
+
+
+def _read_total_layers(root: Path) -> int | None:
+    config_path = root / "config.json"
+    if not config_path.exists():
+        return None
+    config = json.loads(config_path.read_text(encoding="utf-8"))
+    return _layers_from_config(config)
+
+
+def _layers_from_config(config: dict[str, Any]) -> int | None:
+    for key in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"):
+        value = config.get(key)
+        if isinstance(value, int) and value > 0:
+            return value
+
+    text_config = config.get("text_config")
+    if isinstance(text_config, dict):
+        return _layers_from_config(text_config)
+
+    return None
+
+
+def _normalise_relative_file(rel_file: str) -> str:
+    path = Path(rel_file)
+    if path.is_absolute() or ".." in path.parts:
+        raise ValueError(f"unsafe relative file in {INDEX_FILENAME}: {rel_file}")
+    return path.as_posix()

packages/node/meshnet_node/startup.py

@@ -34,6 +34,21 @@ def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int =
     return max(0, ram_mb), "RAM"
 
 
+def _full_model_sources(model_sources: list[dict]) -> list[dict]:
+    """Use tracker full-snapshot URLs for real HF model loading."""
+    full_sources: list[dict] = []
+    for source in model_sources:
+        full_url = source.get("full_url")
+        if isinstance(full_url, str) and full_url:
+            full_sources.append({
+                **source,
+                "url": full_url,
+                "files": [],
+                "type": f"{source.get('type') or 'model-source'}-full",
+            })
+    return full_sources
+
+
 def _hardware_label(device: str, gpu_name: str | None = None) -> str:
     if device == "cuda":
         return "CUDA"
@@ -443,6 +458,16 @@ def run_startup(
                     if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
                         shard_start = net_asgn["shard_start"]
                         shard_end = net_asgn["shard_end"]
+                        full_sources = _full_model_sources(net_asgn.get("model_sources", []))
+                        if full_sources:
+                            cache_dir = download_shard(
+                                model_id.split("/")[-1],
+                                shard_start,
+                                shard_end,
+                                cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
+                                hf_repo=model_id,
+                                model_sources=full_sources,
+                            )
                         print(
                             f"  Tracker found uncovered shard: "
                             f"layers {shard_start}–{shard_end} (of {detected})",
@@ -550,12 +575,24 @@ def run_startup(
         assigned_shard_start: int = net_assignment["shard_start"]
         assigned_shard_end: int = net_assignment["shard_end"]
         assigned_num_layers: int = net_assignment["num_layers"]
+        assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
         print(
             f"  Assigned: {assigned_hf_repo} "
             f"layers {assigned_shard_start}–{assigned_shard_end} "
             f"(of {assigned_num_layers})",
             flush=True,
         )
+        full_sources = _full_model_sources(assigned_model_sources)
+        if full_sources:
+            print("Downloading assigned model snapshot...", flush=True)
+            cache_dir = download_shard(
+                assigned_hf_repo.split("/")[-1],
+                assigned_shard_start,
+                assigned_shard_end,
+                cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
+                hf_repo=assigned_hf_repo,
+                model_sources=full_sources,
+            )
         print("Loading real PyTorch model shard...", flush=True)
         node = TorchNodeServer(
             host=host,
@@ -647,6 +684,7 @@ def run_startup(
     assigned_model: str = assignment.get("model", model)
     hf_repo: str | None = assignment.get("hf_repo")
     peers: list[dict] = assignment.get("peers", [])
+    model_sources: list[dict] = assignment.get("model_sources", [])
     print(f"  Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
 
     # 4. Download shard
@@ -658,6 +696,8 @@ def run_startup(
         dl_kwargs["hf_repo"] = hf_repo
     if peers:
         dl_kwargs["peers"] = peers
+    if model_sources:
+        dl_kwargs["model_sources"] = model_sources
     shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
     shard_checksum = compute_shard_checksum(shard_path)
     print(f"  Cached at: {shard_path}", flush=True)