source/src/bgc_retrieval/checkpoints.py

"""Checkpoint lineage and compatibility validation."""

from __future__ import annotations

from dataclasses import asdict
from pathlib import Path
from typing import Any

import torch

from .artifacts import environment_record, sha256_file, sha256_json
from .model import ModelConfig


def save_checkpoint(
    path: str | Path,
    model: torch.nn.Module,
    model_config: ModelConfig,
    split_path: str | Path,
    input_paths: list[str | Path],
    training_state: dict[str, Any],
    optimizer: torch.optim.Optimizer | None = None,
) -> None:
    output = Path(path)
    if output.exists():
        raise FileExistsError(f"Refusing to overwrite checkpoint: {output}")
    output.parent.mkdir(parents=True, exist_ok=True)
    input_hashes = {str(Path(item).resolve()): sha256_file(item) for item in input_paths}
    payload = {
        "schema_version": 1,
        "model_class": type(model).__name__,
        "model_config": asdict(model_config),
        "model_config_sha256": sha256_json(asdict(model_config)),
        "model_state": model.state_dict(),
        "optimizer_state": optimizer.state_dict() if optimizer else None,
        "split_path": str(Path(split_path).resolve()),
        "split_sha256": sha256_file(split_path),
        "input_sha256": input_hashes,
        "training_state": training_state,
        "environment": environment_record(),
    }
    torch.save(payload, output)


def load_checkpoint(
    path: str | Path,
    model: torch.nn.Module,
    split_path: str | Path,
    map_location: str | torch.device = "cpu",
) -> dict[str, Any]:
    payload = torch.load(path, map_location=map_location)
    if payload.get("schema_version") != 1:
        raise ValueError("Unsupported checkpoint schema")
    if payload["split_sha256"] != sha256_file(split_path):
        raise ValueError("Checkpoint was trained with a different split manifest")
    model.load_state_dict(payload["model_state"])
    return payload