diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index a9bb1b0b8..ae229e3dd 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -67,6 +67,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 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f91e6af42..fddb952a1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -193,6 +193,37 @@ 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) + + 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 555cb19bb..cc3067762 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -260,6 +260,78 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( 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/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 616d5486e..53a1f778b 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -503,6 +503,8 @@ 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", @@ -519,6 +521,8 @@ 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 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 1a1c0c3ae..0595507cc 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -739,6 +739,7 @@ class _DiffusionConfigProjection: dlo_resident_layers: int = Field(default=0, ge=0) host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled" host_weight_runtime_root: str | None = None + host_weight_runtime_validation: str = "manifest_and_metadata" 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,6 +893,7 @@ class _DiffusionConfigProjection: 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, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 06840e55a..ae4ee6fc4 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -483,6 +483,7 @@ 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_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 00764ea46..298fa8d97 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") @@ -831,6 +835,7 @@ 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_limit_gib: float = 0.0 @@ -1241,6 +1246,7 @@ class OmniDiffusionConfig: 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, diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py index 4b9853555..01427909d 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 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, Te 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 @@ class HWRLoaderMixin: HostWeightLeaseCarrier, HostWeightRuntime, HostWeightRuntimeConfig, + IntegrityPolicy, ProductionPolicy, ResolutionOutcome, RuntimeMode, StorageDomainPolicy, + ValidationLevel, ) from vllm_omni.host_weight_runtime.filesystem import detect_storage_class @@ -363,6 +368,11 @@ class HWRLoaderMixin: 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/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 812f4c3e7..bb298b967 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -582,6 +582,7 @@ 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_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 608d605e9..989d47017 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1145,6 +1145,7 @@ class AsyncOmniEngine: "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), + "host_weight_runtime_validation": kwargs.get("host_weight_runtime_validation", "manifest_and_metadata"), "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 ea56acb1a..48354d894 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -770,6 +770,15 @@ class OmniServeCommand(CLISubcommand): "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,