diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index ecd873434..08d4f58e5 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -86,7 +86,9 @@ flowchart TD P -->|"nonretryable failure"| E C --> W{"Post-load publication enabled?"} W -->|"yes"| B2["Publish final model through POST_LOAD_ONLY producer"] - B2 --> R2["Return separate report and optional lease"] + B2 --> CL["Close validated publication lease"] + CL --> R2["Return separate publication report"] + R2 --> D W -->|"no"| D["Keep canonical model"] ``` @@ -144,12 +146,8 @@ sequenceDiagram L->>R: publish_after_load(identity, POST_LOAD_ONLY producer) R->>S: get_or_build(identity, producer) S-->>R: validated HostWeightLease or typed failure - R-->>L: separate publication report and optional lease - opt consume mmap in current startup - L->>X: plan_restore(model, lease) - X-->>L: validation-only restore plan - L->>X: commit() once - end + R->>R: close publication lease + R-->>L: separate publication report end else required or nonretryable failure R-->>L: FAILED report @@ -167,9 +165,9 @@ must be discarded and canonical fallback must construct a fresh model. `publish_after_load()` is synchronous in V1 and accepts only a `POST_LOAD_ONLY` producer. Its report is separate from the terminal resolution, so a publication failure cannot rewrite a successful canonical fallback. A -successful publication transfers a lease to the caller. The loader may close -that lease after warming future startups, or restore from it before the model -enters service to replace private host allocations with shared mmap backing. +successful publication closes the store-returned lease inside the runtime and +warms only future startups. It does not restore, rebind, or otherwise mutate the +canonically loaded model serving the current startup. `allow_local_build` gates producers during pre-load resolution, while `allow_post_load_publish` independently gates this explicit post-load path. diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index d2c4883f3..7268d6a3a 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -150,11 +150,10 @@ only after the canonical loader has established the current model. Post-load publication is synchronous in V1, but its outcome is independent of the already completed resolution. A store failure is returned in a separate publication report and cannot turn a valid canonical model into a failed -resolution. On success, ownership of the validated lease transfers to the -caller. The caller may close it after warming the store, or use a restorer to -rebind a model that has not entered service so the current process can also use -the shared mmap backing. Model mutation and lease teardown remain loader and -transport responsibilities. +resolution. On success, the runtime closes the validated store lease before +returning the report. Post-load publication only warms future startups; it does +not restore, rebind, or otherwise mutate the canonically loaded model serving +the current startup. ## Restoration transaction boundary diff --git a/tests/host_weight_runtime/test_runtime.py b/tests/host_weight_runtime/test_runtime_resolution.py similarity index 75% rename from tests/host_weight_runtime/test_runtime.py rename to tests/host_weight_runtime/test_runtime_resolution.py index be068ebcb..770a35bfa 100644 --- a/tests/host_weight_runtime/test_runtime.py +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -1,10 +1,11 @@ # 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.""" +"""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 @@ -15,6 +16,7 @@ from vllm_omni.host_weight_runtime import ( ComponentIdentity, FailureCode, HostWeightFailure, + HostWeightLease, HostWeightRuntime, HostWeightRuntimeConfig, IntegrityPolicy, @@ -42,7 +44,7 @@ from vllm_omni.host_weight_runtime import ( WeightRepresentation, WeightSourceIdentity, ) -from vllm_omni.host_weight_runtime.filesystem import detect_storage_class +from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -227,7 +229,9 @@ def test_preload_resolution_does_not_run_postload_only_producer(tmp_path: Path) assert resolution.report.attempts[0].action.value == "canonical_fallback" -def test_postload_publication_returns_lease_without_changing_preload_resolution(tmp_path: Path) -> None: +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) @@ -248,26 +252,24 @@ def test_postload_publication_returns_lease_without_changing_preload_resolution( assert producer.calls == 0 publication = runtime.publish_after_load(identity, producer=producer) - assert publication.report.outcome is PostLoadPublicationOutcome.PUBLISHED - assert publication.lease is not None - assert torch.equal( - publication.lease.tensors["weight"], - torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2), - ) - assert publication.report.identity_digest == identity.key - publication.lease.close() + assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED + assert publication.identity_digest == identity.key already_present = runtime.publish_after_load(identity, producer=producer) - assert already_present.report.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT - assert already_present.lease is not None - already_present.lease.close() + 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.report, already_present.report] + 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: @@ -279,11 +281,46 @@ def test_postload_publication_policy_skip_does_not_run_producer(tmp_path: Path) publication = runtime.publish_after_load(identity, producer=producer) - assert publication.report.outcome is PostLoadPublicationOutcome.POLICY_DISABLED - assert publication.lease is None + 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( @@ -296,8 +333,7 @@ def test_disabled_postload_publication_does_not_probe_configured_root(tmp_path: publication = runtime.publish_after_load() - assert publication.report.outcome is PostLoadPublicationOutcome.RUNTIME_DISABLED - assert publication.lease is None + assert publication.outcome is PostLoadPublicationOutcome.RUNTIME_DISABLED assert not root.exists() @@ -343,12 +379,96 @@ def test_postload_publication_failure_is_reported_without_raising( monkeypatch.setattr(runtime.store, "get_or_build", fail_publication) publication = runtime.publish_after_load(identity, producer=producer) - assert publication.report.outcome is PostLoadPublicationOutcome.FAILED - assert publication.report.failure is failure - assert publication.lease is None + 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"), [ diff --git a/vllm_omni/host_weight_runtime/__init__.py b/vllm_omni/host_weight_runtime/__init__.py index 37059d8f9..d870188ef 100644 --- a/vllm_omni/host_weight_runtime/__init__.py +++ b/vllm_omni/host_weight_runtime/__init__.py @@ -46,7 +46,6 @@ from .manifest import ( ) from .outcomes import ( AttemptResult, - HostWeightPublication, HostWeightResolution, PostLoadPublicationOutcome, PostLoadPublicationReport, @@ -93,7 +92,6 @@ __all__ = [ "HostWeightError", "HostWeightFailure", "HostWeightLease", - "HostWeightPublication", "HostWeightResolution", "HostWeightInspector", "HostWeightRuntime", diff --git a/vllm_omni/host_weight_runtime/outcomes.py b/vllm_omni/host_weight_runtime/outcomes.py index adb80f386..28711f636 100644 --- a/vllm_omni/host_weight_runtime/outcomes.py +++ b/vllm_omni/host_weight_runtime/outcomes.py @@ -146,26 +146,8 @@ class PostLoadPublicationReport: raise ValueError("post-load publication failure ownership is inconsistent") -@dataclass(frozen=True) -class HostWeightPublication: - """Separate post-load outcome whose successful lease belongs to the caller.""" - - report: PostLoadPublicationReport - lease: HostWeightLease | None = None - - def __post_init__(self) -> None: - lease_outcomes = { - PostLoadPublicationOutcome.ALREADY_PRESENT, - PostLoadPublicationOutcome.PUBLISHED, - PostLoadPublicationOutcome.JOINED, - } - if (self.report.outcome in lease_outcomes) != (self.lease is not None): - raise ValueError("post-load publication outcome has inconsistent lease ownership") - - __all__ = [ "AttemptResult", - "HostWeightPublication", "HostWeightResolution", "PostLoadPublicationOutcome", "PostLoadPublicationReport", diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py index 7cea035f2..fbc726669 100644 --- a/vllm_omni/host_weight_runtime/runtime.py +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -15,7 +15,6 @@ from .identity import CanonicalJson, WeightArtifactIdentity from .lease import HostWeightLease from .outcomes import ( AttemptResult, - HostWeightPublication, HostWeightResolution, PostLoadPublicationOutcome, PostLoadPublicationReport, @@ -241,12 +240,12 @@ class HostWeightRuntime: identity: WeightArtifactIdentity | None = None, *, producer: WeightProducer | None = None, - ) -> HostWeightPublication: - """Synchronously publish a post-load artifact without revising resolution. + ) -> PostLoadPublicationReport: + """Synchronously warm a future startup without revising resolution. - A successful result transfers one validated lease to the caller. The - loader may use that lease to rebind a model that has not entered service, - or close it when publication is only warming a future startup. + A successful store operation returns a validated lease. The runtime + closes that lease before returning because POST_LOAD_ONLY production + must not mutate or rebind the model serving the current startup. """ started = time.monotonic() operation_id = uuid.uuid4().hex @@ -302,14 +301,36 @@ class HostWeightRuntime: outcome = successful_outcomes.get(result.status) if outcome is not None: assert result.lease is not None + try: + result.lease.close() + except Exception as exc: + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.FAILED, + identity_digest=identity_digest, + elapsed_seconds=time.monotonic() - started, + failure=HostWeightFailure( + stage=ResolutionStage.LIFECYCLE, + code=FailureCode.PUBLICATION_FAILED, + retryable=False, + message="post-load publication lease cleanup failed", + details=CanonicalJson.from_value( + { + "exception_type": type(exc).__name__, + "store_outcome": result.status.value, + } + ), + ), + ) + ) return self._finish_publication( PostLoadPublicationReport( operation_id=operation_id, outcome=outcome, identity_digest=identity_digest, elapsed_seconds=time.monotonic() - started, - ), - lease=result.lease, + ) ) failure = result.failure @@ -439,9 +460,7 @@ class HostWeightRuntime: def _finish_publication( self, report: PostLoadPublicationReport, - *, - lease: HostWeightLease | None = None, - ) -> HostWeightPublication: + ) -> PostLoadPublicationReport: if self._publication_observer is not None: try: self._publication_observer(report) @@ -450,7 +469,7 @@ class HostWeightRuntime: "Host Weight Runtime publication observer failed", extra={"operation_id": report.operation_id}, ) - return HostWeightPublication(report=report, lease=lease) + return report __all__ = ["HostWeightRuntime"]