@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 76139d0801

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 13:33 UTC

[Frontend] Frontend will only attach supported tasks corresponding entrypoints. (#33139)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

wang.yuqi committed Jan 27, 2026 at 20:15 UTC 76139d0801e95bf6c54352650a6f42e181b510be

10 files changed +461 -339

tests/entrypoints/openai/test_transcription_validation_whisper.py

+8 -12

@@ -112,18 +112,14 @@ async def test_long_audio_request(mary_had_lamb, whisper_client):

112 @pytest.mark.asyncio

113 async def test_completion_endpoints(whisper_client):

114 # text to text model

115 - res = await whisper_client.chat.completions.create(

116 - model=MODEL_NAME,

117 - messages=[{"role": "system", "content": "You are a helpful assistant."}],

118 - )

119 - err = res.error

120 - assert err["code"] == 400

121 - assert err["message"] == "The model does not support Chat Completions API"

122 -

123 - res = await whisper_client.completions.create(model=MODEL_NAME, prompt="Hello")

124 - err = res.error

125 - assert err["code"] == 400

126 - assert err["message"] == "The model does not support Completions API"

115 + with pytest.raises(openai.NotFoundError):

116 + await whisper_client.chat.completions.create(

117 + model=MODEL_NAME,

118 + messages=[{"role": "system", "content": "You are a helpful assistant."}],

119 + )

120 +

121 + with pytest.raises(openai.NotFoundError):

122 + await whisper_client.completions.create(model=MODEL_NAME, prompt="Hello")

123

124

125 @pytest.mark.asyncio

tests/entrypoints/openai/test_translation_validation.py

+6 -6

@@ -9,6 +9,7 @@ import json

9 import httpx

10 import librosa

11 import numpy as np

12 +import openai

13 import pytest

14 import pytest_asyncio

15 import soundfile as sf

@@ -52,12 +53,11 @@ async def test_non_asr_model(foscolo, rocm_aiter_fa_attention):

53 model_name, _get_server_args(rocm_aiter_fa_attention)

54 ) as remote_server:

55 client = remote_server.get_async_client()

55 - res = await client.audio.translations.create(

56 - model=model_name, file=foscolo, temperature=0.0

57 - )

58 - err = res.error

59 - assert err["code"] == 400 and not res.text

60 - assert err["message"] == "The model does not support Translations API"

56 +

57 + with pytest.raises(openai.NotFoundError):

58 + await client.audio.translations.create(

59 + model=model_name, file=foscolo, temperature=0.0

60 + )

63 @pytest.mark.asyncio

tests/entrypoints/pooling/classify/test_online.py

+2 -2

@@ -401,7 +401,7 @@ async def test_score(server: RemoteOpenAIServer, model_name: str):

401 "documents": "pong",

402 },

403 )

404 - assert response.json()["error"]["type"] == "BadRequestError"

404 + assert response.json()["detail"] == "Not Found"

405

406

407 @pytest.mark.asyncio

@@ -416,7 +416,7 @@ async def test_rerank(server: RemoteOpenAIServer, model_name: str):

416 "documents": ["pong"],

417 },

418 )

419 - assert response.json()["error"]["type"] == "BadRequestError"

419 + assert response.json()["detail"] == "Not Found"

420

421

422 @pytest.mark.asyncio

vllm/entrypoints/openai/api_server.py

+45 -175

@@ -33,14 +33,10 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send

33 import vllm.envs as envs

34 from vllm.engine.arg_utils import AsyncEngineArgs

35 from vllm.engine.protocol import EngineClient

36 -from vllm.entrypoints.anthropic.serving import AnthropicServingMessages

36 from vllm.entrypoints.chat_utils import load_chat_template

37 from vllm.entrypoints.launcher import serve_http

38 from vllm.entrypoints.logger import RequestLogger

40 -from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer

41 -from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat

39 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args

43 -from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion

40 from vllm.entrypoints.openai.engine.protocol import (

41 ErrorInfo,

42 ErrorResponse,

@@ -50,12 +46,6 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath

46 from vllm.entrypoints.openai.models.serving import (

47 OpenAIServingModels,

48 )

53 -from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses

54 -from vllm.entrypoints.openai.translations.serving import (

55 - OpenAIServingTranscription,

56 - OpenAIServingTranslation,

57 -)

58 -from vllm.entrypoints.serve.disagg.serving import ServingTokens

49 from vllm.entrypoints.serve.elastic_ep.middleware import (

50 ScalingMiddleware,

51 )

@@ -70,6 +60,7 @@ from vllm.entrypoints.utils import (

60 from vllm.exceptions import VLLMValidationError

61 from vllm.logger import init_logger

62 from vllm.reasoning import ReasoningParserManager

63 +from vllm.tasks import POOLING_TASKS, SupportedTask

64 from vllm.tool_parsers import ToolParserManager

65 from vllm.usage.usage_lib import UsageContext

66 from vllm.utils.argparse_utils import FlexibleArgumentParser

@@ -513,7 +504,7 @@ def _log_non_streaming_response(response_body: list) -> None:

504 logger.info("response_body={<binary_data>}")

505

506

516 -def build_app(args: Namespace) -> FastAPI:

507 +def build_app(args: Namespace, supported_tasks: tuple["SupportedTask", ...]) -> FastAPI:

508 if args.disable_fastapi_docs:

509 app = FastAPI(

510 openapi_url=None, docs_url=None, redoc_url=None, lifespan=lifespan

@@ -523,52 +514,44 @@ def build_app(args: Namespace) -> FastAPI:

514 else:

515 app = FastAPI(lifespan=lifespan)

516 app.state.args = args

517 + app.include_router(router)

518 +

519 from vllm.entrypoints.serve import register_vllm_serve_api_routers

520

521 register_vllm_serve_api_routers(app)

529 - from vllm.entrypoints.openai.chat_completion.api_router import (

530 - attach_router as register_chat_api_router,

531 - )

532 -

533 - register_chat_api_router(app)

522

535 - from vllm.entrypoints.openai.responses.api_router import (

536 - attach_router as register_responses_api_router,

537 - )

538 -

539 - register_responses_api_router(app)

540 - from vllm.entrypoints.openai.translations.api_router import (

541 - attach_router as register_translations_api_router,

523 + from vllm.entrypoints.openai.models.api_router import (

524 + attach_router as register_models_api_router,

525 )

526

544 - register_translations_api_router(app)

527 + register_models_api_router(app)

528

546 - from vllm.entrypoints.openai.completion.api_router import (

547 - attach_router as register_completion_api_router,

529 + from vllm.entrypoints.sagemaker.api_router import (

530 + attach_router as register_sagemaker_api_router,

531 )

532

550 - register_completion_api_router(app)

551 - from vllm.entrypoints.anthropic.api_router import (

552 - attach_router as register_anthropic_api_router,

553 - )

533 + register_sagemaker_api_router(app, supported_tasks)

534

555 - register_anthropic_api_router(app)

556 - from vllm.entrypoints.openai.models.api_router import (

557 - attach_router as register_models_api_router,

558 - )

535 + if "generate" in supported_tasks:

536 + from vllm.entrypoints.openai.generate.api_router import (

537 + register_generate_api_routers,

538 + )

539

560 - register_models_api_router(app)

561 - from vllm.entrypoints.sagemaker.routes import register_sagemaker_routes

540 + register_generate_api_routers(app)

541

563 - register_sagemaker_routes(router)

564 - app.include_router(router)

542 + if "transcription" in supported_tasks:

543 + from vllm.entrypoints.openai.translations.api_router import (

544 + attach_router as register_translations_api_router,

545 + )

546

566 - app.root_path = args.root_path

547 + register_translations_api_router(app)

548

568 - from vllm.entrypoints.pooling import register_pooling_api_routers

549 + if any(task in POOLING_TASKS for task in supported_tasks):

550 + from vllm.entrypoints.pooling import register_pooling_api_routers

551

570 - register_pooling_api_routers(app)

552 + register_pooling_api_routers(app, supported_tasks)

553

554 + app.root_path = args.root_path

555 app.add_middleware(

556 CORSMiddleware,

557 allow_origins=args.allowed_origins,

@@ -673,6 +656,7 @@ async def init_app_state(

656 engine_client: EngineClient,

657 state: State,

658 args: Namespace,

659 + supported_tasks: tuple["SupportedTask", ...],

660 ) -> None:

661 vllm_config = engine_client.vllm_config

662

@@ -694,28 +678,9 @@ async def init_app_state(

678 state.log_stats = not args.disable_log_stats

679 state.vllm_config = vllm_config

680 state.args = args

697 - supported_tasks = await engine_client.get_supported_tasks()

698 - logger.info("Supported tasks: %s", supported_tasks)

699 -

681 resolved_chat_template = load_chat_template(args.chat_template)

682

702 - if args.tool_server == "demo":

703 - tool_server: ToolServer | None = DemoToolServer()

704 - assert isinstance(tool_server, DemoToolServer)

705 - await tool_server.init_and_validate()

706 - elif args.tool_server:

707 - tool_server = MCPToolServer()

708 - await tool_server.add_tool_server(args.tool_server)

709 - else:

710 - tool_server = None

711 -

683 # Merge default_mm_loras into the static lora_modules

713 - default_mm_loras = (

714 - vllm_config.lora_config.default_mm_loras

715 - if vllm_config.lora_config is not None

716 - else {}

717 - )

718 -

684 default_mm_loras = (

685 vllm_config.lora_config.default_mm_loras

686 if vllm_config.lora_config is not None

@@ -729,66 +694,6 @@ async def init_app_state(

694 lora_modules=lora_modules,

695 )

696 await state.openai_serving_models.init_static_loras()

732 - state.openai_serving_responses = (

733 - OpenAIServingResponses(

734 - engine_client,

735 - state.openai_serving_models,

736 - request_logger=request_logger,

737 - chat_template=resolved_chat_template,

738 - chat_template_content_format=args.chat_template_content_format,

739 - return_tokens_as_token_ids=args.return_tokens_as_token_ids,

740 - enable_auto_tools=args.enable_auto_tool_choice,

741 - tool_parser=args.tool_call_parser,

742 - tool_server=tool_server,

743 - reasoning_parser=args.structured_outputs_config.reasoning_parser,

744 - enable_prompt_tokens_details=args.enable_prompt_tokens_details,

745 - enable_force_include_usage=args.enable_force_include_usage,

746 - enable_log_outputs=args.enable_log_outputs,

747 - log_error_stack=args.log_error_stack,

748 - )

749 - if "generate" in supported_tasks

750 - else None

751 - )

752 - state.openai_serving_chat = (

753 - OpenAIServingChat(

754 - engine_client,

755 - state.openai_serving_models,

756 - args.response_role,

757 - request_logger=request_logger,

758 - chat_template=resolved_chat_template,

759 - chat_template_content_format=args.chat_template_content_format,

760 - default_chat_template_kwargs=args.default_chat_template_kwargs,

761 - trust_request_chat_template=args.trust_request_chat_template,

762 - return_tokens_as_token_ids=args.return_tokens_as_token_ids,

763 - enable_auto_tools=args.enable_auto_tool_choice,

764 - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,

765 - tool_parser=args.tool_call_parser,

766 - reasoning_parser=args.structured_outputs_config.reasoning_parser,

767 - enable_prompt_tokens_details=args.enable_prompt_tokens_details,

768 - enable_force_include_usage=args.enable_force_include_usage,

769 - enable_log_outputs=args.enable_log_outputs,

770 - enable_log_deltas=args.enable_log_deltas,

771 - log_error_stack=args.log_error_stack,

772 - )

773 - if "generate" in supported_tasks

774 - else None

775 - )

776 - # Warm up chat template processing to avoid first-request latency

777 - if state.openai_serving_chat is not None:

778 - await state.openai_serving_chat.warmup()

779 - state.openai_serving_completion = (

780 - OpenAIServingCompletion(

781 - engine_client,

782 - state.openai_serving_models,

783 - request_logger=request_logger,

784 - return_tokens_as_token_ids=args.return_tokens_as_token_ids,

785 - enable_prompt_tokens_details=args.enable_prompt_tokens_details,

786 - enable_force_include_usage=args.enable_force_include_usage,

787 - log_error_stack=args.log_error_stack,

788 - )

789 - if "generate" in supported_tasks

790 - else None

791 - )

697 state.openai_serving_tokenization = OpenAIServingTokenization(

698 engine_client,

699 state.openai_serving_models,

@@ -798,64 +703,27 @@ async def init_app_state(

703 trust_request_chat_template=args.trust_request_chat_template,

704 log_error_stack=args.log_error_stack,

705 )

801 - state.openai_serving_transcription = (

802 - OpenAIServingTranscription(

803 - engine_client,

804 - state.openai_serving_models,

805 - request_logger=request_logger,

806 - log_error_stack=args.log_error_stack,

807 - enable_force_include_usage=args.enable_force_include_usage,

808 - )

809 - if "transcription" in supported_tasks

810 - else None

811 - )

812 - state.openai_serving_translation = (

813 - OpenAIServingTranslation(

814 - engine_client,

815 - state.openai_serving_models,

816 - request_logger=request_logger,

817 - log_error_stack=args.log_error_stack,

818 - enable_force_include_usage=args.enable_force_include_usage,

706 +

707 + if "generate" in supported_tasks:

708 + from vllm.entrypoints.openai.generate.api_router import init_generate_state

709 +

710 + await init_generate_state(

711 + engine_client, state, args, request_logger, supported_tasks

712 )

820 - if "transcription" in supported_tasks

821 - else None

822 - )

823 - state.anthropic_serving_messages = (

824 - AnthropicServingMessages(

825 - engine_client,

826 - state.openai_serving_models,

827 - args.response_role,

828 - request_logger=request_logger,

829 - chat_template=resolved_chat_template,

830 - chat_template_content_format=args.chat_template_content_format,

831 - return_tokens_as_token_ids=args.return_tokens_as_token_ids,

832 - enable_auto_tools=args.enable_auto_tool_choice,

833 - tool_parser=args.tool_call_parser,

834 - reasoning_parser=args.structured_outputs_config.reasoning_parser,

835 - enable_prompt_tokens_details=args.enable_prompt_tokens_details,

836 - enable_force_include_usage=args.enable_force_include_usage,

713 +

714 + if "transcription" in supported_tasks:

715 + from vllm.entrypoints.openai.translations.api_router import (

716 + init_transcription_state,

717 )

838 - if "generate" in supported_tasks

839 - else None

840 - )

841 - state.serving_tokens = (

842 - ServingTokens(

843 - engine_client,

844 - state.openai_serving_models,

845 - request_logger=request_logger,

846 - return_tokens_as_token_ids=args.return_tokens_as_token_ids,

847 - log_error_stack=args.log_error_stack,

848 - enable_prompt_tokens_details=args.enable_prompt_tokens_details,

849 - enable_log_outputs=args.enable_log_outputs,

850 - force_no_detokenize=args.tokens_only,

718 +

719 + init_transcription_state(

720 + engine_client, state, args, request_logger, supported_tasks

721 )

852 - if "generate" in supported_tasks

853 - else None

854 - )

722

856 - from vllm.entrypoints.pooling import init_pooling_state

723 + if any(task in POOLING_TASKS for task in supported_tasks):

724 + from vllm.entrypoints.pooling import init_pooling_state

725

858 - await init_pooling_state(engine_client, state, args)

726 + init_pooling_state(engine_client, state, args, request_logger, supported_tasks)

727

728 state.enable_server_load_tracking = args.enable_server_load_tracking

729 state.server_load_metrics = 0

@@ -972,9 +840,11 @@ async def run_server_worker(

840 args,

841 client_config=client_config,

842 ) as engine_client:

975 - app = build_app(args)

843 + supported_tasks = await engine_client.get_supported_tasks()

844 + logger.info("Supported tasks: %s", supported_tasks)

845

977 - await init_app_state(engine_client, app.state, args)

846 + app = build_app(args, supported_tasks)

847 + await init_app_state(engine_client, app.state, args, supported_tasks)

848

849 logger.info(

850 "Starting vLLM API server %d on %s",

vllm/entrypoints/openai/generate/__init__.py

vllm/entrypoints/openai/generate/api_router.py new

+166

@@ -0,0 +1,166 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import TYPE_CHECKING

4 +

5 +from fastapi import FastAPI

6 +

7 +if TYPE_CHECKING:

8 + from argparse import Namespace

9 +

10 + from starlette.datastructures import State

11 +

12 + from vllm.engine.protocol import EngineClient

13 + from vllm.entrypoints.logger import RequestLogger

14 + from vllm.tasks import SupportedTask

15 +else:

16 + RequestLogger = object

17 +

18 +

19 +def register_generate_api_routers(app: FastAPI):

20 + from vllm.entrypoints.openai.chat_completion.api_router import (

21 + attach_router as register_chat_api_router,

22 + )

23 +

24 + register_chat_api_router(app)

25 +

26 + from vllm.entrypoints.openai.responses.api_router import (

27 + attach_router as register_responses_api_router,

28 + )

29 +

30 + register_responses_api_router(app)

31 +

32 + from vllm.entrypoints.openai.completion.api_router import (

33 + attach_router as register_completion_api_router,

34 + )

35 +

36 + register_completion_api_router(app)

37 +

38 + from vllm.entrypoints.anthropic.api_router import (

39 + attach_router as register_anthropic_api_router,

40 + )

41 +

42 + register_anthropic_api_router(app)

43 +

44 +

45 +async def init_generate_state(

46 + engine_client: "EngineClient",

47 + state: "State",

48 + args: "Namespace",

49 + request_logger: RequestLogger | None,

50 + supported_tasks: tuple["SupportedTask", ...],

51 +):

52 + from vllm.entrypoints.anthropic.serving import AnthropicServingMessages

53 + from vllm.entrypoints.chat_utils import load_chat_template

54 + from vllm.entrypoints.mcp.tool_server import (

55 + DemoToolServer,

56 + MCPToolServer,

57 + ToolServer,

58 + )

59 + from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat

60 + from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion

61 + from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses

62 + from vllm.entrypoints.serve.disagg.serving import ServingTokens

63 +

64 + if args.tool_server == "demo":

65 + tool_server: ToolServer | None = DemoToolServer()

66 + assert isinstance(tool_server, DemoToolServer)

67 + await tool_server.init_and_validate()

68 + elif args.tool_server:

69 + tool_server = MCPToolServer()

70 + await tool_server.add_tool_server(args.tool_server)

71 + else:

72 + tool_server = None

73 + resolved_chat_template = load_chat_template(args.chat_template)

74 +

75 + state.openai_serving_responses = (

76 + OpenAIServingResponses(

77 + engine_client,

78 + state.openai_serving_models,

79 + request_logger=request_logger,

80 + chat_template=resolved_chat_template,

81 + chat_template_content_format=args.chat_template_content_format,

82 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,

83 + enable_auto_tools=args.enable_auto_tool_choice,

84 + tool_parser=args.tool_call_parser,

85 + tool_server=tool_server,

86 + reasoning_parser=args.structured_outputs_config.reasoning_parser,

87 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,

88 + enable_force_include_usage=args.enable_force_include_usage,

89 + enable_log_outputs=args.enable_log_outputs,

90 + log_error_stack=args.log_error_stack,

91 + )

92 + if "generate" in supported_tasks

93 + else None

94 + )

95 + state.openai_serving_chat = (

96 + OpenAIServingChat(

97 + engine_client,

98 + state.openai_serving_models,

99 + args.response_role,

100 + request_logger=request_logger,

101 + chat_template=resolved_chat_template,

102 + chat_template_content_format=args.chat_template_content_format,

103 + default_chat_template_kwargs=args.default_chat_template_kwargs,

104 + trust_request_chat_template=args.trust_request_chat_template,

105 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,

106 + enable_auto_tools=args.enable_auto_tool_choice,

107 + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,

108 + tool_parser=args.tool_call_parser,

109 + reasoning_parser=args.structured_outputs_config.reasoning_parser,

110 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,

111 + enable_force_include_usage=args.enable_force_include_usage,

112 + enable_log_outputs=args.enable_log_outputs,

113 + enable_log_deltas=args.enable_log_deltas,

114 + log_error_stack=args.log_error_stack,

115 + )

116 + if "generate" in supported_tasks

117 + else None

118 + )

119 + # Warm up chat template processing to avoid first-request latency

120 + if state.openai_serving_chat is not None:

121 + await state.openai_serving_chat.warmup()

122 + state.openai_serving_completion = (

123 + OpenAIServingCompletion(

124 + engine_client,

125 + state.openai_serving_models,

126 + request_logger=request_logger,

127 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,

128 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,

129 + enable_force_include_usage=args.enable_force_include_usage,

130 + log_error_stack=args.log_error_stack,

131 + )

132 + if "generate" in supported_tasks

133 + else None

134 + )

135 + state.anthropic_serving_messages = (

136 + AnthropicServingMessages(

137 + engine_client,

138 + state.openai_serving_models,

139 + args.response_role,

140 + request_logger=request_logger,

141 + chat_template=resolved_chat_template,

142 + chat_template_content_format=args.chat_template_content_format,

143 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,

144 + enable_auto_tools=args.enable_auto_tool_choice,

145 + tool_parser=args.tool_call_parser,

146 + reasoning_parser=args.structured_outputs_config.reasoning_parser,

147 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,

148 + enable_force_include_usage=args.enable_force_include_usage,

149 + )

150 + if "generate" in supported_tasks

151 + else None

152 + )

153 + state.serving_tokens = (

154 + ServingTokens(

155 + engine_client,

156 + state.openai_serving_models,

157 + request_logger=request_logger,

158 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,

159 + log_error_stack=args.log_error_stack,

160 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,

161 + enable_log_outputs=args.enable_log_outputs,

162 + force_no_detokenize=args.tokens_only,

163 + )

164 + if "generate" in supported_tasks

165 + else None

166 + )

vllm/entrypoints/openai/translations/api_router.py

+43 -1

@@ -3,7 +3,7 @@

5 from http import HTTPStatus

6 -from typing import Annotated

6 +from typing import TYPE_CHECKING, Annotated

8 from fastapi import APIRouter, FastAPI, Form, Request

9 from fastapi.responses import JSONResponse, StreamingResponse

@@ -25,6 +25,17 @@ from vllm.entrypoints.utils import (

25 )

26 from vllm.logger import init_logger

28 +if TYPE_CHECKING:

29 + from argparse import Namespace

30 +

31 + from starlette.datastructures import State

32 +

33 + from vllm.engine.protocol import EngineClient

34 + from vllm.entrypoints.logger import RequestLogger

35 + from vllm.tasks import SupportedTask

36 +else:

37 + RequestLogger = object

38 +

39 logger = init_logger(__name__)

41 router = APIRouter()

@@ -115,3 +126,34 @@ async def create_translations(

126

127 def attach_router(app: FastAPI):

128 app.include_router(router)

129 +

130 +

131 +def init_transcription_state(

132 + engine_client: "EngineClient",

133 + state: "State",

134 + args: "Namespace",

135 + request_logger: RequestLogger | None,

136 + supported_tasks: tuple["SupportedTask", ...],

137 +):

138 + state.openai_serving_transcription = (

139 + OpenAIServingTranscription(

140 + engine_client,

141 + state.openai_serving_models,

142 + request_logger=request_logger,

143 + log_error_stack=args.log_error_stack,

144 + enable_force_include_usage=args.enable_force_include_usage,

145 + )

146 + if "transcription" in supported_tasks

147 + else None

148 + )

149 + state.openai_serving_translation = (

150 + OpenAIServingTranslation(

151 + engine_client,

152 + state.openai_serving_models,

153 + request_logger=request_logger,

154 + log_error_stack=args.log_error_stack,

155 + enable_force_include_usage=args.enable_force_include_usage,

156 + )

157 + if "transcription" in supported_tasks

158 + else None

159 + )

vllm/entrypoints/pooling/__init__.py

+31 -17

@@ -11,40 +11,54 @@ if TYPE_CHECKING:

11 from starlette.datastructures import State

13 from vllm.engine.protocol import EngineClient

14 + from vllm.entrypoints.logger import RequestLogger

15 + from vllm.tasks import SupportedTask

16 +else:

17 + RequestLogger = object

18 + SupportedTask = object

16 -def register_pooling_api_routers(app: FastAPI):

17 - from vllm.entrypoints.pooling.classify.api_router import router as classify_router

18 - from vllm.entrypoints.pooling.embed.api_router import router as embed_router

21 +def register_pooling_api_routers(

22 + app: FastAPI, supported_tasks: tuple["SupportedTask", ...]

23 +):

24 from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router

20 - from vllm.entrypoints.pooling.score.api_router import router as score_router

22 - app.include_router(classify_router)

23 - app.include_router(embed_router)

24 - app.include_router(score_router)

26 app.include_router(pooling_router)

28 + if "classify" in supported_tasks:

29 + from vllm.entrypoints.pooling.classify.api_router import (

30 + router as classify_router,

31 + )

32 +

33 + app.include_router(classify_router)

34 +

35 + if "embed" in supported_tasks:

36 + from vllm.entrypoints.pooling.embed.api_router import router as embed_router

37 +

38 + app.include_router(embed_router)

28 -async def init_pooling_state(

29 - engine_client: "EngineClient", state: "State", args: "Namespace"

40 + if "score" in supported_tasks or "embed" in supported_tasks:

41 + from vllm.entrypoints.pooling.score.api_router import router as score_router

42 +

43 + app.include_router(score_router)

44 +

45 +

46 +def init_pooling_state(

47 + engine_client: "EngineClient",

48 + state: "State",

49 + args: "Namespace",

50 + request_logger: RequestLogger | None,

51 + supported_tasks: tuple["SupportedTask", ...],

52 ):

53 from vllm.entrypoints.chat_utils import load_chat_template

32 - from vllm.entrypoints.logger import RequestLogger

54 from vllm.entrypoints.pooling.classify.serving import ServingClassification

55 from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding

56 from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling

57 from vllm.entrypoints.pooling.score.serving import ServingScores

58 from vllm.tasks import POOLING_TASKS

39 - supported_tasks = await engine_client.get_supported_tasks()

40 -

60 resolved_chat_template = load_chat_template(args.chat_template)

43 - if args.enable_log_requests:

44 - request_logger = RequestLogger(max_log_len=args.max_log_len)

45 - else:

46 - request_logger = None

47 -

62 state.openai_serving_pooling = (

63 (

64 OpenAIServingPooling(

vllm/entrypoints/sagemaker/api_router.py new

+160

@@ -0,0 +1,160 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import json

4 +from collections.abc import Awaitable, Callable

5 +from http import HTTPStatus

6 +from typing import Any

7 +

8 +import model_hosting_container_standards.sagemaker as sagemaker_standards

9 +import pydantic

10 +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request

11 +from fastapi.responses import JSONResponse, Response

12 +

13 +from vllm.entrypoints.openai.api_server import base

14 +from vllm.entrypoints.openai.engine.protocol import ErrorResponse

15 +from vllm.entrypoints.openai.engine.serving import OpenAIServing

16 +from vllm.entrypoints.openai.utils import validate_json_request

17 +from vllm.entrypoints.serve.instrumentator.health import health

18 +from vllm.tasks import POOLING_TASKS, SupportedTask

19 +

20 +# TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers

21 +# (requires typing_extensions >= 4.13)

22 +RequestType = Any

23 +GetHandlerFn = Callable[[Request], OpenAIServing | None]

24 +EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]

25 +

26 +

27 +def get_invocation_types(supported_tasks: tuple["SupportedTask", ...]):

28 + # NOTE: Items defined earlier take higher priority

29 + INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []

30 +

31 + if "generate" in supported_tasks:

32 + from vllm.entrypoints.openai.chat_completion.api_router import (

33 + chat,

34 + create_chat_completion,

35 + )

36 + from vllm.entrypoints.openai.chat_completion.protocol import (

37 + ChatCompletionRequest,

38 + )

39 + from vllm.entrypoints.openai.completion.api_router import (

40 + completion,

41 + create_completion,

42 + )

43 + from vllm.entrypoints.openai.completion.protocol import CompletionRequest

44 +

45 + INVOCATION_TYPES += [

46 + (ChatCompletionRequest, (chat, create_chat_completion)),

47 + (CompletionRequest, (completion, create_completion)),

48 + ]

49 +

50 + if "embed" in supported_tasks:

51 + from vllm.entrypoints.pooling.embed.api_router import (

52 + create_embedding,

53 + embedding,

54 + )

55 + from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest

56 +

57 + INVOCATION_TYPES += [

58 + (EmbeddingRequest, (embedding, create_embedding)),

59 + ]

60 +

61 + if "classify" in supported_tasks:

62 + from vllm.entrypoints.pooling.classify.api_router import (

63 + classify,

64 + create_classify,

65 + )

66 + from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest

67 +

68 + INVOCATION_TYPES += [

69 + (ClassificationRequest, (classify, create_classify)),

70 + ]

71 +

72 + if "score" in supported_tasks:

73 + from vllm.entrypoints.pooling.score.api_router import do_rerank, rerank

74 + from vllm.entrypoints.pooling.score.protocol import RerankRequest

75 +

76 + INVOCATION_TYPES += [

77 + (RerankRequest, (rerank, do_rerank)),

78 + ]

79 +

80 + if "score" in supported_tasks or "embed" in supported_tasks:

81 + from vllm.entrypoints.pooling.score.api_router import create_score, score

82 + from vllm.entrypoints.pooling.score.protocol import ScoreRequest

83 +

84 + INVOCATION_TYPES += [

85 + (ScoreRequest, (score, create_score)),

86 + ]

87 +

88 + if any(task in POOLING_TASKS for task in supported_tasks):

89 + from vllm.entrypoints.pooling.pooling.api_router import create_pooling, pooling

90 + from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest

91 +

92 + INVOCATION_TYPES += [

93 + (PoolingRequest, (pooling, create_pooling)),

94 + ]

95 +

96 + return INVOCATION_TYPES

97 +

98 +

99 +def attach_router(app: FastAPI, supported_tasks: tuple["SupportedTask", ...]):

100 + router = APIRouter()

101 +

102 + # NOTE: Construct the TypeAdapters only once

103 + INVOCATION_TYPES = get_invocation_types(supported_tasks)

104 + INVOCATION_VALIDATORS = [

105 + (pydantic.TypeAdapter(request_type), (get_handler, endpoint))

106 + for request_type, (get_handler, endpoint) in INVOCATION_TYPES

107 + ]

108 +

109 + @router.post("/ping", response_class=Response)

110 + @router.get("/ping", response_class=Response)

111 + @sagemaker_standards.register_ping_handler

112 + async def ping(raw_request: Request) -> Response:

113 + """Ping check. Endpoint required for SageMaker"""

114 + return await health(raw_request)

115 +

116 + @router.post(

117 + "/invocations",

118 + dependencies=[Depends(validate_json_request)],

119 + responses={

120 + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},

121 + HTTPStatus.UNSUPPORTED_MEDIA_TYPE.value: {"model": ErrorResponse},

122 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},

123 + },

124 + )

125 + @sagemaker_standards.register_invocation_handler

126 + @sagemaker_standards.stateful_session_manager()

127 + @sagemaker_standards.inject_adapter_id(adapter_path="model")

128 + async def invocations(raw_request: Request):

129 + """For SageMaker, routes requests based on the request type."""

130 + try:

131 + body = await raw_request.json()

132 + except json.JSONDecodeError as e:

133 + raise HTTPException(

134 + status_code=HTTPStatus.BAD_REQUEST.value,

135 + detail=f"JSON decode error: {e}",

136 + ) from e

137 +

138 + valid_endpoints = [

139 + (validator, endpoint)

140 + for validator, (get_handler, endpoint) in INVOCATION_VALIDATORS

141 + if get_handler(raw_request) is not None

142 + ]

143 +

144 + for request_validator, endpoint in valid_endpoints:

145 + try:

146 + request = request_validator.validate_python(body)

147 + except pydantic.ValidationError:

148 + continue

149 +

150 + return await endpoint(request, raw_request)

151 +

152 + type_names = [

153 + t.__name__ if isinstance(t := validator._type, type) else str(t)

154 + for validator, _ in valid_endpoints

155 + ]

156 + msg = f"Cannot find suitable handler for request. Expected one of: {type_names}"

157 + res = base(raw_request).create_error_response(message=msg)

158 + return JSONResponse(content=res.model_dump(), status_code=res.error.code)

159 +

160 + app.include_router(router)

vllm/entrypoints/sagemaker/routes.py deleted

-126

@@ -1,126 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import json

4 -from collections.abc import Awaitable, Callable

5 -from http import HTTPStatus

6 -from typing import Any

7 -

8 -import model_hosting_container_standards.sagemaker as sagemaker_standards

9 -import pydantic

10 -from fastapi import APIRouter, Depends, HTTPException, Request

11 -from fastapi.responses import JSONResponse, Response

12 -

13 -from vllm.entrypoints.openai.api_server import (

14 - base,

15 -)

16 -from vllm.entrypoints.openai.chat_completion.api_router import (

17 - chat,

18 - create_chat_completion,

19 -)

20 -from vllm.entrypoints.openai.chat_completion.protocol import (

21 - ChatCompletionRequest,

22 -)

23 -from vllm.entrypoints.openai.completion.api_router import (

24 - completion,

25 - create_completion,

26 -)

27 -from vllm.entrypoints.openai.completion.protocol import (

28 - CompletionRequest,

29 -)

30 -from vllm.entrypoints.openai.engine.protocol import (

31 - ErrorResponse,

32 -)

33 -from vllm.entrypoints.openai.engine.serving import OpenAIServing

34 -from vllm.entrypoints.openai.utils import validate_json_request

35 -from vllm.entrypoints.pooling.classify.api_router import classify, create_classify

36 -from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest

37 -from vllm.entrypoints.pooling.embed.api_router import create_embedding, embedding

38 -from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest

39 -from vllm.entrypoints.pooling.pooling.api_router import create_pooling, pooling

40 -from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest

41 -from vllm.entrypoints.pooling.score.api_router import (

42 - create_score,

43 - do_rerank,

44 - rerank,

45 - score,

46 -)

47 -from vllm.entrypoints.pooling.score.protocol import RerankRequest, ScoreRequest

48 -from vllm.entrypoints.serve.instrumentator.health import health

49 -

50 -# TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers

51 -# (requires typing_extensions >= 4.13)

52 -RequestType = Any

53 -GetHandlerFn = Callable[[Request], OpenAIServing | None]

54 -EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]

55 -

56 -# NOTE: Items defined earlier take higher priority

57 -INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [

58 - (ChatCompletionRequest, (chat, create_chat_completion)),

59 - (CompletionRequest, (completion, create_completion)),

60 - (EmbeddingRequest, (embedding, create_embedding)),

61 - (ClassificationRequest, (classify, create_classify)),

62 - (ScoreRequest, (score, create_score)),

63 - (RerankRequest, (rerank, do_rerank)),

64 - (PoolingRequest, (pooling, create_pooling)),

65 -]

66 -

67 -# NOTE: Construct the TypeAdapters only once

68 -INVOCATION_VALIDATORS = [

69 - (pydantic.TypeAdapter(request_type), (get_handler, endpoint))

70 - for request_type, (get_handler, endpoint) in INVOCATION_TYPES

71 -]

72 -

73 -

74 -def register_sagemaker_routes(router: APIRouter):

75 - @router.post("/ping", response_class=Response)

76 - @router.get("/ping", response_class=Response)

77 - @sagemaker_standards.register_ping_handler

78 - async def ping(raw_request: Request) -> Response:

79 - """Ping check. Endpoint required for SageMaker"""

80 - return await health(raw_request)

81 -

82 - @router.post(

83 - "/invocations",

84 - dependencies=[Depends(validate_json_request)],

85 - responses={

86 - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},

87 - HTTPStatus.UNSUPPORTED_MEDIA_TYPE.value: {"model": ErrorResponse},

88 - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},

89 - },

90 - )

91 - @sagemaker_standards.register_invocation_handler

92 - @sagemaker_standards.stateful_session_manager()

93 - @sagemaker_standards.inject_adapter_id(adapter_path="model")

94 - async def invocations(raw_request: Request):

95 - """For SageMaker, routes requests based on the request type."""

96 - try:

97 - body = await raw_request.json()

98 - except json.JSONDecodeError as e:

99 - raise HTTPException(

100 - status_code=HTTPStatus.BAD_REQUEST.value,

101 - detail=f"JSON decode error: {e}",

102 - ) from e

103 -

104 - valid_endpoints = [

105 - (validator, endpoint)

106 - for validator, (get_handler, endpoint) in INVOCATION_VALIDATORS

107 - if get_handler(raw_request) is not None

108 - ]

109 -

110 - for request_validator, endpoint in valid_endpoints:

111 - try:

112 - request = request_validator.validate_python(body)

113 - except pydantic.ValidationError:

114 - continue

115 -

116 - return await endpoint(request, raw_request)

117 -

118 - type_names = [

119 - t.__name__ if isinstance(t := validator._type, type) else str(t)

120 - for validator, _ in valid_endpoints

121 - ]

122 - msg = f"Cannot find suitable handler for request. Expected one of: {type_names}"

123 - res = base(raw_request).create_error_response(message=msg)

124 - return JSONResponse(content=res.model_dump(), status_code=res.error.code)

125 -

126 - return router