File size: 2,039 Bytes
8f4f748 | 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 59 60 | """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:
# Leave a failed temporary directory for diagnosis; never remove a
# possibly-active directory from another process.
raise
self._atomic_link("latest", final)
if is_best:
self._atomic_link("best", final)
return final
|