tracker download fix

packages/node/meshnet_node/downloader.py

@@ -213,6 +213,47 @@ def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
     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,
@@ -285,7 +326,14 @@ def download_shard(
         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:
-        print("  download source: HuggingFace", flush=True)
+        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, None)
+    return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, allow_patterns)

packages/node/meshnet_node/safetensors_selection.py

@@ -15,7 +15,7 @@ _LAYER_RE = re.compile(
     r"\.(\d+)(?:\.|$)"
 )
 
-_METADATA_FILENAMES = {
+METADATA_FILENAMES = {
     INDEX_FILENAME,
     "config.json",
     "generation_config.json",
@@ -90,14 +90,32 @@ def select_safetensors_files_for_layers(
 
     inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root)
     selected = _metadata_files(root)
+    selected |= select_files_for_layers_from_index(
+        weight_map, start_layer, end_layer, total_layers=inferred_total_layers
+    )
+    return sorted(selected)
 
+
+def select_files_for_layers_from_index(
+    weight_map: dict[str, str],
+    start_layer: int,
+    end_layer: int,
+    *,
+    total_layers: int | None = None,
+) -> set[str]:
+    """Pure variant of the weight-file selection: takes an already-parsed
+    ``weight_map`` (no local snapshot directory needed), so callers that only
+    have the index fetched over the network — not a full local snapshot — can
+    still compute which shard files they need. Combine the result with
+    ``METADATA_FILENAMES`` for a complete download pattern set.
+    """
+    selected: set[str] = set()
     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):
+        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
             selected.add(_normalise_relative_file(rel_file))
-
-    return sorted(selected)
+    return selected
 
 
 def _tensor_belongs_to_range(
@@ -142,7 +160,7 @@ def _metadata_files(root: Path) -> set[str]:
         if not path.is_file():
             continue
         name = path.name
-        if name in _METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
+        if name in METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
             files.add(name)
     return files
 
@@ -152,10 +170,10 @@ def _read_total_layers(root: Path) -> int | None:
     if not config_path.exists():
         return None
     config = json.loads(config_path.read_text(encoding="utf-8"))
-    return _layers_from_config(config)
+    return layers_from_config_dict(config)
 
 
-def _layers_from_config(config: dict[str, Any]) -> int | None:
+def layers_from_config_dict(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:
@@ -163,7 +181,7 @@ def _layers_from_config(config: dict[str, Any]) -> int | None:
 
     text_config = config.get("text_config")
     if isinstance(text_config, dict):
-        return _layers_from_config(text_config)
+        return layers_from_config_dict(text_config)
 
     return None