235 lines
8.3 KiB
Python
235 lines
8.3 KiB
Python
# 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_enabled() -> None:
|
|
policy = ProductionPolicy(allow_post_load_publish=True)
|
|
|
|
assert policy.allow_local_build
|
|
assert policy.allow_post_load_publish
|
|
|
|
|
|
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 == []
|