[Core] Add UniProc diffusion executor for single-GPU (#6308)

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com> Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
2026-08-22 21:58:21 +08:00
1 parent 72ee535fcd
commit b39f599b09
@@ -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)