Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 90815cc866 [Core] Add opt-in serialized HWR build admission
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-06 16:44:52 +08:00
19 changed files with 399 additions and 148 deletions

No files matched your search

@@ -313,6 +313,16 @@ One process per exact identity owns a build; other workers wait and then acquire
leases for the published artifact. Publication is invisible until all payloads
and metadata are validated, hashed, fsynced, and atomically renamed.
By default, different identities may build concurrently. A domain configured
with `CapacityPolicy(build_admission="serialized", max_store_bytes=...)` admits
only one producer/publication operation at a time. Warm hits remain concurrent,
and admission uses the existing coordination deadline. All workers must agree
on the serialized domain's schema-2 capacity policy; use a new root when opting
in, since in-place migration of an active domain is unsupported. This prevents
competing producers from racing allocation checks but does not provide a
filesystem-wide byte quota or automatic stale-data reclamation. See the
[module capacity contract](../module/host_weight_runtime.md#validation-and-capacity).
`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not
cancel synchronous validation, a producer that has already started, or atomic
publication. A hung in-process producer therefore blocks its owning process and
@@ -289,13 +289,6 @@ store, tensor ownership, or H2D payload. On success, each tensor view copies
directly into the existing rotating HBM block buffers and the two private host
staging slots are not allocated.
`--dlo-host-registration-mode` defaults to `auto`, preserving this behavior.
Select `disabled` to bypass HWR mapping registration and use bounded host staging
regardless of mapping size or budget. This leaves the existing pin-memory policy
for staging buffers unchanged. The equivalent offline/stage configuration field
is `dlo_host_registration_mode`; an explicit disable requires enabled HWR with
no-AllGather DLO.
`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
memory policy, unsupported platform/capability, over-budget mapping, or fully
@@ -303,14 +296,6 @@ rolled-back registration error selects the existing two-slot staging path. A
partial registration that cannot be rolled back aborts startup because closing
the lease would unmap memory still owned by the platform.
Automatic registration logs its entry with device and budget, then each CUDA
region's ordinal/total and byte size before calling `cudaHostRegister`. A
completion message follows each successful call. These identify the last
entered boundary if progress stops; they do not impose a timeout. A synchronous
CUDA call cannot be safely cancelled by a Python timer, and a stuck process
still requires external supervision. The explicit disable option avoids this
registration path; it does not repair the underlying driver hang.
Direct checkpoint mmap remains unchanged and continues to use staging. It may
require loader-owned per-block transforms, while the HWR artifact already
contains final runtime bytes. DLO AllGather never receives an HWR final-layout
+35 -3
View File
@@ -311,9 +311,41 @@ point-in-time evidence only; kernel locks are not a persistent owner registry.
Capacity policy covers all store-owned bytes, including ready, temporary, and
quarantined data. The local writer preflights `max_artifact_bytes`,
`max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the
filesystem supports it. `ENOSPC` remains a normal typed store failure because
concurrent preflight is inherently racy. Automatic eviction and strict
concurrent reservations are not implemented.
filesystem supports it. `ENOSPC` remains a normal typed store failure.
`CapacityPolicy.build_admission` defaults to `concurrent`, preserving parallel
production for different identities and best-effort capacity preflight. An
explicit `serialized` policy, which requires `max_store_bytes`, admits one
cooperative build/publication operation per domain through
`locks/domain-build.lock`. Production lock order becomes domain admission,
per-key build, then artifact. Validated warm hits bypass admission; queued
same-key callers still recheck and join a completed artifact. Admission waiting
uses the caller's existing coordination deadline and returns a retryable build
timeout without interrupting its owner.
The kernel releases admission on process exit. An explicit same-key retry uses
the existing stale-temp recovery; files from unrelated dead builds continue to
count toward capacity. A synchronous producer that hangs still requires
external process supervision. Admission release errors are logged without
revising an already determined publication result.
Concurrent domains retain their exact schema-1 capacity document. Serialized
domains use a schema-2 capacity document with `build_admission=serialized`, so
mixed policies and old readers fail initialization rather than bypass admission.
Opt into serialization using a new root with an agreed policy; in-place policy
migration while old workers may exist is unsupported. For example:
```python
from vllm_omni.host_weight_runtime import CapacityPolicy
capacity = CapacityPolicy(max_store_bytes=192 * 1024**3, build_admission="serialized")
```
Serialization prevents cooperative producers from racing each other's
allocation checks. It is not a filesystem quota: control metadata, finalization,
and outside writers can still affect total usage. Automatic eviction and
parallel byte reservations are not implemented. This policy trades parallel
cold-build throughput for predictable producer admission.
For tmpfs, artifact bytes are host-memory consumption and may consume swap;
they must not be reported operationally as ordinary disk capacity. The store
-33
View File
@@ -193,39 +193,6 @@ def test_resolve_execution_mode_rejects_unknown_execution_type():
omni_config_module._resolve_execution_mode("unknown_execution_type")
@pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str):
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.offloader.base import OffloadConfig
options = {
"enable_distributed_layerwise_offload": True,
"dlo_use_allgather": False,
"host_weight_runtime_mode": "preferred",
"host_weight_runtime_root": str(tmp_path / "unused-store"),
"dlo_host_registration_mode": registration_mode,
}
config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2", cli_overrides=options)
stage = config.stage_by_id(0)
assert isinstance(stage, VllmOmniDiffusionStageConfig)
assert stage.diffusion_config.dlo_host_registration_mode == registration_mode
offline = OmniDiffusionConfig(**options)
transport = OffloadConfig.from_od_config(offline)
assert transport.dlo_host_registration_mode == registration_mode
assert transport.pin_cpu_memory
assert transport.dlo_host_registration_limit_gib == 0.0
assert not (tmp_path / "unused-store").exists()
@pytest.mark.parametrize("registration_mode", ["unknown", None, "disabled"])
def test_hwr_registration_policy_rejects_invalid_or_ineligible_config(registration_mode):
from vllm_omni.diffusion.data import OmniDiffusionConfig
for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection):
with pytest.raises(ValueError, match="dlo_host_registration_mode"):
config_class(dlo_host_registration_mode=registration_mode)
def test_from_pipeline_config_preserves_current_pipeline_config_object():
omni_config = _from_pipeline_key("minicpmo_4_5")
pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
@@ -124,33 +124,6 @@ def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: p
assert runtime.registered == []
@pytest.mark.parametrize("second_result", [0, 7])
def test_registration_logs_entry_before_each_runtime_call(monkeypatch: pytest.MonkeyPatch, second_result: int) -> None:
runtime = _FakeRuntime([0, second_result])
messages: list[str] = []
monkeypatch.setattr(registration_module.logger, "info", lambda message, *args: messages.append(message % args))
original_register = runtime.cudaHostRegister
def register(pointer: int, size: int, flags: int) -> int:
ordinal = len(runtime.registered) + 1
assert messages[-1] == f"Registering HWR host range {ordinal}/2: {size} bytes"
return original_register(pointer, size, flags)
monkeypatch.setattr(runtime, "cudaHostRegister", register)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
regions = (_region("first", 0x1000, 4096), _region("second", 0x3000, 4096))
if second_result:
with pytest.raises(HostRegistrationError, match="error-7"):
CudaHostRegistration.create(regions, max_bytes=None)
assert runtime.unregistered == [0x1000]
else:
registration = CudaHostRegistration.create(regions, max_bytes=None)
assert registration.close() == ()
assert sum(message.startswith("Registering ") for message in messages) == 2
assert sum(message.startswith("Registered ") for message in messages) == (1 if second_result else 2)
def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
runtime = _FakeRuntime([0], read_only_supported=False)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
@@ -96,7 +96,7 @@ class TestDistributedLayerwiseOffloadHook:
hook.offload_layer()
assert not hook.is_materialized
def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch, patched_offload_runtime):
def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
current_block = nn.Linear(2, 2)
next_block = nn.Linear(2, 2)
expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
@@ -1373,14 +1373,10 @@ class TestMmapWeightLoading:
synchronize.assert_called()
@pytest.mark.parametrize("registration_mode", ["auto", "disabled"])
@pytest.mark.parametrize("limit_gib", [0.0, 1.5])
def test_hwr_registration_selection_uses_bounded_staging(
def test_hwr_registration_failure_falls_back_to_bounded_staging(
self,
monkeypatch,
patched_offload_runtime,
registration_mode,
limit_gib,
):
plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
backend = DistributedLayerwiseOffloadBackend(
@@ -1389,8 +1385,6 @@ class TestMmapWeightLoading:
pin_cpu_memory=True,
dp_size=1,
dlo_use_allgather=False,
dlo_host_registration_mode=registration_mode,
dlo_host_registration_limit_gib=limit_gib,
),
torch.device("cpu"),
host_weight_plan=plan,
@@ -1400,14 +1394,12 @@ class TestMmapWeightLoading:
def fail_registration(*args, **kwargs):
del args, kwargs
assert registration_mode == "auto", "disabled mode entered platform registration"
raise HostRegistrationError("registration unavailable in CPU test")
def allocate_unpinned_staging(hooks, resident_group=None):
# Registration selection still observes pin_cpu_memory=True. The
# CPU-only test avoids asking the host for CUDA-pinned allocations.
for hook in hooks:
assert hook.pin_memory, "registration policy disabled staging pinning"
hook.pin_memory = False
return original_staging_allocator(hooks, resident_group)
@@ -1420,7 +1412,6 @@ class TestMmapWeightLoading:
assert backend._using_rank_local_mmap
assert not backend._using_registered_mmap
assert backend._host_registration is None
assert backend.config.pin_cpu_memory
assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
assert not lease.closed
@@ -2639,9 +2630,7 @@ class TestDistributedComponentSelection:
assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
torch.testing.assert_close(block.weight, expected)
def test_multirank_enable_failure_cleanup_skips_restore_collective(
self, monkeypatch, mocker, patched_offload_runtime
):
def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -2822,7 +2811,7 @@ class TestDistributedComponentSelection:
with pytest.raises(ValueError, match="not declared replicated"):
backend.enable(pipeline)
def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self, patched_offload_runtime):
def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
"""Every rank must reject an unsafe encoder group, including stub ranks."""
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
@@ -507,8 +507,6 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
"preferred",
"--host-weight-runtime-root",
"/var/cache/vllm-omni/hwr",
"--dlo-host-registration-mode",
"disabled",
"--dlo-host-registration-limit-gib",
"80",
]
@@ -524,8 +522,6 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
assert engine_args["host_weight_runtime_mode"] == "preferred"
assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
assert engine_args["dlo_host_registration_limit_gib"] == 80
assert explicit_kwargs["dlo_host_registration_mode"] == "disabled"
assert engine_args["dlo_host_registration_mode"] == "disabled"
def test_serve_cli_accepts_diffusion_compile_controls():
@@ -0,0 +1,305 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU process and policy coverage for domain-wide producer admission."""
from __future__ import annotations
import errno
import json
import multiprocessing as mp
import os
import time
from multiprocessing.connection import Connection
from pathlib import Path
from typing import Literal
import pytest
import torch
from tests.host_weight_runtime.test_filesystem_store import FakeProducer, _identity, _make_store, _publish_test_artifact
from vllm_omni.host_weight_runtime import (
BuildRequest,
CapacityPolicy,
FailureCode,
HostWeightError,
HostWeightRuntime,
HostWeightRuntimeConfig,
ProductionMetadata,
ResolutionOutcome,
RuntimeMode,
StoreResult,
StoreStatus,
TensorWriteSpec,
ValidationLevel,
WaitPolicy,
WeightArtifactIdentity,
WeightProductionSpec,
)
from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active
from vllm_omni.host_weight_runtime.protocols import ArtifactWriter
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class AdmissionProducer:
def __init__(self, identity: WeightArtifactIdentity, control: Connection | None = None) -> None:
self._spec = FakeProducer(identity).spec
self.control = control
@property
def spec(self) -> WeightProductionSpec:
return self._spec
def produce(self, writer: ArtifactWriter) -> ProductionMetadata:
action = "produce"
if self.control is not None:
self.control.send("entered")
if not self.control.poll(60):
raise TimeoutError("test did not release producer")
action = self.control.recv()
spec = TensorWriteSpec("payload", (65536,), torch.bfloat16)
with writer.open_tensor_file("weights.safetensors", (spec,)) as output:
output.write_tensor("payload", torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
if action == "crash":
os._exit(27)
return ProductionMetadata("test-producer-v1", "test-restorer-v1")
def _admission_process(root: Path, policy: CapacityPolicy, rank: int, control: Connection) -> None:
try:
store = _make_store(root, capacity=policy)
identity = _identity(tp_rank=rank)
original_lookup = store.lookup
first_lookup = True
def notify_lookup(
identity: WeightArtifactIdentity, *, validation: ValidationLevel, deadline: float | None = None
) -> StoreResult:
nonlocal first_lookup
result = original_lookup(identity, validation=validation, deadline=deadline)
if first_lookup:
# Prove the caller observed a miss before the parent lets the
# active producer publish, independent of process scheduling.
control.send("looked_up")
first_lookup = False
return result
store.lookup = notify_lookup # type: ignore[method-assign]
control.send("ready")
assert control.poll(60) and control.recv() == "start"
result = store.get_or_build(
BuildRequest(identity),
AdmissionProducer(identity, control),
validation=ValidationLevel.FULL_CHECKSUM,
deadline=time.monotonic() + 30,
)
if result.lease is not None:
result.lease.close()
control.send((result.status.value, result.failure.code.value if result.failure is not None else None))
finally:
control.close()
@pytest.mark.parametrize("admission", ["concurrent", "serialized"])
@pytest.mark.parametrize("same_key", [False, True])
def test_cross_identity_admission_and_capacity(
tmp_path: Path, admission: Literal["concurrent", "serialized"], same_key: bool
) -> None:
policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission=admission)
store = _make_store(tmp_path / "store", capacity=policy)
ctx = mp.get_context("spawn")
pairs = [ctx.Pipe() for _ in range(2)]
processes = [
ctx.Process(target=_admission_process, args=(store.root, policy, 0 if same_key else rank, pair[1]))
for rank, pair in enumerate(pairs)
]
try:
for process, pair in zip(processes, pairs, strict=True):
process.start()
pair[1].close()
first, second = [pair[0] for pair in pairs]
for connection in (first, second):
assert connection.poll(60) and connection.recv() == "ready"
first.send("start")
assert first.poll(10) and first.recv() == "looked_up"
assert first.poll(10) and first.recv() == "entered"
second.send("start")
assert second.poll(10) and second.recv() == "looked_up"
if admission == "concurrent" and not same_key:
assert second.poll(10) and second.recv() == "entered"
else:
assert not second.poll(0.2), "a second producer entered while coordination should exclude it"
first.send("produce")
assert first.poll(30) and first.recv() == (StoreStatus.BUILT.value, None)
if same_key:
assert second.poll(30) and second.recv() == (StoreStatus.JOINED.value, None)
else:
if admission == "serialized":
assert second.poll(10) and second.recv() == "entered"
second.send("produce")
assert second.poll(30)
assert second.recv() == (StoreStatus.FAILED.value, FailureCode.STORE_LIMIT_EXCEEDED.value)
assert policy.max_store_bytes is not None
assert store.inspect_domain().store_bytes < policy.max_store_bytes
assert not list(store.tmp_dir.iterdir())
assert {path.name for path in store.artifacts_dir.iterdir()} == {_identity(tp_rank=0).key}
for process in processes:
process.join(10)
assert process.exitcode == 0
finally:
for process in processes:
if process.pid is not None:
if process.is_alive():
process.kill()
process.join(5)
process.close()
for pair in pairs:
for connection in pair:
connection.close()
def test_admission_timeout_preserves_owner_and_warm_hit_bypasses_lock(tmp_path: Path) -> None:
policy = CapacityPolicy(max_store_bytes=512 * 1024, build_admission="serialized")
store = _make_store(tmp_path / "store", capacity=policy)
identity, _ = _publish_test_artifact(store)
other = _identity(tp_rank=1)
lock_path = store.locks_dir / "domain-build.lock"
with FileLock(lock_path, exclusive=True, deadline=None):
hit = store.get_or_build(
BuildRequest(identity),
FakeProducer(identity),
validation=ValidationLevel.FULL_CHECKSUM,
deadline=time.monotonic() + 0.02,
)
assert hit.status is StoreStatus.HIT and hit.lease is not None
hit.lease.close()
timeout = store.get_or_build(
BuildRequest(other),
FakeProducer(other),
validation=ValidationLevel.FULL_CHECKSUM,
deadline=time.monotonic() + 0.02,
)
assert timeout.status is StoreStatus.TIMEOUT
assert timeout.failure is not None and timeout.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
assert timeout.failure.retryable
assert lock_is_active(lock_path)
_publish_test_artifact(store, other)
def test_producer_failure_releases_admission(tmp_path: Path) -> None:
store = _make_store(
tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
)
identity = _identity()
failed = store.get_or_build(
BuildRequest(identity),
FakeProducer(identity, write_mode="incomplete"),
validation=ValidationLevel.FULL_CHECKSUM,
deadline=time.monotonic() + 5,
)
assert failed.status is StoreStatus.FAILED
assert not lock_is_active(store.locks_dir / "domain-build.lock")
_publish_test_artifact(store, identity)
@pytest.mark.parametrize(
("mode", "expected"),
[(RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED)],
)
def test_admission_timeout_obeys_runtime_mode(tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome) -> None:
policy = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
store = _make_store(tmp_path / "store", capacity=policy)
runtime = HostWeightRuntime.from_config(
HostWeightRuntimeConfig(mode=mode, domain=store.domain_policy, capacity=policy, wait=WaitPolicy(0.01))
)
identity = _identity()
with FileLock(store.locks_dir / "domain-build.lock", exclusive=True, deadline=None):
resolution = runtime.resolve(identity, producer=FakeProducer(identity))
assert resolution.report.outcome is expected
failure = resolution.report.attempts[-1].failure
assert failure is not None and failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
def test_crashed_producer_releases_admission_and_same_key_retry_recovers(tmp_path: Path) -> None:
policy = CapacityPolicy(max_store_bytes=192 * 1024, build_admission="serialized")
store = _make_store(tmp_path / "store", capacity=policy)
ctx = mp.get_context("spawn")
parent, child = ctx.Pipe()
process = ctx.Process(target=_admission_process, args=(store.root, policy, 0, child))
process.start()
child.close()
try:
assert parent.poll(60) and parent.recv() == "ready"
parent.send("start")
assert parent.poll(10) and parent.recv() == "looked_up"
assert parent.poll(10) and parent.recv() == "entered"
parent.send("crash")
process.join(10)
assert process.exitcode == 27
finally:
if process.is_alive():
process.kill()
process.join(5)
process.close()
parent.close()
assert not lock_is_active(store.locks_dir / "domain-build.lock")
assert list(store.tmp_dir.iterdir())
identity = _identity(tp_rank=0)
recovered = store.get_or_build(
BuildRequest(identity),
AdmissionProducer(identity),
validation=ValidationLevel.FULL_CHECKSUM,
deadline=time.monotonic() + 10,
)
assert recovered.status is StoreStatus.BUILT and recovered.lease is not None
assert torch.equal(recovered.lease.tensors["payload"], torch.arange(65536, dtype=torch.float32).to(torch.bfloat16))
recovered.lease.close()
assert not list(store.tmp_dir.iterdir())
def test_admission_release_error_preserves_publication(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
store = _make_store(
tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
)
original_close = FileLock.close
def fail_admission_close(lock: FileLock) -> None:
original_close(lock)
if lock.path.name == "domain-build.lock":
raise OSError(errno.EIO, "injected admission release failure")
monkeypatch.setattr(FileLock, "close", fail_admission_close)
_publish_test_artifact(store)
assert not lock_is_active(store.locks_dir / "domain-build.lock")
def test_admission_policy_is_authoritative_and_legacy_document_is_unchanged(tmp_path: Path) -> None:
concurrent = CapacityPolicy(max_store_bytes=16384)
serialized = CapacityPolicy(max_store_bytes=16384, build_admission="serialized")
legacy = _make_store(tmp_path / "legacy", capacity=concurrent)
path = legacy.root / "domain-policy.json"
original = path.read_bytes()
assert json.loads(original) == {
"schema_version": 1,
"policy_version": 1,
"max_artifact_bytes": None,
"max_store_bytes": 16384,
"min_free_bytes": 0,
"eviction": "none",
}
with pytest.raises(HostWeightError, match="incompatible schema"):
_make_store(legacy.root, capacity=serialized)
assert path.read_bytes() == original
strict = _make_store(tmp_path / "serialized", capacity=serialized)
document = json.loads((strict.root / "domain-policy.json").read_bytes())
assert document["schema_version"] == 2 and document["build_admission"] == "serialized"
assert _make_store(strict.root, capacity=serialized).domain_uuid == strict.domain_uuid
with pytest.raises(HostWeightError, match="incompatible schema"):
_make_store(strict.root, capacity=concurrent)
def test_admission_policy_validation() -> None:
with pytest.raises(ValueError, match="requires max_store_bytes"):
CapacityPolicy(build_admission="serialized")
with pytest.raises(ValueError, match="must be concurrent or serialized"):
CapacityPolicy(build_admission="unknown") # type: ignore[arg-type]
-2
View File
@@ -739,7 +739,6 @@ class _DiffusionConfigProjection:
dlo_resident_layers: int = Field(default=0, ge=0)
host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
host_weight_runtime_root: str | None = None
dlo_host_registration_mode: str = "auto"
dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
pin_cpu_memory: bool = True
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -896,7 +895,6 @@ class _DiffusionConfigProjection:
)
self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
limit_gib=self.dlo_host_registration_limit_gib,
mode=self.dlo_host_registration_mode,
enable_dlo=self.enable_distributed_layerwise_offload,
use_allgather=self.dlo_use_allgather,
hwr_mode=self.host_weight_runtime_mode,
-1
View File
@@ -483,7 +483,6 @@ class StageDeployConfig:
dlo_resident_layers: int | None = None
host_weight_runtime_mode: str | None = None
host_weight_runtime_root: str | None = None
dlo_host_registration_mode: str | None = None
dlo_host_registration_limit_gib: float | None = None
# Diffusion-specific debug and observability knobs.
enable_diffusion_pipeline_profiler: bool | None = None
-9
View File
@@ -119,15 +119,8 @@ def validate_dlo_host_registration_options(
enable_dlo: bool,
use_allgather: bool,
hwr_mode: object,
mode: object = "auto",
) -> float:
"""Validate the optional transport budget without probing CUDA or HWR."""
if mode not in ("auto", "disabled"):
raise ValueError("dlo_host_registration_mode must be auto or disabled")
if mode == "disabled" and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
raise ValueError(
"disabled dlo_host_registration_mode requires enabled no-AllGather DLO and Host Weight Runtime"
)
if not isinstance(limit_gib, (int, float, str)):
raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}")
value = float(limit_gib)
@@ -840,7 +833,6 @@ class OmniDiffusionConfig:
host_weight_runtime_root: str | None = None
# Optional per-worker ceiling for registering final-layout HWR mappings.
# Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
dlo_host_registration_mode: str = "auto"
dlo_host_registration_limit_gib: float = 0.0
pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading
@@ -1252,7 +1244,6 @@ class OmniDiffusionConfig:
)
self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
limit_gib=self.dlo_host_registration_limit_gib,
mode=self.dlo_host_registration_mode,
enable_dlo=self.enable_distributed_layerwise_offload,
use_allgather=self.dlo_use_allgather,
hwr_mode=self.host_weight_runtime_mode,
-6
View File
@@ -73,15 +73,12 @@ class OffloadConfig:
# Optional per-worker ceiling for registering an HWR mmap. Zero means no
# additional ceiling; pin_cpu_memory controls whether registration is tried.
dlo_host_registration_limit_gib: float = 0.0
dlo_host_registration_mode: str = "auto"
# ``None`` preserves the model's legacy plan-driven component topology;
# a frozenset is an explicit compact-API selection.
components: frozenset[str] | None = None
dlo_transfers: dict[str, DLOTransfer] | None = None
def __post_init__(self) -> None:
if self.dlo_host_registration_mode not in ("auto", "disabled"):
raise ValueError("dlo_host_registration_mode must be auto or disabled")
if self.components is not None:
self.components = parse_offload_components(self.components)
if self.dlo_transfers is None:
@@ -177,9 +174,7 @@ class OffloadConfig:
dlo_transfers = dict(resolved.transfers)
dlo_resident_layers = resolved.resident_layers
dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT)
registration_mode = getattr(od_config, "dlo_host_registration_mode", "auto")
dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
mode=registration_mode,
limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
enable_dlo=enable_distributed_layerwise_offload,
use_allgather=dit_uses_allgather,
@@ -215,7 +210,6 @@ class OffloadConfig:
dlo_use_allgather=dit_uses_allgather,
dlo_resident_layers=dlo_resident_layers,
dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
dlo_host_registration_mode=registration_mode,
components=components,
dlo_transfers=dlo_transfers,
)
@@ -6,14 +6,12 @@ from __future__ import annotations
import ctypes
import mmap
import time
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Protocol
import torch
from vllm.logger import init_logger
from vllm_omni.host_weight_runtime import MappedHostRegion
@@ -23,8 +21,6 @@ from .host_registration import (
HostRegistrationError,
)
logger = init_logger(__name__)
_CUDA_HOST_REGISTER_READ_ONLY = 0x08
_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
@@ -191,9 +187,7 @@ class CudaHostRegistration:
registered: list[_AddressRange] = []
try:
for index, region in enumerate(mapped, start=1):
logger.info("Registering HWR host range %d/%d: %d bytes", index, len(mapped), region.size)
started = time.perf_counter()
for region in mapped:
error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
if int(error) != 0:
raise HostRegistrationError(
@@ -201,9 +195,6 @@ class CudaHostRegistration:
f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
)
registered.append(region)
logger.info(
"Registered HWR host range %d/%d in %.3f s", index, len(mapped), time.perf_counter() - started
)
except Exception as exc:
rollback_errors: list[str] = []
rollback_failed: list[_AddressRange] = []
@@ -1080,9 +1080,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
lease = self._host_weight_lease
if lease is None:
return False
if self.config.dlo_host_registration_mode == "disabled":
logger.info("HWR mmap registration disabled by transport policy; using bounded host staging")
return False
if not self.config.pin_cpu_memory:
logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
return False
@@ -1092,12 +1089,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
limit_gib = self.config.dlo_host_registration_limit_gib
max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
logger.info(
"Starting HWR mmap registration on %s: %d source range(s), budget_bytes=%s",
self.device,
len(lease.mapped_regions),
max_bytes,
)
started = time.perf_counter()
try:
registration = register_host_mappings(
-1
View File
@@ -582,7 +582,6 @@ class OrchestratorArgs:
dlo_resident_layers: int = 0
host_weight_runtime_mode: str = "disabled"
host_weight_runtime_root: str | None = None
dlo_host_registration_mode: str = "auto"
dlo_host_registration_limit_gib: float = 0.0
boundary_ratio: float | None = None
flow_shift: float | None = None
-1
View File
@@ -1145,7 +1145,6 @@ class AsyncOmniEngine:
"dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
"host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
"host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
"dlo_host_registration_mode": kwargs.get("dlo_host_registration_mode", "auto"),
"dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
"enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
"diffusion_compile_granularity": (
-11
View File
@@ -780,17 +780,6 @@ class OmniServeCommand(CLISubcommand):
"the same persistent path for population and serving."
),
)
omni_config_group.add_argument(
"--dlo-host-registration-mode",
choices=("auto", "disabled"),
default="auto",
help=(
"HWR transport registration policy for eligible no-AllGather DLO. "
"auto attempts registered direct H2D under the existing pin-memory and budget policy; "
"disabled selects bounded staging without registering the HWR mappings. "
"This preserves the pin-memory policy for staging buffers."
),
)
omni_config_group.add_argument(
"--dlo-host-registration-limit-gib",
type=float,
+6
View File
@@ -8,6 +8,7 @@ import math
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Literal
class RuntimeMode(str, Enum):
@@ -66,6 +67,7 @@ class CapacityPolicy:
max_store_bytes: int | None = None
min_free_bytes: int = 0
eviction: str = "none"
build_admission: Literal["concurrent", "serialized"] = "concurrent"
def __post_init__(self) -> None:
for name, value in (
@@ -78,6 +80,10 @@ class CapacityPolicy:
raise ValueError("min_free_bytes must not be negative")
if not isinstance(self.eviction, str) or self.eviction != "none":
raise ValueError("automatic host weight eviction is not implemented")
if self.build_admission not in ("concurrent", "serialized"):
raise ValueError("build_admission must be concurrent or serialized")
if self.build_admission == "serialized" and self.max_store_bytes is None:
raise ValueError("serialized build admission requires max_store_bytes")
@dataclass(frozen=True)
@@ -627,7 +627,7 @@ class FilesystemHostWeightStore:
self.policy_version = 1
def _capacity_document(self, *, policy_version: int) -> dict[str, object]:
return {
document: dict[str, object] = {
"schema_version": 1,
"policy_version": policy_version,
"max_artifact_bytes": self.capacity_policy.max_artifact_bytes,
@@ -635,6 +635,9 @@ class FilesystemHostWeightStore:
"min_free_bytes": self.capacity_policy.min_free_bytes,
"eviction": self.capacity_policy.eviction,
}
if self.capacity_policy.build_admission == "serialized":
document.update(schema_version=2, build_admission="serialized")
return document
def _build_lock_path(self, key: str) -> Path:
return self.locks_dir / f"{key}.build.lock"
@@ -1013,6 +1016,40 @@ class FilesystemHostWeightStore:
if initial.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}:
return initial
if self.capacity_policy.build_admission == "concurrent":
return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
try:
admission = FileLock(self.locks_dir / "domain-build.lock", exclusive=True, deadline=deadline)
except FileLockTimeoutError as exc:
return StoreResult(
StoreStatus.TIMEOUT,
failure=_failure(
ResolutionStage.PRODUCTION, FailureCode.ACTIVE_BUILD_TIMEOUT, str(exc), retryable=True
),
)
except OSError as exc:
return StoreResult(
StoreStatus.FAILED,
failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True),
)
try:
return self._build_after_miss(identity, producer, validation=validation, deadline=deadline)
finally:
try:
admission.close()
except OSError:
# Releasing coordination must not revise the completed build's
# publication outcome or replace its primary failure.
logger.warning("Failed to release domain build admission under %s", self.root, exc_info=True)
def _build_after_miss(
self,
identity: WeightArtifactIdentity,
producer: WeightProducer,
*,
validation: ValidationLevel,
deadline: float,
) -> StoreResult:
try:
build_lock = FileLock(self._build_lock_path(identity.key), exclusive=True, deadline=deadline)
except FileLockTimeoutError as exc: