[Bugfix][Examples] Launch prefill and decode concurrently in NixlPushConnector demo (#55088)
Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com>
junuxyz committed
Sep 12, 2026 at 15:22 UTC
372477d2d9890381dbc794cf66ef86ff989f04e0
1 file changed
+18
-10
examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py
+18
-10
@@ -31,6 +31,7 @@ disagg_proxy_pushconnector_demo.py \
31
"""
32
33
import argparse
34
+import asyncio
35
import contextlib
36
import ipaddress
37
import itertools
@@ -259,17 +260,24 @@ class PushProxy:
260
decode_instance = self.schedule(self.decode_cycler)
261
headers = self._common_headers(request_id)
262
262
- # Fire prefill; we don't read its body but must drain the
263
- # connection so the upstream server can free its slot.
264
- async for _ in self.forward_request(
265
- f"http://{prefill_instance}{path}", prefill_request, headers
266
- ):
267
- continue
263
+ async def drain_prefill():
264
+ async for _ in self.forward_request(
265
+ f"http://{prefill_instance}{path}", prefill_request, headers
266
+ ):
267
+ continue
268
269
- generator = self.forward_request(
270
- f"http://{decode_instance}{path}", decode_request, headers
271
- )
272
- return StreamingResponse(generator, media_type="application/json")
269
+ prefill_task = asyncio.create_task(drain_prefill())
270
+
271
+ async def stream_decode():
272
+ try:
273
+ async for chunk in self.forward_request(
274
+ f"http://{decode_instance}{path}", decode_request, headers
275
+ ):
276
+ yield chunk
277
+ finally:
278
+ await prefill_task
279
+
280
+ return StreamingResponse(stream_decode(), media_type="application/json")
281
282
async def create_completion(self, raw_request: Request):
283
try: