vllm-omni/vllm_omni/diffusion/diffusion_engine.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import asyncio
import concurrent.futures
import inspect
import queue
import threading
import time
from collections.abc import AsyncGenerator, Iterable
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any

import numpy as np
import PIL.Image
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,
    DiffusionRequestAbortedError,
    OmniDiffusionConfig,
)
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
from vllm_omni.diffusion.io_support import (
    get_dummy_run_num_frames,
    image_color_format,
    supports_audio_output,
    supports_multimodal_input,
)
from vllm_omni.diffusion.output_formatter import (
    format_diffusion_outputs,
    format_empty_diffusion_outputs,
    normalize_diffusion_postprocess_output,
)
from vllm_omni.diffusion.registry import (
    DiffusionModelRegistry,
    get_diffusion_post_process_func,
    get_diffusion_pre_process_func,
)
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
from vllm_omni.errors import client_error_from_metadata, is_client_error_status
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt

if TYPE_CHECKING:
    from vllm_omni.outputs import OmniRequestOutput

logger = init_logger(__name__)

_ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds

__all__ = [
    "DiffusionEngine",
    "DiffusionExecutionMode",
    "_RpcTask",
    "_move_tensor_tree_to_cpu",
    "get_dummy_run_num_frames",
    "image_color_format",
    "supports_audio_output",
    "supports_multimodal_input",
]


def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool:
    if func is None:
        return False
    parameters = inspect.signature(func).parameters
    return parameter_name in parameters or any(
        parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
    )


def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) -> type | None:
    if custom_pipeline_args is None:
        return None

    try:
        pipeline_cls = custom_pipeline_args["pipeline_class"]
    except KeyError as exc:
        raise ValueError("custom_pipeline_args must include 'pipeline_class'.") from exc

    if isinstance(pipeline_cls, type):
        return pipeline_cls
    if isinstance(pipeline_cls, str):
        try:
            return resolve_obj_by_qualname(pipeline_cls)
        except (AttributeError, ImportError, ValueError) as exc:
            raise ValueError(f"Failed to resolve custom diffusion pipeline class {pipeline_cls!r}.") from exc
    raise TypeError(
        f"custom_pipeline_args['pipeline_class'] must be a qualified name string or a class, "
        f"got {type(pipeline_cls).__name__}"
    )


def supports_request_batch(od_config: OmniDiffusionConfig) -> bool:
    model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None))
    if model_cls is None:
        model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None))
    if model_cls is None:
        return False
    return bool(getattr(model_cls, "supports_request_batch", False))


def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
    try:
        return max(1, int(getattr(od_config, "max_num_seqs", 1)))
    except (TypeError, ValueError):
        return 1


def _move_tensor_tree_to_cpu(value: object) -> object:
    if isinstance(value, torch.Tensor):
        return value.cpu() if value.device.type != "cpu" else value
    if isinstance(value, dict):
        return {key: _move_tensor_tree_to_cpu(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_move_tensor_tree_to_cpu(item) for item in value]
    if isinstance(value, tuple):
        return tuple(_move_tensor_tree_to_cpu(item) for item in value)
    return value


@dataclass
class _RpcTask:
    """A pending collective_rpc invocation queued for the busy loop."""

    method: str
    args: tuple
    kwargs: dict | None
    deadline: float | None
    unique_reply_rank: int | None
    future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)


class DiffusionExecutionMode(str, Enum):
    REQUEST_BATCH = "request_batch"
    STEP_BATCH = "step_batch"


class DiffusionEngine:
    """The diffusion engine for vLLM-Omni diffusion models."""

    #: Import path of the model runner this engine's workers should build, or
    #: ``None`` for the platform default. Subclasses declare their runner here
    #: (e.g. the AR-Diffusion engine), the worker resolves it through
    #: :meth:`resolve_engine_class` — so engine->runner routing lives on the
    #: engine class itself, and ``od_config.diffusion_model_runner_cls``
    #: remains a pure explicit user override (never mutated by engines).
    default_diffusion_model_runner_cls: str | None = None

    # Class-level default so tests using object.__new__ (without __init__)
    # don't hit AttributeError when _busy_loop accesses self.dp_concurrent.
    dp_concurrent: bool = False

    def __init__(
        self,
        od_config: OmniDiffusionConfig,
        scheduler: BaseScheduler | None = None,
    ):
        """Initialize the diffusion engine.

        Args:
            od_config: The configuration for the diffusion engine.
            scheduler: Optional scheduler override for tests or custom engine
                integrations. When omitted, the engine selects a scheduler
                from the resolved execution mode.
        """
        self.od_config = od_config

        self._init_process_hooks(od_config)
        self.execution_mode = self._resolve_execution_mode(od_config)
        self._init_executor(od_config)
        self._init_scheduler(od_config, scheduler)
        self._init_runtime_state()
        self._init_execute_fn()
        self._log_execution_mode(od_config)

    def _init_process_hooks(self, od_config: OmniDiffusionConfig) -> None:
        self.post_process_func = get_diffusion_post_process_func(od_config)
        self.pre_process_func = get_diffusion_pre_process_func(od_config)
        # Cache whether the model-specific postprocess accepts request-level
        # sampling params so step() can support both legacy and extended hooks.
        self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")

    def _resolve_execution_mode(self, od_config: OmniDiffusionConfig) -> DiffusionExecutionMode:
        self.step_execution = bool(getattr(od_config, "step_execution", False))
        if od_config.streaming_output and not self.step_execution:
            logger.warning("streaming_output=True requires step_execution=True; enabling step execution.")
            od_config.step_execution = True
            self.step_execution = True

        if self.step_execution:
            self.supports_request_batch = False
            return DiffusionExecutionMode.STEP_BATCH

        self.supports_request_batch = supports_request_batch(od_config)
        if not self.supports_request_batch and _max_num_seqs(od_config) > 1:
            raise ValueError(
                f"{getattr(od_config, 'model_class_name', None)!r} does not support request-level batching. "
                "Use max_num_seqs=1 for serial request execution, or choose a pipeline with "
                "supports_request_batch=True."
            )
        return DiffusionExecutionMode.REQUEST_BATCH

    def _init_executor(self, od_config: OmniDiffusionConfig) -> None:
        executor_class = DiffusionExecutor.get_class(od_config)
        self.executor = executor_class(od_config)

    def _init_scheduler(
        self,
        od_config: OmniDiffusionConfig,
        scheduler: BaseScheduler | None = None,
    ) -> None:
        if scheduler is not None:
            self.scheduler = scheduler
        elif self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
            self.scheduler = StepScheduler()
        else:
            self.scheduler = RequestScheduler()
        self.scheduler.initialize(od_config)

    def _init_runtime_state(self) -> None:
        # DP multi-concurrency: allow batching dp_size requests so each
        # worker processes a different request in parallel.  Only enabled
        # for distributed layerwise offload (which shards weights and
        # needs all ranks active simultaneously).  Ordinary DP with a
        # non-batch pipeline should not schedule multiple requests.
        dp_size = 1
        if getattr(self.od_config, "parallel_config", None) is not None:
            dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
        dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
        if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True):
            self.scheduler.max_num_running_reqs = dp_size
            self.dp_concurrent = True
            # Set batch admission wait so concurrent requests accumulate
            # before scheduling.  The scheduler reads this from od_config.
            if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0:
                self.od_config.request_batch_max_wait_ms = 500.0
            logger.info(
                "dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms",
                dp_size,
                self.od_config.request_batch_max_wait_ms,
            )
        else:
            self.dp_concurrent = False
        self.main_loop: asyncio.AbstractEventLoop | None = None
        self.stop_event: threading.Event | None = None
        self.worker_thread: threading.Thread | None = None
        self._loop_started = False
        self._init_lock = asyncio.Lock()
        # _rpc_lock is retained solely as the underlying lock for self._cv,
        # which is used to signal the busy loop. Worker-call serialization is
        # now handled structurally by routing all executor calls through the
        # busy loop rather than via mutual exclusion.
        self._rpc_lock = threading.RLock()
        self._cv = threading.Condition(self._rpc_lock)
        self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {}
        self._closed = False
        self._shutdown_complete = False
        self.abort_queue: queue.Queue[str] = queue.Queue()
        self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()

    def _init_execute_fn(self) -> None:
        if self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
            self.execute_fn = self.executor.execute_step
        else:
            self.execute_fn = self.executor.execute_batch

    def _log_execution_mode(self, od_config: OmniDiffusionConfig) -> None:
        if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
            logger.info(
                "[RequestBatch] engine init max_num_seqs=%s max_wait_ms=%s",
                getattr(od_config, "max_num_seqs", None),
                getattr(od_config, "request_batch_max_wait_ms", None),
            )

    async def _check_and_start_background_loop(self):
        if self._closed:
            raise RuntimeError("DiffusionEngine is closed.")
        if self._loop_started:
            return

        async with self._init_lock:
            # double check, in case of lock queue issue
            if self._closed:
                raise RuntimeError("DiffusionEngine is closed.")
            if self._loop_started:
                return

            self.main_loop = asyncio.get_running_loop()
            self.stop_event = threading.Event()
            self.worker_thread = threading.Thread(target=self._busy_loop)
            self.worker_thread.start()
            self._loop_started = True

    async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[list[OmniRequestOutput], None]:
        await self._check_and_start_background_loop()

        diffusion_engine_start_time = time.perf_counter()

        preprocess_time = 0.0
        if self.pre_process_func is not None:
            preprocess_start_time = time.perf_counter()
            request = self.pre_process_func(request)
            preprocess_time = time.perf_counter() - preprocess_start_time
            logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)

        exec_start_time = time.perf_counter()
        generator = self.async_add_req_and_stream_response(request)
        async for output in generator:
            exec_total_time = time.perf_counter() - exec_start_time
            # Async mode: wait for background D2H/SHM to complete.
            if output.async_output_id:
                fut = self.executor.wait_output_ready(output.async_output_id)
                output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT)
            postprocess_start_time = time.perf_counter()
            formatted_outputs = self.postprocess_output(request, output)
            postprocess_time = time.perf_counter() - postprocess_start_time
            step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
            logger.debug(
                "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, "
                "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
                preprocess_time * 1000,
                exec_total_time * 1000,
                postprocess_time * 1000,
                step_total_ms,
            )
            for request_output in formatted_outputs:
                request_output.metrics.update(
                    {
                        "preprocess_time_ms": preprocess_time * 1000,
                        "diffusion_engine_exec_time_ms": exec_total_time * 1000,
                        "diffusion_engine_total_time_ms": step_total_ms,
                        "postprocess_time_ms": postprocess_time * 1000,
                    }
                )
            yield formatted_outputs

    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
        """Deprecated compatibility wrapper over ``step_streaming()``.

        Use ``step_streaming()`` for new callers. This method drains the
        unified output stream and returns only the final output batch, matching
        the historical non-streaming ``step()`` behavior.
        """
        logger.warning_once(
            "DiffusionEngine.step() is deprecated; use step_streaming() and consume the final output batch instead."
        )
        final_output: list[OmniRequestOutput] | None = None
        async for output in self.step_streaming(request):
            final_output = output
        return final_output or []

    def postprocess_output(
        self,
        request: OmniDiffusionRequest,
        output: DiffusionOutput,
    ) -> list[OmniRequestOutput]:
        """Convert a DiffusionOutput to a list of OmniRequestOutput."""
        if output.aborted:
            raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
        if output.error:
            if is_client_error_status(output.error_status_code):
                raise client_error_from_metadata(
                    output.error,
                    status_code=output.error_status_code,
                    error_type=output.error_type,
                )
            raise RuntimeError(output.error)
        logger.debug("Generation completed successfully.")

        if output.output is None:
            logger.warning("Output is None, returning empty OmniRequestOutput")
            return format_empty_diffusion_outputs(request, finished=output.finished)

        # When CPU offload is enabled, move output to CPU before
        # post-processing to avoid device OOM — model weights may still
        # reside on the device and leave no headroom for intermediates.
        output_data = output.output
        if self.od_config.enable_cpu_offload:
            output_data = _move_tensor_tree_to_cpu(output_data)

        if self.post_process_func is not None:
            # Some video pipelines need request-level controls during
            # postprocess (for example worker-side frame interpolation).
            postprocess_kwargs: dict[str, object] = {}
            if self._post_process_accepts_sampling_params:
                postprocess_kwargs["sampling_params"] = request.sampling_params
            outputs = self.post_process_func(output_data, **postprocess_kwargs)
        else:
            outputs = output_data

        postprocess_output = normalize_diffusion_postprocess_output(outputs)

        return format_diffusion_outputs(
            request=request,
            od_config=self.od_config,
            diffusion_output=output,
            output_data=output_data,
            postprocess_output=postprocess_output,
        )

    def _busy_loop(self):
        while not self.stop_event.is_set():
            self._process_aborts_queue()
            self._process_rpc_queue()

            with self._cv:
                while (
                    not self.scheduler.has_requests()
                    and self._rpc_queue.empty()
                    and self.abort_queue.empty()
                    and not self.stop_event.is_set()
                ):
                    self._cv.wait(timeout=1.0)

                if self.stop_event.is_set():
                    break

                if not self.scheduler.has_requests():
                    # Only RPC / abort work pending; loop back to drain it.
                    continue

                if self.supports_request_batch or self.dp_concurrent:
                    self._wait_for_request_batch_admission_locked()

                sched_output = self.scheduler.schedule()

            if sched_output.is_empty:
                self._emit_finished_outputs(sched_output.finished_req_ids, None)
                continue

            try:
                runner_output: BaseRunnerOutput = self.execute_fn(sched_output)  # pyright: ignore[reportAssignmentType]
            except Exception as exc:
                logger.error(
                    "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
                )
                runner_output = BatchRunnerOutput.from_list(
                    [
                        RunnerOutput(
                            request_id=request_id,
                            step_index=None,
                            finished=True,
                            result=DiffusionOutput.from_exception(exc),
                        )
                        for request_id in sched_output.scheduled_request_ids
                    ]
                )

            self._process_aborts_queue()
            self._process_rpc_queue()
            finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
            self._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, runner_output)

        # Engine is stopping: fail any RPCs still queued so callers don't hang.
        self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))

    def _wait_for_request_batch_admission_locked(self) -> None:
        """Wait for compatible requests to accumulate before scheduling a wave.

        Caller must hold ``self._cv``.
        """
        if self.step_execution or (not self.supports_request_batch and not self.dp_concurrent):
            return

        max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0
        if max_wait_s == 0:
            return

        max_batch = self.scheduler.max_num_running_reqs
        waiting = self.scheduler.num_waiting_requests()
        running = self.scheduler.num_running_requests()

        if running > 0:
            return

        start = time.monotonic()
        deadline = start + max_wait_s
        last_waiting = -1
        stable_since = start
        # For dp_concurrent, use a longer stable window so all dp_size
        # HTTP requests have time to land before scheduling.
        if self.dp_concurrent:
            stable_window_s = min(0.3, max_wait_s / 2.0)
        else:
            stable_window_s = min(0.05, max_wait_s / 5.0)

        while not self.stop_event.is_set():
            waiting = self.scheduler.num_waiting_requests()
            now = time.monotonic()

            if waiting >= max_batch:
                break
            if waiting > 0 and (now - stable_since) >= stable_window_s:
                break
            if now >= deadline:
                break

            if waiting > last_waiting:
                stable_since = now
                last_waiting = waiting

            remaining = deadline - now
            self._cv.wait(timeout=min(remaining, 0.002))

        waited_ms = (time.monotonic() - start) * 1000.0
        final_waiting = self.scheduler.num_waiting_requests()
        if final_waiting > 0:
            logger.info(
                "[RequestBatch] admission wait done waiting=%d max_batch=%d waited_ms=%.1f",
                final_waiting,
                max_batch,
                waited_ms,
            )

    def _process_rpc_queue(self) -> None:
        """Execute pending collective_rpc tasks from the busy-loop thread.

        Running these here means executor calls are naturally serialized
        against execute_fn() without any mutual-exclusion locking.
        """
        while True:
            try:
                task = self._rpc_queue.get_nowait()
            except queue.Empty:
                return

            fut = task.future
            if fut.cancelled() or fut.done():
                continue

            remaining: float | None = None
            if task.deadline is not None:
                remaining = task.deadline - time.monotonic()
                if remaining <= 0:
                    if not fut.done():
                        fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution."))
                    continue

            try:
                result = self.executor.collective_rpc(
                    method=task.method,
                    timeout=remaining,
                    args=task.args,
                    kwargs=task.kwargs,
                    unique_reply_rank=task.unique_reply_rank,
                )
            except BaseException as exc:  # noqa: BLE001 - propagate to caller
                # The future may have been cancelled (e.g. by a sync timeout
                # or asyncio cancellation) while the executor call was
                # running. Setting state on a cancelled/done future raises
                # InvalidStateError, which would kill the busy loop.
                if not fut.done():
                    fut.set_exception(exc)
            else:
                if not fut.done():
                    fut.set_result(result)

    def _fail_pending_rpcs(self, exc: BaseException) -> None:
        while True:
            try:
                task = self._rpc_queue.get_nowait()
            except queue.Empty:
                return
            if not task.future.done():
                task.future.set_exception(exc)

    def _emit_finished_outputs(
        self,
        finished_ids: set[str],
        runner_output: BaseRunnerOutput | None = None,
        missing_result_error: str = "Diffusion execution finished without a final output",
    ) -> None:
        for rid in finished_ids:
            if runner_output is not None:
                _output = runner_output.get_request_output(rid)
            else:
                _output = None
            out = self._finalize_finished_request(rid, _output, missing_result_error)
            self._put_output(rid, out)

    def _emit_outputs(
        self,
        finished_ids: set[str],
        scheduled_request_ids: list[str],
        runner_output: BaseRunnerOutput,
    ) -> None:
        """Emit output chunks for every request through the unified output stream."""
        if self.execution_mode != DiffusionExecutionMode.STEP_BATCH:
            self._emit_finished_outputs(finished_ids, runner_output)
            return

        delivered_finished_req_ids: set[str] = set()

        # finished_ids may have some requests that are not scheduler in this round.
        # First handle this-round requests.
        for request_id in scheduled_request_ids:
            req_output = runner_output.get_request_output(request_id)
            if request_id in finished_ids:
                # This entire request is finished (this is the last chunk)
                out = self._finalize_finished_request(
                    request_id,
                    req_output,
                    missing_result_error="Diffusion step execution finished without a final output.",
                )
                self._put_output(request_id, out)
                delivered_finished_req_ids.add(request_id)
            elif req_output is not None and req_output.result is not None:
                # This is a non-terminal chunk. It is not in scheduler's
                # finished_ids, but it still belongs on the request stream.
                self._put_output(request_id, req_output.result)

        # Then handle other requests that are finished in this round.
        for request_id in finished_ids - delivered_finished_req_ids:
            out = self._finalize_finished_request(
                request_id,
                missing_result_error="Diffusion step request finished without execution output.",
            )
            self._put_output(request_id, out)

    def _has_output_stream(self, request_id: str) -> bool:
        with self._cv:
            return request_id in self._out_streams

    @staticmethod
    def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
        """Resolve the engine class selected by ``config.engine_backend``.

        Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a
        ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy
        config's ``engine_backend``). Kept separate from :meth:`make_engine` so the
        selection is testable without constructing an engine (which runs a dummy forward).

        Args:
            config: The configuration for the diffusion engine.

        Returns:
            The ``DiffusionEngine`` (sub)class to instantiate.
        """
        backend = getattr(config, "engine_backend", "default") or "default"

        if isinstance(backend, type):
            if not issubclass(backend, DiffusionEngine):
                raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.")
            return backend
        if backend == "default":
            return DiffusionEngine
        if isinstance(backend, str):
            try:
                engine_class = resolve_obj_by_qualname(backend)
            except (ImportError, ValueError) as e:
                raise ValueError(
                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}"
                ) from e
            if not issubclass(engine_class, DiffusionEngine):
                raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
            return engine_class
        raise ValueError(f"Unknown engine_backend: {backend!r}")

    @staticmethod
    def make_engine(
        config: OmniDiffusionConfig,
        scheduler: BaseScheduler | None = None,
    ) -> DiffusionEngine:
        """Factory method to create the engine selected by ``config.engine_backend``.

        Args:
            config: The configuration for the diffusion engine.
            scheduler: Optional scheduler override. When omitted, the selected
                engine chooses the scheduler from its execution mode.

        Returns:
            An instance of the resolved ``DiffusionEngine`` (sub)class.
        """
        engine_class = DiffusionEngine.resolve_engine_class(config)
        engine = engine_class(config, scheduler=scheduler)
        engine.run_startup_warmup()
        return engine

    def add_request(self, request: OmniDiffusionRequest) -> str:
        with self._cv:
            if self._closed:
                raise RuntimeError("DiffusionEngine is closed.")
            queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
            request_id = self.scheduler.add_request(request)
            self._out_streams[request_id] = queue
            self._cv.notify_all()

        return request_id

    async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
        with self._cv:
            queue = self._out_streams.get(request_id)
        if queue is None:
            raise RuntimeError(f"Request {request_id} not found in output queue.")
        try:
            while True:
                output: DiffusionOutput = await queue.get()
                yield output
                if output.finished:
                    break
        except Exception as e:
            logger.error(f"Wait for response failed: {e}")
            raise
        finally:
            with self._cv:
                if self._out_streams.get(request_id) is queue:
                    self._out_streams.pop(request_id, None)

    def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]:
        request_id = self.add_request(request)
        return self.get_output_stream(request_id)

    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
        """Deprecated compatibility wrapper over ``async_add_req_and_stream_response()``.

        Use ``async_add_req_and_stream_response()`` for new callers. This
        method drains the unified output stream and returns only the final
        ``DiffusionOutput``, matching the historical non-streaming behavior.
        """
        logger.warning_once(
            "DiffusionEngine.async_add_req_and_wait_for_response() is deprecated; "
            "use async_add_req_and_stream_response() and consume the final output instead."
        )
        final_output: DiffusionOutput | None = None
        async for output in self.async_add_req_and_stream_response(request):
            final_output = output
        if final_output is None:
            raise RuntimeError("Diffusion execution completed without an output.")
        return final_output

    def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
        with self._rpc_lock:
            if self._closed:
                raise RuntimeError("DiffusionEngine is closed.")
            target_request_id = self.scheduler.add_request(request)

            # keep scheduling and executing until the target request is finished
            while True:
                self._process_aborts_queue()
                sched_output = self.scheduler.schedule()
                if sched_output.is_empty:
                    if target_request_id in sched_output.finished_req_ids:
                        return self._finalize_finished_request(target_request_id)
                    if not self.scheduler.has_requests():
                        raise RuntimeError("Diffusion scheduler has no runnable requests.")
                    continue

                # NOTE: add_req_and_wait_for_response() is synchronous, will be only called
                # within _dummy_run, only one request will be scheduled
                request_id = sched_output.scheduled_request_ids[0]
                try:
                    runner_output: BaseRunnerOutput = self.execute_fn(sched_output)  # pyright: ignore[reportAssignmentType]
                except EngineDeadError:
                    raise
                except Exception as exc:
                    logger.error("Execution failed for diffusion request %s", request_id, exc_info=True)
                    runner_output = RunnerOutput(
                        request_id=request_id,
                        step_index=None,
                        finished=True,
                        result=DiffusionOutput.from_exception(exc),
                    )

                self._process_aborts_queue()

                finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)

                # sync func should receive one result
                if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
                    raise ValueError("Sync func should receive one result at one time")
                if target_request_id in finished_req_ids:
                    req_output = runner_output.get_request_output(target_request_id)
                    output = self._finalize_finished_request(
                        target_request_id,
                        runner_output=req_output,
                        missing_result_error="Diffusion execution finished without a final output.",
                    )
                    if output.async_output_id:
                        fut = self.executor.wait_output_ready(output.async_output_id)
                        output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT)
                    return output

    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
        """Start or stop profiling on all diffusion workers.

        Args:
            is_start: True to start profiling, False to stop.
            profile_prefix: Optional prefix for trace filename.
        """
        if is_start:
            if profile_prefix is None:
                profile_prefix = f"diffusion_{int(time.time())}"
            logger.info(f"Starting diffusion profiling with prefix: {profile_prefix}")
        else:
            logger.info("Stopping diffusion profiling...")

        try:
            self.collective_rpc(method="profile", args=(is_start, profile_prefix))
        except Exception as e:
            action = "start" if is_start else "stop"
            logger.error(f"Failed to {action} profiling on workers", exc_info=True)
            if is_start:
                raise RuntimeError(f"Could not {action} profiler: {e}") from e

    def run_startup_warmup(self) -> None:
        dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True)
        # Skip dummy run when AllGather is used with more than 1 rank,
        # because the dummy run sends only 1 request but AllGather requires
        # all ranks to participate simultaneously.  This covers both DP > 1
        # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig).
        pc = getattr(self.od_config, "parallel_config", None)
        dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1
        sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1
        effective_shard_size = max(dp_size, sp_size)
        skip_dummy = (
            getattr(self.od_config, "enable_distributed_layerwise_offload", False)
            and dlo_use_allgather
            and effective_shard_size > 1
        )
        if skip_dummy:
            logger.info(
                "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)",
                dp_size,
                sp_size,
            )
            return
        try:
            self._dummy_run()
        except Exception as e:
            logger.error(f"Dummy run failed: {e}")
            self.close()
            raise e

    def _dummy_run(self):
        """A dummy run to warm up the model."""
        num_inference_steps = 1
        height = 512
        width = 512
        prompt: OmniTextPrompt = {"prompt": "dummy run"}

        supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
        if supports_image_input:
            # Provide a dummy image input if the model supports it
            color_format = image_color_format(self.od_config.model_class_name)
            dummy_image = PIL.Image.new(color_format, (width, height))
            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image

        if supports_audio_input:
            audio_sr = 16000
            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio

        num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
        if num_frames <= 0:
            logger.info("Skipping dummy warmup run (num_frames=0)")
            return
        req = OmniDiffusionRequest(
            prompt=prompt,
            request_id=DUMMY_DIFFUSION_REQUEST_ID,
            sampling_params=OmniDiffusionSamplingParams(
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                num_frames=num_frames,
                # Keep warmup path minimal and robust across text encoders.
                # Some models may fail when warmup implicitly triggers
                # classifier-free guidance with an empty negative prompt.
                guidance_scale=0.0,
                num_outputs_per_prompt=1,
                # Disable CFG for warmup to avoid triggering CFG parallel
                # validation when cfg_parallel_size > 1.
                extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
            ),
        )
        logger.info("dummy run to warm up the model")
        request = self.pre_process_func(req) if self.pre_process_func is not None else req
        output = self.add_req_and_wait_for_response(request)
        if output.error:
            raise RuntimeError(f"Dummy run failed: {output.error}")

    def _submit_rpc(
        self,
        method: str,
        timeout: float | None,
        args: tuple,
        kwargs: dict | None,
        unique_reply_rank: int | None,
    ) -> _RpcTask:
        assert isinstance(method, str), "Only string method names are supported for now"
        deadline = None if timeout is None else time.monotonic() + timeout
        task = _RpcTask(
            method=method,
            args=args,
            kwargs=kwargs,
            deadline=deadline,
            unique_reply_rank=unique_reply_rank,
        )
        with self._cv:
            self._rpc_queue.put(task)
            self._cv.notify_all()
        return task

    def collective_rpc(
        self,
        method: str,
        timeout: float | None = None,
        args: tuple = (),
        kwargs: dict | None = None,
        unique_reply_rank: int | None = None,
    ) -> Any:
        """Call a method on worker processes and get results immediately.

        The call is enqueued and executed by the engine's busy loop between
        scheduler steps, so it is naturally serialized against per-request
        execute_fn() invocations without any explicit mutual-exclusion lock.

        Args:
            method: The method name (str) to execute on workers
            timeout: Optional timeout in seconds
            args: Positional arguments for the method
            kwargs: Keyword arguments for the method
            unique_reply_rank: If set, only get reply from this rank

        Returns:
            Single result if unique_reply_rank is provided, otherwise list of results
        """
        assert isinstance(method, str), "Only string method names are supported for now"

        # If the busy loop hasn't started yet (e.g. during _dummy_run in
        # __init__, or before the first async request after construction),
        # there is no busy-loop thread to drain the RPC queue. Fall back to
        # calling the executor directly, but serialize concurrent callers
        # via self._cv's underlying lock so multiple threads in this window
        # cannot race on the shared broadcast_mq / result_mq transport.
        if not self._loop_started:
            with self._cv:
                # Re-check under the lock: the busy loop may have started
                # between the outer check and acquiring the lock, in which
                # case we should use the queued path for proper ordering.
                if not self._loop_started:
                    return self.executor.collective_rpc(
                        method=method,
                        timeout=timeout,
                        args=args,
                        kwargs=kwargs,
                        unique_reply_rank=unique_reply_rank,
                    )

        task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
        try:
            return task.future.result(timeout=timeout)
        except concurrent.futures.TimeoutError as exc:
            task.future.cancel()
            raise TimeoutError(f"RPC call to {method} timed out.") from exc

    async def async_collective_rpc(
        self,
        method: str,
        timeout: float | None = None,
        args: tuple = (),
        kwargs: dict | None = None,
        unique_reply_rank: int | None = None,
    ) -> Any:
        """Async variant of :meth:`collective_rpc` for event-loop callers.

        Enqueue a task keyed by a future and ``await`` the result without
        blocking the loop.
        """
        await self._check_and_start_background_loop()
        task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
        aio_fut = asyncio.wrap_future(task.future)
        try:
            if timeout is None:
                return await aio_fut
            return await asyncio.wait_for(aio_fut, timeout=timeout)
        except asyncio.TimeoutError as exc:
            task.future.cancel()
            raise TimeoutError(f"RPC call to {method} timed out.") from exc

    def _put_queue_output(
        self,
        queue: asyncio.Queue[DiffusionOutput],
        output: DiffusionOutput,
    ) -> None:
        loop = self.main_loop
        if loop is not None and loop.is_running():
            loop.call_soon_threadsafe(queue.put_nowait, output)
        else:
            queue.put_nowait(output)

    def _put_output(self, request_id: str, output: DiffusionOutput) -> None:
        with self._cv:
            queue = self._out_streams.get(request_id)
        if queue is None:
            return
        self._put_queue_output(queue, output)

    def close(self) -> None:
        pending_streams: list[asyncio.Queue[DiffusionOutput]] = []
        with self._cv:
            if self._closed and self._shutdown_complete:
                return
            if not self._closed:
                self._closed = True
                if self.stop_event is not None:
                    self.stop_event.set()
                pending_streams = list(self._out_streams.values())
                self._out_streams.clear()
                self._cv.notify_all()

        closed_output = DiffusionOutput(error="DiffusionEngine is closed.")
        for stream in pending_streams:
            self._put_queue_output(stream, closed_output)

        worker_thread = self.worker_thread
        if worker_thread is not None:
            if worker_thread.is_alive():
                worker_thread.join(timeout=10)
            if worker_thread.is_alive():
                logger.warning(
                    "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred."
                )
                return
            else:
                self._loop_started = False
        else:
            self._loop_started = False

        self.scheduler.close()
        self.executor.shutdown()
        self._shutdown_complete = True

    def abort(self, request_id: str | Iterable[str]) -> None:
        request_ids = [request_id] if isinstance(request_id, str) else list(request_id)

        with self._cv:
            if self._closed:
                return
            for req_id in request_ids:
                self.abort_queue.put(req_id)
            self._cv.notify_all()

    def _process_aborts_queue(self) -> None:
        with self._cv:
            self._drain_abort_queue()

    def _drain_abort_queue(self) -> None:
        if self.abort_queue.empty():
            return

        request_ids: list[str] = []
        while not self.abort_queue.empty():
            ids = self.abort_queue.get_nowait()
            request_ids.extend((ids,) if isinstance(ids, str) else ids)

        self._abort_requests(request_ids)

    def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
        request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)

        for request_id in dict.fromkeys(request_ids):
            if self.scheduler.get_request_state(request_id) is not None:
                self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)

    def _finalize_finished_request(
        self,
        request_id: str,
        runner_output: RunnerOutput | None = None,
        missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
    ) -> DiffusionOutput:
        state = self.scheduler.get_request_state(request_id)
        popped_state = self.scheduler.pop_request_state(request_id)
        state = state or popped_state

        if state is None:
            raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")

        if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
            # Preserve runner-provided abort details when available.
            if runner_output is not None and runner_output.result is not None and runner_output.result.aborted:
                return runner_output.result
            return DiffusionOutput(
                aborted=True,
                abort_message=f"Request {request_id} aborted.",
            )

        if runner_output is not None and runner_output.result is not None:
            return runner_output.result

        if runner_output is not None and runner_output.async_output_id is not None:
            return DiffusionOutput(async_output_id=runner_output.async_output_id)

        return DiffusionOutput(error=missing_result_error)