OBLITERATUS/obliteratus/persistence_contracts.py

"""Deterministic and transactional contracts for local model checkpoints."""

from __future__ import annotations

import json
import logging
import os
import shutil
import stat
import tempfile
import uuid
from collections.abc import Mapping
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Callable, Iterator, Protocol


logger = logging.getLogger(__name__)

if os.name == "nt":  # pragma: no cover - exercised by the Windows CI lane
    _DIRECTORY_SYNC_FLAGS: int | None = None
else:
    _DIRECTORY_SYNC_FLAGS = os.O_RDONLY | os.O_DIRECTORY


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 _read_json_object(path: Path, artifact_name: str) -> Mapping[str, Any]:
    """Load a required, non-empty JSON object from a checkpoint."""
    if path.is_symlink() or not path.is_file() or path.stat().st_size == 0:
        raise ValueError(f"Checkpoint {artifact_name} is missing or empty: {path}")
    try:
        value = json.loads(path.read_bytes().decode())
    except (OSError, UnicodeError, json.JSONDecodeError) as error:
        raise ValueError(f"Checkpoint {artifact_name} is corrupt: {path}") from error
    if not isinstance(value, Mapping):
        raise ValueError(f"Checkpoint {artifact_name} must contain a JSON object: {path}")
    return value


def _require_nonempty_regular_file(path: Path, artifact_name: str) -> None:
    """Reject absent, empty, or link-backed checkpoint payloads."""
    if (
        path.is_symlink()
        or not path.is_file()
        or not stat.S_ISREG(path.stat().st_mode)
        or path.stat().st_size == 0
    ):
        raise ValueError(f"Checkpoint {artifact_name} is missing or empty: {path}")


def validate_local_checkpoint(
    checkpoint_dir: Path,
    expected_metadata_json: str,
) -> None:
    """Validate the minimum artifacts needed for a local Transformers reload."""
    checkpoint_dir = Path(checkpoint_dir)
    if checkpoint_dir.is_symlink() or not checkpoint_dir.is_dir():
        raise ValueError(f"Checkpoint staging path is not a directory: {checkpoint_dir}")

    expected_metadata = json.loads(expected_metadata_json)
    actual_metadata = _read_json_object(
        checkpoint_dir / "abliteration_metadata.json",
        "metadata",
    )
    if actual_metadata != expected_metadata:
        raise ValueError("Checkpoint metadata does not match the prepared transaction")

    _read_json_object(checkpoint_dir / "config.json", "model config")
    _read_json_object(checkpoint_dir / "tokenizer_config.json", "tokenizer config")

    direct_weights = [
        checkpoint_dir / "model.safetensors",
        checkpoint_dir / "pytorch_model.bin",
    ]
    for weights_path in direct_weights:
        if _path_exists(weights_path):
            _require_nonempty_regular_file(weights_path, "weights")
            return

    index_paths = [
        checkpoint_dir / "model.safetensors.index.json",
        checkpoint_dir / "pytorch_model.bin.index.json",
    ]
    for index_path in index_paths:
        if not _path_exists(index_path):
            continue
        index = _read_json_object(index_path, "weight index")
        weight_map = index.get("weight_map")
        if not isinstance(weight_map, Mapping) or not weight_map:
            raise ValueError(f"Checkpoint weight index has no weight map: {index_path}")
        shard_values = list(weight_map.values())
        if not all(isinstance(name, str) and Path(name).name == name for name in shard_values):
            raise ValueError(f"Checkpoint weight index contains an unsafe shard path: {index_path}")
        shard_names = set(shard_values)
        for shard_name in shard_names:
            _require_nonempty_regular_file(
                checkpoint_dir / shard_name,
                "weight shard",
            )
        return

    raise ValueError(f"Checkpoint has no model weights: {checkpoint_dir}")


def _path_exists(path: Path) -> bool:
    """Return whether a path or dangling symlink occupies ``path``."""
    return path.exists() or path.is_symlink()


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():
        try:
            shutil.rmtree(path)
        except FileNotFoundError:
            pass


def _sync_file(path: Path) -> None:
    """Flush one regular checkpoint artifact to stable storage."""
    mode = path.lstat().st_mode
    if not stat.S_ISREG(mode):
        raise OSError(f"Checkpoint artifact is not a regular file: {path}")
    descriptor = os.open(path, os.O_RDONLY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def _sync_directory(path: Path) -> None:
    """Persist directory entries where the platform exposes directory fsync."""
    if _DIRECTORY_SYNC_FLAGS is None:
        return
    descriptor = os.open(path, _DIRECTORY_SYNC_FLAGS)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def _sync_checkpoint_tree(checkpoint_dir: Path) -> None:
    """Flush checkpoint files and directories from leaves to root."""
    for root, directory_names, file_names in os.walk(
        checkpoint_dir,
        topdown=False,
        followlinks=False,
    ):
        root_path = Path(root)
        for file_name in sorted(file_names):
            _sync_file(root_path / file_name)
        for directory_name in sorted(directory_names):
            directory = root_path / directory_name
            if directory.is_symlink():
                raise OSError(
                    f"Checkpoint directory contains a symlink: {directory}",
                )
            _sync_directory(directory)
        _sync_directory(root_path)


@contextmanager
def _checkpoint_lock(destination: Path) -> Iterator[None]:
    """Serialize commits to one destination across threads and processes."""
    lock_path = destination.with_name(f".{destination.name}.lock")
    with lock_path.open("a+b") as lock_file:
        if os.name == "nt":  # pragma: no cover - exercised by the Windows CI lane
            import msvcrt

            lock_file.seek(0, os.SEEK_END)
            if lock_file.tell() == 0:
                lock_file.write(b"\0")
                lock_file.flush()
            lock_file.seek(0)
            msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1)
            try:
                yield
            finally:
                lock_file.seek(0)
                msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
        else:
            import fcntl

            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)


def _rollback_checkpoint_commit(
    destination: Path,
    staging: Path,
    backup: Path | None,
    *,
    had_destination: bool,
) -> None:
    """Restore the pre-commit state after a failed or cancelled commit."""
    if backup is not None and _path_exists(backup):
        if _path_exists(destination):
            if _path_exists(staging):
                raise RuntimeError(
                    "Checkpoint rollback found both destination and staging; "
                    f"recover the previous checkpoint from {backup}",
                )
            os.replace(destination, staging)
        os.replace(backup, destination)
    elif not had_destination and _path_exists(destination) and not _path_exists(staging):
        os.replace(destination, staging)
    _sync_directory(destination.parent)


def _cleanup_failed_staging(staging: Path) -> None:
    """Remove failed staging without masking the transaction's primary error."""
    try:
        _remove_checkpoint_path(staging)
    except OSError as cleanup_error:
        logger.warning(
            "Checkpoint write failed, and staging directory %s could not be "
            "removed: %s",
            staging,
            cleanup_error,
        )


@contextmanager
def atomic_checkpoint_directory(
    destination: Path,
    *,
    validate: Callable[[Path], None] | None = None,
) -> 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. Concurrent commits to the same destination are serialized.
    Validation and durable file flushes happen before the commit lock is taken.
    If commit or parent-directory sync fails, the previous destination is
    restored. Cancellation and other ``BaseException`` failures follow the same
    rollback path as ordinary exceptions.
    """
    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 validate is not None:
            validate(staging)
        _sync_checkpoint_tree(staging)
        with _checkpoint_lock(destination):
            had_destination = _path_exists(destination)
            try:
                if had_destination:
                    backup = destination.with_name(
                        f".{destination.name}.backup-{uuid.uuid4().hex}",
                    )
                    os.replace(destination, backup)
                    _sync_directory(destination.parent)
                os.replace(staging, destination)
                _sync_directory(destination.parent)
            except BaseException:
                try:
                    _rollback_checkpoint_commit(
                        destination,
                        staging,
                        backup,
                        had_destination=had_destination,
                    )
                except BaseException as restore_error:
                    recovery_path = backup if backup is not None else staging
                    raise RuntimeError(
                        "Checkpoint promotion and rollback both failed; "
                        f"recover the previous checkpoint from {recovery_path}",
                    ) from restore_error
                raise
            if backup is not None:
                try:
                    _remove_checkpoint_path(backup)
                    _sync_directory(destination.parent)
                except OSError as cleanup_error:
                    logger.warning(
                        "Checkpoint promoted, but previous-checkpoint backup %s "
                        "could not be removed durably: %s",
                        backup,
                        cleanup_error,
                    )
    except BaseException:
        _cleanup_failed_staging(staging)
        raise