anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
b0942b9c2357ee4a366fa55964d54239eca38fdd
vllm-omni/tests/entrypoints/openai_api/test_api_server_guards.py
T

822 lines
31 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""API-server surface guards for the OpenAI Omni entrypoint.
Purpose
-------
Freeze the *observable* server contract so helper moves (P0.2), route
extraction (P0.3), and later endpoint-family PRs cannot silently change:
* which routes exist,
* how upstream routes are overridden at assembly time,
* unavailable / missing-handler envelopes,
* shared engine-error JSON fields,
* app-state keys wired by ``omni_init_app_state``.
CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
Common wrong changes these tests are meant to catch
----------------------------------------------------
* Route accidentally removed, renamed, method-changed, or double-registered
after a move (census / assembled-app checks).
* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
* Assembly forgets to strip overridden upstream chat / models / profiler
routes, or also deletes unrelated upstream routes.
* Profiler / Omni router not mounted; storage not started; Omni engine
exception handlers not registered.
* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
* Missing-handler WS/HTTP responses change status, JSON shape, or wording
(clients treat these as capability probes).
* Engine error payloads lose ``request_id`` / ``error_stage_id``.
* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
handler as ``None`` while the route is still registered.
Not covered here (add elsewhere if needed)
------------------------------------------
* Deep per-endpoint success behavior (family UTs).
* Byte-identical helper bodies after a move (P0.2 AST integrity).
* Live GPU / speech byte-identity checks.
"""
from __future__ import annotations
import asyncio
import json
from contextlib import asynccontextmanager
from types import SimpleNamespace
import pytest
from fastapi import FastAPI, HTTPException
from fastapi.routing import APIRoute, APIWebSocketRoute
from fastapi.testclient import TestClient
from starlette.datastructures import State
from starlette.requests import Request
from starlette.websockets import WebSocketDisconnect
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
from vllm_omni.entrypoints.openai import api_server
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# 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.
_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
# Exact Omni router census on main. Update intentionally when routes change.
_EXPECTED_ROUTER_ROUTES = {
("POST", "/v1/chat/completions"),
("POST", "/v1/chat/completions/batch"),
("POST", "/v1/audio/speech"),
("POST", "/v1/audio/speech/batch"),
("POST", "/v1/audio/generate"),
("GET", "/v1/audio/voices"),
("POST", "/v1/audio/voices"),
("DELETE", "/v1/audio/voices/{name}"),
("WEBSOCKET", "/v1/audio/speech/stream"),
("WEBSOCKET", "/v1/video/chat/stream"),
("WEBSOCKET", "/v1/realtime/video"),
("WEBSOCKET", "/v1/realtime"),
("WEBSOCKET", "/v1/realtime/robot/openpi"),
("WEBSOCKET", "/v1/duplex"),
("GET", "/health"),
("GET", "/v1/models"),
("POST", "/v1/images/generations"),
("POST", "/v1/images/edits"),
("POST", "/v1/videos"),
("POST", "/v1/videos/sync"),
("GET", "/v1/videos"),
("GET", "/v1/videos/{video_id}"),
("DELETE", "/v1/videos/{video_id}"),
("GET", "/v1/videos/{video_id}/content"),
("POST", "/v1/omni/sleep"),
("POST", "/v1/omni/wakeup"),
}
_EXPECTED_PROFILER_ROUTES = {
("POST", "/start_profile"),
("POST", "/stop_profile"),
}
# App-state keys that must remain available across refactor phases.
# Presence is not enough: "not wired" is written as ``state.X = None``.
_DIFFUSION_APP_STATE_KEYS = {
"engine_client",
"log_stats",
"args",
"sleeping_stages",
"stage_configs",
"vllm_config",
"diffusion_engine",
"openai_serving_models",
"serving_tokenization",
"openai_serving_chat",
"openai_serving_chat_batch",
"openai_serving_speech",
"openai_serving_audio_generate",
"openai_serving_video",
"openai_streaming_video_output",
"openai_serving_duplex",
"openai_streaming_speech",
"openai_streaming_video",
"openai_serving_realtime_robot",
"enable_server_load_tracking",
"server_load_metrics",
}
_DIFFUSION_MUST_BE_NONE = {
"vllm_config",
"serving_tokenization",
"openai_serving_duplex",
"openai_streaming_speech",
"openai_streaming_video",
}
_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
_MULTISTAGE_APP_STATE_KEYS = {
"engine_client",
"log_stats",
"args",
"sleeping_stages",
"stage_configs",
"vllm_config",
"openai_serving_models",
"serving_tokenization",
"online_renderer",
"openai_serving_chat",
"openai_serving_chat_batch",
"openai_serving_speech",
"openai_serving_audio_generate",
"openai_streaming_speech",
"openai_streaming_video",
"openai_serving_duplex",
"openai_serving_realtime",
"openai_serving_video",
"openai_serving_realtime_robot",
"enable_server_load_tracking",
"server_load_metrics",
}
_MULTISTAGE_MUST_BE_NONE = {
"openai_serving_duplex",
"openai_serving_realtime_robot",
}
_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
def _route_entries(routes) -> list[tuple[str, str]]:
entries: list[tuple[str, str]] = []
for route in routes:
if isinstance(route, APIRoute):
for method in sorted((route.methods or set()) & _HTTP_METHODS):
entries.append((method, route.path))
elif isinstance(route, APIWebSocketRoute):
entries.append(("WEBSOCKET", route.path))
return entries
def _route_manifest(routes) -> set[tuple[str, str]]:
return set(_route_entries(routes))
def _assert_unique_route_entries(routes) -> None:
entries = _route_entries(routes)
duplicates = sorted({item for item in entries if entries.count(item) > 1})
assert not duplicates, f"duplicate route registrations: {duplicates}"
def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
return [
route
for route in routes
if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
]
def _single_http_route(routes, method: str, path: str) -> APIRoute:
matches = _matching_http_routes(routes, method, path)
assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
return matches[0]
def _assert_app_state_snapshot(
state,
*,
expected_keys: set[str],
must_be_wired: set[str],
must_be_none: set[str],
) -> None:
assert must_be_wired.isdisjoint(must_be_none)
assert must_be_wired | must_be_none <= expected_keys
present = {key for key in expected_keys if hasattr(state, key)}
assert present == expected_keys
not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
def _minimal_args(**overrides) -> SimpleNamespace:
args = SimpleNamespace(
tool_parser_plugin=None,
reasoning_parser_plugin=None,
reasoning_parser=None,
structured_outputs_config=SimpleNamespace(reasoning_parser=None),
enable_ssl_refresh=False,
host="127.0.0.1",
port=8000,
uvicorn_log_level="info",
disable_uvicorn_access_log=True,
ssl_keyfile=None,
ssl_certfile=None,
ssl_ca_certs=None,
ssl_cert_reqs=0,
ssl_ciphers=None,
h11_max_incomplete_event_size=None,
h11_max_header_count=None,
log_config_file=None,
disable_access_log_for_endpoints=None,
disable_log_stats=True,
enable_log_requests=False,
max_log_len=None,
served_model_name=None,
model="demo-model",
enable_server_load_tracking=False,
log_error_stack=False,
chat_template=None,
chat_template_content_format="auto",
trust_request_chat_template=False,
tool_server=None,
lora_modules=None,
response_role="assistant",
default_chat_template_kwargs=None,
return_tokens_as_token_ids=False,
enable_auto_tool_choice=False,
exclude_tools_when_tool_choice_none=False,
tool_call_parser=None,
enable_prompt_tokens_details=False,
enable_force_include_usage=False,
enable_log_outputs=False,
enable_log_deltas=False,
tokens_only=False,
shutdown_timeout=0,
stage_configs_path=None,
deploy_config=None,
)
for key, value in overrides.items():
setattr(args, key, value)
return args
def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
return Request(
{
"type": "http",
"method": method,
"path": path,
"headers": [],
"app": app,
"state": {},
}
)
class _FakeSocket:
def __init__(self) -> None:
self.closed = False
def close(self) -> None:
self.closed = True
class _FakeEngineClient:
def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
self.stage_configs = stage_configs if stage_configs is not None else []
self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
self.model_config = SimpleNamespace()
self.input_processor = object()
self.renderer = object()
self.errored = False
self.vllm_config = vllm_config
async def get_vllm_config(self):
return self.vllm_config
async def get_supported_tasks(self) -> tuple[str, ...]:
return ("generate",)
async def get_tokenizer(self):
return SimpleNamespace(chat_template="dummy")
async def check_health(self) -> None:
return None
def _marker(name: str):
return SimpleNamespace(name=name)
def test_router_route_manifest_is_explicit_and_complete() -> None:
"""Lock the Omni + profiler router census.
Fails if a route is added/removed/renamed, a method changes, or a
WebSocket path disappears during helper/router moves. Also fails if the
same public route is registered more than once, which a set manifest alone
would collapse.
"""
# TODO(P0.3): retarget this object-level census when routes move onto family
# routers. The assembled-app census is the client-facing lock that must survive.
_assert_unique_route_entries(api_server.router.routes)
_assert_unique_route_entries(api_server.profiler_router.routes)
assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
def test_router_openapi_paths_cover_http_manifest() -> None:
"""Lock that OpenAPI advertises every HTTP census route/method.
Fails if a decorator still registers a route but OpenAPI loses the path
(or method), e.g. after moving route bodies / response models.
"""
app = FastAPI()
app.include_router(api_server.router)
app.include_router(api_server.profiler_router)
paths = app.openapi().get("paths", {})
http_routes = {
(method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
}
for method, path in sorted(http_routes):
assert path in paths, f"missing OpenAPI path {path}"
assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
@pytest.mark.asyncio
async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
"""Lock worker assembly: override, mount, storage, handlers, full census.
Fails if assembly stops replacing upstream chat/batch/models/profiler,
deletes unrelated upstream routes, forgets Omni/profiler mounts, skips
storage start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError``
handlers.
"""
captured: dict[str, object] = {}
def fake_build_openai_app(args, supported_tasks):
app = FastAPI()
@app.post("/v1/chat/completions")
async def upstream_chat():
return {"owner": "upstream"}
@app.post("/v1/chat/completions/batch")
async def upstream_batch():
return {"owner": "upstream"}
@app.get("/v1/models")
async def upstream_models():
return {"owner": "upstream"}
@app.post("/start_profile")
async def upstream_start_profile():
return {"owner": "upstream"}
@app.post("/stop_profile")
async def upstream_stop_profile():
return {"owner": "upstream"}
@app.get("/v1/upstream-only")
async def upstream_only():
return {"owner": "upstream"}
captured["supported_tasks"] = supported_tasks
return app
@asynccontextmanager
async def fake_build_async_omni(*_args, **_kwargs):
yield _FakeEngineClient(
stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
)
async def fake_omni_init_app_state(engine_client, state, args):
state.engine_client = engine_client
state.initialized_by_omni = True
async def fake_storage_start():
captured["storage_started"] = True
async def fake_serve_http(app, **_kwargs):
captured["served_app"] = app
task = asyncio.create_task(asyncio.sleep(0))
await asyncio.sleep(0)
return task
monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
monkeypatch.setattr(
api_server,
"shutdown_unsupported_routes",
lambda app, restrictions: captured.setdefault("restrictions", restrictions),
)
monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
sock = _FakeSocket()
await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
served_wrapper = captured["served_app"]
# Timestamp middleware wraps the FastAPI app.
assert hasattr(served_wrapper, "_inner")
served_app = served_wrapper._inner
routes = served_app.routes
_assert_unique_route_entries(routes)
manifest = _route_manifest(routes)
assert captured["supported_tasks"] == ("generate",)
assert captured["storage_started"] is True
assert captured["restrictions"] == {}
assert sock.closed is True
assert served_app.state.initialized_by_omni is True
# Upstream override: unrelated route kept; Omni owns overridden pairs once.
# Endpoint identity ensures the surviving route is the Omni handler, not
# the fake upstream handler with the same method/path.
assert ("GET", "/v1/upstream-only") in manifest
assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
assert (
_single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
is api_server.create_batch_chat_completion
)
assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
assert _single_http_route(routes, "POST", "/start_profile").endpoint is api_server.start_profile
assert _single_http_route(routes, "POST", "/stop_profile").endpoint is api_server.stop_profile
# Assembled app includes Omni router HTTP/WS surface + profiler.
for item in _EXPECTED_ROUTER_ROUTES:
assert item in manifest, f"assembled app missing {item}"
assert ("POST", "/start_profile") in manifest
assert ("POST", "/stop_profile") in manifest
assert served_app.exception_handlers[EngineGenerateError]
assert served_app.exception_handlers[EngineDeadError]
# Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
openapi_paths = served_app.openapi().get("paths", {})
assert "/v1/chat/completions" in openapi_paths
assert "/v1/models" in openapi_paths
assert "/health" in openapi_paths
@pytest.mark.asyncio
async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
"""Lock outermost timestamp middleware behavior.
Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
start getting stamped / mutated incorrectly.
"""
captured: dict[str, object] = {}
def fake_build_openai_app(args, supported_tasks):
return FastAPI()
@asynccontextmanager
async def fake_build_async_omni(*_args, **_kwargs):
yield _FakeEngineClient()
async def fake_omni_init_app_state(engine_client, state, args):
state.engine_client = engine_client
async def fake_serve_http(app, **_kwargs):
captured["served_app"] = app
task = asyncio.create_task(asyncio.sleep(0))
await asyncio.sleep(0)
return task
monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
middleware = captured["served_app"]
http_scope = {"type": "http", "state": {}}
ws_scope = {"type": "websocket"}
async def _receive():
return {"type": "http.disconnect"}
async def _send(_message):
return None
called = {"http": False, "websocket": False}
async def inner(scope, receive, send):
called[scope["type"]] = True
middleware._inner = inner
await middleware(http_scope, _receive, _send)
await middleware(ws_scope, _receive, _send)
assert called["http"] is True
assert called["websocket"] is True
assert "request_timestamp" in http_scope["state"]
assert isinstance(http_scope["state"]["request_timestamp"], float)
assert "request_timestamp" not in ws_scope.get("state", {})
@pytest.mark.parametrize(
("path", "payload"),
[
("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
(
"/v1/realtime/robot/openpi",
{"type": "error", "error": "Robot policy not available", "code": "unsupported"},
),
("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
],
)
def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
"""Lock missing-handler WebSocket capability probes.
Fails if unavailable JSON changes (keys/wording/code) or the socket stays
open after the error frame — common when WS wrappers move between modules.
"""
app = FastAPI()
app.include_router(api_server.router)
with TestClient(app) as client:
with client.websocket_connect(path) as websocket:
assert websocket.receive_json() == payload
with pytest.raises(WebSocketDisconnect):
websocket.receive_text()
def test_health_without_engine_returns_stable_unhealthy_response() -> None:
"""Lock ``/health`` with no engine initialized.
Fails if status drifts from 503 or the unhealthy payload shape/reason
string changes during health-route extraction.
"""
app = FastAPI()
response = asyncio.run(api_server.health(_request_for(app, path="/health")))
assert response.status_code == 503
assert json.loads(response.body) == {
"status": "unhealthy",
"reason": "No engine initialized",
}
def test_models_without_handler_returns_empty_openai_list() -> None:
"""Lock ``/v1/models`` when ``openai_serving_models`` is unset.
Fails if the empty listing becomes an error, or the OpenAI list envelope
changes (``object`` / ``data``).
"""
app = FastAPI()
response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
assert response.status_code == 200
assert json.loads(response.body) == {"object": "list", "data": []}
def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
"""Lock images generation when the multi-stage engine is missing.
Fails if the 503 detail string changes — clients/tests use it as a
capability signal when bootstrap failed.
"""
app = FastAPI()
raw_request = _request_for(app, method="POST", path="/v1/images/generations")
request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
with pytest.raises(HTTPException) as exc_info:
asyncio.run(api_server.generate_images(request, raw_request))
assert exc_info.value.status_code == 503
assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
"""Lock images generation when chat serving was not wired.
Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
503 detail changes after images/chat ownership splits.
"""
app = FastAPI()
app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
app.state.stage_configs = app.state.engine_client.stage_configs
app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
app.state.openai_serving_chat = None
app.state.args = SimpleNamespace(max_generated_image_size=None)
raw_request = _request_for(app, method="POST", path="/v1/images/generations")
request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
with pytest.raises(HTTPException) as exc_info:
asyncio.run(api_server.generate_images(request, raw_request))
assert exc_info.value.status_code == 503
assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
def test_speech_without_handler_preserves_not_found_http_error() -> None:
"""Lock speech HTTP when ``openai_serving_speech`` is unset.
Fails if missing speech becomes a different status/detail after TTS route
wrappers move out of ``api_server``.
"""
app = FastAPI()
app.state.openai_serving_speech = None
raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
with pytest.raises(HTTPException) as exc_info:
asyncio.run(api_server.create_speech(request, raw_request))
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "The model does not support Speech API"
def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
"""Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
move (P0.2) or exception-handler wiring changes.
Exercises the envelope through the registered handler, not the private
``_create_engine_error_json_response`` helper, so relocating that helper
in P0.2 is not a false red.
"""
app = FastAPI()
# TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
# leaves api_server (red here is a rename, not a surface regression).
api_server._register_omni_exception_handlers(app)
app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
app.state.server = object()
app.state.args = SimpleNamespace(log_error_stack=False)
req = _request_for(app, method="POST", path="/v1/chat/completions")
req.state.request_metadata = SimpleNamespace(request_id="req-123")
monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
exc = EngineGenerateError("boom")
exc.error_stage_id = 2 # type: ignore[attr-defined]
handler = app.exception_handlers[EngineGenerateError]
response = asyncio.run(handler(req, exc))
payload = json.loads(response.body)
assert "error" in payload
assert payload["error"]["request_id"] == "req-123"
assert payload["error"]["error_stage_id"] == 2
assert response.status_code >= 400
def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
"""Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
Fails if registration is skipped or the handler stops returning OpenAI-style
JSON with ``request_id``.
"""
app = FastAPI()
# TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
# leaves api_server (red here is a rename, not a surface regression).
api_server._register_omni_exception_handlers(app)
app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
app.state.server = object()
app.state.args = SimpleNamespace(log_error_stack=False)
monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
handler = app.exception_handlers[EngineDeadError]
req = _request_for(app)
req.state.request_metadata = SimpleNamespace(request_id="dead-1")
response = asyncio.run(handler(req, EngineDeadError()))
payload = json.loads(response.body)
assert payload["error"]["request_id"] == "dead-1"
assert response.status_code >= 400
@pytest.mark.asyncio
async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
"""Lock pure-diffusion ``app.state`` keys after init, including live vs None.
Fails if bootstrap drops keys route owners still read (e.g. video/speech/
models), wires a handler that must stay None, or leaves a required
handler as None.
"""
stage = SimpleNamespace(engine_args={})
engine = _FakeEngineClient(stage_configs=[stage])
def _for_diffusion_factory(label: str):
@classmethod
def _factory(cls, *args, **kwargs):
return _marker(label)
return _factory
monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
monkeypatch.setattr(
api_server.OmniOpenAIServingAudioGenerate,
"for_diffusion",
_for_diffusion_factory("audio_generate"),
)
monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
monkeypatch.setattr(
api_server.ServingRealtimeRobotOpenPI,
"create_policy_server",
classmethod(lambda cls, *a, **k: _marker("openpi")),
)
state = State()
await api_server.omni_init_app_state(engine, state, _minimal_args())
_assert_app_state_snapshot(
state,
expected_keys=_DIFFUSION_APP_STATE_KEYS,
must_be_wired=_DIFFUSION_MUST_BE_WIRED,
must_be_none=_DIFFUSION_MUST_BE_NONE,
)
assert state.diffusion_engine is engine
@pytest.mark.asyncio
async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
"""Lock multi-stage ``app.state`` keys after init, including live vs None.
Fails if chat/speech/video/realtime/tokenization keys disappear or are
left as None after init refactors — classic “route still registered,
handler never wired” failure mode.
"""
engine = _FakeEngineClient(
stage_configs=[object(), object()],
vllm_config=SimpleNamespace(
lora_config=None,
model_config=SimpleNamespace(),
parallel_config=SimpleNamespace(_api_process_rank=0),
),
)
class _FakeModels:
def __init__(self, *args, **kwargs):
self.base_model_paths = kwargs.get("base_model_paths") or []
async def init_static_loras(self):
return None
class _FakeCtor:
def __init__(self, *args, **kwargs):
pass
def warmup(self):
return None
class _FakeSpeech(_FakeCtor):
async def warmup(self):
return None
monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
state = State()
await api_server.omni_init_app_state(engine, state, _minimal_args())
_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,
)
Reference in new issue View Git Blame Copy Permalink