# 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 _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool: parallel_config = getattr(od_config, "parallel_config", None) dp_size = getattr(parallel_config, "data_parallel_size", 1) return ( dp_size > 1 and getattr(od_config, "enable_distributed_layerwise_offload", False) and getattr(od_config, "dlo_use_allgather", True) ) 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) try: self._init_scheduler(od_config, scheduler) self._init_runtime_state() self._init_execute_fn() self._log_execution_mode(od_config) except Exception: # close() cannot be used because runtime synchronization state may not exist yet. scheduler_to_close = getattr(self, "scheduler", None) if scheduler_to_close is not None: try: scheduler_to_close.close() except Exception: logger.exception("Failed to close Scheduler after DiffusionEngine initialization failed") try: self.executor.shutdown() except Exception: logger.exception("Failed to shut down Executor after DiffusionEngine initialization failed") raise 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 and not _uses_dlo_dp_concurrency(od_config): 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 = getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1) if _uses_dlo_dp_concurrency(self.od_config): self.scheduler.max_num_running_reqs = dp_size self.dp_concurrent = True 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 has_preprocessor = getattr(self, "pre_process_func", None) is not None preprocess_start_time = time.perf_counter() if has_preprocessor else None request = self._prepare_request_for_admission(request) if preprocess_start_time is not None: preprocess_time = time.perf_counter() - preprocess_start_time logger.debug("Pre-processing completed in %.4f seconds", preprocess_time) exec_start_time = time.perf_counter() request_id = self._add_prepared_request(request) generator = self.get_output_stream(request_id) 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 self._wait_for_admission_if_needed_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_admission_if_needed_locked(self) -> None: """Apply scheduler admission policy while holding the engine condition. Caller must hold ``self._cv``. """ start = time.monotonic() decision = self.scheduler.get_admission_wait_decision( now=start, dp_concurrent=self.dp_concurrent, ) if not decision.should_wait: return last_waiting = -1 stable_since = start while not self.stop_event.is_set(): waiting = self.scheduler.num_waiting_requests() now = time.monotonic() if waiting > last_waiting: stable_since = now last_waiting = waiting if self.scheduler.should_end_admission_wait( decision, now=now, stable_since=stable_since, ): break remaining = decision.deadline - now if decision.deadline is not None else 0.002 self._cv.wait(timeout=min(max(remaining, 0.0), 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, decision.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 _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest: """Run model-owned preprocessing once, before entering Engine locks.""" pre_process_func = getattr(self, "pre_process_func", None) if pre_process_func is None: return request return pre_process_func(request) def _add_prepared_request(self, request: OmniDiffusionRequest) -> str: """Admit a request whose model-owned preprocessing is complete.""" 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 def add_request(self, request: OmniDiffusionRequest) -> str: request = self._prepare_request_for_admission(request) return self._add_prepared_request(request) 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: request = self._prepare_request_for_admission(request) 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") output = self.add_req_and_wait_for_response(req) 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)