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 15, 07:15 UTC
[Frontend] Add rerank support to run_batch endpoint (#16278)
Signed-off-by: Pooya Davoodi <pooya.davoodi@parasail.io>
Pooya Davoodi committed
May 31, 2025 at 00:40 UTC
dff80b0e42d5d2deebc0a6b0e3d2f6f06bd01c78
3 files changed
+48
-16
tests/entrypoints/openai/test_run_batch.py
+11
-3
@@ -4,6 +4,8 @@ import json
4
import subprocess
5
import tempfile
6
7
+import pytest
8
+
9
from vllm.entrypoints.openai.protocol import BatchRequestOutput
10
11
# ruff: noqa: E501
@@ -23,9 +25,13 @@ INPUT_EMBEDDING_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "
25
{"custom_id": "request-3", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "Hello world!"}}
26
{"custom_id": "request-4", "method": "POST", "url": "/v1/embeddings", "body": {"model": "NonExistModel", "input": "Hello world!"}}"""
27
26
-INPUT_SCORE_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
28
+INPUT_SCORE_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
29
{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
30
31
+INPUT_RERANK_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
32
+{"custom_id": "request-2", "method": "POST", "url": "/v1/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
33
+{"custom_id": "request-2", "method": "POST", "url": "/v2/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
34
+
35
36
def test_empty_file():
37
with tempfile.NamedTemporaryFile(
@@ -105,11 +111,13 @@ def test_embeddings():
111
BatchRequestOutput.model_validate_json(line)
112
113
108
-def test_score():
114
+@pytest.mark.parametrize("input_batch",
115
+ [INPUT_SCORE_BATCH, INPUT_RERANK_BATCH])
116
+def test_score(input_batch):
117
with tempfile.NamedTemporaryFile(
118
"w") as input_file, tempfile.NamedTemporaryFile(
119
"r") as output_file:
112
- input_file.write(INPUT_SCORE_BATCH)
120
+ input_file.write(input_batch)
121
input_file.flush()
122
proc = subprocess.Popen([
123
"vllm",
vllm/entrypoints/openai/protocol.py
+11
-6
@@ -1481,6 +1481,10 @@ class TranscriptionStreamResponse(OpenAIBaseModel):
1481
usage: Optional[UsageInfo] = Field(default=None)
1482
1483
1484
+BatchRequestInputBody = Union[ChatCompletionRequest, EmbeddingRequest,
1485
+ ScoreRequest, RerankRequest]
1486
+
1487
+
1488
class BatchRequestInput(OpenAIBaseModel):
1489
"""
1490
The per-line object of the batch input file.
@@ -1501,21 +1505,22 @@ class BatchRequestInput(OpenAIBaseModel):
1505
url: str
1506
1507
# The parameters of the request.
1504
- body: Union[ChatCompletionRequest, EmbeddingRequest, ScoreRequest]
1508
+ body: BatchRequestInputBody
1509
1510
@field_validator('body', mode='plain')
1511
@classmethod
1512
def check_type_for_url(cls, value: Any, info: ValidationInfo):
1513
# Use url to disambiguate models
1510
- url = info.data['url']
1514
+ url: str = info.data["url"]
1515
if url == "/v1/chat/completions":
1516
return ChatCompletionRequest.model_validate(value)
1517
if url == "/v1/embeddings":
1518
return TypeAdapter(EmbeddingRequest).validate_python(value)
1515
- if url == "/v1/score":
1519
+ if url.endswith("/score"):
1520
return ScoreRequest.model_validate(value)
1517
- return TypeAdapter(Union[ChatCompletionRequest, EmbeddingRequest,
1518
- ScoreRequest]).validate_python(value)
1521
+ if url.endswith("/rerank"):
1522
+ return RerankRequest.model_validate(value)
1523
+ return TypeAdapter(BatchRequestInputBody).validate_python(value)
1524
1525
1526
class BatchResponseData(OpenAIBaseModel):
@@ -1527,7 +1532,7 @@ class BatchResponseData(OpenAIBaseModel):
1532
1533
# The body of the response.
1534
body: Optional[Union[ChatCompletionResponse, EmbeddingResponse,
1530
- ScoreResponse]] = None
1535
+ ScoreResponse, RerankResponse]] = None
1536
1537
1538
class BatchRequestOutput(OpenAIBaseModel):
vllm/entrypoints/openai/run_batch.py
+26
-7
@@ -21,7 +21,7 @@ from vllm.entrypoints.openai.protocol import (BatchRequestInput,
21
BatchResponseData,
22
ChatCompletionResponse,
23
EmbeddingResponse, ErrorResponse,
24
- ScoreResponse)
24
+ RerankResponse, ScoreResponse)
25
# yapf: enable
26
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
27
from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
@@ -274,8 +274,11 @@ async def run_request(serving_engine_func: Callable,
274
tracker: BatchProgressTracker) -> BatchRequestOutput:
275
response = await serving_engine_func(request.body)
276
277
- if isinstance(response,
278
- (ChatCompletionResponse, EmbeddingResponse, ScoreResponse)):
277
+ if isinstance(
278
+ response,
279
+ (ChatCompletionResponse, EmbeddingResponse, ScoreResponse,
280
+ RerankResponse),
281
+ ):
282
batch_output = BatchRequestOutput(
283
id=f"vllm-{random_uuid()}",
284
custom_id=request.custom_id,
@@ -397,7 +400,7 @@ async def main(args):
400
response_futures.append(
401
run_request(embed_handler_fn, request, tracker))
402
tracker.submitted()
400
- elif request.url == "/v1/score":
403
+ elif request.url.endswith("/score"):
404
score_handler_fn = openai_serving_scores.create_score if \
405
openai_serving_scores is not None else None
406
if score_handler_fn is None:
@@ -411,13 +414,29 @@ async def main(args):
414
response_futures.append(
415
run_request(score_handler_fn, request, tracker))
416
tracker.submitted()
417
+ elif request.url.endswith("/rerank"):
418
+ rerank_handler_fn = openai_serving_scores.do_rerank if \
419
+ openai_serving_scores is not None else None
420
+ if rerank_handler_fn is None:
421
+ response_futures.append(
422
+ make_async_error_request_output(
423
+ request,
424
+ error_msg="The model does not support Rerank API",
425
+ ))
426
+ continue
427
+
428
+ response_futures.append(
429
+ run_request(rerank_handler_fn, request, tracker))
430
+ tracker.submitted()
431
else:
432
response_futures.append(
433
make_async_error_request_output(
434
request,
418
- error_msg=
419
- "Only /v1/chat/completions, /v1/embeddings, and /v1/score "
420
- "are supported in the batch endpoint.",
435
+ error_msg=f"URL {request.url} was used. "
436
+ "Supported endpoints: /v1/chat/completions, /v1/embeddings,"
437
+ " /score, /rerank ."
438
+ "See vllm/entrypoints/openai/api_server.py for supported "
439
+ "score/rerank versions.",
440
))
441
442
with tracker.pbar():