# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Orchestrator for vLLM-Omni multi-stage runtime. Runs inside a background thread with its own asyncio event loop. Owns logical request progression across stage pools and handles stage-to-stage transfer logic. Distributed membership (replica attach/detach, hub monitoring) is handled by :class:`MembershipController`, which is injected optionally. """ from __future__ import annotations import asyncio import os import shutil import time as _time from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass, field from http import HTTPStatus from typing import Any import janus import torch from vllm.config import ModelConfig from vllm.logger import init_logger from vllm.outputs import CompletionOutput, RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, InteractionMessage, OutputMessage, RegisterRemoteReplicaMessage, ShutdownRequestMessage, StageMetricsMessage, StageSubmissionMessage, UnregisterRemoteReplicaMessage, ) from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE, OmniClientError from vllm_omni.metrics import definitions as metric_defs from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: for artifact_dir in artifact_dirs: shutil.rmtree(artifact_dir, ignore_errors=True) # VLLM_OMNI_EVENT_DRIVEN_ORCH=1 switches the orchestration loop (and the # serving-side final-output drain in entrypoints/async_omni.py) from the legacy # 1 ms poll cadence to event-driven wakeups: one reader task per live LLM stage # replica awaits `client.get_output_async()` directly — the same pattern vLLM's # own AsyncLLM output handler uses — and feeds a single serial dispatch queue. # Default off; the legacy poll loop remains the fallback. _EVENT_DRIVEN_ORCH_ENV = "VLLM_OMNI_EVENT_DRIVEN_ORCH" # How often the event-driven loop reconciles its reader-task set against # `available_replica_ids()` (elastic membership, replica eviction) while idle. _ORCH_READER_RECONCILE_INTERVAL_S = 0.5 def _event_driven_orch_enabled() -> bool: return os.environ.get(_EVENT_DRIVEN_ORCH_ENV, "0").strip().lower() in ("1", "true", "yes", "on") def _build_terminal_empty_output( request_id: str, *, final_output_type: str | None, audio_sample_rate: int = 24000, ) -> RequestOutput: """Build a terminal empty output when no downstream stage input exists.""" completion = CompletionOutput( index=0, text="", token_ids=[], cumulative_logprob=None, logprobs=None, finish_reason="stop", stop_reason=None, ) if final_output_type == "audio": completion.multimodal_output = { "audio": torch.zeros((0,), dtype=torch.float32), "sr": audio_sample_rate, } return RequestOutput( request_id=request_id, prompt=None, prompt_token_ids=[], prompt_logprobs=None, outputs=[completion], finished=True, ) def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], params: SamplingParams | PoolingParams, arrival_time: float | None = None, model_config: ModelConfig | None = None, resumable: bool = False, mm_features: list | None = None, ) -> OmniEngineCoreRequest: """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt.""" if arrival_time is None: arrival_time = _time.time() prompt_token_ids = prompt["prompt_token_ids"] sampling_params = None pooling_params = None if isinstance(params, SamplingParams): sampling_params = params.clone() if model_config is not None: remaining = model_config.max_model_len - len(prompt_token_ids) if sampling_params.max_tokens is None: sampling_params.max_tokens = remaining # ``check_stop`` returns early while ``min_tokens`` is unmet, so a # min_tokens wider than the leftover context window never reaches # the length cap: the request stalls once the scheduler can grant # no further tokens. Stage prompts are built here, so clamp now. if sampling_params.min_tokens > remaining: sampling_params.min_tokens = max(0, remaining) else: pooling_params = params.clone() prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds") raw_additional_information = prompt.get("additional_information") model_intermediate_buffer = prompt.get("model_intermediate_buffer") wire_payload: dict[str, Any] | None = None if isinstance(raw_additional_information, dict): wire_payload = dict(raw_additional_information) additional_info_payload = serialize_additional_information( wire_payload, log_prefix=f"build_engine_core_request_from_tokens req={request_id}", ) return OmniEngineCoreRequest( request_id=request_id, prompt_token_ids=prompt_token_ids, mm_features=mm_features, sampling_params=sampling_params, pooling_params=pooling_params, arrival_time=arrival_time, lora_request=getattr(params, "lora_request", None), cache_salt=prompt.get("cache_salt"), data_parallel_rank=None, prompt_embeds=prompt_embeds, resumable=resumable, additional_information=additional_info_payload, model_intermediate_buffer=model_intermediate_buffer if isinstance(model_intermediate_buffer, dict) else None, ) @dataclass class OrchestratorRequestState: """Per-request bookkeeping inside the Orchestrator.""" request_id: str prompt: Any = None sampling_params_list: list[Any] = field(default_factory=list) final_stage_id: int = -1 final_output_stage_ids: set[int] = field(default_factory=set) finished_final_output_stage_ids: set[int] = field(default_factory=set) # Wall-clock timestamp when the client-facing engine request was accepted. request_timestamp: float = 0.0 # Metrics: timestamp when request was submitted to each stage. stage_submit_ts: dict[int, float] = field(default_factory=dict) mm_processor_kwargs: dict | None = None mm_features: list | None = None pd_prefill_multimodal_output: dict[str, Any] | None = None streaming: StreamingInputState = field(default_factory=lambda: StreamingInputState()) # Per-request pipeline timing accumulator (milliseconds) pipeline_timings: dict[str, float] = field(default_factory=dict) # A session-owned request's lifetime belongs to a session (duplex), not to # its own ``finished`` flag: no synthetic terminal, no client emission of # stage-0 segment ends, no terminal re-forward, no cleanup on finish. session_owned: bool = False running_counter_registered: bool = False request_artifact_dirs: set[str] = field(default_factory=set) @dataclass class StreamingSegmentState: """Streaming segment-boundary state for one stage of a request. Every stage of a multi-stage pipeline reaches its own segment boundaries independently, so this is tracked per stage rather than per request. """ # Flag of segment of streaming input finished finished: bool = False # Tokens from the current raw segment boundary. The vLLM output processor # does not guarantee that EngineCoreOutput.new_token_ids survives on the # processed RequestOutput used by the routing layer. token_ids: list[int] = field(default_factory=list) output_metadata: dict[str, Any] = field(default_factory=dict) @dataclass class StreamingInputState: # Flag of streaming input request enabled: bool = False # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll # into this one ``req_state``, while every consumer reads it for a specific # ``stage_id``, so a flat slot is only correct as long as each read happens in # the same iteration as its own stage's write. It is not for a stage that # never records one (the ``poll_diffusion_output`` branch reaches # ``_handle_processed_outputs`` without the raw-output loop) or for any # consumer reading outside the writing iteration. segments: dict[int, StreamingSegmentState] = field(default_factory=dict) # Streaming update prompt length. NOT stage-keyed: its only consumer reads it # through the streaming context on the stage-input-bridge path, not by stage # id. Do not assume the rest of this struct is stage-safe. new_prompt_len_snapshot: int | None = None # Model/bridge-specific runtime states (e.g., thinker->talker) bridge_states: dict[str, Any] = field(default_factory=dict) # Synchronous stage-transition capability installed by the orchestrator # while the downstream input processor consumes upstream token output. source_token_decoder: Callable[..., str] | None = None def segment(self, stage_id: int | None) -> StreamingSegmentState: """Return ``stage_id``'s segment state, or a fresh empty one if unreported. Fresh rather than shared: ``output_metadata`` is handed out by reference. """ segment = self.segments.get(stage_id) return segment if segment is not None else StreamingSegmentState() class OrchestratorBase: """Stage-management loop shared by the turn-based and duplex orchestrators. Runs inside a background thread's asyncio event loop. Subclasses fill the template seams (``_dispatch_message``, ``_background_tasks``, ``_shutdown_extensions``, ``_on_stage_submitted``, ``_intercept_stage_output``, ``_handle_forward_failure``, ``_cleanup_request_ids``). """ # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. _running_counter: OmniRequestCounter | None = None _engines_waiting_counter: OmniRequestCounter | None = None _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], output_async_queue: janus.AsyncQueue[EngineQueueMessage], rpc_async_queue: janus.AsyncQueue[EngineQueueMessage], stage_pools: list[StagePool], *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, engines_waiting_counter: OmniRequestCounter | None = None, transfer_emitter: Any = None, prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue self.rpc_async_queue = rpc_async_queue self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools self.log_stats = log_stats self._prom_metrics = prom_metrics self._stage_replica_waiting: dict[tuple[int, int], int] = {} self._orch_monitor = create_orch_monitor( enabled=enable_orch_monitor, replica_sampler=self._sample_replica_metrics, ) for stage_id, pool in enumerate(self.stage_pools): for replica_id in pool.available_replica_ids(): self._orch_monitor.register_replica(stage_id, replica_id) # PD disaggregation state self._pd_pair: tuple[int, int] | None = None self._pd_bootstrap_addr: str | None = None self._pd_prefill_engine_id: str | None = None self._pd_kv_params: dict[str, Any] = {} if pd_config is not None: self._pd_pair = pd_config.get("pd_pair") self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} self._init_metrics_state( stage_pools, running_counter, transfer_emitter, engines_waiting_counter=engines_waiting_counter, log_stats=log_stats, ) self._cfg_tracker = CfgCompanionTracker() self._stage_input_processors: dict[int, Any] = {} self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._event_driven_orch = _event_driven_orch_enabled() # Distributed membership (optional, injected by DistStageRuntime) self._membership = membership_controller def _init_metrics_state( self, stage_pools: list[StagePool], running_counter: OmniRequestCounter | None, transfer_emitter: Any, engines_waiting_counter: OmniRequestCounter | None = None, log_stats: bool = False, ) -> None: """Wire up all metric-related orchestrator state. Sets ``self._running_counter`` and ``self._transfer_emitter`` (both optional, used by request-add / forward paths), builds the ``(stage_id, replica_id) ↔ engine_idx`` lookup used at record() time, and best-effort constructs the ``OmniPrometheusStatLogger`` wrap that exposes ~37 upstream ``vllm:*`` families with per-(stage, replica) labels. Failure to build the wrap is logged and metrics are simply disabled — orchestrator construction continues so unit tests with a minimal ``vllm_config`` still pass. ``log_stats=False`` short-circuits the wrap entirely so the ~65 upstream ``vllm:*`` families are not registered in the Prometheus default registry at all. The per-step record() path already no-ops on ``scheduler_stats is None`` (which is what the upstream scheduler returns when its own log_stats is False), so this gate is mainly to keep the ``/metrics`` surface clean when the user did not request stats. When stats are enabled, record() must still receive IterationStats even on ticks where scheduler_stats is None; upstream PrometheusStatLogger records token/request metrics from iteration_stats independently of scheduler gauges. """ self._running_counter = running_counter self._engines_waiting_counter = engines_waiting_counter self._transfer_emitter = transfer_emitter # Flat engine_idx ↔ (stage, replica) maps. The reverse map is # consulted at record() time to translate the orchestrator's # (stage_id, replica_id) loop variables into an engine_idx the # underlying PrometheusStatLogger can address. stage_replica_map: dict[int, tuple[str, str]] = {} self._stage_replica_to_engine_idx: dict[tuple[int, int], int] = {} flat_idx = 0 for stage_id, pool in enumerate(stage_pools): for replica_id in range(pool.num_replicas): stage_replica_map[flat_idx] = (str(stage_id), str(replica_id)) self._stage_replica_to_engine_idx[(stage_id, replica_id)] = flat_idx flat_idx += 1 if not log_stats: self._stat_logger = None return vllm_config_for_stats = next( (p.stage_vllm_config for p in stage_pools if p.stage_vllm_config is not None), None, ) if vllm_config_for_stats is None: self._stat_logger = None return try: self._stat_logger = OmniPrometheusStatLogger( vllm_config=vllm_config_for_stats, stage_replica_map=stage_replica_map, ) except Exception: # Minimal vllm_config in unit-test contexts can lack fields the # upstream PrometheusStatLogger expects. Skip wrap rather than # break orchestrator construction. logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled") self._stat_logger = None async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") request_task = asyncio.create_task(self._request_handler(), name="orchestrator-request-handler") output_task = asyncio.create_task( self._orchestration_output_handler(), name="orchestrator-stage-output-handler", ) # Start membership watcher if distributed mode is active. membership_watcher: asyncio.Task[None] | None = None if self._membership is not None: self._membership.install_unregister_handlers( output_queue=self.output_async_queue, cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), replica_removed_callback=self._remove_stage_replica_waiting, ) membership_watcher = self._membership.start() tasks = [request_task, output_task] for index, coro in enumerate(self._background_tasks()): tasks.append(asyncio.create_task(coro, name=f"orchestrator-background-{index}")) if membership_watcher is not None: tasks.append(membership_watcher) try: await asyncio.gather(*tasks) except asyncio.CancelledError: raise except EngineDeadError as e: # The orchestration loop isolates per-replica EngineDeadError # (evict + continue), so this only fires if one escapes outside the # poll handler (e.g. from output post-processing). Treat it as a # teardown signal and let the finally block clean up; do not re-raise. logger.info("[Orchestrator] Engine dead, shutting down: %s", e) # Unblock any collective_rpc callers waiting on rpc_async_queue # instead of leaving them to hit their own timeout. await self.rpc_async_queue.put( ErrorMessage( error=str(e) or "Stage engine died", fatal=True, ) ) except Exception: logger.exception("[Orchestrator] Fatal error in orchestrator tasks") raise finally: self._shutdown_event.set() for task in tasks: if not task.done(): task.cancel() try: await asyncio.gather(*tasks, return_exceptions=True) except Exception: pass await self._shutdown_extensions() if self._membership is not None: await self._membership.drain_tasks(timeout=10.0) self._membership.shutdown() self._orch_monitor.flush() self._shutdown_stages() loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] for task in pending: task.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) # ---- Request handling ---- async def _request_handler(self) -> None: """Read messages from the main thread via request_async_queue.""" while True: msg = await self.request_async_queue.get() msg_type = msg.type if await self._dispatch_message(msg): continue if msg_type == "abort": await self._handle_abort(msg) elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif isinstance(msg, RegisterRemoteReplicaMessage): if self._membership is not None: await self._membership.handle_register(msg.stage_id, msg.replica_id) self._orch_monitor.register_replica(msg.stage_id, msg.replica_id) elif isinstance(msg, UnregisterRemoteReplicaMessage): if self._membership is not None: await self._membership.handle_unregister(msg.stage_id, msg.input_addr) elif isinstance(msg, ShutdownRequestMessage): logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() # Pre-mark stage clients as shutting down to prevent # proc_monitor daemon threads from flagging normal # process exit as EngineDeadError during teardown. for pool in self.stage_pools: for client in pool.clients: if hasattr(client, "_shutting_down"): client._shutting_down = True # Stage teardown runs once in run()'s finally after the # orchestration loop observes _shutdown_event and exits. break else: logger.warning("[Orchestrator] Unknown message type: %s", msg_type) # ---- Template seams filled by the turn-based / duplex subclasses ---- async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: """Handle a subclass-specific request-queue message; return True when consumed.""" del msg return False def _background_tasks(self) -> list[Coroutine[Any, Any, None]]: """Extra coroutines gathered with the request/output handlers in ``run()``.""" return [] async def _shutdown_extensions(self) -> None: """Runs in ``run()``'s ``finally`` after task cancellation, before membership drain.""" def _on_stage_submitted( self, stage_id: int, request_id: str, replica_id: int, req_state: OrchestratorRequestState, ) -> None: """Called after every successful ``submit_initial`` / ``submit_update``.""" del stage_id, request_id, replica_id, req_state async def _intercept_stage_output( self, stage_id: int, replica_id: int, output: Any, req_state: OrchestratorRequestState, stage_metrics: Any, submit_ts: float | None, ) -> bool: """Offer a processed output to the subclass before forwarding. Return True when the output must not be forwarded to the next stage (the subclass consumed it). Client emission for session-owned requests is already suppressed by ``OrchestratorRequestState.session_owned``. """ del stage_id, replica_id, output, req_state, stage_metrics, submit_ts return False async def _handle_forward_failure( self, req_id: str, next_stage_id: int, req_state: OrchestratorRequestState, exc: BaseException, ) -> bool: """Return True when a next-stage input-processing failure was absorbed (else it re-raises).""" del req_id, next_stage_id, req_state, exc return False async def _handle_abort(self, msg: AbortRequestMessage) -> None: """Handle an abort message from the main thread. When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after stage aborts, binding release, and request cleanup complete so the caller can await acknowledgment. Fire-and-forget aborts omit rpc_id. AR final-stage abort outputs (partial tokens) are attached to the result message when ``rpc_id`` is set, or enqueued on ``output_async_queue`` for fire-and-forget aborts. """ request_ids = msg.request_ids error: str | None = None abort_outputs: list[OutputMessage] = [] try: # _cleanup_request_ids is CFG-aware: it expands aborted parents to # their companions and fails a deferred parent whose companion is # aborted before its output arrived. abort_outputs = await self._cleanup_request_ids(list(request_ids), abort=True) logger.info("[Orchestrator] Aborted request(s) %s", request_ids) except Exception as exc: error = str(exc) logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) if msg.rpc_id is None: raise if msg.rpc_id is not None: # Always emit a result when rpc_id is set so CorrelatedRpcClient # waiters are unblocked even on failure. await self.rpc_async_queue.put( AbortResultMessage( rpc_id=msg.rpc_id, success=error is None, error=error, abort_outputs=abort_outputs or None, ) ) elif abort_outputs: for output_msg in abort_outputs: await self.output_async_queue.put(output_msg) async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: """Forward abort requests to all stage pools. Collects final-stage AR abort outputs (partial tokens) while request state is still available for stage-id filtering. Diffusion pools do not contribute abort outputs. """ if not request_ids: return [] abort_outputs: list[OutputMessage] = [] for pool in self.stage_pools: stage_outputs = await pool.abort_requests(request_ids) or [] if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": final_output_type = getattr(pool.stage_client, "final_output_type", None) or "text" for orch_req_id, request_output in stage_outputs: req_state = self.request_states.get(orch_req_id) if req_state is None: continue final_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if pool.stage_id not in final_stage_ids: continue engine_output = OmniRequestOutput.from_stage_output( request_output, request_id=orch_req_id, finished=True, stage_id=pool.stage_id, replica_id=pool.get_bound_replica_id(orch_req_id), final_output_type=final_output_type, ) abort_outputs.append( OutputMessage( request_id=orch_req_id, stage_id=pool.stage_id, replica_id=pool.get_bound_replica_id(orch_req_id), engine_outputs=engine_output, metrics=None, finished=False, stage_submit_ts=req_state.stage_submit_ts.get(pool.stage_id), ) ) pool.release_bindings(request_ids) last_index_by_req: dict[str, int] = {} for index, output_msg in enumerate(abort_outputs): last_index_by_req[output_msg.request_id] = index for index, output_msg in enumerate(abort_outputs): output_msg.finished = index == last_index_by_req[output_msg.request_id] return abort_outputs def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" for pool in self.stage_pools: pool.release_bindings(request_ids) async def _handle_collective_rpc(self, msg: CollectiveRPCRequestMessage) -> None: """Handle a control-plane RPC request from the main thread.""" rpc_id = msg.rpc_id method = msg.method timeout = msg.timeout args = tuple(msg.args) kwargs = dict(msg.kwargs or {}) requested_stage_ids = msg.stage_ids target_pools: list[StagePool] = [] if requested_stage_ids is None: target_pools.extend(self.stage_pools) else: for lid in requested_stage_ids: if not (0 <= lid < self.num_stages): logger.warning("[Orchestrator] collective_rpc: ignoring invalid stage_id %s", lid) continue target_pools.append(self.stage_pools[lid]) results: list[Any] = [] stage_ids: list[int] = [] for pool in target_pools: for replica_id in pool.live_replica_ids(): stage_result = await pool.collective_rpc( replica_id=replica_id, method=method, timeout=timeout, args=args, kwargs=kwargs, ) stage_ids.append(pool.stage_id) results.append(stage_result) await self.rpc_async_queue.put( CollectiveRPCResultMessage( rpc_id=rpc_id, method=method, stage_ids=stage_ids, results=results, ) ) # ---- Orchestration loop ---- def _sample_replica_metrics(self) -> dict[str, tuple[int, int]]: samples: dict[str, tuple[int, int]] = {} for stage_id, pool in enumerate(self.stage_pools): for replica_id in pool.live_replica_ids(): key = replica_key(stage_id, replica_id) samples[key] = pool.replica_monitor_sample(replica_id) return samples async def _orchestration_output_handler(self) -> None: """Poll all stages, handle transfers, send final outputs to main.""" try: if self._event_driven_orch: await self._orchestration_loop_event_driven() else: await self._orchestration_loop() except asyncio.CancelledError: logger.debug("[Orchestrator] _orchestration_output_handler cancelled") return async def _process_llm_stage_outputs( self, stage_id: int, replica_id: int, raw_outputs: Any, raw_terminal_request_ids: set[str], ) -> list[Any]: """Process one raw LLM poll result; returns processed request outputs. Shared by the legacy poll loop and the event-driven dispatcher so both run the exact same per-output handling. Callers own the ``EngineDeadError`` catch, since eviction needs the replica id, and supply ``raw_terminal_request_ids`` — the per-poll accumulator this fills for the caller to drain through ``_finish_raw_terminal_requests`` once routing is done. """ pool = self.stage_pools[stage_id] await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) for eco in raw_outputs.outputs: # Emit kv_wait_s before _handle_kv_ready_raw_outputs' # async_chunk early-return so it lands in all modes. kv_params = getattr(eco, "kv_transfer_params", None) if ( self._prom_metrics is not None and isinstance(kv_params, dict) and (kv_wait_s := kv_params.get("kv_wait_s")) is not None ): self._prom_metrics.observe_kv_wait( kv_params.get("connector_type") or "unknown", float(kv_wait_s), ) req_state = self.request_states.get(getattr(eco, "request_id", None)) if req_state is None or not req_state.streaming.enabled: continue segment_finished = bool(getattr(eco, "is_segment_finished", False)) raw_mm = self._completion_multimodal_output(eco, None) req_state.streaming.segments[stage_id] = StreamingSegmentState( finished=segment_finished, token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []), output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}), ) req_state.streaming.new_prompt_len_snapshot = getattr( eco, "new_prompt_len_snapshot", None, ) if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): raw_terminal_request_ids.add(req_state.request_id) await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state) iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None processed = await pool.process_llm_raw_outputs( replica_id, raw_outputs, iteration_stats=iteration_stats, ) if self._stat_logger is not None and (raw_outputs.scheduler_stats is not None or iteration_stats is not None): self._stat_logger.record( raw_outputs.scheduler_stats, iteration_stats, engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], ) _sched_stats = raw_outputs.scheduler_stats if ( self._prom_metrics is not None and _sched_stats is not None and getattr(_sched_stats, "num_waiting_reqs", None) is not None ): self._update_stage_replica_waiting( stage_id, replica_id, int(_sched_stats.num_waiting_reqs), ) return processed def _absorb_diffusion_metrics(self, stage_id: int, replica_id: int, diffusion_output: Any) -> bool: """Drain a diffusion output's piggybacked metrics. Returns ``True`` when the output carried nothing but metrics and must not be routed. Shared by both orchestration loops: the event-driven poller would otherwise hand a metrics-only sentinel to ``_handle_processed_outputs`` as if it were a real request output. """ output_metrics = getattr(diffusion_output, "metrics", None) if isinstance(output_metrics, dict): n_waiting = output_metrics.pop(metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY, None) if n_waiting is not None: self._update_stage_replica_waiting(stage_id, replica_id, int(n_waiting)) return diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID async def _orchestration_loop(self) -> None: """Poll stage pools and route logical outputs.""" while not self._shutdown_event.is_set(): idle = True for stage_id in range(self.num_stages): pool = self.stage_pools[stage_id] for replica_id in pool.available_replica_ids(): if self._shutdown_event.is_set(): return raw_terminal_request_ids: set[str] = set() # Shared catch so a dead replica on either poll path is # evicted rather than tearing down every stage (#4285). try: if pool.stage_type == "diffusion": diffusion_output = pool.poll_diffusion_output(replica_id) if diffusion_output is None: continue if self._absorb_diffusion_metrics(stage_id, replica_id, diffusion_output): idle = False continue if not is_diffusion_request_started_output(diffusion_output): pool.record_output_timestamps([diffusion_output]) processed = [diffusion_output] else: raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001) if raw_outputs is None: continue processed = await self._process_llm_stage_outputs( stage_id, replica_id, raw_outputs, raw_terminal_request_ids ) except asyncio.CancelledError: raise except EngineDeadError as e: await self._handle_dead_replica(stage_id, replica_id, e) continue except Exception: if self._shutdown_event.is_set(): return logger.exception( "[Orchestrator] Stage-%s replica-%s processing failed", stage_id, replica_id, ) raise await self._handle_processed_outputs(stage_id, replica_id, processed) await self._finish_raw_terminal_requests(stage_id, replica_id, raw_terminal_request_ids) idle = False self._orch_monitor.note_loop(idle=idle) if idle: await asyncio.sleep(0.001) else: await asyncio.sleep(0) async def _orchestration_loop_event_driven(self) -> None: """Event-driven variant of ``_orchestration_loop``. Selected by ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``. One reader task per available LLM stage replica awaits ``client.get_output_async()`` directly — the same pattern vLLM's own ``AsyncLLM`` output handler uses — and feeds a single dispatch queue. This coroutine consumes that queue serially, so routing/handling semantics are identical to the legacy loop; only the 1 ms poll cadence (and its per-tick ``asyncio.wait_for`` task churn) is removed. Diffusion stages keep their nowait-poll contract via a per-pool poller task feeding the same queue. """ ready_q: asyncio.Queue[tuple[str, int, int, Any]] = asyncio.Queue() readers: dict[tuple[int, int], tuple[asyncio.Task, Any]] = {} pollers: dict[int, asyncio.Task] = {} reaped: list[asyncio.Task] = [] async def _llm_replica_reader(stage_id: int, replica_id: int, client: Any) -> None: try: while not self._shutdown_event.is_set(): raw_outputs = await client.get_output_async() # Same keep/drop rule as StagePool._poll_stage_raw: a batch # with no request outputs still carries SchedulerStats on # throttled ticks, and dropping it loses the KV/queue gauges # for that interval. Only a fully empty batch is dropped. A # poll-style client returns one instead of blocking, so keep # the legacy 1 ms cadence for those; a blocking client only # lands here rarely, where 1 ms is irrelevant. if ( not raw_outputs.outputs and raw_outputs.scheduler_stats is None and not raw_outputs.finished_requests ): await asyncio.sleep(0.001) continue await ready_q.put(("llm", stage_id, replica_id, raw_outputs)) except asyncio.CancelledError: raise except BaseException as e: # noqa: BLE001 - routed to the dispatcher await ready_q.put(("error", stage_id, replica_id, e)) async def _diffusion_poller(stage_id: int, pool: StagePool) -> None: try: while not self._shutdown_event.is_set(): got = False for replica_id in pool.available_replica_ids(): try: output = pool.poll_diffusion_output(replica_id) except EngineDeadError as e: await ready_q.put(("error", stage_id, replica_id, e)) continue if output is None: continue await ready_q.put(("diffusion", stage_id, replica_id, output)) got = True await asyncio.sleep(0 if got else 0.001) except asyncio.CancelledError: raise except BaseException as e: # noqa: BLE001 - routed to the dispatcher await ready_q.put(("error", stage_id, -1, e)) def _reconcile_readers() -> None: """Track available replicas: spawn/reap LLM readers and diffusion pollers. Walks ``available_replica_ids()`` for parity with the legacy loop, so a replica evicted by ``_handle_dead_replica`` (or marked unavailable by membership) stops being read here too. Diffusion pollers are reconciled here rather than spawned once at startup: a pool with no live replica yet reports ``stage_type is None``, so a stage whose first replica registers at runtime would otherwise never get a poller and its outputs would never drain. """ live: set[tuple[int, int]] = set() live_pools: set[int] = set() for stage_id in range(self.num_stages): pool = self.stage_pools[stage_id] stage_type = pool.stage_type if stage_type is None: # No live replica yet; nothing to attach to. A later # reconcile picks the stage up once one registers. continue if stage_type == "diffusion": live_pools.add(stage_id) existing_poller = pollers.get(stage_id) if existing_poller is None or existing_poller.done(): pollers[stage_id] = asyncio.create_task( _diffusion_poller(stage_id, pool), name=f"orch-diffusion-poller-s{stage_id}", ) continue for replica_id in pool.available_replica_ids(): client = pool.clients[replica_id] if client is None: continue key = (stage_id, replica_id) live.add(key) existing = readers.get(key) if existing is None: readers[key] = ( asyncio.create_task( _llm_replica_reader(stage_id, replica_id, client), name=f"orch-reader-s{stage_id}r{replica_id}", ), client, ) continue if existing[0].done(): # The reader exited, which means it already queued its # failure. Leave the slot alone until the dispatcher # handles that error and evicts the replica; respawning # now just re-raises the same error against the same # dead client and queues a duplicate eviction. continue if existing[1] is client: continue # Client swapped underneath us: retire the old reader. existing[0].cancel() reaped.append(existing[0]) readers[key] = ( asyncio.create_task( _llm_replica_reader(stage_id, replica_id, client), name=f"orch-reader-s{stage_id}r{replica_id}", ), client, ) for key in list(readers): if key not in live: task, _ = readers.pop(key) task.cancel() reaped.append(task) for stage_id in list(pollers): if stage_id not in live_pools: task = pollers.pop(stage_id) task.cancel() reaped.append(task) _reconcile_readers() logger.info( "[Orchestrator] Event-driven orchestration loop enabled (%s): %d LLM replica readers, %d diffusion pollers", _EVENT_DRIVEN_ORCH_ENV, len(readers), len(pollers), ) shutdown_task = asyncio.create_task(self._shutdown_event.wait(), name="orch-shutdown-wait") pending_get: asyncio.Task | None = None next_reconcile = _time.monotonic() + _ORCH_READER_RECONCILE_INTERVAL_S try: while not self._shutdown_event.is_set(): if pending_get is None: pending_get = asyncio.create_task(ready_q.get(), name="orch-dispatch-get") done, _ = await asyncio.wait( {pending_get, shutdown_task}, timeout=max(0.0, next_reconcile - _time.monotonic()), return_when=asyncio.FIRST_COMPLETED, ) if shutdown_task in done: return # Reconcile on a wall-clock schedule, not only when the queue # goes idle: under continuous traffic `asyncio.wait` returns on # every output and a timeout-only reconcile would never fire, # so a replica that registers at runtime would never get a # reader and its outputs would never drain. if _time.monotonic() >= next_reconcile: next_reconcile = _time.monotonic() + _ORCH_READER_RECONCILE_INTERVAL_S _reconcile_readers() if not done: self._orch_monitor.note_loop(idle=True) continue kind, stage_id, replica_id, payload = pending_get.result() pending_get = None if kind == "error": # replica_id < 0 means the failure was raised by a poller # task itself rather than attributed to one replica, so # there is nothing to evict -- and evict_replica() rejects # an out-of-range id, which would turn a survivable death # into a ValueError out of the dispatcher. if isinstance(payload, EngineDeadError) and replica_id >= 0: # Same per-replica isolation as the legacy loop (#4285): # evict and keep serving. Reconcile immediately so the # dead replica's reader is reaped now rather than at the # next scheduled tick. Skip a replica already evicted by # an earlier error from the same slot -- evict_replica() # would shut a None client down a second time. if replica_id in self.stage_pools[stage_id].live_replica_ids(): await self._handle_dead_replica(stage_id, replica_id, payload) _reconcile_readers() continue if self._shutdown_event.is_set(): return logger.error( "[Orchestrator] Stage-%s replica-%s reader failed: %r", stage_id, replica_id, payload, ) raise payload raw_terminal_request_ids: set[str] = set() try: if kind == "diffusion": pool = self.stage_pools[stage_id] if self._absorb_diffusion_metrics(stage_id, replica_id, payload): self._orch_monitor.note_loop(idle=False) continue if not is_diffusion_request_started_output(payload): pool.record_output_timestamps([payload]) processed = [payload] else: processed = await self._process_llm_stage_outputs( stage_id, replica_id, payload, raw_terminal_request_ids ) except asyncio.CancelledError: raise except EngineDeadError as e: if replica_id in self.stage_pools[stage_id].live_replica_ids(): await self._handle_dead_replica(stage_id, replica_id, e) _reconcile_readers() continue except Exception: if self._shutdown_event.is_set(): return logger.exception( "[Orchestrator] Stage-%s replica-%s processing failed", stage_id, replica_id, ) raise await self._handle_processed_outputs(stage_id, replica_id, processed) await self._finish_raw_terminal_requests(stage_id, replica_id, raw_terminal_request_ids) # Mirrors the legacy loop, which sets idle=False only after a # poll produced outputs that were routed -- an evicted replica # `continue`s above without marking the tick active. self._orch_monitor.note_loop(idle=False) finally: shutdown_task.cancel() if pending_get is not None: pending_get.cancel() reader_tasks = [task for task, _ in readers.values()] poller_tasks = list(pollers.values()) for task in (*reader_tasks, *poller_tasks): task.cancel() # `reaped` holds readers/pollers retired mid-run (client swap, # eviction, stage teardown). They were cancelled but never awaited, # so gather them here too rather than leaving dangling tasks. cleanup = [shutdown_task, *reader_tasks, *poller_tasks, *reaped] if pending_get is not None: cleanup.append(pending_get) await asyncio.gather(*cleanup, return_exceptions=True) async def _handle_processed_outputs(self, stage_id: int, replica_id: int, outputs: list[Any]) -> None: """Route processed stage outputs produced by one stage poll.""" pool = self.stage_pools[stage_id] for output in outputs: req_state = self.request_states.get(output.request_id) if req_state is None: logger.warning( "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)", output.request_id, stage_id, list(self.request_states.keys()), ) continue if getattr(output, "error", None) is not None: await self._handle_stage_error(stage_id, output) continue stage_metrics = None segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished if output.finished or segment_finished: stage_metrics = pool.build_stage_metrics( [output], submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()), request_timestamp=req_state.request_timestamp, replica_id=replica_id, sampling_params=req_state.sampling_params_list[stage_id], ) stage_metrics.pipeline_timings = dict(req_state.pipeline_timings) await self._route_output(stage_id, replica_id, output, req_state, stage_metrics) async def _handle_stage_error(self, stage_id: int, output: Any) -> None: """Emit a frontend-visible error and clean up request state.""" if self._cfg_tracker.is_companion(output.request_id): parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id else: parent_id = output.request_id await self.output_async_queue.put( ErrorMessage( request_id=parent_id, stage_id=stage_id, error=output.error, status_code=getattr(output, "error_status_code", None), error_type=getattr(output, "error_type", None), ) ) await self._cleanup_request_ids( [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], abort=True, release_owners=True, ) def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None: """Update one replica snapshot and expose the stage-wide total.""" self._stage_replica_waiting[(stage_id, replica_id)] = max(int(n_waiting), 0) self._set_stage_waiting_total(stage_id) self._sync_engines_waiting_counter() def _remove_stage_replica_waiting(self, stage_id: int, replica_id: int) -> None: """Drop a dead replica's snapshot and refresh the stage total.""" self._stage_replica_waiting.pop((stage_id, replica_id), None) self._set_stage_waiting_total(stage_id) self._sync_engines_waiting_counter() def _sync_engines_waiting_counter(self) -> None: """Aggregate per-replica waiting into the counter shared with the frontend so engine-queued requests show as waiting, not running.""" counter = self._engines_waiting_counter if counter is not None: counter.value = sum(self._stage_replica_waiting.values()) def _set_stage_waiting_total(self, stage_id: int) -> None: if self._prom_metrics is None: return total = sum( n_waiting for (snapshot_stage_id, _), n_waiting in self._stage_replica_waiting.items() if snapshot_stage_id == stage_id ) self._prom_metrics.set_stage_waiting_requests(stage_id, total) async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None: """Evict a dead stage replica and fail the requests stranded on it (#4285). A dead replica must not tear down the whole server: evict it and keep the orchestrator loop alive so remaining live replicas / stages keep serving. If the stage still has a live replica, fail only the requests bound to the dead replica; requests bound to a surviving replica (or not yet bound) keep running and can be (re)dispatched. If the stage has no live replica left, every request routed through it must fail since none can be served. """ pool = self.stage_pools[stage_id] logger.error( "[Orchestrator] Stage-%s replica-%s is dead; evicting it and failing its in-flight requests: %s", stage_id, replica_id, error, ) pool.evict_replica(replica_id) self._remove_stage_replica_waiting(stage_id, replica_id) stage_has_live = bool(pool.live_replica_ids()) failed_ids: list[str] = [] for req_id, req_state in list(self.request_states.items()): if stage_id not in req_state.stage_submit_ts: continue bound = pool.get_bound_replica_id(req_id) if bound == replica_id or not stage_has_live: await self.output_async_queue.put( ErrorMessage( error=str(error), fatal=True, request_id=req_id, stage_id=stage_id, ) ) failed_ids.append(req_id) # Use the shared cleanup path so the running counter, PD/CFG state, and # bindings across every stage pool are released (not just this pool). # abort=True stops the failed requests' in-flight work on surviving # stages (abort skips dead/evicted bindings), matching the # distributed-membership eviction path. for req_id in failed_ids: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, # A session-owned request has an owner that outlives the stage # request: without this the runner and its admission slot # survive the replica that was serving them, and the client is # told nothing -- the request-scoped ErrorMessage has no # frontend request_states entry to land in. No-op for # turn-based requests, which have no owner. release_owners=True, ) async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: """Fail one request whose target stage has no live replica (#4285). Used at dispatch sites so a fully dead stage fails just that request instead of raising out of the request/orchestration task and tearing down the server. The cleanup releases the running counter, PD/CFG state, and bindings across every stage pool; it is a no-op for a request that was never registered. """ logger.error( "[Orchestrator] req=%s: stage-%d has no live replica; failing the request", req_id, stage_id, ) await self.output_async_queue.put( ErrorMessage( error=f"Stage-{stage_id} has no live replica", fatal=True, request_id=req_id, stage_id=stage_id, ) ) await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, release_owners=True, ) async def _fail_request_client_error( self, req_id: str, stage_id: int, error: str, *, status_code: int = HTTPStatus.BAD_REQUEST.value, error_type: str = DEFAULT_CLIENT_ERROR_TYPE, release_owners: bool = False, ) -> None: """Fail one request with a non-fatal client error (400 by default). The non-fatal counterpart of `_fail_request_dead_stage`: emits a client-error ErrorMessage (default `fatal=False`) so the engine keeps serving, then releases the request's state across every stage pool. """ await self.output_async_queue.put( ErrorMessage( error=error, status_code=status_code, error_type=error_type, request_id=req_id, stage_id=stage_id, ) ) await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, release_owners=release_owners, ) async def _dispatch_or_fail_request( self, dispatch: Callable[[], Awaitable[Any]], *, req_id: str, stage_id: int, operation: str, dispatch_req_id: str | None = None, ) -> bool: """Run a dispatch action, converting stage unavailability into a single-request failure. ``dispatch`` is a thunk rather than a bare coroutine so that replica selection and other argument setup run inside this guard: a synchronous StageUnavailableError raised while building the dispatch is caught here, not left to propagate out of the caller. Replica eviction can interleave with any dispatch, surfacing as StageUnavailableError (no live replica / evicted slot) or EngineDeadError (replica died mid-submit before the poll loop evicted it). Both mean "this request cannot be placed", not "the server is broken": fail the request and keep serving (#4285). An OverflowError from encoding an out-of-range request value is likewise failed as a client error (400) rather than killing the loop. Other unrelated errors propagate. Returns False when the request was failed. ``req_id`` is the request the failure is attributed to; ``dispatch_req_id`` is the id actually being dispatched when it differs (e.g. a CFG companion submitted under its parent's attribution) and is used to locate the replica that died. """ try: await dispatch() return True except StageUnavailableError as e: # No specific replica to evict: the stage already has no live # replica or the chosen slot was evicted. Fail just this request. logger.error( "[Orchestrator] %s dispatch for req=%s hit unavailable stage-%s: %s", operation, req_id, stage_id, e, ) await self._fail_request_dead_stage(req_id, stage_id) return False except EngineDeadError as e: # A live replica died mid-submit before the poll loop evicted it. # Evict it now (and fail every request bound to it) so a burst of # new requests stops round-robining onto the dead slot until the # poll catches up. Capture the replica before failing this request, # since the failure releases its binding. pool = self.stage_pools[stage_id] dead_replica = pool.get_bound_replica_id(dispatch_req_id or req_id) logger.error( "[Orchestrator] %s dispatch for req=%s hit dead stage-%s replica-%s: %s", operation, req_id, stage_id, dead_replica, e, ) await self._fail_request_dead_stage(req_id, stage_id) if dead_replica is not None and dead_replica in pool.live_replica_ids(): await self._handle_dead_replica(stage_id, dead_replica, e) return False except OverflowError as e: # Catch overflow errors in the request payload to ensure that msgpack # encoding can't kill the engine; instead, we just fail this request. # The frontend should generally validate this as well, but this is needed # in case we get values like seed > 2**64-1 from sampling params. logger.warning( "[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s", operation, req_id, stage_id, e, ) await self._fail_request_client_error( req_id, stage_id, f"Request contains a value that cannot be serialized: {e}", ) return False # ---- Shared helpers ---- async def _cleanup_request_ids( self, request_ids: list[str], *, abort: bool = False, release_owners: bool = False, ) -> list[OutputMessage]: """Release pool bindings and logical request state for the given ids. CFG-aware: cleaning a parent releases its tracker state and pulls its companions into the batch; cleaning a companion whose parent is NOT in the batch and whose output never arrived fails that parent (its bundle can never complete). Every teardown path (stage error, abort, replica loss, membership unregister) funnels through here, so tracker state cannot outlive its requests. ``release_owners`` asks the subclass to also release any session that owns these requests (error paths pass True); the base ignores it. Returns: Final-stage AR abort ``OutputMessage`` list when ``abort=True``; otherwise an empty list. """ del release_owners if not request_ids: return [] cleanup_ids = list(dict.fromkeys(request_ids)) batch = set(cleanup_ids) orphaned_parents: dict[str, str] = {} for rid in cleanup_ids: pid = self._cfg_tracker.get_parent_id(rid) if pid is not None and pid not in batch and not self._cfg_tracker.is_companion_done(rid): orphaned_parents.setdefault(pid, rid) for pid, cid in orphaned_parents.items(): deferred = self._cfg_tracker.pop_pending_parent(pid) await self.output_async_queue.put( ErrorMessage( request_id=pid, stage_id=deferred["stage_id"] if deferred is not None else None, error=f"CFG companion {cid} was aborted or lost before its outputs arrived", ) ) batch.add(pid) cleanup_ids.append(pid) for rid in list(cleanup_ids): for cid in self._cfg_tracker.cleanup_parent(rid): if cid not in batch: batch.add(cid) cleanup_ids.append(cid) abort_outputs: list[OutputMessage] = [] if abort: abort_outputs = await self._abort_request_ids(cleanup_ids) self._release_request_bindings(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) if req_state is not None: cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: self._running_counter.decrement() req_state.running_counter_registered = False return abort_outputs async def _apply_raw_terminal_stage_finish( self, stage_id: int, eco: Any, req_state: OrchestratorRequestState, ) -> bool: """Record a session-level finish marker and report whether it was terminal. Streaming segment stops set ``is_segment_finished=True`` and are handled via processed outputs. Session termination (e.g. ``finish_requests`` after ``resumable=False``) emits a terminal ``finish_reason`` with ``is_segment_finished=False``, but vLLM's output processor may remove the request state before that EngineCoreOutput is processed. Only update ``finished_final_output_stage_ids`` here. Request cleanup stays in ``_route_output`` so downstream async-chunk stages can still deliver outputs after stage-0 session end. """ if getattr(eco, "finish_reason", None) is None: return False if getattr(eco, "is_segment_finished", False): return False final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if stage_id not in final_output_stage_ids: return False req_state.finished_final_output_stage_ids.add(stage_id) return True async def _report_duplex_session_request_error( self, stage_id: int, replica_id: int | None, eco: Any, req_state: OrchestratorRequestState, ) -> None: """Turn a scheduler-side error finish of a duplex session request into a request error. A duplex session request is resumable, so its terminal outputs are deliberately not turned into processed outputs (``_finish_raw_terminal_requests`` skips duplex sessions), and vLLM's output processor may already have dropped the request state by the time the raw ``FinishReason.ERROR`` output arrives. Without this the session never learns that its request is gone: the next append re-creates the request from scratch and the session silently loses its context. Deliver the error to the request's consumer instead so the serving layer can fail and close the session. """ if not req_state.session_owned: return if getattr(eco, "finish_reason", None) != FinishReason.ERROR: return if getattr(eco, "is_segment_finished", False): return stop_reason = getattr(eco, "stop_reason", None) error = stop_reason if isinstance(stop_reason, str) and stop_reason else "duplex session request failed" logger.error( "[Orchestrator] Duplex session request %s failed at stage-%s replica-%s: %s", req_state.request_id, stage_id, replica_id, error, ) await self.output_async_queue.put( ErrorMessage( error=error, fatal=False, request_id=req_state.request_id, stage_id=stage_id, ) ) async def _finish_raw_terminal_requests( self, stage_id: int, replica_id: int, request_ids: set[str], ) -> None: """Finish streaming requests whose raw terminal had no processed output.""" pool = self.stage_pools[stage_id] if not pool.final_output: return for request_id in request_ids: req_state = self.request_states.get(request_id) if req_state is None or req_state.session_owned: continue final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if not final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids): continue final_output_type = getattr(pool.stage_client, "final_output_type", None) logger.info( "[Orchestrator] req=%s stage-%s raw terminal produced no processed terminal; returning empty %s output", request_id, stage_id, final_output_type or "text", ) terminal_output = _build_terminal_empty_output( request_id, final_output_type=final_output_type, audio_sample_rate=pool._infer_audio_sample_rate(), ) await self.output_async_queue.put( OutputMessage( request_id=request_id, stage_id=stage_id, replica_id=replica_id, engine_outputs=terminal_output, metrics=None, finished=True, stage_submit_ts=req_state.stage_submit_ts.get(stage_id), ) ) await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)]) def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" companion_request_ids = self._cfg_tracker.get_companion_request_ids(request_id) if not companion_request_ids: return params import copy from vllm_omni.inputs.data import OmniDiffusionSamplingParams if not isinstance(params, OmniDiffusionSamplingParams): return params params = copy.deepcopy(params) params.cfg_kv_request_ids = companion_request_ids return params def _register_running_request(self, req_state: OrchestratorRequestState) -> None: if req_state.running_counter_registered or self._running_counter is None: return self._running_counter.increment() req_state.running_counter_registered = True async def _route_output( self, stage_id: int, replica_id: int, output: Any, req_state: OrchestratorRequestState, stage_metrics: Any, ) -> None: """Route a processed output: send to frontend and/or forward.""" req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished # CFG companion: stash output so parent can bundle [parent, *companions] # into source_outputs for the bridge (e.g. thinker2imagegen). if finished and self._cfg_tracker.is_companion(req_id): self._cfg_tracker.set_companion_output(req_id, output) await self._handle_cfg_companion_ready(req_id) await self._cleanup_request_ids([req_id]) return request_finished = False if finished and self.stage_pools[stage_id].final_output and not segment_finished: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) if req_state.session_owned: # Session-owned outputs are delivered to their session through # ``_intercept_stage_output`` below, never to a request queue. pass elif self.stage_pools[stage_id].final_output: await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=stage_id, replica_id=replica_id, engine_outputs=output, metrics=stage_metrics, finished=request_finished, stage_submit_ts=submit_ts, ) ) elif stage_metrics is not None: await self.output_async_queue.put( StageMetricsMessage( request_id=req_id, stage_id=stage_id, replica_id=replica_id, metrics=stage_metrics, stage_submit_ts=submit_ts, ) ) if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]: kv_params = getattr(output, "kv_transfer_params", None) if kv_params is not None: self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params) req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None) if await self._intercept_stage_output(stage_id, replica_id, output, req_state, stage_metrics, submit_ts): return if ( (finished or segment_finished) and stage_id < req_state.final_stage_id and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) ): if ( finished and self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id) ): self._cfg_tracker.defer_parent(req_id, output, stage_id) else: stage_params = req_state.sampling_params_list[stage_id] final_only_finished = ( req_state.streaming.enabled and finished and getattr(stage_params, "output_kind", None) == RequestOutputKind.FINAL_ONLY ) await self._forward_to_next_stage( req_id, stage_id, output, req_state, src_replica_id=replica_id, is_streaming_session=req_state.streaming.enabled, is_final_update=final_only_finished, ) if ( req_state.streaming.enabled # A failed forward may already have cleaned up this request. and self.request_states.get(req_id) is req_state and finished and not final_only_finished and not req_state.session_owned ): # For streaming sessions, send the terminal (resumable=False) update only on a finish await self._forward_to_next_stage( req_id, stage_id, output, req_state, src_replica_id=replica_id, is_streaming_session=True, is_final_update=True, ) if request_finished and not req_state.session_owned: await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts def _stage_receives_async_chunks(self, stage_id: int) -> bool: """Whether a stage's connector supplies its runtime inputs.""" pool = self.stage_pools[stage_id] model_config = getattr(pool.stage_vllm_config, "model_config", None) return stage_receives_chunks(model_config) def _get_stage_input_processor(self, stage_id: int) -> Any: processor = self._stage_input_processors.get(stage_id) if processor is None: from vllm_omni.engine.stage_init_utils import build_stage0_input_processor processor = build_stage0_input_processor(self.stage_pools[stage_id].stage_vllm_config) self._stage_input_processors[stage_id] = processor return processor def _upgrade_processed_stage_request(self, request: Any, raw_prompt: Any) -> Any: prompt_embeds = getattr(request, "prompt_embeds", None) 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="Orchestrator stage input", ) 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 _next_stage_input_is_tokens(self, next_input: Any) -> bool: return isinstance(next_input, dict) and "prompt_token_ids" in next_input def _build_next_stage_request( self, req_id: str, next_stage_id: int, next_input: Any, params: SamplingParams | PoolingParams, *, mm_features: list | None = None, resumable: bool = False, payload_sender_info: dict[str, Any] | None = None, ) -> Any: next_pool = self.stage_pools[next_stage_id] if self._next_stage_input_is_tokens(next_input): request = build_engine_core_request_from_tokens( request_id=req_id, prompt=next_input, params=params, model_config=next_pool.stage_vllm_config.model_config, mm_features=mm_features, resumable=resumable, ) request.external_req_id = request.request_id request.payload_sender_info = payload_sender_info return request processor = self._get_stage_input_processor(next_stage_id) # A pooling stage is driven by PoolingParams; vLLM validates them against # the stage's supported tasks, so advertise the model's pooling tasks (or # the configured task) instead of generate. if isinstance(params, PoolingParams): model_cfg = next_pool.stage_vllm_config.model_config supported_tasks = tuple(getattr(model_cfg, "supported_tasks", ()) or ()) or (params.task,) else: supported_tasks = ("generate",) request = processor.process_inputs( request_id=req_id, prompt=next_input, params=params, supported_tasks=supported_tasks, arrival_time=_time.time(), resumable=resumable, ) request = self._upgrade_processed_stage_request(request, next_input) request.external_req_id = req_id request.payload_sender_info = payload_sender_info return request @staticmethod def _completion_multimodal_output(output: Any, completion: Any) -> dict[str, Any]: mm_output = getattr(output, "multimodal_output", None) if isinstance(mm_output, dict): return mm_output mm_output = getattr(completion, "multimodal_output", None) if completion is not None else None return mm_output if isinstance(mm_output, dict) else {} @classmethod def _coerce_int_list(cls, value: Any) -> list[int]: if value is None: return [] if hasattr(value, "detach"): try: value = value.detach().cpu().reshape(-1).tolist() except Exception: return [] if not isinstance(value, (list, tuple)): return [] out: list[int] = [] for item in value: token_id = cls._coerce_int(item) if token_id is not None: out.append(token_id) return out @staticmethod def _coerce_int(value: Any) -> int | None: if hasattr(value, "detach"): try: value = value.detach().cpu().reshape(-1) if value.numel() == 0: return None value = value[0].item() except Exception: return None try: return int(value) except (TypeError, ValueError): return None async def _handle_cfg_companion_ready(self, req_id: str) -> None: """Mark a CFG companion as done; if all companions are done, flush deferred parent.""" parent_id = self._cfg_tracker.on_companion_completed(req_id) if parent_id is None: return deferred = self._cfg_tracker.pop_pending_parent(parent_id) if deferred is None: return parent_state = self.request_states.get(parent_id) if parent_state is None: return stage_id = deferred["stage_id"] if (stage_id + 1) in parent_state.stage_submit_ts: return await self._forward_to_next_stage( parent_id, stage_id, deferred["engine_outputs"], parent_state, ) async def _handle_kv_ready_raw_outputs( self, stage_id: int, raw_outputs: EngineCoreOutputs, ) -> None: """Forward split requests once stage-0 KV is ready.""" if self.async_chunk: return for raw_output in raw_outputs.outputs: kv_params = getattr(raw_output, "kv_transfer_params", None) if not (isinstance(kv_params, dict) and kv_params.get("kv_ready")): continue req_id = raw_output.request_id req_state = self.request_states.get(req_id) if req_state is None: continue if self._cfg_tracker.is_companion(req_id): # kv_ready only says the companion's KV hit the connector; its # processed output has not been stashed yet. Counting it as # done here let the parent pass the all_companions_done gate # and dispatch with 0/N companion outputs (degraded CFG). # Done-marking now happens solely on the processed-output path # in _route_output, after set_companion_output. continue if stage_id >= req_state.final_stage_id: continue if (stage_id + 1) in req_state.stage_submit_ts: continue if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): self._cfg_tracker.defer_parent(req_id, raw_output, stage_id) else: await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state) def _build_pd_decode_params(self, req_id: str, sp: Any) -> Any: """Build decode-side sampling params with KV transfer params for PD routing. Clones the sampling params and injects kv_transfer_params that tell the decode engine where to pull the KV cache from (prefill engine's bootstrap addr). """ sp = sp.clone() if sp.extra_args is None: sp.extra_args = {} # Get KV params captured from the prefill output (must include remote_request_id). kv_prefill_params = self._pd_kv_params.pop(req_id, None) if not kv_prefill_params or "remote_request_id" not in kv_prefill_params: raise RuntimeError( f"[Orchestrator][PD] Missing prefill kv_transfer_params.remote_request_id for req={req_id}" ) decode_kv_params: dict[str, Any] = { "transfer_id": f"xfer-{req_id}", } if self._pd_bootstrap_addr: decode_kv_params["remote_bootstrap_addr"] = self._pd_bootstrap_addr if self._pd_prefill_engine_id: decode_kv_params["remote_engine_id"] = self._pd_prefill_engine_id # Overlay params from prefill side (includes remote_request_id set by monkey patch). decode_kv_params.update(kv_prefill_params) # Ensure these flags are set correctly after any overlay. decode_kv_params["do_remote_prefill"] = True decode_kv_params["do_remote_decode"] = False if not decode_kv_params.get("transfer_id"): decode_kv_params["transfer_id"] = f"xfer-{req_id}" sp.extra_args["kv_transfer_params"] = decode_kv_params logger.debug( "[Orchestrator][PD] decode kv_transfer_params for req=%s: %s", req_id, decode_kv_params, ) return sp def _emit_tx_edge( self, *, from_stage: int, from_replica: int, to_stage: int, to_pool: StagePool, request_id: str, tx_ms: float, ) -> None: """Emit per-edge transfer_tx_s + transfer_size_bytes histograms. ``tx_ms`` is the orchestrator-side wall-clock spent in ``next_pool. submit_*`` (serialize + queue submit to the receiving worker). Best- effort size_bytes left at 0 — orchestrator doesn't have a cheap handle on the serialized payload size; a follow-up can plumb that from the connector adapter. """ if self._transfer_emitter is None: return to_replica = to_pool.get_bound_replica_id(request_id) if to_replica is None: return try: self._transfer_emitter.observe_size(from_stage, from_replica, to_stage, to_replica, 0) self._transfer_emitter.observe_tx_time(from_stage, from_replica, to_stage, to_replica, tx_ms / 1000.0) except Exception: logger.debug( "[Orchestrator] transfer_tx emit failed for edge %d->%d req=%s", from_stage, to_stage, request_id, exc_info=True, ) async def _forward_to_next_stage( self, req_id: str, src_stage_id: int, output: Any, req_state: OrchestratorRequestState, *, src_replica_id: int | None = None, is_streaming_session: bool = False, is_final_update: bool = False, ) -> None: """Forward output to the next stage; a dead stage fails only this request.""" await self._dispatch_or_fail_request( lambda: self._forward_to_next_stage_unguarded( req_id, src_stage_id, output, req_state, src_replica_id=src_replica_id, is_streaming_session=is_streaming_session, is_final_update=is_final_update, ), req_id=req_id, stage_id=src_stage_id + 1, operation="inter-stage forward", ) async def _forward_to_next_stage_unguarded( self, req_id: str, src_stage_id: int, output: Any, req_state: OrchestratorRequestState, *, src_replica_id: int | None = None, is_streaming_session: bool = False, is_final_update: bool = False, ) -> None: """Forward output from the current logical stage to the next one.""" next_logical = src_stage_id + 1 next_pool = self.stage_pools[next_logical] if not next_pool.live_replica_ids(): # The downstream stage lost all of its replicas before this in-flight # request could be forwarded. await self._fail_request_dead_stage(req_id, next_logical) return next_client = next_pool.stage_client params = req_state.sampling_params_list[next_logical] source_outputs = [output] next_stage_resumable = is_streaming_session and not is_final_update already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": # Gate: never dispatch with an incomplete CFG bundle. Checked # non-destructively BEFORE popping — a pop-then-redefer would lose # the partial outputs. all_companions_done is trustworthy here # because done-marking happens only after set_companion_output. if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): self._cfg_tracker.defer_parent(req_id, output, src_stage_id) logger.info( "[Orchestrator] req=%s: CFG companion outputs not all stashed yet; re-deferring parent", req_id, ) return # Peek, don't pop: outputs stay stashed until cleanup_parent so a # streaming re-submission bundles the complete set again rather # than an empty one. companion_outputs = self._cfg_tracker.get_companion_outputs(req_id) expected = len(self._cfg_tracker.get_companion_request_ids(req_id)) if expected > len(companion_outputs): # Companions are done but outputs are missing — inconsistent # tracker state (should be unreachable with peek semantics). # Fail the request rather than degrade CFG conditioning. logger.error( "[Orchestrator] req=%s: only %d/%d CFG companion outputs available; " "failing the request instead of dispatching degraded CFG", req_id, len(companion_outputs), expected, ) await self.output_async_queue.put( ErrorMessage( request_id=req_id, stage_id=src_stage_id, error=( f"CFG companion outputs incomplete ({len(companion_outputs)}/{expected}); " "request aborted to avoid degraded CFG conditioning" ), ) ) await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, ) return diffusion_source_outputs = [output, *companion_outputs] if next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() _fn = next_client.custom_process_input_func _extra_kwargs: dict[str, Any] = {} # TODO: replace signature probe with explicit kwarg contract. try: import inspect as _inspect if "sampling_params" in _inspect.signature(_fn).parameters: _extra_kwargs["sampling_params"] = params except (TypeError, ValueError): pass diffusion_prompt = _fn( diffusion_source_outputs, req_state.prompt, requires_multimodal_data, **_extra_kwargs, ) _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d logger.info( "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d", req_id, _dt_ar2d, src_stage_id, next_logical, ) if diffusion_prompt is None: error_output = OmniRequestOutput.from_error( req_id, f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", ) logger.warning( "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " "routing terminal error output", req_id, src_stage_id, next_logical, ) await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=next_logical, engine_outputs=error_output, metrics=None, finished=True, ) ) await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], ) return if isinstance(diffusion_prompt, list): if not diffusion_prompt: error_output = OmniRequestOutput.from_error( req_id, f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", ) logger.warning( "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " "routing terminal error output", req_id, src_stage_id, next_logical, ) await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=next_logical, engine_outputs=error_output, metrics=None, finished=True, ) ) await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], ) return if len(diffusion_prompt) == 1: diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt if already_submitted: replica_id = await next_pool.submit_update(req_id, req_state, diffusion_prompt) else: replica_id = await next_pool.submit_initial( req_id, req_state, diffusion_prompt, submit_kwargs={ "kv_sender_info": self._build_kv_sender_info( list(getattr(next_client, "engine_input_source", None) or [src_stage_id]), request_id=req_id, ) }, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) self._on_stage_submitted( next_logical, req_id, replica_id, req_state, ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, to_stage=next_logical, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, ) return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair: params = self._build_pd_decode_params(req_id, params) # Use the original user prompt for the decode stage (not processed embeddings) original_prompt = req_state.prompt raw_decode_inputs = [original_prompt] if not isinstance(original_prompt, list) else original_prompt decode_inputs: list[dict[str, Any]] = [] for decode_input in raw_decode_inputs: if isinstance(decode_input, dict): decode_inputs.append(decode_input) continue prompt_token_ids = getattr(decode_input, "prompt_token_ids", None) if prompt_token_ids is None: raise TypeError( "[Orchestrator][PD] decode input must be dict or have prompt_token_ids, " f"got {type(decode_input).__name__} for req={req_id}" ) decode_inputs.append({"prompt_token_ids": list(prompt_token_ids)}) for decode_input in decode_inputs: request = build_engine_core_request_from_tokens( request_id=req_id, prompt=decode_input, params=params, model_config=next_pool.stage_vllm_config.model_config, mm_features=req_state.mm_features, resumable=next_stage_resumable, ) request.external_req_id = request.request_id if already_submitted: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) self._on_stage_submitted( next_logical, req_id, replica_id, req_state, ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, to_stage=next_logical, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, ) return if req_state.pd_prefill_multimodal_output is not None: req_state.streaming.bridge_states.setdefault( "pd_prefill_multimodal_output_by_req", {}, )[req_id] = req_state.pd_prefill_multimodal_output previous_decoder = req_state.streaming.source_token_decoder source_processor = self.stage_pools[src_stage_id].output_processor tokenizer = getattr(source_processor, "tokenizer", None) decode = getattr(tokenizer, "decode", None) if callable(decode): req_state.streaming.source_token_decoder = decode try: next_inputs = next_client.process_engine_inputs( source_outputs, req_state.prompt, streaming_context=req_state.streaming, ) except OmniClientError as exc: await self._fail_request_client_error( req_id, next_logical, str(exc), status_code=exc.status_code, error_type=exc.error_type, release_owners=req_state.session_owned, ) return except Exception as exc: logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, next_logical, ) if not await self._handle_forward_failure(req_id, next_logical, req_state, exc): raise return finally: req_state.streaming.source_token_decoder = previous_decoder if not next_inputs: if not getattr(output, "finished", False): logger.debug( "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs", req_id, src_stage_id, next_logical, ) return final_stage_id = req_state.final_stage_id final_pool = self.stage_pools[final_stage_id] final_output_type = getattr(final_pool.stage_client, "final_output_type", None) terminal_output = _build_terminal_empty_output( req_id, final_output_type=final_output_type, audio_sample_rate=final_pool._infer_audio_sample_rate(), ) submit_ts = _time.time() req_state.stage_submit_ts[final_stage_id] = submit_ts logger.info( "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; " "returning empty %s output from final stage-%s", req_id, src_stage_id, next_logical, final_output_type or "text", final_stage_id, ) await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=final_stage_id, replica_id=0, engine_outputs=terminal_output, metrics=None, finished=True, stage_submit_ts=submit_ts, ) ) await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) return # Build and submit requests for each input for next_input in next_inputs: # Only AR thinker stages consume encoder mm_features; downstream # (talker/code2wav/…) must not see them (avoids encoder-cache misses). model_stage = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None) mm_features = req_state.mm_features if model_stage == "thinker" else None request = self._build_next_stage_request( req_id, next_logical, next_input, params=params, mm_features=mm_features, resumable=next_stage_resumable, payload_sender_info=self._build_payload_sender_info(src_stage_id, request_id=req_id), ) if already_submitted: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) self._on_stage_submitted( next_logical, req_id, replica_id, req_state, ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, to_stage=next_logical, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, ) async def _prewarm_async_chunk_stages( self, request_id: str, stage0_request: Any, req_state: OrchestratorRequestState, ) -> bool: """Pre-submit downstream stages for async-chunk mode. Returns False when the request was failed and cleaned up in here, so a caller still holding ``req_state`` stops instead of recording state on an object the cleanup already popped from ``request_states``. """ if req_state.final_stage_id <= 0: return True prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None) if prompt_token_ids is None: # R1.3 of #4855. Skipping the prewarm leaves every downstream stage # unsubmitted, so the request produces no output and no error -- it # just stops. An embeds-only prompt in async-chunk mode is a config # mistake, and the client should hear about it now rather than wait # for the stage-input deadline to notice nothing ever arrived. # # Non-fatal, with a 4xx: `fatal=True` means "the engine is dead", and # the entrypoints act on it -- `OmniBase._handle_output_message` # raises `OmniEngineDeadError` out of `OmniLLM.generate`'s batch loop, # failing every other request in that batch, and the serving layer # calls `terminate_if_errored`. One client's bad prompt must not do # either. This is the same shape `_handle_stage_error` uses for a # request-scoped client error. logger.error( "[Orchestrator] req=%s: async_chunk prewarm needs stage0 prompt_token_ids " "and none were provided; failing the request", request_id, ) await self._fail_request_client_error( request_id, 0, "async_chunk requires prompt_token_ids on the stage-0 request; " "an embeds-only prompt cannot prewarm downstream stages", release_owners=True, ) return False for next_stage_id in range(1, req_state.final_stage_id + 1): next_pool = self.stage_pools[next_stage_id] params = req_state.sampling_params_list[next_stage_id] if not self._stage_receives_async_chunks(next_stage_id): # Outgoing-only stages receive their first real input from the # orchestrator. Pre-submitting a placeholder lets it race and # execute before that conditioning payload arrives. continue req_state.stage_submit_ts[next_stage_id] = _time.time() _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": submit_kwargs = { "kv_sender_info": self._build_kv_sender_info( list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]), request_id=request_id, ) } submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( request_id, req_state, req_state.prompt, submit_kwargs=submit_kwargs, ), req_id=request_id, stage_id=next_stage_id, operation="async-chunk prewarm", ) else: import copy from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length try: next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) except Exception: next_prompt_len = max(1, len(prompt_token_ids)) # A stage may size its own prewarmed placeholder prompt via # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker # whose engine positions must cover a speaker-prompt prefill; # the model validates the value and reports the right one on # mismatch). prewarm_len = getattr( getattr(next_pool.stage_vllm_config, "model_config", None), "hf_config", None, ) prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None) if prewarm_len is not None and int(prewarm_len) > 1: next_prompt_len = int(prewarm_len) original_prompt = req_state.prompt if isinstance(original_prompt, dict): base_input = copy.deepcopy(original_prompt) else: base_input = {} base_input["prompt_token_ids"] = [0] * next_prompt_len base_input["multi_modal_data"] = None base_input["mm_processor_kwargs"] = None downstream_resumable = bool(getattr(stage0_request, "resumable", req_state.streaming.enabled)) request = build_engine_core_request_from_tokens( request_id=request_id, prompt=base_input, params=params, model_config=next_pool.stage_vllm_config.model_config, resumable=downstream_resumable, ) request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( request_id, req_state, request, prompt_text=None, ), req_id=request_id, stage_id=next_stage_id, operation="async-chunk prewarm", ) # Attribute the failure to the stage that failed, not stage 0; the # request is already cleaned up, so stop prewarming. if not submitted: return False # The guard swallows submit_initial's return value; the pool binding # it recorded carries the same replica. bound_replica_id = next_pool.get_bound_replica_id(request_id) self._on_stage_submitted( next_stage_id, request_id, bound_replica_id if bound_replica_id is not None else 0, req_state, ) # async_chunk pre-submit fires per stage edge (N-1 -> N). Source # replica is stage 0's bound replica (single-replica thinker in # all current configs); fall back to 0 if unknown. _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 src_replica = self.stage_pools[next_stage_id - 1].get_bound_replica_id(request_id) self._emit_tx_edge( from_stage=next_stage_id - 1, from_replica=src_replica if src_replica is not None else 0, to_stage=next_stage_id, to_pool=next_pool, request_id=request_id, tx_ms=_tx_ms, ) return True def _build_kv_sender_info( self, sender_stage_ids: list[int], *, request_id: str | None = None, ) -> dict[int, dict[str, Any]] | None: """Build per-request sender info for diffusion KV-transfer receivers.""" sender_infos: dict[int, dict[str, Any]] = {} for sender_stage_id in dict.fromkeys(sender_stage_ids): if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools): continue sender_pool = self.stage_pools[sender_stage_id] sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None if sender_stage is None: sender_stage = sender_pool.stage_client get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) if not callable(get_sender_info): continue sender_info = get_sender_info() if not sender_info: logger.warning( "[Orchestrator] Stage-%s has no KV sender info available", sender_stage_id, ) continue sender_infos[sender_stage_id] = sender_info return sender_infos or None def _build_payload_sender_info( self, sender_stage_id: int, *, request_id: str | None = None, ) -> dict[str, Any] | None: if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools): return None sender_pool = self.stage_pools[sender_stage_id] sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None if sender_stage is None: sender_stage = sender_pool.stage_client get_sender_info = getattr(sender_stage, "get_payload_sender_info", None) if not callable(get_sender_info): return None return get_sender_info() # ---- Shutdown / lifecycle ---- def _shutdown_stages(self) -> None: """Shutdown all stage pools.""" if self._stages_shutdown: return self._stages_shutdown = True total = sum(pool.live_num_replicas for pool in self.stage_pools) logger.info("[Orchestrator] Shutting down all %d client(s)", total) for pool in self.stage_pools: for replica_id in pool.live_replica_ids(): pool.shutdown_replica(replica_id) class Orchestrator(OrchestratorBase): """Turn-based orchestrator: admits ``add_request`` / streaming / companion / interaction messages.""" async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: msg_type = msg.type if msg_type == "add_request": await self._handle_add_request(msg) elif msg_type == "streaming_update": await self._handle_streaming_update(msg) elif msg_type == "add_companion_request": await self._handle_add_companion(msg) elif msg_type == "interaction": await self._handle_interaction(msg) else: return False return True async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: """Handle an add_request message from the main thread.""" stage_id = 0 request_id = msg.request_id prompt = msg.prompt original_prompt = msg.original_prompt sampling_params_list = msg.sampling_params_list if not sampling_params_list: raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") final_stage_id = msg.final_stage_id final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) if not self.stage_pools[stage_id].live_replica_ids(): # Stage 0 lost all replicas between the HTTP-layer errored check and # dispatch. Runs before request state / running counter registration, # so the helper's cleanup is a no-op here. await self._fail_request_dead_stage(request_id, stage_id) return logger.debug( "[Orchestrator] _handle_add_request: stage=%s req=%s " "prompt_type=%s original_prompt_type=%s final_stage=%s " "num_sampling_params=%d", stage_id, request_id, type(prompt).__name__, type(original_prompt).__name__, final_stage_id, len(sampling_params_list), ) req_state = OrchestratorRequestState( request_id=request_id, prompt=original_prompt, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, final_output_stage_ids=final_output_stage_ids, request_timestamp=float(msg.request_timestamp or _time.time()), mm_features=getattr(prompt, "mm_features", None), request_artifact_dirs=set(msg.request_artifact_dirs or ()), ) self.request_states[request_id] = req_state self._register_running_request(req_state) req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) req_state.stage_submit_ts[stage_id] = _time.time() enqueue_ts = msg.enqueue_ts if enqueue_ts > 0: req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 preprocess_ms = msg.preprocess_ms if preprocess_ms > 0: req_state.pipeline_timings["preprocess_ms"] = preprocess_ms if not await self._dispatch_or_fail_request( lambda: self.stage_pools[stage_id].submit_initial( request_id, req_state, prompt, prompt_text=msg.output_prompt_text, ), req_id=request_id, stage_id=stage_id, operation="add_request", ): return if self.async_chunk and stage_id == 0 and final_stage_id > 0: await self._prewarm_async_chunk_stages(request_id, prompt, req_state) async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: """Handle a streaming_update message for an existing request.""" stage_id = 0 request_id = msg.request_id request = msg.prompt final_stage_id = msg.final_stage_id req_state = self.request_states.get(request_id) if req_state is None: # Streaming updates always follow the first-chunk add_request # through the same ordered submission queue, so an unknown id # here can only mean the request already finished or was # aborted (e.g. the client disconnected). Re-adding it would # resurrect a headless session that keeps cycling through the # stages with nobody consuming its outputs (issue #4271). logger.warning( "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", request_id, ) return if msg.sampling_params_list: req_state.sampling_params_list = msg.sampling_params_list req_state.streaming.enabled = True req_state.stage_submit_ts[stage_id] = _time.time() if not await self._dispatch_or_fail_request( lambda: self.stage_pools[stage_id].submit_update( request_id, req_state, request, prompt_text=msg.output_prompt_text, ), req_id=request_id, stage_id=stage_id, operation="streaming_update", ): return if self.async_chunk and stage_id == 0 and final_stage_id > 0: await self._prewarm_async_chunk_stages(request_id, request, req_state) async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: """Handle an add_companion_request message: submit companion to stage 0.""" companion_id = msg.companion_id parent_id = msg.parent_id role = msg.role companion_prompt = msg.prompt sampling_params_list = msg.sampling_params_list parent_state = self.request_states.get(parent_id) if parent_state is None: logger.info( "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", companion_id, role, parent_id, ) return self._cfg_tracker.register_companion(parent_id, role, companion_id) companion_state = OrchestratorRequestState( request_id=companion_id, prompt=companion_prompt, sampling_params_list=sampling_params_list, final_stage_id=0, final_output_stage_ids={0}, request_timestamp=parent_state.request_timestamp, ) self.request_states[companion_id] = companion_state companion_state.stage_submit_ts[0] = _time.time() # A dead-stage failure is attributed to the parent request; its cleanup # also releases this companion. if not await self._dispatch_or_fail_request( lambda: self.stage_pools[0].submit_initial( companion_id, companion_state, companion_prompt, prompt_text=msg.companion_prompt_text, affinity_request_id=parent_id, ), req_id=parent_id, stage_id=0, operation=f"companion {companion_id}", dispatch_req_id=companion_id, ): return logger.info( "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", companion_id, role, parent_id, self.stage_pools[0].get_bound_replica_id(companion_id), ) async def _handle_interaction(self, msg: InteractionMessage) -> None: """Handle a midway interaction for an active streaming diffusion request.""" stage_id = 0 request_id = msg.request_id event_id = msg.interaction.get("event_id") req_state = self.request_states.get(request_id) if req_state is None: logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) await self.output_async_queue.put( ErrorMessage( error=f"No active request for interaction: {request_id}", fatal=False, request_id=request_id, event_id=event_id, stage_id=stage_id, ) ) return try: await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) except Exception as exc: logger.info( "[Orchestrator] Failed interaction for req %s: %s", request_id, exc, exc_info=True, ) await self.output_async_queue.put( ErrorMessage( error=f"Failed interaction for request {request_id}: {exc}", fatal=False, request_id=request_id, event_id=event_id, stage_id=stage_id, ) )