tracker download fix

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