[OpenAI] Propagate PD routing metadata through /v1/responses (#35503)
This commit is contained in:
@@ -1652,6 +1652,18 @@ class ResponsesRequest(BaseModel):
|
|||||||
default=None, description="Cache salt for request caching"
|
default=None, description="Cache salt for request caching"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# For PD disaggregation
|
||||||
|
bootstrap_host: Optional[Union[List[str], str]] = None
|
||||||
|
bootstrap_port: Optional[Union[List[Optional[int]], int]] = None
|
||||||
|
bootstrap_room: Optional[Union[List[int], int]] = None
|
||||||
|
|
||||||
|
# For DP routing — external router assigns a specific DP worker
|
||||||
|
routed_dp_rank: Optional[int] = None
|
||||||
|
# For PD disagg — hint telling decode which prefill DP worker has the KV cache
|
||||||
|
disagg_prefill_dp_rank: Optional[int] = None
|
||||||
|
# Deprecated: use routed_dp_rank instead
|
||||||
|
data_parallel_rank: Optional[int] = None
|
||||||
|
|
||||||
# SGLang sampling extras. ``None`` defers to ``--preferred-sampling-params``.
|
# SGLang sampling extras. ``None`` defers to ``--preferred-sampling-params``.
|
||||||
frequency_penalty: float = 0.0
|
frequency_penalty: float = 0.0
|
||||||
presence_penalty: float = 0.0
|
presence_penalty: float = 0.0
|
||||||
@@ -1669,6 +1681,11 @@ class ResponsesRequest(BaseModel):
|
|||||||
"repetition_penalty": 1.0,
|
"repetition_penalty": 1.0,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@model_validator(mode="before")
|
||||||
|
@classmethod
|
||||||
|
def _handle_deprecated_dp_rank(cls, values):
|
||||||
|
return _migrate_deprecated_dp_rank(values)
|
||||||
|
|
||||||
@model_validator(mode="before")
|
@model_validator(mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
def normalize_reasoning_to_thinking(cls, values):
|
def normalize_reasoning_to_thinking(cls, values):
|
||||||
|
|||||||
@@ -435,6 +435,10 @@ class OpenAIServingResponses(OpenAIServingChat):
|
|||||||
else {}
|
else {}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
effective_routed_dp_rank = self.extract_routed_dp_rank_from_header(
|
||||||
|
raw_request, request.routed_dp_rank
|
||||||
|
)
|
||||||
|
|
||||||
adapted_request = GenerateReqInput(
|
adapted_request = GenerateReqInput(
|
||||||
**prompt_kwargs,
|
**prompt_kwargs,
|
||||||
**logprob_kwargs,
|
**logprob_kwargs,
|
||||||
@@ -464,6 +468,11 @@ class OpenAIServingResponses(OpenAIServingChat):
|
|||||||
session_id=request.session_id,
|
session_id=request.session_id,
|
||||||
extra_key=request.extra_key,
|
extra_key=request.extra_key,
|
||||||
cache_salt=request.cache_salt,
|
cache_salt=request.cache_salt,
|
||||||
|
bootstrap_host=request.bootstrap_host,
|
||||||
|
bootstrap_port=request.bootstrap_port,
|
||||||
|
bootstrap_room=request.bootstrap_room,
|
||||||
|
routed_dp_rank=effective_routed_dp_rank,
|
||||||
|
disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
|
||||||
# background+stream streams on this connection, so don't detach.
|
# background+stream streams on this connection, so don't detach.
|
||||||
background=request.background and not request.stream,
|
background=request.background and not request.stream,
|
||||||
require_reasoning=require_reasoning,
|
require_reasoning=require_reasoning,
|
||||||
|
|||||||
@@ -28,6 +28,25 @@ def _in_progress_response(request: ResponsesRequest) -> ResponsesResponse:
|
|||||||
|
|
||||||
|
|
||||||
class ResponsesRequestTestCase(CustomTestCase):
|
class ResponsesRequestTestCase(CustomTestCase):
|
||||||
|
def test_pd_routing_fields(self):
|
||||||
|
with self.assertWarns(DeprecationWarning):
|
||||||
|
request = ResponsesRequest(
|
||||||
|
model="x",
|
||||||
|
input="hi",
|
||||||
|
bootstrap_host="10.0.0.1",
|
||||||
|
bootstrap_port=8998,
|
||||||
|
bootstrap_room=42,
|
||||||
|
data_parallel_rank=1,
|
||||||
|
disagg_prefill_dp_rank=0,
|
||||||
|
store=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(request.bootstrap_host, "10.0.0.1")
|
||||||
|
self.assertEqual(request.bootstrap_port, 8998)
|
||||||
|
self.assertEqual(request.bootstrap_room, 42)
|
||||||
|
self.assertEqual(request.routed_dp_rank, 1)
|
||||||
|
self.assertEqual(request.disagg_prefill_dp_rank, 0)
|
||||||
|
|
||||||
def test_function_tool_accepted(self):
|
def test_function_tool_accepted(self):
|
||||||
request = ResponsesRequest(
|
request = ResponsesRequest(
|
||||||
model="x",
|
model="x",
|
||||||
|
|||||||
@@ -922,7 +922,7 @@ class EnginePassthroughTestCase(CustomTestCase):
|
|||||||
"""Both flags cross hops with no type contract, and dropping either fails
|
"""Both flags cross hops with no type contract, and dropping either fails
|
||||||
silently."""
|
silently."""
|
||||||
|
|
||||||
def _capture(self, serving, request):
|
def _capture(self, serving, request, raw_request=None):
|
||||||
# Let the real _process_messages run: it is the hop that turns
|
# Let the real _process_messages run: it is the hop that turns
|
||||||
# skip_special_tokens off, so mocking it would make that assertion vacuous.
|
# skip_special_tokens off, so mocking it would make that assertion vacuous.
|
||||||
# chat_template_name=None routes it through the tokenizer's template
|
# chat_template_name=None routes it through the tokenizer's template
|
||||||
@@ -954,9 +954,35 @@ class EnginePassthroughTestCase(CustomTestCase):
|
|||||||
yield context
|
yield context
|
||||||
|
|
||||||
serving._generate_with_builtin_tools = fake_generate
|
serving._generate_with_builtin_tools = fake_generate
|
||||||
asyncio.run(serving.create_responses(request))
|
asyncio.run(serving.create_responses(request, raw_request=raw_request))
|
||||||
return captured
|
return captured
|
||||||
|
|
||||||
|
def test_pd_routing_fields_forwarded_to_engine(self):
|
||||||
|
serving = make_serving()
|
||||||
|
raw_request = Mock(headers={"x-data-parallel-rank": "2"}, state=Mock())
|
||||||
|
|
||||||
|
captured = self._capture(
|
||||||
|
serving,
|
||||||
|
ResponsesRequest(
|
||||||
|
model="x",
|
||||||
|
input="hi",
|
||||||
|
bootstrap_host="10.0.0.1",
|
||||||
|
bootstrap_port=8998,
|
||||||
|
bootstrap_room=42,
|
||||||
|
routed_dp_rank=1,
|
||||||
|
disagg_prefill_dp_rank=0,
|
||||||
|
store=False,
|
||||||
|
),
|
||||||
|
raw_request=raw_request,
|
||||||
|
)
|
||||||
|
|
||||||
|
adapted_request = captured["adapted_request"]
|
||||||
|
self.assertEqual(adapted_request.bootstrap_host, "10.0.0.1")
|
||||||
|
self.assertEqual(adapted_request.bootstrap_port, 8998)
|
||||||
|
self.assertEqual(adapted_request.bootstrap_room, 42)
|
||||||
|
self.assertEqual(adapted_request.routed_dp_rank, 2)
|
||||||
|
self.assertEqual(adapted_request.disagg_prefill_dp_rank, 0)
|
||||||
|
|
||||||
def test_require_reasoning_forwarded_when_reasoning_parser_configured(self):
|
def test_require_reasoning_forwarded_when_reasoning_parser_configured(self):
|
||||||
serving = make_serving()
|
serving = make_serving()
|
||||||
serving.reasoning_parser = "deepseek-r1"
|
serving.reasoning_parser = "deepseek-r1"
|
||||||
|
|||||||
Reference in New Issue
Block a user