106 lines
4.3 KiB
Python
106 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import random_uuid
|
|
|
|
from vllm_omni.entrypoints.duplex.warmup import DUPLEX_WARMUP_CLIENT_WAIT_S
|
|
from vllm_omni.entrypoints.openai.realtime.connection import OpenAIFullDuplexConnection
|
|
from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_QWEN3_OMNI_REALTIME_ARCH = "Qwen3OmniMoeForConditionalGeneration"
|
|
_QWEN3_OMNI_REALTIME_STAGES = {"thinker", "talker", "code2wav"}
|
|
|
|
|
|
def supports_qwen3_omni_realtime(stage_configs: Any) -> bool:
|
|
def stage_arg(stage: Any, name: str) -> Any:
|
|
engine_args = getattr(stage, "engine_args", None)
|
|
return engine_args.get(name) if isinstance(engine_args, Mapping) else getattr(engine_args, name, None)
|
|
|
|
stages = stage_configs or ()
|
|
return (
|
|
len(stages) == len(_QWEN3_OMNI_REALTIME_STAGES)
|
|
and all(stage_arg(stage, "model_arch") == _QWEN3_OMNI_REALTIME_ARCH for stage in stages)
|
|
and {stage_arg(stage, "model_stage") for stage in stages} == _QWEN3_OMNI_REALTIME_STAGES
|
|
)
|
|
|
|
|
|
async def reject_realtime_websocket(websocket: WebSocket, message: str) -> None:
|
|
await websocket.accept()
|
|
await websocket.send_json(
|
|
{
|
|
"event_id": f"evt_{random_uuid()}",
|
|
"type": "error",
|
|
"error": {
|
|
"type": "invalid_request_error",
|
|
"code": "unsupported_model",
|
|
"message": message,
|
|
"param": "model",
|
|
"event_id": None,
|
|
},
|
|
}
|
|
)
|
|
await websocket.close(code=1008)
|
|
|
|
|
|
async def dispatch_generic_realtime_websocket(websocket: WebSocket) -> None:
|
|
"""Fall back to the generic OpenAI Realtime API for non-Qwen3-Omni models."""
|
|
serving = getattr(websocket.app.state, "openai_serving_realtime", None)
|
|
if serving is None:
|
|
await websocket.accept()
|
|
await websocket.send_json({"type": "error", "error": "Realtime API is not available", "code": "unsupported"})
|
|
await websocket.close()
|
|
return
|
|
connection = RealtimeConnection(websocket, serving)
|
|
await connection.handle_connection()
|
|
|
|
|
|
async def dispatch_realtime_websocket(websocket: WebSocket) -> None:
|
|
"""Handle Realtime sessions not already routed to the proprietary duplex
|
|
handler by the caller: Qwen3-Omni gets a conformant full-duplex connection,
|
|
everything else falls back to the generic turn-based Realtime API."""
|
|
state = websocket.app.state
|
|
if not supports_qwen3_omni_realtime(getattr(state, "stage_configs", None)):
|
|
await dispatch_generic_realtime_websocket(websocket)
|
|
return
|
|
|
|
model_name = state.openai_serving_models.base_model_paths[0].name
|
|
# Some OpenAI-compatible clients always send ``model=`` even when the
|
|
# caller leaves model selection to this single-model Omni server.
|
|
requested_model = websocket.query_params.get("model")
|
|
if requested_model and requested_model != model_name:
|
|
await reject_realtime_websocket(websocket, f"Model '{requested_model}' is not available")
|
|
return
|
|
|
|
# Hold real clients until the startup duplex warmup finishes (the warmup
|
|
# connection marks itself with vllm_omni_warmup=1 and passes through).
|
|
warmup_done = getattr(state, "duplex_warmup_done", None)
|
|
if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1":
|
|
try:
|
|
await asyncio.wait_for(warmup_done.wait(), timeout=DUPLEX_WARMUP_CLIENT_WAIT_S)
|
|
except (TimeoutError, asyncio.TimeoutError):
|
|
logger.warning(
|
|
"Duplex warmup still running after %d s; admitting the client anyway.",
|
|
DUPLEX_WARMUP_CLIENT_WAIT_S,
|
|
)
|
|
|
|
tokenizer = await state.engine_client.get_tokenizer()
|
|
connection = OpenAIFullDuplexConnection(
|
|
websocket=websocket,
|
|
engine=state.engine_client,
|
|
model_name=model_name,
|
|
tokenizer=tokenizer,
|
|
tool_call_parser=getattr(state.args, "tool_call_parser", None),
|
|
enable_auto_tool_choice=getattr(state.args, "enable_auto_tool_choice", False),
|
|
)
|
|
await connection.handle_connection()
|