diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 394f7c04f..51b3154a1 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -36,6 +36,21 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: return OmniDiffusionConfig.from_kwargs(**engine_args) +def test_direct_config_extras_are_independent_dicts(): + first = OmniDiffusionConfig(model="test-model") + second = OmniDiffusionConfig(model="test-model") + + assert first.extras.get("prefix_kv_cache_dtype") is None + first.extras["prefix_kv_cache_dtype"] = "fp8" + assert second.extras == {} + + +def test_direct_config_preserves_explicit_extras(): + config = OmniDiffusionConfig(model="test-model", extras={"prefix_kv_cache_dtype": "fp8"}) + + assert config.extras["prefix_kv_cache_dtype"] == "fp8" + + class TestParallelConfigPropagation: """Core regression tests: parallel_config must survive serialization.""" diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index efe6bf9a3..5577a34a3 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -15,7 +15,7 @@ import zmq from vllm.v1.engine.exceptions import EngineDeadError import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module -from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE @@ -35,6 +35,38 @@ from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] +def _signal_worker_ready(_rank, _config, writer, _broadcast, wake_event, *_args): + assert mp.get_start_method() == "spawn" + wake_event.set() + writer.send({"status": "ready", "result_handle": "ready"}) + writer.close() + + +@pytest.mark.skipif("fork" not in mp.get_all_start_methods(), reason="Requires a fork default context") +def test_executor_spawns_workers_with_compatible_events(mocker): + previous_method = mp.get_start_method(allow_none=True) + executor = object.__new__(MultiprocDiffusionExecutor) + executor.od_config = OmniDiffusionConfig(model="test-model", step_execution=True) + broadcast = mocker.Mock() + broadcast.export_handle.return_value = None + mocker.patch.object(executor, "_init_broadcast_queue", return_value=broadcast) + mocker.patch.object(executor, "_init_result_queue") + mocker.patch.object(executor, "_start_worker_monitor") + mocker.patch.object(WorkerProc, "worker_main", new=_signal_worker_ready) + try: + mp.set_start_method("fork", force=True) + executor._init_executor() + for process in executor._processes: + process.join(timeout=30) + assert process.exitcode == 0 + assert all(event.is_set() for event in executor.wake_events) + assert mp.get_start_method() == "fork" + finally: + if getattr(executor, "_finalizer", None) is not None: + executor.shutdown() + mp.set_start_method(previous_method, force=True) + + # ───────────────────────────────────────────── helpers ───────────────────── diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c30529520..ddf55d151 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -15,7 +15,7 @@ import diffusers import huggingface_hub import torch from PIL import Image -from pydantic import Field, model_validator +from pydantic import model_validator from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger @@ -1045,7 +1045,7 @@ class OmniDiffusionConfig: request_batch_max_wait_ms: float = 0.0 # Supplementary model specific parameters - extras: dict[str, Any] = Field(default_factory=dict) + extras: dict[str, Any] = field(default_factory=dict) @property def is_moe(self) -> bool: diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 39e85abcc..7b3400b21 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 @@ -11,6 +14,7 @@ import time import weakref from collections.abc import Callable from dataclasses import dataclass +from multiprocessing.process import BaseProcess from multiprocessing.synchronize import Event from typing import TYPE_CHECKING, Any, cast @@ -79,7 +83,7 @@ class _ExecutorShutdownCleaner: broadcast_mq: MessageQueue | None = None num_workers: int = 0 - processes: list[mp.Process] | None = None + processes: list[BaseProcess] | None = None def __call__(self) -> None: """Clean up background resources.""" @@ -117,7 +121,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): _rpc_wave_id: int = 0 def _init_executor(self) -> None: - self._processes: list[mp.Process] = [] + self._processes: list[BaseProcess] = [] self._closed = False self._is_failed = False self._failure_callbacks: list[Callable[[], None]] = [] @@ -126,7 +130,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): self._rpc_wave_id: int = 0 num_workers = cast(int, self.od_config.num_gpus) - self.wake_events = [mp.Event() for _ in range(num_workers)] + self._mp_context = mp.get_context("spawn") + self.wake_events = [self._mp_context.Event() for _ in range(num_workers)] self._broadcast_mq = self._init_broadcast_queue(num_workers) broadcast_handle = self._broadcast_mq.export_handle() @@ -313,7 +318,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): self, broadcast_handle: Handle, wake_events: list[Event], - ) -> tuple[list[mp.Process], list[Handle]]: + ) -> tuple[list[BaseProcess], list[Handle]]: od_config = self.od_config logger.info("Starting server...") @@ -322,8 +327,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # N-GPU run oversubscribes the host by N x core_count. Honours a # user-provided OMP_NUM_THREADS. set_multiprocessing_worker_envs() - mp.set_start_method("spawn", force=True) - processes = [] + processes: list[BaseProcess] = [] # Extract worker_extension_cls and custom_pipeline_args from od_config worker_extension_cls = od_config.worker_extension_cls @@ -334,9 +338,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): scheduler_pipe_writers = [] for i in range(num_gpus): - reader, writer = mp.Pipe(duplex=False) + reader, writer = self._mp_context.Pipe(duplex=False) scheduler_pipe_writers.append(writer) - process = mp.Process( + process = self._mp_context.Process( target=WorkerProc.worker_main, args=( i, # rank