scripts/checkpoint_io.py

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import torch


SPENT_PREDICTIONS = frozenset({"noise_pred", "audio_noise_pred"})


class DenoiseCheckpointStop(RuntimeError):
    def __init__(self, *, completed_steps: int, total_steps: int):
        self.completed_steps = completed_steps
        self.total_steps = total_steps
        super().__init__(f"stopped after durable checkpoint {completed_steps} of {total_steps}")


def _partial_path(path: Path) -> Path:
    return path.with_name(f"{path.name}.partial-{os.getpid()}")


def _sync_and_replace(partial: Path, target: Path) -> None:
    with partial.open("rb") as handle:
        os.fsync(handle.fileno())
    os.replace(partial, target)
    directory_fd = os.open(target.parent, os.O_RDONLY)
    try:
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)


def atomic_torch_save(payload: Any, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    partial = _partial_path(path)
    try:
        torch.save(payload, partial)
        _sync_and_replace(partial, path)
    finally:
        partial.unlink(missing_ok=True)


def atomic_json_write(payload: Any, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    partial = _partial_path(path)
    try:
        with partial.open("w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2, ensure_ascii=False)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(partial, path)
        directory_fd = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
    finally:
        partial.unlink(missing_ok=True)


def snapshot_to_cpu(value: Any) -> Any:
    if isinstance(value, torch.Tensor):
        return value.detach().to(device="cpu").contiguous().clone()
    if isinstance(value, dict):
        return {key: snapshot_to_cpu(item) for key, item in value.items()}
    if isinstance(value, list):
        return [snapshot_to_cpu(item) for item in value]
    if isinstance(value, tuple):
        return tuple(snapshot_to_cpu(item) for item in value)
    return value


def write_step_checkpoint(
    directory: Path,
    *,
    step_index: int,
    total_steps: int,
    block_state,
    metadata: dict[str, Any],
) -> tuple[Path, Path]:
    if not 0 <= step_index < total_steps:
        raise ValueError(f"invalid checkpoint step {step_index} for total {total_steps}")
    completed_step = step_index + 1
    stem = f"step-{completed_step:03d}-of-{total_steps:03d}"
    state = snapshot_to_cpu(block_state.as_dict())
    for name in SPENT_PREDICTIONS:
        state.pop(name, None)
    payload = {
        "schema_version": 1,
        "completed_step": completed_step,
        "resume_step_index": completed_step,
        "total_steps": total_steps,
        "metadata": dict(metadata),
        "state": state,
    }
    checkpoint = directory / f"{stem}.pt"
    manifest = directory / f"{stem}.json"
    atomic_torch_save(payload, checkpoint)
    atomic_json_write(
        {
            "status": "checkpointed",
            "completed_step": completed_step,
            "resume_step_index": completed_step,
            "total_steps": total_steps,
            "checkpoint": checkpoint.name,
            **metadata,
        },
        manifest,
    )
    atomic_json_write(
        {
            "status": "checkpointed",
            "completed_step": completed_step,
            "resume_step_index": completed_step,
            "total_steps": total_steps,
            "checkpoint": checkpoint.name,
            "manifest": manifest.name,
            **metadata,
        },
        directory / "latest.json",
    )
    return checkpoint, manifest


def install_loop_checkpointing(
    loop_wrapper_cls,
    directory: Path,
    *,
    metadata: dict[str, Any],
    stop_after_steps: int | None = None,
):
    if stop_after_steps is not None and stop_after_steps < 1:
        raise ValueError("stop_after_steps must be positive")
    original_call = loop_wrapper_cls.__call__

    @torch.no_grad()
    def checkpointed_call(self, components, state):
        block_state = self.get_block_state(state)
        total_steps = len(block_state.timesteps)
        with self.progress_bar(total=total_steps) as progress_bar:
            for step_index, timestep in enumerate(block_state.timesteps):
                components, block_state = self.loop_step(
                    components,
                    block_state,
                    i=step_index,
                    t=timestep,
                )
                write_step_checkpoint(
                    directory,
                    step_index=step_index,
                    total_steps=total_steps,
                    block_state=block_state,
                    metadata=metadata,
                )
                progress_bar.update()
                completed_steps = step_index + 1
                if stop_after_steps is not None and completed_steps >= stop_after_steps:
                    self.set_block_state(state, block_state)
                    raise DenoiseCheckpointStop(
                        completed_steps=completed_steps,
                        total_steps=total_steps,
                    )
        self.set_block_state(state, block_state)
        return components, state

    loop_wrapper_cls.__call__ = checkpointed_call
    return original_call