[Diffusion] Register HWR mmap for direct DLO H2D (#6591)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

vllm_omni/diffusion/model_loader/host_weights/source_identity.py

@@ -5,9 +5,13 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import os
 import stat
+import time
+import uuid
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -15,6 +19,7 @@ from pathlib import Path
 import regex as re
 
 from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError
 from vllm_omni.host_weight_runtime.identity import canonical_json
 
 from .contracts import (
@@ -101,6 +106,140 @@ class _FileSnapshot:
         )
 
 
+@dataclass(frozen=True)
+class _FileState:
+    size: int
+    device: int
+    inode: int
+    mtime_ns: int
+    ctime_ns: int
+    symlink_target: str | None
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "size": self.size,
+            "device": self.device,
+            "inode": self.inode,
+            "mtime_ns": self.mtime_ns,
+            "ctime_ns": self.ctime_ns,
+            "symlink_target": self.symlink_target,
+        }
+
+
+def _observe(path: Path) -> _FileState:
+    try:
+        current = path.stat()
+    except OSError as exc:
+        raise ValueError(f"cannot stat canonical weight file {path}") from exc
+    if not stat.S_ISREG(current.st_mode):
+        raise ValueError(f"canonical weight source is not a regular file: {path}")
+    return _FileState(
+        size=current.st_size,
+        device=current.st_dev,
+        inode=current.st_ino,
+        mtime_ns=current.st_mtime_ns,
+        ctime_ns=current.st_ctime_ns,
+        symlink_target=os.readlink(path) if path.is_symlink() else None,
+    )
+
+
+def _sha256_file(path: Path, expected: _FileState) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        before = os.fstat(handle.fileno())
+        while chunk := handle.read(_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+        after = os.fstat(handle.fileno())
+    observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns)
+    completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns)
+    expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns)
+    if observed != expected_stat or completed != expected_stat or _observe(path) != expected:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.SOURCE_CHANGED,
+            f"canonical weight file changed while hashing: {path}",
+        )
+    return digest.hexdigest()
+
+
+class NodeSourceDigestCache:
+    """Share content-backed local-source fingerprints across node workers."""
+
+    def __init__(self, root: Path, *, timeout_seconds: float) -> None:
+        self.entries = root / "source-digests-v1" / "entries"
+        self.locks = root / "source-digests-v1" / "locks"
+        self.timeout_seconds = timeout_seconds
+        self._available = True
+        try:
+            self.entries.mkdir(parents=True, exist_ok=True, mode=0o700)
+            self.locks.mkdir(parents=True, exist_ok=True, mode=0o700)
+        except OSError:
+            self._available = False
+
+    @staticmethod
+    def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None:
+        try:
+            cached = json.loads(entry.read_bytes())
+            if not isinstance(cached, dict):
+                return None
+            cached_state = cached["state"]
+            content_id = cached["content_id"]
+            record_sha256 = cached["record_sha256"]
+        except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError):
+            return None
+        record = {
+            "content_id": content_id,
+            "schema_version": cached.get("schema_version"),
+            "state": cached_state,
+        }
+        if (
+            record["schema_version"] != 1
+            or cached_state != state.to_dict()
+            or not isinstance(content_id, str)
+            or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None
+            or not isinstance(record_sha256, str)
+            or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256
+            or _observe(path) != state
+        ):
+            return None
+        return content_id
+
+    def content_id(self, path: Path, state: _FileState) -> str:
+        if not self._available:
+            return f"sha256:{_sha256_file(path, state)}"
+        key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest()
+        entry = self.entries / f"{key}.json"
+        content_id: str | None = None
+        try:
+            with FileLock(
+                self.locks / f"{key}.lock",
+                exclusive=True,
+                deadline=time.monotonic() + self.timeout_seconds,
+            ):
+                if (cached := self._cached_content_id(entry, path, state)) is not None:
+                    return cached
+                content_id = f"sha256:{_sha256_file(path, state)}"
+                record: dict[str, object] = {
+                    "content_id": content_id,
+                    "schema_version": 1,
+                    "state": state.to_dict(),
+                }
+                record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest()
+                temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+                try:
+                    temporary.write_bytes(canonical_json(record))
+                    os.chmod(temporary, 0o600)
+                    os.replace(temporary, entry)
+                finally:
+                    temporary.unlink(missing_ok=True)
+                return content_id
+        except (FileLockTimeoutError, OSError):
+            # Digest sharing is an optimization. Coordination or cache I/O
+            # failures must not make a readable canonical source unavailable.
+            if content_id is not None and _observe(path) == state:
+                return content_id
+            return f"sha256:{_sha256_file(path, state)}"
+
+
 @dataclass(frozen=True)
 class _SourceSnapshot:
     semantic: CanonicalJson
@@ -129,14 +268,6 @@ class FinalLayoutSourceContext:
             )
 
 
-def _sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        while chunk := handle.read(_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
 def _hf_blob_content_id(
     path: Path,
     *,
@@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str:
     return value
 
 
-def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
+def _snapshot_source(
+    source: PreparedWeightSource,
+    digest_cache: NodeSourceDigestCache | None,
+) -> _SourceSnapshot:
     root = source.resolved_root.expanduser().resolve()
     if not root.is_dir():
         raise ValueError(f"canonical source root is not a directory: {root}")
@@ -210,8 +344,7 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
         if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
             immutable_revision = requested.lower()
 
-    files: list[_FileSnapshot] = []
-    for path in sorted(source.weight_files, key=lambda item: str(item)):
+    def snapshot_file(path: Path) -> _FileSnapshot:
         candidate = path.expanduser()
         if not candidate.is_absolute():
             candidate = root / candidate
@@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
             relative_name = candidate.relative_to(root).as_posix()
         except ValueError as exc:
             raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
-        try:
-            current = candidate.stat()
-        except OSError as exc:
-            raise ValueError(f"cannot stat canonical weight file {candidate}") from exc
-        if not stat.S_ISREG(current.st_mode):
-            raise ValueError(f"canonical weight source is not a regular file: {candidate}")
-        symlink_target = os.readlink(candidate) if candidate.is_symlink() else None
+        current = _observe(candidate)
         content_id = _hf_blob_content_id(
             candidate,
             source=source,
             source_root=root,
         )
         if content_id is None:
-            content_id = f"sha256:{_sha256_file(candidate)}"
-        files.append(
-            _FileSnapshot(
-                path=candidate,
-                relative_name=relative_name,
-                size=current.st_size,
-                device=current.st_dev,
-                inode=current.st_ino,
-                mtime_ns=current.st_mtime_ns,
-                ctime_ns=current.st_ctime_ns,
-                symlink_target=symlink_target,
-                content_id=content_id,
+            content_id = (
+                digest_cache.content_id(candidate, current)
+                if digest_cache is not None
+                else f"sha256:{_sha256_file(candidate, current)}"
             )
+        return _FileSnapshot(
+            path=candidate,
+            relative_name=relative_name,
+            size=current.size,
+            device=current.device,
+            inode=current.inode,
+            mtime_ns=current.mtime_ns,
+            ctime_ns=current.ctime_ns,
+            symlink_target=current.symlink_target,
+            content_id=content_id,
         )
 
+    paths = sorted(source.weight_files, key=lambda item: str(item))
+    with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor:
+        files = list(executor.map(snapshot_file, paths))
+
     file_semantics = [file.semantic_dict() for file in files]
     file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
     revision = immutable_revision or f"content-{file_fingerprint}"
@@ -306,11 +439,12 @@ def resolve_final_layout_source_identity(
     *,
     model_id: str,
     target_names: frozenset[str],
+    digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutSourceContext:
     """Resolve exact source identity for files covering the owned tensors."""
     bindings = _resolve_target_sources(prepared_sources, target_names)
     selected_sources = frozenset(bindings.values())
-    snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources}
+    snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources}
     snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
     source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
     target_bindings = [
@@ -345,6 +479,7 @@ def resolve_final_layout_source_identity(
 
 __all__ = [
     "FinalLayoutSourceContext",
+    "NodeSourceDigestCache",
     "PreparedWeightSource",
     "WeightSourceKind",
     "resolve_final_layout_source_identity",