[Bugfix] Restore Omni health route ownership (#6723)

Signed-off-by: qujing226 <qujing226@gmail.com>
Signed-off-by: Eugene Zhao <113281545+qujing226@users.noreply.github.com>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -351,7 +351,7 @@ 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/profiler,
+    Fails if assembly stops replacing upstream chat/batch/models/health/profiler,
     deletes unrelated upstream routes, forgets Omni/profiler mounts, skips
     storage start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError``
     handlers.
@@ -373,6 +373,10 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
         async def upstream_models():
             return {"owner": "upstream"}
 
+        @app.get("/health")
+        async def upstream_health():
+            return {"owner": "upstream"}
+
         @app.post("/start_profile")
         async def upstream_start_profile():
             return {"owner": "upstream"}
@@ -454,6 +458,7 @@ 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, "GET", "/health").endpoint is api_server.health
     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
 

vllm_omni/entrypoints/openai/api_server.py

@@ -513,6 +513,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         remove_route_from_app(app, "/v1/chat/completions", {"POST"})
         remove_route_from_app(app, "/v1/chat/completions/batch", {"POST"})
         remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
+        remove_route_from_app(app, "/health", {"GET"})
         app.include_router(router)
 
         # OMNI: Override upstream exception handlers with Omni-aware versions
@@ -1759,11 +1760,6 @@ async def duplex_websocket(websocket: WebSocket):
 # Health and Model endpoints for diffusion mode
 
 
-# Remove existing health endpoint if present (from vllm imports)
-# to ensure our handler takes precedence
-_remove_route_from_router(router, "/health")
-
-
 @router.get("/health")
 async def health(raw_request: Request) -> JSONResponse:
     """Health check endpoint that works for both LLM and diffusion modes.