@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4ae77dfd42
Fork 0

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

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

Synced Nov 10, 01:30 UTC

[Frontend][1/n] Make pooling entrypoints request schema consensus | CompletionRequest (#32395)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed Jan 16, 2026 at 14:17 UTC 4ae77dfd42041dc2defe21f6ccf76aecb4478812
22 files changed +629 -594
docs/serving/openai_compatible_server.md
+1 -1
@@ -559,7 +559,7 @@ Our Classification API directly supports Hugging Face sequence-classification mo
559
560 We automatically wrap any other transformer via `as_seq_cls_model()`, which pools on the last token, attaches a `RowParallelLinear` head, and applies a softmax to produce per-class probabilities.
561
562 -Code example: [examples/pooling/classify/openai_classification_client.py](../../examples/pooling/classify/openai_classification_client.py)
562 +Code example: [examples/pooling/classify/classification_online.py](../../examples/pooling/classify/classification_online.py)
563
564 #### Example Requests
565
examples/pooling/classify/classification_online.py new
+67
@@ -0,0 +1,67 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Example Python client for classification API using vLLM API server
4 +NOTE:
5 + start a supported classification model server with `vllm serve`, e.g.
6 + vllm serve jason9693/Qwen2.5-1.5B-apeach
7 +"""
8 +
9 +import argparse
10 +import pprint
11 +
12 +import requests
13 +
14 +headers = {"accept": "application/json", "Content-Type": "application/json"}
15 +
16 +
17 +def parse_args():
18 + parse = argparse.ArgumentParser()
19 + parse.add_argument("--host", type=str, default="localhost")
20 + parse.add_argument("--port", type=int, default=8000)
21 + return parse.parse_args()
22 +
23 +
24 +def main(args):
25 + base_url = f"http://{args.host}:{args.port}"
26 + models_url = base_url + "/v1/models"
27 + classify_url = base_url + "/classify"
28 + tokenize_url = base_url + "/tokenize"
29 +
30 + response = requests.get(models_url, headers=headers)
31 + model = response.json()["data"][0]["id"]
32 +
33 + # /classify can accept str as input
34 + prompts = [
35 + "Hello, my name is",
36 + "The president of the United States is",
37 + "The capital of France is",
38 + "The future of AI is",
39 + ]
40 +
41 + payload = {
42 + "model": model,
43 + "input": prompts,
44 + }
45 + response = requests.post(classify_url, headers=headers, json=payload)
46 + pprint.pprint(response.json())
47 +
48 + # /classify can accept token ids as input
49 + token_ids = []
50 + for prompt in prompts:
51 + response = requests.post(
52 + tokenize_url,
53 + json={"model": model, "prompt": prompt},
54 + )
55 + token_ids.append(response.json()["tokens"])
56 +
57 + payload = {
58 + "model": model,
59 + "input": token_ids,
60 + }
61 + response = requests.post(classify_url, headers=headers, json=payload)
62 + pprint.pprint(response.json())
63 +
64 +
65 +if __name__ == "__main__":
66 + args = parse_args()
67 + main(args)
examples/pooling/classify/openai_classification_client.py deleted
-53
@@ -1,53 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Example Python client for classification API using vLLM API server
4 -NOTE:
5 - start a supported classification model server with `vllm serve`, e.g.
6 - vllm serve jason9693/Qwen2.5-1.5B-apeach
7 -"""
8 -
9 -import argparse
10 -import pprint
11 -
12 -import requests
13 -
14 -
15 -def post_http_request(payload: dict, api_url: str) -> requests.Response:
16 - headers = {"User-Agent": "Test Client"}
17 - response = requests.post(api_url, headers=headers, json=payload)
18 - return response
19 -
20 -
21 -def parse_args():
22 - parse = argparse.ArgumentParser()
23 - parse.add_argument("--host", type=str, default="localhost")
24 - parse.add_argument("--port", type=int, default=8000)
25 - parse.add_argument("--model", type=str, default="jason9693/Qwen2.5-1.5B-apeach")
26 - return parse.parse_args()
27 -
28 -
29 -def main(args):
30 - host = args.host
31 - port = args.port
32 - model_name = args.model
33 -
34 - api_url = f"http://{host}:{port}/classify"
35 - prompts = [
36 - "Hello, my name is",
37 - "The president of the United States is",
38 - "The capital of France is",
39 - "The future of AI is",
40 - ]
41 -
42 - payload = {
43 - "model": model_name,
44 - "input": prompts,
45 - }
46 -
47 - classify_response = post_http_request(payload=payload, api_url=api_url)
48 - pprint.pprint(classify_response.json())
49 -
50 -
51 -if __name__ == "__main__":
52 - args = parse_args()
53 - main(args)
examples/pooling/pooling/vision_language_pooling.py
+1 -1
@@ -16,7 +16,7 @@ from typing import Literal, NamedTuple, TypeAlias, TypedDict, get_args
16 from PIL.Image import Image
17
18 from vllm import LLM, EngineArgs
19 -from vllm.entrypoints.score_utils import ScoreMultiModalParam
19 +from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
20 from vllm.multimodal.utils import fetch_image
21 from vllm.utils.argparse_utils import FlexibleArgumentParser
22
examples/pooling/score/vision_reranker_offline.py
+1 -1
@@ -15,7 +15,7 @@ from pathlib import Path
15 from typing import NamedTuple
16
17 from vllm import LLM, EngineArgs
18 -from vllm.entrypoints.score_utils import ScoreMultiModalParam
18 +from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
19 from vllm.utils.argparse_utils import FlexibleArgumentParser
20
21 TEMPLATE_HOME = Path(__file__).parent / "template"
tests/entrypoints/pooling/classify/test_online.py
+86 -49
@@ -12,6 +12,8 @@ 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
15 +input_text = "This product was excellent and exceeded my expectations"
16 +input_tokens = [1986, 1985, 572, 9073, 323, 33808, 847, 16665]
17
18
19 @pytest.fixture(scope="module")
@@ -29,9 +31,23 @@ def server():
31
32
33 @pytest.mark.parametrize("model_name", [MODEL_NAME])
32 -def test_single_input_classification(server: RemoteOpenAIServer, model_name: str):
33 - input_text = "This product was excellent and exceeded my expectations"
34 +def test_basic(server: RemoteOpenAIServer, model_name: str):
35 + # test /v1/models
36 + response = requests.get(server.url_for("/v1/models"))
37 + served_model = response.json()["data"][0]["id"]
38 + assert served_model == MODEL_NAME
39
40 + # test /tokenize
41 + response = requests.post(
42 + server.url_for("/tokenize"),
43 + json={"model": model_name, "prompt": input_text},
44 + )
45 + assert response.json()["tokens"] == input_tokens
46 +
47 +
48 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
49 +def test_completion_request(server: RemoteOpenAIServer, model_name: str):
50 + # test input: str
51 classification_response = requests.post(
52 server.url_for("classify"),
53 json={"model": model_name, "input": input_text},
@@ -46,35 +62,34 @@ def test_single_input_classification(server: RemoteOpenAIServer, model_name: str
62 assert hasattr(output.data[0], "label")
63 assert hasattr(output.data[0], "probs")
64
49 -
50 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
51 -def test_add_special_tokens_false(server: RemoteOpenAIServer, model_name: str):
52 - response = requests.post(
65 + # test input: list[int]
66 + classification_response = requests.post(
67 server.url_for("classify"),
54 - json={"model": model_name, "input": "hello", "add_special_tokens": False},
68 + json={"model": model_name, "input": input_tokens},
69 )
56 - response.raise_for_status()
57 - ClassificationResponse.model_validate(response.json())
70 +
71 + classification_response.raise_for_status()
72 + output = ClassificationResponse.model_validate(classification_response.json())
73 +
74 + assert output.object == "list"
75 + assert output.model == MODEL_NAME
76 + assert len(output.data) == 1
77 + assert hasattr(output.data[0], "label")
78 + assert hasattr(output.data[0], "probs")
79
80
81 @pytest.mark.parametrize("model_name", [MODEL_NAME])
61 -def test_multiple_inputs_classification(server: RemoteOpenAIServer, model_name: str):
62 - input_texts = [
63 - "The product arrived on time and works perfectly",
64 - "I'm very satisfied with my purchase, would buy again",
65 - "The customer service was helpful and resolved my issue quickly",
66 - "This product broke after one week, terrible quality",
67 - "I'm very disappointed with this purchase, complete waste of money",
68 - "The customer service was rude and unhelpful",
69 - ]
82 +def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
83 + N = 10
84
85 + # test input: list[str]
86 classification_response = requests.post(
87 server.url_for("classify"),
73 - json={"model": model_name, "input": input_texts},
88 + json={"model": model_name, "input": [input_text] * N},
89 )
90 output = ClassificationResponse.model_validate(classification_response.json())
91
77 - assert len(output.data) == len(input_texts)
92 + assert len(output.data) == N
93 for i, item in enumerate(output.data):
94 assert item.index == i
95 assert hasattr(item, "label")
@@ -82,6 +97,44 @@ def test_multiple_inputs_classification(server: RemoteOpenAIServer, model_name:
97 assert len(item.probs) == item.num_classes
98 assert item.label in ["Default", "Spoiled"]
99
100 + # test input: list[list[int]]
101 + classification_response = requests.post(
102 + server.url_for("classify"),
103 + json={"model": model_name, "input": [input_tokens] * N},
104 + )
105 + output = ClassificationResponse.model_validate(classification_response.json())
106 +
107 + assert len(output.data) == N
108 + for i, item in enumerate(output.data):
109 + assert item.index == i
110 + assert hasattr(item, "label")
111 + assert hasattr(item, "probs")
112 + assert len(item.probs) == item.num_classes
113 + assert item.label in ["Default", "Spoiled"]
114 +
115 +
116 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
117 +def test_empty_input_error(server: RemoteOpenAIServer, model_name: str):
118 + classification_response = requests.post(
119 + server.url_for("classify"),
120 + json={"model": model_name, "input": ""},
121 + )
122 +
123 + error = classification_response.json()
124 + assert classification_response.status_code == 400
125 + assert "error" in error
126 +
127 + classification_response = requests.post(
128 + server.url_for("classify"),
129 + json={"model": model_name, "input": []},
130 + )
131 + classification_response.raise_for_status()
132 + output = ClassificationResponse.model_validate(classification_response.json())
133 +
134 + assert output.object == "list"
135 + assert isinstance(output.data, list)
136 + assert len(output.data) == 0
137 +
138
139 @pytest.mark.parametrize("model_name", [MODEL_NAME])
140 def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
@@ -101,11 +154,7 @@ def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
154 assert output.usage.prompt_tokens == 5
155 assert output.usage.total_tokens == 5
156
104 -
105 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
106 -def test_invalid_truncate_prompt_tokens_error(
107 - server: RemoteOpenAIServer, model_name: str
108 -):
157 + # invalid_truncate_prompt_tokens
158 classification_response = requests.post(
159 server.url_for("classify"),
160 json={"model": model_name, "input": "test", "truncate_prompt_tokens": 513},
@@ -117,36 +166,28 @@ def test_invalid_truncate_prompt_tokens_error(
166
167
168 @pytest.mark.parametrize("model_name", [MODEL_NAME])
120 -def test_empty_input_error(server: RemoteOpenAIServer, model_name: str):
121 - classification_response = requests.post(
169 +def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
170 + # FIXME: The add_special_tokens parameter doesn't seem to be working.
171 + response = requests.post(
172 server.url_for("classify"),
123 - json={"model": model_name, "input": ""},
173 + json={"model": model_name, "input": input_text, "add_special_tokens": False},
174 )
175 + response.raise_for_status()
176 + ClassificationResponse.model_validate(response.json())
177
126 - error = classification_response.json()
127 - assert classification_response.status_code == 400
128 - assert "error" in error
129 -
130 -
131 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
132 -def test_batch_classification_empty_list(server: RemoteOpenAIServer, model_name: str):
133 - classification_response = requests.post(
178 + response = requests.post(
179 server.url_for("classify"),
135 - json={"model": model_name, "input": []},
180 + json={"model": model_name, "input": input_text, "add_special_tokens": True},
181 )
137 - classification_response.raise_for_status()
138 - output = ClassificationResponse.model_validate(classification_response.json())
139 -
140 - assert output.object == "list"
141 - assert isinstance(output.data, list)
142 - assert len(output.data) == 0
182 + response.raise_for_status()
183 + ClassificationResponse.model_validate(response.json())
184
185
186 @pytest.mark.asyncio
187 async def test_invocations(server: RemoteOpenAIServer):
188 request_args = {
189 "model": MODEL_NAME,
149 - "input": "This product was excellent and exceeded my expectations",
190 + "input": input_text,
191 }
192
193 classification_response = requests.post(
@@ -175,8 +216,6 @@ async def test_invocations(server: RemoteOpenAIServer):
216 @pytest.mark.asyncio
217 @pytest.mark.parametrize("model_name", [MODEL_NAME])
218 async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
178 - input_text = ["This product was excellent and exceeded my expectations"]
179 -
219 async def get_outputs(use_activation):
220 response = requests.post(
221 server.url_for("classify"),
@@ -237,7 +276,6 @@ async def test_rerank(server: RemoteOpenAIServer, model_name: str):
276 @pytest.mark.asyncio
277 @pytest.mark.parametrize("model_name", [MODEL_NAME])
278 async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
240 - input_text = "This product was excellent and exceeded my expectations"
279 response = requests.post(
280 server.url_for("pooling"),
281 json={
@@ -256,7 +294,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
294 @pytest.mark.parametrize("model_name", [MODEL_NAME])
295 async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
296 task = "token_classify"
259 - input_text = ["This product was excellent and exceeded my expectations"]
297 response = requests.post(
298 server.url_for("pooling"),
299 json={
@@ -282,7 +319,7 @@ async def test_pooling_not_supported(
319 server.url_for("pooling"),
320 json={
321 "model": model_name,
285 - "input": "test",
322 + "input": input_text,
323 "encoding_format": "float",
324 "task": task,
325 },
tests/entrypoints/pooling/embed/test_online.py
+233 -221
@@ -31,7 +31,26 @@ from vllm.utils.serial_utils import (
31 MODEL_NAME = "intfloat/multilingual-e5-small"
32 DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
33 DTYPE = "bfloat16"
34 -
34 +input_text = "The best thing about vLLM is that it supports many different models"
35 +input_tokens = [
36 + 0,
37 + 581,
38 + 2965,
39 + 13580,
40 + 1672,
41 + 81,
42 + 23708,
43 + 594,
44 + 83,
45 + 450,
46 + 442,
47 + 8060,
48 + 7,
49 + 5941,
50 + 12921,
51 + 115774,
52 + 2,
53 +]
54
55 if current_platform.is_rocm():
56 # Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
@@ -79,15 +98,36 @@ def hf_model(hf_runner):
98
99 @pytest.mark.asyncio
100 @pytest.mark.parametrize("model_name", [MODEL_NAME])
82 -async def test_single_embedding(hf_model, client: openai.AsyncOpenAI, model_name: str):
83 - input_texts = [
84 - "The chef prepared a delicious meal.",
85 - ]
101 +async def test_basic(
102 + server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
103 +):
104 + # test /v1/models
105 + response = requests.get(server.url_for("/v1/models"))
106 + model = response.json()["data"][0]["id"]
107 + assert model == MODEL_NAME
108
87 - # test single embedding
109 + models = await client.models.list()
110 + models = models.data
111 + served_model = models[0]
112 + assert served_model.id == MODEL_NAME
113 +
114 + # test /tokenize
115 + response = requests.post(
116 + server.url_for("/tokenize"),
117 + json={"model": model_name, "prompt": input_text},
118 + )
119 + assert response.json()["tokens"] == input_tokens
120 +
121 +
122 +@pytest.mark.asyncio
123 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
124 +async def test_completion_request(
125 + client: openai.AsyncOpenAI, model_name: str, hf_model
126 +):
127 + # test input: str
128 embedding_response = await client.embeddings.create(
129 model=model_name,
90 - input=input_texts,
130 + input=input_text,
131 encoding_format="float",
132 )
133 embeddings = EmbeddingResponse.model_validate(
@@ -98,14 +138,13 @@ async def test_single_embedding(hf_model, client: openai.AsyncOpenAI, model_name
138 assert len(embeddings.data) == 1
139 assert len(embeddings.data[0].embedding) == 384
140 assert embeddings.usage.completion_tokens == 0
101 - assert embeddings.usage.prompt_tokens == 11
102 - assert embeddings.usage.total_tokens == 11
141 + assert embeddings.usage.prompt_tokens == len(input_tokens)
142 + assert embeddings.usage.total_tokens == len(input_tokens)
143
144 vllm_outputs = [d.embedding for d in embeddings.data]
105 - run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
145 + run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
146
107 - # test using token IDs
108 - input_tokens = [1, 1, 1, 1, 1]
147 + # test input: list[int]
148 embedding_response = await client.embeddings.create(
149 model=model_name,
150 input=input_tokens,
@@ -119,19 +158,22 @@ async def test_single_embedding(hf_model, client: openai.AsyncOpenAI, model_name
158 assert len(embeddings.data) == 1
159 assert len(embeddings.data[0].embedding) == 384
160 assert embeddings.usage.completion_tokens == 0
122 - assert embeddings.usage.prompt_tokens == 5
123 - assert embeddings.usage.total_tokens == 5
161 + assert embeddings.usage.prompt_tokens == len(input_tokens)
162 + assert embeddings.usage.total_tokens == len(input_tokens)
163 +
164 + vllm_outputs = [d.embedding for d in embeddings.data]
165 + run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
166
167
168 @pytest.mark.asyncio
169 @pytest.mark.parametrize("model_name", [MODEL_NAME])
128 -async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI, model_name: str):
129 - # test list[str]
130 - input_texts = [
131 - "The cat sat on the mat.",
132 - "A feline was resting on a rug.",
133 - "Stars twinkle brightly in the night sky.",
134 - ]
170 +async def test_completion_request_batched(
171 + client: openai.AsyncOpenAI, model_name: str, hf_model
172 +):
173 + N = 10
174 + input_texts = [input_text] * N
175 +
176 + # test input: list[str]
177 embedding_response = await client.embeddings.create(
178 model=model_name,
179 input=input_texts,
@@ -142,25 +184,19 @@ async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI, model_name:
184 )
185
186 assert embeddings.id is not None
145 - assert len(embeddings.data) == 3
187 + assert len(embeddings.data) == N
188 assert len(embeddings.data[0].embedding) == 384
189 assert embeddings.usage.completion_tokens == 0
148 - assert embeddings.usage.prompt_tokens == 33
149 - assert embeddings.usage.total_tokens == 33
190 + assert embeddings.usage.prompt_tokens == len(input_tokens) * N
191 + assert embeddings.usage.total_tokens == len(input_tokens) * N
192
193 vllm_outputs = [d.embedding for d in embeddings.data]
194 run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
195
196 # test list[list[int]]
155 - input_tokens = [
156 - [4, 5, 7, 9, 20],
157 - [15, 29, 499],
158 - [24, 24, 24, 24, 24],
159 - [25, 32, 64, 77],
160 - ]
197 embedding_response = await client.embeddings.create(
198 model=model_name,
163 - input=input_tokens,
199 + input=[input_tokens] * N,
200 encoding_format="float",
201 )
202 embeddings = EmbeddingResponse.model_validate(
@@ -168,11 +204,14 @@ async def test_batch_embedding(hf_model, client: openai.AsyncOpenAI, model_name:
204 )
205
206 assert embeddings.id is not None
171 - assert len(embeddings.data) == 4
207 + assert len(embeddings.data) == N
208 assert len(embeddings.data[0].embedding) == 384
209 assert embeddings.usage.completion_tokens == 0
174 - assert embeddings.usage.prompt_tokens == 17
175 - assert embeddings.usage.total_tokens == 17
210 + assert embeddings.usage.prompt_tokens == len(input_tokens) * N
211 + assert embeddings.usage.total_tokens == len(input_tokens) * N
212 +
213 + vllm_outputs = [d.embedding for d in embeddings.data]
214 + run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
215
216
217 @pytest.mark.asyncio
@@ -235,9 +274,162 @@ async def test_conversation_embedding(
274
275 @pytest.mark.asyncio
276 @pytest.mark.parametrize("model_name", [MODEL_NAME])
238 -async def test_batch_base64_embedding(
239 - hf_model, client: openai.AsyncOpenAI, model_name: str
240 -):
277 +async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: str):
278 + input_texts = [
279 + "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
280 + ]
281 +
282 + # test single embedding
283 + embedding_response = await client.embeddings.create(
284 + model=model_name, input=input_texts, extra_body={"truncate_prompt_tokens": 10}
285 + )
286 + embeddings = EmbeddingResponse.model_validate(
287 + embedding_response.model_dump(mode="json")
288 + )
289 +
290 + assert embeddings.id is not None
291 + assert len(embeddings.data) == 1
292 + assert len(embeddings.data[0].embedding) == 384
293 + assert embeddings.usage.completion_tokens == 0
294 + assert embeddings.usage.prompt_tokens == 10
295 + assert embeddings.usage.total_tokens == 10
296 +
297 + input_tokens = [
298 + 1,
299 + 24428,
300 + 289,
301 + 18341,
302 + 26165,
303 + 285,
304 + 19323,
305 + 283,
306 + 289,
307 + 26789,
308 + 3871,
309 + 28728,
310 + 9901,
311 + 340,
312 + 2229,
313 + 385,
314 + 340,
315 + 315,
316 + 28741,
317 + 28804,
318 + 2,
319 + ]
320 + embedding_response = await client.embeddings.create(
321 + model=model_name, input=input_tokens, extra_body={"truncate_prompt_tokens": 10}
322 + )
323 + embeddings = EmbeddingResponse.model_validate(
324 + embedding_response.model_dump(mode="json")
325 + )
326 +
327 + assert embeddings.id is not None
328 + assert len(embeddings.data) == 1
329 + assert len(embeddings.data[0].embedding) == 384
330 + assert embeddings.usage.completion_tokens == 0
331 + assert embeddings.usage.prompt_tokens == 10
332 + assert embeddings.usage.total_tokens == 10
333 +
334 + # invalid_truncate_prompt_tokens
335 + input_texts = [
336 + "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
337 + ]
338 +
339 + with pytest.raises(openai.BadRequestError):
340 + response = await client.embeddings.create(
341 + model=model_name,
342 + input=input_texts,
343 + extra_body={"truncate_prompt_tokens": 8193},
344 + )
345 + assert "error" in response.object
346 + assert (
347 + "truncate_prompt_tokens value is greater than max_model_len. "
348 + "Please, select a smaller truncation size." in response.message
349 + )
350 +
351 +
352 +@pytest.mark.asyncio
353 +async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenAI):
354 + request_args = {
355 + "model": MODEL_NAME,
356 + "input": input_text,
357 + "encoding_format": "float",
358 + }
359 +
360 + completion_response = await client.embeddings.create(**request_args)
361 +
362 + invocation_response = requests.post(
363 + server.url_for("invocations"), json=request_args
364 + )
365 + invocation_response.raise_for_status()
366 +
367 + completion_output = completion_response.model_dump()
368 + invocation_output = invocation_response.json()
369 +
370 + assert completion_output.keys() == invocation_output.keys()
371 + for completion_data, invocation_data in zip(
372 + completion_output["data"], invocation_output["data"]
373 + ):
374 + assert completion_data.keys() == invocation_data.keys()
375 + check_embeddings_close(
376 + embeddings_0_lst=[completion_data["embedding"]],
377 + embeddings_1_lst=[invocation_data["embedding"]],
378 + name_0="completion",
379 + name_1="invocation",
380 + )
381 +
382 +
383 +@pytest.mark.asyncio
384 +async def test_invocations_conversation(server: RemoteOpenAIServer):
385 + messages = [
386 + {
387 + "role": "user",
388 + "content": "The cat sat on the mat.",
389 + },
390 + {
391 + "role": "assistant",
392 + "content": "A feline was resting on a rug.",
393 + },
394 + {
395 + "role": "user",
396 + "content": "Stars twinkle brightly in the night sky.",
397 + },
398 + ]
399 +
400 + request_args = {
401 + "model": MODEL_NAME,
402 + "messages": messages,
403 + "encoding_format": "float",
404 + }
405 +
406 + chat_response = requests.post(server.url_for("v1/embeddings"), json=request_args)
407 + chat_response.raise_for_status()
408 +
409 + invocation_response = requests.post(
410 + server.url_for("invocations"), json=request_args
411 + )
412 + invocation_response.raise_for_status()
413 +
414 + chat_output = chat_response.json()
415 + invocation_output = invocation_response.json()
416 +
417 + assert chat_output.keys() == invocation_output.keys()
418 + for chat_data, invocation_data in zip(
419 + chat_output["data"], invocation_output["data"]
420 + ):
421 + assert chat_data.keys() == invocation_data.keys()
422 + check_embeddings_close(
423 + embeddings_0_lst=[chat_data["embedding"]],
424 + embeddings_1_lst=[invocation_data["embedding"]],
425 + name_0="chat",
426 + name_1="invocation",
427 + )
428 +
429 +
430 +@pytest.mark.asyncio
431 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
432 +async def test_base64_embedding(hf_model, client: openai.AsyncOpenAI, model_name: str):
433 input_texts = [
434 "Hello my name is",
435 "The best thing about vLLM is that it supports many different models",
@@ -273,10 +465,7 @@ async def test_batch_base64_embedding(
465 async def test_base64_embed_dtype_and_endianness(
466 server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
467 ):
276 - input_texts = [
277 - "The best thing about vLLM is that it supports many different models",
278 - ]
279 -
468 + input_texts = [input_text] * 3
469 responses_float = await client.embeddings.create(
470 input=input_texts, model=model_name, encoding_format="float"
471 )
@@ -315,10 +504,7 @@ async def test_base64_embed_dtype_and_endianness(
504 async def test_bytes_embed_dtype_and_endianness(
505 server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
506 ):
318 - input_texts = [
319 - "The best thing about vLLM is that it supports many different models",
320 - ]
321 -
507 + input_texts = [input_text] * 3
508 responses_float = await client.embeddings.create(
509 input=input_texts, model=model_name, encoding_format="float"
510 )
@@ -408,15 +594,11 @@ async def test_bytes_only_embed_dtype_and_endianness(
594 async def test_params_not_supported(
595 server: RemoteOpenAIServer, model_name: str, param_name: str
596 ):
411 - input_texts = [
412 - "The best thing about vLLM is that it supports many different models",
413 - ]
414 -
597 responses_base64 = requests.post(
598 server.url_for("/v1/embeddings"),
599 json={
600 "model": model_name,
419 - "input": input_texts,
601 + "input": input_text,
602 "encoding_format": "base64",
603 param_name: f"bad_{param_name}",
604 },
@@ -427,175 +609,9 @@ async def test_params_not_supported(
609 assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
610
611
430 -@pytest.mark.asyncio
431 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
432 -async def test_single_embedding_truncation(client: openai.AsyncOpenAI, model_name: str):
433 - input_texts = [
434 - "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
435 - ]
436 -
437 - # test single embedding
438 - embedding_response = await client.embeddings.create(
439 - model=model_name, input=input_texts, extra_body={"truncate_prompt_tokens": 10}
440 - )
441 - embeddings = EmbeddingResponse.model_validate(
442 - embedding_response.model_dump(mode="json")
443 - )
444 -
445 - assert embeddings.id is not None
446 - assert len(embeddings.data) == 1
447 - assert len(embeddings.data[0].embedding) == 384
448 - assert embeddings.usage.completion_tokens == 0
449 - assert embeddings.usage.prompt_tokens == 10
450 - assert embeddings.usage.total_tokens == 10
451 -
452 - input_tokens = [
453 - 1,
454 - 24428,
455 - 289,
456 - 18341,
457 - 26165,
458 - 285,
459 - 19323,
460 - 283,
461 - 289,
462 - 26789,
463 - 3871,
464 - 28728,
465 - 9901,
466 - 340,
467 - 2229,
468 - 385,
469 - 340,
470 - 315,
471 - 28741,
472 - 28804,
473 - 2,
474 - ]
475 - embedding_response = await client.embeddings.create(
476 - model=model_name, input=input_tokens, extra_body={"truncate_prompt_tokens": 10}
477 - )
478 - embeddings = EmbeddingResponse.model_validate(
479 - embedding_response.model_dump(mode="json")
480 - )
481 -
482 - assert embeddings.id is not None
483 - assert len(embeddings.data) == 1
484 - assert len(embeddings.data[0].embedding) == 384
485 - assert embeddings.usage.completion_tokens == 0
486 - assert embeddings.usage.prompt_tokens == 10
487 - assert embeddings.usage.total_tokens == 10
488 -
489 -
490 -@pytest.mark.asyncio
491 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
492 -async def test_single_embedding_truncation_invalid(
493 - client: openai.AsyncOpenAI, model_name: str
494 -):
495 - input_texts = [
496 - "Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
497 - ]
498 -
499 - with pytest.raises(openai.BadRequestError):
500 - response = await client.embeddings.create(
501 - model=model_name,
502 - input=input_texts,
503 - extra_body={"truncate_prompt_tokens": 8193},
504 - )
505 - assert "error" in response.object
506 - assert (
507 - "truncate_prompt_tokens value is greater than max_model_len. "
508 - "Please, select a smaller truncation size." in response.message
509 - )
510 -
511 -
512 -@pytest.mark.asyncio
513 -async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenAI):
514 - input_texts = [
515 - "The chef prepared a delicious meal.",
516 - ]
517 -
518 - request_args = {
519 - "model": MODEL_NAME,
520 - "input": input_texts,
521 - "encoding_format": "float",
522 - }
523 -
524 - completion_response = await client.embeddings.create(**request_args)
525 -
526 - invocation_response = requests.post(
527 - server.url_for("invocations"), json=request_args
528 - )
529 - invocation_response.raise_for_status()
530 -
531 - completion_output = completion_response.model_dump()
532 - invocation_output = invocation_response.json()
533 -
534 - assert completion_output.keys() == invocation_output.keys()
535 - for completion_data, invocation_data in zip(
536 - completion_output["data"], invocation_output["data"]
537 - ):
538 - assert completion_data.keys() == invocation_data.keys()
539 - check_embeddings_close(
540 - embeddings_0_lst=[completion_data["embedding"]],
541 - embeddings_1_lst=[invocation_data["embedding"]],
542 - name_0="completion",
543 - name_1="invocation",
544 - )
545 -
546 -
547 -@pytest.mark.asyncio
548 -async def test_invocations_conversation(server: RemoteOpenAIServer):
549 - messages = [
550 - {
551 - "role": "user",
552 - "content": "The cat sat on the mat.",
553 - },
554 - {
555 - "role": "assistant",
556 - "content": "A feline was resting on a rug.",
557 - },
558 - {
559 - "role": "user",
560 - "content": "Stars twinkle brightly in the night sky.",
561 - },
562 - ]
563 -
564 - request_args = {
565 - "model": MODEL_NAME,
566 - "messages": messages,
567 - "encoding_format": "float",
568 - }
569 -
570 - chat_response = requests.post(server.url_for("v1/embeddings"), json=request_args)
571 - chat_response.raise_for_status()
572 -
573 - invocation_response = requests.post(
574 - server.url_for("invocations"), json=request_args
575 - )
576 - invocation_response.raise_for_status()
577 -
578 - chat_output = chat_response.json()
579 - invocation_output = invocation_response.json()
580 -
581 - assert chat_output.keys() == invocation_output.keys()
582 - for chat_data, invocation_data in zip(
583 - chat_output["data"], invocation_output["data"]
584 - ):
585 - assert chat_data.keys() == invocation_data.keys()
586 - check_embeddings_close(
587 - embeddings_0_lst=[chat_data["embedding"]],
588 - embeddings_1_lst=[invocation_data["embedding"]],
589 - name_0="chat",
590 - name_1="invocation",
591 - )
592 -
593 -
612 @pytest.mark.asyncio
613 @pytest.mark.parametrize("model_name", [MODEL_NAME])
614 async def test_normalize(server: RemoteOpenAIServer, model_name: str):
597 - input_text = ["The chef prepared a delicious meal."]
598 -
615 async def get_outputs(normalize):
616 request_args = {
617 "model": MODEL_NAME,
@@ -626,8 +642,6 @@ async def test_normalize(server: RemoteOpenAIServer, model_name: str):
642 @pytest.mark.parametrize("model_name", [MODEL_NAME])
643 async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
644 task = "embed"
629 - input_text = ["The chef prepared a delicious meal."]
630 -
645 response = requests.post(
646 server.url_for("pooling"),
647 json={
@@ -648,8 +662,6 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
662 @pytest.mark.parametrize("model_name", [MODEL_NAME])
663 async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
664 task = "token_embed"
651 - input_text = ["The chef prepared a delicious meal."]
652 -
665 response = requests.post(
666 server.url_for("pooling"),
667 json={
@@ -663,7 +675,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
675 poolings = PoolingResponse.model_validate(response.json())
676
677 assert len(poolings.data) == 1
666 - assert len(poolings.data[0].data) == 11
678 + assert len(poolings.data[0].data) == len(input_tokens)
679 assert len(poolings.data[0].data[0]) == 384
680
681
tests/entrypoints/pooling/pooling/test_online.py
+49 -57
@@ -24,6 +24,8 @@ from vllm.utils.serial_utils import (
24
25 MODEL_NAME = "internlm/internlm2-1_8b-reward"
26 DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
27 +input_text = "The chef prepared a delicious meal."
28 +input_tokens = [1, 918, 29981, 10166, 395, 18067, 15265, 281]
29
30
31 @pytest.fixture(scope="module")
@@ -46,30 +48,40 @@ def server():
48 yield remote_server
49
50
49 -@pytest.mark.asyncio
51 @pytest.mark.parametrize("model_name", [MODEL_NAME])
51 -async def test_single_pooling(server: RemoteOpenAIServer, model_name: str):
52 - input_texts = [
53 - "The chef prepared a delicious meal.",
54 - ]
52 +def test_basic(server: RemoteOpenAIServer, model_name: str):
53 + # test /v1/models
54 + response = requests.get(server.url_for("/v1/models"))
55 + served_model = response.json()["data"][0]["id"]
56 + assert served_model == MODEL_NAME
57
56 - # test single pooling
58 + # test /tokenize
59 + response = requests.post(
60 + server.url_for("/tokenize"),
61 + json={"model": model_name, "prompt": input_text},
62 + )
63 + assert response.json()["tokens"] == input_tokens
64 +
65 +
66 +@pytest.mark.asyncio
67 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
68 +def test_completion_request(server: RemoteOpenAIServer, model_name: str):
69 + # test input: str
70 response = requests.post(
71 server.url_for("pooling"),
59 - json={"model": model_name, "input": input_texts, "encoding_format": "float"},
72 + json={"model": model_name, "input": input_text, "encoding_format": "float"},
73 )
74 response.raise_for_status()
75 poolings = PoolingResponse.model_validate(response.json())
76
77 assert poolings.id is not None
78 assert len(poolings.data) == 1
66 - assert len(poolings.data[0].data) == 8
79 + assert len(poolings.data[0].data) == len(input_tokens)
80 assert poolings.usage.completion_tokens == 0
68 - assert poolings.usage.prompt_tokens == 8
69 - assert poolings.usage.total_tokens == 8
81 + assert poolings.usage.prompt_tokens == len(input_tokens)
82 + assert poolings.usage.total_tokens == len(input_tokens)
83
71 - # test using token IDs
72 - input_tokens = [1, 1, 1, 1, 1]
84 + # test input: list[int]
85 response = requests.post(
86 server.url_for("pooling"),
87 json={"model": model_name, "input": input_tokens, "encoding_format": "float"},
@@ -79,21 +91,17 @@ async def test_single_pooling(server: RemoteOpenAIServer, model_name: str):
91
92 assert poolings.id is not None
93 assert len(poolings.data) == 1
82 - assert len(poolings.data[0].data) == 5
94 + assert len(poolings.data[0].data) == len(input_tokens)
95 assert poolings.usage.completion_tokens == 0
84 - assert poolings.usage.prompt_tokens == 5
85 - assert poolings.usage.total_tokens == 5
96 + assert poolings.usage.prompt_tokens == len(input_tokens)
97 + assert poolings.usage.total_tokens == len(input_tokens)
98
99
88 -@pytest.mark.asyncio
100 @pytest.mark.parametrize("model_name", [MODEL_NAME])
90 -async def test_batch_pooling(server: RemoteOpenAIServer, model_name: str):
91 - # test list[str]
92 - input_texts = [
93 - "The cat sat on the mat.",
94 - "A feline was resting on a rug.",
95 - "Stars twinkle brightly in the night sky.",
96 - ]
101 +def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
102 + N = 10
103 + input_texts = [input_text] * N
104 +
105 response = requests.post(
106 server.url_for("pooling"),
107 json={"model": model_name, "input": input_texts, "encoding_format": "float"},
@@ -102,32 +110,30 @@ async def test_batch_pooling(server: RemoteOpenAIServer, model_name: str):
110 poolings = PoolingResponse.model_validate(response.json())
111
112 assert poolings.id is not None
105 - assert len(poolings.data) == 3
106 - assert len(poolings.data[0].data) == 8
113 + assert len(poolings.data) == N
114 + assert len(poolings.data[0].data) == len(input_tokens)
115 assert poolings.usage.completion_tokens == 0
108 - assert poolings.usage.prompt_tokens == 29
109 - assert poolings.usage.total_tokens == 29
116 + assert poolings.usage.prompt_tokens == len(input_tokens) * N
117 + assert poolings.usage.total_tokens == len(input_tokens) * N
118
119 # test list[list[int]]
112 - input_tokens = [
113 - [4, 5, 7, 9, 20],
114 - [15, 29, 499],
115 - [24, 24, 24, 24, 24],
116 - [25, 32, 64, 77],
117 - ]
120 response = requests.post(
121 server.url_for("pooling"),
120 - json={"model": model_name, "input": input_tokens, "encoding_format": "float"},
122 + json={
123 + "model": model_name,
124 + "input": [input_tokens] * N,
125 + "encoding_format": "float",
126 + },
127 )
128 response.raise_for_status()
129 poolings = PoolingResponse.model_validate(response.json())
130
131 assert poolings.id is not None
126 - assert len(poolings.data) == 4
127 - assert len(poolings.data[0].data) == 5
132 + assert len(poolings.data) == N
133 + assert len(poolings.data[0].data) == len(input_tokens)
134 assert poolings.usage.completion_tokens == 0
129 - assert poolings.usage.prompt_tokens == 17
130 - assert poolings.usage.total_tokens == 17
135 + assert poolings.usage.prompt_tokens == len(input_tokens) * N
136 + assert poolings.usage.total_tokens == len(input_tokens) * N
137
138
139 @pytest.mark.asyncio
@@ -259,9 +265,7 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer, model_name: str)
265 async def test_base64_embed_dtype_and_endianness(
266 server: RemoteOpenAIServer, model_name: str
267 ):
262 - input_texts = [
263 - "The best thing about vLLM is that it supports many different models",
264 - ]
268 + input_texts = [input_text] * 3
269
270 url = server.url_for("pooling")
271 float_response = requests.post(
@@ -308,9 +312,7 @@ async def test_base64_embed_dtype_and_endianness(
312 async def test_bytes_embed_dtype_and_endianness(
313 server: RemoteOpenAIServer, model_name: str
314 ):
311 - input_texts = [
312 - "The best thing about vLLM is that it supports many different models",
313 - ]
315 + input_texts = [input_text] * 3
316
317 url = server.url_for("pooling")
318 float_response = requests.post(
@@ -358,9 +360,7 @@ async def test_bytes_embed_dtype_and_endianness(
360 async def test_bytes_only_embed_dtype_and_endianness(
361 server: RemoteOpenAIServer, model_name: str
362 ):
361 - input_texts = [
362 - "The best thing about vLLM is that it supports many different models",
363 - ] * 2
363 + input_texts = [input_text] * 3
364
365 url = server.url_for("pooling")
366 float_response = requests.post(
@@ -414,15 +414,11 @@ async def test_bytes_only_embed_dtype_and_endianness(
414 async def test_params_not_supported(
415 server: RemoteOpenAIServer, model_name: str, param_name: str
416 ):
417 - input_texts = [
418 - "The best thing about vLLM is that it supports many different models",
419 - ]
420 -
417 responses_base64 = requests.post(
418 server.url_for("pooling"),
419 json={
420 "model": model_name,
425 - "input": input_texts,
421 + "input": input_text,
422 "encoding_format": "base64",
423 param_name: f"bad_{param_name}",
424 },
@@ -435,13 +431,9 @@ async def test_params_not_supported(
431
432 @pytest.mark.asyncio
433 async def test_invocations(server: RemoteOpenAIServer):
438 - input_texts = [
439 - "The chef prepared a delicious meal.",
440 - ]
441 -
434 request_args = {
435 "model": MODEL_NAME,
444 - "input": input_texts,
436 + "input": input_text,
437 "encoding_format": "float",
438 }
439
tests/entrypoints/pooling/score/test_online_rerank.py
+19 -5
@@ -13,6 +13,8 @@ from vllm.platforms import current_platform
13
14 MODEL_NAME = "BAAI/bge-reranker-base"
15 DTYPE = "bfloat16"
16 +input_text = "This product was excellent and exceeded my expectations"
17 +input_tokens = [0, 3293, 12996, 509, 40881, 136, 204839, 297, 759, 202702, 2]
18
19
20 @pytest.fixture(scope="module")
@@ -27,6 +29,21 @@ def server():
29 yield remote_server
30
31
32 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
33 +def test_basic(server: RemoteOpenAIServer, model_name: str):
34 + # test /v1/models
35 + response = requests.get(server.url_for("/v1/models"))
36 + served_model = response.json()["data"][0]["id"]
37 + assert served_model == MODEL_NAME
38 +
39 + # test /tokenize
40 + response = requests.post(
41 + server.url_for("/tokenize"),
42 + json={"model": model_name, "prompt": input_text},
43 + )
44 + assert response.json()["tokens"] == input_tokens
45 +
46 +
47 @pytest.mark.parametrize("model_name", [MODEL_NAME])
48 def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
49 query = "What is the capital of France?"
@@ -170,7 +187,6 @@ async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
187 @pytest.mark.asyncio
188 @pytest.mark.parametrize("model_name", [MODEL_NAME])
189 async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
173 - input_text = "This product was excellent and exceeded my expectations"
190 response = requests.post(
191 server.url_for("pooling"),
192 json={
@@ -188,8 +204,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
204 @pytest.mark.asyncio
205 @pytest.mark.parametrize("model_name", [MODEL_NAME])
206 async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
191 - input_text = ["The chef prepared a delicious meal."]
192 -
207 response = requests.post(
208 server.url_for("pooling"),
209 json={"model": model_name, "input": input_text, "encoding_format": "float"},
@@ -198,7 +212,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
212 poolings = PoolingResponse.model_validate(response.json())
213
214 assert len(poolings.data) == 1
201 - assert len(poolings.data[0].data) == 11
215 + assert len(poolings.data[0].data) == len(input_tokens)
216 assert len(poolings.data[0].data[0]) == 1
217
218
@@ -212,7 +226,7 @@ async def test_pooling_not_supported(
226 server.url_for("pooling"),
227 json={
228 "model": model_name,
215 - "input": "test",
229 + "input": input_text,
230 "encoding_format": "float",
231 "task": task,
232 },
tests/entrypoints/pooling/score/test_utils.py
+5 -5
@@ -7,7 +7,7 @@ import pytest
7
8 from vllm.config import ModelConfig
9 from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
10 -from vllm.entrypoints.score_utils import get_score_prompt
10 +from vllm.entrypoints.pooling.score.utils import get_score_prompt
11 from vllm.inputs import TokensPrompt
12 from vllm.tokenizers import get_tokenizer
13
@@ -212,7 +212,7 @@ class TestGetScorePrompt:
212 return_value=mock_model_no_score_template,
213 ),
214 patch(
215 - "vllm.entrypoints.score_utils.apply_hf_chat_template",
215 + "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
216 return_value="test querytest doc",
217 ),
218 ):
@@ -245,7 +245,7 @@ class TestGetScorePrompt:
245 return_value=mock_model_no_score_template,
246 ),
247 patch(
248 - "vllm.entrypoints.score_utils.apply_hf_chat_template",
248 + "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
249 side_effect=ChatTemplateResolutionError("No template"),
250 ),
251 ):
@@ -296,7 +296,7 @@ class TestGetScorePrompt:
296 return_value=mock_model_no_score_template,
297 ),
298 patch(
299 - "vllm.entrypoints.score_utils.apply_hf_chat_template",
299 + "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
300 side_effect=ChatTemplateResolutionError("No template"),
301 ),
302 ):
@@ -331,7 +331,7 @@ class TestGetScorePrompt:
331 return_value=mock_model_with_score_template,
332 ),
333 patch(
334 - "vllm.entrypoints.score_utils.apply_hf_chat_template",
334 + "vllm.entrypoints.pooling.score.utils.apply_hf_chat_template",
335 side_effect=ChatTemplateResolutionError("No template"),
336 ),
337 ):
tests/models/multimodal/pooling/test_jinavl_reranker.py
+1 -1
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import (
10 ChatCompletionContentPartImageParam,
11 ChatCompletionContentPartTextParam,
12 )
13 -from vllm.entrypoints.score_utils import ScoreMultiModalParam
13 +from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
14
15 from ....conftest import HfRunner, VllmRunner
16
vllm/entrypoints/llm.py
+1 -1
@@ -42,7 +42,7 @@ from vllm.entrypoints.chat_utils import (
42 parse_chat_messages,
43 resolve_chat_template_content_format,
44 )
45 -from vllm.entrypoints.score_utils import (
45 +from vllm.entrypoints.pooling.score.utils import (
46 ScoreContentPartParam,
47 ScoreMultiModalParam,
48 _cosine_similarity,
vllm/entrypoints/openai/api_server.py
+4 -58
@@ -54,10 +54,6 @@ from vllm.entrypoints.openai.translations.serving import (
54 OpenAIServingTranscription,
55 OpenAIServingTranslation,
56 )
57 -from vllm.entrypoints.pooling.classify.serving import ServingClassification
58 -from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
59 -from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
60 -from vllm.entrypoints.pooling.score.serving import ServingScores
57 from vllm.entrypoints.serve.disagg.serving import ServingTokens
58 from vllm.entrypoints.serve.elastic_ep.middleware import (
59 ScalingMiddleware,
@@ -73,7 +69,6 @@ from vllm.entrypoints.utils import (
69 from vllm.exceptions import VLLMValidationError
70 from vllm.logger import init_logger
71 from vllm.reasoning import ReasoningParserManager
76 -from vllm.tasks import POOLING_TASKS
72 from vllm.tool_parsers import ToolParserManager
73 from vllm.usage.usage_lib import UsageContext
74 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -761,59 +756,6 @@ async def init_app_state(
756 if "generate" in supported_tasks
757 else None
758 )
764 - state.openai_serving_pooling = (
765 - (
766 - OpenAIServingPooling(
767 - engine_client,
768 - state.openai_serving_models,
769 - supported_tasks=supported_tasks,
770 - request_logger=request_logger,
771 - chat_template=resolved_chat_template,
772 - chat_template_content_format=args.chat_template_content_format,
773 - trust_request_chat_template=args.trust_request_chat_template,
774 - log_error_stack=args.log_error_stack,
775 - )
776 - )
777 - if any(task in POOLING_TASKS for task in supported_tasks)
778 - else None
779 - )
780 - state.openai_serving_embedding = (
781 - OpenAIServingEmbedding(
782 - engine_client,
783 - state.openai_serving_models,
784 - request_logger=request_logger,
785 - chat_template=resolved_chat_template,
786 - chat_template_content_format=args.chat_template_content_format,
787 - trust_request_chat_template=args.trust_request_chat_template,
788 - log_error_stack=args.log_error_stack,
789 - )
790 - if "embed" in supported_tasks
791 - else None
792 - )
793 - state.openai_serving_classification = (
794 - ServingClassification(
795 - engine_client,
796 - state.openai_serving_models,
797 - request_logger=request_logger,
798 - chat_template=resolved_chat_template,
799 - chat_template_content_format=args.chat_template_content_format,
800 - trust_request_chat_template=args.trust_request_chat_template,
801 - log_error_stack=args.log_error_stack,
802 - )
803 - if "classify" in supported_tasks
804 - else None
805 - )
806 - state.openai_serving_scores = (
807 - ServingScores(
808 - engine_client,
809 - state.openai_serving_models,
810 - request_logger=request_logger,
811 - score_template=resolved_chat_template,
812 - log_error_stack=args.log_error_stack,
813 - )
814 - if ("embed" in supported_tasks or "score" in supported_tasks)
815 - else None
816 - )
759 state.openai_serving_tokenization = OpenAIServingTokenization(
760 engine_client,
761 state.openai_serving_models,
@@ -878,6 +820,10 @@ async def init_app_state(
820 else None
821 )
822
823 + from vllm.entrypoints.pooling import init_pooling_state
824 +
825 + await init_pooling_state(engine_client, state, args)
826 +
827 state.enable_server_load_tracking = args.enable_server_load_tracking
828 state.server_load_metrics = 0
829
vllm/entrypoints/pooling/__init__.py
+88
@@ -1,8 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from typing import TYPE_CHECKING
5 +
6 from fastapi import FastAPI
7
8 +if TYPE_CHECKING:
9 + from argparse import Namespace
10 +
11 + from starlette.datastructures import State
12 +
13 + from vllm.engine.protocol import EngineClient
14 +
15
16 def register_pooling_api_routers(app: FastAPI):
17 from vllm.entrypoints.pooling.classify.api_router import router as classify_router
@@ -14,3 +23,82 @@ def register_pooling_api_routers(app: FastAPI):
23 app.include_router(embed_router)
24 app.include_router(score_router)
25 app.include_router(pooling_router)
26 +
27 +
28 +async def init_pooling_state(
29 + engine_client: "EngineClient", state: "State", args: "Namespace"
30 +):
31 + from vllm.entrypoints.logger import RequestLogger
32 + from vllm.entrypoints.pooling.classify.serving import ServingClassification
33 + from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
34 + from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
35 + from vllm.entrypoints.pooling.score.serving import ServingScores
36 + from vllm.entrypoints.utils import process_chat_template
37 + from vllm.tasks import POOLING_TASKS
38 +
39 + supported_tasks = await engine_client.get_supported_tasks()
40 +
41 + vllm_config = engine_client.vllm_config
42 +
43 + resolved_chat_template = await process_chat_template(
44 + args.chat_template, engine_client, vllm_config.model_config
45 + )
46 +
47 + if args.enable_log_requests:
48 + request_logger = RequestLogger(max_log_len=args.max_log_len)
49 + else:
50 + request_logger = None
51 +
52 + state.openai_serving_pooling = (
53 + (
54 + OpenAIServingPooling(
55 + engine_client,
56 + state.openai_serving_models,
57 + supported_tasks=supported_tasks,
58 + request_logger=request_logger,
59 + chat_template=resolved_chat_template,
60 + chat_template_content_format=args.chat_template_content_format,
61 + trust_request_chat_template=args.trust_request_chat_template,
62 + log_error_stack=args.log_error_stack,
63 + )
64 + )
65 + if any(task in POOLING_TASKS for task in supported_tasks)
66 + else None
67 + )
68 + state.openai_serving_embedding = (
69 + OpenAIServingEmbedding(
70 + engine_client,
71 + state.openai_serving_models,
72 + request_logger=request_logger,
73 + chat_template=resolved_chat_template,
74 + chat_template_content_format=args.chat_template_content_format,
75 + trust_request_chat_template=args.trust_request_chat_template,
76 + log_error_stack=args.log_error_stack,
77 + )
78 + if "embed" in supported_tasks
79 + else None
80 + )
81 + state.openai_serving_classification = (
82 + ServingClassification(
83 + engine_client,
84 + state.openai_serving_models,
85 + request_logger=request_logger,
86 + chat_template=resolved_chat_template,
87 + chat_template_content_format=args.chat_template_content_format,
88 + trust_request_chat_template=args.trust_request_chat_template,
89 + log_error_stack=args.log_error_stack,
90 + )
91 + if "classify" in supported_tasks
92 + else None
93 + )
94 + state.openai_serving_scores = (
95 + ServingScores(
96 + engine_client,
97 + state.openai_serving_models,
98 + request_logger=request_logger,
99 + score_template=resolved_chat_template,
100 + log_error_stack=args.log_error_stack,
101 + )
102 + if ("embed" in supported_tasks or "score" in supported_tasks)
103 + else None
104 + )
vllm/entrypoints/pooling/base/__init__.py
vllm/entrypoints/pooling/base/protocol.py new
+46
@@ -0,0 +1,46 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +from typing import Annotated
6 +
7 +from pydantic import Field
8 +
9 +from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
10 +from vllm.utils import random_uuid
11 +
12 +
13 +class PoolingBasicRequestMixin(OpenAIBaseModel):
14 + model: str | None = None
15 + user: str | None = None
16 + truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
17 +
18 + request_id: str = Field(
19 + default_factory=random_uuid,
20 + description=(
21 + "The request_id related to this request. If the caller does "
22 + "not set it, a random_uuid will be generated. This id is used "
23 + "through out the inference process and return in response."
24 + ),
25 + )
26 +
27 + priority: int = Field(
28 + default=0,
29 + description=(
30 + "The priority of the request (lower means earlier handling; "
31 + "default: 0). Any priority other than 0 will raise an error "
32 + "if the served model does not use priority scheduling."
33 + ),
34 + )
35 +
36 +
37 +class CompletionRequestMixin(OpenAIBaseModel):
38 + input: list[int] | list[list[int]] | str | list[str]
39 +
40 + add_special_tokens: bool = Field(
41 + default=True,
42 + description=(
43 + "If true (the default), special tokens (e.g. BOS) will be added to "
44 + "the prompt."
45 + ),
46 + )
vllm/entrypoints/pooling/classify/protocol.py
+7 -51
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import time
5 -from typing import Annotated, Any, TypeAlias
5 +from typing import Any, TypeAlias
6
7 from pydantic import (
8 Field,
@@ -12,39 +12,15 @@ 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.engine.protocol import OpenAIBaseModel, UsageInfo
15 +from vllm.entrypoints.pooling.base.protocol import (
16 + CompletionRequestMixin,
17 + PoolingBasicRequestMixin,
18 +)
19 from vllm.utils import random_uuid
20
21
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 -
22 +class ClassificationCompletionRequest(PoolingBasicRequestMixin, CompletionRequestMixin):
23 # --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 - )
24 softmax: bool | None = Field(
25 default=None,
26 description="softmax will be deprecated, please use use_activation instead.",
@@ -69,11 +45,8 @@ class ClassificationCompletionRequest(OpenAIBaseModel):
45 )
46
47
72 -class ClassificationChatRequest(OpenAIBaseModel):
73 - model: str | None = None
48 +class ClassificationChatRequest(PoolingBasicRequestMixin):
49 messages: list[ChatCompletionMessageParam]
75 - truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
76 - user: str | None = None
50
51 # --8<-- [start:chat-classification-extra-params]
52 add_generation_prompt: bool = Field(
@@ -119,23 +92,6 @@ class ClassificationChatRequest(OpenAIBaseModel):
92 description=("Additional kwargs to pass to the HF processor."),
93 )
94
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 - )
95 softmax: bool | None = Field(
96 default=None,
97 description="softmax will be deprecated, please use use_activation instead.",
vllm/entrypoints/pooling/embed/protocol.py
+9 -50
@@ -1,7 +1,7 @@
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
4 +from typing import Any, TypeAlias
5
6 from pydantic import (
7 Field,
@@ -11,44 +11,22 @@ from pydantic import (
11 from vllm import PoolingParams
12 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
13 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
14 +from vllm.entrypoints.pooling.base.protocol import (
15 + CompletionRequestMixin,
16 + PoolingBasicRequestMixin,
17 +)
18 from vllm.utils import random_uuid
19 from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
20
21
18 -class EmbeddingCompletionRequest(OpenAIBaseModel):
22 +class EmbeddingCompletionRequest(PoolingBasicRequestMixin, CompletionRequestMixin):
23 # Ordered by official OpenAI API documentation
24 # https://platform.openai.com/docs/api-reference/embeddings
21 - model: str | None = None
22 - input: list[int] | list[list[int]] | str | list[str]
25 +
26 encoding_format: EncodingFormat = "float"
27 dimensions: int | None = None
25 - user: str | None = None
26 - truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
28
29 # --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 - )
30 normalize: bool | None = Field(
31 default=None,
32 description="Whether to normalize the embeddings outputs. Default is True.",
@@ -73,20 +51,17 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
51
52 def to_pooling_params(self):
53 return PoolingParams(
76 - truncate_prompt_tokens=self.truncate_prompt_tokens,
54 dimensions=self.dimensions,
55 use_activation=self.normalize,
56 + truncate_prompt_tokens=self.truncate_prompt_tokens,
57 )
58
59
82 -class EmbeddingChatRequest(OpenAIBaseModel):
83 - model: str | None = None
60 +class EmbeddingChatRequest(PoolingBasicRequestMixin):
61 messages: list[ChatCompletionMessageParam]
62
63 encoding_format: EncodingFormat = "float"
64 dimensions: int | None = None
88 - user: str | None = None
89 - truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
65
66 # --8<-- [start:chat-embedding-extra-params]
67 add_generation_prompt: bool = Field(
@@ -137,22 +112,6 @@ class EmbeddingChatRequest(OpenAIBaseModel):
112 default=None,
113 description=("Additional kwargs to pass to the HF processor."),
114 )
140 - priority: int = Field(
141 - default=0,
142 - description=(
143 - "The priority of the request (lower means earlier handling; "
144 - "default: 0). Any priority other than 0 will raise an error "
145 - "if the served model does not use priority scheduling."
146 - ),
147 - )
148 - request_id: str = Field(
149 - default_factory=random_uuid,
150 - description=(
151 - "The request_id related to this request. If the caller does "
152 - "not set it, a random_uuid will be generated. This id is used "
153 - "through out the inference process and return in response."
154 - ),
155 - )
115 normalize: bool | None = Field(
116 default=None,
117 description="Whether to normalize the embeddings outputs. Default is True.",
vllm/entrypoints/pooling/pooling/protocol.py
+2 -10
@@ -10,6 +10,7 @@ from pydantic import (
10 from vllm import PoolingParams
11 from vllm.config.pooler import get_use_activation
12 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
13 +from vllm.entrypoints.pooling.base.protocol import PoolingBasicRequestMixin
14 from vllm.entrypoints.pooling.embed.protocol import (
15 EmbeddingChatRequest,
16 EmbeddingCompletionRequest,
@@ -72,17 +73,8 @@ class PoolingChatRequest(EmbeddingChatRequest):
73 T = TypeVar("T")
74
75
75 -class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
76 - model: str | None = None
77 -
78 - priority: int = Field(default=0)
79 - """
80 - The priority of the request (lower means earlier handling;
81 - default: 0). Any priority other than 0 will raise an error
82 - if the served model does not use priority scheduling.
83 - """
76 +class IOProcessorRequest(PoolingBasicRequestMixin, Generic[T]):
77 data: T
85 -
78 task: PoolingTask = "plugin"
79 encoding_format: EncodingFormat = "float"
80 embed_dtype: EmbedDType = Field(
vllm/entrypoints/pooling/score/protocol.py
+8 -29
@@ -1,7 +1,7 @@
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
4 +from typing import Any
5
6 from pydantic import (
7 BaseModel,
@@ -11,32 +11,24 @@ from pydantic import (
11 from vllm import PoolingParams
12 from vllm.config.pooler import get_use_activation
13 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
14 -from vllm.entrypoints.score_utils import ScoreContentPartParam, ScoreMultiModalParam
14 +from vllm.entrypoints.pooling.base.protocol import PoolingBasicRequestMixin
15 +from vllm.entrypoints.pooling.score.utils import (
16 + ScoreContentPartParam,
17 + ScoreMultiModalParam,
18 +)
19 from vllm.utils import random_uuid
20
21
18 -class ScoreRequest(OpenAIBaseModel):
19 - model: str | None = None
22 +class ScoreRequest(PoolingBasicRequestMixin):
23 text_1: list[str] | str | ScoreMultiModalParam
24 text_2: list[str] | str | ScoreMultiModalParam
22 - truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
25
26 # --8<-- [start:score-extra-params]
25 -
27 mm_processor_kwargs: dict[str, Any] | None = Field(
28 default=None,
29 description=("Additional kwargs to pass to the HF processor."),
30 )
31
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 -
32 softmax: bool | None = Field(
33 default=None,
34 description="softmax will be deprecated, please use use_activation instead.",
@@ -61,29 +53,16 @@ class ScoreRequest(OpenAIBaseModel):
53 )
54
55
64 -class RerankRequest(OpenAIBaseModel):
65 - model: str | None = None
56 +class RerankRequest(PoolingBasicRequestMixin):
57 query: str | ScoreMultiModalParam
58 documents: list[str] | ScoreMultiModalParam
59 top_n: int = Field(default_factory=lambda: 0)
69 - truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
60
61 # --8<-- [start:rerank-extra-params]
72 -
62 mm_processor_kwargs: dict[str, Any] | None = Field(
63 default=None,
64 description=("Additional kwargs to pass to the HF processor."),
65 )
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 -
66 softmax: bool | None = Field(
67 default=None,
68 description="softmax will be deprecated, please use use_activation instead.",
vllm/entrypoints/pooling/score/serving.py
+1 -1
@@ -25,7 +25,7 @@ from vllm.entrypoints.pooling.score.protocol import (
25 ScoreResponse,
26 ScoreResponseData,
27 )
28 -from vllm.entrypoints.score_utils import (
28 +from vllm.entrypoints.pooling.score.utils import (
29 ScoreContentPartParam,
30 ScoreMultiModalParam,
31 _cosine_similarity,
vllm/entrypoints/pooling/score/utils.py renamed