[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>
This commit is contained in:
maithilijoshi20andmaithilijoshi20 authored and GitHub committed 2026-09-02 15:48:01 -07:00
1 parent ca6b4e2112
commit 63ed0fdef4
2 files changed
+28 -33

No files matched your search

+9 -22
View File
@@ -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)