# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import copy import time import weakref from collections.abc import AsyncGenerator, Callable, Iterable, Sequence from typing import Any, TypeVar from vllm.config import VllmConfig from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.config import OmniModelConfig from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length, try_send_via_connector from vllm_omni.distributed.ray_utils.utils import try_close_ray from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni import OmniBase from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load from vllm_omni.entrypoints.utils import ( get_final_stage_id_for_e2e, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams # Internal imports (our code) from vllm_omni.lora.request import LoRARequest from vllm_omni.metrics import OrchestratorAggregator from vllm_omni.outputs import OmniRequestOutput _R = TypeVar("_R") logger = init_logger(__name__) def _weak_close_cleanup_async( stage_list, stage_in_queues, stage_out_queues, ray_pg, output_handler, zmq_ctx=None, inline_engine=None ): """Weak reference cleanup function for AsyncOmni instances.""" if inline_engine is not None: try: inline_engine.close() except Exception as e: logger.warning("Failed to close inline diffusion engine: %s", e) if stage_list: for q in stage_in_queues: try: q.put_nowait(SHUTDOWN_TASK) except Exception as e: logger.warning(f"Failed to send shutdown signal to stage input queue: {e}") close_fn = getattr(q, "close", None) if callable(close_fn): close_fn() for q in stage_out_queues: close_fn = getattr(q, "close", None) if callable(close_fn): close_fn() for stage in stage_list: try: stage.stop_stage_worker() except Exception as e: logger.warning(f"Failed to stop stage worker: {e}") try_close_ray(ray_pg) # Cancel output handler if output_handler is not None: output_handler.cancel() if zmq_ctx is not None: zmq_ctx.term() class AsyncOmni(OmniBase): """Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models. Similar to the Omni class, but provides an asynchronous interface supporting asynchronous LLM and Diffusion models. Args: model: Model name or path to load. **kwargs: Arbitrary keyword arguments. - stage_configs_path: Optional path to YAML file containing stage configurations. If None, configurations are loaded from the model. - log_stats: Whether to enable statistics logging be written to files with stage-specific suffixes. - stage_init_timeout: Per-stage init watchdog (seconds). Measured from when the previous stage finished (possibly a prior Omni run with GPU reuse/overlap) to when the current stage starts to initialize. - shm_threshold_bytes: Threshold in bytes for using shared memory for IPC. Objects larger than this threshold will use shared memory. - worker_backend: Backend for worker processes. Default is "multi_process". - ray_address: Address of Ray cluster for Ray backend, if using Ray backend. - batch_timeout: Timeout in seconds for batching requests within a stage - init_timeout: Timeout in seconds for waiting for all stages to initialize - Additional keyword arguments passed to stage engines. Example: >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B") >>> async for output in async_llm.generate( ... prompt="Hello", ... request_id="req-1", ... sampling_params_list=[SamplingParams(), SamplingParams()] ... ): ... print(output) """ def __init__(self, model: str, **kwargs: dict[str, Any]) -> None: # Pause/resume control attributes self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False # Sleep mode tracking self._is_sleeping: bool = False # Request state tracking self.request_states: dict[str, ClientRequestState] = {} self.output_handler: asyncio.Task | None = None # RPC results storage: {stage_id: {rpc_id: result}} # Used to avoid race condition between output_handler and collective_rpc self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {} # CFG companion → parent request ID mapping for output routing self._companion_to_parent: dict[str, str] = {} super().__init__(model, **kwargs) # Register weak reference cleanup (called on garbage collection) self._weak_finalizer = weakref.finalize( self, _weak_close_cleanup_async, self.stage_list, self._stage_in_queues, self._stage_out_queues, self._ray_pg, self.output_handler, self._zmq_ctx, getattr(self, "_inline_engine", None), ) async def get_supported_tasks(self) -> set[str]: """Return supported tasks based on stage output modalities and capabilities.""" tasks: set[str] = set() if "text" in self.output_modalities or any(stage.is_comprehension for stage in self.stage_list): tasks.add("generate") if "audio" in self.output_modalities: tasks.add("speech") return tasks def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: """Create default diffusion stage configuration.""" # TODO: here is different from the Omni class. We should merge the two in the future. cache_backend = kwargs.get("cache_backend", "none") cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None)) devices = "0" if "parallel_config" in kwargs: parallel_config = kwargs["parallel_config"] num_devices = kwargs["parallel_config"].world_size for i in range(1, num_devices): devices += f",{i}" else: ulysses_degree = kwargs.get("ulysses_degree") or 1 ring_degree = kwargs.get("ring_degree") or 1 sequence_parallel_size = kwargs.get("sequence_parallel_size") tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1 enable_expert_parallel = kwargs.get("enable_expert_parallel") or False cfg_parallel_size = kwargs.get("cfg_parallel_size") or 1 vae_patch_parallel_size = kwargs.get("vae_patch_parallel_size") or 1 use_hsdp = kwargs.get("use_hsdp", False) hsdp_shard_size = kwargs.get("hsdp_shard_size", -1) hsdp_replicate_size = kwargs.get("hsdp_replicate_size", 1) if sequence_parallel_size is None: sequence_parallel_size = ulysses_degree * ring_degree # Calculate num_devices: consider standalone HSDP other_parallel_size = sequence_parallel_size * tensor_parallel_size * cfg_parallel_size if use_hsdp and other_parallel_size == 1 and hsdp_shard_size > 0: # Standalone HSDP: num_devices is determined by HSDP dimensions num_devices = hsdp_shard_size * hsdp_replicate_size else: num_devices = other_parallel_size for i in range(1, num_devices): devices += f",{i}" parallel_config = DiffusionParallelConfig( pipeline_parallel_size=1, data_parallel_size=1, tensor_parallel_size=tensor_parallel_size, enable_expert_parallel=enable_expert_parallel, sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, use_hsdp=use_hsdp, hsdp_shard_size=hsdp_shard_size, hsdp_replicate_size=hsdp_replicate_size, ) default_stage_cfg = [ { "stage_id": 0, "stage_type": "diffusion", "runtime": { "process": True, "devices": devices, "max_batch_size": 1, }, "engine_args": { "parallel_config": parallel_config, "model_class_name": kwargs.get("model_class_name", None), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enforce_eager": kwargs.get("enforce_eager", False), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "quantization": kwargs.get("quantization", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), }, "final_output": True, "final_output_type": "image", } ] default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" return default_stage_cfg def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None: # Store vllm_config received from worker process (may be None for diffusion stages) vllm_config = result.get("vllm_config") if vllm_config is not None: stage.set_vllm_config(vllm_config) tokenizer = result.get("tokenizer") if tokenizer is not None: stage.set_tokenizer(tokenizer) is_tracing_enabled = result.get("is_tracing_enabled") if is_tracing_enabled is not None: stage.set_is_tracing_enabled(is_tracing_enabled) super()._process_stage_ready(stage, stage_id, result) def _wait_for_stages_ready(self, timeout: int = 120) -> None: """Wait for all stages to report readiness.""" super()._wait_for_stages_ready(timeout) for stage in self.stage_list: if stage.vllm_config is not None and stage.tokenizer is not None: try: vllm_config = stage.vllm_config # Initialize input_processor # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config self.input_processor = OmniInputProcessor( vllm_config=vllm_config, ) # Initialize model_config self.model_config = vllm_config.model_config # Initialize io_processor io_processor_plugin = self.model_config.io_processor_plugin self.io_processor = get_io_processor(vllm_config, io_processor_plugin) logger.info( f"[{self._name}] Initialized input_processor, " f"io_processor, and model_config from stage-{stage.stage_id}", ) break except Exception as e: logger.warning( f"[{self._name}] Failed to initialize processors from stage-{stage.stage_id}: {e}", ) # If no LLM stage found, set processors to None if not hasattr(self, "input_processor") or self.input_processor is None: logger.warning( f"[{self._name}] No LLM stage found, processors will not be available. " "This may cause issues with OpenAIServingModels." ) self.input_processor = None self.io_processor = None self.model_config = None def _setup_rpc_result_checkers(self) -> None: """Override base class to use async-friendly RPC result checkers. In the async path the output_handler task drains the output queues and stashes collective_rpc results into ``_rpc_results``. The checker therefore only needs to look there (no queue draining). """ for stage in self.stage_list: sid = stage.stage_id def make_rpc_checker(stage_id: int): def rpc_checker(rpc_id: str) -> dict[str, Any] | None: if stage_id in self._rpc_results and rpc_id in self._rpc_results[stage_id]: return self._rpc_results[stage_id].pop(rpc_id) return None return rpc_checker stage._rpc_result_checker = make_rpc_checker(sid) def shutdown(self): """Shutdown, cleaning up the background proc and IPC. Alias for close() method. Cleans up all stage processes and inter-process communication resources. """ if hasattr(self, "_weak_finalizer"): self._weak_finalizer() async def generate( self, prompt: OmniPromptType, request_id: str, sampling_params_list: Sequence[OmniSamplingParams] | None = None, *, output_modalities: list[str] | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate outputs for the given prompt asynchronously. Coordinates multi-stage pipeline through YAML configuration. Each stage will use AsyncOmniLLM or AsyncOmniDiffusion based on stage_type. Processes the prompt through all stages in the pipeline and yields outputs as they become available. Each stage uses its corresponding sampling parameters from the sampling_params_list. Args: prompt: Prompt to process. Can be a text string, token IDs, or multimodal prompt. request_id: Unique identifier for this request sampling_params_list: List of SamplingParams, one for each stage. Must have the same length as the number of stages. If None, uses default sampling params for each stage. output_modalities: Optional list of output modalities. Yields: OmniRequestOutput objects as they are produced by each stage. Each output contains the stage_id, final_output_type, and the request_output from that stage. Raises: ValueError: If sampling_params_list has incorrect length. """ # Wait until generation is resumed if the engine is paused. async with self._pause_cond: await self._pause_cond.wait_for(lambda: not self._paused) if self._inline_diffusion: async for output in self._generate_inline(prompt, request_id, sampling_params_list, output_modalities): yield output return logger.debug(f"[{self._name}] generate() called") try: # Start output handler on the first call to generate() self._run_output_handler() # TODO: lora_request, trace_headers, priority are not supported yet if sampling_params_list is None: sampling_params_list = self.default_sampling_params_list if len(sampling_params_list) != len(self.stage_list): raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}") # Orchestrator keeps stage objects for input derivation num_stages = len(self.stage_list) # Track per-request start time for end-to-end timing _req_start_ts: dict[int, float] = {} _wall_start_ts: float = time.time() # _last_finish_ts: float = _wall_start_ts # Determine the final stage for E2E stats (highest stage_id with # final_output=True; fallback to last stage) final_stage_id_for_e2e = get_final_stage_id_for_e2e( output_modalities, self.output_modalities, self.stage_list ) # Metrics/aggregation helper metrics = OrchestratorAggregator( num_stages=num_stages, log_stats=self.log_stats, wall_start_ts=_wall_start_ts, final_stage_id_for_e2e=final_stage_id_for_e2e, ) req_state = ClientRequestState(request_id) req_state.metrics = metrics self.request_states[request_id] = req_state # Ensure modalities is in the prompt dict for CFG expansion # (offline path includes it; online serving passes it separately) if isinstance(prompt, dict) and output_modalities and "modalities" not in prompt: prompt["modalities"] = output_modalities # CFG companion tracking (prompt expansion + lifecycle management) cfg = CfgCompanionTracker( prompt_expand_func=getattr(self.stage_list[0], "prompt_expand_func", None), stage0_sampling_params=sampling_params_list[0], ) expanded_companions = cfg.expand_prompts({request_id: prompt}) sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] task = { "request_id": request_id, "engine_inputs": prompt, "sampling_params": sp0, } self.stage_list[0].submit(task) # Submit CFG companion requests to stage-0 if cfg.is_active: for companion_id, companion_prompt in expanded_companions: self._companion_to_parent[companion_id] = request_id companion_task = { "request_id": companion_id, "engine_inputs": companion_prompt, "sampling_params": cfg.stage0_sampling_params, } self.stage_list[0].submit(companion_task) metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time() _req_start_ts[request_id] = time.time() logger.info( f"[{self._name}] Entering scheduling loop: stages={num_stages}, final_stage={final_stage_id_for_e2e}" ) if self.async_chunk: stage_queues = {stage_id: asyncio.Queue() for stage_id in range(num_stages)} req_state.stage_queues = stage_queues async for output in self._process_async_results( request_id, prompt, sampling_params_list, req_state, metrics, final_stage_id_for_e2e, ): yield output else: async for output in self._process_sequential_results( request_id, req_state, metrics, final_stage_id_for_e2e, sampling_params_list, prompt, cfg=cfg, ): yield output logger.debug(f"[{self._name}] Request {request_id} finalized at stage-{final_stage_id_for_e2e}") try: # Finalize E2E metrics if not already done metrics.on_finalize_request( final_stage_id_for_e2e, request_id, _req_start_ts.get(request_id, _wall_start_ts), ) logger.debug(f"[{self._name}] All requests completed") # Summarize and print stats metrics.build_and_log_summary() except Exception as e: logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}") finally: self.request_states.pop(request_id, None) if cfg.is_active: for cid in cfg.get_companion_request_ids(request_id).values(): self._companion_to_parent.pop(cid, None) except (asyncio.CancelledError, GeneratorExit): await self.abort(request_id) logger.info("[AsyncOrchestrator] Request %s aborted.", request_id) raise async def _generate_inline( self, prompt: OmniPromptType, request_id: str, sampling_params_list: Sequence[OmniSamplingParams] | None = None, output_modalities: list[str] | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate using inline diffusion engine (no stage worker subprocess). Eliminates Hop3 IPC overhead by running OmniDiffusion directly in the orchestrator process. The blocking generate() call is offloaded to a thread executor so the asyncio event loop remains responsive. """ _wall_start_ts = time.time() if sampling_params_list is None: sampling_params_list = self.default_sampling_params_list sp0 = sampling_params_list[0] stage = self.stage_list[0] final_stage_id_for_e2e = 0 metrics = OrchestratorAggregator( num_stages=1, log_stats=self.log_stats, wall_start_ts=_wall_start_ts, final_stage_id_for_e2e=final_stage_id_for_e2e, ) metrics.stage_first_ts[0] = time.time() logger.info( "[%s] Inline diffusion generate for request %s", self._name, request_id, ) try: loop = asyncio.get_running_loop() results = await loop.run_in_executor( None, self._inline_engine.generate, prompt, sp0, [request_id], ) for result in results: images = getattr(result, "images", None) or [] finished = getattr(result, "finished", True) output_to_yield = OmniRequestOutput( stage_id=0, final_output_type=stage.final_output_type, request_output=result, images=images, finished=finished, ) metrics.stage_last_ts[0] = time.time() yield output_to_yield try: metrics.on_finalize_request( final_stage_id_for_e2e, request_id, _wall_start_ts, ) metrics.build_and_log_summary() except Exception as e: logger.exception( "[%s] Failed to finalize inline metrics: %s", self._name, e, ) except (asyncio.CancelledError, GeneratorExit): logger.info( "[%s] Inline request %s cancelled.", self._name, request_id, ) raise except Exception as e: logger.exception( "[%s] Inline diffusion failed for request %s: %s", self._name, request_id, e, ) raise async def _process_async_results( self, request_id: str, prompt: Any, sampling_params_list: list[SamplingParams], req_state: ClientRequestState, metrics: OrchestratorAggregator, final_stage_id_for_e2e: int, ) -> AsyncGenerator[OmniRequestOutput, None]: all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)} submit_flag = True _loop_iter = 0 _last_progress_ts = time.time() while not all(all_stages_finished.values()): _loop_iter += 1 for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): if all_stages_finished[stage_id]: continue try: result = req_state.stage_queues[stage_id].get_nowait() except asyncio.QueueEmpty: await asyncio.sleep(0.001) continue _last_progress_ts = time.time() engine_outputs, finished, output_to_yield = self._process_single_result( result, stage, stage_id, metrics, ) if submit_flag and stage_id == 0: submit_flag = False prompt_token_ids = getattr(engine_outputs, "prompt_token_ids", None) if prompt_token_ids is None: prompt_token_ids = [] engine_input = copy.deepcopy(prompt) try: next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) except Exception: raise engine_input["prompt_token_ids"] = [0] * next_prompt_len engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None for _mm_key in ("mm_kwargs", "mm_hashes", "mm_placeholders", "multi_modal_uuids"): engine_input.pop(_mm_key, None) if engine_input.get("type") == "multimodal": engine_input["type"] = "token" for i in range(1, len(self.stage_list)): task = { "request_id": request_id, "engine_inputs": engine_input, "sampling_params": sampling_params_list[i], } self.stage_list[i].submit(task) metrics.stage_first_ts[i] = time.time() all_stages_finished[stage_id] = finished if output_to_yield: yield output_to_yield async def _process_sequential_results( self, request_id: str, req_state: ClientRequestState, metrics: OrchestratorAggregator, final_stage_id_for_e2e: int, sampling_params_list: list[SamplingParams], prompt: Any, cfg: CfgCompanionTracker | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]): cfg_stage0 = stage_id == 0 and cfg is not None and cfg.is_active finished = False while True: if finished and ( not cfg_stage0 or cfg.all_companions_done(request_id) or cfg.is_parent_failed(request_id) ): break result = await req_state.queue.get() if cfg is not None and cfg.is_companion(result.get("request_id", "")): if cfg_stage0: rid = result.get("request_id") if "error" in result: cfg.on_companion_error(rid) else: cfg.on_companion_completed(rid) continue engine_outputs, finished, output_to_yield = self._process_single_result( result, stage, stage_id, metrics, ) if output_to_yield: yield output_to_yield if not isinstance(engine_outputs, list): engine_outputs = [engine_outputs] stage.set_engine_outputs(engine_outputs) # Forward to next stage if there is one next_stage_id = stage_id + 1 if next_stage_id <= final_stage_id_for_e2e: next_stage: OmniStage = self.stage_list[next_stage_id] # Derive inputs for the next stage, record postprocess time with metrics.stage_postprocess_timer(stage_id, request_id): next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt) sp_next: SamplingParams = sampling_params_list[next_stage_id] if cfg is not None and cfg.is_active and not cfg.is_parent_failed(request_id): if isinstance(sp_next, OmniDiffusionSamplingParams): sp_next = copy.deepcopy(sp_next) sp_next.cfg_kv_request_ids = cfg.get_companion_request_ids(request_id) logger.info( "Attaching cfg_kv_request_ids=%s to request %s", sp_next.cfg_kv_request_ids, request_id, ) # Check if we have a connector for this edge connector_key = (str(stage_id), str(next_stage_id)) connector = self.connectors.get(connector_key) sent_via_connector = False if connector: sent_via_connector = try_send_via_connector( connector=connector, stage_id=stage_id, next_stage_id=next_stage_id, req_id=request_id, next_inputs=next_inputs, sampling_params=sp_next, original_prompt=prompt, next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit, metrics=metrics, ) if not sent_via_connector: # Fallback logic removed as we now enforce connector usage. # If no connector is found or send fails, we log an error and raise, # because continuing would cause the request to be silently dropped # and the orchestrator to hang waiting for completion. error_msg = ( f"[{self._name}] Failed to send request {request_id} to stage-{next_stage_id} via connector. " "Configure a connector for this edge or inspect connector logs for details." ) logger.error(error_msg) raise RuntimeError(error_msg) logger.debug(f"[{self._name}] Forwarded request {request_id} to stage-{next_stage_id}") else: logger.debug(f"[{self._name}] Request {request_id} fully completed") def _process_single_result( self, result: dict[str, Any], stage: OmniStage, stage_id: int, metrics: OrchestratorAggregator, ) -> tuple[Any, bool, OmniRequestOutput | None]: """ Process a single result dictionary from a stage. Returns: engine_outputs: The decoded outputs. finished: Whether the stage processing is finished for this request. output_to_yield: An OmniRequestOutput to yield, or None. """ req_id = result.get("request_id") if "error" in result: logger.error( f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}", ) raise RuntimeError(result) engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") if isinstance(engine_outputs, list): engine_outputs = engine_outputs[0] finished = engine_outputs.finished output_to_yield = None if getattr(stage, "final_output", False): # Construct output to yield images = [] if stage.final_output_type == "image": if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images: images = engine_outputs.images elif hasattr(engine_outputs, "images") and engine_outputs.images: images = engine_outputs.images if stage.final_output_type == "image": # Propagate custom_output from inner diffusion output custom_output = {} if isinstance(engine_outputs, OmniRequestOutput): custom_output = engine_outputs.custom_output or {} output_to_yield = OmniRequestOutput( stage_id=stage_id, final_output_type=stage.final_output_type, request_output=engine_outputs, images=images, finished=finished, _custom_output=custom_output, ) else: output_to_yield = OmniRequestOutput( stage_id=stage_id, final_output_type=stage.final_output_type, request_output=engine_outputs, finished=finished, ) # Mark last output time metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) metrics.process_stage_metrics( result=result, stage_type=stage.stage_type, stage_id=stage_id, req_id=req_id, engine_outputs=engine_outputs, finished=finished, final_output_type=stage.final_output_type, output_to_yield=output_to_yield, ) logger.debug( f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing", ) return engine_outputs, finished, output_to_yield def _run_output_handler(self) -> None: if self.output_handler is not None: return stage_list = self.stage_list request_states = self.request_states companion_to_parent = self._companion_to_parent async def output_handler(): try: while True: idle = True for stage_id, stage in enumerate(stage_list): result = stage.try_collect() if result is None: continue idle = False if result.get("type") == "stage_ready": # Only happens when stage is initialized slower than expected, # so we wait for a short time and try again await asyncio.sleep(0.05) continue # Handle collective_rpc results separately to avoid # race condition with the polling in collective_rpc if result.get("type") == "collective_rpc_result": rpc_id = result.get("rpc_id") if rpc_id: if stage_id not in self._rpc_results: self._rpc_results[stage_id] = {} self._rpc_results[stage_id][rpc_id] = result continue req_id = result.get("request_id") req_state = request_states.get(req_id) if req_state is None: parent_id = companion_to_parent.get(req_id) if parent_id is not None: req_state = request_states.get(parent_id) if req_state is None: logger.debug( f"[{self._name}] Request may have been aborted; \ dropping output for req {req_id} at stage-{stage_id}" ) continue if hasattr(req_state, "stage_queues") and stage_id in req_state.stage_queues: await req_state.stage_queues[stage_id].put(result) else: # Fallback to old behavior for compatibility await req_state.queue.put(result) req_state.stage_id = stage_id if idle: await asyncio.sleep(0.001) # Avoid CPU overload when idle else: await asyncio.sleep(0) except Exception as e: logger.exception("AsyncOmni output_handler failed.") for req_state in request_states.values(): error_msg = {"request_id": req_state.request_id, "error": str(e)} # Send error to all stage queues if hasattr(req_state, "stage_queues"): for queue in req_state.stage_queues.values(): await queue.put(error_msg) else: await req_state.queue.put(error_msg) error_msg = {"request_id": req_state.request_id, "error": str(e)} self.output_handler = None # Make possible for restart self.output_handler = asyncio.create_task(output_handler()) @property def is_running(self) -> bool: if self._inline_diffusion: return self._inline_engine is not None return len(self._stage_in_queues) > 0 @property def is_stopped(self) -> bool: return self.errored @property def errored(self) -> bool: return not self.is_running @property def _name(self) -> str: return "AsyncOrchestrator" @property def is_async(self) -> bool: return True @property def dead_error(self) -> BaseException: return EngineDeadError() async def abort(self, request_id: str | Iterable[str]) -> None: if self._inline_diffusion: if self._inline_engine is not None: self._inline_engine.engine.abort(request_id) return None abort_task = {"type": OmniStageTaskType.ABORT, "request_id": request_id} for stage in self.stage_list: stage.submit(abort_task) return None async def get_vllm_config(self) -> VllmConfig: for stage in self.stage_list: if stage.is_comprehension: # Use the vllm_config received from worker process if stage.vllm_config is not None: return stage.vllm_config return None async def get_model_config(self) -> OmniModelConfig: for stage in self.stage_list: if stage.is_comprehension: # Use the vllm_config received from worker process if stage.vllm_config is not None: return stage.vllm_config.model_config return None async def get_input_preprocessor(self) -> InputPreprocessor: return None async def get_tokenizer(self) -> TokenizerLike: for stage in self.stage_list: if stage.is_comprehension: return stage.tokenizer return None async def is_tracing_enabled(self) -> bool: for stage in self.stage_list: if stage.is_comprehension: return stage.is_tracing_enabled return False @property def renderer(self): """Return the renderer from input_processor if available. OMNI: Required by upstream OpenAIServingModels.__init__ which accesses engine_client.renderer. """ if self.input_processor is None: return None return self.input_processor.renderer async def do_log_stats(self) -> None: pass async def check_health(self) -> None: pass async def reset_mm_cache(self) -> None: pass async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool: pass async def collective_rpc( self, method: str | Callable[..., _R], timeout: float | None = None, args: tuple = (), kwargs: dict[str, Any] | None = None, ) -> list[_R]: """Execute an RPC call on all stages asynchronously. Args: method: Name of the method to execute or callable timeout: Optional timeout in seconds args: Positional arguments for the method kwargs: Keyword arguments for the method Returns: List of results from each stage """ self._run_output_handler() # Run synchronous collective_rpc in thread pool to avoid blocking event loop loop = asyncio.get_event_loop() async def run_stage_rpc(stage: OmniStage) -> _R: return await loop.run_in_executor( None, stage.collective_rpc, method, timeout, args, kwargs, ) # Run all stages concurrently results = await asyncio.gather(*[run_stage_rpc(stage) for stage in self.stage_list]) return list(results) async def sleep(self, level: int = 1) -> None: self._is_sleeping = True await self.collective_rpc(method="sleep", args=(level,)) async def wake_up(self, tags: list[str] | None = None) -> None: self._is_sleeping = False await self.collective_rpc(method="wake_up", args=(tags,)) async def is_sleeping(self) -> bool: """Check whether the engine is sleeping""" return getattr(self, "_is_sleeping", False) async def add_lora(self, lora_request: LoRARequest) -> bool: """Load a new LoRA adapter into the engine for future requests.""" result = await self.collective_rpc(method="add_lora", args=(lora_request,)) return result[0][0] async def remove_lora(self, adapter_id: int) -> bool: """Remove a LoRA adapter from the engine.""" result = await self.collective_rpc(method="remove_lora", args=(adapter_id,)) return result[0][0] async def list_loras(self) -> list[int]: """List all loaded LoRA adapter IDs.""" result = await self.collective_rpc(method="list_loras") return result[0][0] async def pin_lora(self, adapter_id: int) -> bool: """Pin a LoRA adapter so it is not evicted from the cache.""" result = await self.collective_rpc(method="pin_lora", args=(adapter_id,)) return result[0][0] async def encode( self, *args, **kwargs, ): """Generate outputs for a request from a pooling model.""" raise NotImplementedError("encode() is not implemented for AsyncOmni") async def start_profile(self, stages: list[int] | None = None) -> None: """Start profiling for specified stages. Async wrapper around the base implementation for API consistency. Args: stages: List of stage IDs to start profiling. If None, starts profiling for all stages that have profiling enabled. Example: >>> await async_omni.start_profile() >>> async for output in async_omni.generate(...): ... pass >>> await async_omni.stop_profile() """ super().start_profile(stages) async def stop_profile(self, stages: list[int] | None = None) -> None: """Stop profiling for specified stages. Async wrapper around the base implementation for API consistency. Args: stages: List of stage IDs to stop profiling. If None, stops profiling for all stages. Example: >>> await async_omni.start_profile() >>> async for output in async_omni.generate(...): ... pass >>> await async_omni.stop_profile() """ super().stop_profile(stages) async def pause_generation( self, *, wait_for_inflight_requests: bool = False, clear_cache: bool = True, ) -> None: """ Pause generation to allow model weight updates. New generation/encoding requests are blocked until resume. Args: wait_for_inflight_requests: When ``True`` waits for in-flight requests to finish before pausing. When ``False`` (default), immediately aborts any in-flight requests. clear_cache: Whether to clear KV cache and prefix cache after draining. Set to ``False`` to preserve cache for faster resume. Default is ``True`` (clear caches). """ async with self._pause_cond: if self._paused: return self._paused = True # Note: AsyncOmni uses a stage-based architecture without a central # output_processor. For now, we simply set the pause flag and let # new requests wait. In-flight requests will complete naturally. # TODO: Implement request abortion for stages if needed. # Clear cache if requested if clear_cache: await self.reset_prefix_cache() await self.reset_mm_cache() async def resume_generation(self) -> None: """Resume generation after :meth:`pause_generation`.""" async with self._pause_cond: self._paused = False self._pause_cond.notify_all() # Wake up all waiting requests async def is_paused(self) -> bool: """Return whether the engine is currently paused.""" async with self._pause_cond: return self._paused