Fix diffusion API initialization (#7868)

Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>

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."""
 

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 ─────────────────────
 
 

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:

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