[Diffusion] Clarify runtime cache module boundary

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/model_loader/runtime_cache_utils.py

@@ -1,28 +1,63 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Deterministic identity helpers for the DLO runtime-weight cache."""
+"""Reusable identity, hashing, and file helpers for the DLO runtime cache."""
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
+import errno
+import fcntl
+import hashlib
+import inspect
 import json
 import os
-from collections.abc import Mapping
+import shutil
+import time
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import TypeAlias
+from typing import Protocol, TypeAlias, TypeVar
 
 import torch
+from torch import nn
 
 JsonScalar: TypeAlias = None | bool | int | float | str
 JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]
 PathIdentity: TypeAlias = str | os.PathLike[str] | os.PathLike[bytes]
+_HASH_CHUNK_BYTES = 64 * 1024**2
+_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
 
 
 class IdentityNormalizationError(ValueError):
     """Raised when a cache identity contains a process-unstable value."""
 
 
+@dataclass(frozen=True)
+class RuntimeTensor:
+    """One final runtime tensor and its manifest identity."""
+
+    name: str
+    tensor: torch.Tensor
+    kind: str
+
+
+class RuntimeCacheError(RuntimeError):
+    """Expected cache failure that should fall back with a stable code."""
+
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+
+
+class _HashUpdater(Protocol):
+    def update(self, data: bytes | bytearray | memoryview) -> object: ...
+
+
+_RuntimeTensorT = TypeVar("_RuntimeTensorT", bound=RuntimeTensor)
+
+
 def _type_identity(value: object) -> str:
     value_type = value if isinstance(value, type) else type(value)
     return f"{value_type.__module__}.{value_type.__qualname__}"
@@ -88,10 +123,163 @@ def canonical_json(value: object) -> bytes:
     ).encode()
 
 
+def implementation_fingerprint(
+    loader_type: type,
+    pipeline: nn.Module,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> str:
+    """Hash relevant loader and module implementations for cache identity."""
+    objects: list[object] = [loader_type, type(pipeline)]
+    load_weights = getattr(type(pipeline), "load_weights", None)
+    if load_weights is not None:
+        objects.append(load_weights)
+    for _, dit_module in dit_modules:
+        objects.append(type(dit_module))
+        post_load = getattr(type(dit_module), "post_load_weights", None)
+        if post_load is not None:
+            objects.append(post_load)
+
+    digest = hashlib.sha256()
+    identities: set[str] = set()
+    for obj in objects:
+        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
+        if identity in identities:
+            continue
+        identities.add(identity)
+        digest.update(identity.encode())
+        try:
+            digest.update(inspect.getsource(obj).encode())
+        except (OSError, TypeError):
+            pass
+    return digest.hexdigest()
+
+
+def tensor_metadata(record: RuntimeTensor) -> dict[str, JsonValue]:
+    """Return stable manifest metadata for one runtime tensor."""
+    tensor = record.tensor
+    return {
+        "kind": record.kind,
+        "shape": list(tensor.shape),
+        "dtype": str(tensor.dtype),
+        "nbytes": tensor.numel() * tensor.element_size(),
+        "layout": "contiguous",
+    }
+
+
+def _update_hash_with_tensor(digest: _HashUpdater, tensor: torch.Tensor) -> None:
+    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
+    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
+        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
+        digest.update(memoryview(chunk.numpy()))
+
+
+def runtime_content_digest(records: Sequence[RuntimeTensor]) -> str:
+    """Hash ordered runtime tensor metadata and contents."""
+    digest = hashlib.sha256()
+    for record in records:
+        digest.update(canonical_json({"name": record.name, **tensor_metadata(record)}))
+        _update_hash_with_tensor(digest, record.tensor)
+    return digest.hexdigest()
+
+
+def file_digest(path: Path) -> str:
+    """Hash one cache file with a bounded read working set."""
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@contextlib.contextmanager
+def exclusive_lock(path: Path, timeout_seconds: float) -> Iterator[None]:
+    """Acquire one non-blocking file lock within a bounded wait."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    deadline = time.monotonic() + timeout_seconds
+    with path.open("a+b") as handle:
+        while True:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                break
+            except OSError as exc:
+                if exc.errno not in (errno.EACCES, errno.EAGAIN):
+                    raise RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise RuntimeCacheError(
+                        "lock_timeout",
+                        f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}",
+                    ) from exc
+                time.sleep(min(0.2, remaining))
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def fsync_file(path: Path) -> None:
+    """Persist one cache file before publication."""
+    with path.open("rb") as handle:
+        os.fsync(handle.fileno())
+
+
+def fsync_directory(path: Path) -> None:
+    """Persist directory entries needed for atomic publication."""
+    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def split_shards(
+    records: Sequence[_RuntimeTensorT],
+    max_shard_bytes: int,
+) -> list[list[_RuntimeTensorT]]:
+    """Split ordered tensors without exceeding the target when possible."""
+    shards: list[list[_RuntimeTensorT]] = []
+    current: list[_RuntimeTensorT] = []
+    current_bytes = 0
+    for record in records:
+        nbytes = record.tensor.numel() * record.tensor.element_size()
+        if current and current_bytes + nbytes > max_shard_bytes:
+            shards.append(current)
+            current = []
+            current_bytes = 0
+        current.append(record)
+        current_bytes += nbytes
+    if current:
+        shards.append(current)
+    return shards
+
+
+def remove_stale_temps(cache_root: Path, cache_key: str) -> None:
+    """Remove abandoned temporary entries for one cache identity."""
+    temp_root = cache_root / ".tmp"
+    if not temp_root.is_dir():
+        return
+    for path in temp_root.glob(f"{cache_key}.*"):
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+
+
 __all__ = [
     "IdentityNormalizationError",
     "JsonValue",
+    "RuntimeCacheError",
+    "RuntimeTensor",
     "canonical_json",
     "canonicalize_existing_local_path",
+    "exclusive_lock",
+    "file_digest",
+    "fsync_directory",
+    "fsync_file",
+    "implementation_fingerprint",
     "normalize_identity",
+    "remove_stale_temps",
+    "runtime_content_digest",
+    "split_shards",
+    "tensor_metadata",
 ]

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -1,14 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Node-local, content-addressed runtime-weight cache for diffusion DLO."""
+"""Build node-local runtime-weight cache plans for diffusion DLO.
+
+``build_runtime_weight_cache_plan`` is the module's single entry point. It
+validates final ordinary-loader DiT tensors, joins or atomically publishes an
+immutable content-addressed entry, validates the mmap result, and returns a
+fail-closed host-weight plan. Reusable identity, hashing, locking, and file
+primitives live in :mod:`runtime_cache_utils`.
+"""
 
 from __future__ import annotations
 
 import contextlib
-import errno
-import fcntl
 import hashlib
-import inspect
 import json
 import math
 import os
@@ -16,7 +20,6 @@ import shutil
 import time
 import uuid
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
 from itertools import chain
 from pathlib import Path
 from typing import Any
@@ -34,9 +37,20 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
 )
 from vllm_omni.diffusion.model_loader.runtime_cache_utils import (
     IdentityNormalizationError,
+    RuntimeCacheError,
+    RuntimeTensor,
     canonical_json,
     canonicalize_existing_local_path,
+    exclusive_lock,
+    file_digest,
+    fsync_directory,
+    fsync_file,
+    implementation_fingerprint,
     normalize_identity,
+    remove_stale_temps,
+    runtime_content_digest,
+    split_shards,
+    tensor_metadata,
 )
 
 logger = init_logger(__name__)
@@ -51,8 +65,6 @@ DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights"
 DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0
 # Avoid one model-sized output file and bound hashing/read working sets.
 DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
-_HASH_CHUNK_BYTES = 64 * 1024**2
-_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
 
 _SUPPORTED_DTYPES = {
     torch.bool,
@@ -71,19 +83,6 @@ for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
         _SUPPORTED_DTYPES.add(_dtype)
 
 
-@dataclass(frozen=True)
-class _RuntimeTensor:
-    name: str
-    tensor: torch.Tensor
-    kind: str
-
-
-class _RuntimeCacheError(RuntimeError):
-    def __init__(self, code: str, message: str):
-        super().__init__(message)
-        self.code = code
-
-
 def default_runtime_cache_root() -> str:
     """Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
     return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT)
@@ -121,7 +120,7 @@ def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Te
 def _collect_runtime_tensors(
     pipeline: nn.Module,
     dit_modules: Sequence[tuple[str, nn.Module]],
-) -> list[_RuntimeTensor]:
+) -> list[RuntimeTensor]:
     """Collect a complete, safely remappable view of final CPU DiT tensors.
 
     Parameters and persistent buffers must resolve to the same objects through
@@ -129,14 +128,14 @@ def _collect_runtime_tensors(
     dtypes, and have no aliases inside or outside the discovered DiT modules.
     The stable name ordering is reused by content hashing and publication.
     """
-    records: dict[str, _RuntimeTensor] = {}
+    records: dict[str, RuntimeTensor] = {}
     for dit_name, dit_module in dit_modules:
         for local_name, tensor in _named_parameters_with_duplicates(dit_module):
             runtime_name = f"{dit_name}.{local_name}"
-            candidate = _RuntimeTensor(runtime_name, tensor, "parameter")
+            candidate = RuntimeTensor(runtime_name, tensor, "parameter")
             existing = records.get(runtime_name)
             if existing is not None and existing.tensor is not tensor:
-                raise _RuntimeCacheError(
+                raise RuntimeCacheError(
                     "ambiguous_ownership",
                     f"multiple DiT tensors resolve to {runtime_name!r}",
                 )
@@ -145,17 +144,17 @@ def _collect_runtime_tensors(
             if not _persistent_buffer(dit_module, local_name):
                 continue
             runtime_name = f"{dit_name}.{local_name}"
-            candidate = _RuntimeTensor(runtime_name, tensor, "buffer")
+            candidate = RuntimeTensor(runtime_name, tensor, "buffer")
             existing = records.get(runtime_name)
             if existing is not None and existing.tensor is not tensor:
-                raise _RuntimeCacheError(
+                raise RuntimeCacheError(
                     "ambiguous_ownership",
                     f"multiple DiT tensors resolve to {runtime_name!r}",
                 )
             records[runtime_name] = candidate
 
     if not records:
-        raise _RuntimeCacheError(
+        raise RuntimeCacheError(
             "ambiguous_ownership",
             "no DiT parameters or persistent buffers were discovered",
         )
@@ -166,37 +165,37 @@ def _collect_runtime_tensors(
         try:
             pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
         except AttributeError as exc:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "ambiguous_ownership",
                 f"DiT tensor {record.name!r} is not owned by the pipeline",
             ) from exc
         if pipeline_tensor is not tensor:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "ambiguous_ownership",
                 f"DiT tensor {record.name!r} does not resolve to the discovered object",
             )
         if tensor.device.type != "cpu" or tensor.is_meta:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
             )
         if tensor.layout != torch.strided:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} uses unsupported layout {tensor.layout}",
             )
         if not tensor.is_contiguous():
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
             )
         if tensor.dtype not in _SUPPORTED_DTYPES:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} uses unsupported dtype {tensor.dtype}",
             )
         if hasattr(tensor, "to_local"):
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} is a distributed tensor",
             )
@@ -204,7 +203,7 @@ def _collect_runtime_tensors(
         tensor_nbytes = tensor.numel() * tensor.element_size()
         storage = tensor.untyped_storage()
         if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_tensor",
                 f"{record.name!r} is a view into a larger storage",
             )
@@ -212,7 +211,7 @@ def _collect_runtime_tensors(
             storage_id = (storage.data_ptr(), storage.nbytes())
             owner = storage_owners.setdefault(storage_id, record.name)
             if owner != record.name:
-                raise _RuntimeCacheError(
+                raise RuntimeCacheError(
                     "unsupported_alias",
                     f"{record.name!r} shares storage with {owner!r}",
                 )
@@ -229,7 +228,7 @@ def _collect_runtime_tensors(
         storage = tensor.untyped_storage()
         owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
         if owner is not None and owner != pipeline_name:
-            raise _RuntimeCacheError(
+            raise RuntimeCacheError(
                 "unsupported_alias",
                 f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
             )
@@ -237,123 +236,6 @@ def _collect_runtime_tensors(
     return [records[name] for name in sorted(records)]
 
 
-def _implementation_fingerprint(
-    loader_type: type, pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]]
-) -> str:
-    objects: list[Any] = [loader_type, type(pipeline)]
-    load_weights = getattr(type(pipeline), "load_weights", None)
-    if load_weights is not None:
-        objects.append(load_weights)
-    for _, dit_module in dit_modules:
-        objects.append(type(dit_module))
-        post_load = getattr(type(dit_module), "post_load_weights", None)
-        if post_load is not None:
-            objects.append(post_load)
-
-    digest = hashlib.sha256()
-    identities: set[str] = set()
-    for obj in objects:
-        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
-        if identity in identities:
-            continue
-        identities.add(identity)
-        digest.update(identity.encode())
-        try:
-            digest.update(inspect.getsource(obj).encode())
-        except (OSError, TypeError):
-            pass
-    return digest.hexdigest()
-
-
-def _tensor_metadata(record: _RuntimeTensor) -> dict[str, Any]:
-    tensor = record.tensor
-    return {
-        "kind": record.kind,
-        "shape": list(tensor.shape),
-        "dtype": str(tensor.dtype),
-        "nbytes": tensor.numel() * tensor.element_size(),
-        "layout": "contiguous",
-    }
-
-
-def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None:
-    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
-    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
-        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
-        digest.update(memoryview(chunk.numpy()))
-
-
-def _runtime_content_digest(records: Sequence[_RuntimeTensor]) -> str:
-    digest = hashlib.sha256()
-    for record in records:
-        digest.update(canonical_json({"name": record.name, **_tensor_metadata(record)}))
-        _update_hash_with_tensor(digest, record.tensor)
-    return digest.hexdigest()
-
-
-def _file_digest(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-@contextlib.contextmanager
-def _exclusive_lock(path: Path, timeout_seconds: float):
-    path.parent.mkdir(parents=True, exist_ok=True)
-    deadline = time.monotonic() + timeout_seconds
-    with path.open("a+b") as handle:
-        while True:
-            try:
-                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
-                break
-            except OSError as exc:
-                if exc.errno not in (errno.EACCES, errno.EAGAIN):
-                    raise _RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
-                remaining = deadline - time.monotonic()
-                if remaining <= 0:
-                    raise _RuntimeCacheError(
-                        "lock_timeout",
-                        f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}",
-                    ) from exc
-                time.sleep(min(0.2, remaining))
-        try:
-            yield
-        finally:
-            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
-
-
-def _fsync_file(path: Path) -> None:
-    with path.open("rb") as handle:
-        os.fsync(handle.fileno())
-
-
-def _fsync_directory(path: Path) -> None:
-    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
-    try:
-        os.fsync(fd)
-    finally:
-        os.close(fd)
-
-
-def _split_shards(records: Sequence[_RuntimeTensor], max_shard_bytes: int) -> list[list[_RuntimeTensor]]:
-    shards: list[list[_RuntimeTensor]] = []
-    current: list[_RuntimeTensor] = []
-    current_bytes = 0
-    for record in records:
-        nbytes = record.tensor.numel() * record.tensor.element_size()
-        if current and current_bytes + nbytes > max_shard_bytes:
-            shards.append(current)
-            current = []
-            current_bytes = 0
-        current.append(record)
-        current_bytes += nbytes
-    if current:
-        shards.append(current)
-    return shards
-
-
 def _publish_entry(
     *,
     cache_root: Path,
@@ -361,14 +243,14 @@ def _publish_entry(
     cache_key: str,
     layout_identity: dict[str, Any],
     content_digest: str,
-    records: Sequence[_RuntimeTensor],
+    records: Sequence[RuntimeTensor],
     max_shard_bytes: int,
 ) -> None:
     total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
     free_bytes = shutil.disk_usage(cache_root).free
     required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
     if free_bytes < required_bytes:
-        raise _RuntimeCacheError(
+        raise RuntimeCacheError(
             "insufficient_disk",
             f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
         )
@@ -378,7 +260,7 @@ def _publish_entry(
     temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
     temp_dir.mkdir()
     try:
-        shards = _split_shards(records, max_shard_bytes)
+        shards = split_shards(records, max_shard_bytes)
         tensor_manifest: dict[str, dict[str, Any]] = {}
         file_manifest: dict[str, dict[str, Any]] = {}
         shard_count = len(shards)
@@ -390,14 +272,14 @@ def _publish_entry(
                 str(path),
                 metadata={"format": "pt", "dlo_runtime_cache_schema": str(RUNTIME_CACHE_SCHEMA_VERSION)},
             )
-            _fsync_file(path)
+            fsync_file(path)
             file_manifest[filename] = {
                 "size": path.stat().st_size,
-                "sha256": _file_digest(path),
+                "sha256": file_digest(path),
             }
             for record in shard:
                 tensor_manifest[record.name] = {
-                    **_tensor_metadata(record),
+                    **tensor_metadata(record),
                     "file": filename,
                     "storage_key": record.name,
                 }
@@ -418,14 +300,14 @@ def _publish_entry(
             json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
             handle.flush()
             os.fsync(handle.fileno())
-        _fsync_directory(temp_dir)
+        fsync_directory(temp_dir)
 
         invalid_dir: Path | None = None
         if entry_dir.exists():
             invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
             os.replace(entry_dir, invalid_dir)
         os.replace(temp_dir, entry_dir)
-        _fsync_directory(entry_dir.parent)
+        fsync_directory(entry_dir.parent)
         if invalid_dir is not None:
             shutil.rmtree(invalid_dir, ignore_errors=True)
     except Exception:
@@ -439,45 +321,45 @@ def _validate_entry(
     cache_key: str,
     layout_identity: dict[str, Any],
     content_digest: str,
-    expected_records: Sequence[_RuntimeTensor],
+    expected_records: Sequence[RuntimeTensor],
 ) -> HostWeightPlan:
     manifest_path = entry_dir / "manifest.json"
     try:
         with manifest_path.open(encoding="utf-8") as handle:
             manifest = json.load(handle)
     except (OSError, ValueError) as exc:
-        raise _RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
+        raise RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
 
     if manifest.get("schema_version") != RUNTIME_CACHE_SCHEMA_VERSION:
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch")
     if manifest.get("cache_key") != cache_key:
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache key mismatch")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache key mismatch")
     if manifest.get("layout_identity") != layout_identity:
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch")
     if manifest.get("runtime_content_sha256") != content_digest:
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch")
 
-    expected_metadata = {record.name: _tensor_metadata(record) for record in expected_records}
+    expected_metadata = {record.name: tensor_metadata(record) for record in expected_records}
     tensor_manifest = manifest.get("tensors")
     file_manifest = manifest.get("files")
     if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict):
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections")
     if set(tensor_manifest) != set(expected_metadata):
-        raise _RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT")
+        raise RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT")
 
     for filename, metadata in file_manifest.items():
         if Path(filename).name != filename or not isinstance(metadata, dict):
-            raise _RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}")
+            raise RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}")
         path = entry_dir / filename
         try:
             stat = path.stat()
         except OSError as exc:
-            raise _RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
-        if stat.st_size != metadata.get("size") or _file_digest(path) != metadata.get("sha256"):
-            raise _RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
+            raise RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
+        if stat.st_size != metadata.get("size") or file_digest(path) != metadata.get("sha256"):
+            raise RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
 
     bindings: dict[str, TensorBinding] = {}
-    mapped_records: list[_RuntimeTensor] = []
+    mapped_records: list[RuntimeTensor] = []
     with contextlib.ExitStack() as stack:
         handles = {
             filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
@@ -487,33 +369,33 @@ def _validate_entry(
         for record in expected_records:
             stored = tensor_manifest.get(record.name)
             if not isinstance(stored, dict):
-                raise _RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
+                raise RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
             filename = stored.get("file")
             storage_key = stored.get("storage_key")
             if filename not in handles or not isinstance(storage_key, str):
-                raise _RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
+                raise RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
             actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]}
             if actual_metadata != expected_metadata[record.name]:
-                raise _RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
+                raise RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
             expected_keys_by_file[filename].add(storage_key)
             try:
                 tensor = handles[filename].get_tensor(storage_key)
             except Exception as exc:
-                raise _RuntimeCacheError(
+                raise RuntimeCacheError(
                     "entry_rejected",
                     f"cannot map {record.name!r} from {filename}: {exc}",
                 ) from exc
             if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype:
-                raise _RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
-            mapped_records.append(_RuntimeTensor(record.name, tensor, record.kind))
+                raise RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
+            mapped_records.append(RuntimeTensor(record.name, tensor, record.kind))
             bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename))
 
         for filename, expected_keys in expected_keys_by_file.items():
             if set(handles[filename].keys()) != expected_keys:
-                raise _RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
+                raise RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
 
-        if _runtime_content_digest(mapped_records) != content_digest:
-            raise _RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
+        if runtime_content_digest(mapped_records) != content_digest:
+            raise RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
 
     return HostWeightPlan(
         backing_kind="runtime_cache",
@@ -523,17 +405,6 @@ def _validate_entry(
     )
 
 
-def _remove_stale_temps(cache_root: Path, cache_key: str) -> None:
-    temp_root = cache_root / ".tmp"
-    if not temp_root.is_dir():
-        return
-    for path in temp_root.glob(f"{cache_key}.*"):
-        if path.is_dir():
-            shutil.rmtree(path, ignore_errors=True)
-        else:
-            path.unlink(missing_ok=True)
-
-
 def build_runtime_weight_cache_plan(
     pipeline: nn.Module,
     *,
@@ -556,7 +427,20 @@ def build_runtime_weight_cache_plan(
     cfg_parallel_size: int,
     pipeline_parallel_size: int,
 ) -> HostWeightPlanResult:
-    """Publish or join one immutable cache entry for final CPU DiT tensors."""
+    """Build the runtime-cache host-weight plan after ordinary model loading.
+
+    This is the module's main entry point. It first rejects unsupported runtime
+    layouts, then discovers and validates the final CPU DiT tensors and derives
+    a content-addressed identity from their metadata, bytes, loader inputs, and
+    parallelism guards. A valid existing entry is joined immediately; otherwise
+    one process publishes the entry under a per-identity file lock and peers
+    validate it after the lock is released.
+
+    Success returns bindings to immutable safetensors mmap files. Any expected
+    incompatibility, lock, publication, or validation failure returns a
+    fail-closed result with no plan and a stable fallback code, leaving the
+    already-loaded pipeline tensors intact.
+    """
     if load_format == "dummy":
         return HostWeightPlanResult(
             None,
@@ -588,7 +472,7 @@ def build_runtime_weight_cache_plan(
 
     try:
         records = _collect_runtime_tensors(pipeline, dit_modules)
-        content_digest = _runtime_content_digest(records)
+        content_digest = runtime_content_digest(records)
         layout_identity = {
             "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
             "model": (canonicalize_existing_local_path(model_identity) if model_identity is not None else None),
@@ -596,7 +480,7 @@ def build_runtime_weight_cache_plan(
             "runtime_dtype": str(runtime_dtype),
             "load_format": load_format,
             "loader_inputs": normalize_identity(loader_inputs),
-            "loader_implementation_sha256": _implementation_fingerprint(loader_type, pipeline, dit_modules),
+            "loader_implementation_sha256": implementation_fingerprint(loader_type, pipeline, dit_modules),
             "components": [name for name, _ in dit_modules],
             "tensor_parallel_size": tensor_parallel_size,
             "tensor_parallel_rank": tensor_parallel_rank,
@@ -620,12 +504,12 @@ def build_runtime_weight_cache_plan(
                 )
                 logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
                 return HostWeightPlanResult(plan)
-            except _RuntimeCacheError as exc:
+            except RuntimeCacheError as exc:
                 logger.warning("DLO runtime cache entry %s will be rebuilt: %s", cache_key[:12], exc)
 
         lock_path = root / ".locks" / f"{cache_key}.lock"
-        with _exclusive_lock(lock_path, lock_timeout_seconds):
-            _remove_stale_temps(root, cache_key)
+        with exclusive_lock(lock_path, lock_timeout_seconds):
+            remove_stale_temps(root, cache_key)
             if entry_dir.is_dir():
                 try:
                     plan = _validate_entry(
@@ -637,7 +521,7 @@ def build_runtime_weight_cache_plan(
                     )
                     logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
                     return HostWeightPlanResult(plan)
-                except _RuntimeCacheError:
+                except RuntimeCacheError:
                     pass
 
             _publish_entry(
@@ -663,7 +547,7 @@ def build_runtime_weight_cache_plan(
                 entry_dir,
             )
             return HostWeightPlanResult(plan)
-    except _RuntimeCacheError as exc:
+    except RuntimeCacheError as exc:
         return HostWeightPlanResult(None, str(exc), exc.code)
     except IdentityNormalizationError as exc:
         return HostWeightPlanResult(None, str(exc), "unstable_identity")