vllm-omni/vllm_omni/entrypoints/duplex_request_client.py

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

from __future__ import annotations

import asyncio
import time
from collections.abc import Callable, MutableMapping
from typing import Protocol

from vllm_omni.engine.duplex.contracts import (
    duplex_data_plane_request_info,
    duplex_resource_request_id,
)
from vllm_omni.engine.duplex.lease import DuplexLeaseActivity
from vllm_omni.engine.duplex.messages import DuplexFence
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.outputs.duplex import get_duplex_output_decision


class DuplexEnginePort(Protocol):
    """Engine control methods required by the experimental request client.

    The protocol keeps request preregistration and output collection testable
    without coupling this component to ``AsyncOmniEngine``.
    """

    async def open_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ...

    async def append_duplex_input_async(self, session_id: str, **kwargs) -> dict[str, object]: ...

    async def signal_duplex_turn_async(self, session_id: str, **kwargs) -> dict[str, object]: ...

    async def close_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ...

    async def touch_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ...

    async def resume_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ...


class DuplexRequestOutputPort:
    """Narrow output-processing surface owned by the AsyncOmni entrypoint."""

    def __init__(
        self,
        *,
        request_states: MutableMapping[str, ClientRequestState],
        num_stages: int,
        log_stats: bool,
        start_output_handler: Callable[[], object],
        process_single_result: Callable[..., OmniRequestOutput | None],
    ) -> None:
        self.request_states = request_states
        self.num_stages = num_stages
        self.log_stats = log_stats
        self._start_output_handler = start_output_handler
        self._process_single_result = process_single_result

    def start_output_handler(self) -> None:
        self._start_output_handler()

    def process_single_result(self, *args, **kwargs) -> OmniRequestOutput | None:
        return self._process_single_result(*args, **kwargs)


class DuplexRequestClient:
    """Owns duplex request identity, preregistration and output collection."""

    def __init__(self, engine: DuplexEnginePort, output_port: DuplexRequestOutputPort) -> None:
        self.engine = engine
        self.output_port = output_port

    @staticmethod
    def request_info(result: dict[str, object]) -> tuple[str | None, int | None]:
        return duplex_data_plane_request_info(result)

    async def open(
        self,
        session_id: str,
        *,
        session_mode: str,
        capabilities: dict[str, object] | None,
        session_config: dict[str, object] | None,
        runtime_config: dict[str, object] | None,
        fence: DuplexFence,
        timeout: float | None,
    ) -> dict[str, object]:
        # Lifecycle expiry messages share the engine output channel. Start its
        # sole dispatcher before open so an idle session can be reaped even if
        # it never submits an append request.
        self.output_port.start_output_handler()
        kwargs: dict[str, object] = {
            "session_mode": session_mode,
            "capabilities": capabilities,
            "session_config": session_config,
            "fence": fence,
            "timeout": timeout,
        }
        if runtime_config is not None:
            kwargs["runtime_config"] = runtime_config
        return await self.engine.open_duplex_session_async(session_id, **kwargs)

    async def append(
        self,
        session_id: str,
        *,
        mode: str,
        payload: object,
        operation_id: str | None,
        final: bool,
        expected_epoch: int | None,
        fence: DuplexFence,
        timeout: float | None,
        collect_outputs: bool,
    ) -> dict[str, object]:
        kwargs: dict[str, object] = {
            "mode": mode,
            "payload": payload,
            "final": final,
            "expected_epoch": expected_epoch,
            "timeout": timeout,
            "fence": fence,
        }
        if operation_id is not None:
            kwargs["operation_id"] = operation_id

        expected_request_id = duplex_resource_request_id(fence, "stage0")
        request_state = self.output_port.request_states.get(expected_request_id)
        created_request_state = request_state is None
        if request_state is None:
            request_state = ClientRequestState(expected_request_id)
            self.output_port.request_states[expected_request_id] = request_state
        if request_state.metrics is None:
            wall_start_ts = time.time()
            request_state.metrics = OrchestratorMetrics(
                self.output_port.num_stages,
                self.output_port.log_stats,
                wall_start_ts,
                max(0, self.output_port.num_stages - 1),
            )
            request_state.request_arrival_ts = wall_start_ts
        self.output_port.start_output_handler()
        try:
            result = await self.engine.append_duplex_input_async(session_id, **kwargs)
        except BaseException:
            if created_request_state and request_state.queue.empty():
                self.output_port.request_states.pop(expected_request_id, None)
            raise

        request_id, response_stage_id = duplex_data_plane_request_info(result)
        if request_id is None:
            if created_request_state and request_state.queue.empty():
                self.output_port.request_states.pop(expected_request_id, None)
            return result
        if request_id != expected_request_id:
            if created_request_state and request_state.queue.empty():
                self.output_port.request_states.pop(expected_request_id, None)
            raise RuntimeError(
                f"duplex data-plane request id mismatch: expected {expected_request_id!r}, got {request_id!r}"
            )
        if not collect_outputs:
            return result
        outputs = await self.collect_outputs(
            request_id,
            request_state,
            response_stage_id=response_stage_id,
            timeout=timeout,
        )
        if outputs:
            result = dict(result)
            result["data_plane_outputs"] = outputs
        return result

    async def collect_registered_outputs(
        self,
        request_id: str,
        *,
        response_stage_id: int | None,
        timeout: float | None,
    ) -> list[OmniRequestOutput]:
        self.output_port.start_output_handler()
        request_state = self.output_port.request_states.get(request_id)
        if request_state is None:
            return []
        return await self.collect_outputs(
            request_id,
            request_state,
            response_stage_id=response_stage_id,
            timeout=timeout,
        )

    async def signal(
        self,
        session_id: str,
        *,
        event: str,
        fence: DuplexFence,
        next_fence: DuplexFence | None,
        session_config: dict[str, object] | None,
        runtime_config: dict[str, object] | None,
        timeout: float | None,
    ) -> dict[str, object]:
        kwargs: dict[str, object] = {"event": event, "fence": fence, "timeout": timeout}
        if next_fence is not None:
            kwargs["next_fence"] = next_fence
        if session_config is not None:
            kwargs["session_config"] = session_config
        if runtime_config is not None:
            kwargs["runtime_config"] = runtime_config
        result = await self.engine.signal_duplex_turn_async(session_id, **kwargs)
        if event in {"barge_in", "input.cancel", "response.cancel"}:
            self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None)
        return result

    async def close(
        self,
        session_id: str,
        *,
        reason: str,
        fence: DuplexFence,
        timeout: float | None,
    ) -> dict[str, object]:
        result = await self.engine.close_duplex_session_async(
            session_id,
            reason=reason,
            fence=fence,
            timeout=timeout,
        )
        self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None)
        return result

    async def touch(
        self,
        session_id: str,
        *,
        fence: DuplexFence,
        activity: DuplexLeaseActivity,
        timeout: float | None,
    ) -> dict[str, object]:
        return await self.engine.touch_duplex_session_async(
            session_id,
            fence=fence,
            activity=activity,
            timeout=timeout,
        )

    async def resume(
        self,
        session_id: str,
        *,
        fence: DuplexFence,
        expected_lease_generation: int,
        timeout: float | None,
    ) -> dict[str, object]:
        return await self.engine.resume_duplex_session_async(
            session_id,
            fence=fence,
            expected_lease_generation=expected_lease_generation,
            timeout=timeout,
        )

    async def collect_outputs(
        self,
        request_id: str,
        request_state: ClientRequestState,
        *,
        response_stage_id: int | None,
        timeout: float | None,
    ) -> list[OmniRequestOutput]:
        deadline = None if timeout is None else time.monotonic() + timeout
        wall_start_ts = request_state.request_arrival_ts or time.time()
        final_stage_id = response_stage_id if response_stage_id is not None else max(0, self.output_port.num_stages - 1)
        num_stages = max(self.output_port.num_stages, final_stage_id + 1)
        metrics = getattr(request_state, "metrics", None) or OrchestratorMetrics(
            num_stages,
            self.output_port.log_stats,
            wall_start_ts,
            final_stage_id,
        )
        stage_event_cursor = metrics.stage_event_cursor(request_id)
        request_start_ts = {request_id: wall_start_ts}
        outputs: list[OmniRequestOutput] = []
        while True:
            remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
            if remaining == 0.0:
                break
            try:
                message = await asyncio.wait_for(request_state.queue.get(), timeout=remaining)
            except asyncio.TimeoutError:
                break
            if isinstance(message, ErrorMessage):
                raise RuntimeError(message.error)
            if not isinstance(message, OutputMessage):
                continue
            engine_outputs = message.engine_outputs
            is_response_stage = response_stage_id is None or message.stage_id >= response_stage_id
            is_direct_response = self.is_direct_response(engine_outputs)
            output_to_collect = None
            if isinstance(engine_outputs, OmniRequestOutput):
                output_to_collect = engine_outputs
                if message.metrics is not None:
                    existing_metrics = dict(output_to_collect.metrics)
                    metrics.process_stage_metrics(
                        result={"metrics": message.metrics},
                        stage_type="",
                        stage_id=message.stage_id,
                        req_id=request_id,
                        engine_outputs=engine_outputs,
                        finished=message.finished,
                        final_output_type=engine_outputs.final_output_type,
                        output_to_yield=output_to_collect,
                        event_cursor=stage_event_cursor,
                    )
                    generated_metrics = dict(output_to_collect.metrics)
                    existing_stage_metrics = existing_metrics.get("stage_metrics")
                    generated_stage_metrics = generated_metrics.get("stage_metrics")
                    if isinstance(existing_stage_metrics, dict) and isinstance(generated_stage_metrics, dict):
                        generated_metrics["stage_metrics"] = {
                            **existing_stage_metrics,
                            **generated_stage_metrics,
                        }
                    output_to_collect.metrics = {
                        **existing_metrics,
                        **generated_metrics,
                    }
            elif is_response_stage:
                output_to_collect = self.output_port.process_single_result(
                    message,
                    message.stage_id,
                    metrics,
                    request_start_ts,
                    wall_start_ts,
                    final_stage_id,
                    stage_event_cursor,
                )
            if output_to_collect is not None and (is_response_stage or is_direct_response):
                if message.finished:
                    output_to_collect.finished = True
                outputs.append(output_to_collect)
            if message.finished or outputs:
                break
        return outputs

    @classmethod
    def is_direct_response(cls, output: object) -> bool:
        multimodal_output = cls.multimodal_output(output)
        if not multimodal_output:
            return False
        return multimodal_output.get("duplex_direct_response") is True or multimodal_output.get(
            "duplex_native_decision"
        ) in {"listen", "speak"}

    @classmethod
    def multimodal_output(cls, output: object) -> dict[str, object]:
        if isinstance(output, OmniRequestOutput):
            decision = get_duplex_output_decision(output)
            if decision is not None:
                return dict(decision.metadata)
            private_output = getattr(output, "_multimodal_output", None)
            multimodal_output = output.multimodal_output
            if isinstance(multimodal_output, dict) and multimodal_output:
                return multimodal_output
            inner_output = getattr(output, "request_output", None)
            if inner_output is not None and inner_output is not output:
                inner_multimodal_output = cls.multimodal_output(inner_output)
                if inner_multimodal_output:
                    return inner_multimodal_output
            if isinstance(private_output, dict):
                return private_output
            return multimodal_output if isinstance(multimodal_output, dict) else {}
        multimodal_output = getattr(output, "multimodal_output", None)
        if isinstance(multimodal_output, dict):
            return multimodal_output
        outputs = getattr(output, "outputs", None)
        completion = outputs[0] if isinstance(outputs, list) and outputs else None
        multimodal_output = getattr(completion, "multimodal_output", None) if completion is not None else None
        return multimodal_output if isinstance(multimodal_output, dict) else {}


__all__ = ["DuplexEnginePort", "DuplexRequestClient", "DuplexRequestOutputPort"]