diff --git a/tests/entrypoints/openai/test_profiler_endpoints.py b/tests/entrypoints/openai/test_profiler_endpoints.py new file mode 100644 index 000000000..02e48e19a --- /dev/null +++ b/tests/entrypoints/openai/test_profiler_endpoints.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm_omni.entrypoints.openai import api_server + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("path", "engine_method"), + [ + ("/start_profile", "start_profile"), + ("/stop_profile", "stop_profile"), + ], +) +def test_profiler_endpoint_forwards_selected_stages(path: str, engine_method: str) -> None: + engine_client = SimpleNamespace( + start_profile=AsyncMock(), + stop_profile=AsyncMock(), + ) + + app = FastAPI() + app.include_router(api_server.profiler_router) + app.state.engine_client = engine_client + + response = TestClient(app).post(path, json={"stages": [1]}) + + assert response.status_code == 200 + assert response.json() == {"status": "SUCCESS"} + getattr(engine_client, engine_method).assert_awaited_once_with(stages=[1]) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index ed1c65e13..b9f2e5c71 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """API-server surface guards for the OpenAI Omni entrypoint. Purpose @@ -20,8 +20,8 @@ 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 upstream ``/v1/chat/completions`` (batch) / - ``/v1/models``, or also deletes unrelated upstream routes. +* 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. @@ -351,9 +351,10 @@ def test_router_openapi_paths_cover_http_manifest() -> None: 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, deletes - unrelated upstream routes, forgets Omni/profiler mounts, skips storage - start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers. + 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] = {} @@ -372,6 +373,14 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout 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"} @@ -437,6 +446,8 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout 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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0adc9468e..6838e31ae 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -540,6 +540,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None if _should_enable_profiler_endpoints(stage_configs): logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!") + _remove_route_from_app(app, "/start_profile", frozenset({"POST"})) + _remove_route_from_app(app, "/stop_profile", frozenset({"POST"})) app.include_router(profiler_router) vllm_config = await _get_vllm_config(engine_client)