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, 04:45 UTC
[PD] let toy proxy handle /chat/completions (#19730)
Signed-off-by: Linkun <github@lkchen.net>
lkchen committed
Jun 25, 2025 at 12:17 UTC
4734704b30efff60feebf5f524ca9a3f5f8ac23a
1 file changed
+15
-7
tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
+15
-7
@@ -196,8 +196,7 @@ async def stream_service_response(client_info: dict, endpoint: str,
196
yield chunk
197
198
199
-@app.post("/v1/completions")
200
-async def handle_completions(request: Request):
199
+async def _handle_completions(api: str, request: Request):
200
try:
201
req_data = await request.json()
202
request_id = str(uuid.uuid4())
@@ -206,9 +205,8 @@ async def handle_completions(request: Request):
205
prefill_client_info = get_next_client(request.app, 'prefill')
206
207
# Send request to prefill service
209
- response = await send_request_to_service(prefill_client_info,
210
- "/completions", req_data,
211
- request_id)
208
+ response = await send_request_to_service(prefill_client_info, api,
209
+ req_data, request_id)
210
211
# Extract the needed fields
212
response_json = response.json()
@@ -224,7 +222,7 @@ async def handle_completions(request: Request):
222
# Stream response from decode service
223
async def generate_stream():
224
async for chunk in stream_service_response(decode_client_info,
227
- "/completions",
225
+ api,
226
req_data,
227
request_id=request_id):
228
yield chunk
@@ -237,12 +235,22 @@ async def handle_completions(request: Request):
235
import traceback
236
exc_info = sys.exc_info()
237
print("Error occurred in disagg prefill proxy server"
240
- " - completions endpoint")
238
+ f" - {api} endpoint")
239
print(e)
240
print("".join(traceback.format_exception(*exc_info)))
241
raise
242
243
244
+@app.post("/v1/completions")
245
+async def handle_completions(request: Request):
246
+ return await _handle_completions("/completions", request)
247
+
248
+
249
+@app.post("/v1/chat/completions")
250
+async def handle_chat_completions(request: Request):
251
+ return await _handle_completions("/chat/completions", request)
252
+
253
+
254
@app.get("/healthcheck")
255
async def healthcheck():
256
"""Simple endpoint to check if the server is running."""