[Frontend] Scale Omni serving across API processes (#6923)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -42,7 +42,10 @@ from __future__ import annotations
 
 import asyncio
 import json
+import signal
+from argparse import Namespace
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 import pytest
@@ -59,6 +62,66 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+
+@dataclass
+class _FakeParallelConfig:
+    _api_process_rank: int
+
+
+@dataclass
+class _FakeVllmConfig:
+    parallel_config: _FakeParallelConfig
+
+
+def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker):
+    args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}])
+    death_signal = mocker.patch.object(api_server, "set_death_signal")
+
+    with pytest.raises(RuntimeError, match="client configuration"):
+        api_server.run_omni_api_server_worker_proc(
+            "127.0.0.1:8000",
+            mocker.Mock(),
+            args,
+            {"client_count": 2, "client_index": -1},
+        )
+
+    death_signal.assert_called_once_with(signal.SIGTERM)
+
+
+@pytest.mark.asyncio
+async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
+    captured: dict[str, object] = {}
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **kwargs):
+        captured.update(kwargs)
+        yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1)))
+
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI())
+    monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0))
+    monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None)
+
+    async def fake_storage_start() -> None:
+        return None
+
+    monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
+
+    async def fake_serve_http(app, **_kwargs):
+        assert app.state.api_server_count == 2
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    args = _minimal_args(model="dummy")
+    config = {"client_count": 2, "client_index": 1, "stage_addresses": {}}
+
+    await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config)
+    assert captured["client_config"] is config
+
+
 # FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
 # would invent ("HEAD", path) entries for every GET and fail the current census.
 # TODO: extend this set and the expected census if @router.head / @router.options appear.
@@ -626,6 +689,48 @@ async def test_realtime_route_defaults_to_configured_duplex_handler(
     assert calls == [expected_handler]
 
 
+@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime", "/v1/realtime?duplex=1"])
+def test_multi_api_duplex_reconnect_rejected_on_each_worker(path: str) -> None:
+    class SessionHandler:
+        async def handle_realtime_session(self, _websocket) -> None:
+            pytest.fail("Multi-API requests must not access process-local sessions")
+
+    class PendingWarmup:
+        def is_set(self) -> bool:
+            raise AssertionError("Multi-API rejection must not wait for duplex warmup")
+
+    # Separate apps represent workers with independent frontend registries.
+    # Opening on A then reconnecting to B must reject at the same boundary.
+    for worker_index in (0, 1):
+        app = FastAPI()
+        app.state.api_server_count = 2
+        app.state.api_server_index = worker_index
+        app.state.openai_serving_duplex = SessionHandler()
+        app.state.duplex_warmup_done = PendingWarmup()
+        app.include_router(api_server.router)
+        with TestClient(app) as client, client.websocket_connect(path) as websocket:
+            assert websocket.receive_json()["code"] == "multi_api_duplex_unsupported"
+            with pytest.raises(WebSocketDisconnect) as exc:
+                websocket.receive_text()
+            assert exc.value.code == 1008
+
+
+@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"])
+def test_single_api_duplex_session_still_reaches_handler(path: str) -> None:
+    class SessionHandler:
+        async def handle_realtime_session(self, websocket) -> None:
+            await websocket.accept()
+            await websocket.send_json({"type": "session.created"})
+            await websocket.close()
+
+    app = FastAPI()
+    app.state.api_server_count = 1
+    app.state.openai_serving_duplex = SessionHandler()
+    app.include_router(api_server.router)
+    with TestClient(app) as client, client.websocket_connect(path) as websocket:
+        assert websocket.receive_json() == {"type": "session.created"}
+
+
 def test_health_without_engine_returns_stable_unhealthy_response() -> None:
     """Lock ``/health`` with no engine initialized.
 
@@ -715,6 +820,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
     assert exc_info.value.detail == "The model does not support Speech API"
 
 
+@pytest.mark.asyncio
+async def test_multi_api_rejects_runtime_voice_upload() -> None:
+    app = FastAPI()
+    app.state.api_server_count = 2
+    raw_request = _request_for(app, method="POST", path="/v1/audio/voices")
+
+    with pytest.raises(HTTPException) as exc_info:
+        await api_server.upload_voice(
+            raw_request,
+            audio_sample=None,
+            speaker_embedding=None,
+            consent="consent-id",
+            name="probe",
+        )
+
+    assert exc_info.value.status_code == 409
+    assert "process-local frontend state" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_multi_api_rejects_sleep_route() -> None:
+    app = FastAPI()
+    app.state.api_server_count = 2
+    raw_request = _request_for(app, method="POST", path="/v1/omni/sleep")
+
+    with pytest.raises(HTTPException) as exc_info:
+        await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request)
+
+    assert exc_info.value.status_code == 409
+    assert "process-local frontend state" in exc_info.value.detail
+
+
 def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
     """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
 
@@ -937,3 +1074,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
         must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
         must_be_none=_MULTISTAGE_MUST_BE_NONE,
     )
+
+
+@pytest.mark.parametrize("count", [None, 0, -1, "2", True])
+def test_process_local_guard_rejects_unknown_topology(count):
+    app = FastAPI()
+    if count is not None:
+        app.state.api_server_count = count
+    # CLI args cannot substitute for missing/invalid actual runtime state.
+    app.state.args = Namespace(api_server_count=1)
+    with pytest.raises(HTTPException) as exc:
+        api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep")
+    assert exc.value.status_code == 503
+
+
+@pytest.mark.parametrize(
+    "requested,config,expected",
+    [(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)],
+)
+def test_api_topology_resolved_from_launch(requested, config, expected):
+    assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected
+
+
+@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})])
+def test_api_topology_rejects_inconsistent_launch(requested, config):
+    with pytest.raises(ValueError):
+        api_server._resolve_api_server_count(Namespace(api_server_count=requested), config)
+
+
+@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)])
+@pytest.mark.parametrize(
+    "method,path",
+    [
+        ("POST", "/v1/videos"),
+        ("GET", "/v1/videos"),
+        ("GET", "/v1/videos/job"),
+        ("GET", "/v1/videos/job/content"),
+        ("DELETE", "/v1/videos/job"),
+    ],
+)
+def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path):
+    from fastapi.testclient import TestClient
+
+    app = FastAPI()
+    if count is not None:
+        app.state.api_server_count = count
+    app.include_router(api_server.router)
+
+    async def forbidden_form():
+        pytest.fail("video form processing must not run before topology rejection")
+
+    app.dependency_overrides[api_server._parse_video_form] = forbidden_form
+    with TestClient(app) as client:
+        response = client.request(method, path)
+    assert response.status_code == status