[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>
This commit is contained in:
1 parent
4c773109b1
commit
339761045f
6 files changed
+185
-69
No files matched your search
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+142
-22
@@ -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"),
|
||||
[
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user