giant/giant/data/setup_cache.py

"""Sidecar cache for `giant train`'s setup stage (vocab maps, event-id split
index, process maps, normalizer stats).

The setup stage scans the full training dataset before a single epoch runs
(see giant/pipeline.py:run_train_job); on multi-hundred-million-row datasets
that scan is itself expensive, and it's pure waste to repeat when the same
`data` path is reused across runs (hyperparameter sweeps via `dwarf
hparam-scan`, repeated manual training attempts, ...). This module persists
those scan outputs to a JSON file next to `data`, validated by a file
fingerprint + fixed dimension constants + a manually-bumped format version
before reuse — see `load`/`save`.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from giant import config
from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
from giant.data.loader import event_id_offset, load_event_ids
from giant.data.transforms import Normalizer, sorted_membership

# Bump manually on a change to the data-encoding semantics (e.g. a future
# energy_simplex_encode bugfix) that doesn't also move one of _DIMS below —
# a dims change already hard-invalidates on its own.
# v2: event_id is now offset per-file (see loader.event_id_offset) to avoid
# cross-file collisions, so a v1 sidecar's event_index/normalizers were
# computed against collided ids and must not be reused.
# v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
# energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
# v2 sidecar has no such grid to fall back on, so it must be recomputed.
_CACHE_FORMAT_VERSION = 3

_DIMS = {
    "COND_DIM": COND_DIM,
    "X_DIM": X_DIM,
    "K_MAX": K_MAX,
    "PARTICLE_PHYS_DIM": PARTICLE_PHYS_DIM,
    "SEC_SLOT_DIM": SEC_SLOT_DIM,
}

# Resolution of the stored energy-quantile summary (see NormalizerEntry).
# Only a handful of quantile *levels* (one per EnergyRouter expert) are ever
# consumed (see pipeline.py), so a dense fixed grid of quantile values is
# enough to reconstruct any level via interpolation (energy_quantile_at) —
# at roughly 1/100th the storage of the raw 100k-value reservoir sample it
# replaces, with negligible loss of resolution for that use.
ENERGY_QUANTILE_LEVELS = 1001


def energy_quantiles_from_sample(sample: np.ndarray) -> np.ndarray:
    """Collapse a raw reservoir sample into the fixed grid stored on disk."""
    if sample.size == 0:
        return np.empty(0, dtype=np.float32)
    levels = np.linspace(0.0, 1.0, ENERGY_QUANTILE_LEVELS)
    return np.quantile(sample, levels).astype(np.float32)


def energy_quantile_at(energy_quantiles: np.ndarray, levels: np.ndarray) -> np.ndarray:
    """Interpolate quantile values at arbitrary probability `levels` from the
    stored grid (e.g. `np.linspace(0, 1, n_experts)` for router centers)."""
    grid_levels = np.linspace(0.0, 1.0, len(energy_quantiles))
    return np.interp(levels, grid_levels, energy_quantiles).astype(np.float32)


def sidecar_path(data: str | Path) -> Path:
    """The cache sidecar for `data`, always a sibling of `data` itself.

    A directory `data` gets a sidecar *next to* it (not inside), since the
    directory may be a shared/read-only dataset mount, and other code globs
    `*.parquet` directly inside it.
    """
    p = Path(data)
    return p.parent / f"{p.name}.giant_train_cache.json"


def fingerprint_files(files: list[Path]) -> list[list]:
    """`[[resolved_path_str, size, mtime_ns], ...]`, in `files` order (not sorted).

    Order must be preserved rather than normalized (e.g. sorted): file scan
    order affects `build_process_map_from_files`'s tie-breaking (see
    tests/test_loader.py), so the cached fingerprint has to reflect the same
    order `find_parquet_files` produced.
    """
    out = []
    for f in files:
        resolved = Path(f).resolve()
        st = resolved.stat()
        out.append([str(resolved), st.st_size, st.st_mtime_ns])
    return out


def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
    # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
    # spurious cache misses between runs with the "same" val_fraction.
    return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"


@dataclass
class NormalizerEntry:
    cond_norm: Normalizer
    tgt_norm: Normalizer
    sec_phys_norm: Normalizer
    n_train_steps: int
    energy_quantiles: np.ndarray
    """Fixed ENERGY_QUANTILE_LEVELS-point quantile grid of the raw (pre-
    normalization) pre-step energy column — see energy_quantiles_from_sample
    / energy_quantile_at."""

    def to_json(self) -> dict:
        return {
            "cond_norm": self.cond_norm.to_dict(),
            "tgt_norm": self.tgt_norm.to_dict(),
            "sec_phys_norm": self.sec_phys_norm.to_dict(),
            "n_train_steps": self.n_train_steps,
            "energy_quantiles": np.asarray(
                self.energy_quantiles, dtype=np.float32
            ).tolist(),
        }

    @classmethod
    def from_json(cls, d: dict) -> "NormalizerEntry":
        return cls(
            cond_norm=Normalizer.from_dict(d["cond_norm"]),
            tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
            sec_phys_norm=Normalizer.from_dict(d["sec_phys_norm"]),
            n_train_steps=int(d["n_train_steps"]),
            energy_quantiles=np.array(d["energy_quantiles"], dtype=np.float32),
        )


@dataclass
class SetupCache:
    fingerprint: list
    git_hash: str = field(default_factory=config.git_hash)
    vocab: tuple[dict[int, int], dict[str, int]] | None = None
    event_index: tuple[np.ndarray, np.ndarray] | None = None
    proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
    normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)

    @classmethod
    def empty(cls, files: list[Path]) -> "SetupCache":
        return cls(fingerprint=fingerprint_files(files))

    def to_json(self) -> dict:
        d: dict = {
            "format_version": _CACHE_FORMAT_VERSION,
            "dims": dict(_DIMS),
            "git_hash": self.git_hash,
            "fingerprint": self.fingerprint,
            "proc_maps": {str(k): v for k, v in self.proc_maps.items()},
            "normalizers": {k: v.to_json() for k, v in self.normalizers.items()},
        }
        if self.vocab is not None:
            pdg_map, mat_map = self.vocab
            d["vocab"] = {
                "pdg_map": {str(k): v for k, v in pdg_map.items()},
                "mat_map": dict(mat_map),
            }
        if self.event_index is not None:
            unique_ids, counts = self.event_index
            d["event_index"] = {
                "event_ids": np.asarray(unique_ids).tolist(),
                "counts": np.asarray(counts).tolist(),
            }
        return d

    @classmethod
    def from_json(cls, d: dict) -> "SetupCache":
        vocab = None
        if "vocab" in d:
            pdg_map = {int(k): v for k, v in d["vocab"]["pdg_map"].items()}
            mat_map = dict(d["vocab"]["mat_map"])
            vocab = (pdg_map, mat_map)
        event_index = None
        if "event_index" in d:
            event_index = (
                np.array(d["event_index"]["event_ids"], dtype=np.int64),
                np.array(d["event_index"]["counts"], dtype=np.int64),
            )
        proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
        normalizers = {
            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
        }
        return cls(
            fingerprint=d["fingerprint"],
            git_hash=d.get("git_hash", "unknown"),
            vocab=vocab,
            event_index=event_index,
            proc_maps=proc_maps,
            normalizers=normalizers,
        )

    def merge(self, other: "SetupCache") -> "SetupCache":
        """Union of both caches; `other`'s populated fields win on a shared key.

        Used by `save` to combine freshly-computed sections with whatever a
        concurrent writer already persisted, so two runs against the same
        dataset with different (e.g.) val_fraction don't clobber each
        other's normalizer entries.
        """
        return SetupCache(
            fingerprint=other.fingerprint,
            git_hash=other.git_hash,
            vocab=other.vocab if other.vocab is not None else self.vocab,
            event_index=(
                other.event_index if other.event_index is not None else self.event_index
            ),
            proc_maps={**self.proc_maps, **other.proc_maps},
            normalizers={**self.normalizers, **other.normalizers},
        )


def load(
    data: str | Path, files: list[Path], echo=lambda *a, **k: None
) -> SetupCache | None:
    """Load and validate the sidecar for `data`; `None` on any miss (never raises).

    A missing file, corrupt JSON, format-version mismatch, dimension-constant
    mismatch, or file-fingerprint mismatch are all clean misses. A git-hash
    mismatch alone is a soft warning only (see
    `config.warn_if_git_hash_mismatch`) — most commits to this repo don't
    touch data-encoding semantics, so hard-invalidating on every one would
    defeat the cache.
    """
    path = sidecar_path(data)
    if not path.exists():
        return None
    try:
        raw = json.loads(path.read_text())
    except (json.JSONDecodeError, OSError) as exc:
        echo(f"setup cache: {path} is corrupt ({exc}) — ignoring")
        return None

    try:
        if raw.get("format_version") != _CACHE_FORMAT_VERSION:
            echo("setup cache: format version changed — ignoring stale cache")
            return None
        if raw.get("dims") != _DIMS:
            echo(
                "setup cache: model dimension constants changed — ignoring stale cache"
            )
            return None
        fp = fingerprint_files(files)
        if raw.get("fingerprint") != fp:
            echo("setup cache: input files changed — ignoring stale cache")
            return None
        cache = SetupCache.from_json(raw)
    except (KeyError, TypeError, ValueError) as exc:
        echo(f"setup cache: {path} is malformed ({exc}) — ignoring")
        return None

    config.warn_if_git_hash_mismatch({"meta": {"git_hash": cache.git_hash}}, path)
    return cache


def save(
    data: str | Path,
    files: list[Path],
    sections: SetupCache,
    echo=lambda *a, **k: None,
) -> None:
    """Merge `sections` into the on-disk sidecar and write it atomically.

    Best-effort: any OSError (permission denied on a read-only mount, disk
    full, ...) is caught, echoed as a warning, and swallowed — a failure to
    cache must never fail training.
    """
    path = sidecar_path(data)
    tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
    try:
        base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
        merged = base.merge(sections)
        payload = json.dumps(merged.to_json(), separators=(",", ":"))
        tmp.write_text(payload)
        os.replace(tmp, path)
    except OSError as exc:
        echo(
            f"setup cache: could not write {path} ({exc}) — continuing without caching"
        )
        try:
            tmp.unlink(missing_ok=True)
        except OSError:
            pass


def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
    """Unique event ids + per-event row (step) counts, across all `files`."""
    if not files:
        return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
    all_ids = np.concatenate(
        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
    )
    unique_ids, counts = np.unique(all_ids, return_counts=True)
    return unique_ids, counts


def n_train_steps_for_split(
    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
) -> int:
    """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.

    `train_events_arr` must be ascending and duplicate-free (as produced by
    `np.array(sorted(train_events))` in giant/pipeline.py).
    """
    mask = sorted_membership(unique_ids, train_events_arr)
    return int(counts[mask].sum())