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 11, 05:00 UTC
[Bugfix] fix DP-aware routing in OpenAI API requests (#29002)
Signed-off-by: inkcherry <mingzhi.liu@amd.com>
inkcherry committed
Dec 19, 2025 at 01:50 UTC
500f26e6d35d9743167bc4908c01ca356e543836
7 files changed
+68
tests/entrypoints/openai/test_chat_error.py
+1
@@ -76,6 +76,7 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
76
lora_request,
77
trace_headers,
78
priority,
79
+ data_parallel_rank,
80
):
81
return dict(engine_prompt), {}
82
tests/entrypoints/openai/test_completion_error.py
+1
@@ -73,6 +73,7 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion:
73
lora_request,
74
trace_headers,
75
priority,
76
+ data_parallel_rank,
77
):
78
return dict(engine_prompt), {}
79
tests/entrypoints/openai/test_serving_chat.py
+1
@@ -396,6 +396,7 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
396
lora_request,
397
trace_headers,
398
priority,
399
+ data_parallel_rank,
400
):
401
return dict(engine_prompt), {}
402
tests/v1/engine/test_async_llm.py
+61
@@ -11,6 +11,13 @@ from vllm import SamplingParams
11
from vllm.assets.image import ImageAsset
12
from vllm.config import VllmConfig
13
from vllm.engine.arg_utils import AsyncEngineArgs
14
+from vllm.entrypoints.openai.protocol import (
15
+ ChatCompletionRequest,
16
+ ChatCompletionResponse,
17
+ ErrorResponse,
18
+)
19
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
20
+from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
21
from vllm.inputs import PromptType
22
from vllm.outputs import RequestOutput
23
from vllm.platforms import current_platform
@@ -484,6 +491,60 @@ async def test_dp_rank_argument():
491
pass
492
493
494
+@pytest.mark.asyncio(scope="module")
495
+async def test_header_dp_rank_argument():
496
+ with ExitStack() as after:
497
+ with set_default_torch_num_threads(1):
498
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
499
+ after.callback(engine.shutdown)
500
+
501
+ MODEL_NAME = "test-model"
502
+ BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
503
+
504
+ # Create models first
505
+ models = OpenAIServingModels(
506
+ engine_client=engine,
507
+ base_model_paths=BASE_MODEL_PATHS,
508
+ )
509
+
510
+ # Create serving chat instance
511
+ serving_chat = OpenAIServingChat(
512
+ engine_client=engine,
513
+ models=models,
514
+ response_role="assistant",
515
+ chat_template=None,
516
+ chat_template_content_format="auto",
517
+ request_logger=None,
518
+ )
519
+ # Create a chat completion request
520
+ req = ChatCompletionRequest(
521
+ model=MODEL_NAME,
522
+ messages=[{"role": "user", "content": TEXT_PROMPT}],
523
+ max_tokens=100,
524
+ temperature=1.0,
525
+ seed=33,
526
+ )
527
+ # Test 1: Valid DP rank (0)
528
+ mock_raw_request = MagicMock()
529
+ mock_raw_request.headers = {"X-data-parallel-rank": "0"}
530
+ mock_raw_request.state = MagicMock()
531
+
532
+ # Should succeed with valid rank
533
+ response = await serving_chat.create_chat_completion(req, mock_raw_request)
534
+ assert isinstance(response, ChatCompletionResponse), (
535
+ "Expected a ChatCompletionResponse for valid DP rank"
536
+ )
537
+
538
+ # Test 2: Out-of-range DP rank (1)
539
+ mock_raw_request.headers = {"X-data-parallel-rank": "1"}
540
+
541
+ # should return ErrorResponse for out-of-range rank
542
+ response2 = await serving_chat.create_chat_completion(req, mock_raw_request)
543
+ assert isinstance(response2, ErrorResponse), (
544
+ "Expected an ErrorResponse for out-of-range DP rank"
545
+ )
546
+
547
+
548
@pytest.mark.asyncio
549
async def test_check_health():
550
"""Test that check_health returns normally for healthy engine
vllm/entrypoints/openai/serving_chat.py
+1
@@ -381,6 +381,7 @@ class OpenAIServingChat(OpenAIServing):
381
lora_request=lora_request,
382
trace_headers=trace_headers,
383
priority=request.priority,
384
+ data_parallel_rank=data_parallel_rank,
385
)
386
387
generator = self.engine_client.generate(
vllm/entrypoints/openai/serving_completion.py
+1
@@ -230,6 +230,7 @@ class OpenAIServingCompletion(OpenAIServing):
230
lora_request=lora_request,
231
trace_headers=trace_headers,
232
priority=request.priority,
233
+ data_parallel_rank=data_parallel_rank,
234
)
235
236
generator = self.engine_client.generate(
vllm/entrypoints/openai/serving_engine.py
+2
@@ -1231,6 +1231,7 @@ class OpenAIServing:
1231
lora_request: LoRARequest | None,
1232
trace_headers: Mapping[str, str] | None,
1233
priority: int,
1234
+ data_parallel_rank: int | None = None,
1235
) -> tuple[EngineCoreRequest, dict[str, Any]]:
1236
"""Use the Processor to process inputs for AsyncLLM."""
1237
tokenization_kwargs: dict[str, Any] = {}
@@ -1246,6 +1247,7 @@ class OpenAIServing:
1247
tokenization_kwargs=tokenization_kwargs,
1248
trace_headers=trace_headers,
1249
priority=priority,
1250
+ data_parallel_rank=data_parallel_rank,
1251
)
1252
return engine_request, tokenization_kwargs
1253