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 12, 07:15 UTC
[Frontend] Resettle pooling entrypoints (#29634)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Dec 1, 2025 at 15:30 UTC
62de4f4257255bff553a61f475212e3e9454706c
39 files changed
+1264
-1067
.github/CODEOWNERS
+1
@@ -149,6 +149,7 @@ mkdocs.yaml @hmellor
149
/examples/*/pooling/ @noooop
150
/tests/models/*/pooling* @noooop
151
/tests/entrypoints/pooling @noooop
152
+/vllm/entrypoints/pooling @aarnphm @chaunceyjiang @noooop
153
/vllm/config/pooler.py @noooop
154
/vllm/pooling_params.py @noooop
155
/vllm/model_executor/layers/pooler.py @noooop
docs/design/io_processor_plugins.md
+1
-1
@@ -77,7 +77,7 @@ The `parse_request` method is used for validating the user prompt and converting
77
The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.
78
The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.
79
The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
80
-The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/openai/serving_pooling.py).
80
+The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).
81
82
An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/online_serving/pooling/prithvi_geospatial_mae.py](../../examples/online_serving/pooling/prithvi_geospatial_mae.py)) and offline ([examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py](../../examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py)) inference examples.
83
docs/serving/openai_compatible_server.md
+5
-5
@@ -351,7 +351,7 @@ The following extra parameters are supported by default:
351
??? code
352
353
```python
354
- --8<-- "vllm/entrypoints/openai/protocol.py:embedding-extra-params"
354
+ --8<-- "vllm/entrypoints/pooling/embed/protocol.py:embedding-extra-params"
355
```
356
357
For chat-like input (i.e. if `messages` is passed), these extra parameters are supported instead:
@@ -359,7 +359,7 @@ For chat-like input (i.e. if `messages` is passed), these extra parameters are s
359
??? code
360
361
```python
362
- --8<-- "vllm/entrypoints/openai/protocol.py:chat-embedding-extra-params"
362
+ --8<-- "vllm/entrypoints/pooling/embed/protocol.py:chat-embedding-extra-params"
363
```
364
365
### Transcriptions API
@@ -629,7 +629,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
629
The following extra parameters are supported:
630
631
```python
632
---8<-- "vllm/entrypoints/openai/protocol.py:classification-extra-params"
632
+--8<-- "vllm/entrypoints/pooling/classify/protocol.py:classification-extra-params"
633
```
634
635
### Score API
@@ -834,7 +834,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
834
The following extra parameters are supported:
835
836
```python
837
---8<-- "vllm/entrypoints/openai/protocol.py:score-extra-params"
837
+--8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
838
```
839
840
### Re-rank API
@@ -915,7 +915,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
915
The following extra parameters are supported:
916
917
```python
918
---8<-- "vllm/entrypoints/openai/protocol.py:rerank-extra-params"
918
+--8<-- "vllm/entrypoints/pooling/score/protocol.py:rerank-extra-params"
919
```
920
921
## Ray Serve LLM
tests/entrypoints/openai/test_run_batch.py
+1
-1
@@ -7,7 +7,7 @@ import tempfile
7
8
import pytest
9
10
-from vllm.entrypoints.openai.protocol import BatchRequestOutput
10
+from vllm.entrypoints.openai.run_batch import BatchRequestOutput
11
12
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
13
tests/entrypoints/pooling/classify/test_online.py
+2
-1
@@ -7,7 +7,8 @@ import torch
7
import torch.nn.functional as F
8
9
from tests.utils import RemoteOpenAIServer
10
-from vllm.entrypoints.openai.protocol import ClassificationResponse, PoolingResponse
10
+from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
11
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
12
13
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
14
DTYPE = "float32" # Use float32 to avoid NaN issue
tests/entrypoints/pooling/classify/test_online_vision.py
+1
-1
@@ -7,7 +7,7 @@ import pytest
7
import requests
8
9
from tests.utils import RemoteOpenAIServer
10
-from vllm.entrypoints.openai.protocol import ClassificationResponse
10
+from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
11
12
VLM_MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
13
MAXIMUM_VIDEOS = 1
tests/entrypoints/pooling/embed/test_online.py
+2
-4
@@ -15,10 +15,8 @@ import torch.nn.functional as F
15
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
16
from tests.models.utils import check_embeddings_close
17
from tests.utils import RemoteOpenAIServer
18
-from vllm.entrypoints.openai.protocol import (
19
- EmbeddingResponse,
20
- PoolingResponse,
21
-)
18
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
19
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
20
from vllm.platforms import current_platform
21
from vllm.transformers_utils.tokenizer import get_tokenizer
22
from vllm.utils.serial_utils import (
tests/entrypoints/pooling/embed/test_online_dimensions.py
+1
-1
@@ -11,7 +11,7 @@ from tests.conftest import HfRunner
11
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
12
from tests.models.utils import EmbedModelInfo
13
from tests.utils import RemoteOpenAIServer
14
-from vllm.entrypoints.openai.protocol import EmbeddingResponse
14
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
15
from vllm.platforms import current_platform
16
17
if current_platform.is_rocm():
tests/entrypoints/pooling/embed/test_online_long_text.py
+1
-1
@@ -15,7 +15,7 @@ import pytest
15
import pytest_asyncio
16
17
from tests.utils import RemoteOpenAIServer
18
-from vllm.entrypoints.openai.protocol import EmbeddingResponse
18
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
19
from vllm.platforms import current_platform
20
21
if current_platform.is_rocm():
tests/entrypoints/pooling/embed/test_online_vision.py
+1
-1
@@ -8,7 +8,7 @@ import requests
8
from transformers import AutoProcessor
9
10
from tests.utils import VLLM_PATH, RemoteOpenAIServer
11
-from vllm.entrypoints.openai.protocol import EmbeddingResponse
11
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
12
from vllm.multimodal.utils import encode_image_base64, fetch_image
13
14
MODEL_NAME = "TIGER-Lab/VLM2Vec-Full"
tests/entrypoints/pooling/pooling/test_online.py
+1
-1
@@ -11,7 +11,7 @@ import torch
11
12
from tests.models.utils import check_embeddings_close
13
from tests.utils import RemoteOpenAIServer
14
-from vllm.entrypoints.openai.protocol import PoolingResponse
14
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
15
from vllm.transformers_utils.tokenizer import get_tokenizer
16
from vllm.utils.serial_utils import (
17
EMBED_DTYPE_TO_TORCH_DTYPE,
tests/entrypoints/pooling/score/test_online_rerank.py
+2
-1
@@ -7,7 +7,8 @@ import torch
7
import torch.nn.functional as F
8
9
from tests.utils import RemoteOpenAIServer
10
-from vllm.entrypoints.openai.protocol import PoolingResponse, RerankResponse
10
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
11
+from vllm.entrypoints.pooling.score.protocol import RerankResponse
12
from vllm.platforms import current_platform
13
14
if current_platform.is_rocm():
tests/entrypoints/pooling/score/test_online_score.py
+1
-1
@@ -9,7 +9,7 @@ import torch.nn.functional as F
9
from torch import tensor
10
11
from tests.utils import RemoteOpenAIServer
12
-from vllm.entrypoints.openai.protocol import ScoreResponse
12
+from vllm.entrypoints.pooling.score.protocol import ScoreResponse
13
from vllm.platforms import current_platform
14
15
if current_platform.is_rocm():
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
+4
-1
@@ -18,7 +18,10 @@ from einops import rearrange
18
from terratorch.datamodules import Sen1Floods11NonGeoDataModule
19
20
from vllm.config import VllmConfig
21
-from vllm.entrypoints.openai.protocol import IOProcessorRequest, IOProcessorResponse
21
+from vllm.entrypoints.pooling.pooling.protocol import (
22
+ IOProcessorRequest,
23
+ IOProcessorResponse,
24
+)
25
from vllm.inputs.data import PromptType
26
from vllm.logger import init_logger
27
from vllm.outputs import PoolingRequestOutput
tests/plugins_tests/test_io_processor_plugins.py
+1
-1
@@ -7,7 +7,7 @@ import requests
7
8
from tests.utils import RemoteOpenAIServer
9
from vllm.config import VllmConfig
10
-from vllm.entrypoints.openai.protocol import IOProcessorResponse
10
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
11
from vllm.plugins.io_processors import get_io_processor
12
13
MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
vllm/entrypoints/openai/api_server.py
+11
-296
@@ -14,7 +14,7 @@ import socket
14
import tempfile
15
import uuid
16
from argparse import Namespace
17
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
17
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable
18
from contextlib import asynccontextmanager
19
from http import HTTPStatus
20
from typing import Annotated, Any, Literal
@@ -54,29 +54,16 @@ from vllm.entrypoints.openai.orca_metrics import metrics_header
54
from vllm.entrypoints.openai.protocol import (
55
ChatCompletionRequest,
56
ChatCompletionResponse,
57
- ClassificationRequest,
58
- ClassificationResponse,
57
CompletionRequest,
58
CompletionResponse,
59
DetokenizeRequest,
60
DetokenizeResponse,
63
- EmbeddingBytesResponse,
64
- EmbeddingRequest,
65
- EmbeddingResponse,
61
ErrorInfo,
62
ErrorResponse,
63
GenerateRequest,
64
GenerateResponse,
70
- IOProcessorResponse,
71
- PoolingBytesResponse,
72
- PoolingRequest,
73
- PoolingResponse,
74
- RerankRequest,
75
- RerankResponse,
65
ResponsesRequest,
66
ResponsesResponse,
78
- ScoreRequest,
79
- ScoreResponse,
67
StreamingResponsesResponse,
68
TokenizeRequest,
69
TokenizeResponse,
@@ -86,17 +73,13 @@ from vllm.entrypoints.openai.protocol import (
73
TranslationResponse,
74
)
75
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
89
-from vllm.entrypoints.openai.serving_classification import ServingClassification
76
from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
91
-from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
77
from vllm.entrypoints.openai.serving_engine import OpenAIServing
78
from vllm.entrypoints.openai.serving_models import (
79
BaseModelPath,
80
OpenAIServingModels,
81
)
97
-from vllm.entrypoints.openai.serving_pooling import OpenAIServingPooling
82
from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
99
-from vllm.entrypoints.openai.serving_score import ServingScores
83
from vllm.entrypoints.openai.serving_tokenization import OpenAIServingTokenization
84
from vllm.entrypoints.openai.serving_tokens import ServingTokens
85
from vllm.entrypoints.openai.serving_transcription import (
@@ -104,6 +87,11 @@ from vllm.entrypoints.openai.serving_transcription import (
87
OpenAIServingTranslation,
88
)
89
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
90
+from vllm.entrypoints.openai.utils import validate_json_request
91
+from vllm.entrypoints.pooling.classify.serving import ServingClassification
92
+from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
93
+from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
94
+from vllm.entrypoints.pooling.score.serving import ServingScores
95
from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
96
from vllm.entrypoints.utils import (
97
cli_env_setup,
@@ -254,15 +242,6 @@ async def build_async_engine_client_from_engine_args(
242
async_llm.shutdown()
243
244
257
-async def validate_json_request(raw_request: Request):
258
- content_type = raw_request.headers.get("content-type", "").lower()
259
- media_type = content_type.split(";", maxsplit=1)[0]
260
- if media_type != "application/json":
261
- raise RequestValidationError(
262
- errors=["Unsupported Media Type: Only 'application/json' is allowed"]
263
- )
264
-
265
-
245
router = APIRouter()
246
247
@@ -324,26 +303,6 @@ def completion(request: Request) -> OpenAIServingCompletion | None:
303
return request.app.state.openai_serving_completion
304
305
327
-def pooling(request: Request) -> OpenAIServingPooling | None:
328
- return request.app.state.openai_serving_pooling
329
-
330
-
331
-def embedding(request: Request) -> OpenAIServingEmbedding | None:
332
- return request.app.state.openai_serving_embedding
333
-
334
-
335
-def score(request: Request) -> ServingScores | None:
336
- return request.app.state.openai_serving_scores
337
-
338
-
339
-def classify(request: Request) -> ServingClassification | None:
340
- return request.app.state.openai_serving_classification
341
-
342
-
343
-def rerank(request: Request) -> ServingScores | None:
344
- return request.app.state.openai_serving_scores
345
-
346
-
306
def tokenization(request: Request) -> OpenAIServingTokenization:
307
return request.app.state.openai_serving_tokenization
308
@@ -817,166 +776,6 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
776
return StreamingResponse(content=generator, media_type="text/event-stream")
777
778
820
-@router.post(
821
- "/v1/embeddings",
822
- dependencies=[Depends(validate_json_request)],
823
- responses={
824
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
825
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
826
- },
827
-)
828
-@with_cancellation
829
-@load_aware_call
830
-async def create_embedding(
831
- request: EmbeddingRequest,
832
- raw_request: Request,
833
-):
834
- handler = embedding(raw_request)
835
- if handler is None:
836
- return base(raw_request).create_error_response(
837
- message="The model does not support Embeddings API"
838
- )
839
-
840
- try:
841
- generator = await handler.create_embedding(request, raw_request)
842
- except Exception as e:
843
- raise HTTPException(
844
- status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
845
- ) from e
846
-
847
- if isinstance(generator, ErrorResponse):
848
- return JSONResponse(
849
- content=generator.model_dump(), status_code=generator.error.code
850
- )
851
- elif isinstance(generator, EmbeddingResponse):
852
- return JSONResponse(content=generator.model_dump())
853
- elif isinstance(generator, EmbeddingBytesResponse):
854
- return StreamingResponse(
855
- content=generator.body,
856
- headers={"metadata": generator.metadata},
857
- media_type=generator.media_type,
858
- )
859
-
860
- assert_never(generator)
861
-
862
-
863
-@router.post(
864
- "/pooling",
865
- dependencies=[Depends(validate_json_request)],
866
- responses={
867
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
868
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
869
- },
870
-)
871
-@with_cancellation
872
-@load_aware_call
873
-async def create_pooling(request: PoolingRequest, raw_request: Request):
874
- handler = pooling(raw_request)
875
- if handler is None:
876
- return base(raw_request).create_error_response(
877
- message="The model does not support Pooling API"
878
- )
879
- try:
880
- generator = await handler.create_pooling(request, raw_request)
881
- except Exception as e:
882
- raise HTTPException(
883
- status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
884
- ) from e
885
- if isinstance(generator, ErrorResponse):
886
- return JSONResponse(
887
- content=generator.model_dump(), status_code=generator.error.code
888
- )
889
- elif isinstance(generator, (PoolingResponse, IOProcessorResponse)):
890
- return JSONResponse(content=generator.model_dump())
891
- elif isinstance(generator, PoolingBytesResponse):
892
- return StreamingResponse(
893
- content=generator.body,
894
- headers={"metadata": generator.metadata},
895
- media_type=generator.media_type,
896
- )
897
-
898
- assert_never(generator)
899
-
900
-
901
-@router.post("/classify", dependencies=[Depends(validate_json_request)])
902
-@with_cancellation
903
-@load_aware_call
904
-async def create_classify(request: ClassificationRequest, raw_request: Request):
905
- handler = classify(raw_request)
906
- if handler is None:
907
- return base(raw_request).create_error_response(
908
- message="The model does not support Classification API"
909
- )
910
-
911
- try:
912
- generator = await handler.create_classify(request, raw_request)
913
- except Exception as e:
914
- raise HTTPException(
915
- status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
916
- ) from e
917
- if isinstance(generator, ErrorResponse):
918
- return JSONResponse(
919
- content=generator.model_dump(), status_code=generator.error.code
920
- )
921
-
922
- elif isinstance(generator, ClassificationResponse):
923
- return JSONResponse(content=generator.model_dump())
924
-
925
- assert_never(generator)
926
-
927
-
928
-@router.post(
929
- "/score",
930
- dependencies=[Depends(validate_json_request)],
931
- responses={
932
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
933
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
934
- },
935
-)
936
-@with_cancellation
937
-@load_aware_call
938
-async def create_score(request: ScoreRequest, raw_request: Request):
939
- handler = score(raw_request)
940
- if handler is None:
941
- return base(raw_request).create_error_response(
942
- message="The model does not support Score API"
943
- )
944
-
945
- try:
946
- generator = await handler.create_score(request, raw_request)
947
- except Exception as e:
948
- raise HTTPException(
949
- status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
950
- ) from e
951
- if isinstance(generator, ErrorResponse):
952
- return JSONResponse(
953
- content=generator.model_dump(), status_code=generator.error.code
954
- )
955
- elif isinstance(generator, ScoreResponse):
956
- return JSONResponse(content=generator.model_dump())
957
-
958
- assert_never(generator)
959
-
960
-
961
-@router.post(
962
- "/v1/score",
963
- dependencies=[Depends(validate_json_request)],
964
- responses={
965
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
966
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
967
- },
968
-)
969
-@with_cancellation
970
-@load_aware_call
971
-async def create_score_v1(request: ScoreRequest, raw_request: Request):
972
- logger.warning(
973
- "To indicate that Score API is not part of standard OpenAI API, we "
974
- "have moved it to `/score`. Please update your client accordingly."
975
- )
976
-
977
- return await create_score(request, raw_request)
978
-
979
-
779
@router.post(
780
"/v1/audio/transcriptions",
781
responses={
@@ -1055,70 +854,6 @@ async def create_translations(
854
return StreamingResponse(content=generator, media_type="text/event-stream")
855
856
1058
-@router.post(
1059
- "/rerank",
1060
- dependencies=[Depends(validate_json_request)],
1061
- responses={
1062
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
1063
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
1064
- },
1065
-)
1066
-@with_cancellation
1067
-@load_aware_call
1068
-async def do_rerank(request: RerankRequest, raw_request: Request):
1069
- handler = rerank(raw_request)
1070
- if handler is None:
1071
- return base(raw_request).create_error_response(
1072
- message="The model does not support Rerank (Score) API"
1073
- )
1074
- try:
1075
- generator = await handler.do_rerank(request, raw_request)
1076
- except Exception as e:
1077
- raise HTTPException(
1078
- status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
1079
- ) from e
1080
- if isinstance(generator, ErrorResponse):
1081
- return JSONResponse(
1082
- content=generator.model_dump(), status_code=generator.error.code
1083
- )
1084
- elif isinstance(generator, RerankResponse):
1085
- return JSONResponse(content=generator.model_dump())
1086
-
1087
- assert_never(generator)
1088
-
1089
-
1090
-@router.post(
1091
- "/v1/rerank",
1092
- dependencies=[Depends(validate_json_request)],
1093
- responses={
1094
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
1095
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
1096
- },
1097
-)
1098
-@with_cancellation
1099
-async def do_rerank_v1(request: RerankRequest, raw_request: Request):
1100
- logger.warning_once(
1101
- "To indicate that the rerank API is not part of the standard OpenAI"
1102
- " API, we have located it at `/rerank`. Please update your client "
1103
- "accordingly. (Note: Conforms to JinaAI rerank API)"
1104
- )
1105
-
1106
- return await do_rerank(request, raw_request)
1107
-
1108
-
1109
-@router.post(
1110
- "/v2/rerank",
1111
- dependencies=[Depends(validate_json_request)],
1112
- responses={
1113
- HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
1114
- HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
1115
- },
1116
-)
1117
-@with_cancellation
1118
-async def do_rerank_v2(request: RerankRequest, raw_request: Request):
1119
- return await do_rerank(request, raw_request)
1120
-
1121
-
857
if envs.VLLM_SERVER_DEV_MODE:
858
logger.warning(
859
"SECURITY WARNING: Development endpoints are enabled! "
@@ -1285,30 +1020,6 @@ async def is_scaling_elastic_ep(raw_request: Request):
1020
return JSONResponse({"is_scaling_elastic_ep": _scaling_elastic_ep})
1021
1022
1288
-# TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
1289
-# (requires typing_extensions >= 4.13)
1290
-RequestType = Any
1291
-GetHandlerFn = Callable[[Request], OpenAIServing | None]
1292
-EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
1293
-
1294
-# NOTE: Items defined earlier take higher priority
1295
-INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [
1296
- (ChatCompletionRequest, (chat, create_chat_completion)),
1297
- (CompletionRequest, (completion, create_completion)),
1298
- (EmbeddingRequest, (embedding, create_embedding)),
1299
- (ClassificationRequest, (classify, create_classify)),
1300
- (ScoreRequest, (score, create_score)),
1301
- (RerankRequest, (rerank, do_rerank)),
1302
- (PoolingRequest, (pooling, create_pooling)),
1303
-]
1304
-
1305
-# NOTE: Construct the TypeAdapters only once
1306
-INVOCATION_VALIDATORS = [
1307
- (pydantic.TypeAdapter(request_type), (get_handler, endpoint))
1308
- for request_type, (get_handler, endpoint) in INVOCATION_TYPES
1309
-]
1310
-
1311
-
1023
@router.post(
1024
"/inference/v1/generate",
1025
dependencies=[Depends(validate_json_request)],
@@ -1653,12 +1364,16 @@ def build_app(args: Namespace) -> FastAPI:
1364
from vllm.entrypoints.sagemaker.routes import register_sagemaker_routes
1365
1366
register_sagemaker_routes(router)
1656
-
1367
app.include_router(router)
1368
+
1369
app.root_path = args.root_path
1370
1371
mount_metrics(app)
1372
1373
+ from vllm.entrypoints.pooling import register_pooling_api_routers
1374
+
1375
+ register_pooling_api_routers(app)
1376
+
1377
app.add_middleware(
1378
CORSMiddleware,
1379
allow_origins=args.allowed_origins,
vllm/entrypoints/openai/protocol.py
+1
-709
@@ -6,7 +6,7 @@
6
import json
7
import time
8
from http import HTTPStatus
9
-from typing import Annotated, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar
9
+from typing import Annotated, Any, ClassVar, Literal, TypeAlias
10
11
import regex as re
12
import torch
@@ -48,14 +48,6 @@ from openai.types.responses.response_reasoning_item import (
48
)
49
from openai_harmony import Message as OpenAIHarmonyMessage
50
51
-from vllm.config.pooler import get_use_activation
52
-from vllm.tasks import PoolingTask
53
-from vllm.utils.serial_utils import (
54
- EmbedDType,
55
- EncodingFormat,
56
- Endianness,
57
-)
58
-
51
# Backward compatibility for OpenAI client versions
52
try: # For older openai versions (< 1.100.0)
53
from openai.types.responses import ResponseTextConfig
@@ -70,19 +62,14 @@ from pydantic import (
62
BaseModel,
63
ConfigDict,
64
Field,
73
- TypeAdapter,
65
ValidationError,
75
- ValidationInfo,
66
field_serializer,
77
- field_validator,
67
model_validator,
68
)
69
70
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam, make_tool_call_id
82
-from vllm.entrypoints.score_utils import ScoreContentPartParam, ScoreMultiModalParam
71
from vllm.logger import init_logger
72
from vllm.logprobs import Logprob
85
-from vllm.pooling_params import PoolingParams
73
from vllm.sampling_params import (
74
BeamSearchParams,
75
RequestOutputKind,
@@ -1345,401 +1332,6 @@ class CompletionRequest(OpenAIBaseModel):
1332
return data
1333
1334
1348
-class EmbeddingCompletionRequest(OpenAIBaseModel):
1349
- # Ordered by official OpenAI API documentation
1350
- # https://platform.openai.com/docs/api-reference/embeddings
1351
- model: str | None = None
1352
- input: list[int] | list[list[int]] | str | list[str]
1353
- encoding_format: EncodingFormat = "float"
1354
- dimensions: int | None = None
1355
- user: str | None = None
1356
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1357
-
1358
- # --8<-- [start:embedding-extra-params]
1359
- add_special_tokens: bool = Field(
1360
- default=True,
1361
- description=(
1362
- "If true (the default), special tokens (e.g. BOS) will be added to "
1363
- "the prompt."
1364
- ),
1365
- )
1366
- priority: int = Field(
1367
- default=0,
1368
- description=(
1369
- "The priority of the request (lower means earlier handling; "
1370
- "default: 0). Any priority other than 0 will raise an error "
1371
- "if the served model does not use priority scheduling."
1372
- ),
1373
- )
1374
- request_id: str = Field(
1375
- default_factory=random_uuid,
1376
- description=(
1377
- "The request_id related to this request. If the caller does "
1378
- "not set it, a random_uuid will be generated. This id is used "
1379
- "through out the inference process and return in response."
1380
- ),
1381
- )
1382
- normalize: bool | None = Field(
1383
- default=None,
1384
- description="Whether to normalize the embeddings outputs. Default is True.",
1385
- )
1386
- embed_dtype: EmbedDType = Field(
1387
- default="float32",
1388
- description=(
1389
- "What dtype to use for encoding. Default to using float32 for base64 "
1390
- "encoding to match the OpenAI python client behavior. "
1391
- "This parameter will affect base64 and binary_response."
1392
- ),
1393
- )
1394
- endianness: Endianness = Field(
1395
- default="native",
1396
- description=(
1397
- "What endianness to use for encoding. Default to using native for "
1398
- "base64 encoding to match the OpenAI python client behavior."
1399
- "This parameter will affect base64 and binary_response."
1400
- ),
1401
- )
1402
- # --8<-- [end:embedding-extra-params]
1403
-
1404
- def to_pooling_params(self):
1405
- return PoolingParams(
1406
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1407
- dimensions=self.dimensions,
1408
- normalize=self.normalize,
1409
- )
1410
-
1411
-
1412
-class EmbeddingChatRequest(OpenAIBaseModel):
1413
- model: str | None = None
1414
- messages: list[ChatCompletionMessageParam]
1415
-
1416
- encoding_format: EncodingFormat = "float"
1417
- dimensions: int | None = None
1418
- user: str | None = None
1419
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1420
-
1421
- # --8<-- [start:chat-embedding-extra-params]
1422
- add_generation_prompt: bool = Field(
1423
- default=False,
1424
- description=(
1425
- "If true, the generation prompt will be added to the chat template. "
1426
- "This is a parameter used by chat template in tokenizer config of the "
1427
- "model."
1428
- ),
1429
- )
1430
-
1431
- add_special_tokens: bool = Field(
1432
- default=False,
1433
- description=(
1434
- "If true, special tokens (e.g. BOS) will be added to the prompt "
1435
- "on top of what is added by the chat template. "
1436
- "For most models, the chat template takes care of adding the "
1437
- "special tokens so this should be set to false (as is the "
1438
- "default)."
1439
- ),
1440
- )
1441
- chat_template: str | None = Field(
1442
- default=None,
1443
- description=(
1444
- "A Jinja template to use for this conversion. "
1445
- "As of transformers v4.44, default chat template is no longer "
1446
- "allowed, so you must provide a chat template if the tokenizer "
1447
- "does not define one."
1448
- ),
1449
- )
1450
- chat_template_kwargs: dict[str, Any] | None = Field(
1451
- default=None,
1452
- description=(
1453
- "Additional keyword args to pass to the template renderer. "
1454
- "Will be accessible by the chat template."
1455
- ),
1456
- )
1457
- mm_processor_kwargs: dict[str, Any] | None = Field(
1458
- default=None,
1459
- description=("Additional kwargs to pass to the HF processor."),
1460
- )
1461
- priority: int = Field(
1462
- default=0,
1463
- description=(
1464
- "The priority of the request (lower means earlier handling; "
1465
- "default: 0). Any priority other than 0 will raise an error "
1466
- "if the served model does not use priority scheduling."
1467
- ),
1468
- )
1469
- request_id: str = Field(
1470
- default_factory=random_uuid,
1471
- description=(
1472
- "The request_id related to this request. If the caller does "
1473
- "not set it, a random_uuid will be generated. This id is used "
1474
- "through out the inference process and return in response."
1475
- ),
1476
- )
1477
- normalize: bool | None = Field(
1478
- default=None,
1479
- description="Whether to normalize the embeddings outputs. Default is True.",
1480
- )
1481
- embed_dtype: EmbedDType = Field(
1482
- default="float32",
1483
- description=(
1484
- "What dtype to use for encoding. Default to using float32 for base64 "
1485
- "encoding to match the OpenAI python client behavior. "
1486
- "This parameter will affect base64 and binary_response."
1487
- ),
1488
- )
1489
- endianness: Endianness = Field(
1490
- default="native",
1491
- description=(
1492
- "What endianness to use for encoding. Default to using native for "
1493
- "base64 encoding to match the OpenAI python client behavior."
1494
- "This parameter will affect base64 and binary_response."
1495
- ),
1496
- )
1497
- # --8<-- [end:chat-embedding-extra-params]
1498
-
1499
- @model_validator(mode="before")
1500
- @classmethod
1501
- def check_generation_prompt(cls, data):
1502
- if data.get("continue_final_message") and data.get("add_generation_prompt"):
1503
- raise ValueError(
1504
- "Cannot set both `continue_final_message` and "
1505
- "`add_generation_prompt` to True."
1506
- )
1507
- return data
1508
-
1509
- def to_pooling_params(self):
1510
- return PoolingParams(
1511
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1512
- dimensions=self.dimensions,
1513
- normalize=self.normalize,
1514
- )
1515
-
1516
-
1517
-EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
1518
-
1519
-
1520
-class PoolingCompletionRequest(EmbeddingCompletionRequest):
1521
- task: PoolingTask | None = None
1522
- softmax: bool | None = Field(
1523
- default=None,
1524
- description="softmax will be deprecated, please use use_activation instead.",
1525
- )
1526
- activation: bool | None = Field(
1527
- default=None,
1528
- description="activation will be deprecated, please use use_activation instead.",
1529
- )
1530
- use_activation: bool | None = Field(
1531
- default=None,
1532
- description="Whether to use activation for classification outputs. "
1533
- "If it is a classify or token_classify task, the default is True; "
1534
- "for other tasks, this value should be None.",
1535
- )
1536
-
1537
- def to_pooling_params(self):
1538
- return PoolingParams(
1539
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1540
- dimensions=self.dimensions,
1541
- normalize=self.normalize,
1542
- use_activation=get_use_activation(self),
1543
- )
1544
-
1545
-
1546
-class PoolingChatRequest(EmbeddingChatRequest):
1547
- task: PoolingTask | None = None
1548
- softmax: bool | None = Field(
1549
- default=None,
1550
- description="softmax will be deprecated, please use use_activation instead.",
1551
- )
1552
- activation: bool | None = Field(
1553
- default=None,
1554
- description="activation will be deprecated, please use use_activation instead.",
1555
- )
1556
- use_activation: bool | None = Field(
1557
- default=None,
1558
- description="Whether to use activation for classification outputs. "
1559
- "If it is a classify or token_classify task, the default is True; "
1560
- "for other tasks, this value should be None.",
1561
- )
1562
-
1563
- def to_pooling_params(self):
1564
- return PoolingParams(
1565
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1566
- dimensions=self.dimensions,
1567
- normalize=self.normalize,
1568
- use_activation=get_use_activation(self),
1569
- )
1570
-
1571
-
1572
-T = TypeVar("T")
1573
-
1574
-
1575
-class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1576
- model: str | None = None
1577
-
1578
- priority: int = Field(default=0)
1579
- """
1580
- The priority of the request (lower means earlier handling;
1581
- default: 0). Any priority other than 0 will raise an error
1582
- if the served model does not use priority scheduling.
1583
- """
1584
- data: T
1585
-
1586
- task: PoolingTask = "plugin"
1587
- encoding_format: EncodingFormat = "float"
1588
- embed_dtype: EmbedDType = Field(
1589
- default="float32",
1590
- description=(
1591
- "What dtype to use for encoding. Default to using float32 for base64 "
1592
- "encoding to match the OpenAI python client behavior. "
1593
- "This parameter will affect base64 and binary_response."
1594
- ),
1595
- )
1596
- endianness: Endianness = Field(
1597
- default="native",
1598
- description=(
1599
- "What endianness to use for encoding. Default to using native for "
1600
- "base64 encoding to match the OpenAI python client behavior."
1601
- "This parameter will affect base64 and binary_response."
1602
- ),
1603
- )
1604
-
1605
- def to_pooling_params(self):
1606
- return PoolingParams()
1607
-
1608
-
1609
-class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
1610
- request_id: str | None = None
1611
- """
1612
- The request_id associated with this response
1613
- """
1614
- created_at: int = Field(default_factory=lambda: int(time.time()))
1615
-
1616
- data: T
1617
- """
1618
- When using plugins IOProcessor plugins, the actual output is generated
1619
- by the plugin itself. Hence, we use a generic type for the response data
1620
- """
1621
-
1622
-
1623
-PoolingRequest: TypeAlias = (
1624
- PoolingCompletionRequest | PoolingChatRequest | IOProcessorRequest
1625
-)
1626
-
1627
-
1628
-class ScoreRequest(OpenAIBaseModel):
1629
- model: str | None = None
1630
- text_1: list[str] | str | ScoreMultiModalParam
1631
- text_2: list[str] | str | ScoreMultiModalParam
1632
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1633
-
1634
- # --8<-- [start:score-extra-params]
1635
-
1636
- mm_processor_kwargs: dict[str, Any] | None = Field(
1637
- default=None,
1638
- description=("Additional kwargs to pass to the HF processor."),
1639
- )
1640
-
1641
- priority: int = Field(
1642
- default=0,
1643
- description=(
1644
- "The priority of the request (lower means earlier handling; "
1645
- "default: 0). Any priority other than 0 will raise an error "
1646
- "if the served model does not use priority scheduling."
1647
- ),
1648
- )
1649
-
1650
- softmax: bool | None = Field(
1651
- default=None,
1652
- description="softmax will be deprecated, please use use_activation instead.",
1653
- )
1654
-
1655
- activation: bool | None = Field(
1656
- default=None,
1657
- description="activation will be deprecated, please use use_activation instead.",
1658
- )
1659
-
1660
- use_activation: bool | None = Field(
1661
- default=None,
1662
- description="Whether to use activation for classification outputs. "
1663
- "Default is True.",
1664
- )
1665
- # --8<-- [end:score-extra-params]
1666
-
1667
- def to_pooling_params(self):
1668
- return PoolingParams(
1669
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1670
- use_activation=get_use_activation(self),
1671
- )
1672
-
1673
-
1674
-class RerankRequest(OpenAIBaseModel):
1675
- model: str | None = None
1676
- query: str | ScoreMultiModalParam
1677
- documents: list[str] | ScoreMultiModalParam
1678
- top_n: int = Field(default_factory=lambda: 0)
1679
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1680
-
1681
- # --8<-- [start:rerank-extra-params]
1682
-
1683
- mm_processor_kwargs: dict[str, Any] | None = Field(
1684
- default=None,
1685
- description=("Additional kwargs to pass to the HF processor."),
1686
- )
1687
-
1688
- priority: int = Field(
1689
- default=0,
1690
- description=(
1691
- "The priority of the request (lower means earlier handling; "
1692
- "default: 0). Any priority other than 0 will raise an error "
1693
- "if the served model does not use priority scheduling."
1694
- ),
1695
- )
1696
-
1697
- softmax: bool | None = Field(
1698
- default=None,
1699
- description="softmax will be deprecated, please use use_activation instead.",
1700
- )
1701
-
1702
- activation: bool | None = Field(
1703
- default=None,
1704
- description="activation will be deprecated, please use use_activation instead.",
1705
- )
1706
-
1707
- use_activation: bool | None = Field(
1708
- default=None,
1709
- description="Whether to use activation for classification outputs. "
1710
- "Default is True.",
1711
- )
1712
- # --8<-- [end:rerank-extra-params]
1713
-
1714
- def to_pooling_params(self):
1715
- return PoolingParams(
1716
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1717
- use_activation=get_use_activation(self),
1718
- )
1719
-
1720
-
1721
-class RerankDocument(BaseModel):
1722
- text: str | None = None
1723
- multi_modal: ScoreContentPartParam | None = None
1724
-
1725
-
1726
-class RerankResult(BaseModel):
1727
- index: int
1728
- document: RerankDocument
1729
- relevance_score: float
1730
-
1731
-
1732
-class RerankUsage(BaseModel):
1733
- total_tokens: int
1734
-
1735
-
1736
-class RerankResponse(OpenAIBaseModel):
1737
- id: str
1738
- model: str
1739
- usage: RerankUsage
1740
- results: list[RerankResult]
1741
-
1742
-
1335
class CompletionLogProbs(OpenAIBaseModel):
1336
text_offset: list[int] = Field(default_factory=list)
1337
token_logprobs: list[float | None] = Field(default_factory=list)
@@ -1809,229 +1401,6 @@ class CompletionStreamResponse(OpenAIBaseModel):
1401
usage: UsageInfo | None = Field(default=None)
1402
1403
1812
-class EmbeddingResponseData(OpenAIBaseModel):
1813
- index: int
1814
- object: str = "embedding"
1815
- embedding: list[float] | str
1816
-
1817
-
1818
-class EmbeddingResponse(OpenAIBaseModel):
1819
- id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
1820
- object: str = "list"
1821
- created: int = Field(default_factory=lambda: int(time.time()))
1822
- model: str
1823
- data: list[EmbeddingResponseData]
1824
- usage: UsageInfo
1825
-
1826
-
1827
-class EmbeddingBytesResponse(OpenAIBaseModel):
1828
- body: list[bytes]
1829
- metadata: str
1830
- media_type: str = "application/octet-stream"
1831
-
1832
-
1833
-class PoolingResponseData(OpenAIBaseModel):
1834
- index: int
1835
- object: str = "pooling"
1836
- data: list[list[float]] | list[float] | str
1837
-
1838
-
1839
-class PoolingResponse(OpenAIBaseModel):
1840
- id: str = Field(default_factory=lambda: f"pool-{random_uuid()}")
1841
- object: str = "list"
1842
- created: int = Field(default_factory=lambda: int(time.time()))
1843
- model: str
1844
- data: list[PoolingResponseData]
1845
- usage: UsageInfo
1846
-
1847
-
1848
-class PoolingBytesResponse(OpenAIBaseModel):
1849
- body: list[bytes]
1850
- metadata: str
1851
- media_type: str = "application/octet-stream"
1852
-
1853
-
1854
-class ScoreResponseData(OpenAIBaseModel):
1855
- index: int
1856
- object: str = "score"
1857
- score: float
1858
-
1859
-
1860
-class ScoreResponse(OpenAIBaseModel):
1861
- id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
1862
- object: str = "list"
1863
- created: int = Field(default_factory=lambda: int(time.time()))
1864
- model: str
1865
- data: list[ScoreResponseData]
1866
- usage: UsageInfo
1867
-
1868
-
1869
-class ClassificationCompletionRequest(OpenAIBaseModel):
1870
- model: str | None = None
1871
- input: list[str] | str
1872
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1873
- user: str | None = None
1874
-
1875
- # --8<-- [start:classification-extra-params]
1876
- priority: int = Field(
1877
- default=0,
1878
- description=(
1879
- "The priority of the request (lower means earlier handling; "
1880
- "default: 0). Any priority other than 0 will raise an error "
1881
- "if the served model does not use priority scheduling."
1882
- ),
1883
- )
1884
- add_special_tokens: bool = Field(
1885
- default=True,
1886
- description=(
1887
- "If true (the default), special tokens (e.g. BOS) will be added to "
1888
- "the prompt."
1889
- ),
1890
- )
1891
- request_id: str = Field(
1892
- default_factory=random_uuid,
1893
- description=(
1894
- "The request_id related to this request. If the caller does "
1895
- "not set it, a random_uuid will be generated. This id is used "
1896
- "through out the inference process and return in response."
1897
- ),
1898
- )
1899
- softmax: bool | None = Field(
1900
- default=None,
1901
- description="softmax will be deprecated, please use use_activation instead.",
1902
- )
1903
-
1904
- activation: bool | None = Field(
1905
- default=None,
1906
- description="activation will be deprecated, please use use_activation instead.",
1907
- )
1908
-
1909
- use_activation: bool | None = Field(
1910
- default=None,
1911
- description="Whether to use activation for classification outputs. "
1912
- "Default is True.",
1913
- )
1914
- # --8<-- [end:classification-extra-params]
1915
-
1916
- def to_pooling_params(self):
1917
- return PoolingParams(
1918
- truncate_prompt_tokens=self.truncate_prompt_tokens,
1919
- use_activation=get_use_activation(self),
1920
- )
1921
-
1922
-
1923
-class ClassificationChatRequest(OpenAIBaseModel):
1924
- model: str | None = None
1925
- messages: list[ChatCompletionMessageParam]
1926
- truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
1927
- user: str | None = None
1928
-
1929
- # --8<-- [start:chat-classification-extra-params]
1930
- add_generation_prompt: bool = Field(
1931
- default=False,
1932
- description=(
1933
- "If true, the generation prompt will be added to the chat template. "
1934
- "This is a parameter used by chat template in tokenizer config of the "
1935
- "model."
1936
- ),
1937
- )
1938
-
1939
- add_special_tokens: bool = Field(
1940
- default=False,
1941
- description=(
1942
- "If true, special tokens (e.g. BOS) will be added to the prompt "
1943
- "on top of what is added by the chat template. "
1944
- "For most models, the chat template takes care of adding the "
1945
- "special tokens so this should be set to false (as is the "
1946
- "default)."
1947
- ),
1948
- )
1949
-
1950
- chat_template: str | None = Field(
1951
- default=None,
1952
- description=(
1953
- "A Jinja template to use for this conversion. "
1954
- "As of transformers v4.44, default chat template is no longer "
1955
- "allowed, so you must provide a chat template if the tokenizer "
1956
- "does not define one."
1957
- ),
1958
- )
1959
-
1960
- chat_template_kwargs: dict[str, Any] | None = Field(
1961
- default=None,
1962
- description=(
1963
- "Additional keyword args to pass to the template renderer. "
1964
- "Will be accessible by the chat template."
1965
- ),
1966
- )
1967
-
1968
- mm_processor_kwargs: dict[str, Any] | None = Field(
1969
- default=None,
1970
- description=("Additional kwargs to pass to the HF processor."),
1971
- )
1972
-
1973
- priority: int = Field(
1974
- default=0,
1975
- description=(
1976
- "The priority of the request (lower means earlier handling; "
1977
- "default: 0). Any priority other than 0 will raise an error "
1978
- "if the served model does not use priority scheduling."
1979
- ),
1980
- )
1981
-
1982
- request_id: str = Field(
1983
- default_factory=random_uuid,
1984
- description=(
1985
- "The request_id related to this request. If the caller does "
1986
- "not set it, a random_uuid will be generated. This id is used "
1987
- "through out the inference process and return in response."
1988
- ),
1989
- )
1990
- softmax: bool | None = Field(
1991
- default=None,
1992
- description="softmax will be deprecated, please use use_activation instead.",
1993
- )
1994
-
1995
- activation: bool | None = Field(
1996
- default=None,
1997
- description="activation will be deprecated, please use use_activation instead.",
1998
- )
1999
-
2000
- use_activation: bool | None = Field(
2001
- default=None,
2002
- description="Whether to use activation for classification outputs. "
2003
- "Default is True.",
2004
- )
2005
- # --8<-- [end:chat-classification-extra-params]
2006
-
2007
- def to_pooling_params(self):
2008
- return PoolingParams(
2009
- truncate_prompt_tokens=self.truncate_prompt_tokens,
2010
- use_activation=get_use_activation(self),
2011
- )
2012
-
2013
-
2014
-ClassificationRequest: TypeAlias = (
2015
- ClassificationCompletionRequest | ClassificationChatRequest
2016
-)
2017
-
2018
-
2019
-class ClassificationData(OpenAIBaseModel):
2020
- index: int
2021
- label: str | None
2022
- probs: list[float]
2023
- num_classes: int
2024
-
2025
-
2026
-class ClassificationResponse(OpenAIBaseModel):
2027
- id: str = Field(default_factory=lambda: f"classify-{random_uuid()}")
2028
- object: str = "list"
2029
- created: int = Field(default_factory=lambda: int(time.time()))
2030
- model: str
2031
- data: list[ClassificationData]
2032
- usage: UsageInfo
2033
-
2034
-
1404
class FunctionCall(OpenAIBaseModel):
1405
name: str
1406
arguments: str
@@ -2409,83 +1778,6 @@ StreamingResponsesResponse: TypeAlias = (
1778
| ResponseCodeInterpreterCallCompletedEvent
1779
)
1780
2412
-BatchRequestInputBody: TypeAlias = (
2413
- ChatCompletionRequest | EmbeddingRequest | ScoreRequest | RerankRequest
2414
-)
2415
-
2416
-
2417
-class BatchRequestInput(OpenAIBaseModel):
2418
- """
2419
- The per-line object of the batch input file.
2420
-
2421
- NOTE: Currently only the `/v1/chat/completions` endpoint is supported.
2422
- """
2423
-
2424
- # A developer-provided per-request id that will be used to match outputs to
2425
- # inputs. Must be unique for each request in a batch.
2426
- custom_id: str
2427
-
2428
- # The HTTP method to be used for the request. Currently only POST is
2429
- # supported.
2430
- method: str
2431
-
2432
- # The OpenAI API relative URL to be used for the request. Currently
2433
- # /v1/chat/completions is supported.
2434
- url: str
2435
-
2436
- # The parameters of the request.
2437
- body: BatchRequestInputBody
2438
-
2439
- @field_validator("body", mode="plain")
2440
- @classmethod
2441
- def check_type_for_url(cls, value: Any, info: ValidationInfo):
2442
- # Use url to disambiguate models
2443
- url: str = info.data["url"]
2444
- if url == "/v1/chat/completions":
2445
- return ChatCompletionRequest.model_validate(value)
2446
- if url == "/v1/embeddings":
2447
- return TypeAdapter(EmbeddingRequest).validate_python(value)
2448
- if url.endswith("/score"):
2449
- return ScoreRequest.model_validate(value)
2450
- if url.endswith("/rerank"):
2451
- return RerankRequest.model_validate(value)
2452
- return TypeAdapter(BatchRequestInputBody).validate_python(value)
2453
-
2454
-
2455
-class BatchResponseData(OpenAIBaseModel):
2456
- # HTTP status code of the response.
2457
- status_code: int = 200
2458
-
2459
- # An unique identifier for the API request.
2460
- request_id: str
2461
-
2462
- # The body of the response.
2463
- body: (
2464
- ChatCompletionResponse
2465
- | EmbeddingResponse
2466
- | ScoreResponse
2467
- | RerankResponse
2468
- | None
2469
- ) = None
2470
-
2471
-
2472
-class BatchRequestOutput(OpenAIBaseModel):
2473
- """
2474
- The per-line object of the batch output and error files
2475
- """
2476
-
2477
- id: str
2478
-
2479
- # A developer-provided per-request id that will be used to match outputs to
2480
- # inputs.
2481
- custom_id: str
2482
-
2483
- response: BatchResponseData | None
2484
-
2485
- # For requests that failed with a non-HTTP error, this will contain more
2486
- # information on the cause of the failure.
2487
- error: Any | None
2488
-
1781
1782
class TokenizeCompletionRequest(OpenAIBaseModel):
1783
model: str | None = None
vllm/entrypoints/openai/run_batch.py
+92
-8
@@ -7,29 +7,35 @@ from argparse import Namespace
7
from collections.abc import Awaitable, Callable
8
from http import HTTPStatus
9
from io import StringIO
10
+from typing import Any, TypeAlias
11
12
import aiohttp
13
import torch
14
from prometheus_client import start_http_server
15
+from pydantic import TypeAdapter, field_validator
16
+from pydantic_core.core_schema import ValidationInfo
17
from tqdm import tqdm
18
19
from vllm.engine.arg_utils import AsyncEngineArgs, optional_type
20
from vllm.engine.protocol import EngineClient
21
from vllm.entrypoints.logger import RequestLogger
22
from vllm.entrypoints.openai.protocol import (
20
- BatchRequestInput,
21
- BatchRequestOutput,
22
- BatchResponseData,
23
+ ChatCompletionRequest,
24
ChatCompletionResponse,
24
- EmbeddingResponse,
25
ErrorResponse,
26
- RerankResponse,
27
- ScoreResponse,
26
+ OpenAIBaseModel,
27
)
28
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
30
-from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
29
from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
32
-from vllm.entrypoints.openai.serving_score import ServingScores
30
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest, EmbeddingResponse
31
+from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
32
+from vllm.entrypoints.pooling.score.protocol import (
33
+ RerankRequest,
34
+ RerankResponse,
35
+ ScoreRequest,
36
+ ScoreResponse,
37
+)
38
+from vllm.entrypoints.pooling.score.serving import ServingScores
39
from vllm.logger import init_logger
40
from vllm.reasoning import ReasoningParserManager
41
from vllm.utils import random_uuid
@@ -39,6 +45,84 @@ from vllm.version import __version__ as VLLM_VERSION
45
logger = init_logger(__name__)
46
47
48
+BatchRequestInputBody: TypeAlias = (
49
+ ChatCompletionRequest | EmbeddingRequest | ScoreRequest | RerankRequest
50
+)
51
+
52
+
53
+class BatchRequestInput(OpenAIBaseModel):
54
+ """
55
+ The per-line object of the batch input file.
56
+
57
+ NOTE: Currently only the `/v1/chat/completions` endpoint is supported.
58
+ """
59
+
60
+ # A developer-provided per-request id that will be used to match outputs to
61
+ # inputs. Must be unique for each request in a batch.
62
+ custom_id: str
63
+
64
+ # The HTTP method to be used for the request. Currently only POST is
65
+ # supported.
66
+ method: str
67
+
68
+ # The OpenAI API relative URL to be used for the request. Currently
69
+ # /v1/chat/completions is supported.
70
+ url: str
71
+
72
+ # The parameters of the request.
73
+ body: BatchRequestInputBody
74
+
75
+ @field_validator("body", mode="plain")
76
+ @classmethod
77
+ def check_type_for_url(cls, value: Any, info: ValidationInfo):
78
+ # Use url to disambiguate models
79
+ url: str = info.data["url"]
80
+ if url == "/v1/chat/completions":
81
+ return ChatCompletionRequest.model_validate(value)
82
+ if url == "/v1/embeddings":
83
+ return TypeAdapter(EmbeddingRequest).validate_python(value)
84
+ if url.endswith("/score"):
85
+ return ScoreRequest.model_validate(value)
86
+ if url.endswith("/rerank"):
87
+ return RerankRequest.model_validate(value)
88
+ return TypeAdapter(BatchRequestInputBody).validate_python(value)
89
+
90
+
91
+class BatchResponseData(OpenAIBaseModel):
92
+ # HTTP status code of the response.
93
+ status_code: int = 200
94
+
95
+ # An unique identifier for the API request.
96
+ request_id: str
97
+
98
+ # The body of the response.
99
+ body: (
100
+ ChatCompletionResponse
101
+ | EmbeddingResponse
102
+ | ScoreResponse
103
+ | RerankResponse
104
+ | None
105
+ ) = None
106
+
107
+
108
+class BatchRequestOutput(OpenAIBaseModel):
109
+ """
110
+ The per-line object of the batch output and error files
111
+ """
112
+
113
+ id: str
114
+
115
+ # A developer-provided per-request id that will be used to match outputs to
116
+ # inputs.
117
+ custom_id: str
118
+
119
+ response: BatchResponseData | None
120
+
121
+ # For requests that failed with a non-HTTP error, this will contain more
122
+ # information on the cause of the failure.
123
+ error: Any | None
124
+
125
+
126
def make_arg_parser(parser: FlexibleArgumentParser):
127
parser.add_argument(
128
"-i",
vllm/entrypoints/openai/serving_engine.py
+22
-13
@@ -18,6 +18,28 @@ from pydantic import ConfigDict, TypeAdapter
18
from starlette.datastructures import Headers
19
from typing_extensions import TypeIs
20
21
+from vllm.entrypoints.pooling.classify.protocol import (
22
+ ClassificationChatRequest,
23
+ ClassificationCompletionRequest,
24
+ ClassificationRequest,
25
+ ClassificationResponse,
26
+)
27
+from vllm.entrypoints.pooling.embed.protocol import (
28
+ EmbeddingChatRequest,
29
+ EmbeddingCompletionRequest,
30
+ EmbeddingRequest,
31
+ EmbeddingResponse,
32
+)
33
+from vllm.entrypoints.pooling.pooling.protocol import (
34
+ IOProcessorRequest,
35
+ PoolingResponse,
36
+)
37
+from vllm.entrypoints.pooling.score.protocol import (
38
+ RerankRequest,
39
+ ScoreRequest,
40
+ ScoreResponse,
41
+)
42
+
43
if sys.version_info >= (3, 12):
44
from typing import TypedDict
45
else:
@@ -45,29 +67,16 @@ from vllm.entrypoints.openai.protocol import (
67
ChatCompletionNamedToolChoiceParam,
68
ChatCompletionRequest,
69
ChatCompletionResponse,
48
- ClassificationChatRequest,
49
- ClassificationCompletionRequest,
50
- ClassificationRequest,
51
- ClassificationResponse,
70
CompletionRequest,
71
CompletionResponse,
72
DetokenizeRequest,
55
- EmbeddingChatRequest,
56
- EmbeddingCompletionRequest,
57
- EmbeddingRequest,
58
- EmbeddingResponse,
73
ErrorInfo,
74
ErrorResponse,
75
FunctionCall,
76
FunctionDefinition,
77
GenerateRequest,
78
GenerateResponse,
65
- IOProcessorRequest,
66
- PoolingResponse,
67
- RerankRequest,
79
ResponsesRequest,
69
- ScoreRequest,
70
- ScoreResponse,
80
TokenizeChatRequest,
81
TokenizeCompletionRequest,
82
TokenizeResponse,
vllm/entrypoints/openai/utils.py
+12
@@ -2,6 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import TypeVar
4
5
+from fastapi import Request
6
+from fastapi.exceptions import RequestValidationError
7
+
8
from vllm.entrypoints.openai.protocol import (
9
ChatCompletionRequest,
10
ChatCompletionResponseChoice,
@@ -35,3 +38,12 @@ def maybe_filter_parallel_tool_calls(
38
]
39
40
return choice
41
+
42
+
43
+async def validate_json_request(raw_request: Request):
44
+ content_type = raw_request.headers.get("content-type", "").lower()
45
+ media_type = content_type.split(";", maxsplit=1)[0]
46
+ if media_type != "application/json":
47
+ raise RequestValidationError(
48
+ errors=["Unsupported Media Type: Only 'application/json' is allowed"]
49
+ )
vllm/entrypoints/pooling/__init__.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from fastapi import FastAPI
5
+
6
+
7
+def register_pooling_api_routers(app: FastAPI):
8
+ from vllm.entrypoints.pooling.classify.api_router import router as classify_router
9
+ from vllm.entrypoints.pooling.embed.api_router import router as embed_router
10
+ from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router
11
+ from vllm.entrypoints.pooling.score.api_router import router as score_router
12
+
13
+ app.include_router(classify_router)
14
+ app.include_router(embed_router)
15
+ app.include_router(score_router)
16
+ app.include_router(pooling_router)
vllm/entrypoints/pooling/classify/__init__.py
vllm/entrypoints/pooling/classify/api_router.py
new
+50
@@ -0,0 +1,50 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from http import HTTPStatus
4
+
5
+from fastapi import APIRouter, Depends, HTTPException, Request
6
+from starlette.responses import JSONResponse
7
+from typing_extensions import assert_never
8
+
9
+from vllm.entrypoints.openai.protocol import ErrorResponse
10
+from vllm.entrypoints.openai.utils import validate_json_request
11
+from vllm.entrypoints.pooling.classify.protocol import (
12
+ ClassificationRequest,
13
+ ClassificationResponse,
14
+)
15
+from vllm.entrypoints.pooling.classify.serving import ServingClassification
16
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
17
+
18
+router = APIRouter()
19
+
20
+
21
+def classify(request: Request) -> ServingClassification | None:
22
+ return request.app.state.openai_serving_classification
23
+
24
+
25
+@router.post("/classify", dependencies=[Depends(validate_json_request)])
26
+@with_cancellation
27
+@load_aware_call
28
+async def create_classify(request: ClassificationRequest, raw_request: Request):
29
+ handler = classify(raw_request)
30
+ if handler is None:
31
+ base_server = raw_request.app.state.openai_serving_tokenization
32
+ return base_server.create_error_response(
33
+ message="The model does not support Classification API"
34
+ )
35
+
36
+ try:
37
+ generator = await handler.create_classify(request, raw_request)
38
+ except Exception as e:
39
+ raise HTTPException(
40
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
41
+ ) from e
42
+ if isinstance(generator, ErrorResponse):
43
+ return JSONResponse(
44
+ content=generator.model_dump(), status_code=generator.error.code
45
+ )
46
+
47
+ elif isinstance(generator, ClassificationResponse):
48
+ return JSONResponse(content=generator.model_dump())
49
+
50
+ assert_never(generator)
vllm/entrypoints/pooling/classify/protocol.py
new
+181
@@ -0,0 +1,181 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import time
5
+from typing import Annotated, Any, TypeAlias
6
+
7
+from pydantic import (
8
+ Field,
9
+)
10
+
11
+from vllm import PoolingParams
12
+from vllm.config.pooler import get_use_activation
13
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
14
+from vllm.entrypoints.openai.protocol import OpenAIBaseModel, UsageInfo
15
+from vllm.utils import random_uuid
16
+
17
+
18
+class ClassificationCompletionRequest(OpenAIBaseModel):
19
+ model: str | None = None
20
+ input: list[str] | str
21
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
22
+ user: str | None = None
23
+
24
+ # --8<-- [start:classification-extra-params]
25
+ priority: int = Field(
26
+ default=0,
27
+ description=(
28
+ "The priority of the request (lower means earlier handling; "
29
+ "default: 0). Any priority other than 0 will raise an error "
30
+ "if the served model does not use priority scheduling."
31
+ ),
32
+ )
33
+ add_special_tokens: bool = Field(
34
+ default=True,
35
+ description=(
36
+ "If true (the default), special tokens (e.g. BOS) will be added to "
37
+ "the prompt."
38
+ ),
39
+ )
40
+ request_id: str = Field(
41
+ default_factory=random_uuid,
42
+ description=(
43
+ "The request_id related to this request. If the caller does "
44
+ "not set it, a random_uuid will be generated. This id is used "
45
+ "through out the inference process and return in response."
46
+ ),
47
+ )
48
+ softmax: bool | None = Field(
49
+ default=None,
50
+ description="softmax will be deprecated, please use use_activation instead.",
51
+ )
52
+
53
+ activation: bool | None = Field(
54
+ default=None,
55
+ description="activation will be deprecated, please use use_activation instead.",
56
+ )
57
+
58
+ use_activation: bool | None = Field(
59
+ default=None,
60
+ description="Whether to use activation for classification outputs. "
61
+ "Default is True.",
62
+ )
63
+ # --8<-- [end:classification-extra-params]
64
+
65
+ def to_pooling_params(self):
66
+ return PoolingParams(
67
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
68
+ use_activation=get_use_activation(self),
69
+ )
70
+
71
+
72
+class ClassificationChatRequest(OpenAIBaseModel):
73
+ model: str | None = None
74
+ messages: list[ChatCompletionMessageParam]
75
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
76
+ user: str | None = None
77
+
78
+ # --8<-- [start:chat-classification-extra-params]
79
+ add_generation_prompt: bool = Field(
80
+ default=False,
81
+ description=(
82
+ "If true, the generation prompt will be added to the chat template. "
83
+ "This is a parameter used by chat template in tokenizer config of the "
84
+ "model."
85
+ ),
86
+ )
87
+
88
+ add_special_tokens: bool = Field(
89
+ default=False,
90
+ description=(
91
+ "If true, special tokens (e.g. BOS) will be added to the prompt "
92
+ "on top of what is added by the chat template. "
93
+ "For most models, the chat template takes care of adding the "
94
+ "special tokens so this should be set to false (as is the "
95
+ "default)."
96
+ ),
97
+ )
98
+
99
+ chat_template: str | None = Field(
100
+ default=None,
101
+ description=(
102
+ "A Jinja template to use for this conversion. "
103
+ "As of transformers v4.44, default chat template is no longer "
104
+ "allowed, so you must provide a chat template if the tokenizer "
105
+ "does not define one."
106
+ ),
107
+ )
108
+
109
+ chat_template_kwargs: dict[str, Any] | None = Field(
110
+ default=None,
111
+ description=(
112
+ "Additional keyword args to pass to the template renderer. "
113
+ "Will be accessible by the chat template."
114
+ ),
115
+ )
116
+
117
+ mm_processor_kwargs: dict[str, Any] | None = Field(
118
+ default=None,
119
+ description=("Additional kwargs to pass to the HF processor."),
120
+ )
121
+
122
+ priority: int = Field(
123
+ default=0,
124
+ description=(
125
+ "The priority of the request (lower means earlier handling; "
126
+ "default: 0). Any priority other than 0 will raise an error "
127
+ "if the served model does not use priority scheduling."
128
+ ),
129
+ )
130
+
131
+ request_id: str = Field(
132
+ default_factory=random_uuid,
133
+ description=(
134
+ "The request_id related to this request. If the caller does "
135
+ "not set it, a random_uuid will be generated. This id is used "
136
+ "through out the inference process and return in response."
137
+ ),
138
+ )
139
+ softmax: bool | None = Field(
140
+ default=None,
141
+ description="softmax will be deprecated, please use use_activation instead.",
142
+ )
143
+
144
+ activation: bool | None = Field(
145
+ default=None,
146
+ description="activation will be deprecated, please use use_activation instead.",
147
+ )
148
+
149
+ use_activation: bool | None = Field(
150
+ default=None,
151
+ description="Whether to use activation for classification outputs. "
152
+ "Default is True.",
153
+ )
154
+ # --8<-- [end:chat-classification-extra-params]
155
+
156
+ def to_pooling_params(self):
157
+ return PoolingParams(
158
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
159
+ use_activation=get_use_activation(self),
160
+ )
161
+
162
+
163
+ClassificationRequest: TypeAlias = (
164
+ ClassificationCompletionRequest | ClassificationChatRequest
165
+)
166
+
167
+
168
+class ClassificationData(OpenAIBaseModel):
169
+ index: int
170
+ label: str | None
171
+ probs: list[float]
172
+ num_classes: int
173
+
174
+
175
+class ClassificationResponse(OpenAIBaseModel):
176
+ id: str = Field(default_factory=lambda: f"classify-{random_uuid()}")
177
+ object: str = "list"
178
+ created: int = Field(default_factory=lambda: int(time.time()))
179
+ model: str
180
+ data: list[ClassificationData]
181
+ usage: UsageInfo
vllm/entrypoints/pooling/classify/serving.py
renamed
+7
-5
@@ -13,11 +13,6 @@ from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
13
from vllm.entrypoints.logger import RequestLogger
14
from vllm.entrypoints.openai.protocol import (
15
ChatCompletionRequest,
16
- ClassificationChatRequest,
17
- ClassificationCompletionRequest,
18
- ClassificationData,
19
- ClassificationRequest,
20
- ClassificationResponse,
16
ErrorResponse,
17
UsageInfo,
18
)
@@ -27,6 +22,13 @@ from vllm.entrypoints.openai.serving_engine import (
22
ServeContext,
23
)
24
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
25
+from vllm.entrypoints.pooling.classify.protocol import (
26
+ ClassificationChatRequest,
27
+ ClassificationCompletionRequest,
28
+ ClassificationData,
29
+ ClassificationRequest,
30
+ ClassificationResponse,
31
+)
32
from vllm.entrypoints.renderer import RenderConfig
33
from vllm.logger import init_logger
34
from vllm.outputs import ClassificationOutput, PoolingRequestOutput
vllm/entrypoints/pooling/embed/__init__.py
vllm/entrypoints/pooling/embed/api_router.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from http import HTTPStatus
4
+
5
+from fastapi import APIRouter, Depends, HTTPException, Request
6
+from fastapi.responses import JSONResponse, StreamingResponse
7
+from typing_extensions import assert_never
8
+
9
+from vllm.entrypoints.openai.protocol import ErrorResponse
10
+from vllm.entrypoints.openai.utils import validate_json_request
11
+from vllm.entrypoints.pooling.embed.protocol import (
12
+ EmbeddingBytesResponse,
13
+ EmbeddingRequest,
14
+ EmbeddingResponse,
15
+)
16
+from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
17
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
18
+
19
+router = APIRouter()
20
+
21
+
22
+def embedding(request: Request) -> OpenAIServingEmbedding | None:
23
+ return request.app.state.openai_serving_embedding
24
+
25
+
26
+@router.post(
27
+ "/v1/embeddings",
28
+ dependencies=[Depends(validate_json_request)],
29
+ responses={
30
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
31
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
32
+ },
33
+)
34
+@with_cancellation
35
+@load_aware_call
36
+async def create_embedding(
37
+ request: EmbeddingRequest,
38
+ raw_request: Request,
39
+):
40
+ handler = embedding(raw_request)
41
+ if handler is None:
42
+ base_server = raw_request.app.state.openai_serving_tokenization
43
+ return base_server.create_error_response(
44
+ message="The model does not support Embeddings API"
45
+ )
46
+
47
+ try:
48
+ generator = await handler.create_embedding(request, raw_request)
49
+ except Exception as e:
50
+ raise HTTPException(
51
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
52
+ ) from e
53
+
54
+ if isinstance(generator, ErrorResponse):
55
+ return JSONResponse(
56
+ content=generator.model_dump(), status_code=generator.error.code
57
+ )
58
+ elif isinstance(generator, EmbeddingResponse):
59
+ return JSONResponse(content=generator.model_dump())
60
+ elif isinstance(generator, EmbeddingBytesResponse):
61
+ return StreamingResponse(
62
+ content=generator.body,
63
+ headers={"metadata": generator.metadata},
64
+ media_type=generator.media_type,
65
+ )
66
+
67
+ assert_never(generator)
vllm/entrypoints/pooling/embed/protocol.py
new
+208
@@ -0,0 +1,208 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+from typing import Annotated, Any, TypeAlias
5
+
6
+from pydantic import (
7
+ Field,
8
+ model_validator,
9
+)
10
+
11
+from vllm import PoolingParams
12
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
13
+from vllm.entrypoints.openai.protocol import OpenAIBaseModel, UsageInfo
14
+from vllm.utils import random_uuid
15
+from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
16
+
17
+
18
+class EmbeddingCompletionRequest(OpenAIBaseModel):
19
+ # Ordered by official OpenAI API documentation
20
+ # https://platform.openai.com/docs/api-reference/embeddings
21
+ model: str | None = None
22
+ input: list[int] | list[list[int]] | str | list[str]
23
+ encoding_format: EncodingFormat = "float"
24
+ dimensions: int | None = None
25
+ user: str | None = None
26
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
27
+
28
+ # --8<-- [start:embedding-extra-params]
29
+ add_special_tokens: bool = Field(
30
+ default=True,
31
+ description=(
32
+ "If true (the default), special tokens (e.g. BOS) will be added to "
33
+ "the prompt."
34
+ ),
35
+ )
36
+ priority: int = Field(
37
+ default=0,
38
+ description=(
39
+ "The priority of the request (lower means earlier handling; "
40
+ "default: 0). Any priority other than 0 will raise an error "
41
+ "if the served model does not use priority scheduling."
42
+ ),
43
+ )
44
+ request_id: str = Field(
45
+ default_factory=random_uuid,
46
+ description=(
47
+ "The request_id related to this request. If the caller does "
48
+ "not set it, a random_uuid will be generated. This id is used "
49
+ "through out the inference process and return in response."
50
+ ),
51
+ )
52
+ normalize: bool | None = Field(
53
+ default=None,
54
+ description="Whether to normalize the embeddings outputs. Default is True.",
55
+ )
56
+ embed_dtype: EmbedDType = Field(
57
+ default="float32",
58
+ description=(
59
+ "What dtype to use for encoding. Default to using float32 for base64 "
60
+ "encoding to match the OpenAI python client behavior. "
61
+ "This parameter will affect base64 and binary_response."
62
+ ),
63
+ )
64
+ endianness: Endianness = Field(
65
+ default="native",
66
+ description=(
67
+ "What endianness to use for encoding. Default to using native for "
68
+ "base64 encoding to match the OpenAI python client behavior."
69
+ "This parameter will affect base64 and binary_response."
70
+ ),
71
+ )
72
+ # --8<-- [end:embedding-extra-params]
73
+
74
+ def to_pooling_params(self):
75
+ return PoolingParams(
76
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
77
+ dimensions=self.dimensions,
78
+ normalize=self.normalize,
79
+ )
80
+
81
+
82
+class EmbeddingChatRequest(OpenAIBaseModel):
83
+ model: str | None = None
84
+ messages: list[ChatCompletionMessageParam]
85
+
86
+ encoding_format: EncodingFormat = "float"
87
+ dimensions: int | None = None
88
+ user: str | None = None
89
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
90
+
91
+ # --8<-- [start:chat-embedding-extra-params]
92
+ add_generation_prompt: bool = Field(
93
+ default=False,
94
+ description=(
95
+ "If true, the generation prompt will be added to the chat template. "
96
+ "This is a parameter used by chat template in tokenizer config of the "
97
+ "model."
98
+ ),
99
+ )
100
+
101
+ add_special_tokens: bool = Field(
102
+ default=False,
103
+ description=(
104
+ "If true, special tokens (e.g. BOS) will be added to the prompt "
105
+ "on top of what is added by the chat template. "
106
+ "For most models, the chat template takes care of adding the "
107
+ "special tokens so this should be set to false (as is the "
108
+ "default)."
109
+ ),
110
+ )
111
+ chat_template: str | None = Field(
112
+ default=None,
113
+ description=(
114
+ "A Jinja template to use for this conversion. "
115
+ "As of transformers v4.44, default chat template is no longer "
116
+ "allowed, so you must provide a chat template if the tokenizer "
117
+ "does not define one."
118
+ ),
119
+ )
120
+ chat_template_kwargs: dict[str, Any] | None = Field(
121
+ default=None,
122
+ description=(
123
+ "Additional keyword args to pass to the template renderer. "
124
+ "Will be accessible by the chat template."
125
+ ),
126
+ )
127
+ mm_processor_kwargs: dict[str, Any] | None = Field(
128
+ default=None,
129
+ description=("Additional kwargs to pass to the HF processor."),
130
+ )
131
+ priority: int = Field(
132
+ default=0,
133
+ description=(
134
+ "The priority of the request (lower means earlier handling; "
135
+ "default: 0). Any priority other than 0 will raise an error "
136
+ "if the served model does not use priority scheduling."
137
+ ),
138
+ )
139
+ request_id: str = Field(
140
+ default_factory=random_uuid,
141
+ description=(
142
+ "The request_id related to this request. If the caller does "
143
+ "not set it, a random_uuid will be generated. This id is used "
144
+ "through out the inference process and return in response."
145
+ ),
146
+ )
147
+ normalize: bool | None = Field(
148
+ default=None,
149
+ description="Whether to normalize the embeddings outputs. Default is True.",
150
+ )
151
+ embed_dtype: EmbedDType = Field(
152
+ default="float32",
153
+ description=(
154
+ "What dtype to use for encoding. Default to using float32 for base64 "
155
+ "encoding to match the OpenAI python client behavior. "
156
+ "This parameter will affect base64 and binary_response."
157
+ ),
158
+ )
159
+ endianness: Endianness = Field(
160
+ default="native",
161
+ description=(
162
+ "What endianness to use for encoding. Default to using native for "
163
+ "base64 encoding to match the OpenAI python client behavior."
164
+ "This parameter will affect base64 and binary_response."
165
+ ),
166
+ )
167
+ # --8<-- [end:chat-embedding-extra-params]
168
+
169
+ @model_validator(mode="before")
170
+ @classmethod
171
+ def check_generation_prompt(cls, data):
172
+ if data.get("continue_final_message") and data.get("add_generation_prompt"):
173
+ raise ValueError(
174
+ "Cannot set both `continue_final_message` and "
175
+ "`add_generation_prompt` to True."
176
+ )
177
+ return data
178
+
179
+ def to_pooling_params(self):
180
+ return PoolingParams(
181
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
182
+ dimensions=self.dimensions,
183
+ normalize=self.normalize,
184
+ )
185
+
186
+
187
+EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
188
+
189
+
190
+class EmbeddingResponseData(OpenAIBaseModel):
191
+ index: int
192
+ object: str = "embedding"
193
+ embedding: list[float] | str
194
+
195
+
196
+class EmbeddingResponse(OpenAIBaseModel):
197
+ id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
198
+ object: str = "list"
199
+ created: int = Field(default_factory=lambda: int(time.time()))
200
+ model: str
201
+ data: list[EmbeddingResponseData]
202
+ usage: UsageInfo
203
+
204
+
205
+class EmbeddingBytesResponse(OpenAIBaseModel):
206
+ body: list[bytes]
207
+ metadata: str
208
+ media_type: str = "application/octet-stream"
vllm/entrypoints/pooling/embed/serving.py
renamed
+8
-6
@@ -13,12 +13,6 @@ from vllm.engine.protocol import EngineClient
13
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
14
from vllm.entrypoints.logger import RequestLogger
15
from vllm.entrypoints.openai.protocol import (
16
- EmbeddingBytesResponse,
17
- EmbeddingChatRequest,
18
- EmbeddingCompletionRequest,
19
- EmbeddingRequest,
20
- EmbeddingResponse,
21
- EmbeddingResponseData,
16
ErrorResponse,
17
UsageInfo,
18
)
@@ -29,6 +23,14 @@ from vllm.entrypoints.openai.serving_engine import (
23
TextTokensPrompt,
24
)
25
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
26
+from vllm.entrypoints.pooling.embed.protocol import (
27
+ EmbeddingBytesResponse,
28
+ EmbeddingChatRequest,
29
+ EmbeddingCompletionRequest,
30
+ EmbeddingRequest,
31
+ EmbeddingResponse,
32
+ EmbeddingResponseData,
33
+)
34
from vllm.entrypoints.renderer import RenderConfig
35
from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
36
from vllm.logger import init_logger
vllm/entrypoints/pooling/pooling/__init__.py
vllm/entrypoints/pooling/pooling/api_router.py
new
+63
@@ -0,0 +1,63 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from http import HTTPStatus
4
+
5
+from fastapi import APIRouter, Depends, HTTPException, Request
6
+from fastapi.responses import JSONResponse, StreamingResponse
7
+from typing_extensions import assert_never
8
+
9
+from vllm.entrypoints.openai.protocol import ErrorResponse
10
+from vllm.entrypoints.openai.utils import validate_json_request
11
+from vllm.entrypoints.pooling.pooling.protocol import (
12
+ IOProcessorResponse,
13
+ PoolingBytesResponse,
14
+ PoolingRequest,
15
+ PoolingResponse,
16
+)
17
+from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
18
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
19
+
20
+router = APIRouter()
21
+
22
+
23
+def pooling(request: Request) -> OpenAIServingPooling | None:
24
+ return request.app.state.openai_serving_pooling
25
+
26
+
27
+@router.post(
28
+ "/pooling",
29
+ dependencies=[Depends(validate_json_request)],
30
+ responses={
31
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
32
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
33
+ },
34
+)
35
+@with_cancellation
36
+@load_aware_call
37
+async def create_pooling(request: PoolingRequest, raw_request: Request):
38
+ handler = pooling(raw_request)
39
+ if handler is None:
40
+ base_server = raw_request.app.state.openai_serving_tokenization
41
+ return base_server.create_error_response(
42
+ message="The model does not support Pooling API"
43
+ )
44
+ try:
45
+ generator = await handler.create_pooling(request, raw_request)
46
+ except Exception as e:
47
+ raise HTTPException(
48
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
49
+ ) from e
50
+ if isinstance(generator, ErrorResponse):
51
+ return JSONResponse(
52
+ content=generator.model_dump(), status_code=generator.error.code
53
+ )
54
+ elif isinstance(generator, (PoolingResponse, IOProcessorResponse)):
55
+ return JSONResponse(content=generator.model_dump())
56
+ elif isinstance(generator, PoolingBytesResponse):
57
+ return StreamingResponse(
58
+ content=generator.body,
59
+ headers={"metadata": generator.metadata},
60
+ media_type=generator.media_type,
61
+ )
62
+
63
+ assert_never(generator)
vllm/entrypoints/pooling/pooling/protocol.py
new
+148
@@ -0,0 +1,148 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+from typing import Generic, TypeAlias, TypeVar
5
+
6
+from pydantic import (
7
+ Field,
8
+)
9
+
10
+from vllm import PoolingParams
11
+from vllm.config.pooler import get_use_activation
12
+from vllm.entrypoints.openai.protocol import OpenAIBaseModel, UsageInfo
13
+from vllm.entrypoints.pooling.embed.protocol import (
14
+ EmbeddingChatRequest,
15
+ EmbeddingCompletionRequest,
16
+)
17
+from vllm.tasks import PoolingTask
18
+from vllm.utils import random_uuid
19
+from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
20
+
21
+
22
+class PoolingCompletionRequest(EmbeddingCompletionRequest):
23
+ task: PoolingTask | None = None
24
+ softmax: bool | None = Field(
25
+ default=None,
26
+ description="softmax will be deprecated, please use use_activation instead.",
27
+ )
28
+ activation: bool | None = Field(
29
+ default=None,
30
+ description="activation will be deprecated, please use use_activation instead.",
31
+ )
32
+ use_activation: bool | None = Field(
33
+ default=None,
34
+ description="Whether to use activation for classification outputs. "
35
+ "If it is a classify or token_classify task, the default is True; "
36
+ "for other tasks, this value should be None.",
37
+ )
38
+
39
+ def to_pooling_params(self):
40
+ return PoolingParams(
41
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
42
+ dimensions=self.dimensions,
43
+ normalize=self.normalize,
44
+ use_activation=get_use_activation(self),
45
+ )
46
+
47
+
48
+class PoolingChatRequest(EmbeddingChatRequest):
49
+ task: PoolingTask | None = None
50
+ softmax: bool | None = Field(
51
+ default=None,
52
+ description="softmax will be deprecated, please use use_activation instead.",
53
+ )
54
+ activation: bool | None = Field(
55
+ default=None,
56
+ description="activation will be deprecated, please use use_activation instead.",
57
+ )
58
+ use_activation: bool | None = Field(
59
+ default=None,
60
+ description="Whether to use activation for classification outputs. "
61
+ "If it is a classify or token_classify task, the default is True; "
62
+ "for other tasks, this value should be None.",
63
+ )
64
+
65
+ def to_pooling_params(self):
66
+ return PoolingParams(
67
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
68
+ dimensions=self.dimensions,
69
+ normalize=self.normalize,
70
+ use_activation=get_use_activation(self),
71
+ )
72
+
73
+
74
+T = TypeVar("T")
75
+
76
+
77
+class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
78
+ model: str | None = None
79
+
80
+ priority: int = Field(default=0)
81
+ """
82
+ The priority of the request (lower means earlier handling;
83
+ default: 0). Any priority other than 0 will raise an error
84
+ if the served model does not use priority scheduling.
85
+ """
86
+ data: T
87
+
88
+ task: PoolingTask = "plugin"
89
+ encoding_format: EncodingFormat = "float"
90
+ embed_dtype: EmbedDType = Field(
91
+ default="float32",
92
+ description=(
93
+ "What dtype to use for encoding. Default to using float32 for base64 "
94
+ "encoding to match the OpenAI python client behavior. "
95
+ "This parameter will affect base64 and binary_response."
96
+ ),
97
+ )
98
+ endianness: Endianness = Field(
99
+ default="native",
100
+ description=(
101
+ "What endianness to use for encoding. Default to using native for "
102
+ "base64 encoding to match the OpenAI python client behavior."
103
+ "This parameter will affect base64 and binary_response."
104
+ ),
105
+ )
106
+
107
+ def to_pooling_params(self):
108
+ return PoolingParams()
109
+
110
+
111
+class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
112
+ request_id: str | None = None
113
+ """
114
+ The request_id associated with this response
115
+ """
116
+ created_at: int = Field(default_factory=lambda: int(time.time()))
117
+
118
+ data: T
119
+ """
120
+ When using plugins IOProcessor plugins, the actual output is generated
121
+ by the plugin itself. Hence, we use a generic type for the response data
122
+ """
123
+
124
+
125
+PoolingRequest: TypeAlias = (
126
+ PoolingCompletionRequest | PoolingChatRequest | IOProcessorRequest
127
+)
128
+
129
+
130
+class PoolingResponseData(OpenAIBaseModel):
131
+ index: int
132
+ object: str = "pooling"
133
+ data: list[list[float]] | list[float] | str
134
+
135
+
136
+class PoolingResponse(OpenAIBaseModel):
137
+ id: str = Field(default_factory=lambda: f"pool-{random_uuid()}")
138
+ object: str = "list"
139
+ created: int = Field(default_factory=lambda: int(time.time()))
140
+ model: str
141
+ data: list[PoolingResponseData]
142
+ usage: UsageInfo
143
+
144
+
145
+class PoolingBytesResponse(OpenAIBaseModel):
146
+ body: list[bytes]
147
+ metadata: str
148
+ media_type: str = "application/octet-stream"
vllm/entrypoints/pooling/pooling/serving.py
renamed
+5
-3
@@ -16,6 +16,11 @@ from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
16
from vllm.entrypoints.logger import RequestLogger
17
from vllm.entrypoints.openai.protocol import (
18
ErrorResponse,
19
+ UsageInfo,
20
+)
21
+from vllm.entrypoints.openai.serving_engine import OpenAIServing
22
+from vllm.entrypoints.openai.serving_models import OpenAIServingModels
23
+from vllm.entrypoints.pooling.pooling.protocol import (
24
IOProcessorRequest,
25
IOProcessorResponse,
26
PoolingBytesResponse,
@@ -24,10 +29,7 @@ from vllm.entrypoints.openai.protocol import (
29
PoolingRequest,
30
PoolingResponse,
31
PoolingResponseData,
27
- UsageInfo,
32
)
29
-from vllm.entrypoints.openai.serving_engine import OpenAIServing
30
-from vllm.entrypoints.openai.serving_models import OpenAIServingModels
33
from vllm.entrypoints.renderer import RenderConfig
34
from vllm.entrypoints.utils import _validate_truncation_size
35
from vllm.logger import init_logger
vllm/entrypoints/pooling/score/__init__.py
vllm/entrypoints/pooling/score/api_router.py
new
+149
@@ -0,0 +1,149 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from http import HTTPStatus
4
+
5
+from fastapi import APIRouter, Depends, HTTPException, Request
6
+from fastapi.responses import JSONResponse
7
+from typing_extensions import assert_never
8
+
9
+from vllm.entrypoints.openai.protocol import ErrorResponse
10
+from vllm.entrypoints.openai.utils import validate_json_request
11
+from vllm.entrypoints.pooling.score.protocol import (
12
+ RerankRequest,
13
+ RerankResponse,
14
+ ScoreRequest,
15
+ ScoreResponse,
16
+)
17
+from vllm.entrypoints.pooling.score.serving import ServingScores
18
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
19
+from vllm.logger import init_logger
20
+
21
+router = APIRouter()
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+def score(request: Request) -> ServingScores | None:
27
+ return request.app.state.openai_serving_scores
28
+
29
+
30
+def rerank(request: Request) -> ServingScores | None:
31
+ return request.app.state.openai_serving_scores
32
+
33
+
34
+@router.post(
35
+ "/score",
36
+ dependencies=[Depends(validate_json_request)],
37
+ responses={
38
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
39
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
40
+ },
41
+)
42
+@with_cancellation
43
+@load_aware_call
44
+async def create_score(request: ScoreRequest, raw_request: Request):
45
+ handler = score(raw_request)
46
+ if handler is None:
47
+ base_server = raw_request.app.state.openai_serving_tokenization
48
+ return base_server.create_error_response(
49
+ message="The model does not support Score API"
50
+ )
51
+
52
+ try:
53
+ generator = await handler.create_score(request, raw_request)
54
+ except Exception as e:
55
+ raise HTTPException(
56
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
57
+ ) from e
58
+ if isinstance(generator, ErrorResponse):
59
+ return JSONResponse(
60
+ content=generator.model_dump(), status_code=generator.error.code
61
+ )
62
+ elif isinstance(generator, ScoreResponse):
63
+ return JSONResponse(content=generator.model_dump())
64
+
65
+ assert_never(generator)
66
+
67
+
68
+@router.post(
69
+ "/v1/score",
70
+ dependencies=[Depends(validate_json_request)],
71
+ responses={
72
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
73
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
74
+ },
75
+)
76
+@with_cancellation
77
+@load_aware_call
78
+async def create_score_v1(request: ScoreRequest, raw_request: Request):
79
+ logger.warning(
80
+ "To indicate that Score API is not part of standard OpenAI API, we "
81
+ "have moved it to `/score`. Please update your client accordingly."
82
+ )
83
+
84
+ return await create_score(request, raw_request)
85
+
86
+
87
+@router.post(
88
+ "/rerank",
89
+ dependencies=[Depends(validate_json_request)],
90
+ responses={
91
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
92
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
93
+ },
94
+)
95
+@with_cancellation
96
+@load_aware_call
97
+async def do_rerank(request: RerankRequest, raw_request: Request):
98
+ handler = rerank(raw_request)
99
+ if handler is None:
100
+ base_server = raw_request.app.state.openai_serving_tokenization
101
+ return base_server.create_error_response(
102
+ message="The model does not support Rerank (Score) API"
103
+ )
104
+ try:
105
+ generator = await handler.do_rerank(request, raw_request)
106
+ except Exception as e:
107
+ raise HTTPException(
108
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
109
+ ) from e
110
+ if isinstance(generator, ErrorResponse):
111
+ return JSONResponse(
112
+ content=generator.model_dump(), status_code=generator.error.code
113
+ )
114
+ elif isinstance(generator, RerankResponse):
115
+ return JSONResponse(content=generator.model_dump())
116
+
117
+ assert_never(generator)
118
+
119
+
120
+@router.post(
121
+ "/v1/rerank",
122
+ dependencies=[Depends(validate_json_request)],
123
+ responses={
124
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
125
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
126
+ },
127
+)
128
+@with_cancellation
129
+async def do_rerank_v1(request: RerankRequest, raw_request: Request):
130
+ logger.warning_once(
131
+ "To indicate that the rerank API is not part of the standard OpenAI"
132
+ " API, we have located it at `/rerank`. Please update your client "
133
+ "accordingly. (Note: Conforms to JinaAI rerank API)"
134
+ )
135
+
136
+ return await do_rerank(request, raw_request)
137
+
138
+
139
+@router.post(
140
+ "/v2/rerank",
141
+ dependencies=[Depends(validate_json_request)],
142
+ responses={
143
+ HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
144
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
145
+ },
146
+)
147
+@with_cancellation
148
+async def do_rerank_v2(request: RerankRequest, raw_request: Request):
149
+ return await do_rerank(request, raw_request)
vllm/entrypoints/pooling/score/protocol.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+from typing import Annotated, Any
5
+
6
+from pydantic import (
7
+ BaseModel,
8
+ Field,
9
+)
10
+
11
+from vllm import PoolingParams
12
+from vllm.config.pooler import get_use_activation
13
+from vllm.entrypoints.openai.protocol import OpenAIBaseModel, UsageInfo
14
+from vllm.entrypoints.score_utils import ScoreContentPartParam, ScoreMultiModalParam
15
+from vllm.utils import random_uuid
16
+
17
+
18
+class ScoreRequest(OpenAIBaseModel):
19
+ model: str | None = None
20
+ text_1: list[str] | str | ScoreMultiModalParam
21
+ text_2: list[str] | str | ScoreMultiModalParam
22
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
23
+
24
+ # --8<-- [start:score-extra-params]
25
+
26
+ mm_processor_kwargs: dict[str, Any] | None = Field(
27
+ default=None,
28
+ description=("Additional kwargs to pass to the HF processor."),
29
+ )
30
+
31
+ priority: int = Field(
32
+ default=0,
33
+ description=(
34
+ "The priority of the request (lower means earlier handling; "
35
+ "default: 0). Any priority other than 0 will raise an error "
36
+ "if the served model does not use priority scheduling."
37
+ ),
38
+ )
39
+
40
+ softmax: bool | None = Field(
41
+ default=None,
42
+ description="softmax will be deprecated, please use use_activation instead.",
43
+ )
44
+
45
+ activation: bool | None = Field(
46
+ default=None,
47
+ description="activation will be deprecated, please use use_activation instead.",
48
+ )
49
+
50
+ use_activation: bool | None = Field(
51
+ default=None,
52
+ description="Whether to use activation for classification outputs. "
53
+ "Default is True.",
54
+ )
55
+ # --8<-- [end:score-extra-params]
56
+
57
+ def to_pooling_params(self):
58
+ return PoolingParams(
59
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
60
+ use_activation=get_use_activation(self),
61
+ )
62
+
63
+
64
+class RerankRequest(OpenAIBaseModel):
65
+ model: str | None = None
66
+ query: str | ScoreMultiModalParam
67
+ documents: list[str] | ScoreMultiModalParam
68
+ top_n: int = Field(default_factory=lambda: 0)
69
+ truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
70
+
71
+ # --8<-- [start:rerank-extra-params]
72
+
73
+ mm_processor_kwargs: dict[str, Any] | None = Field(
74
+ default=None,
75
+ description=("Additional kwargs to pass to the HF processor."),
76
+ )
77
+
78
+ priority: int = Field(
79
+ default=0,
80
+ description=(
81
+ "The priority of the request (lower means earlier handling; "
82
+ "default: 0). Any priority other than 0 will raise an error "
83
+ "if the served model does not use priority scheduling."
84
+ ),
85
+ )
86
+
87
+ softmax: bool | None = Field(
88
+ default=None,
89
+ description="softmax will be deprecated, please use use_activation instead.",
90
+ )
91
+
92
+ activation: bool | None = Field(
93
+ default=None,
94
+ description="activation will be deprecated, please use use_activation instead.",
95
+ )
96
+
97
+ use_activation: bool | None = Field(
98
+ default=None,
99
+ description="Whether to use activation for classification outputs. "
100
+ "Default is True.",
101
+ )
102
+ # --8<-- [end:rerank-extra-params]
103
+
104
+ def to_pooling_params(self):
105
+ return PoolingParams(
106
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
107
+ use_activation=get_use_activation(self),
108
+ )
109
+
110
+
111
+class RerankDocument(BaseModel):
112
+ text: str | None = None
113
+ multi_modal: ScoreContentPartParam | None = None
114
+
115
+
116
+class RerankResult(BaseModel):
117
+ index: int
118
+ document: RerankDocument
119
+ relevance_score: float
120
+
121
+
122
+class RerankUsage(BaseModel):
123
+ total_tokens: int
124
+
125
+
126
+class RerankResponse(OpenAIBaseModel):
127
+ id: str
128
+ model: str
129
+ usage: RerankUsage
130
+ results: list[RerankResult]
131
+
132
+
133
+class ScoreResponseData(OpenAIBaseModel):
134
+ index: int
135
+ object: str = "score"
136
+ score: float
137
+
138
+
139
+class ScoreResponse(OpenAIBaseModel):
140
+ id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
141
+ object: str = "list"
142
+ created: int = Field(default_factory=lambda: int(time.time()))
143
+ model: str
144
+ data: list[ScoreResponseData]
145
+ usage: UsageInfo
vllm/entrypoints/pooling/score/serving.py
renamed
+5
-3
@@ -11,6 +11,11 @@ from vllm.engine.protocol import EngineClient
11
from vllm.entrypoints.logger import RequestLogger
12
from vllm.entrypoints.openai.protocol import (
13
ErrorResponse,
14
+ UsageInfo,
15
+)
16
+from vllm.entrypoints.openai.serving_engine import OpenAIServing
17
+from vllm.entrypoints.openai.serving_models import OpenAIServingModels
18
+from vllm.entrypoints.pooling.score.protocol import (
19
RerankDocument,
20
RerankRequest,
21
RerankResponse,
@@ -19,10 +24,7 @@ from vllm.entrypoints.openai.protocol import (
24
ScoreRequest,
25
ScoreResponse,
26
ScoreResponseData,
22
- UsageInfo,
27
)
24
-from vllm.entrypoints.openai.serving_engine import OpenAIServing
25
-from vllm.entrypoints.openai.serving_models import OpenAIServingModels
28
from vllm.entrypoints.score_utils import (
29
ScoreContentPartParam,
30
ScoreMultiModalParam,
vllm/entrypoints/sagemaker/routes.py
+48
-2
@@ -1,7 +1,9 @@
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
@@ -9,12 +11,56 @@ from fastapi import APIRouter, Depends, HTTPException, Request
11
from fastapi.responses import JSONResponse, Response
12
13
from vllm.entrypoints.openai.api_server import (
12
- INVOCATION_VALIDATORS,
14
base,
15
+ chat,
16
+ completion,
17
+ create_chat_completion,
18
+ create_completion,
19
health,
20
validate_json_request,
21
)
17
-from vllm.entrypoints.openai.protocol import ErrorResponse
22
+from vllm.entrypoints.openai.protocol import (
23
+ ChatCompletionRequest,
24
+ CompletionRequest,
25
+ ErrorResponse,
26
+)
27
+from vllm.entrypoints.openai.serving_engine import OpenAIServing
28
+from vllm.entrypoints.pooling.classify.api_router import classify, create_classify
29
+from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
30
+from vllm.entrypoints.pooling.embed.api_router import create_embedding, embedding
31
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
32
+from vllm.entrypoints.pooling.pooling.api_router import create_pooling, pooling
33
+from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
34
+from vllm.entrypoints.pooling.score.api_router import (
35
+ create_score,
36
+ do_rerank,
37
+ rerank,
38
+ score,
39
+)
40
+from vllm.entrypoints.pooling.score.protocol import RerankRequest, ScoreRequest
41
+
42
+# TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
43
+# (requires typing_extensions >= 4.13)
44
+RequestType = Any
45
+GetHandlerFn = Callable[[Request], OpenAIServing | None]
46
+EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
47
+
48
+# NOTE: Items defined earlier take higher priority
49
+INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [
50
+ (ChatCompletionRequest, (chat, create_chat_completion)),
51
+ (CompletionRequest, (completion, create_completion)),
52
+ (EmbeddingRequest, (embedding, create_embedding)),
53
+ (ClassificationRequest, (classify, create_classify)),
54
+ (ScoreRequest, (score, create_score)),
55
+ (RerankRequest, (rerank, do_rerank)),
56
+ (PoolingRequest, (pooling, create_pooling)),
57
+]
58
+
59
+# NOTE: Construct the TypeAdapters only once
60
+INVOCATION_VALIDATORS = [
61
+ (pydantic.TypeAdapter(request_type), (get_handler, endpoint))
62
+ for request_type, (get_handler, endpoint) in INVOCATION_TYPES
63
+]
64
65
66
def register_sagemaker_routes(router: APIRouter):
vllm/plugins/io_processors/interface.py
+1
-1
@@ -6,7 +6,7 @@ from collections.abc import AsyncGenerator, Sequence
6
from typing import Any, Generic, TypeVar
7
8
from vllm.config import VllmConfig
9
-from vllm.entrypoints.openai.protocol import IOProcessorResponse
9
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
10
from vllm.inputs.data import PromptType
11
from vllm.outputs import PoolingRequestOutput
12
from vllm.pooling_params import PoolingParams