Skip to content

[PD] let toy proxy handle /chat/completions #19730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 25, 2025
Merged

[PD] let toy proxy handle /chat/completions #19730

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,8 +196,7 @@ async def stream_service_response(client_info: dict, endpoint: str,
yield chunk


@app.post("/v1/completions")
async def handle_completions(request: Request):
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
request_id = str(uuid.uuid4())
Expand All @@ -206,9 +205,8 @@ async def handle_completions(request: Request):
prefill_client_info = get_next_client(request.app, 'prefill')

# Send request to prefill service
response = await send_request_to_service(prefill_client_info,
"/completions", req_data,
request_id)
response = await send_request_to_service(prefill_client_info, api,
req_data, request_id)

# Extract the needed fields
response_json = response.json()
Expand All @@ -224,7 +222,7 @@ async def handle_completions(request: Request):
# Stream response from decode service
async def generate_stream():
async for chunk in stream_service_response(decode_client_info,
"/completions",
api,
req_data,
request_id=request_id):
yield chunk
Expand All @@ -237,12 +235,22 @@ async def generate_stream():
import traceback
exc_info = sys.exc_info()
print("Error occurred in disagg prefill proxy server"
" - completions endpoint")
f" - {api} endpoint")
Comment on lines 237 to +238
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This error message has been updated to be dynamic using an f-string, which is a good improvement.

For more robust error handling in server applications, consider using the existing logger instance (available as logger in this file, initialized on line 16) instead of print. For example, using logger.error() or logger.exception().

If logger.exception(f"Error occurred in disagg prefill proxy server - {api} endpoint: {e!r}") were used here, it would also capture and log the full traceback automatically. This could simplify the subsequent print(e) (line 239) and print(traceback.format_exception(...)) (line 240) calls, though those specific lines are outside the scope of this particular change.

        logger.error(f"Error occurred in disagg prefill proxy server - {api} endpoint. Exception: {e!r}")

print(e)
print("".join(traceback.format_exception(*exc_info)))
raise


@app.post("/v1/completions")
async def handle_completions(request: Request):
return await _handle_completions("/completions", request)


@app.post("/v1/chat/completions")
async def handle_chat_completions(request: Request):
return await _handle_completions("/chat/completions", request)


@app.get("/healthcheck")
async def healthcheck():
"""Simple endpoint to check if the server is running."""
Expand Down