diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 3c44da8b3..e2b5d708a 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/protocol/chat_completion.py b/vllm_omni/entrypoints/openai/protocol/chat_completion.py index d0c83f56f..cdc93f672 100644 --- a/vllm_omni/entrypoints/openai/protocol/chat_completion.py +++ b/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): diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f35008260..46a91fa9f 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 223557aab..8cff9dd3c 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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