dual billing; tracker to node model sharing

packages/tracker/meshnet_tracker/model_files.py

@@ -0,0 +1,172 @@
+"""Helpers for serving layer-scoped model files from tracker-local snapshots."""
+
+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 snapshot_dir_for_repo(models_dir: Path, repo_id: str) -> Path | None:
+    """Return the most likely local HF snapshot directory for *repo_id*."""
+    candidates = [
+        models_dir / repo_id,
+        models_dir / repo_id.replace("/", "--"),
+        models_dir / f"models--{repo_id.replace('/', '--')}",
+    ]
+    for candidate in candidates:
+        if (candidate / "snapshots").is_dir():
+            snapshots = sorted(p for p in (candidate / "snapshots").iterdir() if p.is_dir())
+            if snapshots:
+                return snapshots[-1]
+        if candidate.is_dir():
+            return candidate
+    return None
+
+
+def files_for_layer_range(snapshot_dir: Path, shard_start: int, shard_end: int) -> list[str]:
+    """Select files needed to load a conservative safetensors shard subset."""
+    return select_safetensors_files_for_layers(snapshot_dir, shard_start, shard_end)
+
+
+def select_safetensors_files_for_layers(
+    model_dir: str | Path,
+    start_layer: int,
+    end_layer: int,
+    *,
+    total_layers: int | None = None,
+) -> list[str]:
+    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:
+        return sorted(p.name for p in root.glob("*.safetensors") if p.is_file())
+
+    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(rel for rel in selected if (root / rel).is_file())
+
+
+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)
+    return int(match.group(1)) if match else None
+
+
+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()