"""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}")