262 lines
11 KiB
Python
262 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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:
|
|
# The worker runs inline, so a pipeline's OmniClientError lands
|
|
# here; from_exception keeps its 4xx status for the engine/API.
|
|
runner_outputs.append(
|
|
RunnerOutput(
|
|
request_id=new_req.request_id,
|
|
step_index=None,
|
|
finished=True,
|
|
result=DiffusionOutput.from_exception(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)
|