a35fbe22ea* Studio: serialize non-streaming responses once and pool the proxy client Two safe latency wins on the OpenAI/Anthropic-compatible endpoints that leave the streaming generation paths untouched (they keep Connection: close and max_keepalive_connections=0 so a client disconnect still stops GPU decode). 1. Non-streaming responses used JSONResponse(content=model.model_dump()), which builds a dict and then re-runs json.dumps. Serialize once with model.model_dump_json() via a small _model_json_response helper. The body is byte-identical (nulls preserved), about 3x faster to encode in a microbench. 2. The non-streaming completions and embeddings proxies built a fresh httpx.AsyncClient per request. Route them through one pooled client (core/inference/llama_http) closed on shutdown; streaming generation keeps its own per-request close-only client. About 5x faster per call to the local llama-server in a microbench. The existing API-monitor tests for the non-streaming completions, embeddings and passthrough paths now patch nonstreaming_client instead of httpx.AsyncClient to match the pooled client, so they stay deterministic. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make the pooled non-streaming client per event loop Review follow-up on the shared httpx client. It was a single module-global instance, which has two lifecycle problems the per-request client did not: 1. After aclose() in lifespan shutdown, nonstreaming_client() kept handing back the closed client, so a second lifespan in the same process (repeated TestClient, embedded restart) failed with "client has been closed". 2. An httpx client binds its transport to the loop it first runs on, so reuse from another loop could raise "Event loop is closed". Hold one client per running loop in a WeakKeyDictionary, recreate when missing or closed, and close all on shutdown. Single-loop production is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""_model_json_response produces the same body as JSONResponse(model.model_dump())."""
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Optional
|
|
|
|
from fastapi.responses import JSONResponse
|
|
from pydantic import BaseModel
|
|
|
|
import routes.inference as inference_route
|
|
from core.inference import llama_http
|
|
|
|
|
|
class _Usage(BaseModel):
|
|
prompt_tokens: int = 3
|
|
completion_tokens: int = 5
|
|
details: Optional[dict] = None
|
|
|
|
|
|
class _Choice(BaseModel):
|
|
index: int = 0
|
|
text: str = "hello"
|
|
logprobs: Optional[dict] = None
|
|
|
|
|
|
class _Resp(BaseModel):
|
|
id: str = "chatcmpl-abc"
|
|
object: str = "chat.completion"
|
|
created: int = 1700000000
|
|
model: str = "unsloth/SmolLM2-135M-Instruct-GGUF"
|
|
choices: list[_Choice] = [_Choice()]
|
|
usage: _Usage = _Usage()
|
|
system_fingerprint: Optional[str] = None
|
|
|
|
|
|
def _old_body(model) -> bytes:
|
|
# What the previous code emitted: dict -> Starlette json.dumps.
|
|
return JSONResponse(content = model.model_dump()).body
|
|
|
|
|
|
def test_body_matches_old_jsonresponse():
|
|
model = _Resp()
|
|
resp = inference_route._model_json_response(model)
|
|
# Same decoded JSON (key order is irrelevant once parsed), nulls preserved.
|
|
assert json.loads(resp.body) == json.loads(_old_body(model))
|
|
assert json.loads(resp.body)["system_fingerprint"] is None # null kept, not dropped
|
|
|
|
|
|
def test_media_type_and_status():
|
|
resp = inference_route._model_json_response(_Resp(), status_code = 200)
|
|
assert resp.media_type == "application/json"
|
|
assert resp.status_code == 200
|
|
err = inference_route._model_json_response(_Resp(), status_code = 503)
|
|
assert err.status_code == 503
|
|
|
|
|
|
def test_pooled_client_reused_within_loop_and_recreated_after_close():
|
|
async def _scenario():
|
|
a = llama_http.nonstreaming_client()
|
|
b = llama_http.nonstreaming_client()
|
|
assert a is b # reused within one loop
|
|
await llama_http.aclose()
|
|
assert a.is_closed
|
|
c = llama_http.nonstreaming_client() # must not return the closed client
|
|
assert c is not a and not c.is_closed
|
|
await llama_http.aclose()
|
|
|
|
asyncio.run(_scenario())
|
|
|
|
|
|
def test_pooled_client_is_per_event_loop():
|
|
clients = []
|
|
# Each asyncio.run uses a fresh loop; the pooled client must not leak across.
|
|
for _ in range(2):
|
|
|
|
async def _grab():
|
|
clients.append(llama_http.nonstreaming_client())
|
|
await llama_http.aclose()
|
|
|
|
asyncio.run(_grab())
|
|
assert clients[0] is not clients[1]
|