vllm-omni/vllm_omni/diffusion/model_loader/host_weights/source_identity.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Canonical diffusion weight-source identity and one-load replacement guard."""

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

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 (
    FinalLayoutContractCode,
    FinalLayoutContractError,
    ImplementationIdentity,
)

_HASH_CHUNK_BYTES = 8 * 1024**2
_IMMUTABLE_REVISION_RE = re.compile(r"(?:[0-9a-fA-F]{40}|[0-9a-fA-F]{64})\Z")


class WeightSourceKind(str, Enum):
    LOCAL_PATH = "local_path"
    HUGGING_FACE_HUB = "hugging_face_hub"


@dataclass(frozen=True)
class PreparedWeightSource:
    """Canonical source files and adapter ABI selected by the ordinary loader."""

    model_or_path: str
    subfolder: str | None
    requested_revision: str | None
    prefix: str
    resolved_root: Path
    weight_files: tuple[Path, ...]
    use_safetensors: bool
    checkpoint_adapter: ImplementationIdentity | None = None
    source_kind: WeightSourceKind = WeightSourceKind.LOCAL_PATH

    def __post_init__(self) -> None:
        if not self.model_or_path or not isinstance(self.prefix, str):
            raise ValueError("prepared source model and prefix must be valid strings")
        if not isinstance(self.resolved_root, Path):
            raise ValueError("prepared source root must use pathlib.Path")
        if not isinstance(self.weight_files, tuple) or not self.weight_files:
            raise ValueError("prepared source must contain an immutable tuple of weight files")
        if any(not isinstance(path, Path) for path in self.weight_files):
            raise ValueError("prepared source weight files must use pathlib.Path")
        if not isinstance(self.use_safetensors, bool):
            raise ValueError("prepared source safetensors mode must be a boolean")
        if self.checkpoint_adapter is not None and not isinstance(
            self.checkpoint_adapter,
            ImplementationIdentity,
        ):
            raise ValueError("checkpoint adapter must use ImplementationIdentity")
        if not isinstance(self.source_kind, WeightSourceKind):
            raise ValueError("prepared source kind must use WeightSourceKind")


@dataclass(frozen=True)
class _FileSnapshot:
    path: Path
    relative_name: str
    size: int
    device: int
    inode: int
    mtime_ns: int
    ctime_ns: int
    symlink_target: str | None
    content_id: str

    def semantic_dict(self) -> dict[str, object]:
        return {
            "relative_name": self.relative_name,
            "size": self.size,
            "content_id": self.content_id,
        }

    def unchanged(self) -> bool:
        try:
            current = self.path.stat()
            symlink_target = os.readlink(self.path) if self.path.is_symlink() else None
        except OSError:
            return False
        return (
            current.st_size == self.size
            and current.st_dev == self.device
            and current.st_ino == self.inode
            and current.st_mtime_ns == self.mtime_ns
            and current.st_ctime_ns == self.ctime_ns
            and symlink_target == self.symlink_target
        )


@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
    revision: str
    files: tuple[_FileSnapshot, ...]

    def unchanged(self) -> bool:
        return all(file.unchanged() for file in self.files)


@dataclass(frozen=True)
class FinalLayoutSourceContext:
    """Resolved source identity plus process-local replacement observations."""

    identity: WeightSourceIdentity
    snapshots: tuple[_SourceSnapshot, ...]

    def sources_unchanged(self) -> bool:
        return all(snapshot.unchanged() for snapshot in self.snapshots)

    def ensure_sources_unchanged(self) -> None:
        if not self.sources_unchanged():
            raise FinalLayoutContractError(
                FinalLayoutContractCode.SOURCE_CHANGED,
                "canonical source changed after final-layout identity resolution",
            )


def _hf_blob_content_id(
    path: Path,
    *,
    source: PreparedWeightSource,
    source_root: Path,
) -> str | None:
    if source.source_kind is not WeightSourceKind.HUGGING_FACE_HUB or not path.is_symlink():
        return None
    if Path(source.model_or_path).expanduser().exists():
        return None

    parts = source_root.parts
    snapshot_indexes = [index for index, part in enumerate(parts[:-1]) if part == "snapshots"]
    if not snapshot_indexes:
        return None
    snapshot_index = snapshot_indexes[-1]
    revision = parts[snapshot_index + 1]
    if _IMMUTABLE_REVISION_RE.fullmatch(revision) is None:
        return None
    repo_root = Path(*parts[:snapshot_index])
    expected_repo_name = f"models--{source.model_or_path.replace('/', '--')}"
    if repo_root.name != expected_repo_name:
        return None
    snapshot_root = repo_root / "snapshots" / revision
    try:
        path.relative_to(snapshot_root)
    except ValueError:
        return None
    blobs_root = repo_root / "blobs"
    try:
        resolved_target = path.resolve(strict=True)
        resolved_blobs_root = blobs_root.resolve(strict=True)
    except OSError:
        return None
    if resolved_target.parent != resolved_blobs_root:
        return None
    blob_name = resolved_target.name
    if _IMMUTABLE_REVISION_RE.fullmatch(blob_name) is None:
        return None
    return f"immutable-blob:{blob_name.lower()}"


def _snapshot_revision(path: Path) -> str | None:
    parts = path.absolute().parts
    for index, part in enumerate(parts[:-1]):
        if part != "snapshots":
            continue
        candidate = parts[index + 1]
        if _IMMUTABLE_REVISION_RE.fullmatch(candidate) is not None:
            return candidate.lower()
    return None


def _logical_model_id(value: str) -> str:
    candidate = Path(value).expanduser()
    try:
        if candidate.exists():
            return str(candidate.resolve())
    except OSError:
        pass
    return value


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}")
    immutable_revision = _snapshot_revision(root)
    if immutable_revision is None and source.requested_revision is not None:
        requested = source.requested_revision.strip()
        if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
            immutable_revision = requested.lower()

    def snapshot_file(path: Path) -> _FileSnapshot:
        candidate = path.expanduser()
        if not candidate.is_absolute():
            candidate = root / candidate
        candidate = candidate.absolute()
        try:
            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
        current = _observe(candidate)
        content_id = _hf_blob_content_id(
            candidate,
            source=source,
            source_root=root,
        )
        if content_id is None:
            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}"
    semantic = CanonicalJson.from_value(
        {
            "checkpoint_adapter": (
                source.checkpoint_adapter.to_dict() if source.checkpoint_adapter is not None else None
            ),
            "files": file_semantics,
            "model_or_path": _logical_model_id(source.model_or_path),
            "prefix": source.prefix,
            "resolved_revision": revision,
            "source_kind": source.source_kind.value,
            "subfolder": source.subfolder,
            "use_safetensors": source.use_safetensors,
        }
    )
    return _SourceSnapshot(semantic=semantic, revision=revision, files=tuple(files))


def _resolve_target_sources(
    prepared_sources: Sequence[PreparedWeightSource],
    target_names: frozenset[str],
) -> dict[str, PreparedWeightSource]:
    """Bind every target to one source using deterministic longest-prefix wins."""
    if not target_names:
        raise FinalLayoutContractError(
            FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
            "final-layout source identity requires at least one owned tensor",
        )

    bindings: dict[str, PreparedWeightSource] = {}
    for target_name in sorted(target_names):
        matches = [source for source in prepared_sources if target_name.startswith(source.prefix)]
        if not matches:
            raise FinalLayoutContractError(
                FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
                f"no canonical weight source covers final-layout tensor {target_name!r}",
            )
        longest_prefix = max(len(source.prefix) for source in matches)
        winners = [source for source in matches if len(source.prefix) == longest_prefix]
        if len(winners) != 1:
            candidates = sorted(
                f"{source.model_or_path}:{source.subfolder or ''}:{source.prefix}" for source in winners
            )
            raise FinalLayoutContractError(
                FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
                f"multiple equally specific canonical sources cover {target_name!r}: {candidates}",
            )
        bindings[target_name] = winners[0]
    return bindings


def resolve_final_layout_source_identity(
    prepared_sources: Sequence[PreparedWeightSource],
    *,
    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, 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 = [
        {
            "source_fingerprint": hashlib.sha256(snapshot_by_source[source].semantic.encoded).hexdigest(),
            "source_prefix": source.prefix,
            "target_name": target_name,
        }
        for target_name, source in sorted(bindings.items())
    ]
    source_document = {
        "sources": source_semantics,
        "target_bindings": target_bindings,
    }
    source_fingerprint = hashlib.sha256(canonical_json(source_document)).hexdigest()
    unique_revisions = tuple(sorted({snapshot.revision for snapshot in snapshots}))
    aggregate_revision = (
        unique_revisions[0]
        if len(unique_revisions) == 1
        else f"aggregate-{hashlib.sha256(canonical_json(unique_revisions)).hexdigest()}"
    )
    return FinalLayoutSourceContext(
        identity=WeightSourceIdentity(
            model_id=_logical_model_id(model_id),
            revision=aggregate_revision,
            fingerprint=source_fingerprint,
            metadata=CanonicalJson.from_value(source_document),
        ),
        snapshots=snapshots,
    )


__all__ = [
    "FinalLayoutSourceContext",
    "NodeSourceDigestCache",
    "PreparedWeightSource",
    "WeightSourceKind",
    "resolve_final_layout_source_identity",
]