anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Frontend] Share regular and duplex chat initialization (#7646)

Browse Source 
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
This commit is contained in:
NATURE authored and GitHub committed 2026-10-07 10:49:34 +00:00
1 parent de253412b0
commit b21df3bcb6
3 files changed
+116 -69
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/openai_api/test_api_server_guards.py
+26 -6
View File
@@ -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])
tests/entrypoints/openai_api/test_duplex_api_server.py
+51 -6
View File
@@ -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)
vllm_omni/entrypoints/openai/api_server.py
+39 -57
View File
@@ -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,