[OpenAI] Propagate PD routing metadata through /v1/responses (#35503)

test/registered/unit/entrypoints/openai/test_responses_protocol.py

@@ -28,6 +28,25 @@ def _in_progress_response(request: ResponsesRequest) -> ResponsesResponse:
 
 
 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):
         request = ResponsesRequest(
             model="x",

test/registered/unit/entrypoints/openai/test_serving_responses.py

@@ -922,7 +922,7 @@ class EnginePassthroughTestCase(CustomTestCase):
     """Both flags cross hops with no type contract, and dropping either fails
     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
         # skip_special_tokens off, so mocking it would make that assertion vacuous.
         # chat_template_name=None routes it through the tokenizer's template
@@ -954,9 +954,35 @@ class EnginePassthroughTestCase(CustomTestCase):
             yield context
 
         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
 
+    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):
         serving = make_serving()
         serving.reasoning_parser = "deepseek-r1"