vllm-omni/vllm_omni/engine/duplex_orchestrator.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""DuplexOrchestrator: the duplex sibling of ``Orchestrator``.

It hosts the ``DuplexSessionManager`` (admission, leases, session runners) on
the orchestrator loop and implements the ``DuplexStagePort`` the runners use
to submit resumable Stage0 requests. Generic stage forwarding, prewarm and
cleanup stay in ``OrchestratorBase``; this class only fills the template seams
and applies the session-owned policy.
"""

from __future__ import annotations

import time as _time
from collections.abc import Coroutine
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from vllm.logger import init_logger

from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex.contracts import (
    DuplexFence,
    DuplexOutputContext,
    DuplexRequestIdentity,
    DuplexStagePort,
    DuplexStageRequestContext,
    DuplexStageSubmission,
    DuplexStageSubmissionResult,
)
from vllm_omni.engine.duplex.plugin import DuplexModelPlugin
from vllm_omni.engine.duplex.session.engine_session import DuplexFenceMismatchError
from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
from vllm_omni.engine.messages import EngineQueueMessage, OutputMessage
from vllm_omni.engine.orchestrator import (
    Orchestrator,
    OrchestratorRequestState,
    build_engine_core_request_from_tokens,
)

if TYPE_CHECKING:
    from vllm.config import ModelConfig
    from vllm.outputs import RequestOutput

    from vllm_omni.metrics.stats import StageRequestStats

logger = init_logger(__name__)


@dataclass
class DuplexOrchestratorRequestState(OrchestratorRequestState):
    """Request bookkeeping for a session-owned (resumable duplex) stage request."""

    session_id: str = ""
    fence: DuplexFence | None = None
    stage_fences: dict[int, DuplexFence] = field(default_factory=dict)
    config_generation: int = -1


class DuplexOrchestrator(Orchestrator, DuplexStagePort):
    """Stage management for a duplex deployment; owns one ``DuplexSessionManager``.

    Extends the turn-based orchestrator rather than sitting beside it, so one
    engine serves both a duplex session and an ordinary request. The
    dependency runs duplex -> turn-based and never the reverse, which is what
    keeps ``Orchestrator`` free of duplex code.
    """

    def __init__(
        self,
        # Any: ``*args`` / ``**kwargs`` are forwarded verbatim to ``OrchestratorBase.__init__``.
        *args: Any,
        plugin: DuplexModelPlugin,
        duplex_session_config: DuplexSessionRuntimeConfig,
        model_config: ModelConfig | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.plugin = plugin
        self.duplex_session_config = duplex_session_config
        self.session_manager = DuplexSessionManager(
            plugin=plugin,
            stage_port=self,
            output_sink=self.output_async_queue,
            result_sink=self.rpc_async_queue,
            runtime_config=self.duplex_session_config,
            model_config=model_config,
        )

    # ------------------------------------------------------------------ #
    # Template seams                                                     #
    # ------------------------------------------------------------------ #

    async def _dispatch_message(self, msg: EngineQueueMessage) -> bool:
        if self.session_manager.accepts(msg):
            self.session_manager.dispatch(msg)
            return True
        # Not a session message: it is an ordinary turn-based request.
        return await super()._dispatch_message(msg)

    # Any: ``Coroutine``'s send/yield parameters, as in the ``OrchestratorBase`` seam.
    def _background_tasks(self) -> list[Coroutine[Any, Any, None]]:
        return [self.session_manager.reaper_loop(self._shutdown_event)]

    async def _shutdown_extensions(self) -> None:
        await self.session_manager.shutdown()

    def _on_stage_submitted(
        self,
        stage_id: int,
        request_id: str,
        replica_id: int,
        req_state: OrchestratorRequestState,
    ) -> None:
        del replica_id
        if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None:
            return
        runner = self.session_manager.runner_for_request_id(request_id)
        if runner is None:
            return
        fence = req_state.fence
        req_state.stage_fences[stage_id] = fence
        try:
            runner.session.bind_stage_request(stage_id, request_id, fence=fence)
        except DuplexFenceMismatchError:
            # The session already advanced past this request's epoch (cancel
            # raced the submit); the cancel path aborts the stale request id.
            logger.debug("[DuplexOrchestrator] stale stage binding ignored for %s stage-%s", request_id, stage_id)
            return
        req_state.stage_submit_ts[stage_id] = _time.time()
        self._register_running_request(req_state)
        self.session_manager.register_request(request_id, req_state.session_id)

    async def _intercept_stage_output(
        self,
        stage_id: int,
        replica_id: int,
        output: RequestOutput,
        req_state: OrchestratorRequestState,
        stage_metrics: StageRequestStats | None,
        submit_ts: float | None,
    ) -> bool:
        del replica_id, submit_ts
        if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None:
            return False
        request_id = output.request_id
        runner = self.session_manager.runner_for_request_id(request_id)
        if runner is None:
            # Session gone: nothing may forward or reach a client.
            return True
        segment = req_state.streaming.segment(stage_id)
        context = DuplexOutputContext(
            identity=DuplexRequestIdentity(
                session_id=req_state.session_id,
                fence=req_state.stage_fences.get(stage_id, req_state.fence),
            ),
            final_stage_id=req_state.final_stage_id,
            segment_finished=req_state.streaming.enabled and segment.finished,
            segment_token_ids=tuple(segment.token_ids),
            segment_output_metadata=segment.output_metadata,
        )
        return runner.on_stage_output(
            stage_id,
            output,
            stage_metrics,
            request_id=request_id,
            context=context,
        )

    async def _handle_forward_failure(
        self,
        req_id: str,
        next_stage_id: int,
        req_state: OrchestratorRequestState,
        exc: BaseException,
    ) -> bool:
        if not req_state.session_owned:
            return False
        runner = self.session_manager.runner_for_request_id(req_id)
        if runner is not None:
            runner.on_stage_failure(next_stage_id, exc)
        await self._cleanup_request_ids(
            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
            abort=True,
            release_owners=True,
        )
        return True

    async def _cleanup_request_ids(
        self,
        request_ids: list[str],
        *,
        abort: bool = False,
        release_owners: bool = False,
    ) -> list[OutputMessage]:
        if not request_ids:
            return []
        cleanup_ids = list(dict.fromkeys(request_ids))
        closing_session_ids: list[str] = []
        if release_owners:
            closed_sessions = self.session_manager.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(
                    "[DuplexOrchestrator] 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:
            outputs = await super()._cleanup_request_ids(cleanup_ids, abort=abort)
        except BaseException:
            if closing_session_ids:
                self.session_manager.defer_request_cleanups(closing_session_ids)
            raise
        if closing_session_ids:
            self.session_manager.finalize_closed_sessions(closing_session_ids)
        for request_id in cleanup_ids:
            self.session_manager.unregister_request(request_id)
        return outputs

    # ------------------------------------------------------------------ #
    # DuplexStagePort                                                    #
    # ------------------------------------------------------------------ #

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

    def sampling_defaults(self) -> tuple[object, ...]:
        defaults = []
        for pool in self.stage_pools:
            client = pool.stage_client
            if client is None:
                # Every replica of this stage was evicted (engine core died); a
                # session cannot be admitted until the deployment is restarted.
                raise RuntimeError(f"stage {pool.stage_id} has no live replica")
            defaults.append(client.default_sampling_params)
        return tuple(defaults)

    @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")
        current_model_turn_id = duplex_state.get("model_turn_id")
        if (
            not isinstance(current_model_turn_id, int)
            or previous_epoch != context.fence.epoch
            or current_model_turn_id < context.fence.turn_id
        ):
            # A safety boundary can close a response before the model emits
            # its normal turn_eos. Catch up the engine-owned identity when the
            # next fenced append starts.
            duplex_state["model_turn_id"] = context.fence.turn_id
        duplex_state.update(
            {
                "session_id": context.session_id,
                "fence": context.fence,
                "epoch": context.fence.epoch,
                "turn_id": context.fence.turn_id,
                "session_config": dict(context.session_config),
                "runtime_config": dict(context.runtime_config),
            }
        )

    def ensure_request(self, context: DuplexStageRequestContext) -> None:
        request_state = self.request_states.get(context.request_id)
        if request_state is None:
            request_state = DuplexOrchestratorRequestState(
                request_id=context.request_id,
                prompt=None,
                sampling_params_list=list(context.sampling_params),
                final_stage_id=context.final_stage_id,
                session_owned=True,
                session_id=context.session_id,
                fence=context.fence,
                config_generation=context.config_generation,
            )
            request_state.streaming.enabled = True
            self.request_states[context.request_id] = request_state
        elif isinstance(request_state, DuplexOrchestratorRequestState):
            if request_state.config_generation != context.config_generation:
                request_state.sampling_params_list = list(context.sampling_params)
                request_state.config_generation = context.config_generation
            request_state.session_id = context.session_id
            request_state.fence = context.fence
        else:
            raise RuntimeError(f"request {context.request_id} is not session-owned")
        self._sync_bridge_state(request_state, context)
        self.session_manager.register_request(context.request_id, context.session_id)

    async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult:
        context = submission.context
        request_state = self.request_states.get(context.request_id)
        if not isinstance(request_state, DuplexOrchestratorRequestState):
            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; the runner turns this into an append
                    # error instead of writing onto an orphaned object.
                    raise RuntimeError(
                        f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted"
                    )
        request_state.stage_fences[context.stage_id] = context.fence
        request_state.stage_submit_ts[context.stage_id] = _time.time()
        self._register_running_request(request_state)
        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)

    async def abort_requests(self, request_ids: list[str]) -> None:
        if request_ids:
            await self._abort_request_ids(list(dict.fromkeys(request_ids)))


__all__ = ["DuplexOrchestrator", "DuplexOrchestratorRequestState"]