[Refactor] Update OpenAI API server and related modules for improved structure and functionality

- Refactored import statements in api_server.py to enhance module organization, including changes to the routing and server setup.
- Updated the method for removing routes to ensure compatibility with the new app structure.
- Adjusted the handling of chat completion requests to improve error responses and streamline processing.
- Modified serving modules to align with the new import paths and enhance code clarity.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/entrypoints/openai/api_server.py

@@ -13,7 +13,7 @@ from http import HTTPStatus
 from typing import Any
 
 import vllm.envs as envs
-from fastapi import Depends, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from starlette.datastructures import State
 from starlette.routing import Route
@@ -21,30 +21,30 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving_messages import AnthropicServingMessages
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
-from vllm.entrypoints.openai.api_server import (
-    base,
-    build_app,
-    load_log_config,
-    router,
-    setup_server,
-)
-from vllm.entrypoints.openai.orca_metrics import metrics_header
-from vllm.entrypoints.openai.protocol import (
+from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
+from vllm.entrypoints.openai.api_server import build_app as build_openai_app
+from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
+from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
-    ErrorResponse,
-    ModelCard,
-    ModelList,
-    ModelPermission,
 )
 
 # yapf conflicts with isort for this block
 # yapf: disable
 # yapf: enable
-from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
-from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
-from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
-from vllm.entrypoints.openai.serving_transcription import (
+from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.openai.models.protocol import (
+    BaseModelPath,
+    ModelCard,
+    ModelList,
+    ModelPermission,
+)
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.openai.orca_metrics import metrics_header
+from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
+from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config
+from vllm.entrypoints.openai.translations.serving import (
     OpenAIServingTranscription,
     OpenAIServingTranslation,
 )
@@ -55,7 +55,6 @@ from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
 from vllm.entrypoints.pooling.score.serving import ServingScores
 from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
-from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.utils import (
     load_aware_call,
     process_chat_template,
@@ -84,29 +83,24 @@ from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 
 logger = init_logger(__name__)
+router = APIRouter()
 
 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
 
 
-def _remove_route_from_router(router_obj, path: str, methods: set[str] | None = None):
-    """Remove a route from the router by path and optionally by methods.
+def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
+    """Remove a route from the app by path and optionally by methods.
 
-    This is needed because vllm's api_server registers routes when imported,
-    and we need to override some routes (like /v1/chat/completions) with
-    omni-specific implementations.
+    OMNI: used to override upstream /v1/chat/completions with omni behavior.
     """
     routes_to_remove = []
-    for route in router_obj.routes:
+    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:
-        router_obj.routes.remove(route)
-
-
-# Remove vllm's /v1/chat/completions route so we can register our own omni version
-_remove_route_from_router(router, "/v1/chat/completions", {"POST"})
+        app.routes.remove(route)
 
 
 class _DiffusionServingModels:
@@ -153,7 +147,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     # Add process-specific prefix to stdout and stderr.
     decorate_logs("APIServer")
 
-    listen_address, sock = setup_server(args)
+    listen_address, sock = setup_openai_server(args)
 
     # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -170,7 +164,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
 
     # Load logging config for uvicorn if specified
-    log_config = load_log_config(args.log_config_file)
+    log_config = get_uvicorn_log_config(args)
     if log_config is not None:
         uvicorn_kwargs["log_config"] = log_config
 
@@ -178,7 +172,17 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         args,
         client_config=client_config,
     ) as engine_client:
-        app = build_app(args)
+        supported_tasks: tuple[str, ...]
+        if hasattr(engine_client, "get_supported_tasks"):
+            supported_tasks = tuple(await engine_client.get_supported_tasks())
+        else:
+            supported_tasks = ("generate",)
+        if not supported_tasks:
+            supported_tasks = ("generate",)
+
+        app = build_openai_app(args, supported_tasks)
+        _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
+        app.include_router(router)
 
         await omni_init_app_state(engine_client, app.state, args)
 
@@ -365,6 +369,8 @@ async def omni_init_app_state(
         state.vllm_config = None
         state.diffusion_engine = engine_client
         state.openai_serving_models = _DiffusionServingModels(base_model_paths)
+        # OMNI: tokenization endpoints are not supported in pure diffusion mode.
+        state.openai_serving_tokenization = None
 
         # Use for_diffusion method to create chat handler
         state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion(
@@ -666,14 +672,6 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
 
-# Remove the original /v1/chat/completions route before registering our own
-# This prevents duplicate route registration warnings in FastAPI logs.
-for route in router.routes[:]:
-    if hasattr(route, "path") and route.path == "/v1/chat/completions":
-        router.routes.remove(route)
-        break
-
-
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
@@ -690,7 +688,13 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "")
     handler = Omnichat(raw_request)
     if handler is None:
-        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
+        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Chat Completions API",
+            )
+        return base_server.create_error_response(message="The model does not support Chat Completions API")
     try:
         generator = await handler.create_chat_completion(request, raw_request)
     except Exception as e:
@@ -756,7 +760,13 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
 async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request):
     handler = Omnispeech(raw_request)
     if handler is None:
-        return base(raw_request).create_error_response(message="The model does not support Speech API")
+        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Speech API",
+            )
+        return base_server.create_error_response(message="The model does not support Speech API")
     try:
         return await handler.create_speech(request, raw_request)
     except Exception as e:

vllm_omni/entrypoints/openai/protocol/chat_completion.py

@@ -1,6 +1,4 @@
-from vllm.entrypoints.openai.protocol import (
-    ChatCompletionStreamResponse,
-)
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionStreamResponse
 
 
 class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse):

vllm_omni/entrypoints/openai/serving_chat.py

@@ -30,11 +30,7 @@ from vllm.entrypoints.chat_utils import (
     make_tool_call_id,
     resolve_chat_template_content_format,
 )
-from vllm.entrypoints.openai.parser.harmony_utils import (
-    get_streamable_parser_for_assistant,
-    parse_chat_output,
-)
-from vllm.entrypoints.openai.protocol import (
+from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
     ChatCompletionResponse,
@@ -42,6 +38,9 @@ from vllm.entrypoints.openai.protocol import (
     ChatCompletionResponseStreamChoice,
     ChatCompletionStreamResponse,
     ChatMessage,
+)
+from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
+from vllm.entrypoints.openai.engine.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
     DeltaToolCall,
@@ -51,15 +50,15 @@ from vllm.entrypoints.openai.protocol import (
     FunctionDefinition,
     PromptTokenUsageInfo,
     RequestResponseMetadata,
-    ResponsesRequest,
     ToolCall,
     UsageInfo,
 )
-from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
-from vllm.entrypoints.openai.serving_engine import (
-    ChatLikeRequest,
-    clamp_prompt_logprobs,
+from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs
+from vllm.entrypoints.openai.parser.harmony_utils import (
+    get_streamable_parser_for_assistant,
+    parse_chat_output,
 )
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
 from vllm.entrypoints.utils import should_include_usage
 from vllm.inputs.data import PromptType, TokensPrompt
@@ -167,7 +166,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             model_name = self.models.model_name(lora_request)
 
-            tokenizer = await self.engine_client.get_tokenizer()
+            renderer = self.renderer
+            tokenizer = renderer.get_tokenizer()
+            if tokenizer is None:
+                tokenizer = await self.engine_client.get_tokenizer()
 
             tool_parser = self.tool_parser
 
@@ -224,7 +226,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     engine_prompts,
                 ) = await self._preprocess_chat(
                     request,
-                    tokenizer,
+                    renderer,
                     request.messages,
                     chat_template=request.chat_template or self.chat_template,
                     chat_template_content_format=self.chat_template_content_format,
@@ -323,7 +325,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     async def _preprocess_chat(
         self,
         request: ChatLikeRequest | ResponsesRequest,
-        tokenizer: TokenizerLike,
+        renderer: Any,
         messages: list[ChatCompletionMessageParam],
         chat_template: str | None,
         chat_template_content_format: ChatTemplateContentFormatOption,
@@ -341,6 +343,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         list[TokensPrompt],
     ]:
         model_config = self.model_config
+        tokenizer = renderer.get_tokenizer() if renderer is not None else None
 
         resolved_content_format = resolve_chat_template_content_format(
             chat_template,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -3,7 +3,7 @@ from typing import Any
 
 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.openai.serving_engine import OpenAIServing
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid