anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Clarify runtime cache module boundary

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-19 11:21:59 +08:00
1 parent 6751c1f404
commit 6b7b5368f2
2 files changed
+279 -207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/model_loader/runtime_weight_cache.py
+88 -204
View File
@@ -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")