diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md new file mode 100644 index 00000000000..8998111518f --- /dev/null +++ b/benchmarks/host_weight_runtime/README.md @@ -0,0 +1,57 @@ +# Host weight dependency memory diagnostic + +`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from +access to one cached tensor view. It imports neither vLLM nor HWR and uses a +synthetic float32 payload controlled by `--tensor-elements` (default: 64). +Tensor size can affect retention: a negative result for 16 elements does not +rule out growth for 64 elements. The selected size is recorded in the JSON +arguments and used for both file creation and the correctness check. +HWR itself calls `get_tensor()` when +acquiring a lease, then exposes cached views; this probe does not measure HWR +requests or repeated lease acquisition. + +Run one mode per fresh process in the same Python environment and CPU affinity: + +```bash +git rev-parse HEAD > revision.txt +git status --short > working-tree.txt +git diff > working-tree.patch +# Select an allowed CPU from your process affinity; 0 is only an example. +CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json +``` + +For a size comparison, run 16 and 64 elements twice each in fresh processes, +with cached reuse as a control at both sizes: + +```bash +for repetition in 1 2; do + for elements in 16 64; do + for mode in get_tensor reuse; do + CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode "$mode" --tensor-elements "$elements" \ + > "$mode-$elements-$repetition.json" + done + done +done +``` + +The diagnostic sets Torch to one CPU thread, warms up before the baseline, +and reports preparation separately from timed loop work. Checkpoints run GC +and report Linux private/anonymous memory, RSS, and open descriptors. It checks +that transient tensor objects and the final cached view are released. Linux +`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc +support fail explicitly. Keep raw JSON and the repository snapshot with results. + +Compare the final loop sample with the warmed baseline, then inspect the closed +sample. Private-memory growth after Python tensor objects disappear is evidence +of retention in the dependency/allocator path, not proof of a leak's root cause. +RSS can include file-backed pages; it is not a private-memory metric. Samples +also include small diagnostic bookkeeping allocations. Do not assert a portable +memory threshold or turn this synthetic loop into a per-request service estimate. + +Do not drop shared caches, trim allocators, or add serving-time GC to make the +numbers look smaller. A dependency replacement or version change requires its +own reproducible comparison and correctness/lifecycle validation. diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py new file mode 100644 index 00000000000..55795e7ef80 --- /dev/null +++ b/benchmarks/host_weight_runtime/safetensors_retention.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import platform +import tempfile +import time +import weakref +from pathlib import Path + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]: + gc.collect() + counters = {} + for line in Path("/proc/self/smaps_rollup").read_text().splitlines(): + fields = line.split() + if len(fields) == 3 and fields[2] == "kB": + counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024 + return { + "phase": phase, + "iterations": iterations, + "loop_seconds": loop_seconds, + "rss_bytes": counters["Rss"], + "private_bytes": counters["Private_Clean"] + counters["Private_Dirty"], + "anonymous_bytes": counters["Anonymous"], + "fd_count": len(tuple(Path("/proc/self/fd").iterdir())), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True) + parser.add_argument( + "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor" + ) + parser.add_argument("--iterations", type=positive_int, default=1_000_000) + parser.add_argument("--sample-every", type=positive_int, default=100_000) + parser.add_argument("--warmup", type=positive_int, default=1000) + args = parser.parse_args() + if not Path("/proc/self/smaps_rollup").is_file(): + parser.error("requires Linux /proc/self/smaps_rollup") + + preparation_start = time.monotonic() + import safetensors + import torch + from safetensors import safe_open + from safetensors.torch import save_file + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + samples = [sample("before_file", 0, 0.0)] + with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory: + path = Path(directory) / "synthetic.safetensors" + save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path)) + with safe_open(path, framework="pt", device="cpu") as reader: + cached = reader.get_tensor("weight") + assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32)) + for _ in range(args.warmup): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + del tensor + samples.append(sample("baseline", 0, 0.0)) + preparation_seconds = time.monotonic() - preparation_start + completed = 0 + loop_seconds = 0.0 + while completed < args.iterations: + count = min(args.sample_every, args.iterations - completed) + start = time.monotonic() + for _ in range(count): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + last_tensor = weakref.ref(tensor) + del tensor + loop_seconds += time.monotonic() - start + completed += count + snapshot = sample("loop", completed, loop_seconds) + snapshot["last_tensor_alive"] = last_tensor() is not None + if args.mode == "get_tensor": + assert last_tensor() is None, "diagnostic retained a transient tensor" + samples.append(snapshot) + del cached + del reader + samples.append(sample("closed", completed, loop_seconds)) + assert last_tensor() is None, "last tensor survived reader/view teardown" + + print( + json.dumps( + { + "schema_version": 1, + "arguments": vars(args), + "python": platform.python_version(), + "torch": torch.__version__, + "safetensors": safetensors.__version__, + "kernel": platform.release(), + "hostname": platform.node(), + "pid": os.getpid(), + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "torch_threads": torch.get_num_threads(), + "pythonmalloc": os.environ.get("PYTHONMALLOC", "default"), + "preparation_seconds": preparation_seconds, + "samples": samples, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index a9bb1b0b838..28a01974802 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -42,6 +42,11 @@ V1 does not include: ## Motivation and use cases +For investigating CPU memory retained by dependency tensor materialization, +see the [standalone safetensors diagnostic](../../../benchmarks/host_weight_runtime/README.md). +It distinguishes repeated dependency calls from reuse of cached views and does +not establish per-request HWR leakage. + Model loading can create the same final host representation repeatedly. This is especially expensive when loading performs checkpoint decoding, tensor renaming, TP slicing, quantization, packing, or scale construction before GPU @@ -67,6 +72,26 @@ views and mapped ranges. A separate transport still decides whether to use registered mmap, private pinned staging, synchronous copies, or asynchronous H2D transfer. +## Payload validation at startup + +Eligible diffusion loaders accept `--host-weight-runtime-validation` (or the +`host_weight_runtime_validation` offline/stage configuration field): + +- `manifest_and_metadata` is the existing default. It validates identity, + metadata, tensor structure, and sizes, but does not detect payload-only + corruption. +- `full_checksum` reads and hashes every payload on warm acquisition before + restoring weights. This adds startup I/O/CPU work proportional to artifact + bytes; it is not a per-inference check. + +For example, add `--host-weight-runtime-validation full_checksum` to an enabled +HWR deployment. Preferred mode falls back to canonical loading and can publish +a replacement after detecting corruption; required mode fails startup. Readers +can use different validation levels in the same domain without changing artifact +identity. Neither level authenticates an untrusted writer or prevents mutation +after validation. Checksumming does not imply automatic recovery of a running +service. + ## Resolution behavior The loader resolves the immutable canonical source and computes the exact @@ -313,8 +338,25 @@ 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. -`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not -cancel synchronous validation, a producer that has already started, or atomic +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 domain-initialization and lookup/build +lock acquisition. Store construction and each later resolution or publication +operation have separate budgets from the same wait policy, rather than one +end-to-end startup deadline. A domain-init timeout follows the retryable domain +failure policy below. After contention ends, a fresh runtime construction can +retry initialization; the timed-out runtime retains its original failure. + +The coordination budget does not cancel filesystem I/O, synchronous validation, +a producer that has already started, or atomic publication. A hung in-process producer therefore blocks its owning process and must be handled by external process supervision. Enforceable producer cancellation requires a future process-isolated producer contract. diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 730d5547252..315a1defce8 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -289,6 +289,13 @@ 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 @@ -296,6 +303,14 @@ 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 diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index ac92711fc4d..b060ed0ac9a 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -136,8 +136,16 @@ so only when it is marked retryable; unsupported capabilities, identity collisions, producer failures, and publication failures remain visible as startup failures. -`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it -is not a hard wall-clock deadline for in-process validation, production, or +`coordination_timeout_seconds` bounds acquisition of the domain-initialization +lock and lookup/build locks. Store construction and each later resolution or +publication operation receive separate coordination budgets from the same +`WaitPolicy`; they do not share an end-to-end startup deadline. A domain-init +lock timeout is a retryable domain failure: preferred mode may fall back and +required mode fails. Retrying initialization requires constructing a fresh +runtime after contention ends; a waiter never cancels or unlocks the owner. + +The coordination budget is not a hard wall-clock deadline for filesystem I/O, +in-process validation, production, or atomic publication. Once this V1 implementation becomes the producer, the synchronous producer and publication run to completion. A hung producer blocks its owning process and requires external process supervision; enforceable @@ -299,6 +307,43 @@ Explicit cleanup uses the locked move and retries exact-key `.cleanup.*` tombstones left by an interrupted removal; a later build performs the same tombstone reconciliation before producing replacement content. +Failure-quarantined entries remain until explicitly selected for removal; +ordinary `cleanup(identity)` does not remove that diagnostic history. +`FilesystemHostWeightStore.cleanup_quarantined(storage_name)` accepts one +quarantined inventory basename and preserves the current artifact, deny marker, +and other failure-quarantined entries. It takes the same nonblocking build and +exclusive artifact locks, so an active builder or lease returns a typed refusal. +Malformed names are rejected before mutation, and symlinks are not followed. + +Selected entries move through the existing `.cleanup.*` tombstone protocol. +Retrying the original name reconciles pending cleanup tombstones for that key, +even if the original entry disappeared during an earlier attempt. A missing +selection is otherwise an idempotent success. Existing cleanup intents for the +key may be completed along with the selection; other quarantine history is +retained. Parent synchronization is retried even when an earlier deletion +removed the last tombstone before its sync failed. + +For a configured store, operators can inspect and deliberately select an entry: + +```python +from vllm_omni.host_weight_runtime import ArtifactInventoryState, HostWeightError + +for entry in store.inspect_domain().inventory: + if entry.state is ArtifactInventoryState.QUARANTINED: + print(entry.storage_name, entry.size_bytes) + +selected_storage_name = input("Quarantined entry to remove: ") +failure = store.cleanup_quarantined(selected_storage_name) +if failure is not None: + raise HostWeightError(failure) +``` + +Use the same authoritative domain and capacity policy as the service. Review +inventory sizes and inspect again after cleanup before explicitly retrying a +build or startup. These are logical-byte, point-in-time observations; concurrent +activity can affect the difference and physical disk/RAM reclamation is not +implied. This operation adds no automatic retention policy or service retry. + An operational failure while inspecting a noncooperative competing publication is a retryable storage failure. The competing entry remains authoritative and is not mislabeled as corrupt or quarantined merely because its manifest could @@ -311,9 +356,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 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f91e6af42b9..e41e127900d 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -193,6 +193,74 @@ def test_resolve_execution_mode_rejects_unknown_execution_type(): omni_config_module._resolve_execution_mode("unknown_execution_type") +@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"]) +def test_hwr_validation_survives_config_projection(validation: str): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + override = StageDeployConfig(stage_id=0, host_weight_runtime_validation=validation) + config = _from_pipeline_key( + "dreamzero", + deploy_config_path="dreamzero_tp1_cfg2", + cli_overrides={"host_weight_runtime_validation": validation}, + ) + stage = config.stage_by_id(0) + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.host_weight_runtime_validation == validation + projection = omni_config_module._DiffusionConfigProjection.from_kwargs( + host_weight_runtime_validation=override.host_weight_runtime_validation + ) + assert projection.host_weight_runtime_validation == validation + assert OmniDiffusionConfig(host_weight_runtime_validation=validation).host_weight_runtime_validation == validation + assert omni_config_module._DiffusionConfigProjection().host_weight_runtime_validation == "manifest_and_metadata" + assert OmniDiffusionConfig().host_weight_runtime_validation == "manifest_and_metadata" + + +@pytest.mark.parametrize("validation", ["fs_verity", "unknown", None]) +def test_hwr_validation_rejects_unsupported_values(validation): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection): + with pytest.raises(ValueError, match="host_weight_runtime_validation"): + config_class(host_weight_runtime_validation=validation) + + +@pytest.mark.parametrize("registration_mode", ["auto", "disabled"]) +@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"]) +def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str, validation: 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, + "host_weight_runtime_validation": validation, + } + 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 + assert stage.diffusion_config.host_weight_runtime_validation == validation + offline = OmniDiffusionConfig(**options) + assert offline.host_weight_runtime_validation == validation + 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") diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 555cb19bb4d..cc306776242 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -260,6 +260,78 @@ def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]: assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 +@pytest.mark.parametrize("validation", ["manifest_and_metadata", "full_checksum"]) +@pytest.mark.parametrize("mode", ["preferred", "required"]) +def test_hwr_loader_payload_integrity_policy(tmp_path: Path, monkeypatch, validation: str, mode: str): + from vllm_omni.diffusion.offloader.startup import take_offload_startup_state + + canonical = tmp_path / "canonical" + canonical.mkdir() + expected = torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2) + save_file({"weight": expected}, str(canonical / "model.safetensors")) + root = tmp_path / "store" + + def make_loader(runtime_mode: str): + config = _hwr_config(canonical, root, mode=runtime_mode) + config.host_weight_runtime_validation = validation + loader = DiffusersPipelineLoader(LoadConfig(), config) + pipeline = _HWRPipeline(canonical) + monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline) + return loader, pipeline + + def close_plan(pipeline): + state = take_offload_startup_state(pipeline) + if state is not None and state.host_weight_plan is not None: + carrier = state.host_weight_plan.lease_carrier + if carrier is not None: + carrier.close() + + cold_loader, cold = make_loader("preferred") + cold_loader.load_model(load_device="cpu", device=torch.device("cpu")) + assert cold.load_count == 1 + close_plan(cold) + payloads = list((root / "artifacts").glob("*/*.safetensors")) + assert len(payloads) == 1 + payload = payloads[0] + data = bytearray(payload.read_bytes()) + data[-1] ^= 1 # Keep the header and file length valid; change only a BF16 payload byte. + original_mode = payload.stat().st_mode + payload.chmod(0o600) + payload.write_bytes(data) + payload.chmod(original_mode) + + loader, pipeline = make_loader(mode) + try: + if validation == "full_checksum" and mode == "required": + with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"): + loader.load_model(load_device="cpu", device=torch.device("cpu")) + assert pipeline.load_count == 0 + assert loader.take_host_weight_plan() is None + else: + loader.load_model(load_device="cpu", device=torch.device("cpu")) + if validation == "manifest_and_metadata": + assert pipeline.load_count == 0 + assert not torch.equal(pipeline.transformer.weight, expected) + else: + assert pipeline.load_count == 1 + assert torch.equal(pipeline.transformer.weight, expected) + finally: + close_plan(pipeline) + + if validation == "full_checksum": + if mode == "required": + assert list((root / "deny").iterdir()) + else: + assert list((root / "quarantine").iterdir()) + warm_loader, warm = make_loader(mode) + try: + warm_loader.load_model(load_device="cpu", device=torch.device("cpu")) + assert warm.load_count == 0 + assert torch.equal(warm.transformer.weight, expected) + finally: + close_plan(warm) + + def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): cfg = SimpleNamespace( lora_backend="distill", diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index a3efbb40a12..bbed5a3e3a1 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -124,6 +124,33 @@ 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) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 05d7d93b583..0a838a36952 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -96,7 +96,7 @@ def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime hook.offload_layer() assert not hook.is_materialized - def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch): + def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch, patched_offload_runtime): 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,10 +1373,14 @@ def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runti synchronize.assert_called() - def test_hwr_registration_failure_falls_back_to_bounded_staging( + @pytest.mark.parametrize("registration_mode", ["auto", "disabled"]) + @pytest.mark.parametrize("limit_gib", [0.0, 1.5]) + def test_hwr_registration_selection_uses_bounded_staging( self, monkeypatch, patched_offload_runtime, + registration_mode, + limit_gib, ): plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback") backend = DistributedLayerwiseOffloadBackend( @@ -1385,6 +1389,8 @@ def test_hwr_registration_failure_falls_back_to_bounded_staging( 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, @@ -1394,12 +1400,14 @@ def test_hwr_registration_failure_falls_back_to_bounded_staging( 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) @@ -1412,6 +1420,7 @@ def allocate_unpinned_staging(hooks, resident_group=None): 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 @@ -2630,7 +2639,9 @@ def fail_second_hook(*args, **kwargs): 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): + def test_multirank_enable_failure_cleanup_skips_restore_collective( + self, monkeypatch, mocker, patched_offload_runtime + ): backend = DistributedLayerwiseOffloadBackend( OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, @@ -2811,7 +2822,7 @@ def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_ with pytest.raises(ValueError, match="not declared replicated"): backend.enable(pipeline) - def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): + def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self, patched_offload_runtime): """Every rank must reject an unsafe encoder group, including stub ranks.""" backend = DistributedLayerwiseOffloadBackend( OffloadConfig( diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index efe6bf9a320..ade8d8fa725 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -4,8 +4,10 @@ import asyncio import multiprocessing as mp import queue +import signal import threading import time +import weakref from types import SimpleNamespace from unittest.mock import MagicMock, Mock @@ -1330,6 +1332,7 @@ def test_worker_joins_share_one_global_deadline(self, monkeypatch): class FakeProcess: def __init__(self, name): self.name = name + self.pid = 123 self.alive = True self.terminated = False self.join_timeouts = [] @@ -1358,6 +1361,132 @@ def terminate(self): assert first.terminated and second.terminated assert not first.is_alive() and not second.is_alive() + @pytest.mark.parametrize("cooperative", [False, True]) + def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + ctx = mp.get_context("fork") + ready, stop = ctx.Event(), ctx.Event() + + def run_worker(): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + ready.set() + stop.wait() + + dead = ctx.Process(target=lambda: None) + worker = ctx.Process(target=run_worker) + monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2) + monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1) + monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False) + dead.start() + worker.start() + try: + dead.join(5) + assert not dead.is_alive() + assert ready.wait(5), "worker did not install its signal handler" + mq = Mock() + if cooperative: + mq.enqueue.side_effect = lambda *args, **kwargs: stop.set() + else: + mq.enqueue.side_effect = OSError("shutdown queue unavailable") + cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker]) + + cleaner() + + assert not worker.is_alive() + assert worker.exitcode == (0 if cooperative else -signal.SIGKILL) + assert cleaner.processes == [] + assert cleaner.broadcast_mq is None + cleaner() + finally: + for proc in (dead, worker): + if proc.is_alive(): + proc.kill() + proc.join(5) + proc.close() + + @pytest.mark.parametrize("failed_action", ["terminate", "kill"]) + def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + first, second = Mock(pid=101), Mock(pid=102) + first.name, second.name = "first", "second" + for proc in (first, second): + proc.is_alive.side_effect = [True, True, True, False] + getattr(first, failed_action).side_effect = OSError("signal failed") + first.join.side_effect = [OSError("join failed"), None, None] + monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134]) + monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic)) + cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second]) + + cleaner() + + first.kill.assert_called_once() + second.kill.assert_called_once() + assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5] + assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1] + assert cleaner.processes == [] + + def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + survivor = Mock(pid=123) + survivor.name = "surviving-worker" + survivor.is_alive.return_value = True + cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor]) + executor, _, _ = _make_executor() + executor._shutdown_cleaner = cleaner + executor._processes = [survivor] + executor._finalizer = weakref.finalize(executor, cleaner) + executor._pump_stop = threading.Event() + executor._futures_lock = threading.RLock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._batch_split_map = {} + log_error = Mock() + monkeypatch.setattr(executor_module.logger, "error", log_error) + + executor.shutdown() + + assert not executor._finalizer.alive + assert executor._shutdown_cleaner is cleaner + assert executor._processes == [survivor] + log_error.assert_called_once() + assert log_error.call_args.args[1] == [("surviving-worker", 123)] + survivor.is_alive.return_value = False + + executor.shutdown() + executor.shutdown() + + assert executor._shutdown_cleaner is None + assert executor._processes == [] + + def test_concurrent_cleanup_does_not_repeat_process_operations(self): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + joining, release = threading.Event(), threading.Event() + proc = Mock(pid=123) + proc.is_alive.return_value = True + + def join(timeout): + joining.set() + assert release.wait(5) + proc.is_alive.return_value = False + + proc.join.side_effect = join + cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc]) + thread = threading.Thread(target=cleaner) + thread.start() + try: + assert joining.wait(5) + cleaner() + proc.join.assert_called_once() + finally: + release.set() + thread.join(5) + assert not thread.is_alive() + assert cleaner.processes == [] + # ───────── monitor thread & death sentinel integration tests ───────── @@ -1406,7 +1535,7 @@ def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self): executor._result_mq = None executor._shutdown_cleaner = None # Use a no-op so shutdown() doesn't crash on None resources. - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) # ------------------------------------------------------------------ # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates # over them, so they need to exist even when constructed via __new__. @@ -1441,7 +1570,7 @@ def test_worker_monitor_noop_when_already_closed(self): executor._broadcast_mq = None executor._result_mq = None executor._shutdown_cleaner = None - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) proc = _make_short_lived_process() executor._processes = [proc] diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 616d5486eb4..ed9fcff2908 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -503,10 +503,14 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", + "--host-weight-runtime-validation", + "full_checksum", "--host-weight-runtime-mode", "preferred", "--host-weight-runtime-root", "/var/cache/vllm-omni/hwr", + "--dlo-host-registration-mode", + "disabled", "--dlo-host-registration-limit-gib", "80", ] @@ -519,9 +523,13 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80 + assert explicit_kwargs["host_weight_runtime_validation"] == "full_checksum" + assert engine_args["host_weight_runtime_validation"] == "full_checksum" 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(): diff --git a/tests/host_weight_runtime/test_build_admission.py b/tests/host_weight_runtime/test_build_admission.py new file mode 100644 index 00000000000..8bd06bbb6cf --- /dev/null +++ b/tests/host_weight_runtime/test_build_admission.py @@ -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] diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py index 7fe716c5ddd..d4051fd0e81 100644 --- a/tests/host_weight_runtime/test_filesystem_store.py +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -440,6 +440,22 @@ def _corrupt_test_payload(artifact: Path) -> None: os.chmod(artifact, 0o555) +def _quarantine_and_rebuild(store: FilesystemHostWeightStore, identity: WeightArtifactIdentity) -> Path: + before = set(store.quarantine_dir.iterdir()) + _corrupt_test_payload(store.artifacts_dir / identity.key) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + (quarantined,) = set(store.quarantine_dir.iterdir()) - before + return quarantined + + def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None: try: identity = _identity() @@ -2496,6 +2512,199 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path: assert not list(free_limited.artifacts_dir.iterdir()) +def test_cleanup_quarantined_preserves_current_artifact_other_entries_and_deny(tmp_path: Path) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + first = _quarantine_and_rebuild(store, identity) + second = _quarantine_and_rebuild(store, identity) + ready_inode = artifact.stat().st_ino + lock_inodes = {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()} + + assert store.cleanup_quarantined(first.name) is None + assert set(store.quarantine_dir.iterdir()) == {second} + assert artifact.stat().st_ino == ready_inode + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.lease is not None + assert torch.equal(hit.lease.tensors["layer.weight"], _weight()) + hit.lease.close() + + _corrupt_test_payload(artifact) + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.INVALID + deny = store.deny_dir / f"{identity.key}.json" + deny_bytes = deny.read_bytes() + assert store.cleanup_quarantined(second.name) is None + assert deny.read_bytes() == deny_bytes + assert artifact.stat().st_ino == ready_inode + assert {path.name: path.stat().st_ino for path in store.locks_dir.iterdir()} == lock_inodes + assert store.cleanup_quarantined(first.name) is None + + +@pytest.mark.parametrize("active", ["build", "lease"]) +def test_cleanup_quarantined_refuses_active_key(tmp_path: Path, active: str) -> None: + store = _make_store(tmp_path / "store") + identity, _ = _publish_test_artifact(store) + quarantined = _quarantine_and_rebuild(store, identity) + owner: FileLock | HostWeightLease + if active == "build": + owner = FileLock(store._build_lock_path(identity.key), exclusive=True, deadline=None) + expected = FailureCode.ACTIVE_BUILD_TIMEOUT + else: + lease = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).lease + assert lease is not None + owner = lease + expected = FailureCode.ACTIVE_LEASE + try: + failure = store.cleanup_quarantined(quarantined.name) + assert failure is not None and failure.code is expected and failure.retryable + assert quarantined.is_dir() + finally: + owner.close() + assert store.cleanup_quarantined(quarantined.name) is None + + +@pytest.mark.parametrize("recovery", ["cleanup", "build"]) +def test_cleanup_quarantined_retries_post_move_removal_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, recovery: str +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + quarantined = _quarantine_and_rebuild(store, identity) + original_rmtree = shutil.rmtree + + def fail_removal(path: str | os.PathLike[str]) -> None: + if Path(path).parent == store.quarantine_dir: + raise OSError(errno.EIO, "injected removal failure") + original_rmtree(path) + + with monkeypatch.context() as fault: + fault.setattr(shutil, "rmtree", fail_removal) + failure = store.cleanup_quarantined(quarantined.name) + assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable + assert not quarantined.exists() + assert len(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) == 1 + + if recovery == "cleanup": + assert store.cleanup_quarantined(quarantined.name) is None + else: + # A real build (not a warm hit) reconciles the same tombstone protocol. + _corrupt_test_payload(artifact) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert result.status is StoreStatus.BUILT + assert result.lease is not None + assert torch.equal(result.lease.tensors["layer.weight"], _weight()) + result.lease.close() + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +@pytest.mark.parametrize("after_removal", [False, True]) +def test_cleanup_quarantined_retries_parent_sync_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, after_removal: bool +) -> None: + store = _make_store(tmp_path / "store") + identity, _ = _publish_test_artifact(store) + quarantined = _quarantine_and_rebuild(store, identity) + original_sync = filesystem_store_module._fsync_directory + + def fail_sync(path: Path) -> None: + if path == store.quarantine_dir and ( + not after_removal or not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + ): + raise OSError(errno.EIO, "injected quarantine sync failure") + original_sync(path) + + with monkeypatch.context() as fault: + fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync) + failure = store.cleanup_quarantined(quarantined.name) + assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED and failure.retryable + assert not quarantined.exists() + assert bool(list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))) is not after_removal + + # An already absent selection must still retry the failed durability step. + with monkeypatch.context() as fault: + fault.setattr(filesystem_store_module, "_fsync_directory", fail_sync) + assert store.cleanup_quarantined(quarantined.name) is not None + assert store.cleanup_quarantined(quarantined.name) is None + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +@pytest.mark.parametrize("entry_kind", ["symlink", "file"]) +def test_cleanup_quarantined_rejects_non_directory_entry(tmp_path: Path, entry_kind: str) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + entry = store.quarantine_dir / f"{identity.key}.validation_failure.test" + if entry_kind == "symlink": + entry.symlink_to(artifact, target_is_directory=True) + else: + entry.write_text("not an artifact directory") + failure = store.cleanup_quarantined(entry.name) + assert failure is not None and failure.code is FailureCode.QUARANTINE_FAILED + assert not failure.retryable + assert entry.exists() + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.lease is not None + assert torch.equal(hit.lease.tensors["layer.weight"], _weight()) + hit.lease.close() + + +@pytest.mark.parametrize( + "name", + ["../escape", "/tmp/escape", "f" * 64, "g" * 64 + ".x", "f" * 64 + ".", "f" * 64 + ".x/../y", "f" * 64 + ".x\0"], +) +def test_cleanup_quarantined_rejects_invalid_selection_before_mutation(tmp_path: Path, name: str) -> None: + store = _make_store(tmp_path / "store") + locks = set(store.locks_dir.iterdir()) + with pytest.raises(ValueError, match="inventory basename"): + store.cleanup_quarantined(name) + assert set(store.locks_dir.iterdir()) == locks + + +def test_cleanup_quarantined_restores_bounded_store_recovery_headroom(tmp_path: Path) -> None: + store = _make_store(tmp_path / "store", capacity=CapacityPolicy(max_store_bytes=16384)) + identity = _identity() + for _ in range(32): + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + if result.status is StoreStatus.FAILED: + break + assert result.lease is not None + result.lease.close() + _corrupt_test_payload(store.artifacts_dir / identity.key) + assert result.failure is not None and result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED + quarantined = [ + entry for entry in store.inspect_domain().inventory if entry.state is ArtifactInventoryState.QUARANTINED + ] + assert len(quarantined) > 1 + assert store.cleanup(identity) is None + still_full = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert still_full.failure is not None and still_full.failure.code is FailureCode.STORE_LIMIT_EXCEEDED + + assert store.cleanup_quarantined(quarantined[0].storage_name) is None + recovered = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert recovered.status is StoreStatus.BUILT + assert recovered.lease is not None + assert torch.equal(recovered.lease.tensors["layer.weight"], _weight()) + recovered.lease.close() + + def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None: root = tmp_path / "store" _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096)) diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py index 21bab37dee4..cd2de28165d 100644 --- a/tests/host_weight_runtime/test_runtime_resolution.py +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -4,6 +4,10 @@ from __future__ import annotations +import multiprocessing as mp +import threading +import time +from multiprocessing.connection import Connection from pathlib import Path from typing import cast @@ -15,6 +19,7 @@ CanonicalJson, ComponentIdentity, FailureCode, + HostWeightError, HostWeightFailure, HostWeightLease, HostWeightLeaseCarrier, @@ -40,12 +45,14 @@ StoreStatus, TensorWriteSpec, ValidationLevel, + WaitPolicy, WeightArtifactIdentity, WeightProductionSpec, WeightRepresentation, WeightSourceIdentity, ) from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -518,6 +525,107 @@ def test_retryable_store_initialization_failure_obeys_runtime_mode( assert resolution.report.attempts[0].failure.retryable +def _initialize_contended_runtime(config: HostWeightRuntimeConfig, output: Connection) -> None: + try: + output.send("starting") + started = time.monotonic() + runtime = HostWeightRuntime.from_config(config) + resolution = runtime.resolve(_identity()) + try: + output.send((time.monotonic() - started, resolution.report)) + finally: + if resolution.lease is not None: + resolution.lease.close() + finally: + output.close() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_domain_lock_timeout_and_fresh_runtime_recovery( + tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome +) -> None: + config = HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path), wait=WaitPolicy(0.1)) + initial = HostWeightRuntime.from_config(config) + built = initial.resolve(_identity(), producer=CountingProducer(_identity())) + assert built.lease is not None + built.lease.close() + metadata = {name: (tmp_path / name).read_bytes() for name in ("domain.json", "domain-policy.json")} + lock_path = tmp_path / "locks" / "domain-init.lock" + ctx = mp.get_context("spawn") + reader, writer = ctx.Pipe(duplex=False) + process = ctx.Process(target=_initialize_contended_runtime, args=(config, writer)) + + with FileLock(lock_path, exclusive=True, deadline=None): + process.start() + writer.close() + try: + # Import/bootstrap is outside the short coordination budget. + assert reader.poll(60), "waiter did not start" + assert reader.recv() == "starting" + assert reader.poll(5), "domain initialization did not respect the coordination timeout" + elapsed, report = reader.recv() + assert 0.1 <= elapsed < 5 + assert report.outcome is expected + failure = report.attempts[0].failure + assert failure is not None + assert failure.stage is ResolutionStage.DOMAIN + assert failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failure.retryable + assert "timed out waiting for" in failure.message + assert "domain-init.lock" in failure.message + assert lock_is_active(lock_path), "waiter must not unlock the owner" + process.join(5) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.kill() + process.join(5) + process.close() + reader.close() + + assert {name: (tmp_path / name).read_bytes() for name in metadata} == metadata + recovered = HostWeightRuntime.from_config(config).resolve(_identity()) + assert recovered.report.outcome is ResolutionOutcome.LOCAL_HIT + assert recovered.lease is not None + recovered.lease.close() + + +def test_direct_store_waits_for_domain_lock_released_before_deadline(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + initial = HostWeightRuntime.from_config(config) + assert initial.store is not None + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None) as owner: + release = threading.Timer(0.1, owner.close) + release.start() + try: + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(5)) + finally: + release.join(5) + assert store.root == tmp_path + + +def test_direct_store_domain_timeout_is_typed(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + HostWeightRuntime.from_config(config) + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None): + with pytest.raises(HostWeightError) as caught: + FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(0.01)) + assert caught.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert caught.value.failure.retryable + assert lock_is_active(lock_path) + assert not lock_is_active(lock_path) + + @pytest.mark.parametrize( ("stage", "code"), [ diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 1a1c0c3ae8e..9d0108bf178 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -739,6 +739,8 @@ class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without 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 + host_weight_runtime_validation: str = "manifest_and_metadata" + 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" @@ -892,9 +894,11 @@ def __post_init__(self) -> None: validate_host_weight_runtime_options( mode=self.host_weight_runtime_mode, root=self.host_weight_runtime_root, + validation=self.host_weight_runtime_validation, ) 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, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 06840e55ac1..cd90e4c490a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -483,6 +483,8 @@ class StageDeployConfig: dlo_resident_layers: int | None = None host_weight_runtime_mode: str | None = None host_weight_runtime_root: str | None = None + host_weight_runtime_validation: 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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 00764ea46d8..ce11fbe3b1f 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -100,7 +100,9 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any] return normalized -def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: +def validate_host_weight_runtime_options( + *, mode: object, root: object, validation: object = "manifest_and_metadata" +) -> None: """Validate HWR policy without touching the configured storage domain. Filesystem locality and store construction belong to the eligible loader @@ -109,6 +111,8 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: """ if mode not in {"disabled", "preferred", "required"}: raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") + if validation not in ("manifest_and_metadata", "full_checksum"): + raise ValueError("host_weight_runtime_validation must be manifest_and_metadata or full_checksum") if mode != "disabled" and (not isinstance(root, str) or not root.strip()): raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") @@ -119,8 +123,15 @@ 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) @@ -831,8 +842,10 @@ class OmniDiffusionConfig: # existing loader/storage path. host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None + host_weight_runtime_validation: str = "manifest_and_metadata" # 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 @@ -1241,9 +1254,11 @@ def __post_init__(self): validate_host_weight_runtime_options( mode=self.host_weight_runtime_mode, root=self.host_weight_runtime_root, + validation=self.host_weight_runtime_validation, ) 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, diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 39e85abcc5c..57b23656cf8 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import concurrent.futures @@ -10,7 +13,7 @@ import time import weakref from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from multiprocessing.synchronize import Event from typing import TYPE_CHECKING, Any, cast @@ -42,6 +45,7 @@ _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0)) _WORKER_SHUTDOWN_GRACE_S = 15.0 _WORKER_TERMINATE_GRACE_S = 5.0 +_WORKER_KILL_GRACE_S = 5.0 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0 @@ -80,33 +84,61 @@ class _ExecutorShutdownCleaner: broadcast_mq: MessageQueue | None = None num_workers: int = 0 processes: list[mp.Process] | None = None + _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) def __call__(self) -> None: """Clean up background resources.""" + # The worker monitor and an explicit shutdown may race. A retry must + # not signal/join the same Process objects while cleanup is in flight. + if not self._lock.acquire(blocking=False): + return + try: + self._cleanup() + finally: + self._lock.release() + + def _cleanup(self) -> None: if self.broadcast_mq is not None: try: for _ in range(self.num_workers): self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0) - self.broadcast_mq = None except Exception as exc: logger.warning("Failed to send shutdown signal: %s", exc) + finally: + self.broadcast_mq = None if self.processes: - join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S - for proc in self.processes: - if not proc.is_alive(): - continue - proc.join(max(0.0, join_deadline - time.monotonic())) - alive = [proc for proc in self.processes if proc.is_alive()] - for proc in alive: - logger.warning("Terminating diffusion worker %s after timeout", proc.name) - proc.terminate() + for action, grace in ( + (None, _WORKER_SHUTDOWN_GRACE_S), + ("terminate", _WORKER_TERMINATE_GRACE_S), + ("kill", _WORKER_KILL_GRACE_S), + ): + if not alive: + break + if action is not None: + for proc in alive: + try: + logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid) + getattr(proc, action)() + except OSError: + logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid) - terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S - for proc in alive: - proc.join(max(0.0, terminate_deadline - time.monotonic())) + deadline = time.monotonic() + grace + for proc in alive: + try: + proc.join(max(0.0, deadline - time.monotonic())) + except OSError: + logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid) + alive = [proc for proc in alive if proc.is_alive()] + + self.processes = alive + if alive: + logger.error( + "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry", + [(proc.name, proc.pid) for proc in alive], + ) class MultiprocDiffusionExecutor(DiffusionExecutor): @@ -1006,8 +1038,12 @@ def check_health(self) -> None: def shutdown(self) -> None: self._closed = True self._pump_stop.set() + cleaner = self._shutdown_cleaner try: - self._finalizer() + if self._finalizer.alive: + self._finalizer() + elif cleaner is not None: + cleaner() finally: pump_threads = getattr(self, "_result_pump_threads", []) for thread in pump_threads: @@ -1031,5 +1067,6 @@ def shutdown(self) -> None: self._rpc_futures.clear() self._output_futures.clear() self._batch_split_map.clear() - self._shutdown_cleaner = None - self._processes = [] + self._processes = (cleaner.processes or []) if cleaner is not None else [] + if not self._processes: + self._shutdown_cleaner = None diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py index 4b98535558d..01427909d91 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_loader.py +++ b/vllm_omni/diffusion/model_loader/host_weight_loader.py @@ -26,6 +26,7 @@ logger = init_logger(__name__) if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext from vllm_omni.diffusion.model_loader.host_weights.source_identity import ( NodeSourceDigestCache, @@ -40,6 +41,8 @@ class _HWRCommitError(RuntimeError): class HWRLoaderMixin: """Optional final-layout HWR behavior shared by diffusion loaders.""" + od_config: OmniDiffusionConfig + @staticmethod def _identity_value(value: object) -> object: """Convert config objects into deterministic identity metadata.""" @@ -319,10 +322,12 @@ def _resolve_hwr( HostWeightLeaseCarrier, HostWeightRuntime, HostWeightRuntimeConfig, + IntegrityPolicy, ProductionPolicy, ResolutionOutcome, RuntimeMode, StorageDomainPolicy, + ValidationLevel, ) from vllm_omni.host_weight_runtime.filesystem import detect_storage_class @@ -363,6 +368,11 @@ def _resolve_hwr( HostWeightRuntimeConfig( mode=mode, domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + integrity=IntegrityPolicy( + local_lookup=ValidationLevel( + getattr(self.od_config, "host_weight_runtime_validation", "manifest_and_metadata") + ), + ), production=ProductionPolicy( allow_local_build=False, allow_post_load_publish=True, diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index bb047a7ebc0..c1afb2c17e1 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -73,12 +73,15 @@ 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: @@ -174,7 +177,9 @@ def from_od_config(cls, od_config: OmniDiffusionConfig) -> "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, @@ -210,6 +215,7 @@ def from_od_config(cls, od_config: OmniDiffusionConfig) -> "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, ) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index 418b5d9fb0d..b5347b6c4f4 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -6,12 +6,14 @@ 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 @@ -21,6 +23,8 @@ HostRegistrationError, ) +logger = init_logger(__name__) + _CUDA_HOST_REGISTER_READ_ONLY = 0x08 _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113 @@ -187,7 +191,9 @@ def create( registered: list[_AddressRange] = [] try: - for region in mapped: + 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() error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY) if int(error) != 0: raise HostRegistrationError( @@ -195,6 +201,9 @@ def create( 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] = [] diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e626d87a543..f04060e0b6e 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1080,6 +1080,9 @@ def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bo 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 @@ -1089,6 +1092,12 @@ def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bo 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( diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 812f4c3e7d2..7f918ece33e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -582,6 +582,8 @@ class OrchestratorArgs: dlo_resident_layers: int = 0 host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None + host_weight_runtime_validation: str = "manifest_and_metadata" + dlo_host_registration_mode: str = "auto" dlo_host_registration_limit_gib: float = 0.0 boundary_ratio: float | None = None flow_shift: float | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 608d605e9b1..4c32297ab8b 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1145,6 +1145,8 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: "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"), + "host_weight_runtime_validation": kwargs.get("host_weight_runtime_validation", "manifest_and_metadata"), + "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": ( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ea56acb1af3..01961acc69f 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -770,6 +770,15 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu "preferred first." ), ) + omni_config_group.add_argument( + "--host-weight-runtime-validation", + choices=("manifest_and_metadata", "full_checksum"), + default="manifest_and_metadata", + help=( + "HWR warm-acquisition validation. The default checks metadata, not payload corruption; " + "full_checksum reads and hashes every payload before restoring weights." + ), + ) omni_config_group.add_argument( "--host-weight-runtime-root", type=str, @@ -780,6 +789,17 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu "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, diff --git a/vllm_omni/host_weight_runtime/config.py b/vllm_omni/host_weight_runtime/config.py index 0bacd69d200..6ab6c7de771 100644 --- a/vllm_omni/host_weight_runtime/config.py +++ b/vllm_omni/host_weight_runtime/config.py @@ -8,6 +8,7 @@ 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 @@ def __post_init__(self) -> None: 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) diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 0f01e36dac5..34403626676 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -23,7 +23,7 @@ import torch from safetensors import SafetensorError, safe_open -from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel +from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel, WaitPolicy from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json from ..inspection import ( @@ -429,6 +429,8 @@ def __init__( domain: StorageDomainPolicy, capacity: CapacityPolicy, integrity: IntegrityPolicy, + *, + wait: WaitPolicy = WaitPolicy(), ) -> None: self.domain_policy = domain self.capacity_policy = capacity @@ -464,7 +466,7 @@ def __init__( ) ) try: - self._initialize_domain() + self._initialize_domain(deadline=time.monotonic() + wait.coordination_timeout_seconds) except HostWeightError: raise except OSError as exc: @@ -485,7 +487,7 @@ def __init__( ) ) from exc - def _initialize_domain(self) -> None: + def _initialize_domain(self, *, deadline: float) -> None: try: self.filesystem_type = detect_filesystem_type(self.root) detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type) @@ -545,7 +547,7 @@ def _initialize_domain(self) -> None: ) ) - with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None): + with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=deadline): domain_path = self.root / _DOMAIN_FILE if domain_path.exists(): domain_metadata = _read_json_file(domain_path) @@ -627,7 +629,7 @@ def _initialize_domain(self) -> None: 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 +637,9 @@ def _capacity_document(self, *, policy_version: int) -> dict[str, object]: "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 +1018,40 @@ def get_or_build( 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: @@ -1247,8 +1286,10 @@ def _check_capacity(self, projected_artifact_bytes: int, additional_bytes: int) retryable=True, ) ) - store_bytes = self._tree_bytes(self.root) - if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes: + if ( + policy.max_store_bytes is not None + and self._tree_bytes(self.root) + additional_bytes > policy.max_store_bytes + ): raise HostWeightError( _failure( ResolutionStage.CAPACITY, @@ -1466,43 +1507,45 @@ def _tree_bytes(root: Path) -> int: total += file_stat.st_size return total - def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None: - """Explicitly remove an inactive artifact without waiting for leases.""" + @contextlib.contextmanager + def _cleanup_locks(self, key: str) -> Iterator[None]: + """Exclude cooperative builders and leases without waiting for them.""" try: - try: - build_lock = FileLock( - self._build_lock_path(identity.key), - exclusive=True, - deadline=None, - nonblocking=True, - ) - except FileLockTimeoutError: - return _failure( + build_lock = FileLock(self._build_lock_path(key), exclusive=True, deadline=None, nonblocking=True) + except FileLockTimeoutError as exc: + raise HostWeightError( + _failure( ResolutionStage.LIFECYCLE, FailureCode.ACTIVE_BUILD_TIMEOUT, - f"artifact {identity.key} has an active builder", + f"artifact {key} has an active builder", retryable=True, ) - with build_lock: - try: - artifact_lock = FileLock( - self._artifact_lock_path(identity.key), - exclusive=True, - deadline=None, - nonblocking=True, - ) - except FileLockTimeoutError: - return _failure( + ) from exc + with build_lock: + try: + artifact_lock = FileLock(self._artifact_lock_path(key), exclusive=True, deadline=None, nonblocking=True) + except FileLockTimeoutError as exc: + raise HostWeightError( + _failure( ResolutionStage.LIFECYCLE, FailureCode.ACTIVE_LEASE, - f"artifact {identity.key} has an active lease", + f"artifact {key} has an active lease", retryable=True, ) - with artifact_lock: - self._repair_quarantine_transitions_locked(identity.key) - self._quarantine_locked(identity.key, reason="cleanup") - self._remove_cleanup_tombstones_locked(identity.key) - self._remove_deny_locked(identity.key) + ) from exc + with artifact_lock: + yield + + def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None: + """Explicitly remove an inactive artifact without waiting for leases.""" + try: + with self._cleanup_locks(identity.key): + self._repair_quarantine_transitions_locked(identity.key) + self._quarantine_locked(identity.key, reason="cleanup") + self._remove_cleanup_tombstones_locked(identity.key) + self._remove_deny_locked(identity.key) + except HostWeightError as exc: + return exc.failure except OSError as exc: return _failure( ResolutionStage.LIFECYCLE, @@ -1512,6 +1555,50 @@ def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None: ) return None + def cleanup_quarantined(self, storage_name: str) -> HostWeightFailure | None: + """Remove one quarantine inventory entry, preserving the current artifact. + + A retry also finishes earlier cleanup tombstones for the same key. + Active builders or leases are reported rather than interrupted. + """ + key, separator, suffix = storage_name.partition(".") + if ( + not separator + or not suffix + or "\0" in storage_name + or _ARTIFACT_KEY_RE.fullmatch(key) is None + or Path(storage_name).name != storage_name + ): + raise ValueError("quarantined storage_name must be a single inventory basename with an artifact key") + try: + with self._cleanup_locks(key): + entry = self.quarantine_dir / storage_name + exists = _path_exists_without_following(entry) + if exists and not stat.S_ISDIR(entry.lstat().st_mode): + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.QUARANTINE_FAILED, + f"quarantine entry is not a directory: {storage_name}", + ) + self._repair_quarantine_transitions_locked(key) + if exists and not storage_name.startswith(f"{key}.cleanup."): + tombstone = self.quarantine_dir / f"{key}.cleanup.{uuid.uuid4().hex}" + _move_artifact_locked(entry, tombstone) + self._remove_cleanup_tombstones_locked(key) + # A previous attempt may have removed the tombstone but failed + # its parent sync. Retry that sync even when no entry remains. + _fsync_directory(self.quarantine_dir) + except HostWeightError as exc: + return exc.failure + except OSError as exc: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.QUARANTINE_FAILED, + f"failed to remove quarantined entry {storage_name}: {exc}", + retryable=True, + ) + return None + def inspect_domain(self) -> DomainInspection: usage = shutil.disk_usage(self.root) inventory: list[ArtifactInventoryEntry] = [] diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py index fbc726669cf..17ababce9eb 100644 --- a/vllm_omni/host_weight_runtime/runtime.py +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -70,7 +70,7 @@ def from_config( from .filesystem import FilesystemHostWeightStore # noqa: PLC0415 try: - store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity) + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=config.wait) except HostWeightError as exc: return cls( config,