vllm-omni/vllm_omni/engine/orchestrator.py

# 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 time as _time
from collections.abc import Awaitable, Callable
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
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.metrics.stats import IterationStats

from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
from vllm_omni.engine.membership_controller import MembershipController
from vllm_omni.engine.messages import (
    AbortRequestMessage,
    AbortResultMessage,
    AddCompanionRequestMessage,
    CollectiveRPCRequestMessage,
    CollectiveRPCResultMessage,
    EngineQueueMessage,
    ErrorMessage,
    InteractionMessage,
    OutputMessage,
    RegisterRemoteReplicaMessage,
    ShutdownRequestMessage,
    StageMetricsMessage,
    StageSubmissionMessage,
    UnregisterRemoteReplicaMessage,
)
from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key
from vllm_omni.engine.serialization import serialize_additional_information
from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError
from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE
from vllm_omni.metrics import definitions as metric_defs
from vllm_omni.metrics.prometheus import OmniRequestCounter
from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
from vllm_omni.outputs import OmniRequestOutput

logger = init_logger(__name__)

if TYPE_CHECKING:
    from vllm_omni.experimental.fullduplex.engine.contracts import (
        DuplexControlPlanePort,
        DuplexOutputContext,
        DuplexOutputDecision,
        DuplexRequestIdentity,
        DuplexRuntimeExtension,
        DuplexStageRequestContext,
        DuplexStageSubmission,
        DuplexStageSubmissionResult,
    )
    from vllm_omni.experimental.fullduplex.engine.duplex_session import (
        DuplexSessionRuntimeManager,
        DuplexSessionRuntimeState,
    )
    from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence


def _build_terminal_empty_output(
    request_id: str,
    *,
    final_output_type: str | None,
    audio_sample_rate: int = 24000,
) -> RequestOutput:
    """Build a terminal empty output when no downstream stage input exists."""
    completion = CompletionOutput(
        index=0,
        text="",
        token_ids=[],
        cumulative_logprob=None,
        logprobs=None,
        finish_reason="stop",
        stop_reason=None,
    )
    if final_output_type == "audio":
        completion.multimodal_output = {
            "audio": torch.zeros((0,), dtype=torch.float32),
            "sr": audio_sample_rate,
        }
    return RequestOutput(
        request_id=request_id,
        prompt=None,
        prompt_token_ids=[],
        prompt_logprobs=None,
        outputs=[completion],
        finished=True,
    )


def build_engine_core_request_from_tokens(
    request_id: str,
    prompt: dict[str, Any],
    params: SamplingParams | PoolingParams,
    arrival_time: float | None = None,
    model_config: ModelConfig | None = None,
    resumable: bool = False,
    mm_features: list | None = None,
) -> OmniEngineCoreRequest:
    """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt."""
    if arrival_time is None:
        arrival_time = _time.time()

    prompt_token_ids = prompt["prompt_token_ids"]

    sampling_params = None
    pooling_params = None
    if isinstance(params, SamplingParams):
        sampling_params = params.clone()
        if model_config is not None:
            remaining = model_config.max_model_len - len(prompt_token_ids)
            if sampling_params.max_tokens is None:
                sampling_params.max_tokens = remaining
            # ``check_stop`` returns early while ``min_tokens`` is unmet, so a
            # min_tokens wider than the leftover context window never reaches
            # the length cap: the request stalls once the scheduler can grant
            # no further tokens. Stage prompts are built here, so clamp now.
            if sampling_params.min_tokens > remaining:
                sampling_params.min_tokens = max(0, remaining)
    else:
        pooling_params = params.clone()

    prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds")
    raw_additional_information = prompt.get("additional_information")
    model_intermediate_buffer = prompt.get("model_intermediate_buffer")
    wire_payload: dict[str, Any] | None = None
    if isinstance(raw_additional_information, dict):
        wire_payload = dict(raw_additional_information)
    additional_info_payload = serialize_additional_information(
        wire_payload,
        log_prefix=f"build_engine_core_request_from_tokens req={request_id}",
    )

    return OmniEngineCoreRequest(
        request_id=request_id,
        prompt_token_ids=prompt_token_ids,
        mm_features=mm_features,
        sampling_params=sampling_params,
        pooling_params=pooling_params,
        arrival_time=arrival_time,
        lora_request=getattr(params, "lora_request", None),
        cache_salt=prompt.get("cache_salt"),
        data_parallel_rank=None,
        prompt_embeds=prompt_embeds,
        resumable=resumable,
        additional_information=additional_info_payload,
        model_intermediate_buffer=model_intermediate_buffer if isinstance(model_intermediate_buffer, dict) else None,
    )


@dataclass
class OrchestratorRequestState:
    """Per-request bookkeeping inside the Orchestrator."""

    request_id: str
    prompt: Any = None
    sampling_params_list: list[Any] = field(default_factory=list)
    final_stage_id: int = -1
    final_output_stage_ids: set[int] = field(default_factory=set)
    finished_final_output_stage_ids: set[int] = field(default_factory=set)

    # Wall-clock timestamp when the client-facing engine request was accepted.
    request_timestamp: float = 0.0

    # Metrics: timestamp when request was submitted to each stage.
    stage_submit_ts: dict[int, float] = field(default_factory=dict)
    mm_processor_kwargs: dict | None = None
    mm_features: list | None = None
    pd_prefill_multimodal_output: dict[str, Any] | None = None

    streaming: StreamingInputState = field(default_factory=lambda: StreamingInputState())

    # Per-request pipeline timing accumulator (milliseconds)
    pipeline_timings: dict[str, float] = field(default_factory=dict)
    duplex_identity: DuplexRequestIdentity | None = None
    duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict)
    duplex_config_generation: int = -1
    running_counter_registered: bool = False


@dataclass
class StreamingInputState:
    # Flag of streaming input request
    enabled: bool = False
    # Flag of segment of streaming input finished
    segment_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.
    segment_token_ids: list[int] = field(default_factory=list)
    segment_output_metadata: dict[str, Any] = field(default_factory=dict)
    # Streaming update prompt length
    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


class _OrchestratorDuplexStagePort:
    """Adapts generic stage pools to the model-neutral duplex control plane."""

    def __init__(
        self,
        *,
        stage_pools: list[StagePool],
        request_states: dict[str, OrchestratorRequestState],
        running_counter: OmniRequestCounter | None,
        cleanup_request_ids: Callable[..., Any],
        async_chunk: bool,
        prewarm_async_chunk_stages: Callable[
            [str, Any, OrchestratorRequestState],
            Awaitable[bool],
        ],
    ) -> None:
        self._stage_pools = stage_pools
        self._request_states = request_states
        self._running_counter = running_counter
        self._cleanup_request_ids = cleanup_request_ids
        self._async_chunk = async_chunk
        self._prewarm_async_chunk_stages = prewarm_async_chunk_stages

    @property
    def stage_count(self) -> int:
        return len(self._stage_pools)

    def sampling_defaults(self) -> tuple[object, ...]:
        return tuple(pool.stage_client.default_sampling_params for pool in self._stage_pools)

    @staticmethod
    def _sync_bridge_state(
        request_state: OrchestratorRequestState,
        context: DuplexStageRequestContext,
    ) -> None:
        duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {})
        if not isinstance(duplex_state, dict):
            duplex_state = {}
            request_state.streaming.bridge_states["duplex"] = duplex_state
        previous_epoch = duplex_state.get("epoch")
        if not isinstance(duplex_state.get("model_turn_id"), int) or previous_epoch != context.fence.epoch:
            duplex_state["model_turn_id"] = context.fence.turn_id
        duplex_state.update(
            {
                "session_id": context.session_id,
                "fence": context.fence,
                "incarnation": context.fence.incarnation,
                "epoch": context.fence.epoch,
                "turn_id": context.fence.turn_id,
                "response_seq": context.fence.response_seq,
                "session_config": dict(context.session_config),
                "runtime_config": dict(context.runtime_config),
            }
        )

    def ensure_request(self, context: DuplexStageRequestContext) -> None:
        from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity

        request_state = self._request_states.get(context.request_id)
        if request_state is None:
            request_state = OrchestratorRequestState(
                request_id=context.request_id,
                prompt=None,
                sampling_params_list=list(context.sampling_params),
                final_stage_id=context.final_stage_id,
                duplex_config_generation=context.config_generation,
            )
            request_state.streaming.enabled = True
            self._request_states[context.request_id] = request_state
        elif request_state.duplex_config_generation != context.config_generation:
            request_state.sampling_params_list = list(context.sampling_params)
            request_state.duplex_config_generation = context.config_generation
        request_state.duplex_identity = DuplexRequestIdentity(
            session_id=context.session_id,
            fence=context.fence,
        )
        self._sync_bridge_state(request_state, context)

    async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult:
        from vllm_omni.experimental.fullduplex.engine.contracts import DuplexStageSubmissionResult

        context = submission.context
        request_state = self._request_states.get(context.request_id)
        if request_state is None:
            raise RuntimeError(f"duplex request was not preregistered: {context.request_id}")
        request = build_engine_core_request_from_tokens(
            request_id=context.request_id,
            prompt=dict(submission.prompt),
            params=context.stage_sampling_params,
            model_config=self._stage_pools[context.stage_id].stage_vllm_config.model_config,
            resumable=True,
        )
        request.external_req_id = request.request_id
        pool = self._stage_pools[context.stage_id]
        if submission.already_submitted:
            replica_id = await pool.submit_update(context.request_id, request_state, request)
        else:
            replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None)
            if self._async_chunk and context.stage_id == 0:
                prewarmed = await self._prewarm_async_chunk_stages(
                    context.request_id,
                    request,
                    request_state,
                )
                if not prewarmed:
                    # The prewarm already failed the request, aborted it, and
                    # popped its state. Falling through would write fences and
                    # a submit timestamp onto that orphaned object and
                    # re-increment the running counter the cleanup just
                    # released, then hand the control plane a success result
                    # for a request that no longer exists. Raise instead:
                    # ``DuplexControlPlane.handle_append`` turns this into an
                    # error result for the append.
                    raise RuntimeError(
                        f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted"
                    )
        request_state.duplex_stage_fences[context.stage_id] = context.fence
        request_state.stage_submit_ts[context.stage_id] = _time.time()
        if not request_state.running_counter_registered and self._running_counter is not None:
            self._running_counter.increment()
            request_state.running_counter_registered = True
        return DuplexStageSubmissionResult(
            request_id=context.request_id,
            stage_id=context.stage_id,
            replica_id=replica_id,
        )

    async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None:
        await self._cleanup_request_ids(request_ids, abort=abort)


class Orchestrator:
    """Runs inside a background thread's asyncio event loop."""

    # 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
    _transfer_emitter: Any = None
    _prom_metrics: Any = None
    _stat_logger: OmniPrometheusStatLogger | None = None
    duplex_control_plane: DuplexControlPlanePort | None = None

    def __init__(
        self,
        request_async_queue: janus.AsyncQueue[EngineQueueMessage],
        output_async_queue: janus.AsyncQueue[EngineQueueMessage],
        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
        stage_pools: list[StagePool],
        *,
        async_chunk: bool = False,
        pd_config: dict[str, Any] | None = None,
        membership_controller: MembershipController | None = None,
        running_counter: OmniRequestCounter | None = None,
        transfer_emitter: Any = None,
        prom_metrics: Any = None,
        log_stats: bool = False,
        enable_orch_monitor: bool = False,
        duplex_runtime_extension: DuplexRuntimeExtension | None = None,
        enable_duplex_control: bool = False,
        duplex_session_config: DuplexSessionRuntimeConfig | None = None,
    ) -> None:
        self.request_async_queue = request_async_queue
        self.output_async_queue = output_async_queue
        self.rpc_async_queue = rpc_async_queue

        self.async_chunk = bool(async_chunk)
        self.num_stages = len(stage_pools)
        self.stage_pools: list[StagePool] = stage_pools
        self.log_stats = log_stats
        self._prom_metrics = prom_metrics
        self._stage_replica_waiting: dict[tuple[int, int], int] = {}
        self._orch_monitor = create_orch_monitor(
            enabled=enable_orch_monitor,
            replica_sampler=self._sample_replica_metrics,
        )
        for stage_id, pool in enumerate(self.stage_pools):
            for replica_id in pool.available_replica_ids():
                self._orch_monitor.register_replica(stage_id, replica_id)

        # PD disaggregation state
        self._pd_pair: tuple[int, int] | None = None
        self._pd_bootstrap_addr: str | None = None
        self._pd_prefill_engine_id: str | None = None
        self._pd_kv_params: dict[str, Any] = {}
        if pd_config is not None:
            self._pd_pair = pd_config.get("pd_pair")
            self._pd_bootstrap_addr = pd_config.get("bootstrap_addr")
            self._pd_prefill_engine_id = pd_config.get("prefill_engine_id")
        self.request_states: dict[str, OrchestratorRequestState] = {}
        self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats)

        self._cfg_tracker = CfgCompanionTracker()
        self._stage_input_processors: dict[int, Any] = {}

        self.duplex_control_plane: DuplexControlPlanePort | None = None
        self._duplex_reaper_interval_s = 1.0
        if enable_duplex_control:
            from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import DuplexControlPlane
            from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseConfig

            runtime_session_config = duplex_session_config or DuplexSessionRuntimeConfig()
            self._duplex_reaper_interval_s = runtime_session_config.reaper_interval_s

            self.duplex_control_plane = DuplexControlPlane(
                extension=duplex_runtime_extension,
                stage_port=_OrchestratorDuplexStagePort(
                    stage_pools=self.stage_pools,
                    request_states=self.request_states,
                    running_counter=self._running_counter,
                    cleanup_request_ids=self._cleanup_request_ids,
                    async_chunk=self.async_chunk,
                    prewarm_async_chunk_stages=self._prewarm_async_chunk_stages,
                ),
                result_sink=self.rpc_async_queue,
                lifecycle_sink=self.output_async_queue,
                lease_config=DuplexLeaseConfig(
                    idle_ttl_s=runtime_session_config.idle_ttl_s,
                    disconnect_grace_s=runtime_session_config.disconnect_grace_s,
                ),
                max_sessions=runtime_session_config.max_sessions,
                completed_append_limit=runtime_session_config.completed_append_cache_size,
            )

        self._shutdown_event = asyncio.Event()
        self._stages_shutdown = False

        # 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,
        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._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

    @property
    def duplex_sessions(self) -> DuplexSessionRuntimeManager:
        if self.duplex_control_plane is None:
            raise RuntimeError("duplex control plane is disabled")
        return self.duplex_control_plane.sessions

    def _require_duplex_control_plane(self) -> DuplexControlPlanePort:
        if self.duplex_control_plane is None:
            raise RuntimeError("duplex control plane is disabled")
        return self.duplex_control_plane

    async def run(self) -> None:
        """Main entry point for the Orchestrator event loop."""
        logger.info("[Orchestrator] Starting event loop")

        request_task = asyncio.create_task(self._request_handler(), name="orchestrator-request-handler")
        output_task = asyncio.create_task(
            self._orchestration_output_handler(),
            name="orchestrator-stage-output-handler",
        )

        # Start membership watcher if distributed mode is active.
        membership_watcher: asyncio.Task[None] | None = None
        if self._membership is not None:
            self._membership.install_unregister_handlers(
                output_queue=self.output_async_queue,
                cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True),
                replica_removed_callback=self._remove_stage_replica_waiting,
            )
            membership_watcher = self._membership.start()

        tasks = [request_task, output_task]
        if self.duplex_control_plane is not None:
            tasks.append(asyncio.create_task(self._duplex_reaper_loop(), name="orchestrator-duplex-reaper"))
        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
            if self.duplex_control_plane is not None:
                await self.duplex_control_plane.shutdown()

            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 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 self.duplex_control_plane is not None and self.duplex_control_plane.accepts(msg):
                self.duplex_control_plane.dispatch(msg)
            elif msg_type == "abort":
                await self._handle_abort(msg)
            elif msg_type == "interaction":
                await self._handle_interaction(msg)
            elif msg_type == "collective_rpc":
                await self._handle_collective_rpc(msg)
            elif isinstance(msg, RegisterRemoteReplicaMessage):
                if self._membership is not None:
                    await self._membership.handle_register(msg.stage_id, msg.replica_id)
                    self._orch_monitor.register_replica(msg.stage_id, msg.replica_id)
            elif isinstance(msg, UnregisterRemoteReplicaMessage):
                if self._membership is not None:
                    await self._membership.handle_unregister(msg.stage_id, msg.input_addr)
            elif isinstance(msg, ShutdownRequestMessage):
                logger.info("[Orchestrator] Received shutdown signal")
                self._shutdown_event.set()
                # Pre-mark stage clients as shutting down to prevent
                # proc_monitor daemon threads from flagging normal
                # process exit as EngineDeadError during teardown.
                for pool in self.stage_pools:
                    for client in pool.clients:
                        if hasattr(client, "_shutting_down"):
                            client._shutting_down = True
                # Stage teardown runs once in run()'s finally after the
                # orchestration loop observes _shutdown_event and exits.
                break
            else:
                logger.warning("[Orchestrator] Unknown message type: %s", msg_type)

    async def _duplex_reaper_loop(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                await asyncio.wait_for(
                    self._shutdown_event.wait(),
                    timeout=self._duplex_reaper_interval_s,
                )
            except TimeoutError:
                plane = self.duplex_control_plane
                if plane is not None:
                    try:
                        await plane.reap_expired()
                    except Exception:
                        logger.exception("[Orchestrator] Duplex expiry cleanup failed; retrying on next tick")

    async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
        """Handle an add_request message from the main thread."""
        stage_id = 0
        request_id = msg.request_id
        prompt = msg.prompt
        original_prompt = msg.original_prompt
        sampling_params_list = msg.sampling_params_list
        if not sampling_params_list:
            raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
        final_stage_id = msg.final_stage_id
        final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id])

        if not self.stage_pools[stage_id].live_replica_ids():
            # Stage 0 lost all replicas between the HTTP-layer errored check and
            # dispatch. Runs before request state / running counter registration,
            # so the helper's cleanup is a no-op here.
            await self._fail_request_dead_stage(request_id, stage_id)
            return

        logger.debug(
            "[Orchestrator] _handle_add_request: stage=%s req=%s "
            "prompt_type=%s original_prompt_type=%s final_stage=%s "
            "num_sampling_params=%d",
            stage_id,
            request_id,
            type(prompt).__name__,
            type(original_prompt).__name__,
            final_stage_id,
            len(sampling_params_list),
        )

        req_state = OrchestratorRequestState(
            request_id=request_id,
            prompt=original_prompt,
            sampling_params_list=sampling_params_list,
            final_stage_id=final_stage_id,
            final_output_stage_ids=final_output_stage_ids,
            request_timestamp=float(msg.request_timestamp or _time.time()),
            mm_features=getattr(prompt, "mm_features", None),
        )
        self.request_states[request_id] = req_state
        self._register_running_request(req_state)
        req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
        req_state.stage_submit_ts[stage_id] = _time.time()
        enqueue_ts = msg.enqueue_ts
        if enqueue_ts > 0:
            req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0
        preprocess_ms = msg.preprocess_ms
        if preprocess_ms > 0:
            req_state.pipeline_timings["preprocess_ms"] = preprocess_ms
        if not await self._dispatch_or_fail_request(
            lambda: self.stage_pools[stage_id].submit_initial(
                request_id,
                req_state,
                prompt,
                prompt_text=msg.output_prompt_text,
            ),
            req_id=request_id,
            stage_id=stage_id,
            operation="add_request",
        ):
            return

        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
            await self._prewarm_async_chunk_stages(request_id, prompt, req_state)

    async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None:
        """Handle a streaming_update message for an existing request."""
        stage_id = 0
        request_id = msg.request_id
        request = msg.prompt
        final_stage_id = msg.final_stage_id
        req_state = self.request_states.get(request_id)
        if req_state is None:
            # Streaming updates always follow the first-chunk add_request
            # through the same ordered submission queue, so an unknown id
            # here can only mean the request already finished or was
            # aborted (e.g. the client disconnected). Re-adding it would
            # resurrect a headless session that keeps cycling through the
            # stages with nobody consuming its outputs (issue #4271).
            logger.warning(
                "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)",
                request_id,
            )
            return

        if msg.sampling_params_list:
            req_state.sampling_params_list = msg.sampling_params_list

        req_state.streaming.enabled = True
        req_state.stage_submit_ts[stage_id] = _time.time()
        if not await self._dispatch_or_fail_request(
            lambda: self.stage_pools[stage_id].submit_update(
                request_id,
                req_state,
                request,
                prompt_text=msg.output_prompt_text,
            ),
            req_id=request_id,
            stage_id=stage_id,
            operation="streaming_update",
        ):
            return

        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
            await self._prewarm_async_chunk_stages(request_id, request, req_state)

    async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None:
        """Handle an add_companion_request message: submit companion to stage 0."""
        companion_id = msg.companion_id
        parent_id = msg.parent_id
        role = msg.role
        companion_prompt = msg.prompt
        sampling_params_list = msg.sampling_params_list

        parent_state = self.request_states.get(parent_id)
        if parent_state is None:
            logger.info(
                "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active",
                companion_id,
                role,
                parent_id,
            )
            return

        self._cfg_tracker.register_companion(parent_id, role, companion_id)

        companion_state = OrchestratorRequestState(
            request_id=companion_id,
            prompt=companion_prompt,
            sampling_params_list=sampling_params_list,
            final_stage_id=0,
            final_output_stage_ids={0},
            request_timestamp=parent_state.request_timestamp,
        )
        self.request_states[companion_id] = companion_state
        companion_state.stage_submit_ts[0] = _time.time()
        # A dead-stage failure is attributed to the parent request; its cleanup
        # also releases this companion.
        if not await self._dispatch_or_fail_request(
            lambda: self.stage_pools[0].submit_initial(
                companion_id,
                companion_state,
                companion_prompt,
                prompt_text=msg.companion_prompt_text,
                affinity_request_id=parent_id,
            ),
            req_id=parent_id,
            stage_id=0,
            operation=f"companion {companion_id}",
            dispatch_req_id=companion_id,
        ):
            return

        logger.info(
            "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)",
            companion_id,
            role,
            parent_id,
            self.stage_pools[0].get_bound_replica_id(companion_id),
        )

    async def _handle_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.
        """
        request_ids = msg.request_ids
        error: str | None = None
        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.
            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,
                )
            )

    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,
                )
            )

    async def _abort_request_ids(self, request_ids: list[str]) -> None:
        """Forward abort requests to all stage pools."""
        if not request_ids:
            return
        for pool in self.stage_pools:
            await pool.abort_requests(request_ids)
            pool.release_bindings(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:
            for replica_id in pool.live_replica_ids():
                stage_result = await pool.collective_rpc(
                    replica_id=replica_id,
                    method=method,
                    timeout=timeout,
                    args=args,
                    kwargs=kwargs,
                )
                stage_ids.append(pool.stage_id)
                results.append(stage_result)

        await self.rpc_async_queue.put(
            CollectiveRPCResultMessage(
                rpc_id=rpc_id,
                method=method,
                stage_ids=stage_ids,
                results=results,
            )
        )

    # ---- Orchestration loop ----

    def _sample_replica_metrics(self) -> dict[str, tuple[int, int]]:
        samples: dict[str, tuple[int, int]] = {}
        for stage_id, pool in enumerate(self.stage_pools):
            for replica_id in pool.live_replica_ids():
                key = replica_key(stage_id, replica_id)
                samples[key] = pool.replica_monitor_sample(replica_id)
        return samples

    async def _orchestration_output_handler(self) -> None:
        """Poll all stages, handle transfers, send final outputs to main."""
        try:
            await self._orchestration_loop()
        except asyncio.CancelledError:
            logger.debug("[Orchestrator] _orchestration_output_handler cancelled")
            return

    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

                    # 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

                            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))

                            if diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID:
                                idle = False
                                continue

                            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

                            await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs)
                            for eco in raw_outputs.outputs:
                                # Emit kv_wait_s before _handle_kv_ready_raw_outputs'
                                # async_chunk early-return so it lands in all modes.
                                kv_params = getattr(eco, "kv_transfer_params", None)
                                if (
                                    self._prom_metrics is not None
                                    and isinstance(kv_params, dict)
                                    and (kv_wait_s := kv_params.get("kv_wait_s")) is not None
                                ):
                                    self._prom_metrics.observe_kv_wait(
                                        kv_params.get("connector_type") or "unknown",
                                        float(kv_wait_s),
                                    )
                                req_state = self.request_states.get(getattr(eco, "request_id", None))
                                if req_state is None or not req_state.streaming.enabled:
                                    continue
                                req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
                                req_state.streaming.segment_token_ids = (
                                    self._coerce_int_list(getattr(eco, "new_token_ids", None))
                                    if req_state.streaming.segment_finished
                                    else []
                                )
                                raw_mm = self._completion_multimodal_output(eco, None)
                                req_state.streaming.segment_output_metadata = (
                                    dict(raw_mm)
                                    if req_state.streaming.segment_finished and isinstance(raw_mm, dict)
                                    else {}
                                )
                                req_state.streaming.new_prompt_len_snapshot = getattr(
                                    eco,
                                    "new_prompt_len_snapshot",
                                    None,
                                )
                                if req_state.streaming.enabled:
                                    await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state)
                            iteration_stats = (
                                IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
                            )
                            processed = await pool.process_llm_raw_outputs(
                                replica_id,
                                raw_outputs,
                                iteration_stats=iteration_stats,
                            )
                            if self._stat_logger is not None and (
                                raw_outputs.scheduler_stats is not None or iteration_stats is not None
                            ):
                                self._stat_logger.record(
                                    raw_outputs.scheduler_stats,
                                    iteration_stats,
                                    engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)],
                                )
                            _sched_stats = raw_outputs.scheduler_stats
                            if (
                                self._prom_metrics is not None
                                and _sched_stats is not None
                                and getattr(_sched_stats, "num_waiting_reqs", None) is not None
                            ):
                                self._update_stage_replica_waiting(
                                    stage_id,
                                    replica_id,
                                    int(_sched_stats.num_waiting_reqs),
                                )
                    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)
                    idle = False

            self._orch_monitor.note_loop(idle=idle)
            if idle:
                await asyncio.sleep(0.001)
            else:
                await asyncio.sleep(0)

    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_finished
            if output.finished or segment_finished:
                stage_metrics = pool.build_stage_metrics(
                    [output],
                    submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()),
                    request_timestamp=req_state.request_timestamp,
                    replica_id=replica_id,
                    sampling_params=req_state.sampling_params_list[stage_id],
                )
                stage_metrics.pipeline_timings = dict(req_state.pipeline_timings)

            await self._route_output(stage_id, replica_id, output, req_state, stage_metrics)

    async def _handle_stage_error(self, stage_id: int, output: Any) -> None:
        """Emit a frontend-visible error and clean up request state."""
        if self._cfg_tracker.is_companion(output.request_id):
            parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id
        else:
            parent_id = output.request_id
        await self.output_async_queue.put(
            ErrorMessage(
                request_id=parent_id,
                stage_id=stage_id,
                error=output.error,
                status_code=getattr(output, "error_status_code", None),
                error_type=getattr(output, "error_type", None),
            )
        )
        await self._cleanup_request_ids(
            [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)],
            abort=True,
            close_duplex_sessions=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)

    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)

    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,
            )

    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,
        )

    async def _fail_request_client_error(
        self,
        req_id: str,
        stage_id: int,
        error: str,
        *,
        close_duplex_sessions: bool = False,
    ) -> None:
        """Fail one request with a non-fatal 400 (bad input, engine survives).

        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=HTTPStatus.BAD_REQUEST.value,
                error_type=DEFAULT_CLIENT_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,
            close_duplex_sessions=close_duplex_sessions,
        )

    async def _dispatch_or_fail_request(
        self,
        dispatch: Callable[[], Awaitable[Any]],
        *,
        req_id: str,
        stage_id: int,
        operation: str,
        dispatch_req_id: str | None = None,
    ) -> bool:
        """Run a dispatch action, converting stage unavailability into a
        single-request failure.

        ``dispatch`` is a thunk rather than a bare coroutine so that replica
        selection and other argument setup run inside this guard: a synchronous
        StageUnavailableError raised while building the dispatch is caught here,
        not left to propagate out of the caller.

        Replica eviction can interleave with any dispatch, surfacing as
        StageUnavailableError (no live replica / evicted slot) or
        EngineDeadError (replica died mid-submit before the poll loop evicted
        it). Both mean "this request cannot be placed", not "the server is
        broken": fail the request and keep serving (#4285). An OverflowError
        from encoding an out-of-range request value is likewise failed as a
        client error (400) rather than killing the loop. Other unrelated errors
        propagate. Returns False when the request was failed.

        ``req_id`` is the request the failure is attributed to; ``dispatch_req_id``
        is the id actually being dispatched when it differs (e.g. a CFG companion
        submitted under its parent's attribution) and is used to locate the
        replica that died.
        """
        try:
            await dispatch()
            return True
        except StageUnavailableError as e:
            # No specific replica to evict: the stage already has no live
            # replica or the chosen slot was evicted. Fail just this request.
            logger.error(
                "[Orchestrator] %s dispatch for req=%s hit unavailable stage-%s: %s",
                operation,
                req_id,
                stage_id,
                e,
            )
            await self._fail_request_dead_stage(req_id, stage_id)
            return False
        except EngineDeadError as e:
            # A live replica died mid-submit before the poll loop evicted it.
            # Evict it now (and fail every request bound to it) so a burst of
            # new requests stops round-robining onto the dead slot until the
            # poll catches up. Capture the replica before failing this request,
            # since the failure releases its binding.
            pool = self.stage_pools[stage_id]
            dead_replica = pool.get_bound_replica_id(dispatch_req_id or req_id)
            logger.error(
                "[Orchestrator] %s dispatch for req=%s hit dead stage-%s replica-%s: %s",
                operation,
                req_id,
                stage_id,
                dead_replica,
                e,
            )
            await self._fail_request_dead_stage(req_id, stage_id)
            if dead_replica is not None and dead_replica in pool.live_replica_ids():
                await self._handle_dead_replica(stage_id, dead_replica, e)
            return False
        except OverflowError as e:
            # Catch overflow errors in the request payload to ensure that msgpack
            # encoding can't kill the engine; instead, we just fail this request.
            # The frontend should generally validate this as well, but this is needed
            # in case we get values like seed > 2**64-1 from sampling params.
            logger.warning(
                "[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s",
                operation,
                req_id,
                stage_id,
                e,
            )
            await self._fail_request_client_error(
                req_id,
                stage_id,
                f"Request contains a value that cannot be serialized: {e}",
            )
            return False

    # ---- Shared helpers ----

    async def _cleanup_request_ids(
        self,
        request_ids: list[str],
        *,
        abort: bool = False,
        close_duplex_sessions: bool = False,
    ) -> None:
        """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.
        """
        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)
        closing_session_ids: list[str] = []
        if close_duplex_sessions and self.duplex_control_plane is not None:
            closed_sessions = self.duplex_control_plane.close_sessions_for_request_ids(
                cleanup_ids,
                abort=abort,
                cleanup_in_progress=True,
            )
            closing_session_ids.extend(closed_sessions)
            for session_id, stale_request_ids in closed_sessions.items():
                logger.info(
                    "[Orchestrator] closed duplex session %s while cleaning failed request ids %s",
                    session_id,
                    stale_request_ids,
                )
                cleanup_ids.extend(stale_request_ids)
            cleanup_ids = list(dict.fromkeys(cleanup_ids))

        try:
            if abort:
                await self._abort_request_ids(cleanup_ids)
            self._release_request_bindings(cleanup_ids)
            for request_id in cleanup_ids:
                self._pd_kv_params.pop(request_id, None)
                req_state = self.request_states.pop(request_id, None)
                if req_state is not None and req_state.running_counter_registered and self._running_counter is not None:
                    self._running_counter.decrement()
                    req_state.running_counter_registered = False
        except BaseException:
            if closing_session_ids and self.duplex_control_plane is not None:
                self.duplex_control_plane.defer_request_cleanups(closing_session_ids)
            raise
        if closing_session_ids and self.duplex_control_plane is not None:
            self.duplex_control_plane.finalize_closed_sessions(closing_session_ids)

    async def _apply_raw_terminal_stage_finish(
        self,
        stage_id: int,
        eco: Any,
        req_state: OrchestratorRequestState,
    ) -> None:
        """Record session-level finish markers dropped by the streaming output processor.

        Streaming segment stops set ``is_segment_finished=True`` and are handled
        via processed outputs. Session termination (e.g. ``finish_requests`` after
        ``resumable=False``) emits a terminal ``finish_reason`` with
        ``is_segment_finished=False``, but vLLM's output processor may remove the
        request state before that EngineCoreOutput is processed.

        Only update ``finished_final_output_stage_ids`` here. Request cleanup stays
        in ``_route_output`` so downstream async-chunk stages can still deliver
        outputs after stage-0 session end.
        """
        if getattr(eco, "finish_reason", None) is None:
            return
        if getattr(eco, "is_segment_finished", False):
            return

        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
        req_state.finished_final_output_stage_ids.add(stage_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 _duplex_session_for_req_state(self, req_state: OrchestratorRequestState) -> DuplexSessionRuntimeState | None:
        if self.duplex_control_plane is None:
            return None
        return self.duplex_control_plane.session_for_identity(req_state.duplex_identity)

    def _record_duplex_stage_submission(
        self,
        stage_id: int,
        request_id: str,
        replica_id: int,
        req_state: OrchestratorRequestState,
    ) -> None:
        del replica_id
        identity = req_state.duplex_identity
        session = self._duplex_session_for_req_state(req_state)
        if identity is None or session is None:
            return
        req_state.duplex_stage_fences[stage_id] = identity.fence
        session.bind_stage_request(stage_id, request_id, fence=identity.fence)
        req_state.stage_submit_ts[stage_id] = _time.time()
        self._register_running_request(req_state)

    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)
        # CFG companion: stash output so parent can bundle [parent, *companions]
        # into source_outputs for the bridge (e.g. thinker2imagegen).
        if finished and self._cfg_tracker.is_companion(req_id):
            self._cfg_tracker.set_companion_output(req_id, output)
            await self._handle_cfg_companion_ready(req_id)
            await self._cleanup_request_ids([req_id])
            return

        request_finished = False
        if (
            finished
            and self.stage_pools[stage_id].final_output
            and not (req_state.streaming.enabled and req_state.streaming.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)
        # Duplex stage-0 segment boundaries are not client-visible outputs:
        # direct decisions are emitted by the model runtime extension below,
        # while spoken content flows through the next stage. Forwarding
        # the raw stage-0 output as well injects one cumulative-text,
        # no-audio message per unit that every downstream consumer must
        # filter out again (the official implementation returns exactly one
        # result per audio chunk).
        is_duplex_stage0_segment = (
            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished
        )
        if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment:
            await self.output_async_queue.put(
                OutputMessage(
                    request_id=req_id,
                    stage_id=stage_id,
                    replica_id=replica_id,
                    engine_outputs=output,
                    metrics=stage_metrics,
                    finished=(
                        request_finished
                        or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished)
                    ),
                    stage_submit_ts=submit_ts,
                )
            )
        elif stage_metrics is not None:
            await self.output_async_queue.put(
                StageMetricsMessage(
                    request_id=req_id,
                    stage_id=stage_id,
                    replica_id=replica_id,
                    metrics=stage_metrics,
                    stage_submit_ts=submit_ts,
                )
            )

        if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]:
            kv_params = getattr(output, "kv_transfer_params", None)
            if kv_params is not None:
                self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params)
            req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None)

        duplex_output_decision = self._duplex_output_decision(stage_id, output, req_state)
        if duplex_output_decision is not None:
            await self._emit_duplex_direct_output(
                stage_id,
                req_id,
                output,
                duplex_output_decision,
                stage_metrics,
                submit_ts,
            )
            return

        if (
            (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished))
            and stage_id < req_state.final_stage_id
            and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))
            and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled)
        ):
            if (
                finished
                and self._cfg_tracker.has_companions(req_id)
                and not self._cfg_tracker.all_companions_done(req_id)
            ):
                self._cfg_tracker.defer_parent(req_id, output, stage_id)
            else:
                stage_params = req_state.sampling_params_list[stage_id]
                final_only_finished = (
                    req_state.streaming.enabled
                    and finished
                    and getattr(stage_params, "output_kind", None) == RequestOutputKind.FINAL_ONLY
                )
                await self._forward_to_next_stage(
                    req_id,
                    stage_id,
                    output,
                    req_state,
                    src_replica_id=replica_id,
                    is_streaming_session=req_state.streaming.enabled,
                    is_final_update=final_only_finished,
                )
                if (
                    req_state.streaming.enabled
                    and finished
                    and not final_only_finished
                    and not self._is_duplex_session_request(req_state)
                ):
                    # For streaming sessions, send the terminal (resumable=False) update only on a finish
                    await self._forward_to_next_stage(
                        req_id,
                        stage_id,
                        output,
                        req_state,
                        src_replica_id=replica_id,
                        is_streaming_session=True,
                        is_final_update=True,
                    )

        if request_finished and not self._is_duplex_session_request(req_state):
            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,
    ) -> 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
            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
        return request

    @staticmethod
    def _duplex_output_context(
        req_state: OrchestratorRequestState,
        *,
        stage_id: int | None = None,
    ) -> DuplexOutputContext | None:
        identity = req_state.duplex_identity
        if identity is None:
            return None
        from vllm_omni.experimental.fullduplex.engine.contracts import (
            DuplexOutputContext,
            DuplexRequestIdentity,
        )

        fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence
        return DuplexOutputContext(
            identity=DuplexRequestIdentity(
                session_id=identity.session_id,
                fence=fence,
            ),
            final_stage_id=req_state.final_stage_id,
            segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished,
            segment_token_ids=tuple(req_state.streaming.segment_token_ids),
            segment_output_metadata=req_state.streaming.segment_output_metadata,
        )

    @staticmethod
    def _is_duplex_session_request(req_state: OrchestratorRequestState) -> bool:
        return req_state.duplex_identity is not None

    @staticmethod
    def _duplex_fence_for_req_state(
        req_state: OrchestratorRequestState,
        *,
        stage_id: int | None = None,
    ) -> DuplexFence | None:
        context = Orchestrator._duplex_output_context(req_state, stage_id=stage_id)
        return context.identity.fence if context is not None else None

    def _duplex_output_decision(
        self,
        stage_id: int,
        output: Any,
        req_state: OrchestratorRequestState,
    ) -> DuplexOutputDecision | None:
        if self.duplex_control_plane is None:
            return None
        context = self._duplex_output_context(req_state, stage_id=stage_id)
        decision = self.duplex_control_plane.decide_output(
            stage_id,
            output,
            context,
        )
        return decision

    async def _emit_duplex_direct_output(
        self,
        stage_id: int,
        req_id: str,
        output: Any,
        decision: DuplexOutputDecision,
        stage_metrics: Any,
        submit_ts: float | None,
    ) -> None:
        action = getattr(decision.action, "value", decision.action)
        if action != "direct_response":
            raise ValueError(f"Unsupported duplex output action: {action}")
        from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision

        engine_output = attach_duplex_output_decision(
            OmniRequestOutput.from_stage_output(
                output,
                request_id=req_id,
                finished=True,
                stage_id=stage_id,
                final_output_type=decision.final_output_type,
            ),
            decision,
        )
        await self.output_async_queue.put(
            OutputMessage(
                request_id=req_id,
                stage_id=stage_id,
                engine_outputs=engine_output,
                metrics=stage_metrics,
                finished=True,
                stage_submit_ts=submit_ts,
            )
        )

    @staticmethod
    def _completion_multimodal_output(output: Any, completion: Any) -> dict[str, Any]:
        mm_output = getattr(output, "multimodal_output", None)
        if isinstance(mm_output, dict):
            return mm_output
        mm_output = getattr(completion, "multimodal_output", None) if completion is not None else None
        return mm_output if isinstance(mm_output, dict) else {}

    @classmethod
    def _coerce_int_list(cls, value: Any) -> list[int]:
        if value is None:
            return []
        if hasattr(value, "detach"):
            try:
                value = value.detach().cpu().reshape(-1).tolist()
            except Exception:
                return []
        if not isinstance(value, (list, tuple)):
            return []
        out: list[int] = []
        for item in value:
            token_id = cls._coerce_int(item)
            if token_id is not None:
                out.append(token_id)
        return out

    @staticmethod
    def _coerce_int(value: Any) -> int | None:
        if hasattr(value, "detach"):
            try:
                value = value.detach().cpu().reshape(-1)
                if value.numel() == 0:
                    return None
                value = value[0].item()
            except Exception:
                return None
        try:
            return int(value)
        except (TypeError, ValueError):
            return None

    async def _handle_cfg_companion_ready(self, req_id: str) -> None:
        """Mark a CFG companion as done; if all companions are done, flush deferred parent."""
        parent_id = self._cfg_tracker.on_companion_completed(req_id)
        if parent_id is None:
            return

        deferred = self._cfg_tracker.pop_pending_parent(parent_id)
        if deferred is None:
            return

        parent_state = self.request_states.get(parent_id)
        if parent_state is None:
            return

        stage_id = deferred["stage_id"]
        if (stage_id + 1) in parent_state.stage_submit_ts:
            return

        await self._forward_to_next_stage(
            parent_id,
            stage_id,
            deferred["engine_outputs"],
            parent_state,
        )

    async def _handle_kv_ready_raw_outputs(
        self,
        stage_id: int,
        raw_outputs: EngineCoreOutputs,
    ) -> None:
        """Forward split requests once stage-0 KV is ready."""
        if self.async_chunk:
            return

        for raw_output in raw_outputs.outputs:
            kv_params = getattr(raw_output, "kv_transfer_params", None)
            if not (isinstance(kv_params, dict) and kv_params.get("kv_ready")):
                continue

            req_id = raw_output.request_id
            req_state = self.request_states.get(req_id)
            if req_state is None:
                continue
            if self._cfg_tracker.is_companion(req_id):
                # kv_ready only says the companion's KV hit the connector; its
                # processed output has not been stashed yet. Counting it as
                # done here let the parent pass the all_companions_done gate
                # and dispatch with 0/N companion outputs (degraded CFG).
                # Done-marking now happens solely on the processed-output path
                # in _route_output, after set_companion_output.
                continue
            if stage_id >= req_state.final_stage_id:
                continue
            if (stage_id + 1) in req_state.stage_submit_ts:
                continue

            if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id):
                self._cfg_tracker.defer_parent(req_id, raw_output, stage_id)
            else:
                await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state)

    def _build_pd_decode_params(self, req_id: str, sp: Any) -> Any:
        """Build decode-side sampling params with KV transfer params for PD routing.

        Clones the sampling params and injects kv_transfer_params that tell the
        decode engine where to pull the KV cache from (prefill engine's bootstrap addr).
        """
        sp = sp.clone()
        if sp.extra_args is None:
            sp.extra_args = {}

        # Get KV params captured from the prefill output (must include remote_request_id).
        kv_prefill_params = self._pd_kv_params.pop(req_id, None)
        if not kv_prefill_params or "remote_request_id" not in kv_prefill_params:
            raise RuntimeError(
                f"[Orchestrator][PD] Missing prefill kv_transfer_params.remote_request_id for req={req_id}"
            )

        decode_kv_params: dict[str, Any] = {
            "transfer_id": f"xfer-{req_id}",
        }

        if self._pd_bootstrap_addr:
            decode_kv_params["remote_bootstrap_addr"] = self._pd_bootstrap_addr

        if self._pd_prefill_engine_id:
            decode_kv_params["remote_engine_id"] = self._pd_prefill_engine_id

        # Overlay params from prefill side (includes remote_request_id set by monkey patch).
        decode_kv_params.update(kv_prefill_params)

        # Ensure these flags are set correctly after any overlay.
        decode_kv_params["do_remote_prefill"] = True
        decode_kv_params["do_remote_decode"] = False
        if not decode_kv_params.get("transfer_id"):
            decode_kv_params["transfer_id"] = f"xfer-{req_id}"

        sp.extra_args["kv_transfer_params"] = decode_kv_params

        logger.debug(
            "[Orchestrator][PD] decode kv_transfer_params for req=%s: %s",
            req_id,
            decode_kv_params,
        )
        return sp

    def _emit_tx_edge(
        self,
        *,
        from_stage: int,
        from_replica: int,
        to_stage: int,
        to_pool: StagePool,
        request_id: str,
        tx_ms: float,
    ) -> None:
        """Emit per-edge transfer_tx_s + transfer_size_bytes histograms.

        ``tx_ms`` is the orchestrator-side wall-clock spent in ``next_pool.
        submit_*`` (serialize + queue submit to the receiving worker). Best-
        effort size_bytes left at 0 — orchestrator doesn't have a cheap handle
        on the serialized payload size; a follow-up can plumb that from the
        connector adapter.
        """
        if self._transfer_emitter is None:
            return
        to_replica = to_pool.get_bound_replica_id(request_id)
        if to_replica is None:
            return
        try:
            self._transfer_emitter.observe_size(from_stage, from_replica, to_stage, to_replica, 0)
            self._transfer_emitter.observe_tx_time(from_stage, from_replica, to_stage, to_replica, tx_ms / 1000.0)
        except Exception:
            logger.debug(
                "[Orchestrator] transfer_tx emit failed for edge %d->%d req=%s",
                from_stage,
                to_stage,
                request_id,
                exc_info=True,
            )

    async def _forward_to_next_stage(
        self,
        req_id: str,
        src_stage_id: int,
        output: Any,
        req_state: OrchestratorRequestState,
        *,
        src_replica_id: int | None = None,
        is_streaming_session: bool = False,
        is_final_update: bool = False,
    ) -> None:
        """Forward output to the next stage; a dead stage fails only this request."""
        await self._dispatch_or_fail_request(
            lambda: self._forward_to_next_stage_unguarded(
                req_id,
                src_stage_id,
                output,
                req_state,
                src_replica_id=src_replica_id,
                is_streaming_session=is_streaming_session,
                is_final_update=is_final_update,
            ),
            req_id=req_id,
            stage_id=src_stage_id + 1,
            operation="inter-stage forward",
        )

    async def _forward_to_next_stage_unguarded(
        self,
        req_id: str,
        src_stage_id: int,
        output: Any,
        req_state: OrchestratorRequestState,
        *,
        src_replica_id: int | None = None,
        is_streaming_session: bool = False,
        is_final_update: bool = False,
    ) -> None:
        """Forward output from the current logical stage to the next one."""
        next_logical = src_stage_id + 1
        next_pool = self.stage_pools[next_logical]
        if not next_pool.live_replica_ids():
            # The downstream stage lost all of its replicas before this in-flight
            # request could be forwarded.
            await self._fail_request_dead_stage(req_id, next_logical)
            return
        next_client = next_pool.stage_client
        params = req_state.sampling_params_list[next_logical]
        source_outputs = [output]
        next_stage_resumable = is_streaming_session and not is_final_update
        already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
        requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
        _t_submit_start = _time.perf_counter()

        if next_pool.stage_type == "diffusion":
            # Gate: never dispatch with an incomplete CFG bundle. Checked
            # non-destructively BEFORE popping — a pop-then-redefer would lose
            # the partial outputs. all_companions_done is trustworthy here
            # because done-marking happens only after set_companion_output.
            if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id):
                self._cfg_tracker.defer_parent(req_id, output, src_stage_id)
                logger.info(
                    "[Orchestrator] req=%s: CFG companion outputs not all stashed yet; re-deferring parent",
                    req_id,
                )
                return
            # Peek, don't pop: outputs stay stashed until cleanup_parent so a
            # streaming re-submission bundles the complete set again rather
            # than an empty one.
            companion_outputs = self._cfg_tracker.get_companion_outputs(req_id)
            expected = len(self._cfg_tracker.get_companion_request_ids(req_id))
            if expected > len(companion_outputs):
                # Companions are done but outputs are missing — inconsistent
                # tracker state (should be unreachable with peek semantics).
                # Fail the request rather than degrade CFG conditioning.
                logger.error(
                    "[Orchestrator] req=%s: only %d/%d CFG companion outputs available; "
                    "failing the request instead of dispatching degraded CFG",
                    req_id,
                    len(companion_outputs),
                    expected,
                )
                await self.output_async_queue.put(
                    ErrorMessage(
                        request_id=req_id,
                        stage_id=src_stage_id,
                        error=(
                            f"CFG companion outputs incomplete ({len(companion_outputs)}/{expected}); "
                            "request aborted to avoid degraded CFG conditioning"
                        ),
                    )
                )
                await self._cleanup_request_ids(
                    [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                    abort=True,
                )
                return
            diffusion_source_outputs = [output, *companion_outputs]
            if next_client.custom_process_input_func is not None:
                _t_ar2d = _time.perf_counter()
                _fn = next_client.custom_process_input_func
                _extra_kwargs: dict[str, Any] = {}
                # TODO: replace signature probe with explicit kwarg contract.
                try:
                    import inspect as _inspect

                    if "sampling_params" in _inspect.signature(_fn).parameters:
                        _extra_kwargs["sampling_params"] = params
                except (TypeError, ValueError):
                    pass
                diffusion_prompt = _fn(
                    diffusion_source_outputs,
                    req_state.prompt,
                    requires_multimodal_data,
                    **_extra_kwargs,
                )
                _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
                req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
                logger.info(
                    "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
                    req_id,
                    _dt_ar2d,
                    src_stage_id,
                    next_logical,
                )
                if diffusion_prompt is None:
                    error_output = OmniRequestOutput.from_error(
                        req_id,
                        f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
                    )
                    logger.warning(
                        "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
                        "routing terminal error output",
                        req_id,
                        src_stage_id,
                        next_logical,
                    )
                    await self.output_async_queue.put(
                        OutputMessage(
                            request_id=req_id,
                            stage_id=next_logical,
                            engine_outputs=error_output,
                            metrics=None,
                            finished=True,
                        )
                    )
                    await self._cleanup_request_ids(
                        [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                    )
                    return
                if isinstance(diffusion_prompt, list):
                    if not diffusion_prompt:
                        error_output = OmniRequestOutput.from_error(
                            req_id,
                            f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
                        )
                        logger.warning(
                            "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
                            "routing terminal error output",
                            req_id,
                            src_stage_id,
                            next_logical,
                        )
                        await self.output_async_queue.put(
                            OutputMessage(
                                request_id=req_id,
                                stage_id=next_logical,
                                engine_outputs=error_output,
                                metrics=None,
                                finished=True,
                            )
                        )
                        await self._cleanup_request_ids(
                            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                        )
                        return
                    if len(diffusion_prompt) == 1:
                        diffusion_prompt = diffusion_prompt[0]
            else:
                diffusion_prompt = req_state.prompt

            if already_submitted:
                replica_id = await next_pool.submit_update(req_id, req_state, diffusion_prompt)
            else:
                replica_id = await next_pool.submit_initial(
                    req_id,
                    req_state,
                    diffusion_prompt,
                    submit_kwargs={
                        "kv_sender_info": self._build_kv_sender_info(
                            list(getattr(next_client, "engine_input_source", None) or [src_stage_id]),
                            request_id=req_id,
                        )
                    },
                    params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
                )
            self._record_duplex_stage_submission(
                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._record_duplex_stage_submission(
                    next_logical,
                    req_id,
                    replica_id,
                    req_state,
                )

            req_state.stage_submit_ts[next_logical] = _time.time()
            _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
            self._emit_tx_edge(
                from_stage=src_stage_id,
                from_replica=src_replica_id if src_replica_id is not None else 0,
                to_stage=next_logical,
                to_pool=next_pool,
                request_id=req_id,
                tx_ms=_tx_ms,
            )
            return

        if req_state.pd_prefill_multimodal_output is not None:
            req_state.streaming.bridge_states.setdefault(
                "pd_prefill_multimodal_output_by_req",
                {},
            )[req_id] = req_state.pd_prefill_multimodal_output

        previous_decoder = req_state.streaming.source_token_decoder
        source_processor = self.stage_pools[src_stage_id].output_processor
        tokenizer = getattr(source_processor, "tokenizer", None)
        decode = getattr(tokenizer, "decode", None)
        if callable(decode):
            req_state.streaming.source_token_decoder = decode

        try:
            next_inputs = next_client.process_engine_inputs(
                source_outputs,
                req_state.prompt,
                streaming_context=req_state.streaming,
            )
        except Exception:
            logger.exception(
                "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s",
                req_id,
                next_logical,
            )
            raise
        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,
            )

            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._record_duplex_stage_submission(
                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",
                close_duplex_sessions=True,
            )
            return False

        for next_stage_id in range(1, req_state.final_stage_id + 1):
            next_pool = self.stage_pools[next_stage_id]
            params = req_state.sampling_params_list[next_stage_id]
            if not self._stage_receives_async_chunks(next_stage_id):
                # Outgoing-only stages receive their first real input from the
                # orchestrator. Pre-submitting a placeholder lets it race and
                # execute before that conditioning payload arrives.
                continue

            req_state.stage_submit_ts[next_stage_id] = _time.time()
            _t_submit_start = _time.perf_counter()

            if next_pool.stage_type == "diffusion":
                submit_kwargs = {
                    "kv_sender_info": self._build_kv_sender_info(
                        list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
                        request_id=request_id,
                    )
                }
                submitted = await self._dispatch_or_fail_request(
                    lambda: next_pool.submit_initial(
                        request_id,
                        req_state,
                        req_state.prompt,
                        submit_kwargs=submit_kwargs,
                    ),
                    req_id=request_id,
                    stage_id=next_stage_id,
                    operation="async-chunk prewarm",
                )
            else:
                import copy

                from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length

                try:
                    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
                except Exception:
                    next_prompt_len = max(1, len(prompt_token_ids))

                original_prompt = req_state.prompt
                if isinstance(original_prompt, dict):
                    base_input = copy.deepcopy(original_prompt)
                else:
                    base_input = {}

                base_input["prompt_token_ids"] = [0] * next_prompt_len
                base_input["multi_modal_data"] = None
                base_input["mm_processor_kwargs"] = None
                downstream_resumable = bool(getattr(stage0_request, "resumable", req_state.streaming.enabled))
                request = build_engine_core_request_from_tokens(
                    request_id=request_id,
                    prompt=base_input,
                    params=params,
                    model_config=next_pool.stage_vllm_config.model_config,
                    resumable=downstream_resumable,
                )
                request.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._record_duplex_stage_submission(
                next_stage_id,
                request_id,
                bound_replica_id if bound_replica_id is not None else 0,
                req_state,
            )

            # async_chunk pre-submit fires per stage edge (N-1 -> N). Source
            # replica is stage 0's bound replica (single-replica thinker in
            # all current configs); fall back to 0 if unknown.
            _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
            src_replica = self.stage_pools[next_stage_id - 1].get_bound_replica_id(request_id)
            self._emit_tx_edge(
                from_stage=next_stage_id - 1,
                from_replica=src_replica if src_replica is not None else 0,
                to_stage=next_stage_id,
                to_pool=next_pool,
                request_id=request_id,
                tx_ms=_tx_ms,
            )

        return True

    def _build_kv_sender_info(
        self,
        sender_stage_ids: list[int],
        *,
        request_id: str | None = None,
    ) -> dict[int, dict[str, Any]] | None:
        """Build per-request sender info for diffusion KV-transfer receivers."""
        sender_infos: dict[int, dict[str, Any]] = {}
        for sender_stage_id in dict.fromkeys(sender_stage_ids):
            if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools):
                continue

            sender_pool = self.stage_pools[sender_stage_id]
            sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None
            if sender_stage is None:
                sender_stage = sender_pool.stage_client
            get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
            if not callable(get_sender_info):
                continue

            sender_info = get_sender_info()
            if not sender_info:
                logger.warning(
                    "[Orchestrator] Stage-%s has no KV sender info available",
                    sender_stage_id,
                )
                continue

            sender_infos[sender_stage_id] = sender_info

        return sender_infos or None

    # ---- 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)