""" Async Omni Engine for vLLM-Omni multi-stage runtime. AsyncOmniEngine in the caller's thread is a thin proxy that communicates with the Orchestrator (running in a background thread) via janus queues. """ from __future__ import annotations import asyncio import concurrent.futures import dataclasses import json import queue import threading import time import uuid import weakref from collections.abc import Mapping, Sequence from dataclasses import asdict from typing import Any, Literal, cast import janus import torch from omegaconf import OmegaConf from vllm import envs as vllm_envs from vllm.engine.arg_utils import EngineArgs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.config.stage_config import strip_parent_engine_args from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config from vllm_omni.diffusion.diffusion_engine import supports_audio_output from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.messages import ( AbortRequestMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, ShutdownRequestMessage, StageSubmissionMessage, ) from vllm_omni.engine.orchestrator import Orchestrator from vllm_omni.engine.serialization import ( deserialize_additional_information, serialize_additional_information, ) from vllm_omni.engine.stage_client import StageClient from vllm_omni.engine.stage_init_utils import build_stage0_input_processor from vllm_omni.engine.stage_pool import StagePool from vllm_omni.engine.stage_runtime import ( StageRuntimeInfo, create_stage_runtime, ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter logger = init_logger(__name__) _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 # ============================================================================ # Parent-EngineArgs field-routing contracts (consumed by # AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set). # ============================================================================ # Fields that must survive the "equal to default → strip" filter because # diffusion stages need them even when equal to vllm's default value # (e.g. colocate worker setup relies on worker_extension_cls being forwarded). _PARENT_ARGS_KEEP: frozenset[str] = frozenset( { "worker_extension_cls", "allowed_local_media_path", "allowed_media_domains", # Legacy stage-config YAMLs may intentionally leave parallel or # distributed knobs unspecified at the stage level and rely on # top-level CLI values to fill them in during the per-stage merge. # Keep these fields so stages that omit them can inherit CLI values, # while stages with explicit YAML values still win because the legacy # stage-config loader prefers stage-local engine args. "tensor_parallel_size", } ) # Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that # must never leak into per-stage EngineArgs (``stage_configs_path`` would # trigger the ``create_model_config`` guard). _PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) # Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so # their presence as an override is never a surprise — suppress the # "override ignored" warning for these. _PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"}) def _inject_global_id(target: Any, request_id: str) -> None: """Inject global_request_id into a prompt dict's additional_information.""" if isinstance(target, dict): if "additional_information" not in target: target["additional_information"] = {} if target["additional_information"] is None: target["additional_information"] = {} if isinstance(target["additional_information"], dict): target["additional_information"]["global_request_id"] = [str(request_id)] def _upgrade_to_omni_request( request: EngineCoreRequest, raw_prompt: Any, ) -> EngineCoreRequest: """Restore omni-only fields omitted by upstream InputProcessor.""" prompt_embeds = request.prompt_embeds additional_information = None if isinstance(raw_prompt, dict): if prompt_embeds is None: raw_prompt_embeds = raw_prompt.get("prompt_embeds") if isinstance(raw_prompt_embeds, torch.Tensor): prompt_embeds = raw_prompt_embeds additional_information = serialize_additional_information( raw_prompt.get("additional_information"), log_prefix="AsyncOmniEngine", ) if prompt_embeds is None and additional_information is None: return request return OmniEngineCoreRequest.from_request( request, prompt_embeds=prompt_embeds, additional_information=additional_information, ) def _apply_omni_final_stage_metadata( request: EngineCoreRequest, final_stage_id: int, ) -> EngineCoreRequest: """Tag EngineCoreRequest so OmniARScheduler can skip DiT KV when final_stage_id is 0.""" merged: dict[str, Any] = {} if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: merged = deserialize_additional_information(request.additional_information) merged["omni_final_stage_id"] = final_stage_id payload = serialize_additional_information(merged) return OmniEngineCoreRequest.from_request( request, additional_information=payload, ) def _weak_shutdown_async_omni_engine( orchestrator_thread: threading.Thread | None, request_queue: janus.Queue[EngineQueueMessage] | None, output_queue: janus.Queue[EngineQueueMessage] | None, rpc_output_queue: janus.Queue[EngineQueueMessage] | None, ) -> None: """Best-effort orchestrator cleanup for GC finalization.""" try: if request_queue is not None: request_queue.sync_q.put_nowait(ShutdownRequestMessage()) except Exception: pass try: if orchestrator_thread is not None and orchestrator_thread.is_alive(): orchestrator_thread.join() except Exception: pass for q in (request_queue, output_queue, rpc_output_queue): try: if q is not None: q.close() except Exception: pass class AsyncOmniEngine: """Thin proxy that launches an Orchestrator in a background thread. All stage clients, input/output processors, and stage-to-stage transfer logic live inside the Orchestrator coroutine (running in its own thread with a dedicated asyncio event loop). This class communicates with it via janus queues (sync side for callers, async side for orchestrator). Args: model: Model name or path init_timeout: Total timeout waiting for orchestrator startup (seconds). stage_init_timeout: Timeout for stage initialization (seconds) **kwargs: Additional arguments """ # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when stage-init / forward paths touch these attrs. _log_stats: bool = False _coordinator_runtime: Any = None _transfer_emitter: Any = None def __init__( self, model: str, stage_init_timeout: int = 300, init_timeout: int = 600, diffusion_batch_size: int = 1, single_stage_mode: bool = False, transfer_emitter: Any = None, log_stats: bool = False, tokenizer: str | None = None, **kwargs: Any, ) -> None: self.model = model self.tokenizer = tokenizer self.diffusion_batch_size = diffusion_batch_size # Cached by get_diffusion_od_config(). self._diffusion_od_config_view: Any = None startup_timeout = int(init_timeout) # Forwarded into Orchestrator so its _forward_to_next_stage path can # emit per-edge transfer_tx_s / transfer_size_bytes histograms. # Optional: when None, Orchestrator silently skips TX emit (existing # RX path still works via OrchestratorAggregator). self._transfer_emitter = transfer_emitter # Drives upstream EngineCore + scheduler stats production. When False # the engine skips SchedulerStats / IterationStats; the per-(stage, # replica) vllm:* wrap stays registered but reads zero. Respects the # --log-stats CLI flag set by the user via OmniBase. self._log_stats = log_stats logger.info(f"[AsyncOmniEngine] Initializing with model {model}") # ------------------------------------------------------------------ # # Single-stage mode detection # # ------------------------------------------------------------------ # # Single-stage mode is enabled when the caller explicitly passes # # single_stage_mode=True, or when a stage_id is provided in the args. # _stage_id_kwarg = kwargs.get("stage_id") if isinstance(_stage_id_kwarg, int) and not single_stage_mode: single_stage_mode = True self.single_stage_mode: bool = single_stage_mode self._single_stage_id_filter: int | None = ( int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None ) self._omni_master_address: str | None = kwargs.get("omni_master_address") self._omni_master_port: int | None = kwargs.get("omni_master_port") # New omni-coordinator flags. Consumed only in single_stage_mode. # ``omni_dp_size_local`` is process-local: each invocation (head and # every headless) launches that many replicas for its own stage. self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) if self._omni_dp_size_local < 1: raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) if self._omni_heartbeat_timeout <= 0: raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") if single_stage_mode: logger.info( "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", self._single_stage_id_filter, self._omni_master_address, self._omni_master_port, ) self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) self.num_stages = len(self.stage_configs) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None self.prompt_expand_func: Any | None = None self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[OmniSamplingParams] = [] self.stage_metadata: list[StageRuntimeInfo] = [] # Janus queues are constructed eagerly here (not deferred to the # orchestrator thread) so the master server's ROUTER thread always # sees a non-None ``self.request_queue`` when on_register fires. # ``async_q`` lazily binds to whatever event loop first awaits on # it (the orchestrator loop), so cross-thread use stays correct. self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self._shutdown_called = False self._weak_finalizer: weakref.finalize | None = None self._rpc_lock = threading.Lock() self._running_counter = OmniRequestCounter() logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") # Launch orchestrator background thread startup_future: concurrent.futures.Future = concurrent.futures.Future() self.orchestrator_thread = threading.Thread( target=self._bootstrap_orchestrator, args=( stage_init_timeout, startup_future, ), daemon=True, name="orchestrator", ) self.orchestrator_thread.start() self._wait_for_orchestrator_init(startup_future, startup_timeout) # Stage runtime fields are assigned directly on self by the bootstrap thread. self._weak_finalizer = weakref.finalize( self, _weak_shutdown_async_omni_engine, self.orchestrator_thread, self.request_queue, self.output_queue, self.rpc_output_queue, ) logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") def get_diffusion_od_config(self) -> Any: """Expose the diffusion ``model_class_name`` to client-side model-extras. The worker holds the full config; here we just resolve the pipeline class name from the model config (cached). ``model_class_name`` may be ``None``. """ if self._diffusion_od_config_view is None: from types import SimpleNamespace from vllm_omni.diffusion.data import resolve_model_class_name self._diffusion_od_config_view = SimpleNamespace(model_class_name=resolve_model_class_name(self.model)) return self._diffusion_od_config_view def _initialize_stages(self, stage_init_timeout: int) -> None: """Initialize stage clients/processors via StageRuntime and assign to self.""" self._runtime = create_stage_runtime( stage_configs=self.stage_configs, model=self.model, config_path=self.config_path, single_stage_mode=self.single_stage_mode, stage_init_timeout=stage_init_timeout, diffusion_batch_size=self.diffusion_batch_size, async_chunk=self.async_chunk, tokenizer=self.tokenizer, single_stage_id_filter=self._single_stage_id_filter, omni_master_address=self._omni_master_address, omni_master_port=self._omni_master_port, omni_dp_size_local=self._omni_dp_size_local, omni_heartbeat_timeout=self._omni_heartbeat_timeout, omni_lb_policy=self._omni_lb_policy, request_queue=self.request_queue, ) self._runtime.initialize() self.num_stages = len(self.stage_configs) self.stage_pools = self._runtime.stage_pools self.stage_clients = [ cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None ] self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] self.output_processors = [pool.output_processor for pool in self.stage_pools] self.input_processor = ( build_stage0_input_processor(self.stage_vllm_configs[0]) if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None else None ) self.prompt_expand_func = next( ( getattr(client, "prompt_expand_func", None) for client in self.stage_clients if getattr(client, "prompt_expand_func", None) is not None ), None, ) self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] self.stage_metadata = [ StageRuntimeInfo( final_output=client.final_output, final_output_type=client.final_output_type, stage_type=client.stage_type, model_stage=getattr(client, "model_stage", None), ) for client in self.stage_clients ] supported_tasks: set[str] = set() if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): supported_tasks.add("generate") if any(meta.final_output_type == "audio" for meta in self.stage_metadata): supported_tasks.add("speech") self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) def _bootstrap_orchestrator( self, stage_init_timeout: int, startup_future: concurrent.futures.Future, ) -> None: """Create loop, initialize stages, then run Orchestrator.""" loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def _run_orchestrator() -> None: self._initialize_stages(stage_init_timeout) pd_config = self._detect_pd_config() membership_controller = self._runtime.create_membership_controller() orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, rpc_async_queue=self.rpc_output_queue.async_q, stage_pools=self.stage_pools, async_chunk=self.async_chunk, pd_config=pd_config, membership_controller=membership_controller, running_counter=self._running_counter, transfer_emitter=self._transfer_emitter, log_stats=self._log_stats, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) await orchestrator.run() try: loop.run_until_complete(_run_orchestrator()) except Exception as e: if not startup_future.done(): wrapped = RuntimeError(f"Orchestrator initialization failed: {e}") wrapped.__cause__ = e startup_future.set_exception(wrapped) logger.exception("[AsyncOmniEngine] Orchestrator thread crashed") error_text = str(e) or "Orchestrator thread crashed" try: error_msg = ErrorMessage(error=error_text, fatal=True) if self.output_queue is not None: self.output_queue.sync_q.put_nowait(error_msg) if self.rpc_output_queue is not None: self.rpc_output_queue.sync_q.put_nowait(error_msg) except Exception: pass raise finally: try: pending = [task for task in asyncio.all_tasks(loop) if not task.done()] for task in pending: task.cancel() if pending: loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) loop.run_until_complete(loop.shutdown_asyncgens()) if hasattr(loop, "shutdown_default_executor"): loop.run_until_complete(loop.shutdown_default_executor()) except Exception: logger.exception("[AsyncOmniEngine] Failed during orchestrator loop cleanup") finally: asyncio.set_event_loop(None) loop.close() def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: """ Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. """ deadline = time.monotonic() + startup_timeout while True: remaining = deadline - time.monotonic() if remaining <= 0: self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout") raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") try: startup_future.result( timeout=min(remaining, _STARTUP_POLL_INTERVAL_S), ) break except concurrent.futures.TimeoutError: if not self.orchestrator_thread.is_alive(): self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") if startup_future.done(): startup_future.result() # re-raises the real exception raise RuntimeError("Orchestrator thread died during startup") except Exception: self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") raise # ---- request helpers ---- @staticmethod def _iter_multimodal_items(value: Any) -> list[Any]: if value is None: return [] if isinstance(value, list): return value return [value] def _ensure_stage_replica_mm_uuids( self, prompt: Any, *, stage_id: int, replica_id: int, ) -> None: """Make multimodal processor-cache keys local to a stage replica. vLLM's frontend multimodal sender cache is process-global, while each vllm-omni stage replica owns a separate EngineCore receiver cache. If two requests with the same image are routed to different stage-0 replicas, a plain content hash can make the sender omit the tensor for a replica that has never received it. Prefixing user/content UUIDs with the selected replica keeps cache reuse within the receiver that owns it. """ if not isinstance(prompt, dict): return mm_data = prompt.get("multi_modal_data") if not isinstance(mm_data, dict) or not mm_data: return from vllm.multimodal.hasher import MultiModalHasher existing_uuids = prompt.get("multi_modal_uuids") if not isinstance(existing_uuids, dict): existing_uuids = {} model_id = str(getattr(self, "model", "")) scoped_uuids: dict[str, list[str | None]] = dict(existing_uuids) for modality, raw_items in mm_data.items(): items = self._iter_multimodal_items(raw_items) if not items: continue modality_existing = existing_uuids.get(modality) if not isinstance(modality_existing, list): modality_existing = [modality_existing] if modality_existing is not None else [] modality_uuids: list[str | None] = [] for idx, item in enumerate(items): user_uuid = modality_existing[idx] if idx < len(modality_existing) else None if user_uuid is not None: base_uuid = str(user_uuid) elif item is None: base_uuid = None else: base_uuid = MultiModalHasher.hash_kwargs( model_id=model_id, **{modality: item}, ) if base_uuid is None: modality_uuids.append(None) else: modality_uuids.append(f"stage{stage_id}:rep{replica_id}:{base_uuid}") scoped_uuids[modality] = modality_uuids if scoped_uuids: prompt["multi_modal_uuids"] = scoped_uuids @staticmethod def _stage_pool_replica_count(stage_pool: Any) -> int: try: live_num_replicas = getattr(stage_pool, "live_num_replicas", None) if live_num_replicas is not None: return int(live_num_replicas) except Exception: pass try: live_replica_ids = getattr(stage_pool, "live_replica_ids", None) if callable(live_replica_ids): return len(live_replica_ids()) except Exception: pass try: clients = getattr(stage_pool, "clients", None) if clients is not None: return sum(1 for client in clients if client is not None) except Exception: pass return int(getattr(stage_pool, "num_replicas", 1) or 1) @staticmethod def _stage_pool_is_distributed(stage_pool: Any) -> bool: try: is_distributed = getattr(stage_pool, "is_distributed", None) if is_distributed is not None: return bool(is_distributed() if callable(is_distributed) else is_distributed) except Exception: pass return getattr(stage_pool, "_hub", None) is not None def _scope_stage0_multimodal_cache_to_replica( self, request_id: str, prompt: Any, ) -> int | None: stage_pools = getattr(self, "stage_pools", None) if isinstance(prompt, EngineCoreRequest) or not stage_pools: return None stage0_pool = stage_pools[0] # TODO: Currently only supports the ar -> dit process. # Future scenarios (e.g., dit -> ar) need to be added, which will require modifications here. if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1: return None prompts = prompt if isinstance(prompt, list) else [prompt] if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts): return None if self._stage_pool_is_distributed(stage0_pool): preselect_replica_id = getattr(stage0_pool, "preselect_replica_id", None) if not callable(preselect_replica_id): logger.debug( "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing " "without preselect support req=%s", request_id, ) return None replica_id = preselect_replica_id(request_id) if replica_id is None: logger.debug( "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing " "because no serviceable replica is available yet req=%s", request_id, ) return None else: replica_id = stage0_pool.select_replica_id(request_id) for p in prompts: self._ensure_stage_replica_mm_uuids( p, stage_id=0, replica_id=replica_id, ) logger.debug( "[AsyncOmniEngine] Scoped multimodal cache keys to stage-0 replica-%s for req=%s", replica_id, request_id, ) return replica_id def _build_add_request_message( self, request_id: str, prompt: EngineCoreRequest | PromptType, prompt_text: str | None = None, sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, final_output_stage_ids: Sequence[int] | None = None, arrival_time: float | None = None, lora_request: Any = None, tokenization_kwargs: dict[str, Any] | None = None, trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, resumable: bool = False, message_type: Literal["add_request", "streaming_update"] = "add_request", ) -> StageSubmissionMessage: """Build an add_request message after stage-0 preprocessing.""" request_timestamp = float(arrival_time) if arrival_time is not None else time.time() effective_sampling_params_list: list[OmniSamplingParams] = ( list(cast(Sequence[OmniSamplingParams], sampling_params_list)) if sampling_params_list is not None else list(self.default_sampling_params_list) ) if not effective_sampling_params_list: raise ValueError( f"Missing sampling params for stage 0. Got {len(effective_sampling_params_list)} stage params." ) params = effective_sampling_params_list[0] # Keep the original prompt for downstream stages (they need the raw # dict, e.g. for multi_modal_data). original_prompt = prompt preselected_stage0_replica: int | None = None stage_type = self.stage_metadata[0].stage_type output_prompt_text: Any = None _preprocess_ms = 0.0 if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): # Inject global_request_id into the raw prompt. if isinstance(prompt, dict): _inject_global_id(prompt, request_id) elif isinstance(prompt, list): for item in prompt: _inject_global_id(item, request_id) preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica( request_id, prompt, ) # Full input processing (tokenization, multimodal, etc.) _t_preprocess = time.perf_counter() try: request = self.input_processor.process_inputs( request_id=request_id, prompt=prompt, params=params, supported_tasks=self.supported_tasks, arrival_time=arrival_time, lora_request=lora_request, tokenization_kwargs=tokenization_kwargs, trace_headers=trace_headers, priority=priority, data_parallel_rank=data_parallel_rank, resumable=resumable, ) except Exception: if preselected_stage0_replica is not None and self.stage_pools: self.stage_pools[0].release_binding(request_id) raise _preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0 # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have # additional_information field in the prompt. request = _upgrade_to_omni_request(request, prompt) if reasoning_ended is not None: request.reasoning_ended = reasoning_ended # Restore external_req_id to the original user-facing request_id. # InputProcessor.process_inputs() renames request_id to an internal # UUID (saving the original in external_req_id), but then overwrites # external_req_id with the new internal ID. We need external_req_id # to match the key used in Orchestrator.request_states so that # output routing (output.request_id lookup) can find the req_state. request.external_req_id = request_id request = _apply_omni_final_stage_metadata(request, final_stage_id) # Registration with stage 0's output processor is deferred to the # orchestrator thread (see Orchestrator._handle_add_request), which # now routes admission through StagePool.submit_initial(). output_prompt_text = prompt_text if output_prompt_text is None and isinstance(original_prompt, dict): output_prompt_text = original_prompt.get("prompt") prompt = request return StageSubmissionMessage( type=message_type, request_id=request_id, prompt=prompt, original_prompt=original_prompt, output_prompt_text=output_prompt_text, sampling_params_list=effective_sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=list(final_output_stage_ids) if final_output_stage_ids is not None else None, preprocess_ms=_preprocess_ms, request_timestamp=request_timestamp, enqueue_ts=time.perf_counter(), ) def _enqueue_cfg_companions( self, parent_id: str, original_prompt: Any, stage0_params: Any, sampling_params_list: list[Any], ) -> None: """Expand prompt into CFG companions, process through InputProcessor, and enqueue.""" try: expanded = self.prompt_expand_func(original_prompt, stage0_params) except Exception: logger.exception("[AsyncOmniEngine] prompt_expand_func failed for req %s", parent_id) return if not expanded: return for ep in expanded: cid = f"{parent_id}{ep.request_id_suffix}" companion_prompt = ep.prompt companion_params, companion_spl = ep.apply_overrides(stage0_params, sampling_params_list) if isinstance(companion_prompt, dict): _inject_global_id(companion_prompt, cid) request = self.input_processor.process_inputs( request_id=cid, prompt=companion_prompt, params=companion_params, supported_tasks=self.supported_tasks, ) request.external_req_id = cid # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes # admission through StagePool.submit_initial(..., affinity_request_id=...). self.request_queue.sync_q.put( AddCompanionRequestMessage( companion_id=cid, parent_id=parent_id, role=ep.role, prompt=request, companion_prompt_text=companion_prompt, sampling_params_list=companion_spl, ) ) logger.info( "[AsyncOmniEngine] CFG expansion for req %s: %d companions", parent_id, len(expanded), ) @staticmethod def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None: if cache_backend == "cache_dit": return { "Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.24, "max_continuous_cached_steps": 3, "enable_taylorseer": False, "taylorseer_order": 1, "scm_steps_mask_policy": None, "scm_steps_policy": "dynamic", } if cache_backend == "tea_cache": return { "rel_l1_thresh": 0.2, } if cache_backend == "mag_cache": return { "mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } if cache_backend in ("step_cache"): return { "step_cache_dit_enabled": True, "velocity_sim_thresholds": [0.95, 0.93], "velocity_skip_countdowns": [4, 2], "step_cache_dit_min_history": 2, "step_cache_dit_max_history": 2, } return None @staticmethod def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None: if isinstance(cache_config, str): try: cache_config = json.loads(cache_config) except json.JSONDecodeError: logger.warning("Invalid cache_config JSON, using defaults.") cache_config = None if cache_config is None and cache_backend not in (None, "", "none"): cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend) return cache_config def _detect_pd_config(self) -> dict[str, Any] | None: """Detect PD (Prefill-Decode) disaggregation config from stage_configs. Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. """ pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) if pd_pair is None: return None prefill_idx, decode_idx = pd_pair # Extract bootstrap address from prefill stage engine_args bootstrap_addr: str | None = None try: prefill_cfg = self.stage_configs[prefill_idx] ea = getattr(prefill_cfg, "engine_args", None) kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None if kv_cfg is not None: port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" bootstrap_addr = f"http://{kv_ip}:{port}" except Exception as exc: logger.warning("[AsyncOmniEngine] Could not extract PD bootstrap address: %s", exc) logger.info( "[AsyncOmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", prefill_idx, decode_idx, bootstrap_addr, ) prefill_engine_id: str | None = None try: prefill_client = self.stage_clients[prefill_idx] kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) prefill_engine_id = getattr(kv_cfg, "engine_id", None) except Exception as exc: logger.warning("[AsyncOmniEngine] Could not extract prefill engine_id: %s", exc) return { "pd_pair": (prefill_idx, decode_idx), "bootstrap_addr": bootstrap_addr, "prefill_engine_id": prefill_engine_id, } @staticmethod def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: """Create a default single-stage diffusion config from kwargs.""" # We temporally create a default config for diffusion stage. # In the future, we should merge the default config with the user-provided config. normalized_kwargs = dict(kwargs) default_sampling_params = normalized_kwargs.get("default_sampling_params") if isinstance(default_sampling_params, str): try: default_sampling_params = json.loads(default_sampling_params) except json.JSONDecodeError: logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") default_sampling_params = None if not isinstance(default_sampling_params, dict): default_sampling_params = None stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} if normalized_kwargs.get("dtype") is None: normalized_kwargs["dtype"] = "auto" # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): if not isinstance(normalized_kwargs["dtype"], torch.dtype): raise TypeError( f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}" ) normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") cache_backend = normalized_kwargs.get("cache_backend", "none") cache_config = AsyncOmniEngine._normalize_cache_config( cache_backend, normalized_kwargs.get("cache_config", None), ) parallel_config = normalized_kwargs.get("parallel_config") if isinstance(parallel_config, dict): parallel_config = DiffusionParallelConfig.from_dict(parallel_config) if parallel_config is None: ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 ring_degree = normalized_kwargs.get("ring_degree") or 1 ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1 tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False use_hsdp = normalized_kwargs.get("use_hsdp", False) hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) if sequence_parallel_size is None: sequence_parallel_size = ulysses_degree * ring_degree parallel_config = DiffusionParallelConfig( pipeline_parallel_size=pipeline_parallel_size, data_parallel_size=data_parallel_size, 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, ulysses_mode=ulysses_mode, 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, ) num_devices = max(1, int(parallel_config.world_size)) devices = ",".join(str(i) for i in range(num_devices)) model_class_name = kwargs.get("model_class_name", None) final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image" attention_config = None if ( kwargs.get("diffusion_attention_config") is not None or kwargs.get("diffusion_attention_backend") is not None ): attention_config = parse_attention_config( kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), ) stage_engine_args = { "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, "model_class_name": kwargs.get("model_class_name", None), "model_config": kwargs.get("model_config", None), "additional_config": kwargs.get("additional_config", None), "step_execution": kwargs.get("step_execution", False), "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": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "boundary_ratio": kwargs.get("boundary_ratio", None), "flow_shift": kwargs.get("flow_shift", None), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), "distributed_executor_backend": ( "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"] ), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None), "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), "extras": { "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), }, **( { "profiler_config": asdict(kwargs["profiler_config"]) if hasattr(kwargs["profiler_config"], "__dataclass_fields__") else kwargs["profiler_config"] } if kwargs.get("profiler_config") is not None else {} ), } # Only set dtype if it was already explicitly passed and normalized if "dtype" in normalized_kwargs: stage_engine_args["dtype"] = normalized_kwargs["dtype"] # New split fields for diffusers adapter kwargs. if kwargs.get("diffusers_load_kwargs") is not None: stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] if kwargs.get("diffusers_call_kwargs") is not None: stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] default_stage_cfg = [ { "stage_id": 0, "stage_type": "diffusion", "runtime": { "process": True, "devices": devices, }, "engine_args": stage_engine_args, "default_sampling_params": stage_default_sampling_params, "final_output": True, "final_output_type": final_output_type, } ] default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" return default_stage_cfg @staticmethod def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]: """Remove parent ``EngineArgs`` fields from *kwargs*. When ``stage_configs_path`` is set, per-stage engine args are defined in the YAML. Top-level single-engine fields (``compilation_config``, ``tensor_parallel_size``, …) must not leak into per-stage configs via the ``base_engine_args`` merge in ``load_stage_configs_from_yaml`` — they can cause type errors (e.g. ``compilation_config`` as a JSON string rejected by ``VllmConfig``) or silently override YAML values. Logs a warning for any parent field whose value differs from the dataclass default, so users know their explicit overrides are ignored. See the module-level ``_PARENT_ARGS_*`` constants for the routing contracts this method enforces. """ parent_fields: dict[str, dataclasses.Field] = {f.name: f for f in dataclasses.fields(EngineArgs)} result, overridden = strip_parent_engine_args( kwargs, parent_fields=parent_fields, keep_keys=_PARENT_ARGS_KEEP, strip_keys=_PARENT_ARGS_STRIP, no_warn_keys=_PARENT_ARGS_NO_WARN, ) if overridden: logger.warning( "stage_configs_path is set — the following top-level engine " "args are ignored (per-stage YAML takes precedence): %s", ", ".join(sorted(overridden)), ) return result def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) explicit_stage_configs = kwargs.pop("stage_configs", None) if explicit_stage_configs is not None: logger.warning( "`stage_configs` is not part of the public API. " "Ignoring it and resolving stages from stage_configs_path/model factory." ) if stage_configs_path is not None: base_kwargs = self._strip_single_engine_args(kwargs) else: base_kwargs = kwargs # Parse --stage-overrides JSON string if provided stage_overrides = None if stage_overrides_json: if isinstance(stage_overrides_json, str): try: stage_overrides = json.loads(stage_overrides_json) except json.JSONDecodeError as exc: raise ValueError( f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}" ) from exc else: stage_overrides = stage_overrides_json config_path, stage_configs = load_and_resolve_stage_configs( model, stage_configs_path, base_kwargs, default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, ) # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config. for cfg in stage_configs: try: if not hasattr(cfg, "engine_args") or cfg.engine_args is None: cfg.engine_args = OmegaConf.create({}) global_sleep_mode = kwargs.get("enable_sleep_mode") if global_sleep_mode is not None: if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None: cfg.engine_args.enable_sleep_mode = global_sleep_mode if getattr(cfg, "stage_type", None) != "diffusion": continue if not hasattr(cfg, "engine_args") or cfg.engine_args is None: cfg.engine_args = OmegaConf.create({}) additional_config = kwargs.get("additional_config") if additional_config is not None: current_additional_config = getattr(cfg.engine_args, "additional_config", None) if current_additional_config in (None, {}): cfg.engine_args.additional_config = additional_config if kwargs.get("lora_path") is not None: if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None: cfg.engine_args.lora_path = kwargs["lora_path"] lora_scale = kwargs.get("lora_scale") if lora_scale is None: # Backwards compatibility for older callers. lora_scale = kwargs.get("static_lora_scale") if lora_scale is not None: if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None: cfg.engine_args.lora_scale = lora_scale if ( kwargs.get("diffusion_attention_config") is not None or kwargs.get("diffusion_attention_backend") is not None ): has_stage_attention = ( hasattr(cfg.engine_args, "diffusion_attention_config") and cfg.engine_args.diffusion_attention_config is not None ) if not has_stage_attention: cfg.engine_args.diffusion_attention_config = parse_attention_config( kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), ) quantization_config = kwargs.get("diffusion_quantization_config") if quantization_config is not None: if ( not hasattr(cfg.engine_args, "quantization_config") or cfg.engine_args.quantization_config is None ): cfg.engine_args.quantization_config = quantization_config # Inject profiler flags for diffusion stages for profiler_key in ( "enable_diffusion_pipeline_profiler", "enable_ar_profiler", ): val = kwargs.get(profiler_key) if val: if not hasattr(cfg.engine_args, profiler_key) or not getattr( cfg.engine_args, profiler_key, False ): setattr(cfg.engine_args, profiler_key, val) quantization = kwargs.get("quantization") if quantization is not None: if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None: cfg.engine_args.quantization = quantization diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype") if diffusion_kv_cache_dtype is not None: if ( not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype") or cfg.engine_args.diffusion_kv_cache_dtype is None ): cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps") if diffusion_kv_cache_skip_steps is not None: if ( not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps") or cfg.engine_args.diffusion_kv_cache_skip_steps is None ): cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers") if diffusion_kv_cache_skip_layers is not None: if ( not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers") or cfg.engine_args.diffusion_kv_cache_skip_layers is None ): cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers except Exception as e: logger.warning("Failed to inject LoRA config for stage: %s", e) return config_path, stage_configs # ==================== Public API ==================== def add_request( self, request_id: str, prompt: EngineCoreRequest | PromptType, prompt_text: str | None = None, sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, final_output_stage_ids: Sequence[int] | None = None, arrival_time: float | None = None, lora_request: Any = None, tokenization_kwargs: dict[str, Any] | None = None, trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, resumable: bool = False, ) -> None: """Process stage-0 input locally, then send to the Orchestrator. Input processing and output processor registration happen here in the caller's thread, avoiding a queue + coroutine-switch round-trip. The Orchestrator receives a ready-to-submit OmniEngineCoreRequest. """ msg = self._build_add_request_message( request_id=request_id, prompt=prompt, prompt_text=prompt_text, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=final_output_stage_ids, arrival_time=arrival_time, lora_request=lora_request, tokenization_kwargs=tokenization_kwargs, trace_headers=trace_headers, priority=priority, data_parallel_rank=data_parallel_rank, reasoning_ended=reasoning_ended, resumable=resumable, ) self.request_queue.sync_q.put(msg) # CFG companion expansion: create and enqueue companion requests # so the AR stage also generates their KV caches. if self.prompt_expand_func is not None and final_stage_id > 0: original_prompt = msg.original_prompt effective_spl = msg.sampling_params_list stage0_params = effective_spl[0] if effective_spl else None if stage0_params is not None: self._enqueue_cfg_companions(request_id, original_prompt, stage0_params, effective_spl) async def add_request_async( self, request_id: str, prompt: EngineCoreRequest | PromptType, prompt_text: str | None = None, sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, final_output_stage_ids: Sequence[int] | None = None, arrival_time: float | None = None, lora_request: Any = None, tokenization_kwargs: dict[str, Any] | None = None, trace_headers: Mapping[str, str] | None = None, priority: int = 0, data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, resumable: bool = False, ) -> None: """Async add_request API.""" self.add_request( request_id=request_id, prompt=prompt, prompt_text=prompt_text, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=final_output_stage_ids, arrival_time=arrival_time, lora_request=lora_request, tokenization_kwargs=tokenization_kwargs, trace_headers=trace_headers, priority=priority, data_parallel_rank=data_parallel_rank, reasoning_ended=reasoning_ended, resumable=resumable, ) def add_streaming_update( self, request_id: str, prompt: EngineCoreRequest | PromptType, prompt_text: str | None = None, sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, final_output_stage_ids: Sequence[int] | None = None, arrival_time: float | None = None, *, resumable: bool = True, ) -> None: """Send an incremental streaming update for an existing request.""" msg = self._build_add_request_message( request_id=request_id, prompt=prompt, prompt_text=prompt_text, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=final_output_stage_ids, arrival_time=arrival_time, resumable=resumable, message_type="streaming_update", ) self.request_queue.sync_q.put(msg) async def add_streaming_update_async( self, request_id: str, prompt: EngineCoreRequest | PromptType, prompt_text: str | None = None, sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, final_output_stage_ids: Sequence[int] | None = None, arrival_time: float | None = None, *, resumable: bool = True, ) -> None: """Async wrapper for add_streaming_update().""" self.add_streaming_update( request_id=request_id, prompt=prompt, prompt_text=prompt_text, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=final_output_stage_ids, arrival_time=arrival_time, resumable=resumable, ) def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: """Read one output message from the Orchestrator output queue.""" try: return self.output_queue.sync_q.get(timeout=timeout) except queue.Empty: if not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") return None async def try_get_output_async(self) -> EngineQueueMessage | None: """Async read from the Orchestrator output queue.""" try: return self.output_queue.sync_q.get_nowait() except queue.Empty: if not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") return None def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: """Get cached metadata for a stage.""" return self.stage_metadata[stage_id] def abort(self, request_ids: list[str]) -> None: """Send abort message to the Orchestrator.""" if self.request_queue is None: raise RuntimeError("request_queue is not initialized") self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) async def abort_async(self, request_ids: list[str]) -> None: """Async abort API.""" self.abort(request_ids) def collective_rpc( self, method: str, timeout: float | None = None, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, stage_ids: list[int] | None = None, ) -> list[Any]: """Send a control RPC to the Orchestrator and wait for aggregated results. This uses a dedicated RPC output queue so control-plane messages do not race with the normal request output polling loop. """ rpc_id = uuid.uuid4().hex msg = CollectiveRPCRequestMessage( rpc_id=rpc_id, method=method, timeout=timeout, args=tuple(args), kwargs=kwargs or {}, stage_ids=stage_ids, ) with self._rpc_lock: self.request_queue.sync_q.put(msg) deadline = None if timeout is None else time.monotonic() + timeout while True: remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) try: result_msg = self.rpc_output_queue.sync_q.get(timeout=remaining) except queue.Empty as exc: raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") from exc if isinstance(result_msg, ErrorMessage): raise RuntimeError(result_msg.error) if not isinstance(result_msg, CollectiveRPCResultMessage): logger.warning( "[AsyncOmniEngine] Dropping unexpected rpc queue message type=%s", getattr(result_msg, "type", type(result_msg).__name__), ) continue if result_msg.rpc_id != rpc_id: logger.warning( "[AsyncOmniEngine] Dropping mismatched rpc result rpc_id=%s expected=%s", result_msg.rpc_id, rpc_id, ) continue return list(result_msg.results) async def collective_rpc_async( self, method: str, timeout: float | None = None, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, stage_ids: list[int] | None = None, ) -> list[Any]: """Async wrapper around collective_rpc().""" loop = asyncio.get_running_loop() return await loop.run_in_executor( None, lambda: self.collective_rpc( method=method, timeout=timeout, args=args, kwargs=kwargs, stage_ids=stage_ids, ), ) def is_alive(self) -> bool: """Whether the orchestrator thread is alive.""" return bool(self.orchestrator_thread.is_alive()) def shutdown(self) -> None: """Send shutdown message and wait for the Orchestrator thread to exit.""" if getattr(self, "_shutdown_called", False): return self._shutdown_called = True finalizer = getattr(self, "_weak_finalizer", None) if finalizer is not None and finalizer.alive: finalizer.detach() logger.info("[AsyncOmniEngine] Shutting down Orchestrator") if self.request_queue is not None: self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) if self.is_alive(): self.orchestrator_thread.join() for q in (self.request_queue, self.output_queue, self.rpc_output_queue): try: q.close() except Exception: pass if hasattr(self, "_runtime") and self._runtime is not None: try: self._runtime.shutdown() except Exception: logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime") # ── Release CuMem allocator memory pool ────────────────────────────── # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory # Management) allocator holds model weights in a singleton memory pool # that lives in the parent process. Killing the engine-core subprocess # does NOT release this pool — the weights stay resident on the GPU # and can cause CUDA OOM for subsequent engine instances (especially # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). # # Discard mode (level=2) is correct at shutdown: there is no benefit to # keeping a CPU backup when the engine is being torn down. try: from vllm.device_allocator.cumem import CuMemAllocator, cumem_available if cumem_available: allocator = CuMemAllocator.get_instance() # Sleep at level 2 discards all pool memory from the GPU # without creating CPU backups — cheapest and fastest. allocator.sleep() logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") except Exception: pass def _try_shutdown(self, *args, **kwargs) -> None: try: self.shutdown() except Exception: logger.exception(*args, **kwargs)