365 lines
12 KiB
Python
365 lines
12 KiB
Python
# 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"
|