From b21df3bcb6761dcb5f6489ccfa9c65dd1a6a1c0f Mon Sep 17 00:00:00 2001 From: NATURE Date: Wed, 7 Oct 2026 18:49:34 +0800 Subject: [PATCH] [Frontend] Share regular and duplex chat initialization (#7646) Signed-off-by: natureofnature --- .../openai_api/test_api_server_guards.py | 32 +++++-- .../openai_api/test_duplex_api_server.py | 57 +++++++++-- vllm_omni/entrypoints/openai/api_server.py | 96 ++++++++----------- 3 files changed, 116 insertions(+), 69 deletions(-) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 7bae57d36..c1282e060 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -1053,7 +1053,8 @@ async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch, tmp @pytest.mark.asyncio -async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None: +@pytest.mark.parametrize("supported_tasks", [("generate",), ("embed",)]) +async def test_multistage_app_state_key_snapshot(monkeypatch, mocker, tmp_path, supported_tasks) -> None: """Lock multi-stage ``app.state`` keys after init, including live vs None. Fails if chat/speech/video/realtime/tokenization keys disappear or are @@ -1068,6 +1069,7 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None: parallel_config=SimpleNamespace(_api_process_rank=0), ), ) + mocker.patch.object(engine, "get_supported_tasks", return_value=supported_tasks) class _FakeModels: def __init__(self, *args, **kwargs): @@ -1078,10 +1080,12 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None: class _FakeCtor: def __init__(self, *args, **kwargs): - pass + self.args = args + self.kwargs = kwargs + self.warmup_calls = 0 def warmup(self): - return None + self.warmup_calls += 1 deploy = tmp_path / "deploy.yaml" deploy.write_text("speech_cache:\n resolve_max_bytes: 1024\n speaker_max_bytes: 8\n") @@ -1119,17 +1123,33 @@ async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None: monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor) state = State() - await api_server.omni_init_app_state(engine, state, _minimal_args()) + await api_server.omni_init_app_state(engine, state, _minimal_args(log_error_stack=True)) assert speech_kwargs["speech_cache_config"].resolve_max_bytes == 1024 assert speech_kwargs["speech_cache_config"].resolve_max_entries == 2048 assert speech_kwargs["speech_cache_config"].speaker_max_bytes == 8 + disabled = ( + set() + if "generate" in supported_tasks + else {"openai_serving_chat", "openai_serving_chat_batch", "openai_streaming_video"} + ) _assert_app_state_snapshot( state, expected_keys=_MULTISTAGE_APP_STATE_KEYS, - must_be_wired=_MULTISTAGE_MUST_BE_WIRED, - must_be_none=_MULTISTAGE_MUST_BE_NONE, + must_be_wired=_MULTISTAGE_MUST_BE_WIRED - disabled, + must_be_none=_MULTISTAGE_MUST_BE_NONE | disabled, ) + assert state.online_renderer.kwargs["log_error_stack"] is True + assert state.online_renderer.warmup_calls == 1 + for service in (state.openai_serving_chat, state.openai_serving_chat_batch): + if "generate" in supported_tasks: + assert service.kwargs["online_renderer"] is state.online_renderer + else: + assert service is None + if "generate" in supported_tasks: + assert state.openai_serving_responses.args[2] is state.online_renderer + else: + assert state.openai_serving_responses is None @pytest.mark.parametrize("count", [None, 0, -1, "2", True]) diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py index 653c32321..3f764d074 100644 --- a/tests/entrypoints/openai_api/test_duplex_api_server.py +++ b/tests/entrypoints/openai_api/test_duplex_api_server.py @@ -166,6 +166,7 @@ def _minimal_args(**overrides) -> SimpleNamespace: enable_force_include_usage=False, enable_log_outputs=False, enable_log_deltas=False, + log_error_stack=False, ) for key, value in overrides.items(): setattr(args, key, value) @@ -316,7 +317,8 @@ def test_minicpmo_serving_profiles(monkeypatch, deploy_name, expected): @pytest.mark.asyncio -async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> None: +@pytest.mark.parametrize("log_error_stack", [False, True]) +async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch, mocker, log_error_stack) -> None: """Lock the duplex ``app.state``: the two session-backed surfaces are live, every turn-based service is None. Fails if a turn-based service is wired into a duplex server (its route @@ -326,12 +328,12 @@ async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> the turn-based chat service. """ monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) - monkeypatch.setattr(api_server, "OnlineRenderer", lambda **kwargs: SimpleNamespace(**kwargs)) + renderer_ctor = mocker.patch.object(api_server, "OnlineRenderer") monkeypatch.setattr(api_server, "OmniOpenAIServingChat", lambda **kwargs: SimpleNamespace(**kwargs)) engine = _FakeDuplexOmni() state = State() - await api_server.omni_init_app_state(engine, state, _minimal_args()) + await api_server.omni_init_app_state(engine, state, _minimal_args(log_error_stack=log_error_stack)) present = {key for key in _DUPLEX_APP_STATE_KEYS if hasattr(state, key)} assert present == _DUPLEX_APP_STATE_KEYS @@ -341,6 +343,10 @@ async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> assert not unexpectedly_set, f"turn-based services wired into a duplex server: {unexpectedly_set}" assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) assert state.openai_serving_chat is not None + renderer_ctor.assert_called_once() + assert renderer_ctor.call_args.kwargs["log_error_stack"] is log_error_stack + assert state.openai_serving_chat.online_renderer is state.online_renderer is renderer_ctor.return_value + state.online_renderer.warmup.assert_called_once_with() assert state.engine_client is engine assert state.vllm_config is engine._vllm_config @@ -498,16 +504,27 @@ async def test_warmup_gate_lets_the_warmup_connection_and_plain_servers_through( @pytest.mark.asyncio -async def test_a_model_that_does_not_declare_chat_completions_does_not_get_the_route(monkeypatch) -> None: - """The decision stays the model's: no capability, no chat service, and the route says so.""" +@pytest.mark.parametrize( + "supports_chat, supported_tasks", [(False, ("generate",)), (True, ())], ids=["no-capability", "no-generation"] +) +async def test_duplex_chat_initialization_respects_model_gates( + monkeypatch, mocker, supports_chat, supported_tasks +) -> None: + """Neither gate may construct or warm up a chat renderer.""" monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + renderer_ctor = mocker.patch.object(api_server, "OnlineRenderer") + chat_ctor = mocker.patch.object(api_server, "OmniOpenAIServingChat") engine = _FakeDuplexOmni() - engine._capabilities = DuplexCapabilities(supports_chat_completions=False) + engine._capabilities = DuplexCapabilities(supports_chat_completions=supports_chat) + mocker.patch.object(engine, "get_supported_tasks", return_value=supported_tasks) state = State() await api_server.omni_init_app_state(engine, state, _minimal_args()) assert state.openai_serving_chat is None + assert state.openai_serving_chat_batch is None + renderer_ctor.assert_not_called() + chat_ctor.assert_not_called() assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) @@ -535,3 +552,31 @@ def test_turn_deployment_rejects_explicit_duplex_without_falling_back_to_stt(fla with pytest.raises(WebSocketDisconnect) as exc: websocket.receive_text() assert exc.value.code == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "explicit_template, tokenizer_template, tokenizer_error, expected", + [ + ("explicit", None, False, "explicit"), + (None, "tokenizer", False, None), + (None, None, False, "model-json"), + (None, None, True, "model-json"), + ], +) +async def test_shared_chat_template_resolution( + mocker, explicit_template, tokenizer_template, tokenizer_error, expected +): + engine = _FakeDuplexOmni() + tokenizer = mocker.patch.object( + engine, + "get_tokenizer", + return_value=SimpleNamespace(chat_template=tokenizer_template), + side_effect=RuntimeError("tokenizer unavailable") if tokenizer_error else None, + ) + mocker.patch.object(api_server, "load_chat_template", return_value=explicit_template) + model_template = mocker.patch.object(api_server, "_load_model_chat_template_json", return_value="model-json") + + assert await api_server._resolve_chat_template(engine, _minimal_args()) == expected + assert tokenizer.call_count == (0 if explicit_template is not None else 1) + assert model_template.call_count == (1 if expected == "model-json" else 0) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 553da3fe6..cda13f31f 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -737,16 +737,37 @@ async def _init_duplex_chat( if "generate" not in supported_tasks: return None + resolved_chat_template = await _resolve_chat_template(engine_client, args) + _init_chat_services( + engine_client, state, args, request_logger, resolved_chat_template, enable_chat=True, enable_batch=False + ) + return state.openai_serving_chat + + +async def _resolve_chat_template(engine_client: EngineClient, args: Namespace) -> str | None: resolved_chat_template = load_chat_template(args.chat_template) if resolved_chat_template is None: try: tokenizer = await engine_client.get_tokenizer() except Exception as exc: - logger.debug("Could not inspect tokenizer chat_template before duplex chat init: %s", exc) + logger.debug("Could not inspect tokenizer chat_template before serving init: %s", exc) tokenizer = None if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: resolved_chat_template = _load_model_chat_template_json(args.model) + return resolved_chat_template + +def _init_chat_services( + engine_client: EngineClient, + state: State, + args: Namespace, + request_logger: RequestLogger | None, + resolved_chat_template: str | None, + *, + enable_chat: bool, + enable_batch: bool, +) -> None: + """Share chat setup while callers retain capability and endpoint gating.""" state.online_renderer = OnlineRenderer( model_config=engine_client.model_config, renderer=engine_client.renderer, @@ -759,8 +780,9 @@ async def _init_duplex_chat( tool_parser=args.tool_call_parser, reasoning_parser=args.structured_outputs_config.reasoning_parser, default_chat_template_kwargs=args.default_chat_template_kwargs, + log_error_stack=args.log_error_stack, ) - return OmniOpenAIServingChat( + chat_kwargs = dict( engine_client=engine_client, models=state.openai_serving_models, response_role=args.response_role, @@ -780,6 +802,12 @@ async def _init_duplex_chat( enable_log_outputs=args.enable_log_outputs, enable_log_deltas=args.enable_log_deltas, ) + state.openai_serving_chat = OmniOpenAIServingChat(**chat_kwargs) if enable_chat else None + state.openai_serving_chat_batch = ( + OmniOpenAIServingChatBatch(**chat_kwargs) if enable_chat and enable_batch else None + ) + # Warm up template processing, not model inference, before serving requests. + state.online_renderer.warmup() async def omni_init_app_state( @@ -923,15 +951,7 @@ async def omni_init_app_state( supported_tasks = set(await engine_client.get_supported_tasks()) logger.info("Supported tasks: %s", supported_tasks) - resolved_chat_template = load_chat_template(args.chat_template) - if resolved_chat_template is None: - try: - tokenizer = await engine_client.get_tokenizer() - except Exception as exc: - logger.debug("Could not inspect tokenizer chat_template before serving init: %s", exc) - tokenizer = None - if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: - resolved_chat_template = _load_model_chat_template_json(args.model) + resolved_chat_template = await _resolve_chat_template(engine_client, args) chat_template_config = ChatTemplateConfig( chat_template=resolved_chat_template, @@ -1003,21 +1023,14 @@ async def omni_init_app_state( ) await state.openai_serving_models.init_static_loras() - # NOTE: kept aligned with upstream `init_app_state`: - # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream). - state.online_renderer = OnlineRenderer( - model_config=engine_client.model_config, - renderer=engine_client.renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - default_chat_template_kwargs=args.default_chat_template_kwargs, - log_error_stack=args.log_error_stack, + _init_chat_services( + engine_client, + state, + args, + request_logger, + resolved_chat_template, + enable_chat="generate" in supported_tasks, + enable_batch=True, ) state.openai_serving_responses = ( @@ -1041,37 +1054,6 @@ async def omni_init_app_state( else None ) - _chat_kwargs = dict( - engine_client=engine_client, - models=state.openai_serving_models, - response_role=args.response_role, - online_renderer=state.online_renderer, - request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=args.default_chat_template_kwargs, - trust_request_chat_template=args.trust_request_chat_template, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - enable_prompt_tokens_details=args.enable_prompt_tokens_details, - enable_force_include_usage=args.enable_force_include_usage, - enable_log_outputs=args.enable_log_outputs, - enable_log_deltas=args.enable_log_deltas, - ) - - state.openai_serving_chat = OmniOpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None - state.openai_serving_chat_batch = ( - OmniOpenAIServingChatBatch(**_chat_kwargs) if "generate" in supported_tasks else None - ) - - # Warm up chat template processing to avoid first-request latency - # Upstream f5ffc59b6a removed OpenAIServingChat.warmup() and moved the - # warmup onto the renderer (OnlineRenderer.warmup()); mirror upstream. - state.online_renderer.warmup() - state.openai_serving_completion = ( OpenAIServingCompletion( engine_client,