| """Race-safe immutable checkpoint management. |
| |
| Checkpoint directories are written under a temporary name and atomically |
| renamed once complete. ``best`` and ``latest`` are atomically replaced symlinks. |
| No checkpoint is deleted while a run is active; pruning, if ever requested, |
| must happen only after all consumers have stopped. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import uuid |
| from pathlib import Path |
| from typing import Any, Callable |
|
|
|
|
| class AtomicCheckpointManager: |
| def __init__(self, root: Path): |
| self.root = root.resolve() |
| self.checkpoints = self.root / "checkpoints" |
| self.checkpoints.mkdir(parents=True, exist_ok=True) |
|
|
| def _atomic_link(self, link_name: str, target: Path) -> None: |
| link = self.root / link_name |
| temp = self.root / f".{link_name}.tmp.{os.getpid()}.{uuid.uuid4().hex}" |
| relative_target = os.path.relpath(target, self.root) |
| os.symlink(relative_target, temp) |
| os.replace(temp, link) |
|
|
| def save( |
| self, |
| tag: str, |
| writer: Callable[[Path], None], |
| metadata: dict[str, Any], |
| *, |
| is_best: bool, |
| ) -> Path: |
| final = self.checkpoints / tag |
| if final.exists(): |
| raise FileExistsError(f"immutable checkpoint already exists: {final}") |
| temp = self.checkpoints / f".{tag}.tmp.{os.getpid()}.{uuid.uuid4().hex}" |
| temp.mkdir(parents=False) |
| try: |
| writer(temp) |
| with (temp / "checkpoint_metadata.json").open("w", encoding="utf-8") as f: |
| json.dump(metadata, f, ensure_ascii=False, indent=2) |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(temp, final) |
| except Exception: |
| |
| |
| raise |
|
|
| self._atomic_link("latest", final) |
| if is_best: |
| self._atomic_link("best", final) |
| return final |
|
|
|
|
