giant/giant/config.py

import random
import subprocess
import sys
import tomllib
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch

DEFAULT_CONFIG: dict = {
    "train": {
        "mode": "flow",
        "epochs": 100,
        "batch_size": 4096,
        "lr": 3e-4,
        "val_fraction": 0.1,
        "num_workers": 4,
        "seed": 0,
        "validate_every": 10,
        "validate_steps": 10,
        "warmup_epochs": 5,
    },
    "model": {
        "hidden_dim": 256,
        "n_blocks": 6,
        "emb_dim": 16,
        "dropout": 0.1,
    },
}


def git_hash() -> str:
    try:
        return (
            subprocess.check_output(
                ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
            )
            .decode()
            .strip()
        )
    except Exception:
        return "unknown"


def auto_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# Calibration point for estimate_batch_size(training=True): hidden_dim=512,
# n_blocks=6, batch_size=131072 measured at ~8 GiB VRAM. Activation memory is
# assumed to scale linearly with batch_size * hidden_dim * n_blocks (the
# ResBlock stack dominates), so this is a rough estimate rather than a
# guaranteed bound.
_REF_BYTES = 8 * 1024**3
_REF_BATCH_SIZE = 131072
_REF_HIDDEN_DIM = 512
_REF_N_BLOCKS = 6

# Calibration point for estimate_batch_size(training=False): inference has no
# backward graph or optimizer state, so its memory footprint is much smaller
# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
# MiB VRAM.
_REF_BYTES_PREDICT = 2037 * 1024**2
_REF_BATCH_SIZE_PREDICT = 65536
_REF_HIDDEN_DIM_PREDICT = 1024
_REF_N_BLOCKS_PREDICT = 8


def estimate_batch_size(
    hidden_dim: int,
    n_blocks: int,
    device: torch.device,
    safety_factor: float = 0.8,
    min_batch_size: int = 1024,
    training: bool = True,
) -> int:
    """Estimate a batch size that fits in the free memory on `device`.

    Only supported on CUDA devices, which expose a free/total memory query;
    other backends (cpu, mps) raise ValueError. Pass `training=False` for
    inference (e.g. `predict`), which uses a much lower per-sample memory
    calibration since there's no backward graph or optimizer state.
    """
    if device.type != "cuda":
        raise ValueError(
            f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
        )
    device_index = (
        device.index if device.index is not None else torch.cuda.current_device()
    )
    free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
    if training:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES,
            _REF_BATCH_SIZE,
            _REF_HIDDEN_DIM,
            _REF_N_BLOCKS,
        )
    else:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES_PREDICT,
            _REF_BATCH_SIZE_PREDICT,
            _REF_HIDDEN_DIM_PREDICT,
            _REF_N_BLOCKS_PREDICT,
        )
    bytes_per_unit = ref_bytes / (ref_batch_size * ref_hidden_dim * ref_n_blocks)
    bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
    batch_size = int(free_bytes * safety_factor / bytes_per_sample)
    batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
    return batch_size


def load_toml(path: Path) -> dict:
    with open(path, "rb") as f:
        return tomllib.load(f)


def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
    """Warn (don't fail) if a config.toml's [meta].git_hash predates the current checkout.

    A config saved by a previous run may have been produced by code that has
    since changed, so its hyperparameters might not mean what they used to —
    surface that as a heads-up rather than blocking the rerun.
    """
    file_hash = file_cfg.get("meta", {}).get("git_hash")
    current_hash = git_hash()
    if not file_hash or file_hash == "unknown" or current_hash == "unknown":
        return
    if file_hash != current_hash:
        print(
            f"warning: {config_path} was generated at git commit {file_hash}, "
            f"but the current checkout is at {current_hash} — hyperparameters "
            "may not match the code that originally produced this config",
            file=sys.stderr,
        )


def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
    """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.

    Training writes config.toml into the same out_dir as its checkpoints, so a
    checkpoint loaded later (for `predict` or `giant.analysis`) can be
    cross-checked the same way `--config` loading is, without the caller having
    to pass the toml path explicitly. Silently does nothing if no config.toml
    is found alongside the checkpoint.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return
    warn_if_git_hash_mismatch(load_toml(config_path), config_path)


def merge_cli_overrides(
    defaults: dict,
    config_path: Path | None,
    train_overrides: dict,
    model_overrides: dict,
) -> dict:
    """Resolve config as defaults -> TOML file -> explicit CLI flags."""
    cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])}
    if config_path is not None:
        file_cfg = load_toml(config_path)
        for section in ("train", "model"):
            cfg[section].update(file_cfg.get(section, {}))
        warn_if_git_hash_mismatch(file_cfg, config_path)
    cfg["train"].update(train_overrides)
    cfg["model"].update(model_overrides)
    return cfg


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
    lines = []
    for section, values in cfg.items():
        lines.append(f"[{section}]")
        for k, v in values.items():
            lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")
        lines.append("")

    lines.append("[meta]")
    for k, v in meta.items():
        lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")

    (out_dir / "config.toml").write_text("\n".join(lines))


def build_run_meta(
    data: Path,
    seed: int,
    n_pdg_codes: int,
    n_materials: int,
    n_train_events: int,
    n_val_events: int,
    n_train_steps: int,
) -> dict:
    return {
        "git_hash": git_hash(),
        "seed": seed,
        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
        "data_path": str(data),
        "n_pdg_codes": n_pdg_codes,
        "n_materials": n_materials,
        "n_train_events": n_train_events,
        "n_val_events": n_val_events,
        "n_train_steps": n_train_steps,
    }