Files
vllm-omni/tests/host_weight_runtime/test_runtime_resolution.py
T
Hongsheng Liu 339761045f [Core] Make post-load publication warm-only
Close publication leases inside Host Weight Runtime, return only the publication report, and cover the new terminal mappings. Rename the runtime test module to avoid the pytest basename collision reported in #6436.

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-21 12:18:12 +08:00

610 lines
21 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Runtime resolution and publication-policy tests for Host Weight Runtime."""
from __future__ import annotations
from pathlib import Path
from typing import cast
import pytest
import torch
from vllm_omni.host_weight_runtime import (
AdaptationIdentity,
CanonicalJson,
ComponentIdentity,
FailureCode,
HostWeightFailure,
HostWeightLease,
HostWeightRuntime,
HostWeightRuntimeConfig,
IntegrityPolicy,
LookupPhase,
PostLoadPublicationOutcome,
PostLoadPublicationReport,
ProducerIdentity,
ProductionMetadata,
ProductionPolicy,
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 FilesystemHostWeightStore, 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"
def test_postload_publication_closes_lease_and_only_warms_future_resolution(
tmp_path: Path,
) -> None:
publications: list[PostLoadPublicationReport] = []
identity = _identity()
producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
runtime = HostWeightRuntime.from_config(
HostWeightRuntimeConfig(
mode=RuntimeMode.PREFERRED,
domain=_domain(tmp_path / "store"),
production=ProductionPolicy(
allow_local_build=False,
allow_post_load_publish=True,
),
),
publication_observer=publications.append,
)
preload = runtime.resolve(identity, producer=producer)
assert preload.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
assert producer.calls == 0
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
assert publication.identity_digest == identity.key
already_present = runtime.publish_after_load(identity, producer=producer)
assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT
warm = runtime.resolve(identity)
assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT
assert warm.lease is not None
assert torch.equal(
warm.lease.tensors["weight"],
torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2),
)
warm.lease.close()
assert producer.calls == 1
assert publications == [publication, already_present]
assert isinstance(runtime.store, FilesystemHostWeightStore)
assert runtime.store.cleanup(identity) is None
def test_postload_publication_policy_skip_does_not_run_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"))
)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.POLICY_DISABLED
assert producer.calls == 0
def test_postload_publication_reports_lease_cleanup_failure(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> 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"),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
assert runtime.store is not None
class FailingLease:
def close(self) -> None:
raise RuntimeError("injected lease cleanup failure")
def return_failing_lease(*_args: object, **_kwargs: object) -> StoreResult:
return StoreResult(StoreStatus.BUILT, lease=cast(HostWeightLease, FailingLease()))
monkeypatch.setattr(runtime.store, "get_or_build", return_failing_lease)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.FAILED
assert publication.failure is not None
assert publication.failure.stage is ResolutionStage.LIFECYCLE
assert publication.failure.code is FailureCode.PUBLICATION_FAILED
assert publication.failure.details.to_value() == {
"exception_type": "RuntimeError",
"store_outcome": "built",
}
def test_disabled_postload_publication_does_not_probe_configured_root(tmp_path: Path) -> None:
root = tmp_path / "must-not-exist"
runtime = HostWeightRuntime.from_config(
HostWeightRuntimeConfig(
mode=RuntimeMode.DISABLED,
domain=_domain(root),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
publication = runtime.publish_after_load()
assert publication.outcome is PostLoadPublicationOutcome.RUNTIME_DISABLED
assert not root.exists()
def test_postload_publication_rejects_preload_producer(tmp_path: Path) -> None:
identity = _identity()
runtime = HostWeightRuntime.from_config(
HostWeightRuntimeConfig(
mode=RuntimeMode.PREFERRED,
domain=_domain(tmp_path / "store"),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
with pytest.raises(ValueError, match="POST_LOAD_ONLY"):
runtime.publish_after_load(identity, producer=CountingProducer(identity))
def test_postload_publication_failure_is_reported_without_raising(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> 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"),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
assert runtime.store is not None
failure = HostWeightFailure(
stage=ResolutionStage.CAPACITY,
code=FailureCode.ENOSPC,
retryable=True,
message="injected post-load publication failure",
details=CanonicalJson.empty(),
)
def fail_publication(*_args: object, **_kwargs: object) -> StoreResult:
return StoreResult(StoreStatus.FAILED, failure=failure)
monkeypatch.setattr(runtime.store, "get_or_build", fail_publication)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.FAILED
assert publication.failure is failure
assert producer.calls == 0
def test_postload_publication_reports_store_initialization_failure(tmp_path: Path) -> None:
root = tmp_path / "not-a-directory"
root.write_text("occupied", encoding="utf-8")
identity = _identity()
producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
runtime = HostWeightRuntime.from_config(
HostWeightRuntimeConfig(
mode=RuntimeMode.PREFERRED,
domain=_domain(root),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.FAILED
assert publication.failure is not None
assert publication.failure.code is FailureCode.DOMAIN_UNAVAILABLE
assert publication.failure.retryable
assert producer.calls == 0
def test_postload_publication_maps_joined_and_closes_lease(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> 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"),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
assert runtime.store is not None
class ClosingLease:
closed = False
def close(self) -> None:
self.closed = True
lease = ClosingLease()
def return_joined_lease(*_args: object, **_kwargs: object) -> StoreResult:
return StoreResult(StoreStatus.JOINED, lease=cast(HostWeightLease, lease))
monkeypatch.setattr(runtime.store, "get_or_build", return_joined_lease)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.JOINED
assert lease.closed
def test_postload_publication_synthesizes_failure_for_unexpected_store_status(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> 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"),
production=ProductionPolicy(allow_post_load_publish=True),
)
)
assert runtime.store is not None
def return_unexpected_miss(*_args: object, **_kwargs: object) -> StoreResult:
return StoreResult(StoreStatus.MISS)
monkeypatch.setattr(runtime.store, "get_or_build", return_unexpected_miss)
publication = runtime.publish_after_load(identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.FAILED
assert publication.failure is not None
assert publication.failure.stage is ResolutionStage.PRODUCTION
assert publication.failure.code is FailureCode.PRODUCER_FAILED
assert not publication.failure.retryable
@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"