diff --git a/vllm_omni/diffusion/model_loader/runtime_cache_utils.py b/vllm_omni/diffusion/model_loader/runtime_cache_utils.py index 431343137..6859e65e2 100644 --- a/vllm_omni/diffusion/model_loader/runtime_cache_utils.py +++ b/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", ] diff --git a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py index 96dfc4d7f..bde659f4d 100644 --- a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py +++ b/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")