mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
"""Deterministic and transactional contracts for local model checkpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import uuid
|
|
from collections.abc import Mapping
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any, Iterator, Protocol
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SizedTensor(Protocol):
|
|
"""Structural subset used to estimate a serialized state dictionary."""
|
|
|
|
def numel(self) -> int: ...
|
|
|
|
def element_size(self) -> int: ...
|
|
|
|
|
|
def state_dict_size_bytes(state_dict: Mapping[str, SizedTensor]) -> int:
|
|
"""Return the exact unsharded tensor payload size for a state dictionary."""
|
|
return sum(value.numel() * value.element_size() for value in state_dict.values())
|
|
|
|
|
|
def required_checkpoint_bytes(payload_bytes: int) -> int:
|
|
"""Return the checkpoint capacity requirement with ten percent headroom."""
|
|
if payload_bytes < 0:
|
|
raise ValueError("checkpoint payload size cannot be negative")
|
|
return (payload_bytes * 11 + 9) // 10
|
|
|
|
|
|
def ensure_checkpoint_capacity(
|
|
free_bytes: int,
|
|
payload_bytes: int,
|
|
) -> None:
|
|
"""Reject a checkpoint write that cannot satisfy the headroom policy."""
|
|
if free_bytes < required_checkpoint_bytes(payload_bytes):
|
|
raise OSError(
|
|
f"Insufficient disk space: {free_bytes / 1e9:.1f} GB free, "
|
|
f"need ~{payload_bytes / 1e9:.1f} GB. "
|
|
f"Try a different --output-dir on a larger filesystem.",
|
|
)
|
|
|
|
|
|
def serialize_checkpoint_metadata(metadata: Mapping[str, Any]) -> str:
|
|
"""Serialize checkpoint metadata before any model artifact is written."""
|
|
return json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False)
|
|
|
|
|
|
def _remove_checkpoint_path(path: Path) -> None:
|
|
"""Remove a staging/backup path without following directory symlinks."""
|
|
if path.is_symlink() or path.is_file():
|
|
path.unlink(missing_ok=True)
|
|
elif path.exists():
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
|
|
@contextmanager
|
|
def atomic_checkpoint_directory(destination: Path) -> Iterator[Path]:
|
|
"""Yield a sibling staging directory and atomically promote it on success.
|
|
|
|
An existing checkpoint is moved to a uniquely named backup immediately
|
|
before promotion. If promotion fails, that backup is restored. Exceptions
|
|
while writing only remove the staging directory and leave the destination
|
|
untouched.
|
|
"""
|
|
destination = Path(destination)
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
staging = Path(
|
|
tempfile.mkdtemp(
|
|
prefix=f".{destination.name or 'checkpoint'}.staging-",
|
|
dir=destination.parent,
|
|
),
|
|
)
|
|
backup: Path | None = None
|
|
try:
|
|
yield staging
|
|
if destination.exists() or destination.is_symlink():
|
|
backup = destination.with_name(
|
|
f".{destination.name}.backup-{uuid.uuid4().hex}",
|
|
)
|
|
os.replace(destination, backup)
|
|
try:
|
|
os.replace(staging, destination)
|
|
except Exception as promotion_error:
|
|
if backup is not None and (backup.exists() or backup.is_symlink()):
|
|
try:
|
|
os.replace(backup, destination)
|
|
except Exception as restore_error:
|
|
raise RuntimeError(
|
|
"Checkpoint promotion and rollback both failed; "
|
|
f"recover the previous checkpoint from {backup}",
|
|
) from restore_error
|
|
raise promotion_error
|
|
if backup is not None:
|
|
try:
|
|
_remove_checkpoint_path(backup)
|
|
except OSError as cleanup_error:
|
|
logger.warning(
|
|
"Checkpoint promoted, but previous-checkpoint backup %s "
|
|
"could not be removed: %s",
|
|
backup,
|
|
cleanup_error,
|
|
)
|
|
except Exception:
|
|
_remove_checkpoint_path(staging)
|
|
raise
|