diff --git a/docs/api/README.md b/docs/api/README.md index 2266a5241..f65cbb525 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -5,7 +5,6 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] -- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][] - [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][] - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][] - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][] diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 81392b201..967d0cc6d 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -168,7 +168,6 @@ vllm_omni/ tests/ │ └── arg_utils.py │ └── test_arg_utils.py ⬜ │ ├── entrypoints/ → ├── entrypoints/ -│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅ │ ├── stage_utils.py │ ├── test_stage_utils.py ✅ │ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve) │ │ └── ... │ │ └── test_*.py ⬜ diff --git a/docs/contributing/ci/tests_style.md b/docs/contributing/ci/tests_style.md index 0b07c5ffe..8b10cf4cc 100644 --- a/docs/contributing/ci/tests_style.md +++ b/docs/contributing/ci/tests_style.md @@ -73,7 +73,6 @@ vllm_omni/ tests/ │ └── arg_utils.py │ └── test_arg_utils.py ⬜ │ ├── entrypoints/ → ├── entrypoints/ -│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅ │ ├── stage_utils.py │ ├── test_stage_utils.py ✅ │ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve) │ │ └── ... │ │ └── test_*.py ⬜ diff --git a/docs/design/module/async_omni_architecture.md b/docs/design/module/async_omni_architecture.md index 59275c556..92b13a3da 100644 --- a/docs/design/module/async_omni_architecture.md +++ b/docs/design/module/async_omni_architecture.md @@ -69,7 +69,7 @@ [5] Orchestrator._orchestration_loop (loop) -> poll stage output - llm stage: await get_output_async() - - diffusion stage: get_diffusion_output_async() + - diffusion stage: get_diffusion_output_nowait() -> (llm stage) output_processors[i].process_outputs(...) -> _route_output(...) -> if finished and not final_stage and non-async-chunk: @@ -112,7 +112,7 @@ sequenceDiagram ORCH->>S0: add_request_async loop poll route forward - ORCH->>S0: get_output_async / get_diffusion_output_async + ORCH->>S0: get_output_async / get_diffusion_output_nowait ORCH->>ORCH: _route_output alt need forward to next stage ORCH->>SN: add_request_async diff --git a/tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py b/tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py index e5c738726..d5f82f893 100644 --- a/tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py +++ b/tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py @@ -509,7 +509,7 @@ def test_diffusion_batching_async_explicit_batch(model_name: str): all prompts in a single engine call and returns a single combined result. The list-prompt path routes through the orchestrator's - ``add_batch_request_async`` → ``AsyncOmniDiffusion.generate_batch`` + ``add_batch_request_async`` → ``AsyncOmni.generate_batch`` and yields ONE ``OmniRequestOutput`` with ALL images combined. """ diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 28b44e9bd..9f47fd449 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -113,6 +113,7 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): executor_class=object, engine_manager=object(), coordinator=object(), + proc=None, addresses=types.SimpleNamespace( inputs=["inproc://input"], outputs=["inproc://output"], diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 7d2a67e73..d68143dae 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -106,7 +106,7 @@ def test_encode_image_base64(): class MockGenerationResult: - """Mock result object from AsyncOmniDiffusion.generate()""" + """Mock result object from AsyncOmni.generate()""" def __init__(self, images): self.images = images diff --git a/tests/entrypoints/test_async_omni_diffusion.py b/tests/entrypoints/test_async_omni_diffusion.py deleted file mode 100644 index c8aaae4f9..000000000 --- a/tests/entrypoints/test_async_omni_diffusion.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import threading -from concurrent.futures import ThreadPoolExecutor -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -import vllm_omni.diffusion.stage_diffusion_client as stage_diffusion_client_module -from vllm_omni.diffusion.data import DiffusionRequestAbortedError -from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient -from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_get_diffusion_od_config_returns_direct_config(): - diffusion = object.__new__(AsyncOmniDiffusion) - diffusion.od_config = object() - - assert diffusion.get_diffusion_od_config() is diffusion.od_config - - -def test_async_omni_diffusion_generate_aborts_engine_on_cancel(): - async def run_test(): - started = threading.Event() - release = threading.Event() - abort = Mock() - - def step(request): - del request - started.set() - release.wait(timeout=5) - return [SimpleNamespace(request_id="req-1")] - - diffusion = object.__new__(AsyncOmniDiffusion) - diffusion.engine = SimpleNamespace(step=step, abort=abort) - diffusion._executor = ThreadPoolExecutor(max_workers=1) - - task = asyncio.create_task( - diffusion.generate( - prompt="hello", - sampling_params=OmniDiffusionSamplingParams(), - request_id="req-1", - ) - ) - try: - assert await asyncio.to_thread(started.wait, 1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - finally: - release.set() - diffusion._executor.shutdown(wait=True) - - abort.assert_called_once_with("req-1") - - asyncio.run(run_test()) - - -def test_stage_diffusion_client_abort_requests_forwards_to_engine(): - async def run_test(): - aborted_request_ids: list[list[str]] = [] - - async def abort(request_ids): - aborted_request_ids.append(request_ids) - - client = object.__new__(StageDiffusionClient) - client._engine = SimpleNamespace(abort=abort) - client._tasks = {} - - task = asyncio.create_task(asyncio.sleep(60)) - client._tasks["req-1"] = task - - await client.abort_requests_async(["req-1", "req-2"]) - - with pytest.raises(asyncio.CancelledError): - await task - assert client._tasks == {} - assert aborted_request_ids == [["req-1", "req-2"]] - - asyncio.run(run_test()) - - -def test_stage_diffusion_client_run_treats_abort_as_normal_path(monkeypatch): - async def run_test(): - async def generate(prompt, sampling_params, request_id): - del prompt, sampling_params - raise DiffusionRequestAbortedError(f"Request {request_id} aborted.") - - info = Mock() - exception = Mock() - monkeypatch.setattr(stage_diffusion_client_module.logger, "info", info) - monkeypatch.setattr(stage_diffusion_client_module.logger, "exception", exception) - - client = object.__new__(StageDiffusionClient) - client.stage_id = 3 - client._engine = SimpleNamespace(generate=generate) - client._output_queue = asyncio.Queue() - client._tasks = {"req-1": object()} - - await client._run("req-1", "prompt", OmniDiffusionSamplingParams()) - - assert client._output_queue.empty() - assert client._tasks == {} - info.assert_called_once() - exception.assert_not_called() - - asyncio.run(run_test()) diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index d3d7b3aff..9aafc1cf1 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -31,6 +31,12 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]: import numpy as np tensor = tensor.detach().cpu().contiguous() + original_dtype = tensor.dtype + # NumPy does not support bfloat16; promote to float32 for the SHM + # transfer and record the original dtype so _tensor_from_shm can + # convert back. The round-trip is lossless for bfloat16 values. + if original_dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) arr = tensor.numpy() nbytes = arr.nbytes shm = shared_memory.SharedMemory(create=True, size=nbytes) @@ -40,7 +46,7 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]: "__tensor_shm__": True, "name": shm.name, "shape": list(tensor.shape), - "torch_dtype": str(tensor.dtype), + "torch_dtype": str(original_dtype), "numpy_dtype": str(arr.dtype), "nbytes": nbytes, } @@ -59,6 +65,13 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: np_dtype = np.dtype(handle["numpy_dtype"]) arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]]) tensor = torch.from_numpy(arr.copy()) + # Restore the original dtype if it differs from the numpy-compatible + # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16). + torch_dtype_str = handle.get("torch_dtype", "") + if torch_dtype_str: + original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None) + if original_dtype is not None and tensor.dtype != original_dtype: + tensor = tensor.to(original_dtype) finally: shm.close() shm.unlink() diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index ddad2f9f3..db13f99aa 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -1,20 +1,30 @@ """Stage Diffusion Client for vLLM-Omni multi-stage runtime. -Wraps AsyncOmniDiffusion to expose the same interface the Orchestrator -expects from any stage client. +Spawns StageDiffusionProc in a subprocess and communicates via ZMQ +(PUSH/PULL) to expose the same interface the Orchestrator expects +from any stage client. """ from __future__ import annotations import asyncio import time +import uuid +from dataclasses import fields, is_dataclass from typing import TYPE_CHECKING, Any +import zmq from vllm.logger import init_logger -from vllm_omni.diffusion.data import DiffusionRequestAbortedError -from vllm_omni.engine.stage_init_utils import StageMetadata -from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion +from vllm_omni.diffusion.stage_diffusion_proc import ( + complete_diffusion_handshake, + spawn_diffusion_proc, +) +from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, +) +from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -25,11 +35,12 @@ logger = init_logger(__name__) class StageDiffusionClient: - """Wraps AsyncOmniDiffusion for use inside the Orchestrator. + """Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator. Exposes the same attributes and async methods the Orchestrator uses on StageEngineCoreClient, but routes execution through - DiffusionEngine instead of vLLM EngineCore. + a StageDiffusionProc subprocess instead of running the diffusion + engine in-process. """ stage_type: str = "diffusion" @@ -48,56 +59,137 @@ class StageDiffusionClient: self.custom_process_input_func = metadata.custom_process_input_func self.engine_input_source = metadata.engine_input_source - self._engine = AsyncOmniDiffusion(model=model, od_config=od_config, batch_size=batch_size) + # Spawn StageDiffusionProc subprocess and wait for READY. + proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config) + complete_diffusion_handshake(proc, handshake_address) + self._proc = proc + + # ZMQ sockets (sync) for communicating with the subprocess. + self._zmq_ctx = zmq.Context() + self._request_socket = self._zmq_ctx.socket(zmq.PUSH) + self._request_socket.connect(request_address) + self._response_socket = self._zmq_ctx.socket(zmq.PULL) + self._response_socket.connect(response_address) + + self._encoder = OmniMsgpackEncoder() + self._decoder = OmniMsgpackDecoder() + + # Buffers for demultiplexing response messages. self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + self._rpc_results: dict[str, Any] = {} + self._pending_rpcs: set[str] = set() self._tasks: dict[str, asyncio.Task] = {} + self._shutting_down = False logger.info("[StageDiffusionClient] Stage-%s initialized (batch_size=%d)", self.stage_id, batch_size) + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _drain_responses(self) -> None: + """Non-blocking drain of all available responses from the subprocess.""" + while True: + try: + raw = self._response_socket.recv(zmq.NOBLOCK) + except zmq.Again: + break + + msg = self._decoder.decode(raw) + msg_type = msg.get("type") + + if msg_type == "result": + self._output_queue.put_nowait(msg["output"]) + elif msg_type == "rpc_result": + self._rpc_results[msg["rpc_id"]] = msg["result"] + elif msg_type == "error": + req_id = msg.get("request_id") + rpc_id = msg.get("rpc_id") + error_msg = msg.get("error") + logger.error( + "[StageDiffusionClient] Stage-%s subprocess error for %s: %s", + self.stage_id, + rpc_id or req_id, + error_msg, + ) + # Route RPC errors so collective_rpc_async can unblock. + if rpc_id is not None and rpc_id in self._pending_rpcs: + self._rpc_results[rpc_id] = { + "error": True, + "reason": error_msg, + } + + # Fields that are subprocess-local and cannot be serialized across + # process boundaries. They are recreated in the subprocess with + # their default values. + _NON_SERIALIZABLE_FIELDS = frozenset( + { + "generator", # torch.Generator — recreated from seed + "modules", # model components — loaded in subprocess + } + ) + + @staticmethod + def _sampling_params_to_dict(sampling_params: Any) -> dict[str, Any]: + """Convert sampling params to a plain dict for serialization. + + Uses ``dataclasses.fields`` + ``getattr`` instead of ``asdict`` + to avoid deep-copying large tensors, and skips fields that + cannot cross process boundaries. + + When a ``torch.Generator`` is present but ``seed`` is not set, + the generator's initial seed is extracted so the subprocess can + recreate an equivalent generator via ``diffusion_model_runner``. + """ + if is_dataclass(sampling_params) and not isinstance(sampling_params, type): + result = { + f.name: getattr(sampling_params, f.name) + for f in fields(sampling_params) + if f.name not in StageDiffusionClient._NON_SERIALIZABLE_FIELDS + } + elif not isinstance(sampling_params, dict): + raise TypeError(f"sampling_params is not a dict but {sampling_params.__class__.__name__}") + else: + result = { + k: v for k, v in sampling_params.items() if k not in StageDiffusionClient._NON_SERIALIZABLE_FIELDS + } + + # Preserve the generator's seed across the process boundary so + # the subprocess can recreate deterministic random state. + if result.get("seed") is None: + generator = ( + getattr(sampling_params, "generator", None) + if not isinstance(sampling_params, dict) + else sampling_params.get("generator") + ) + if generator is not None: + if isinstance(generator, list) and generator: + generator = generator[0] + if hasattr(generator, "initial_seed"): + result["seed"] = generator.initial_seed() + + return result + + # ------------------------------------------------------------------ + # Public API (matches the interface the Orchestrator expects) + # ------------------------------------------------------------------ + async def add_request_async( self, request_id: str, prompt: OmniPromptType, sampling_params: OmniDiffusionSamplingParams, ) -> None: - task = asyncio.create_task( - self._run(request_id, prompt, sampling_params), - name=f"diffusion-{request_id}", + self._request_socket.send( + self._encoder.encode( + { + "type": "add_request", + "request_id": request_id, + "prompt": prompt, + "sampling_params": self._sampling_params_to_dict(sampling_params), + } + ) ) - self._tasks[request_id] = task - - async def _run( - self, - request_id: str, - prompt: OmniPromptType, - sampling_params: OmniDiffusionSamplingParams, - ) -> None: - try: - result = await self._engine.generate(prompt, sampling_params, request_id) - await self._output_queue.put(result) - except asyncio.CancelledError: - logger.info( - "[StageDiffusionClient] Stage-%s req=%s cancelled", - self.stage_id, - request_id, - ) - raise - except DiffusionRequestAbortedError as e: - logger.info( - "[StageDiffusionClient] Stage-%s req=%s aborted: %s", - self.stage_id, - request_id, - e, - ) - except Exception as e: - logger.exception( - "[StageDiffusionClient] Stage-%s req=%s failed: %s", - self.stage_id, - request_id, - e, - ) - finally: - self._tasks.pop(request_id, None) # TODO(Long): Temporary solution to boost performance of diffusion stages. # Remove this after scheduling algorithm is implemented @@ -126,12 +218,16 @@ class StageDiffusionClient: sampling_params: OmniDiffusionSamplingParams, ) -> None: try: - result = await self._engine.generate_batch( - prompts, - sampling_params, - request_id, + self._request_socket.send( + self._encoder.encode( + { + "type": "add_batch_request", + "request_id": request_id, + "prompts": prompts, + "sampling_params": self._sampling_params_to_dict(sampling_params), + } + ) ) - await self._output_queue.put(result) except Exception as e: logger.exception( "[StageDiffusionClient] Stage-%s batch req=%s failed: %s", @@ -142,18 +238,24 @@ class StageDiffusionClient: finally: self._tasks.pop(request_id, None) - def get_diffusion_output_async(self) -> OmniRequestOutput | None: + def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: + self._drain_responses() try: return self._output_queue.get_nowait() except asyncio.QueueEmpty: + if not self._shutting_down and self._proc is not None and not self._proc.is_alive(): + raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {self._proc.exitcode})") return None async def abort_requests_async(self, request_ids: list[str]) -> None: - for rid in request_ids: - task = self._tasks.pop(rid, None) - if task: - task.cancel() - await self._engine.abort(request_ids) + self._request_socket.send( + self._encoder.encode( + { + "type": "abort", + "request_ids": list(request_ids), + } + ) + ) async def collective_rpc_async( self, @@ -162,60 +264,66 @@ class StageDiffusionClient: args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> Any: - """Best-effort control RPC shim for diffusion stages. - - TODO(AsyncOmni): add dedicated wrappers on AsyncOmniDiffusion for the - remaining control APIs instead of reaching into its underlying engine. - """ - kwargs = kwargs or {} - - # Handle profile method: inject stage_id into profile_prefix for diffusion stages + """Forward control RPCs to the diffusion subprocess.""" + # Inject a default profile_prefix that includes stage_id when profiling. if method == "profile": - target = getattr(self._engine, method, None) - if target is None: - return { - "supported": False, - "todo": True, - "reason": f"AsyncOmniDiffusion.{method} is not implemented", - } - # Extract is_start and profile_prefix from args - is_start = args[0] if args else True - profile_prefix = args[1] if len(args) > 1 else None - # Generate profile_prefix with stage_id if starting and no prefix provided + args_list = list(args) + is_start = args_list[0] if args_list else True + profile_prefix = args_list[1] if len(args_list) > 1 else None if is_start and profile_prefix is None: profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" - result = target(is_start, profile_prefix) - if timeout is not None: - return await asyncio.wait_for(result, timeout=timeout) - return await result + if len(args_list) > 1: + args_list[1] = profile_prefix + else: + args_list.append(profile_prefix) + args = tuple(args_list) - if method in {"add_lora", "remove_lora", "list_loras", "pin_lora"}: - target = getattr(self._engine, method, None) - if target is None: - return { - "supported": False, - "todo": True, - "reason": f"AsyncOmniDiffusion.{method} is not implemented", + kwargs = kwargs or {} + rpc_id = uuid.uuid4().hex + self._pending_rpcs.add(rpc_id) + + self._request_socket.send( + self._encoder.encode( + { + "type": "collective_rpc", + "rpc_id": rpc_id, + "method": method, + "timeout": timeout, + "args": list(args), + "kwargs": kwargs, } - result = target(*args, **kwargs) - if timeout is not None: - return await asyncio.wait_for(result, timeout=timeout) - return await result - - # Fall back to collective RPC for other methods - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._engine._executor, - self._engine.engine.collective_rpc, - method, - timeout, - args, - kwargs, - None, + ) ) + deadline = time.monotonic() + timeout if timeout else None + # Wait for the matching RPC response, buffering result messages. + try: + while True: + self._drain_responses() + if rpc_id in self._rpc_results: + return self._rpc_results.pop(rpc_id) + if self._proc is not None and not self._proc.is_alive(): + raise RuntimeError( + f"StageDiffusionProc died while waiting for " + f"collective_rpc '{method}' (exit code {self._proc.exitcode})" + ) + if deadline and time.monotonic() > deadline: + raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s") + await asyncio.sleep(0.01) + finally: + self._pending_rpcs.discard(rpc_id) + def shutdown(self) -> None: - for task in self._tasks.values(): - task.cancel() - self._tasks.clear() - self._engine.close() + self._shutting_down = True + try: + self._request_socket.send(self._encoder.encode({"type": "shutdown"})) + except Exception: + pass + + if self._proc is not None and self._proc.is_alive(): + self._proc.join(timeout=10) + terminate_alive_proc(self._proc) + + self._request_socket.close(linger=0) + self._response_socket.close(linger=0) + self._zmq_ctx.term() diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py new file mode 100644 index 000000000..8677da037 --- /dev/null +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -0,0 +1,604 @@ +"""Subprocess entry point for the diffusion engine. + +StageDiffusionProc runs DiffusionEngine in a child process, +communicating with StageDiffusionClient via ZMQ (PUSH/PULL). +""" + +from __future__ import annotations + +import asyncio +import signal +import time +from concurrent.futures import ThreadPoolExecutor +from multiprocessing.process import BaseProcess +from typing import TYPE_CHECKING, Any + +import msgspec +import zmq +import zmq.asyncio +from vllm.logger import init_logger +from vllm.transformers_utils.config import get_hf_file_to_dict +from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx +from vllm.utils.system_utils import get_mp_context +from vllm.v1.utils import shutdown + +from vllm_omni.diffusion.data import DiffusionRequestAbortedError, TransformerConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + +_HANDSHAKE_POLL_TIMEOUT_S = 600 + + +class StageDiffusionProc: + """Subprocess entry point for diffusion inference. + + Manages DiffusionEngine lifecycle, async request processing, + and ZMQ-based communication with StageDiffusionClient. + """ + + def __init__(self, model: str, od_config: OmniDiffusionConfig) -> None: + self._model = model + self._od_config = od_config + self._engine: DiffusionEngine | None = None + self._executor: ThreadPoolExecutor | None = None + self._closed = False + + # ------------------------------------------------------------------ + # Initialization + # ------------------------------------------------------------------ + + def initialize(self) -> None: + """Enrich config, create DiffusionEngine and thread pool.""" + self._enrich_config() + self._engine = DiffusionEngine.make_engine(self._od_config) + self._executor = ThreadPoolExecutor(max_workers=1) + logger.info("StageDiffusionProc initialized with model: %s", self._model) + + def _enrich_config(self) -> None: + """Load model metadata from HuggingFace and populate od_config fields. + + Diffusers-style models expose ``model_index.json`` with ``_class_name``. + Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``, + so we fall back to reading that and mapping model_type manually. + """ + od_config = self._od_config + + try: + config_dict = get_hf_file_to_dict("model_index.json", od_config.model) + if config_dict is not None: + if od_config.model_class_name is None: + od_config.model_class_name = config_dict.get("_class_name", None) + od_config.update_multimodal_support() + + tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model) + od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + else: + raise FileNotFoundError("model_index.json not found") + except (AttributeError, OSError, ValueError, FileNotFoundError): + cfg = get_hf_file_to_dict("config.json", od_config.model) + if cfg is None: + raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}") + + od_config.tf_model_config = TransformerConfig.from_dict(cfg) + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + od_config.model_class_name = "BagelPipeline" + od_config.tf_model_config = TransformerConfig() + od_config.update_multimodal_support() + elif model_type == "nextstep": + if od_config.model_class_name is None: + od_config.model_class_name = "NextStep11Pipeline" + od_config.tf_model_config = TransformerConfig() + od_config.update_multimodal_support() + elif architectures and len(architectures) == 1: + od_config.model_class_name = architectures[0] + else: + raise + + # ------------------------------------------------------------------ + # Request processing + # ------------------------------------------------------------------ + + def _reconstruct_sampling_params(self, sampling_params_dict: dict) -> OmniDiffusionSamplingParams: + """Reconstruct OmniDiffusionSamplingParams from a dict, handling LoRA.""" + lora_req = sampling_params_dict.get("lora_request") + if lora_req is not None: + from vllm.lora.request import LoRARequest + + if not isinstance(lora_req, LoRARequest): + sampling_params_dict["lora_request"] = msgspec.convert(lora_req, LoRARequest) + + return OmniDiffusionSamplingParams(**sampling_params_dict) + + async def _process_request( + self, + request_id: str, + prompt: Any, + sampling_params_dict: dict, + ) -> OmniRequestOutput: + """Build a diffusion request and run DiffusionEngine.step().""" + sampling_params = self._reconstruct_sampling_params(sampling_params_dict) + + request = OmniDiffusionRequest( + prompts=[prompt], + sampling_params=sampling_params, + request_ids=[request_id], + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + result = results[0] + if not result.request_id: + result.request_id = request_id + return result + + async def _process_batch_request( + self, + request_id: str, + prompts: list[Any], + sampling_params_dict: dict, + ) -> OmniRequestOutput: + """Build a batched diffusion request and run DiffusionEngine.step(). + + All prompts are processed in a single step() call. The per-prompt + results are merged into one :class:`OmniRequestOutput` whose + ``images`` list contains every generated image, matching the + contract expected by the orchestrator and tests. + """ + sampling_params = self._reconstruct_sampling_params(sampling_params_dict) + + request = OmniDiffusionRequest( + prompts=prompts, + sampling_params=sampling_params, + request_ids=[request_id] * len(prompts), + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + + # Merge per-prompt results into a single combined output. + all_images: list = [] + merged_mm: dict[str, Any] = {} + merged_metrics: dict[str, Any] = {} + merged_durations: dict[str, float] = {} + peak_mem = 0.0 + latents = None + final_output_type = "image" + + for r in results: + all_images.extend(r.images) + merged_mm.update(r._multimodal_output) + merged_metrics.update(r.metrics) + merged_durations.update(r.stage_durations) + peak_mem = max(peak_mem, r.peak_memory_mb) + if latents is None and r.latents is not None: + latents = r.latents + if r.final_output_type != "image": + final_output_type = r.final_output_type + + return OmniRequestOutput.from_diffusion( + request_id=request_id, + images=all_images, + prompt=prompts[0] if len(prompts) == 1 else None, + metrics=merged_metrics, + latents=latents, + multimodal_output=merged_mm or None, + final_output_type=final_output_type, + stage_durations=merged_durations, + peak_memory_mb=peak_mem, + ) + + # ------------------------------------------------------------------ + # Collective RPC dispatch + # ------------------------------------------------------------------ + + async def _handle_collective_rpc( + self, + method: str, + timeout: float | None, + args: tuple, + kwargs: dict, + ) -> Any: + """Dispatch collective RPC calls to DiffusionEngine. + + LoRA methods remap arguments and post-process results to match + the contract that ``AsyncOmni`` provides. + """ + loop = asyncio.get_running_loop() + + if method == "profile": + is_start = args[0] if args else True + profile_prefix = args[1] if len(args) > 1 else None + return await loop.run_in_executor( + self._executor, + self._engine.profile, + is_start, + profile_prefix, + ) + + if method == "add_lora": + # Reconstruct LoRARequest after IPC if needed. + lora_request = args[0] if args else kwargs.get("lora_request") + if lora_request is not None: + from vllm.lora.request import LoRARequest + + if not isinstance(lora_request, LoRARequest): + lora_request = msgspec.convert(lora_request, LoRARequest) + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "add_lora", + timeout, + (), + {"lora_request": lora_request}, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "remove_lora": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "remove_lora", + timeout, + args, + kwargs or {}, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "list_loras": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "list_loras", + timeout, + (), + {}, + None, + ) + if not isinstance(results, list): + return results or [] + merged: set[int] = set() + for part in results: + merged.update(part or []) + return sorted(merged) + + if method == "pin_lora": + lora_id = args[0] if args else kwargs.get("adapter_id") + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "pin_lora", + timeout, + (), + {"adapter_id": lora_id}, + None, + ) + return all(results) if isinstance(results, list) else results + + # Fall back to DiffusionEngine.collective_rpc for all other methods + # (e.g. worker extension RPCs like "test_extension_name"). + return await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + method, + timeout, + args, + kwargs or {}, + None, + ) + + # ------------------------------------------------------------------ + # ZMQ event loop + # ------------------------------------------------------------------ + + async def run_loop( + self, + request_address: str, + response_address: str, + ) -> None: + """Async event loop handling ZMQ messages from StageDiffusionClient.""" + ctx = zmq.asyncio.Context() + + request_socket = ctx.socket(zmq.PULL) + request_socket.bind(request_address) + + response_socket = ctx.socket(zmq.PUSH) + response_socket.bind(response_address) + + encoder = OmniMsgpackEncoder() + decoder = OmniMsgpackDecoder() + + tasks: dict[str, asyncio.Task] = {} + + async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None: + """Process a single diffusion request and send the response.""" + try: + result = await self._process_request(request_id, prompt, sampling_params_dict) + await response_socket.send(encoder.encode({"type": "result", "output": result})) + except DiffusionRequestAbortedError as e: + logger.info( + "request_id: %s aborted: %s", + request_id, + str(e), + ) + except Exception as e: + logger.exception("Diffusion request %s failed: %s", request_id, e) + await response_socket.send( + encoder.encode( + { + "type": "error", + "request_id": request_id, + "error": str(e), + } + ) + ) + finally: + tasks.pop(request_id, None) + + try: + while True: + raw = await request_socket.recv() + msg = decoder.decode(raw) + msg_type = msg.get("type") + + if msg_type == "add_request": + request_id = msg["request_id"] + task = asyncio.create_task( + _dispatch_request( + request_id, + msg["prompt"], + msg["sampling_params"], + ) + ) + tasks[request_id] = task + + elif msg_type == "add_batch_request": + request_id = msg["request_id"] + + async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None: + try: + result = await self._process_batch_request(rid, prompts, sp_dict) + await response_socket.send(encoder.encode({"type": "result", "output": result})) + except DiffusionRequestAbortedError as e: + logger.info( + "request_id: %s aborted: %s", + rid, + str(e), + ) + except Exception as e: + logger.exception("Batch diffusion request %s failed: %s", rid, e) + await response_socket.send( + encoder.encode( + { + "type": "error", + "request_id": rid, + "error": str(e), + } + ) + ) + finally: + tasks.pop(rid, None) + + task = asyncio.create_task( + _dispatch_batch( + request_id, + msg["prompts"], + msg["sampling_params"], + ) + ) + tasks[request_id] = task + + elif msg_type == "abort": + for rid in msg.get("request_ids", []): + task = tasks.pop(rid, None) + if task: + task.cancel() + self._engine.abort(rid) + + elif msg_type == "collective_rpc": + rpc_id = msg["rpc_id"] + try: + result = await self._handle_collective_rpc( + msg["method"], + msg.get("timeout"), + tuple(msg.get("args", ())), + msg.get("kwargs", {}), + ) + await response_socket.send( + encoder.encode( + { + "type": "rpc_result", + "rpc_id": rpc_id, + "result": result, + } + ) + ) + except Exception as e: + logger.exception("Collective RPC %s failed: %s", msg["method"], e) + await response_socket.send( + encoder.encode( + { + "type": "error", + "rpc_id": rpc_id, + "error": str(e), + } + ) + ) + + elif msg_type == "shutdown": + break + + finally: + for task in tasks.values(): + task.cancel() + if tasks: + await asyncio.gather(*tasks.values(), return_exceptions=True) + + request_socket.close() + response_socket.close() + ctx.term() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def close(self) -> None: + """Release engine and thread pool resources.""" + if self._closed: + return + self._closed = True + + if self._engine is not None: + try: + self._engine.close() + except Exception as e: + logger.warning("Error closing diffusion engine: %s", e) + + if self._executor is not None: + try: + self._executor.shutdown(wait=False) + except Exception as e: + logger.warning("Error shutting down executor: %s", e) + + # ------------------------------------------------------------------ + # Subprocess entry point + # ------------------------------------------------------------------ + + @classmethod + def run_diffusion_proc( + cls, + model: str, + od_config: OmniDiffusionConfig, + handshake_address: str, + request_address: str, + response_address: str, + ) -> None: + """Entry point for the diffusion subprocess.""" + shutdown_requested = False + + def signal_handler(signum: int, frame: Any) -> None: + nonlocal shutdown_requested + if not shutdown_requested: + shutdown_requested = True + raise SystemExit() + + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + + proc = cls(model, od_config) + try: + proc.initialize() + + # Send READY via handshake socket + handshake_ctx = zmq.Context() + handshake_socket = handshake_ctx.socket(zmq.DEALER) + handshake_socket.connect(handshake_address) + handshake_socket.send(msgspec.msgpack.encode({"status": "READY"})) + handshake_socket.close() + handshake_ctx.term() + + # Run async event loop + asyncio.run(proc.run_loop(request_address, response_address)) + + except SystemExit: + logger.debug("StageDiffusionProc exiting.") + raise + except Exception: + logger.exception("StageDiffusionProc encountered a fatal error.") + raise + finally: + proc.close() + + +# -- Free functions for backward compatibility with StageDiffusionClient ------ + + +def spawn_diffusion_proc( + model: str, + od_config: OmniDiffusionConfig, +) -> tuple[BaseProcess, str, str, str]: + """Spawn a StageDiffusionProc subprocess. + + Returns ``(proc, handshake_address, request_address, response_address)``. + """ + handshake_address = get_open_zmq_ipc_path() + request_address = get_open_zmq_ipc_path() + response_address = get_open_zmq_ipc_path() + + ctx = get_mp_context() + proc = ctx.Process( + target=StageDiffusionProc.run_diffusion_proc, + name="StageDiffusionProc", + kwargs={ + "model": model, + "od_config": od_config, + "handshake_address": handshake_address, + "request_address": request_address, + "response_address": response_address, + }, + ) + proc.start() + # Wait for the process to become alive before returning. + deadline = time.monotonic() + 10 + while not proc.is_alive(): + if proc.exitcode is not None: + raise RuntimeError(f"StageDiffusionProc failed to start (exit code {proc.exitcode})") + if time.monotonic() > deadline: + raise TimeoutError("StageDiffusionProc did not become alive within 10s") + time.sleep(0.01) + return proc, handshake_address, request_address, response_address + + +def complete_diffusion_handshake( + proc: BaseProcess, + handshake_address: str, +) -> None: + """Wait for the diffusion subprocess to signal READY. + + On failure the process is terminated before re-raising. + """ + try: + _perform_diffusion_handshake(proc, handshake_address) + except Exception: + shutdown([proc]) + raise + + +def _perform_diffusion_handshake( + proc: BaseProcess, + handshake_address: str, +) -> None: + """Run the handshake with the diffusion subprocess.""" + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + poller = zmq.Poller() + poller.register(handshake_socket, zmq.POLLIN) + poller.register(proc.sentinel, zmq.POLLIN) + + timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000 + while True: + events = dict(poller.poll(timeout=timeout_ms)) + if not events: + raise TimeoutError("Timed out waiting for READY from StageDiffusionProc") + if handshake_socket in events: + identity, raw = handshake_socket.recv_multipart() + msg = msgspec.msgpack.decode(raw) + if msg.get("status") == "READY": + return + raise RuntimeError(f"Expected READY, got: {msg}") + if proc.exitcode is not None: + raise RuntimeError(f"StageDiffusionProc died during handshake (exit code {proc.exitcode})") diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index c998870ce..c987106fe 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -32,7 +32,6 @@ from vllm.logger import init_logger from vllm.tokenizers import cached_tokenizer_from_config from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor -from vllm.v1.engine.utils import get_engine_zmq_addresses, launch_core_engines from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.distributed.omni_connectors.utils.initialization import ( @@ -45,6 +44,10 @@ from vllm_omni.engine.orchestrator import Orchestrator from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_engine_core_proc import ( + complete_stage_handshake, + spawn_stage_core, +) from vllm_omni.engine.stage_init_utils import ( StartedLlmStage, acquire_device_locks, @@ -334,21 +337,17 @@ class AsyncOmniEngine: engine_args_dict, stage_init_timeout, ) - addresses = get_engine_zmq_addresses(vllm_config) - launch_cm = launch_core_engines( + addresses, proc, handshake_address = spawn_stage_core( vllm_config=vllm_config, executor_class=executor_class, log_stats=False, - addresses=addresses, ) - engine_manager, coordinator, addresses = launch_cm.__enter__() started_stage = StartedLlmStage( stage_id=metadata.stage_id, metadata=metadata, vllm_config=vllm_config, executor_class=executor_class, - engine_manager=engine_manager, - coordinator=coordinator, + proc=proc, addresses=addresses, ) finally: @@ -358,7 +357,7 @@ class AsyncOmniEngine: current_omni_platform.set_device_control_env_var(previous_visible_devices) logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) - launch_cm.__exit__(None, None, None) + complete_stage_handshake(proc, handshake_address, addresses, vllm_config) logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) assert started_stage is not None return started_stage @@ -389,11 +388,9 @@ class AsyncOmniEngine: executor_class=started.executor_class, metadata=started.metadata, client_addresses=client_addresses, - engine_manager=started.engine_manager, - coordinator=started.coordinator, + proc=started.proc, ) - started.engine_manager = None - started.coordinator = None + started.proc = None except Exception: close_started_llm_stage(started) raise diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4a85a2c6c..8ea9a5096 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -241,7 +241,7 @@ class Orchestrator: # the output format in the future to simplify the processing logic in Orchestrator. stage_client = self.stage_clients[stage_id] if stage_client.stage_type == "diffusion": - output = stage_client.get_diffusion_output_async() + output = stage_client.get_diffusion_output_nowait() if output is not None: idle = False req_state = self.request_states.get(output.request_id) diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 284cc2d31..e08ce7801 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -25,13 +25,13 @@ logger = init_logger(__name__) class StageEngineCoreClient(AsyncMPClient): """Stage async client that inherits from vLLM's AsyncMPClient. - Fully reuses AsyncMPClient.__init__ for: + Fully reuses AsyncMPClient for: - ZMQ setup, sockets - - launch_core_engines() -> EngineCoreProc - outputs_queue, output_queue_task - - All utility methods (shutdown, get_output_async, abort_requests_async, etc.) + - All utility methods (get_output_async, abort_requests_async, etc.) - This is the async version of StageMPClient, designed for use with AsyncOmniEngine. + The subprocess is spawned externally via ``spawn_stage_core`` / + ``complete_stage_handshake`` from *stage_engine_core_proc.py*. """ def __init__( @@ -40,6 +40,7 @@ class StageEngineCoreClient(AsyncMPClient): executor_class: type, log_stats: bool = False, client_addresses: dict[str, str] | None = None, + proc: Any = None, client_count: int = 1, client_index: int = 0, *, @@ -53,6 +54,11 @@ class StageEngineCoreClient(AsyncMPClient): engine args building, device locking) is done by the Orchestrator via helpers in stage_init_utils.py. This constructor just stores metadata and calls super().__init__(). + + The subprocess is spawned externally via ``spawn_stage_core`` / + ``complete_stage_handshake`` (see *stage_engine_core_proc.py*). + The resulting ``proc`` handle is passed in so this client can + manage the process lifecycle on shutdown. """ # -------- Stage metadata (public fields used at runtime) -------- if metadata is not None: @@ -69,6 +75,7 @@ class StageEngineCoreClient(AsyncMPClient): self.model_stage = metadata.model_stage self.engine_outputs: Any = None + self._proc = proc logger.info( "[StageEngineCoreClient] Stage-%s initializing EngineCore", @@ -83,10 +90,6 @@ class StageEngineCoreClient(AsyncMPClient): client_count=client_count, client_index=client_index, ) - if engine_manager is not None: - self.resources.engine_manager = engine_manager - if coordinator is not None: - self.resources.coordinator = coordinator except Exception: logger.exception( "[StageEngineCoreClient] Stage-%s EngineCore init failed", @@ -173,3 +176,12 @@ class StageEngineCoreClient(AsyncMPClient): args=args, kwargs=kwargs, ) + + def shutdown(self) -> None: + """Shutdown ZMQ connections and the subprocess.""" + super().shutdown() + if self._proc is not None and self._proc.is_alive(): + self._proc.terminate() + self._proc.join(timeout=5) + if self._proc.is_alive(): + self._proc.kill() diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py new file mode 100644 index 000000000..05d8f107c --- /dev/null +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -0,0 +1,206 @@ +""" +Stage Core Process for vLLM-Omni V1 architecture. + +StageEngineCoreProc inherits from vLLM's EngineCoreProc and runs the engine core +busy loop in a subprocess, communicating with StageEngineCoreClient via ZMQ. +""" + +from __future__ import annotations + +import signal +from multiprocessing.process import BaseProcess +from typing import TYPE_CHECKING, Any + +import msgspec +import zmq +from vllm.logger import init_logger +from vllm.transformers_utils.config import ( + maybe_register_config_serialize_by_value, +) +from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx +from vllm.utils.system_utils import ( + decorate_logs, + get_mp_context, + set_process_title, +) +from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.engine.utils import ( + EngineHandshakeMetadata, + EngineZmqAddresses, + get_engine_zmq_addresses, +) +from vllm.v1.utils import shutdown + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.executor import Executor + +logger = init_logger(__name__) + +_HANDSHAKE_POLL_TIMEOUT_S = 600 + + +class StageEngineCoreProc(EngineCoreProc): + """Stage-specific engine core process for vLLM-Omni. + + Inherits from EngineCoreProc and provides its own ``run_stage_core`` + entry point for launching in a subprocess. Does **not** delegate to + ``EngineCoreProc.run_engine_core()``. + """ + + @staticmethod + def run_stage_core( + *args: Any, + dp_rank: int = 0, + local_dp_rank: int = 0, + **kwargs: Any, + ) -> None: + """Launch StageEngineCoreProc busy loop in background process.""" + shutdown_requested = False + maybe_register_config_serialize_by_value() + + def signal_handler(signum: int, frame: Any) -> None: + nonlocal shutdown_requested + if not shutdown_requested: + shutdown_requested = True + raise SystemExit() + + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + + engine_core: StageEngineCoreProc | None = None + try: + vllm_config: VllmConfig = kwargs["vllm_config"] + parallel_config = vllm_config.parallel_config + + set_process_title(f"StageEngineCoreProc_DP{dp_rank}") + decorate_logs() + + # the current vllm-omni does not support data parallelism, + # so we set the data parallel size to 1. + # [TODO] support data parallelism in the future. + # https://github.com/vllm-project/vllm-omni/issues/984 + parallel_config.data_parallel_size = 1 + parallel_config.data_parallel_size_local = 1 + parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = dp_rank + + engine_core = StageEngineCoreProc( + *args, + engine_index=dp_rank, + **kwargs, + ) + engine_core.run_busy_loop() + + except SystemExit: + logger.debug("StageEngineCoreProc exiting.") + raise + except Exception: + if engine_core is None: + logger.exception("StageEngineCoreProc failed to start.") + else: + logger.exception("StageEngineCoreProc encountered a fatal error.") + engine_core._send_engine_dead() + raise + finally: + if engine_core is not None: + engine_core.shutdown() + + +def spawn_stage_core( + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool = False, +) -> tuple[EngineZmqAddresses, BaseProcess, str]: + """Spawn a *StageEngineCoreProc* subprocess without performing the handshake. + + Must be called while the correct device env vars are set (e.g. under + the stage-launch lock). Call ``complete_stage_handshake`` afterwards. + + Returns ``(addresses, process, handshake_address)``. + """ + addresses = get_engine_zmq_addresses(vllm_config) + handshake_address = get_open_zmq_ipc_path() + + ctx = get_mp_context() + proc = ctx.Process( + target=StageEngineCoreProc.run_stage_core, + name="StageEngineCoreProc", + kwargs={ + "vllm_config": vllm_config, + "local_client": True, + "handshake_address": handshake_address, + "executor_class": executor_class, + "log_stats": log_stats, + "dp_rank": 0, + "local_dp_rank": 0, + }, + ) + proc.start() + return addresses, proc, handshake_address + + +def complete_stage_handshake( + proc: BaseProcess, + handshake_address: str, + addresses: EngineZmqAddresses, + vllm_config: VllmConfig, +) -> None: + """Perform the HELLO/INIT/READY handshake with an already-spawned proc. + + On failure the process is terminated before re-raising. + """ + try: + _perform_handshake(proc, handshake_address, addresses, vllm_config) + except Exception: + shutdown([proc]) + raise + + +def _perform_handshake( + proc: BaseProcess, + handshake_address: str, + addresses: EngineZmqAddresses, + vllm_config: VllmConfig, +) -> None: + """Run the HELLO / INIT / READY handshake with the subprocess.""" + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + poller = zmq.Poller() + poller.register(handshake_socket, zmq.POLLIN) + poller.register(proc.sentinel, zmq.POLLIN) + + identity, msg = _recv(poller, handshake_socket, proc, "HELLO") + if msg.get("status") != "HELLO": + raise RuntimeError(f"Expected HELLO, got: {msg}") + + init_payload = EngineHandshakeMetadata( + addresses=addresses, + parallel_config={}, + ) + handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)]) + + identity, msg = _recv(poller, handshake_socket, proc, "READY") + if msg.get("status") != "READY": + raise RuntimeError(f"Expected READY, got: {msg}") + num_gpu_blocks = msg.get("num_gpu_blocks") + if num_gpu_blocks is not None: + vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks + + +def _recv( + poller: zmq.Poller, + handshake_socket: zmq.Socket, + proc: BaseProcess, + expected: str, +) -> tuple[bytes, dict]: + """Wait for one handshake message; raise if the process dies first.""" + timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000 + while True: + events = dict(poller.poll(timeout=timeout_ms)) + if not events: + raise TimeoutError(f"Timed out waiting for {expected} from StageEngineCoreProc") + if handshake_socket in events: + identity, raw = handshake_socket.recv_multipart() + return identity, msgspec.msgpack.decode(raw) + if proc.exitcode is not None: + raise RuntimeError(f"StageEngineCoreProc died during {expected} (exit code {proc.exitcode})") diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9c246ce6e..6e8137206 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -75,6 +75,14 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s return model +def terminate_alive_proc(proc, timeout=5): + if proc.is_alive(): + proc.terminate() + proc.join(timeout=timeout) + if proc.is_alive(): + proc.kill() + + def resolve_worker_cls(engine_args: dict[str, Any]) -> None: """Resolve worker_cls from worker_type for non-diffusion stages.""" worker_type = engine_args.get("worker_type", None) @@ -121,8 +129,7 @@ class StartedLlmStage: metadata: Any vllm_config: Any executor_class: type - engine_manager: Any - coordinator: Any + proc: Any addresses: Any @@ -446,7 +453,7 @@ def initialize_diffusion_stage( metadata: Extracted stage metadata. batch_size: Maximum number of requests to batch together in the diffusion engine. Passed through to ``StageDiffusionClient`` - and ultimately to ``AsyncOmniDiffusion``. + and ultimately to ``AsyncOmni``. """ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient @@ -461,23 +468,17 @@ def initialize_diffusion_stage( def close_started_llm_stage(started: StartedLlmStage) -> None: - """Close managers owned by a launched stage that never attached.""" - resources = ( - ("engine manager", started.engine_manager), - ("coordinator", started.coordinator), - ) - for resource_name, resource in resources: - if resource is None: - continue - try: - resource.close() - except Exception as cleanup_error: - logger.warning( - "[stage_init] Failed to close launched %s for stage %s: %s", - resource_name, - started.stage_id, - cleanup_error, - ) + """Terminate the subprocess owned by a launched stage that never attached.""" + if started.proc is None: + return + try: + terminate_alive_proc(started.proc) + except Exception as cleanup_error: + logger.warning( + "[stage_init] Failed to terminate process for stage %s: %s", + started.stage_id, + cleanup_error, + ) def finalize_initialized_stages( diff --git a/vllm_omni/entrypoints/__init__.py b/vllm_omni/entrypoints/__init__.py index d0830df96..7b09adf93 100644 --- a/vllm_omni/entrypoints/__init__.py +++ b/vllm_omni/entrypoints/__init__.py @@ -6,11 +6,9 @@ vLLM-Omni entrypoints module. """ from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion from vllm_omni.entrypoints.omni import Omni __all__ = [ "AsyncOmni", - "AsyncOmniDiffusion", "Omni", ] diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py deleted file mode 100644 index 558ef96cb..000000000 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ /dev/null @@ -1,473 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Async entrypoint for vLLM-Omni diffusion model inference. - -Provides an asynchronous interface for running diffusion models, -enabling concurrent request handling and streaming generation. -""" - -import asyncio -import uuid -import weakref -from collections.abc import AsyncGenerator, Iterable -from concurrent.futures import ThreadPoolExecutor -from typing import Any - -from vllm.logger import init_logger -from vllm.transformers_utils.config import get_hf_file_to_dict - -from vllm_omni.diffusion.data import ( - DiffusionRequestAbortedError, - OmniDiffusionConfig, - TransformerConfig, -) -from vllm_omni.diffusion.diffusion_engine import DiffusionEngine -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType -from vllm_omni.lora.request import LoRARequest -from vllm_omni.outputs import OmniRequestOutput - -logger = init_logger(__name__) - - -def _weak_close_async_omni_diffusion(engine: DiffusionEngine, executor: ThreadPoolExecutor) -> None: - """Best-effort diffusion cleanup for GC finalization.""" - try: - engine.close() - except Exception: - pass - try: - executor.shutdown(wait=False) - except Exception: - pass - - -class AsyncOmniDiffusion: - """Async entry point for vLLM-Omni diffusion model inference. - - This class provides an asynchronous interface for running diffusion models, - enabling concurrent request handling. It wraps the DiffusionEngine and - provides async methods for image generation. - - Args: - model: Model name or path to load - od_config: Optional OmniDiffusionConfig. If not provided, it will be - created from kwargs - **kwargs: Additional keyword arguments passed to OmniDiffusionConfig - - Example: - >>> async_diffusion = AsyncOmniDiffusion(model="Qwen/Qwen-Image") - >>> result = await async_diffusion.generate( - ... prompt="A beautiful sunset over the ocean", - ... request_id="req-1", - ... ) - >>> print(result.images) - """ - - def __init__( - self, - model: str, - od_config: OmniDiffusionConfig | None = None, - batch_size: int = 1, - **kwargs: Any, - ): - self.model = model - - # Set batch size (default 1 for backward compatibility) - self._batch_size = max(1, batch_size) - - # Capture stage info from kwargs before they might be filtered out - stage_id = kwargs.get("stage_id") - engine_input_source = kwargs.get("engine_input_source") - cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None) - - # Build config - if od_config is None: - od_config = OmniDiffusionConfig.from_kwargs(model=model, **kwargs) - elif isinstance(od_config, dict): - # If config is dict, check it too (priority to kwargs if both exist) - if stage_id is None: - stage_id = od_config.get("stage_id") - if engine_input_source is None: - engine_input_source = od_config.get("engine_input_source") - od_config = OmniDiffusionConfig.from_kwargs(**od_config) - - self.od_config = od_config - - # Inject stage info into omni_kv_config if present - if stage_id is not None: - self.od_config.omni_kv_config.setdefault("stage_id", stage_id) - if engine_input_source is not None: - self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source) - - # Diffusers-style models expose `model_index.json` with `_class_name`. - # Non-diffusers models (e.g. Bagel, NextStep) only have `config.json`, - # so we fall back to reading that and mapping model_type manually. - try: - config_dict = get_hf_file_to_dict("model_index.json", od_config.model) - if config_dict is not None: - if od_config.model_class_name is None: - od_config.model_class_name = config_dict.get("_class_name", None) - od_config.update_multimodal_support() - - tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model) - od_config.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) - else: - raise FileNotFoundError("model_index.json not found") - except (AttributeError, OSError, ValueError, FileNotFoundError): - cfg = get_hf_file_to_dict("config.json", od_config.model) - if cfg is None: - raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}") - - od_config.set_tf_model_config(TransformerConfig.from_dict(cfg)) - model_type = cfg.get("model_type") - architectures = cfg.get("architectures") or [] - # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually - if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: - od_config.model_class_name = "BagelPipeline" - od_config.tf_model_config = TransformerConfig() - od_config.update_multimodal_support() - elif model_type == "nextstep": - if od_config.model_class_name is None: - od_config.model_class_name = "NextStep11Pipeline" - od_config.tf_model_config = TransformerConfig() - od_config.update_multimodal_support() - elif architectures and len(architectures) == 1: - od_config.model_class_name = architectures[0] - else: - raise - - if cfg_kv_collect_func is not None: - od_config.cfg_kv_collect_func = cfg_kv_collect_func - - # Initialize engine - self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) - - # Thread pool for running sync engine in async context - self._executor = ThreadPoolExecutor(max_workers=1) - self._closed = False - self._weak_finalizer = weakref.finalize( - self, - _weak_close_async_omni_diffusion, - self.engine, - self._executor, - ) - - logger.info("AsyncOmniDiffusion initialized with model: %s, batch_size: %d", model, self._batch_size) - - # ------------------------------------------------------------------ - # batch_size property - # ------------------------------------------------------------------ - - @property - def batch_size(self) -> int: - """Return the configured batch size for request batching.""" - return self._batch_size - - @batch_size.setter - def batch_size(self, value: int) -> None: - if not isinstance(value, int) or value < 1: - raise ValueError("batch_size must be a positive integer") - self._batch_size = value - - # ------------------------------------------------------------------ - # Public batch generation API - # ------------------------------------------------------------------ - - async def generate_batch( - self, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - request_id: str | None = None, - lora_request: LoRARequest | None = None, - ) -> OmniRequestOutput: - """Generate images from multiple prompts in a single engine call. - - Batches the given prompts into **one** ``DiffusionEngine.step()`` - call and returns a single ``OmniRequestOutput`` containing all - generated images. Called by ``StageDiffusionClient._run_batch`` - when the orchestrator receives a list-prompt request. - - Args: - prompts: List of text prompts describing the desired images. - sampling_params: Shared sampling parameters for all prompts. - request_id: Optional unique identifier. Auto-generated when *None*. - lora_request: Optional LoRA adapter to apply. - - Returns: - A single ``OmniRequestOutput`` with all images combined. - """ - if request_id is None: - request_id = f"diff-batch-{uuid.uuid4().hex[:8]}" - return await self._generate_batch(prompts, sampling_params, request_id, lora_request) - - # ------------------------------------------------------------------ - # Internal batch generation - # ------------------------------------------------------------------ - - async def _generate_batch( - self, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - request_id: str, - lora_request: LoRARequest | None = None, - ) -> OmniRequestOutput: - """Generate images from multiple prompts in a single engine call.""" - if not prompts: - return OmniRequestOutput(request_id=request_id, images=[], final_output_type="image") - - if sampling_params.guidance_scale: - sampling_params.guidance_scale_provided = True - - if lora_request is not None: - sampling_params.lora_request = lora_request - - request = OmniDiffusionRequest( - prompts=prompts, - sampling_params=sampling_params, - request_ids=[f"{request_id}-{i}" for i in range(len(prompts))], - ) - - logger.debug("Starting batch generation for %d prompts, request_id=%s", len(prompts), request_id) - - loop = asyncio.get_event_loop() - try: - results = await loop.run_in_executor( - self._executor, - self.engine.step, - request, - ) - except Exception as e: - logger.error("Batch generation failed for request %s: %s", request_id, e) - raise RuntimeError(f"Diffusion batch generation failed: {e}") from e - - # Combine all per-prompt results into a single OmniRequestOutput - all_images = [] - for result in results: - all_images.extend(result.images) - - return OmniRequestOutput( - request_id=request_id, - images=all_images, - final_output_type="image", - finished=True, - ) - - def get_diffusion_od_config(self) -> OmniDiffusionConfig: - """Return the diffusion config used by this engine.""" - return self.od_config - - # ------------------------------------------------------------------ - # Public generate API - # ------------------------------------------------------------------ - - async def generate( - self, - prompt: OmniPromptType, - sampling_params: OmniDiffusionSamplingParams, - request_id: str | None = None, - lora_request: LoRARequest | None = None, - ) -> OmniRequestOutput: - """Generate images asynchronously from a single text prompt. - - For batched generation (multiple prompts in one engine call), use - :meth:`generate_batch` instead. This method always processes - exactly one prompt per call. - - Args: - prompt: Text prompt describing the desired image - sampling_params: Sampling parameters - request_id: Optional unique identifier for tracking the request - lora_request: Optional LoRA adapter to apply - - Returns: - OmniRequestOutput containing generated images - - Raises: - RuntimeError: If generation fails - """ - if request_id is None: - request_id = f"diff-{uuid.uuid4().hex[:16]}" - if sampling_params.guidance_scale: - sampling_params.guidance_scale_provided = True - - if lora_request is not None: - sampling_params.lora_request = lora_request - - request = OmniDiffusionRequest( - prompts=[prompt], - sampling_params=sampling_params, - request_ids=[request_id], - ) - - logger.debug("Starting generation for request %s", request_id) - - loop = asyncio.get_event_loop() - try: - result = await loop.run_in_executor( - self._executor, - self.engine.step, - request, - ) - result = result[0] - except asyncio.CancelledError: - self.engine.abort(request_id) - raise - except DiffusionRequestAbortedError: - raise - except Exception as e: - logger.error("Generation failed for request %s: %s", request_id, e) - raise RuntimeError(f"Diffusion generation failed: {e}") from e - - if not result.request_id: - result.request_id = request_id - return result - - async def generate_stream( - self, - prompt: str, - request_id: str | None = None, - **kwargs: Any, - ) -> AsyncGenerator[OmniRequestOutput, None]: - """Generate images with streaming progress updates. - - Currently, diffusion models don't support true streaming, so this - yields a single result after generation completes. Future implementations - may support step-by-step progress updates. - - Args: - prompt: Text prompt describing the desired image - request_id: Optional unique identifier for tracking the request - **kwargs: Additional generation parameters - - Yields: - OmniRequestOutput with generation progress/results - """ - result = await self.generate(prompt=prompt, request_id=request_id, **kwargs) - yield result - - def close(self) -> None: - """Close the engine and release resources. - - Should be called when done using the AsyncOmniDiffusion instance. - """ - if self._closed: - return - self._closed = True - - finalizer = getattr(self, "_weak_finalizer", None) - if finalizer is not None and finalizer.alive: - finalizer.detach() - - try: - self.engine.close() - except Exception as e: - logger.warning("Error closing diffusion engine: %s", e) - - try: - self._executor.shutdown(wait=False) - except Exception as e: - logger.warning("Error shutting down executor: %s", e) - - logger.info("AsyncOmniDiffusion closed") - - def shutdown(self) -> None: - """Alias for close() method.""" - self.close() - - async def abort(self, request_id: str | Iterable[str]) -> None: - """Abort a request.""" - self.engine.abort(request_id) - - @property - def is_running(self) -> bool: - """Check if the engine is running.""" - return not self._closed - - @property - def is_stopped(self) -> bool: - """Check if the engine is stopped.""" - return self._closed - - async def remove_lora(self, adapter_id: int) -> bool: - """Remove a LoRA""" - loop = asyncio.get_event_loop() - results = await loop.run_in_executor( - self._executor, - self.engine.collective_rpc, - "remove_lora", - None, - (adapter_id,), - {}, - None, - ) - return all(results) if isinstance(results, list) else results - - async def add_lora(self, lora_request: LoRARequest) -> bool: - """Add a LoRA adapter""" - loop = asyncio.get_event_loop() - results = await loop.run_in_executor( - self._executor, - self.engine.collective_rpc, - "add_lora", - None, - (), - {"lora_request": lora_request}, - None, - ) - return all(results) if isinstance(results, list) else results - - async def list_loras(self) -> list[int]: - """List all registered LoRA adapter IDs.""" - loop = asyncio.get_event_loop() - results = await loop.run_in_executor( - self._executor, - self.engine.collective_rpc, - "list_loras", - None, - (), - {}, - None, - ) - # collective_rpc returns list from workers; flatten unique ids - if not isinstance(results, list): - return results or [] - merged: set[int] = set() - for part in results: - merged.update(part or []) - return sorted(merged) - - async def pin_lora(self, lora_id: int) -> bool: - """Prevent an adapter from being evicted.""" - loop = asyncio.get_event_loop() - results = await loop.run_in_executor( - self._executor, - self.engine.collective_rpc, - "pin_lora", - None, - (), - {"adapter_id": lora_id}, - None, - ) - return all(results) if isinstance(results, list) else results - - async def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: - """Start or stop profiling for the diffusion model. - - Args: - is_start: True to start profiling, False to stop. - profile_prefix: Optional prefix for trace filename (vLLM compat). - - Note: - Matches vLLM's worker.profile() signature for consistency. - Traces are saved automatically via on_trace_ready callback. - """ - loop = asyncio.get_event_loop() - await loop.run_in_executor( - self._executor, - self.engine.profile, - is_start, - profile_prefix, - ) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 527947be9..35f56516c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -6,7 +6,7 @@ import uuid from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime, timedelta, timezone from io import BytesIO -from typing import TYPE_CHECKING, Any, Final, Optional, cast +from typing import Any, Final, cast import jinja2 import torch @@ -89,9 +89,6 @@ from vllm_omni.entrypoints.openai.utils import parse_lora_request from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput -if TYPE_CHECKING: - from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion - logger = init_logger(__name__) @@ -107,13 +104,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Diffusion mode attributes _diffusion_mode: bool = False - _diffusion_engine: Optional["AsyncOmniDiffusion"] = None + _diffusion_engine: AsyncOmni | None = None _diffusion_model_name: str = "" @classmethod def for_diffusion( cls, - diffusion_engine: "AsyncOmniDiffusion", + diffusion_engine: AsyncOmni, model_name: str, ) -> "OmniOpenAIServingChat": """Create a chat serving instance for diffusion models. @@ -2153,7 +2150,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if resolution is not None: gen_params.resolution = resolution - # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni). + # Parse per-request LoRA. if lora_body and isinstance(lora_body, dict): try: lora_req, lora_scale = parse_lora_request(lora_body) @@ -2187,26 +2184,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) # Generate image - # Handle both AsyncOmniDiffusion (returns OmniRequestOutput) and AsyncOmni (returns AsyncGenerator) - if isinstance(self._diffusion_engine, AsyncOmni): - diffusion_engine = cast(AsyncOmni, self._diffusion_engine) - result = None - async for output in diffusion_engine.generate( - prompt=gen_prompt, - sampling_params_list=[gen_params], # Pass as single-stage params - request_id=request_id, - ): - result = output - if result is None: - return self._create_error_response("No output generated from AsyncOmni") - else: - # AsyncOmniDiffusion: direct call - diffusion_engine = cast(AsyncOmniDiffusion, self._diffusion_engine) - result = await diffusion_engine.generate( - prompt=gen_prompt, - sampling_params=gen_params, - request_id=request_id, - ) + diffusion_engine = cast(AsyncOmni, self._diffusion_engine) + result = None + async for output in diffusion_engine.generate( + prompt=gen_prompt, + sampling_params_list=[gen_params], # Pass as single-stage params + request_id=request_id, + ): + result = output + if result is None: + return self._create_error_response("No output generated from AsyncOmni") # Extract images from result # Handle nested OmniRequestOutput structure where images might be in request_output images = getattr(result.request_output, "images", [])