File size: 1,943 Bytes
c87881a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """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
|