dual billing; tracker to node model sharing

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()