anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
main
vllm-omni/vllm_omni/engine/orchestrator.py
T

3172 lines
141 KiB
Python
Raw Permalink Blame History

# 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 TYPE_CHECKING, 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.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane
from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY
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.errors import NativeKVHandoffError
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
if TYPE_CHECKING:
from vllm_omni.engine import OmniEngineCoreOutput
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 is off except for pipelines with an explicit validated default.
_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(*, default: bool = False) -> bool:
value = os.environ.get(_EVENT_DRIVEN_ORCH_ENV)
if value is None:
return default
return value.strip().lower() in ("1", "true", "yes", "on")
def _event_driven_orch_default_for_pipeline(pipeline_model_type: str | None) -> bool:
"""Return whether a pipeline has a validated event-driven default."""
return pipeline_model_type == "qwen3_tts"
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)
finished_stage_ids: set[int] = field(default_factory=set)
pending_final_output: OutputMessage | None = None
# 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
# Duplex sentence TTS already submitted this output. Do not also run the
# legacy full-text handoff for the same stage result.
skip_legacy_stage_forward: bool = False
running_counter_registered: bool = False
request_artifact_dirs: set[str] = field(default_factory=set)
native_kv_transfer_id: str | None = None
# First-frame delivery and final-stage outputs use different channels.
upstream_first_audio: bool = False
pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None
pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list)
@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) if stage_id is not None else None
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
_transfer_release_tasks: set[asyncio.Task] = set()
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,
event_driven_orch_default: 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] = {}
# Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC.
self._transfer_release_tasks: set[asyncio.Task] = set()
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(default=event_driven_orch_default)
# 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:
async def cleanup_member_requests(ids: list[str]) -> None:
await self._cleanup_request_ids(ids, abort=True)
self._membership.install_unregister_handlers(
output_queue=self.output_async_queue,
cleanup_callback=cleanup_member_requests,
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 client is not None and 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_stage_transfer_resources(self, request_ids: list[str]) -> None:
"""Drop each stage's inter-stage transfer resources for finished requests.
This is the only point that knows every stage is done with the request,
so it is the only safe place to reclaim segments a consumer never
drained. Scheduled rather than awaited: reclaim is best-effort and must
not add RPC latency to request teardown.
"""
if not request_ids:
return
async def _run() -> None:
results = await asyncio.gather(
*(pool.release_request_resources(request_ids) for pool in self.stage_pools),
return_exceptions=True,
)
for result in results:
if isinstance(result, Exception):
logger.warning("[Orchestrator] release transfer resources failed: %s", result)
try:
task = asyncio.get_running_loop().create_task(_run())
self._transfer_release_tasks.add(task)
task.add_done_callback(self._transfer_release_tasks.discard)
except RuntimeError:
logger.warning(
"[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s",
request_ids,
)
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:
clear_sender = pool.stage_type != "diffusion" and (
method in ("reset_mm_cache", "sleep", "release_kv_cache_memory")
or (method == "pause_scheduler" and kwargs.get("clear_cache", args[1] if len(args) > 1 else True))
)
for replica_id in pool.live_replica_ids():
try:
if clear_sender:
processor = self._stage_input_processors.get(pool.stage_id)
if processor is not None:
clear = processor.renderer.clear_mm_cache_async()
if timeout is None:
await clear
else:
await asyncio.wait_for(clear, timeout=timeout)
clear_sender = False
stage_result = await pool.collective_rpc(
replica_id=replica_id,
method=method,
timeout=timeout,
args=args,
kwargs=kwargs,
)
except Exception as exc:
if (
method not in ("pause_scheduler", "resume_scheduler")
and method not in StagePool._CACHE_RESET_METHODS
):
raise
# Administrative failures must reach the caller without
# terminating unrelated stages. _engine_core_rpc raises
# on this result; remaining replicas still receive the RPC.
stage_result = {"supported": False, "error": f"{type(exc).__name__}: {exc}"}
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._route_upstream_first_audio(stage_id, replica_id, raw_outputs)
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),
)
request_id = getattr(eco, "request_id", None)
req_state = self.request_states.get(request_id) if request_id is not None else None
if req_state is None:
continue
if (
getattr(eco, "finish_reason", None) == FinishReason.ERROR
and not getattr(eco, "is_segment_finished", False)
and not req_state.session_owned
):
reason = getattr(eco, "stop_reason", None)
await self._handle_stage_error(
stage_id,
eco,
error=reason if isinstance(reason, str) and reason else "Stage request failed",
)
continue
if req_state.streaming.enabled:
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):
key = (stage_id, replica_id)
if key not in self._stage_replica_to_engine_idx:
engine_idx = len(self._stage_replica_to_engine_idx)
self._stat_logger.register_replica(engine_idx, str(stage_id), str(replica_id))
self._stage_replica_to_engine_idx[key] = engine_idx
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
async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_outputs: Any) -> None:
"""Order explicitly marked first audio ahead of a codec that skipped it.
The sender and codec run independently. Keep overtaking codec outputs,
including terminal outputs, on the request until its first PCM arrives.
Cancellation and stage errors use ordinary request cleanup.
"""
kept = []
final_output = self.stage_pools[stage_id].final_output
for eco in raw_outputs.outputs:
mm = getattr(eco, "multimodal_output", None)
is_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_KEY, False))
requires_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_REQUIRED_KEY, False))
request_id = getattr(eco, "request_id", None)
req_state = self.request_states.get(request_id) if request_id is not None else None
if is_first:
assert isinstance(mm, dict)
# Only the dedicated first-frame channel is intercepted.
if (
req_state is None
or req_state.upstream_first_audio
or req_state.pending_upstream_first_audio is not None
or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id
):
continue
audio = mm.get("model_outputs")
if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
await self._handle_stage_error(stage_id, eco, error="Invalid or empty first-audio payload")
continue
req_state.pending_upstream_first_audio = (stage_id, eco)
await self._flush_upstream_first_audio(req_state)
continue
if (
req_state is not None
and final_output
and not req_state.upstream_first_audio
and (requires_first or req_state.pending_first_audio_outputs)
and getattr(eco, "finish_reason", None) != FinishReason.ERROR
):
req_state.pending_first_audio_outputs.append((stage_id, replica_id, eco))
continue
kept.append(eco)
raw_outputs.outputs = kept
async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) -> None:
"""Feed upstream PCM through the codec's normal output-kind contract.
The source can finish its first frame while downstream prewarm is
still selecting/registering a replica. Retain that frame until the
codec owns an output-processor state, then accumulate it before any
codec output that omitted it. This preserves FINAL_ONLY/CUMULATIVE
waveforms and downstream consumers such as forced alignment.
"""
from vllm_omni.engine import OmniEngineCoreOutputs
pending = req_state.pending_upstream_first_audio
if pending is None or self.request_states.get(req_state.request_id) is not req_state:
return
source_stage, first_output = pending
codec_stage = source_stage if self.stage_pools[source_stage].final_output else source_stage + 1
pool = self.stage_pools[codec_stage]
replica_id = pool.get_bound_replica_id(req_state.request_id)
if replica_id is None or req_state.request_id not in pool.output_processor.request_states:
return
req_state.pending_upstream_first_audio = None
req_state.upstream_first_audio = True
try:
processed = await pool.process_llm_raw_outputs(replica_id, OmniEngineCoreOutputs(outputs=[first_output]))
await self._handle_processed_outputs(codec_stage, replica_id, processed)
except EngineDeadError as error:
await self._handle_dead_replica(codec_stage, replica_id, error)
return
except Exception as error:
# A promised first frame must either enter the codec accumulator
# or fail its request; leaving it pending would stall completion.
logger.exception("First-audio output processing failed for %s", req_state.request_id)
await self._handle_stage_error(codec_stage, first_output, error=str(error))
return
if self.request_states.get(req_state.request_id) is not req_state:
return
ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, []
for pending_stage, pending_replica, output in ready:
if self.request_states.get(req_state.request_id) is not req_state:
break
terminal: set[str] = set()
processed = await self._process_llm_stage_outputs(
pending_stage, pending_replica, OmniEngineCoreOutputs(outputs=[output]), terminal
)
await self._handle_processed_outputs(pending_stage, pending_replica, processed)
await self._finish_raw_terminal_requests(pending_stage, pending_replica, terminal)
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 the explicit ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` value or the
pipeline default computed at engine initialization. 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))
self._orch_monitor.set_dispatch_queue_size(ready_q.qsize())
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))
self._orch_monitor.set_dispatch_queue_size(ready_q.qsize())
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
# Record the queue after dequeue so the gauge reflects work
# still waiting for dispatch rather than producer-side depth.
self._orch_monitor.set_dispatch_queue_size(ready_q.qsize())
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, *, error: str | None = None) -> 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=error if error is not None else 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 NativeKVHandoffError as e:
await self._fail_request_client_error(
req_id,
stage_id,
str(e),
status_code=HTTPStatus.BAD_GATEWAY.value,
error_type="NativeKVHandoffError",
release_owners=True,
)
return False
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)
self._release_stage_transfer_resources(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", set()))
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.
Record raw stage completion here. Client completion is
resolved after this raw batch passes through the output processor, so a
real processed output wins over the swallowed-terminal fallback.
"""
if getattr(eco, "finish_reason", None) is None:
return False
if getattr(eco, "is_segment_finished", False):
return False
req_state.finished_stage_ids.add(stage_id)
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 True
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 requests whose raw terminal had no processed output."""
pool = self.stage_pools[stage_id]
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
pending = req_state.pending_final_output
if pending is not None:
if set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids):
req_state.pending_final_output = None
await self.output_async_queue.put(pending)
await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)])
# A real terminal output is pending; do not replace it with
# the swallowed-output fallback before upstream completion.
continue
if not pool.final_output:
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
if finished and not segment_finished:
req_state.finished_stage_ids.add(stage_id)
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 through the session interceptor.
pass
elif self.stage_pools[stage_id].final_output:
message = 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,
)
if (
request_finished
and self.async_chunk
and not req_state.streaming.enabled
and not set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids)
):
# Direct chunk delivery can finish the decoder before the
# upstream engine's terminal output/metrics reaches this loop.
# Keep request state until those messages are routed, otherwise
# the frontend publishes incomplete usage and loses stop reasons.
req_state.pending_final_output = message
request_finished = False
else:
await self.output_async_queue.put(message)
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 req_state.skip_legacy_stage_forward
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,
)
pending = req_state.pending_final_output
if pending is not None and set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids):
req_state.pending_final_output = None
await self.output_async_queue.put(pending)
request_finished = 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 is not None and 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
submit_kwargs = self._diffusion_submit_kwargs(req_id, src_stage_id, next_client, req_state, output)
payload_sender_info = self._build_payload_sender_info(src_stage_id, request_id=req_id)
if payload_sender_info is not None:
submit_kwargs["payload_sender_info"] = payload_sender_info
if already_submitted:
replica_id = await next_pool.submit_update(
req_id, req_state, diffusion_prompt, submit_kwargs=submit_kwargs
)
else:
replica_id = await next_pool.submit_initial(
req_id,
req_state,
diffusion_prompt,
submit_kwargs=submit_kwargs,
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:
assert next_client is not None
process_inputs = getattr(next_client, "process_engine_inputs")
next_inputs = process_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 = self._diffusion_submit_kwargs(
request_id,
next_stage_id - 1,
next_pool.stage_client,
req_state,
)
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:
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):
from vllm_omni.engine.request_snapshot import copy_request_snapshot
base_input = copy_request_snapshot(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,
)
if req_state.pending_upstream_first_audio is not None or req_state.upstream_first_audio:
await self._flush_upstream_first_audio(req_state)
if self.request_states.get(request_id) is not req_state:
return False
# 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 _maybe_attach_native_kv_transfer_params(
self,
req_state: OrchestratorRequestState,
prompt: Any,
) -> None:
source_stage_id = 0
if source_stage_id + 1 > req_state.final_stage_id:
return
config = getattr(self.stage_pools[source_stage_id].stage_vllm_config, "kv_transfer_config", None)
if getattr(config, "kv_role", None) != "kv_producer":
return
if self.stage_pools[source_stage_id + 1].stage_type != "diffusion":
return
from vllm_omni.diffusion.diffusion_kv.kv_connector import build_source_kv_transfer_params, mint_transfer_id
transfer_id = mint_transfer_id(req_state.request_id)
req_state.native_kv_transfer_id = transfer_id
params = build_source_kv_transfer_params(
transfer_id=transfer_id,
remote_engine_id=None,
remote_bootstrap_addr=None,
)
for sampling in (req_state.sampling_params_list[source_stage_id], prompt.sampling_params):
sampling.extra_args = {**(sampling.extra_args or {}), "kv_transfer_params": params}
def _diffusion_submit_kwargs(
self,
request_id: str,
source_stage_id: int,
next_client: Any,
req_state: OrchestratorRequestState,
output: Any = None,
) -> dict[str, Any]:
source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [source_stage_id])
if req_state.native_kv_transfer_id is None:
return {"kv_sender_info": self._build_kv_sender_info(source_stage_ids, request_id=request_id)}
if output is None:
raise NativeKVHandoffError("Native KV handoff requires one completed AR source")
from vllm_omni.diffusion.diffusion_kv.kv_connector import (
bootstrap_addr_from_kv_transfer_config,
build_target_kv_transfer_params,
)
source = self.stage_pools[source_stage_id].get_bound_client(request_id)
config = getattr(getattr(source, "vllm_config", None), "kv_transfer_config", None)
if source is None or config is None:
raise NativeKVHandoffError(
f"Native KV handoff for {request_id}: bound AR replica or its KV configuration is unavailable"
)
params = getattr(output, "kv_transfer_params", None)
if not params or "num_transfer_tokens" not in params:
raise NativeKVHandoffError("AR source completed without native KV transfer metadata")
return {
"kv_transfer_params": build_target_kv_transfer_params(
source_params=params,
remote_engine_id=config.engine_id,
remote_bootstrap_addr=bootstrap_addr_from_kv_transfer_config(config),
)
}
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
def _native_mrv2_receiver_stage(self, final_stage_id: int) -> int | None:
"""First downstream stage up to ``final_stage_id`` that receives on MRv2's native data plane."""
for stage_id in range(1, min(final_stage_id, len(self.stage_pools) - 1) + 1):
vllm_config = getattr(self.stage_pools[stage_id], "stage_vllm_config", None)
model_config = getattr(vllm_config, "model_config", None)
if uses_native_mrv2_data_plane(
model_config,
use_v2_model_runner=bool(getattr(model_config, "use_v2_model_runner", False)),
):
return stage_id
return None
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
if getattr(prompt, "resumable", False):
mrv2_stage_id = self._native_mrv2_receiver_stage(final_stage_id)
if mrv2_stage_id is not None:
# Streaming-input sessions replace downstream prompts through
# the V1 chunk adapter, which MRv2's native data plane lacks.
await self._fail_request_client_error(
request_id,
stage_id,
f"streaming (resumable) input is not supported: stage {mrv2_stage_id} runs on "
"model_runner v2, which supports turn-based requests only; use model_runner: v1 "
"for that stage",
)
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._maybe_attach_native_kv_transfer_params(req_state, prompt)
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,
)
)
Reference in new issue View Git Blame Copy Permalink