Blame for src/utils/ckpt_cache.py

"""Local NVMe staging + loading helpers for slow-filesystem-backed checkpoints.

Why this module exists
----------------------

Repeated checkpoint reads from network or FUSE filesystems can be slow.
Training initialization and inference share these helpers to stage checkpoint
files in a local cache and reuse them across runs.

What it does
------------

* :func:`stage_ckpt_to_local` optionally copies a checkpoint file or its
  regular sibling files to a local cache root and returns the resulting path.
  Callers can choose their own ``torch.load`` semantics.
* :func:`load_ckpt` is the convenience wrapper used by training and inference:
  stages, then ``torch.load(..., mmap=True, weights_only=True)`` with a
  graceful fallback chain and explicit timing logs.

Both honour the same env vars:

* ``WAN_CKPT_LOCAL_CACHE``           - cache root (default
  ``/tmp/wan-ckpt-cache``).
* ``WAN_CKPT_LOCAL_CACHE_ENABLE``    - set to ``0`` / ``false`` to bypass
  staging entirely.
* ``WAN_CKPT_LOCAL_CACHE_BUDGET_GIB`` - LRU eviction budget in GiB
  (default 200).
* ``WAN_CKPT_REMOTE_PREFIXES``       - optional comma-separated list of
  path prefixes whose checkpoints should be staged. When **unset**,
  staging applies to *every* source path outside the cache root.
  When set, only paths starting
  with one of the prefixes are staged - useful if some of your
  checkpoints already live on fast local disk and you want to skip the
  copy for those.

Concurrency
-----------

Multiple ranks / processes on the same node may all enter staging at
once. We serialise per-checkpoint via ``fcntl.flock`` on a sibling
lockfile and gate completion on a ``.stage_done`` sentinel file inside
the staged dir. Whoever wins the race performs the copy; everyone else
hits the fast path on wakeup.

Failure handling
----------------

Every external dependency has a fall-through branch that logs a warning
and returns the *original* path. The caller's load will then succeed,
just slowly. This is intentional: a flaky cache must never block a run.
"""

from __future__ import annotations

import logging
import os
import time

import torch

try:
    import fcntl
except ImportError:  # Windows has no flock; checkpoints then load from their own path.
    fcntl = None

__all__ = [
    "stage_ckpt_to_local",
    "load_ckpt",
]

# --- Configuration (env-driven) ------------------------------------------

_CKPT_LOCAL_CACHE_DIR = os.environ.get("WAN_CKPT_LOCAL_CACHE", "/tmp/wan-ckpt-cache")
_CKPT_LOCAL_CACHE_ENABLE = os.environ.get("WAN_CKPT_LOCAL_CACHE_ENABLE", "1").lower() not in (
    "0",
    "false",
    "no",
)
# Cache budget in GiB; configurable through the environment.
_CKPT_LOCAL_CACHE_BUDGET_BYTES = int(
    float(os.environ.get("WAN_CKPT_LOCAL_CACHE_BUDGET_GIB", "200")) * (1024**3)
)

# Optional whitelist: only paths under one of these prefixes are staged.
# Configure via the comma-separated ``WAN_CKPT_REMOTE_PREFIXES`` env var
# (e.g. ``WAN_CKPT_REMOTE_PREFIXES=/mnt/s3/,/mnt/nfs/``). When unset,
# *every* source path is eligible for staging (we still skip paths that
# are already inside the cache dir, see ``_should_stage`` below).
_REMOTE_FUSE_PREFIX = tuple(
    p for p in os.environ.get("WAN_CKPT_REMOTE_PREFIXES", "").split(",") if p
)


def _should_stage(path: str) -> bool:
    """Return True iff ``path`` should be copied to the local cache.

    Rules:

    * If staging is disabled via env var -> False.
    * If ``WAN_CKPT_REMOTE_PREFIXES`` is set, only paths under one of the
      configured prefixes are staged. This preserves the original opt-in
      whitelist behaviour for users who want it.
    * Otherwise (default), every path is eligible *except* paths that
      already live inside the local cache directory - copying ``cache/x``
      to ``cache/y`` would be pointless and would also break the
      "directory hash -> staged subdir" mapping.
    """
    if fcntl is None or not _CKPT_LOCAL_CACHE_ENABLE:
        return False
    if _REMOTE_FUSE_PREFIX:
        return path.startswith(_REMOTE_FUSE_PREFIX)
    # Default: stage everything that isn't already in the cache root.
    try:
        path_abs = os.path.realpath(path)
        cache_abs = os.path.realpath(_CKPT_LOCAL_CACHE_DIR)
    except OSError:
        return True
    return not path_abs.startswith(cache_abs.rstrip("/") + "/")


def _get_logger(logger):
    """Return a usable logger.

    Callers in ``inference.py`` already have a structured logger; the
    trainer/model code just imports ``logging`` at module level. Accepting
    ``None`` lets us slot into both without forcing every call site to
    construct one.
    """
    return logger if logger is not None else logging.getLogger(__name__)


# --- Public API ----------------------------------------------------------


def stage_ckpt_to_local(path: str, *, logger=None, copy_siblings: bool = True) -> str:
    """Mirror a slow-filesystem-backed checkpoint to a local cache; return new path.

    Staging is enabled by default (``WAN_CKPT_LOCAL_CACHE_ENABLE=1``) and
    applies to every source path unless ``WAN_CKPT_REMOTE_PREFIXES`` is
    set, in which case only paths under one of the configured prefixes
    are staged. Paths that already live inside the cache root are
    returned as-is to avoid self-copies. See :func:`_should_stage` for
    the full predicate.

    When staging applies:

    1. With ``copy_siblings=True`` (the trainer default), hash the parent
       directory and copy every regular sibling so warm-start EMA lookup sees
       the complete checkpoint. With ``False`` (used by inference), hash the
       parent plus selected filename and copy only that file.
    3. Write ``.stage_done`` to mark completion; subsequent calls hit the
       fast path and ``utime`` the sentinel so LRU reflects "most
       recently *used*" rather than "most recently *fetched*".
    4. After a successful stage, evict oldest staged dirs until the cache
       size is back under :data:`_CKPT_LOCAL_CACHE_BUDGET_BYTES`.

    On any error (copy failure, disk full, ...) returns the original
    ``path`` so the caller's ``torch.load`` still works.
    """
    log = _get_logger(logger)

    if not _should_stage(path):
        return path

    import hashlib
    import shutil

    src_dir = os.path.dirname(path)
    fname = os.path.basename(path)
    scope_key = src_dir if copy_siblings else f"{src_dir}\0{fname}"
    dir_key = hashlib.sha1(scope_key.encode("utf-8")).hexdigest()[:16]
    scope_suffix = (
        os.path.basename(src_dir) if copy_siblings else f"{os.path.basename(src_dir)}__{fname}"
    )
    cache_subdir = os.path.join(_CKPT_LOCAL_CACHE_DIR, f"{dir_key}__{scope_suffix}")
    sentinel = os.path.join(cache_subdir, ".stage_done")
    lock_path = os.path.join(_CKPT_LOCAL_CACHE_DIR, f".{dir_key}.lock")
    staged_path = os.path.join(cache_subdir, fname)

    try:
        os.makedirs(_CKPT_LOCAL_CACHE_DIR, exist_ok=True)
    except OSError as e:
        log.warning(
            f"[ckpt-cache] cannot create cache dir {_CKPT_LOCAL_CACHE_DIR} "
            f"({e}); falling back to source path."
        )
        return path

    # Fast path: directory already staged.
    if os.path.isfile(sentinel) and os.path.isfile(staged_path):
        try:
            os.utime(sentinel, None)  # bump LRU mtime
        except OSError:
            pass
        log.info(f"[ckpt-cache] hit: {staged_path}")
        return staged_path

    log.info(
        f"[ckpt-cache] staging {src_dir} -> {cache_subdir} "
        f"(once per ckpt; reused on subsequent runs)"
    )
    _t = time.perf_counter()
    os.makedirs(cache_subdir, exist_ok=True)
    with open(lock_path, "w") as lf:
        fcntl.flock(lf.fileno(), fcntl.LOCK_EX)
        try:
            # Re-check after lock acquisition: a sibling rank may have
            # finished while we were waiting.
            if os.path.isfile(sentinel) and os.path.isfile(staged_path):
                log.info(
                    "[ckpt-cache] another process finished staging while "
                    f"we waited; using {staged_path}"
                )
            else:
                try:
                    if copy_siblings:
                        for entry in os.scandir(src_dir):
                            if not entry.is_file(follow_symlinks=False):
                                continue
                            shutil.copy2(entry.path, os.path.join(cache_subdir, entry.name))
                    else:
                        shutil.copy2(path, staged_path)
                except OSError as e:
                    log.warning(
                        f"[ckpt-cache] copy failed ({e}); falling back to source path {path}"
                    )
                    return path
                if not os.path.isfile(staged_path):
                    log.warning(
                        f"[ckpt-cache] expected staged file {staged_path} "
                        f"missing after copy; falling back to {path}"
                    )
                    return path
                with open(sentinel, "w") as f:
                    f.write(f"staged {src_dir} -> {cache_subdir}\n")
        finally:
            fcntl.flock(lf.fileno(), fcntl.LOCK_UN)

    dt = time.perf_counter() - _t
    try:
        sz = os.path.getsize(staged_path) / (1024**3)
        rate = f", {sz * 1024 / dt:.0f} MiB/s" if dt > 0 else ""
        log.info(f"[ckpt-cache] staged {sz:.1f} GiB in {dt:.1f}s{rate} -> {staged_path}")
    except OSError:
        log.info(f"[ckpt-cache] staged in {dt:.1f}s -> {staged_path}")

    _evict_ckpt_cache_lru(logger=log, keep_subdir=cache_subdir)
    return staged_path


def load_ckpt(
    path: str,
    *,
    logger=None,
    label: str = "ckpt",
    map_location: str = "cpu",
) -> dict:
    """Stage to local NVMe (if remote) and ``torch.load`` with mmap fallback.

    Convenience wrapper for callers that want the full inference-style
    fast path:

    - ``mmap=True`` so tensor storages are demand-paged from the file
      rather than copied into Python memory up-front.
    - ``weights_only=True`` skips arbitrary pickle execution. Falls back
      to ``weights_only=False`` if the file contains non-tensor objects
      (older checkpoints occasionally do), then to a non-mmap full-read
      if mmap fails for any reason.

    Callers that need to control the load flags themselves (e.g. trainer
    code that wants a full-read for predictable timing) should call
    :func:`stage_ckpt_to_local` directly and run their own ``torch.load``.
    """
    log = _get_logger(logger)

    path = stage_ckpt_to_local(path, logger=log, copy_siblings=False)

    file_size_mb: float | None = None
    try:
        file_size_mb = os.path.getsize(path) / (1024 * 1024)
    except OSError:
        pass

    size_str = f" ({file_size_mb:.1f} MiB)" if file_size_mb is not None else ""
    log.info(f"[ckpt] Loading {label} from {path}{size_str}")
    t0 = time.perf_counter()

    try:
        blob = torch.load(path, map_location=map_location, mmap=True, weights_only=True)
        mode = "mmap+weights_only"
    except Exception as e_safe:  # noqa: BLE001
        log.warning(
            f"[ckpt] weights_only=True load failed ({type(e_safe).__name__}: "
            f"{e_safe}); retrying without it."
        )
        try:
            blob = torch.load(path, map_location=map_location, mmap=True, weights_only=False)
            mode = "mmap"
        except Exception as e_mmap:  # noqa: BLE001
            log.warning(
                f"[ckpt] mmap=True load failed ({type(e_mmap).__name__}: "
                f"{e_mmap}); retrying with full read."
            )
            blob = torch.load(path, map_location=map_location, weights_only=False)
            mode = "full-read"

    dt = time.perf_counter() - t0
    rate = f", {file_size_mb / dt:.1f} MiB/s" if file_size_mb is not None and dt > 0 else ""
    log.info(f"[ckpt] {label} loaded in {dt:.2f}s [{mode}]{rate}")
    return blob


# --- Internal: LRU eviction ----------------------------------------------


def _evict_ckpt_cache_lru(*, logger, keep_subdir: str | None = None) -> None:
    """LRU-evict staged ckpt dirs until total size <= budget.

    Order is by ``.stage_done`` mtime, oldest first. ``keep_subdir`` is
    never evicted (typically the one we just staged). Best-effort: any
    I/O error during scanning or removal is logged and ignored.
    """
    try:
        entries: list[tuple[float, str, int]] = []  # (mtime, dir, bytes)
        with os.scandir(_CKPT_LOCAL_CACHE_DIR) as it:
            for de in it:
                if not de.is_dir(follow_symlinks=False):
                    continue
                sub = de.path
                sentinel = os.path.join(sub, ".stage_done")
                if not os.path.isfile(sentinel):
                    continue  # in-progress or unrelated; skip
                try:
                    mt = os.path.getmtime(sentinel)
                except OSError:
                    continue
                total = 0
                for root, _, files in os.walk(sub):
                    for fn in files:
                        try:
                            total += os.path.getsize(os.path.join(root, fn))
                        except OSError:
                            pass
                entries.append((mt, sub, total))
    except FileNotFoundError:
        return

    total_bytes = sum(e[2] for e in entries)
    if total_bytes <= _CKPT_LOCAL_CACHE_BUDGET_BYTES:
        return

    import shutil as _shutil

    entries.sort(key=lambda e: e[0])  # oldest first
    budget = _CKPT_LOCAL_CACHE_BUDGET_BYTES
    over = total_bytes - budget
    logger.info(
        f"[ckpt-cache] LRU eviction: {total_bytes / 1024**3:.1f} GiB > "
        f"{budget / 1024**3:.1f} GiB budget, freeing "
        f"{over / 1024**3:.1f} GiB"
    )
    for _mt, sub, sz in entries:
        if total_bytes <= budget:
            break
        if keep_subdir is not None and os.path.realpath(sub) == os.path.realpath(keep_subdir):
            continue
        try:
            _shutil.rmtree(sub)
            total_bytes -= sz
            logger.info(f"[ckpt-cache]   evicted {sub} ({sz / 1024**3:.1f} GiB)")
        except OSError as e:
            logger.warning(f"[ckpt-cache]   failed to evict {sub}: {e}")