From b39f599b09a9ec06cfef05213273de7ed6626b7f Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Sat, 22 Aug 2026 15:58:21 +0200 Subject: [PATCH] [Core] Add UniProc diffusion executor for single-GPU (#6308) Signed-off-by: Rahul Steiger Signed-off-by: Rahul Steiger Co-authored-by: Hongsheng Liu Co-authored-by: Rahul Steiger --- docs/configuration/stage_configs.md | 6 +- tests/diffusion/test_uniproc_executor.py | 262 ++++++++++++++++++ .../test_async_omni_collective_rpc.py | 106 +++---- .../test_async_omni_qwen_image_generate.py | 17 ++ .../test_async_omni_diffusion_config.py | 2 +- vllm_omni/config/omni_config.py | 2 +- vllm_omni/diffusion/data.py | 6 +- vllm_omni/diffusion/executor/abstract.py | 15 +- .../diffusion/executor/uniproc_executor.py | 259 +++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 4 +- 10 files changed, 619 insertions(+), 60 deletions(-) create mode 100644 tests/diffusion/test_uniproc_executor.py create mode 100644 vllm_omni/diffusion/executor/uniproc_executor.py diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index cb056f640..1afa16d54 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -74,8 +74,10 @@ For fields whose deploy default is `null`, the deploy layer contributes no override. The effective value may still come from a platform section, an explicit CLI or stage override, or the downstream vLLM engine default. -Note: for diffusion path, `distributed_executor_backend` currently defaults to -`mp`, and `ray` / `external_launcher` are not fully supported yet. +Note: for the diffusion path, an omitted `distributed_executor_backend` selects +`uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output +segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker +subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet. ### Stage fields diff --git a/tests/diffusion/test_uniproc_executor.py b/tests/diffusion/test_uniproc_executor.py new file mode 100644 index 000000000..dfc1d2eb1 --- /dev/null +++ b/tests/diffusion/test_uniproc_executor.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Single-GPU (uniproc) diffusion executor: selection and RPC behaviour.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor +from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeODConfig: + def __init__(self, num_gpus: int = 1, backend: str | None = None) -> None: + self.num_gpus = num_gpus + self.distributed_executor_backend = backend + self.worker_extension_cls = None + self.custom_pipeline_args = None + + +def _od_config(num_gpus: int = 1, backend: str | None = None) -> _FakeODConfig: + return _FakeODConfig(num_gpus=num_gpus, backend=backend) + + +def test_single_gpu_defaults_to_uniproc(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1)) is UniProcDiffusionExecutor + + +def test_explicit_mp_is_honored_on_single_gpu(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="mp")) is MultiprocDiffusionExecutor + + +def test_multi_gpu_defaults_to_multiproc(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=2)) is MultiprocDiffusionExecutor + + +def test_explicit_uni_backend_is_honored(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="uni")) is UniProcDiffusionExecutor + + +@pytest.fixture +def executor(monkeypatch: pytest.MonkeyPatch): + """A ``UniProcDiffusionExecutor`` with a mocked worker (no model load).""" + worker = MagicMock() + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker"), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname", + MagicMock(return_value=object), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase", + MagicMock(return_value=worker), + ) + return UniProcDiffusionExecutor(_od_config(num_gpus=1)), worker + + +def test_rejects_multi_gpu_config(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="unused"), + ) + with pytest.raises(ValueError, match="single GPU only"): + UniProcDiffusionExecutor(_od_config(num_gpus=4)) + + +def test_shutdown_is_safe_on_partially_constructed_executor(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="unused"), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname", + MagicMock(return_value=object), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase", + MagicMock(side_effect=RuntimeError("model load blew up")), + ) + constructed = UniProcDiffusionExecutor.__new__(UniProcDiffusionExecutor) + with pytest.raises(RuntimeError, match="model load blew up"): + constructed.__init__(_od_config(num_gpus=1)) + + constructed.shutdown() + + +def test_collective_rpc_calls_worker_directly(executor): + exec_, worker = executor + worker.execute_method.return_value = "result" + + out = exec_.collective_rpc("some_method", args=(1, 2), kwargs={"k": "v"}, unique_reply_rank=0) + + assert out == "result" + worker.execute_method.assert_called_once_with("some_method", 1, 2, k="v") + + +def test_collective_rpc_returns_list_when_no_reply_rank(executor): + exec_, worker = executor + worker.execute_method.return_value = "result" + + assert exec_.collective_rpc("some_method") == ["result"] + + +def test_collective_rpc_propagates_worker_exceptions(executor): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + +def test_recoverable_worker_failure_keeps_the_executor_alive(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("CUDA out of memory") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: True) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is False + died.assert_not_called() + exec_.check_health() + + +def test_poisoned_cuda_context_marks_the_executor_dead(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("an illegal memory access was encountered") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: False) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is True + died.assert_called_once_with() + with pytest.raises(EngineDeadError): + exec_.check_health() + + +def test_failure_is_latched_and_callbacks_fire_once(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: False) + died = MagicMock() + exec_.register_failure_callback(died) + + for _ in range(3): + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + died.assert_called_once_with() + + +def test_device_probe_is_skipped_when_cuda_was_never_initialized(executor): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + if not torch.cuda.is_initialized(): + assert exec_.is_dead is False + + +def test_device_probe_skipped_when_accelerator_unavailable(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: False) + sync = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", sync) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + sync.assert_not_called() + assert exec_.is_dead is False + + +def test_device_probe_skipped_when_device_not_initialized(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu")) + monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: False), raising=False) + sync = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", sync) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + sync.assert_not_called() + assert exec_.is_dead is False + + +def test_device_probe_latches_sticky_fault_on_npu(executor, monkeypatch): + """NPU must not short-circuit the probe the way a CUDA-only guard did.""" + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu")) + monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: True), raising=False) + monkeypatch.setattr( + torch.accelerator, + "synchronize", + MagicMock(side_effect=RuntimeError("NPU context poisoned")), + ) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is True + died.assert_called_once_with() + + +def test_check_health_ok_then_dead(executor): + exec_, _ = executor + exec_.check_health() + + exec_._is_failed = True + with pytest.raises(EngineDeadError): + exec_.check_health() + + +def test_shutdown_is_idempotent_and_closes_executor(executor, monkeypatch): + exec_, worker = executor + collect = MagicMock() + empty_cache = MagicMock() + monkeypatch.setattr("vllm_omni.diffusion.executor.uniproc_executor.gc.collect", collect) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.is_available", + lambda: True, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.empty_cache", + empty_cache, + ) + exec_.register_failure_callback(MagicMock()) + + exec_.shutdown() + exec_.shutdown() + + worker.shutdown.assert_called_once() + assert exec_.driver_worker is None + assert exec_._failure_callbacks == [] + collect.assert_called_once_with() + empty_cache.assert_called_once_with() + with pytest.raises(RuntimeError, match="closed"): + exec_.collective_rpc("some_method", unique_reply_rank=0) diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py index 04e83caeb..a11c9b7e4 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py @@ -166,8 +166,8 @@ async def test_sleep_memory_reclaimed_custom_pipeline(): The fix moves custom_pipeline init outside the CUDA context so all weights go through the caching allocator and are therefore fully reclaimed by - sleep(level=1). A non-zero ``CuMemAllocator.get_current_usage()`` after - sleep is the direct signal that the bypass is still occurring. + sleep(level=1). The tracked allocations must transition to the allocator's + asleep state and report physical reclamation. """ with ExitStack() as after: engine = AsyncOmni( @@ -187,62 +187,72 @@ async def test_sleep_memory_reclaimed_custom_pipeline(): free_before, total = torch.cuda.mem_get_info() used_before_gib = (total - free_before) / 1024**3 - # Measure CuMemAllocator-tracked usage before sleep. In inline mode - # the worker runs in a thread pool inside this process, so the allocator - # singleton is shared and can be read directly. + # Measure CuMemAllocator-tracked usage before sleep. In inline / uni + # mode the worker runs in this process, so the allocator singleton is + # shared and can be read directly. allocator = None tracked_before = 0 + tracked_ptrs: set[int] = set() try: from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() - tracked_before = allocator.get_current_usage() + tracked_ptrs = {ptr for ptr, data in allocator.pointer_to_data.items() if not data.is_asleep} + tracked_before = sum(allocator.pointer_to_data[ptr].handle[1] for ptr in tracked_ptrs) except Exception: pass - # Put the engine to sleep; all weights should be offloaded via the pool. - acks = await engine.sleep(level=1) - await asyncio.sleep(0.5) # allow the CUDA driver to settle - torch.accelerator.synchronize() + slept = False + try: + # Put the engine to sleep; this engine's weights should be + # offloaded via the pool. + acks = await engine.sleep(level=1) + slept = True + await asyncio.sleep(0.5) # allow the CUDA driver to settle + torch.accelerator.synchronize() - # Measure after sleep. - free_after, _ = torch.cuda.mem_get_info() - used_after_gib = (total - free_after) / 1024**3 - drop_gib = used_before_gib - used_after_gib + # Measure after sleep. + free_after, _ = torch.cuda.mem_get_info() + used_after_gib = (total - free_after) / 1024**3 + drop_gib = used_before_gib - used_after_gib - # --- Primary assertion: allocator reports zero tracked memory. --- - # If this fails it means weights were allocated outside the CuMem pool - # (safetensors direct-to-GPU bypass) — the exact regression this test - # is designed to catch. - if allocator is not None: - tracked_after = allocator.get_current_usage() - assert tracked_after == 0, ( - f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB " - f"after sleep(level=1) on custom_pipeline path " - f"(was {tracked_before / 1024**3:.3f} GiB before sleep). " - "Weights were allocated outside the CuMem pool via the " - "safetensors direct-to-GPU fast path — loader-context fix " - "may not be applied." + # get_current_usage() includes unmapped sleeping handles, so it + # does not decrease on sleep. Verify that allocations mapped + # before the call instead transitioned to the asleep state. + if allocator is not None: + assert tracked_before > 0, ( + "Expected custom_pipeline weights to use the CuMem pool, " + "but no mapped allocations were tracked before sleep." + ) + still_awake = [ + ptr + for ptr in tracked_ptrs + if ptr in allocator.pointer_to_data and not allocator.pointer_to_data[ptr].is_asleep + ] + assert not still_awake, ( + f"{len(still_awake)} CuMem allocation(s) remained mapped " + "after sleep(level=1) on the custom_pipeline path." + ) + + # Physical VRAM or ACK freed_bytes confirms reclamation at the + # driver level. + total_freed_bytes = sum( + (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0)) + for ack in acks + if ack is not None + ) + freed_gib = total_freed_bytes / 1024**3 + assert freed_gib > 0 or drop_gib > 0, ( + f"Expected GPU memory to be reclaimed after sleep(level=1) on " + f"custom_pipeline + enable_sleep_mode=True. " + f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, " + f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB." ) - # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms - # reclamation at the driver level. - total_freed_bytes = sum( - (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0)) - for ack in acks - if ack is not None - ) - freed_gib = total_freed_bytes / 1024**3 - assert freed_gib > 0 or drop_gib > 0, ( - f"Expected GPU memory to be reclaimed after sleep(level=1) on " - f"custom_pipeline + enable_sleep_mode=True. " - f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, " - f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB." - ) - - # Engine must report it is sleeping. - assert await engine.is_sleeping() - - # Wake up and confirm the engine is functional again. - await engine.wake_up() - assert not await engine.is_sleeping() + assert await engine.is_sleeping() + await engine.wake_up() + slept = False + assert not await engine.is_sleeping() + finally: + if slept: + await engine.wake_up() diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index b24f364e4..12ad4f1b3 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -17,9 +17,12 @@ from functools import lru_cache import numpy as np import pytest +import torch.distributed as dist from transformers import AutoTokenizer from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -131,6 +134,19 @@ async def _generate_once( return last_output +def _assert_live_uniproc_worker(engine: AsyncOmni) -> None: + """Default single-GPU path must use the in-process executor, not mp IPC.""" + clients = engine.engine.stage_clients + assert len(clients) == 1 + client = clients[0] + assert isinstance(client, InlineStageDiffusionClient) + executor = client._engine.executor + assert isinstance(executor, UniProcDiffusionExecutor) + assert executor.driver_worker is not None + assert dist.is_initialized() + assert dist.get_world_size() == 1 + + def _assert_valid_image_output(output: OmniRequestOutput) -> None: assert output.final_output_type == "image" assert output.images, "Expected at least one generated image" @@ -158,6 +174,7 @@ async def test_async_omni_generate(): ) after.callback(engine.shutdown) + _assert_live_uniproc_worker(engine) output = await _generate_once( engine, "a beautiful sunset over the ocean with vibrant orange and purple clouds " diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1d5d58011..28667227d 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -225,7 +225,7 @@ def test_default_stage_config_whitelist_none_fallback(): engine_args = stage_cfg["engine_args"] assert engine_args["trust_remote_code"] is False - assert engine_args["distributed_executor_backend"] == "mp" + assert engine_args["distributed_executor_backend"] is None assert engine_args["dtype"] == "auto" assert engine_args["enforce_eager"] is False diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index b60f574a1..5a1dbb12f 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -663,7 +663,7 @@ class _DiffusionConfigProjection: dtype: Any = "auto" trust_remote_code: bool = False revision: str | None = None - distributed_executor_backend: str = "mp" + distributed_executor_backend: str | None = None dist_timeout: int | None = None nccl_port: int | None = None master_port: int | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5dde3be73..9adda521a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -700,8 +700,10 @@ class OmniDiffusionConfig: # in both directions. See docs/features/session_state_manager.md. enable_session_state_manager: bool = False - # Distributed executor backend - distributed_executor_backend: str = "mp" + # Distributed executor backend. ``None`` means auto: ``"uni"`` at + # ``num_gpus == 1``, ``"mp"`` otherwise. Explicit ``"mp"`` keeps a + # worker subprocess (process isolation and RPC timeouts). + distributed_executor_backend: str | None = None nccl_port: int | None = None # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class`` diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index a491c3391..f15db04bd 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -24,10 +24,15 @@ class DiffusionExecutor(ABC): def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]: executor_class: type[DiffusionExecutor] distributed_executor_backend = od_config.distributed_executor_backend - # Keep backward-compatible behavior for callers/configs that omit this - # field and rely on the historical diffusion default backend. + # Mirror vLLM's `world_size == 1 -> "uni"` default + # (vllm/config/parallel.py). A single-GPU diffusion deployment has + # nothing to distribute: spawning a worker only adds MessageQueues, + # /dev/shm output segments, and a second model load. Explicit "mp" + # keeps process isolation and RPC timeouts. Multi-GPU still defaults + # to "mp". if distributed_executor_backend is None: - distributed_executor_backend = "mp" + num_gpus = od_config.num_gpus or 1 + distributed_executor_backend = "uni" if num_gpus == 1 else "mp" if isinstance(distributed_executor_backend, type): if not issubclass(distributed_executor_backend, DiffusionExecutor): @@ -42,6 +47,10 @@ class DiffusionExecutor(ABC): from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor executor_class = MultiprocDiffusionExecutor + elif distributed_executor_backend == "uni": + from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor + + executor_class = UniProcDiffusionExecutor elif distributed_executor_backend == "external_launcher": raise NotImplementedError("external_launcher backend is not yet supported.") elif isinstance(distributed_executor_backend, str): diff --git a/vllm_omni/diffusion/executor/uniproc_executor.py b/vllm_omni/diffusion/executor/uniproc_executor.py new file mode 100644 index 000000000..8eb098323 --- /dev/null +++ b/vllm_omni/diffusion/executor/uniproc_executor.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""In-process diffusion executor for single-GPU deployments. + +Counterpart to :class:`MultiprocDiffusionExecutor` for ``num_gpus == 1``. +The worker is constructed and driven in the engine process, so none of the +multiproc machinery is created: + +* no spawned worker subprocess (and no second model load / CUDA context) +* no ``MessageQueue`` shared-memory ring buffers or zmq ``ipc://`` sockets +* no POSIX ``/dev/shm`` segments for output tensors (``ipc.py`` pack/unpack) + +Mirrors vLLM's ``UniProcExecutor`` (``vllm/v1/executor/uniproc_executor.py``). + +``timeout`` on :meth:`collective_rpc` is accepted for interface parity but is +not enforced: the call runs on the calling thread, so a hung worker blocks it. +The multiproc executor enforces deadlines via its result-queue dequeue. +""" + +from __future__ import annotations + +import gc +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput + from vllm_omni.diffusion.worker.utils import BaseRunnerOutput + +logger = init_logger(__name__) + + +class UniProcDiffusionExecutor(DiffusionExecutor): + """Runs a single diffusion worker inline, with no IPC of any kind.""" + + def _init_executor(self) -> None: + from vllm_omni.diffusion.worker.diffusion_worker import WorkerWrapperBase + + self._closed = False + self._is_failed = False + self.driver_worker = None + self._failure_callbacks: list[Callable[[], None]] = [] + self._warned_about_timeout = False + + num_gpus = self.od_config.num_gpus or 1 + if num_gpus != 1: + raise ValueError( + f"UniProcDiffusionExecutor supports a single GPU only, got num_gpus={num_gpus}. " + 'Use distributed_executor_backend="mp" for multi-GPU deployments.' + ) + + # The worker still initializes a one-rank distributed environment so + # that group lookups (`get_dp_group()` and friends) resolve. At world + # size 1, vLLM's GroupCoordinator skips the device communicator and + # the shared-memory MessageQueue, and every collective short-circuits. + worker_cls_path = current_omni_platform.get_diffusion_worker_cls() + self.driver_worker = WorkerWrapperBase( + gpu_id=0, + od_config=self.od_config, + worker_extension_cls=self.od_config.worker_extension_cls, + custom_pipeline_args=getattr(self.od_config, "custom_pipeline_args", None), + base_worker_class=resolve_obj_by_qualname(worker_cls_path), + ) + logger.info("Diffusion worker initialized in-process (uniproc executor)") + + @property + def is_dead(self) -> bool: + return self._closed or self._is_failed + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("DiffusionExecutor is closed.") + if self.driver_worker is None: + raise RuntimeError("Diffusion worker is not initialized.") + + def collective_rpc( + self, + method: str, + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None, + unique_reply_rank: int | None = None, + exec_all_ranks: bool = False, + ) -> Any: + """Invoke ``method`` on the single in-process worker. + + ``unique_reply_rank`` / ``exec_all_ranks`` only select which ranks + participate, which is degenerate at world size 1. The return shape + still matches :class:`MultiprocDiffusionExecutor` (a bare result when a + reply rank is named, a one-element list otherwise). + """ + self._ensure_open() + if timeout is not None and not self._warned_about_timeout: + self._warned_about_timeout = True + logger.info( + "Diffusion RPC deadlines are not enforced by the in-process executor: " + "the worker runs on the calling thread, so a hung call blocks it." + ) + try: + result = self.driver_worker.execute_method(method, *(args or ()), **(kwargs or {})) + except Exception: + if not self._device_is_usable(): + self._mark_failed() + raise + return result if unique_reply_rank is not None else [result] + + def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.sched.interface import validate_new_request_data_identity + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + self._ensure_open() + runner_outputs: list[RunnerOutput] = [] + for new_req in scheduler_output.scheduled_new_reqs: + validate_new_request_data_identity(new_req) + try: + args: tuple = (new_req.req, self.od_config, scheduler_output.kv_prefetch_job) + if new_req.diffusion_kv_metadata is not None: + args += (new_req.diffusion_kv_metadata,) + result = self.collective_rpc( + "execute_model", + args=args, + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, DiffusionOutput): + raise RuntimeError(f"Unexpected response type: {type(result)!r}") + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=result, + ) + ) + except Exception as exc: + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=DiffusionOutput(error=str(exc)), + ) + ) + return BatchRunnerOutput.from_list(runner_outputs) + + def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput + + self._ensure_open() + if len(scheduler_output.scheduled_new_reqs) <= 1: + return self.execute_request(scheduler_output) + + result = self.collective_rpc( + "execute_model_batch", + args=(scheduler_output, self.od_config), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, BatchRunnerOutput): + raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}") + return result + + def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.worker.utils import BaseRunnerOutput + + self._ensure_open() + result = self.collective_rpc( + "execute_stepwise", + args=(scheduler_output,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if isinstance(result, BaseRunnerOutput): + return result + raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}") + + def _device_is_usable(self) -> bool: + """Whether the accelerator context survived the failure we just caught. + + Running the worker inline means a worker fault no longer kills a + process, so nothing else notices it. Recoverable faults (OOM, an + invalid request) leave the context intact and must not take the + deployment down. A sticky fault (illegal memory access, ECC error) + poisons the context: every later launch fails, so the engine needs + to be restarted. ``synchronize()`` re-raises a pending sticky error + on the registered accelerator (CUDA, ROCm, NPU). + """ + if not torch.accelerator.is_available(): + return True + accelerator = torch.accelerator.current_accelerator() + device_mod = getattr(torch, accelerator.type, None) if accelerator is not None else None + is_initialized = getattr(device_mod, "is_initialized", None) + if callable(is_initialized) and not is_initialized(): + return True + try: + torch.accelerator.synchronize() + except Exception as exc: + logger.error("Device context is unusable after a worker failure: %s", exc) + return False + return True + + def _mark_failed(self) -> None: + if self._is_failed: + return + self._is_failed = True + logger.error("Diffusion worker failed unrecoverably in-process; marking the executor dead.") + for callback in self._failure_callbacks: + try: + callback() + except Exception: + logger.exception("Diffusion executor failure callback raised") + + def register_failure_callback(self, callback: Callable[[], None]) -> None: + """Register a callback invoked when the inline worker fails fatally. + + The multiproc executor fires these from its process monitor. There is + no process to monitor here, so they are fired from ``collective_rpc`` + instead — without them a poisoned device context would leave + ``check_health`` reporting healthy while every request fails. + """ + self._failure_callbacks.append(callback) + + def check_health(self) -> None: + if self._is_failed: + raise EngineDeadError() + self._ensure_open() + + def shutdown(self) -> None: + if self._closed: + return + self._closed = True + worker = self.driver_worker + self.driver_worker = None + self._failure_callbacks.clear() + if worker is not None: + try: + worker.shutdown() + except Exception as exc: + logger.warning("Diffusion worker shutdown encountered an error: %s", exc) + # Unlike the multiprocess executor, tearing down an inline worker does + # not exit its process. Drop the final model reference and return cached + # allocations to the driver so a subsequent engine can use the device. + del worker + gc.collect() + try: + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + except Exception as exc: + logger.warning("Failed to release accelerator cache during diffusion worker shutdown: %s", exc) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index d30db03cb..12236294e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1065,9 +1065,7 @@ class AsyncOmniEngine: "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), - "distributed_executor_backend": ( - "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"] - ), + "distributed_executor_backend": kwargs.get("distributed_executor_backend"), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32),