evidence/source/checkpoint_fs.py

"""Crash-safe filesystem transactions for Wisp checkpoint directories."""

import ctypes
import errno
import hashlib
import json
import math
import os
import re
import secrets
import shutil
import stat
import sys


RENAME_SWAP = 0x00000002
RENAME_EXCL = 0x00000004
CHECKPOINT_FILENAMES = (
    "master.safetensors",
    "meta.json",
    "optimizer.safetensors",
)


def _require_real_directory(path, label):
    if os.path.islink(path) or not os.path.isdir(path):
        raise ValueError(f"{label} is not a real directory: {path}")


def _fsync_directory(path):
    descriptor = os.open(path, os.O_RDONLY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def _same_identity(left, right):
    return (
        left.st_dev,
        left.st_ino,
        left.st_mode,
        left.st_size,
        left.st_mtime_ns,
        left.st_ctime_ns,
    ) == (
        right.st_dev,
        right.st_ino,
        right.st_mode,
        right.st_size,
        right.st_mtime_ns,
        right.st_ctime_ns,
    )


def _stable_regular_file(path, label, capture_bytes=False):
    before_path = os.lstat(path)
    if not stat.S_ISREG(before_path.st_mode):
        raise ValueError(f"{label} is not a real regular file: {path}")
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(path, flags)
    digest = hashlib.sha256()
    blocks = [] if capture_bytes else None
    try:
        before_fd = os.fstat(descriptor)
        if not stat.S_ISREG(before_fd.st_mode) or not _same_identity(
            before_path, before_fd
        ):
            raise ValueError(f"{label} changed while it was opened: {path}")
        while True:
            block = os.read(descriptor, 1024 * 1024)
            if not block:
                break
            digest.update(block)
            if blocks is not None:
                blocks.append(block)
        after_fd = os.fstat(descriptor)
    finally:
        os.close(descriptor)
    after_path = os.lstat(path)
    if not (
        _same_identity(before_fd, after_fd)
        and _same_identity(before_fd, after_path)
    ):
        raise ValueError(f"{label} changed while it was read: {path}")
    return {
        "mode": stat.S_IMODE(before_fd.st_mode),
        "size": before_fd.st_size,
        "sha256": digest.hexdigest(),
        "bytes": b"".join(blocks) if blocks is not None else None,
    }


def _reject_duplicate_object_keys(pairs):
    value = {}
    for key, item in pairs:
        if key in value:
            raise ValueError(f"duplicate JSON key: {key}")
        value[key] = item
    return value


def _finite_json_float(raw):
    value = float(raw)
    if not math.isfinite(value):
        raise ValueError(f"non-finite JSON number: {raw}")
    return value


def _parse_checkpoint_meta(raw, label):
    try:
        value = json.loads(
            raw,
            object_pairs_hook=_reject_duplicate_object_keys,
            parse_float=_finite_json_float,
            parse_constant=lambda token: (_ for _ in ()).throw(
                ValueError(f"non-finite JSON number: {token}")
            ),
        )
    except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as exc:
        raise ValueError(f"{label} is not strict JSON") from exc
    legacy_fields = {
        "step",
        "config",
        "model_args",
        "optimizer_state_included",
    }
    exact_fields = set(value) if isinstance(value, dict) else set()
    if exact_fields not in (legacy_fields, legacy_fields | {"train_sampler"}):
        raise ValueError(f"{label} has an unknown root field set")
    if (
        not isinstance(value["step"], int)
        or isinstance(value["step"], bool)
        or value["step"] <= 0
        or not isinstance(value["config"], dict)
        or not isinstance(value["model_args"], dict)
        or value["optimizer_state_included"] is not True
        or (
            "train_sampler" in value
            and not isinstance(value["train_sampler"], dict)
        )
    ):
        raise ValueError(f"{label} has an invalid training checkpoint shape")
    return value


def _checkpoint_manifest(path, label):
    before = os.lstat(path)
    if not stat.S_ISDIR(before.st_mode):
        raise ValueError(f"{label} is not a real directory: {path}")
    names = tuple(sorted(os.listdir(path)))
    if names != tuple(sorted(CHECKPOINT_FILENAMES)):
        raise ValueError(f"{label} does not have the exact checkpoint file set")
    files = {}
    metadata = None
    for name in CHECKPOINT_FILENAMES:
        result = _stable_regular_file(
            os.path.join(path, name),
            f"{label} {name}",
            capture_bytes=name == "meta.json",
        )
        files[name] = {
            "mode": result["mode"],
            "size": result["size"],
            "sha256": result["sha256"],
        }
        if name == "meta.json":
            metadata = _parse_checkpoint_meta(
                result["bytes"],
                f"{label} meta.json",
            )
    after = os.lstat(path)
    if not _same_identity(before, after):
        raise ValueError(f"{label} changed while it was inspected: {path}")
    return {"files": files, "metadata": metadata}


def _copy_regular_file(source, destination, label):
    source_flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    source_descriptor = os.open(source, source_flags)
    destination_descriptor = None
    try:
        source_stat = os.fstat(source_descriptor)
        if not stat.S_ISREG(source_stat.st_mode):
            raise ValueError(f"{label} source is not a regular file")
        destination_descriptor = os.open(
            destination,
            os.O_WRONLY | os.O_CREAT | os.O_EXCL,
            stat.S_IMODE(source_stat.st_mode),
        )
        os.fchmod(destination_descriptor, stat.S_IMODE(source_stat.st_mode))
        while True:
            block = os.read(source_descriptor, 1024 * 1024)
            if not block:
                break
            offset = 0
            while offset < len(block):
                offset += os.write(destination_descriptor, block[offset:])
        os.fsync(destination_descriptor)
    finally:
        if destination_descriptor is not None:
            os.close(destination_descriptor)
        os.close(source_descriptor)


def _rename_directory_no_replace(staging, target):
    if sys.platform != "darwin":
        raise RuntimeError(
            "exclusive checkpoint directory publication requires macOS"
        )
    libc = ctypes.CDLL(None, use_errno=True)
    renamex_np = libc.renamex_np
    renamex_np.argtypes = [
        ctypes.c_char_p,
        ctypes.c_char_p,
        ctypes.c_uint,
    ]
    renamex_np.restype = ctypes.c_int
    result = renamex_np(
        os.fsencode(staging),
        os.fsencode(target),
        RENAME_EXCL,
    )
    if result != 0:
        error = ctypes.get_errno()
        raise OSError(error, os.strerror(error), f"{staging} -> {target}")


def _remove_snapshot_recovery_staging(staging):
    """Remove only a private staging tree with checkpoint-shaped contents."""
    _require_real_directory(staging, "snapshot recovery staging")
    names = tuple(sorted(os.listdir(staging)))
    unexpected = set(names) - set(CHECKPOINT_FILENAMES)
    if unexpected:
        raise ValueError(
            "snapshot recovery staging contains an unowned entry: "
            + sorted(unexpected)[0]
        )
    for name in names:
        path = os.path.join(staging, name)
        info = os.lstat(path)
        if not stat.S_ISREG(info.st_mode):
            raise ValueError(
                "snapshot recovery staging contains a non-file entry: "
                + name
            )
    for name in names:
        os.unlink(os.path.join(staging, name))
    os.rmdir(staging)


def _create_snapshot_recovery_staging(snapshot_checkpoint):
    """Create a private, invocation-owned sibling staging directory."""
    parent = os.path.dirname(snapshot_checkpoint)
    basename = os.path.basename(snapshot_checkpoint)
    for _ in range(128):
        token = secrets.token_hex(16)
        staging = os.path.join(
            parent,
            f".{basename}.recovery-{token}.partial",
        )
        try:
            os.mkdir(staging, 0o700)
        except FileExistsError:
            continue
        os.chmod(staging, 0o700)
        return staging
    raise RuntimeError("could not allocate private snapshot recovery staging")


def ensure_snapshot_checkpoint(
    resume_checkpoint,
    snapshot_checkpoint,
    expected_step,
    expected_metadata,
):
    """
    Make a missing immortal snapshot an exact copy of the resume checkpoint.

    A differing existing snapshot is evidence from another training state and
    is never replaced. Publication is exclusive, so even a racing creator
    cannot be overwritten.
    """
    if (
        not isinstance(expected_step, int)
        or isinstance(expected_step, bool)
        or expected_step <= 0
        or not isinstance(expected_metadata, dict)
    ):
        raise ValueError("snapshot recovery inputs are malformed")
    resume_checkpoint = os.path.abspath(resume_checkpoint)
    snapshot_checkpoint = os.path.abspath(snapshot_checkpoint)
    if resume_checkpoint == snapshot_checkpoint:
        raise ValueError("resume and snapshot checkpoint paths must differ")
    parent = os.path.dirname(snapshot_checkpoint)
    _require_real_directory(parent, "snapshot checkpoint parent")

    source_manifest = _checkpoint_manifest(
        resume_checkpoint,
        "resume checkpoint",
    )
    if (
        source_manifest["metadata"]["step"] != expected_step
        or source_manifest["metadata"] != expected_metadata
    ):
        raise ValueError("resume checkpoint metadata differs from loaded state")

    if os.path.lexists(snapshot_checkpoint):
        snapshot_manifest = _checkpoint_manifest(
            snapshot_checkpoint,
            "immortal snapshot checkpoint",
        )
        if snapshot_manifest != source_manifest:
            raise ValueError(
                "refusing to replace a differing immortal snapshot checkpoint"
            )
        _fsync_directory(parent)
        return False

    staging = _create_snapshot_recovery_staging(snapshot_checkpoint)
    published = False
    try:
        for name in CHECKPOINT_FILENAMES:
            _copy_regular_file(
                os.path.join(resume_checkpoint, name),
                os.path.join(staging, name),
                f"snapshot recovery {name}",
            )
        fsync_checkpoint_tree(staging)
        if _checkpoint_manifest(staging, "staged snapshot checkpoint") != (
            source_manifest
        ):
            raise ValueError("staged snapshot differs from resume checkpoint")
        if _checkpoint_manifest(
            resume_checkpoint,
            "resume checkpoint after snapshot copy",
        ) != source_manifest:
            raise ValueError("resume checkpoint changed during snapshot copy")
        try:
            _rename_directory_no_replace(staging, snapshot_checkpoint)
            published = True
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise
        snapshot_manifest = _checkpoint_manifest(
            snapshot_checkpoint,
            "published immortal snapshot checkpoint",
        )
        if snapshot_manifest != source_manifest:
            raise ValueError(
                "published immortal snapshot differs from resume checkpoint"
            )
        _fsync_directory(parent)
        return published
    finally:
        if os.path.lexists(staging):
            _remove_snapshot_recovery_staging(staging)


def _active_snapshot_marker_count(raw_log, step, resume_log_path):
    marker = json.dumps({"step": step, "snapshot": True}).encode("ascii")
    resume_prefix = f"resumed from {resume_log_path} at step ".encode("ascii")
    active_count = 0
    for line in raw_log.splitlines():
        if line.startswith(resume_prefix):
            raw_step = line[len(resume_prefix):]
            if re.fullmatch(rb"0|[1-9][0-9]*", raw_step):
                resume_step = int(raw_step)
                if resume_step < step:
                    active_count = 0
        elif line == marker:
            active_count += 1
    return active_count


def prepare_snapshot_boundary_recovery(
    resume_checkpoint,
    snapshot_checkpoint,
    train_log_path,
    resume_log_path,
    expected_step,
    expected_metadata,
):
    """
    Verify/recreate a snapshot and return its exact missing stdout marker.

    Resume lines below the snapshot step abandon all earlier markers at that
    step. At most one marker may remain active; repeated recovery at the same
    boundary therefore emits nothing.
    """
    ensure_snapshot_checkpoint(
        resume_checkpoint,
        snapshot_checkpoint,
        expected_step,
        expected_metadata,
    )
    if os.path.lexists(train_log_path):
        result = _stable_regular_file(
            train_log_path,
            "training log used for snapshot recovery",
            capture_bytes=True,
        )
        raw_log = result["bytes"]
    else:
        raw_log = b""
    active_count = _active_snapshot_marker_count(
        raw_log,
        expected_step,
        resume_log_path,
    )
    if active_count > 1:
        raise ValueError(
            "training log has duplicate active snapshot markers at resume step"
        )
    if active_count == 1:
        return None
    return json.dumps(
        {"step": expected_step, "snapshot": True}
    ).encode("ascii")


def fsync_checkpoint_tree(path):
    """Flush every staged file and its directory before publication."""
    _require_real_directory(path, "staged checkpoint")
    for name in sorted(os.listdir(path)):
        item = os.path.join(path, name)
        if os.path.islink(item) or not os.path.isfile(item):
            raise ValueError(
                f"staged checkpoint contains a non-file entry: {name}"
            )
        descriptor = os.open(item, os.O_RDONLY)
        try:
            os.fsync(descriptor)
        finally:
            os.close(descriptor)
    _fsync_directory(path)


def _swap_directories_macos(left, right):
    if sys.platform != "darwin":
        raise RuntimeError(
            "atomic checkpoint directory exchange requires macOS"
        )
    libc = ctypes.CDLL(None, use_errno=True)
    renamex_np = libc.renamex_np
    renamex_np.argtypes = [
        ctypes.c_char_p,
        ctypes.c_char_p,
        ctypes.c_uint,
    ]
    renamex_np.restype = ctypes.c_int
    result = renamex_np(
        os.fsencode(left),
        os.fsencode(right),
        RENAME_SWAP,
    )
    if result != 0:
        error = ctypes.get_errno()
        raise OSError(
            error,
            os.strerror(error),
            f"{left} <-> {right}",
        )


def install_checkpoint(staging, target):
    """Publish a complete staged checkpoint without removing the live path."""
    staging = os.path.abspath(staging)
    target = os.path.abspath(target)
    parent = os.path.dirname(target)
    previous = target + ".prev"
    _require_real_directory(staging, "staged checkpoint")
    fsync_checkpoint_tree(staging)

    if os.path.lexists(target):
        _require_real_directory(target, "current checkpoint")
        if os.path.lexists(previous):
            if os.path.islink(previous) or not os.path.isdir(previous):
                raise ValueError(
                    f"previous checkpoint is not a real directory: {previous}"
                )
            shutil.rmtree(previous)
        _swap_directories_macos(staging, target)
        os.rename(staging, previous)
    else:
        os.rename(staging, target)
    _fsync_directory(parent)
    return target


def recover_checkpoint(target):
    """Recover the last canonical checkpoint from a legacy rename gap."""
    target = os.path.abspath(target)
    if os.path.lexists(target):
        _require_real_directory(target, "current checkpoint")
        return False
    previous = target + ".prev"
    if not os.path.lexists(previous):
        return False
    _require_real_directory(previous, "previous checkpoint")
    try:
        os.rename(previous, target)
    except OSError as exc:
        if exc.errno == errno.ENOENT and os.path.isdir(target):
            return False
        raise
    _fsync_directory(os.path.dirname(target))
    return True