[Core] Add Host Weight Runtime foundation (#6419)

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

tests/host_weight_runtime/test_contracts.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Contract tests for semantic identity and immutable artifact schemas."""
+
+from __future__ import annotations
+
+import ast
+import math
+from pathlib import Path
+from typing import get_type_hints
+
+import pytest
+
+from vllm_omni.host_weight_runtime import (
+    AdaptationIdentity,
+    CanonicalJson,
+    ComponentIdentity,
+    FileManifestEntry,
+    IdentityValidationError,
+    ManifestValidationError,
+    ProducerIdentity,
+    ProductionPolicy,
+    PublicationMarker,
+    RuntimeWeightLayout,
+    TensorKind,
+    TensorManifestEntry,
+    WeightArtifactIdentity,
+    WeightManifest,
+    WeightRepresentation,
+    WeightRestorePlan,
+    WeightRestorer,
+    WeightSourceIdentity,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_SHA_A = "a" * 64
+_SHA_B = "b" * 64
+_SHA_C = "c" * 64
+
+
+def _identity(
+    *,
+    tp_size: int = 1,
+    tp_rank: int = 0,
+    sp_size: int = 1,
+    adaptation: AdaptationIdentity | None = None,
+    source_metadata: CanonicalJson | None = None,
+) -> WeightArtifactIdentity:
+    return WeightArtifactIdentity(
+        schema_version=1,
+        source=WeightSourceIdentity(
+            model_id="test-org/test-model",
+            revision="0123456789abcdef",
+            fingerprint="source-files-sha256",
+            metadata=source_metadata or CanonicalJson.empty(),
+        ),
+        component=ComponentIdentity(name="transformer", ownership="diffusion.dit"),
+        representation=WeightRepresentation(name="runtime-bf16", dtype="torch.bfloat16"),
+        layout=RuntimeWeightLayout(
+            name="final-module-layout",
+            tensor_parallel_size=tp_size,
+            tensor_parallel_rank=tp_rank,
+            sequence_parallel_size=sp_size,
+            sequence_parallel_backend="ulysses" if sp_size > 1 else "none",
+        ),
+        adaptation=adaptation or AdaptationIdentity(),
+        producer=ProducerIdentity(
+            producer_id="test.final-layout",
+            version="1",
+            implementation_fingerprint="implementation-sha256",
+            manifest_schema="test-producer-v1",
+            restorer_schema="test-restorer-v1",
+        ),
+    )
+
+
+def test_canonical_metadata_is_a_stable_immutable_snapshot() -> None:
+    items: list[object] = [1, {"value": 2}]
+    mutable = {"z": items, "a": True}
+    snapshot = CanonicalJson.from_value(mutable)
+    items.append(3)
+
+    assert snapshot.encoded == b'{"a":true,"z":[1,{"value":2}]}'
+    assert snapshot.to_value() == {"a": True, "z": [1, {"value": 2}]}
+    assert CanonicalJson.from_value({"a": True, "z": [1, {"value": 2}]}) == snapshot
+
+
+@pytest.mark.parametrize("value", [math.inf, math.nan, object(), {1: "non-string-key"}])
+def test_canonical_metadata_rejects_process_specific_or_non_json_values(value: object) -> None:
+    with pytest.raises(IdentityValidationError):
+        CanonicalJson.from_value(value)
+
+
+@pytest.mark.parametrize("encoded", [b"not-json", b'{"b":1, "a":2}', b'{"value":NaN}'])
+def test_canonical_metadata_cannot_bypass_stable_encoding(encoded: bytes) -> None:
+    with pytest.raises(IdentityValidationError):
+        CanonicalJson(encoded)
+
+
+def test_identity_round_trip_and_hash_are_deterministic() -> None:
+    first = _identity(source_metadata=CanonicalJson.from_value({"files": ["b", "a"], "size": 3}))
+    second = _identity(source_metadata=CanonicalJson.from_value({"size": 3, "files": ["b", "a"]}))
+
+    assert first == second
+    assert first.key == second.key
+    assert WeightArtifactIdentity.from_dict(first.to_dict()) == first
+    assert len(first.key) == 64
+
+
+def test_identity_separates_semantic_parallel_layout_and_static_adaptation() -> None:
+    base = _identity()
+    tp0 = _identity(tp_size=2, tp_rank=0)
+    tp1 = _identity(tp_size=2, tp_rank=1)
+    sp2 = _identity(sp_size=2)
+    merged_lora = _identity(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256"))
+
+    assert len({item.key for item in (base, tp0, tp1, sp2, merged_lora)}) == 5
+    assert "data_parallel" not in base.layout.to_dict()
+    assert "sequence_parallel_rank" not in sp2.layout.to_dict()
+
+
+def test_identity_schema_rejects_unknown_fields() -> None:
+    document = _identity().to_dict()
+    document["transfer_mode"] = "registered"
+
+    with pytest.raises(IdentityValidationError, match="unexpected"):
+        WeightArtifactIdentity.from_dict(document)
+
+
+def test_identity_schema_does_not_coerce_field_types() -> None:
+    document = _identity().to_dict()
+    document["schema_version"] = "1"
+
+    with pytest.raises(IdentityValidationError, match="must be an integer"):
+        WeightArtifactIdentity.from_dict(document)
+
+
+def test_restoration_contract_has_a_validation_only_one_shot_commit_boundary() -> None:
+    assert get_type_hints(WeightRestorePlan.commit)["return"] is type(None)
+    assert "one-shot" in (WeightRestorePlan.__doc__ or "")
+    assert "without mutating" in (WeightRestorer.__doc__ or "")
+
+
+def test_post_load_publication_policy_is_explicitly_unsupported() -> None:
+    with pytest.raises(ValueError, match="post-load publication is not implemented"):
+        ProductionPolicy(allow_post_load_publish=True)
+
+
+def test_manifest_and_publication_marker_round_trip_exactly() -> None:
+    identity = _identity()
+    manifest = WeightManifest(
+        schema_version=1,
+        artifact_key=identity.key,
+        identity=identity,
+        producer_schema="test-producer-v1",
+        restorer_schema="test-restorer-v1",
+        tensors=(
+            TensorManifestEntry(
+                name="layer.weight",
+                kind=TensorKind.PARAMETER,
+                role="weight",
+                file_name="weights.safetensors",
+                storage_key="layer.weight",
+                shape=(2, 4),
+                dtype="torch.bfloat16",
+                stride=(4, 1),
+                nbytes=16,
+                sha256=_SHA_A,
+            ),
+        ),
+        files=(FileManifestEntry(file_name="weights.safetensors", size=128, sha256=_SHA_B),),
+        artifact_content_sha256=_SHA_C,
+        format_metadata=CanonicalJson.from_value({"quantization": "none"}),
+    )
+    marker = PublicationMarker(
+        schema_version=1,
+        publication_id="publication-id",
+        manifest_sha256=_SHA_A,
+        artifact_content_sha256=_SHA_C,
+        integrity_mode="full_checksum",
+    )
+
+    assert WeightManifest.from_dict(manifest.to_dict()) == manifest
+    assert PublicationMarker.from_dict(marker.to_dict()) == marker
+
+    document = manifest.to_dict()
+    document["unknown"] = True
+    with pytest.raises(ManifestValidationError, match="unexpected"):
+        WeightManifest.from_dict(document)
+
+    malformed = manifest.to_dict()
+    tensors = malformed["tensors"]
+    assert isinstance(tensors, list) and isinstance(tensors[0], dict)
+    tensors[0]["nbytes"] = "16"
+    with pytest.raises(ManifestValidationError, match="must be an integer"):
+        WeightManifest.from_dict(malformed)
+
+
+@pytest.mark.parametrize("file_name", ["", ".", "..", "../weights.safetensors", "nested/weights.safetensors"])
+def test_manifest_rejects_unsafe_payload_names(file_name: str) -> None:
+    with pytest.raises(ManifestValidationError, match="safe relative basename"):
+        FileManifestEntry(file_name=file_name, size=1, sha256=_SHA_A)
+
+
+def test_package_source_keeps_the_neutral_dependency_boundary() -> None:
+    package_root = Path(__file__).resolve().parents[2] / "vllm_omni" / "host_weight_runtime"
+    assert package_root.is_dir()
+    forbidden = (
+        "vllm_omni.diffusion",
+        "vllm_omni.distributed",
+        "torch.cuda",
+        "torch.distributed",
+    )
+    violations: list[str] = []
+    for source_path in package_root.rglob("*.py"):
+        tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path))
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                violations.extend(
+                    f"{source_path.name}:{node.lineno}:{alias.name}"
+                    for alias in node.names
+                    if alias.name.startswith(forbidden)
+                )
+            elif isinstance(node, ast.ImportFrom) and node.level == 0:
+                module = node.module or ""
+                if module.startswith(forbidden) or (
+                    module == "torch" and any(alias.name in {"cuda", "distributed"} for alias in node.names)
+                ):
+                    violations.append(f"{source_path.name}:{node.lineno}:{module}")
+
+    assert violations == []

tests/host_weight_runtime/test_runtime.py

@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Resolution-policy tests for the loader-adjacent Host Weight Runtime."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from vllm_omni.host_weight_runtime import (
+    AdaptationIdentity,
+    CanonicalJson,
+    ComponentIdentity,
+    FailureCode,
+    HostWeightFailure,
+    HostWeightRuntime,
+    HostWeightRuntimeConfig,
+    IntegrityPolicy,
+    LookupPhase,
+    ProducerIdentity,
+    ProductionMetadata,
+    ProductionSourceMode,
+    RemoteImportPolicy,
+    RemoteOnMiss,
+    ResolutionOutcome,
+    ResolutionReport,
+    ResolutionStage,
+    RuntimeMode,
+    RuntimeWeightLayout,
+    StorageDomainPolicy,
+    StoreResult,
+    StoreStatus,
+    TensorWriteSpec,
+    ValidationLevel,
+    WeightArtifactIdentity,
+    WeightProductionSpec,
+    WeightRepresentation,
+    WeightSourceIdentity,
+)
+from vllm_omni.host_weight_runtime.filesystem import detect_storage_class
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _identity() -> WeightArtifactIdentity:
+    return WeightArtifactIdentity(
+        schema_version=1,
+        source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"),
+        component=ComponentIdentity("transformer", "diffusion.dit"),
+        representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"),
+        layout=RuntimeWeightLayout("final-module-layout"),
+        adaptation=AdaptationIdentity(),
+        producer=ProducerIdentity(
+            "test.final-layout",
+            "1",
+            "implementation-sha256",
+            "test-producer-v1",
+            "test-restorer-v1",
+        ),
+    )
+
+
+class CountingProducer:
+    def __init__(
+        self,
+        identity: WeightArtifactIdentity,
+        *,
+        lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE,
+    ) -> None:
+        self.calls = 0
+        self._spec = WeightProductionSpec(
+            producer_id="test.final-layout",
+            outputs=(identity,),
+            source_mode=ProductionSourceMode.FINALIZED_MODEL,
+            lookup_phase=lookup_phase,
+        )
+
+    @property
+    def spec(self) -> WeightProductionSpec:
+        return self._spec
+
+    def produce(self, writer: object) -> ProductionMetadata:
+        self.calls += 1
+        spec = TensorWriteSpec("weight", (2, 2), torch.bfloat16)
+        with writer.open_tensor_file("weights.safetensors", (spec,)) as output:  # type: ignore[attr-defined]
+            output.write_tensor("weight", torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))
+        return ProductionMetadata(
+            producer_schema="test-producer-v1",
+            restorer_schema="test-restorer-v1",
+            format_metadata=CanonicalJson.from_value({"source": "test"}),
+        )
+
+
+def _domain(root: Path) -> StorageDomainPolicy:
+    return StorageDomainPolicy(root=root, storage_class=detect_storage_class(root.parent))
+
+
+def test_disabled_mode_does_not_probe_or_create_configured_root(tmp_path: Path) -> None:
+    root = tmp_path / "must-not-exist"
+    runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.DISABLED, domain=_domain(root)))
+
+    resolution = runtime.resolve()
+
+    assert resolution.report.outcome is ResolutionOutcome.CANONICAL_DIRECT
+    assert resolution.report.attempts == ()
+    assert resolution.lease is None
+    assert not root.exists()
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+        (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+    ],
+)
+def test_local_miss_uses_configured_fallback_policy(
+    tmp_path: Path,
+    mode: RuntimeMode,
+    expected: ResolutionOutcome,
+) -> None:
+    runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path / mode.value)))
+
+    resolution = runtime.resolve(_identity())
+
+    assert resolution.report.outcome is expected
+    assert resolution.lease is None
+    assert [attempt.result.value for attempt in resolution.report.attempts] == ["miss", "skipped"]
+    expected_action = "canonical_fallback" if mode is RuntimeMode.PREFERRED else "fail_startup"
+    assert [attempt.action.value for attempt in resolution.report.attempts] == [expected_action, expected_action]
+
+
+def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp_path: Path) -> None:
+    reports: list[ResolutionReport] = []
+    identity = _identity()
+    producer = CountingProducer(identity)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")),
+        observer=reports.append,
+    )
+
+    cold = runtime.resolve(identity, producer=producer)
+    assert cold.report.outcome is ResolutionOutcome.LOCAL_PRODUCTION
+    assert cold.lease is not None
+    assert torch.equal(
+        cold.lease.tensors["weight"],
+        torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2),
+    )
+    cold.lease.close()
+    assert cold.report.attempts[0].action.value == "try_producer"
+
+    warm = runtime.resolve(identity)
+    assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert warm.lease is not None
+    warm.lease.close()
+
+    assert producer.calls == 1
+    assert reports == [cold.report, warm.report]
+    assert all(report.resolution_id for report in reports)
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        (RuntimeMode.PREFERRED, ResolutionOutcome.FAILED),
+        (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+    ],
+)
+def test_remote_required_is_explicitly_unsupported_and_skips_local_production(
+    tmp_path: Path,
+    mode: RuntimeMode,
+    expected: ResolutionOutcome,
+) -> None:
+    identity = _identity()
+    producer = CountingProducer(identity)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(
+            mode=mode,
+            domain=_domain(tmp_path / mode.value),
+            remote=RemoteImportPolicy(
+                on_local_miss=RemoteOnMiss.REQUIRE,
+                providers=("future-mooncake",),
+            ),
+        )
+    )
+
+    resolution = runtime.resolve(identity, producer=producer)
+
+    assert resolution.report.outcome is expected
+    assert producer.calls == 0
+    assert [attempt.stage.value for attempt in resolution.report.attempts] == ["lookup", "remote"]
+    assert resolution.report.attempts[-1].failure is not None
+    assert resolution.report.attempts[-1].failure.code.value == "unsupported"
+    assert resolution.report.attempts[-1].action.value == "fail_startup"
+
+
+def test_observer_failure_cannot_change_resolution_outcome(tmp_path: Path) -> None:
+    def broken_observer(_report: object) -> None:
+        raise RuntimeError("metrics backend unavailable")
+
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")),
+        observer=broken_observer,
+    )
+
+    resolution = runtime.resolve(_identity())
+
+    assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+
+def test_preload_resolution_does_not_run_postload_only_producer(tmp_path: Path) -> None:
+    identity = _identity()
+    producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store"))
+    )
+
+    resolution = runtime.resolve(identity, producer=producer)
+
+    assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+    assert producer.calls == 0
+    assert resolution.report.attempts[0].action.value == "canonical_fallback"
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [
+        (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+        (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+    ],
+)
+def test_retryable_store_initialization_failure_obeys_runtime_mode(
+    tmp_path: Path,
+    mode: RuntimeMode,
+    expected: ResolutionOutcome,
+) -> None:
+    root = tmp_path / "not-a-directory"
+    root.write_text("occupied", encoding="utf-8")
+    runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(root)))
+
+    resolution = runtime.resolve(_identity())
+
+    assert resolution.report.outcome is expected
+    assert resolution.lease is None
+    assert len(resolution.report.attempts) == 1
+    assert resolution.report.attempts[0].failure is not None
+    assert resolution.report.attempts[0].failure.code.value == "domain_unavailable"
+    assert resolution.report.attempts[0].failure.retryable
+
+
+@pytest.mark.parametrize(
+    ("stage", "code"),
+    [
+        (ResolutionStage.PRODUCTION, FailureCode.PRODUCER_UNSUPPORTED),
+        (ResolutionStage.VALIDATION, FailureCode.IDENTITY_COLLISION),
+        (ResolutionStage.LIFECYCLE, FailureCode.PUBLICATION_FAILED),
+    ],
+)
+def test_nonretryable_post_initialization_failure_cannot_fall_back_in_preferred_mode(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    stage: ResolutionStage,
+    code: FailureCode,
+) -> None:
+    identity = _identity()
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / code.value))
+    )
+    assert runtime.store is not None
+
+    def fail_production(
+        _request: object,
+        _producer: object,
+        *,
+        validation: object,
+        deadline: float,
+    ) -> StoreResult:
+        del validation, deadline
+        return StoreResult(
+            StoreStatus.FAILED,
+            failure=HostWeightFailure(
+                stage=stage,
+                code=code,
+                retryable=False,
+                message="injected nonretryable production failure",
+                details=CanonicalJson.empty(),
+            ),
+        )
+
+    monkeypatch.setattr(runtime.store, "get_or_build", fail_production)
+    resolution = runtime.resolve(identity, producer=CountingProducer(identity))
+
+    assert resolution.report.outcome is ResolutionOutcome.FAILED
+    assert resolution.report.attempts[-1].action.value == "fail_startup"
+    assert resolution.report.attempts[-1].failure is not None
+    assert resolution.report.attempts[-1].failure.code is code
+
+
+def test_retryable_post_initialization_failure_can_fall_back_in_preferred_mode(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    identity = _identity()
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "retryable"))
+    )
+    assert runtime.store is not None
+
+    def fail_production(
+        _request: object,
+        _producer: object,
+        *,
+        validation: object,
+        deadline: float,
+    ) -> StoreResult:
+        del validation, deadline
+        return StoreResult(
+            StoreStatus.FAILED,
+            failure=HostWeightFailure(
+                stage=ResolutionStage.CAPACITY,
+                code=FailureCode.ENOSPC,
+                retryable=True,
+                message="injected recoverable capacity failure",
+                details=CanonicalJson.empty(),
+            ),
+        )
+
+    monkeypatch.setattr(runtime.store, "get_or_build", fail_production)
+    resolution = runtime.resolve(identity, producer=CountingProducer(identity))
+
+    assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+    assert resolution.report.attempts[-1].action.value == "canonical_fallback"
+
+
+def test_nonretryable_store_configuration_failure_remains_visible_in_preferred_mode(tmp_path: Path) -> None:
+    root = tmp_path / "untrusted"
+    root.mkdir(mode=0o700)
+    root.chmod(0o777)
+    runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(root)))
+
+    resolution = runtime.resolve(_identity())
+
+    assert resolution.report.outcome is ResolutionOutcome.FAILED
+    assert resolution.report.attempts[0].action.value == "fail_startup"
+    assert resolution.report.attempts[0].failure is not None
+    assert not resolution.report.attempts[0].failure.retryable
+
+
+def test_unsupported_integrity_policy_fails_instead_of_disguising_configuration_error(tmp_path: Path) -> None:
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(
+            mode=RuntimeMode.PREFERRED,
+            domain=_domain(tmp_path / "store"),
+            integrity=IntegrityPolicy(local_lookup=ValidationLevel.FS_VERITY),
+        )
+    )
+
+    resolution = runtime.resolve(_identity())
+
+    assert resolution.report.outcome is ResolutionOutcome.FAILED
+    assert resolution.report.attempts[0].failure is not None
+    assert resolution.report.attempts[0].failure.code.value == "unsupported"