neuron-tai/packages/tracker/meshnet_tracker/model_files.py

"""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):
            rel = _normalise_relative_file(rel_file)
            if not (root / rel).is_file():
                # A required weight file is missing from this local snapshot: the
                # snapshot is incomplete for this layer range, so it must not be
                # advertised as a source at all. Silently dropping just this file
                # (as before) made a partial snapshot look "fully cached" to
                # downstream clients, which then skipped the real download and
                # only discovered the missing weights at model-load time.
                return []
            selected.add(rel)

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