[Bugfix] Honor stage selection in online profiler endpoints (#6609)

Signed-off-by: qujing226 <qujing226@gmail.com>

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])

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:

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)