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,
    },
    "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")


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,
    }