@dataclass
class ModelSelectionCheckpoint:
"""On-disk checkpoint for resumable model selection runs."""
version: int = CHECKPOINT_VERSION
run_fingerprint: str = ""
created_at: str = ""
updated_at: str = ""
completed_combinations: list[str] = field(default_factory=list)
summaries_by_key: dict[str, dict[str, str | float | None]] = field(
default_factory=dict
)
singleton_scores_by_key: dict[str, float] = field(default_factory=dict)
stages: dict[str, StageState] = field(default_factory=dict)
failures: list[FailureRecord] = field(default_factory=list)
checkpoint_fusion_pairs: int | None = None
checkpoint_fusion_triples: int | None = None
def to_json_dict(self) -> dict[str, Any]:
return {
"version": self.version,
"run_fingerprint": self.run_fingerprint,
"created_at": self.created_at,
"updated_at": self.updated_at,
"checkpoint_fusion_pairs": self.checkpoint_fusion_pairs,
"checkpoint_fusion_triples": self.checkpoint_fusion_triples,
"completed_combinations": sorted(self.completed_combinations),
"summaries_by_key": {
k: dict(v) for k, v in sorted(self.summaries_by_key.items())
},
"singleton_scores_by_key": {
k: float(v) for k, v in sorted(self.singleton_scores_by_key.items())
},
"stages": dict(sorted(self.stages.items())),
"failures": [
{"combination_key": f.combination_key, "error": f.error, "at": f.at}
for f in self.failures
],
}
@staticmethod
def from_json_dict(data: dict[str, Any]) -> ModelSelectionCheckpoint:
if int(data.get("version", -1)) != CHECKPOINT_VERSION:
raise ValueError(
f"Unsupported checkpoint version: {data.get('version')!r}; "
f"expected {CHECKPOINT_VERSION}"
)
failures_raw = data.get("failures") or []
failures = [
FailureRecord(
combination_key=str(fr["combination_key"]),
error=str(fr["error"]),
at=str(fr["at"]),
)
for fr in failures_raw
]
cfp = data.get("checkpoint_fusion_pairs")
cft = data.get("checkpoint_fusion_triples")
return ModelSelectionCheckpoint(
version=int(data["version"]),
run_fingerprint=str(data["run_fingerprint"]),
created_at=str(data.get("created_at", "")),
updated_at=str(data.get("updated_at", "")),
completed_combinations=list(data.get("completed_combinations") or []),
summaries_by_key=dict(data.get("summaries_by_key") or {}),
singleton_scores_by_key={
str(k): float(v)
for k, v in (data.get("singleton_scores_by_key") or {}).items()
},
stages=dict(data.get("stages") or {}),
failures=failures,
checkpoint_fusion_pairs=int(cfp) if cfp is not None else None,
checkpoint_fusion_triples=int(cft) if cft is not None else None,
)