Fix diffusion API initialization (#7868)

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

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