[Bugfix] Fix vLLM serve error response imports (#6707)

Signed-off-by: maithilijoshi20 <maithilijoshi2003@gmail.com>
Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>
Co-authored-by: maithilijoshi20 <maithilijoshi2003@gmail.com>

vllm_omni/config/endpoint_policy.py

@@ -2,20 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Endpoint restriction policy for omni pipelines."""
 
+from collections.abc import Set
 from dataclasses import dataclass
 from enum import Enum
 from typing import NamedTuple
 
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-
-# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
-# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
-# serve.exception_handling and re-exports it from the package root.
-try:
-    from vllm.entrypoints.serve import create_error_response
-except ImportError:
-    from vllm.entrypoints.serve.utils.error_response import create_error_response
+from starlette.routing import Route
+from vllm.entrypoints.serve.exception_handling.error_response import create_error_response
 
 
 class RouteTarget(NamedTuple):
@@ -58,6 +53,21 @@ def build_rejection_handler(reason: str):
     return rejection_handler
 
 
+def remove_route_from_app(
+    app: FastAPI,
+    path: str,
+    methods: Set[str],
+) -> None:
+    """Remove routes matching a path and one of the given HTTP methods."""
+    routes_to_remove = [
+        route
+        for route in app.routes
+        if isinstance(route, Route) and route.path == path and route.methods is not None and route.methods & methods
+    ]
+    for route in routes_to_remove:
+        app.routes.remove(route)
+
+
 def shutdown_unsupported_routes(
     app: FastAPI,
     endpoint_restrictions: tuple[EndpointRestriction, ...],
@@ -65,12 +75,10 @@ def shutdown_unsupported_routes(
     """Given an initialized FastAPI server instance and a set of model specific endpoint
     restrictions, remove the restricted routes and patch a handler that returns 400.
     """
-    from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
-
     for end_restrict in endpoint_restrictions:
         capability = end_restrict.capability
         # Remove the route from the app
-        _remove_route_from_app(app, capability.path, capability.methods)
+        remove_route_from_app(app, capability.path, capability.methods)
 
         # Patch the bad request error with the model specific
         # reason for shutting down this endpoint

vllm_omni/entrypoints/openai/api_server.py

@@ -29,7 +29,6 @@ from fastapi.responses import FileResponse, JSONResponse, Response, StreamingRes
 from PIL import Image
 from pydantic import BaseModel, Field
 from starlette.datastructures import State
-from starlette.routing import Route
 from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
@@ -101,7 +100,10 @@ from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
-from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
+from vllm_omni.config.endpoint_policy import (
+    remove_route_from_app,
+    shutdown_unsupported_routes,
+)
 from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
@@ -283,21 +285,6 @@ async def _get_vllm_config(engine_client: EngineClient) -> Any:
     return getattr(engine_client, "vllm_config", None)
 
 
-def _remove_route_from_app(app, path: str, methods: frozenset[str] | None = None):
-    """Remove a route from the app by path and optionally by methods.
-
-    OMNI: used to override upstream /v1/chat/completions with omni behavior.
-    """
-    routes_to_remove = []
-    for route in app.routes:
-        if isinstance(route, Route) and route.path == path:
-            if methods is None or (hasattr(route, "methods") and route.methods & methods):
-                routes_to_remove.append(route)
-
-    for route in routes_to_remove:
-        app.routes.remove(route)
-
-
 def _register_omni_exception_handlers(app) -> None:
     """Override upstream vLLM exception handlers with Omni-aware versions.
 
@@ -523,9 +510,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         app = build_openai_app(args, supported_tasks)
 
         # OMNI: Remove upstream routes that we override with omni-specific handlers
-        _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, "/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
         app.include_router(router)
 
         # OMNI: Override upstream exception handlers with Omni-aware versions
@@ -547,8 +534,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"}))
+            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)