ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: serialize non-streaming responses once and pool the proxy client (#6393)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-06-17 22:38:02 -07:00 • committed by GitHub
commit a35fbe22ea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 195 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

56
studio/backend/core/inference/llama_http.py Normal file
View file

@ -0,0 +1,56 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared pooled httpx.AsyncClient for NON-streaming calls to the local llama-server.
Streaming generation must NOT use this. It relies on ``Connection: close`` and
``max_keepalive_connections=0`` so a client disconnect tears down the upstream
socket and stops GPU decode (PR #5749). This pooled client is only for short
request/response proxy calls (non-streaming completions, embeddings) where
reusing a connection removes per-request setup cost. Per-request ``timeout`` is
still passed at each call site.
"""
from __future__ import annotations
import asyncio
import weakref
import httpx
_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32)
def _new_client() -> httpx.AsyncClient:
try:
return httpx.AsyncClient(limits = _LIMITS)
except Exception:
# Mirror external_provider: an unsupported env proxy scheme can raise.
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
# One client per running event loop: an httpx client binds its transport to the
# loop it first runs on, so a single global instance breaks across a lifespan
# restart or a second test loop. Weak keys let a finished loop drop its client.
_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = (
weakref.WeakKeyDictionary()
)
def nonstreaming_client() -> httpx.AsyncClient:
loop = asyncio.get_running_loop()
client = _clients.get(loop)
if client is None or client.is_closed:
client = _new_client()
_clients[loop] = client
return client
async def aclose() -> None:
clients = list(_clients.values())
_clients.clear()
for client in clients:
try:
await client.aclose()
except Exception:
pass

4
studio/backend/main.py
View file

@ -465,6 +465,10 @@ async def lifespan(app: FastAPI):
app.state.bootstrap_password = storage.get_bootstrap_password()
yield
from core.inference.llama_http import aclose as _close_llama_http
await _close_llama_http()
await run_lifespan_shutdown(
terminate_hub_downloads,
clear_unsloth_compiled_cache,

74
studio/backend/routes/inference.py
View file

@ -878,6 +878,7 @@ from state.tool_approvals import resolve_tool_decision
from core.inference.key_exchange import decrypt_api_key
from core.inference.api_monitor import api_monitor
from core.inference.llama_http import nonstreaming_client
from core.inference.providers import get_provider_info, get_base_url
from core.inference.external_provider import ExternalProviderClient
from core.inference.chat_templates import resolve_effective_chat_template_override
@ -1922,6 +1923,19 @@ def get_llama_cpp_backend() -> LlamaCppBackend:
return _llama_cpp_backend
def _model_json_response(model, status_code: int = 200) -> Response:
"""Serialize a pydantic response once via pydantic-core.
Equivalent body to ``JSONResponse(content = model.model_dump())`` but
avoids the dict round-trip plus Starlette's second ``json.dumps``.
"""
return Response(
content = model.model_dump_json(),
media_type = "application/json",
status_code = status_code,
)
@router.post("/load", response_model = LoadResponse)
async def load_model(
request: LoadRequest,
@ -4318,7 +4332,7 @@ async def openai_chat_completions(
)
],
)
return JSONResponse(content = response.model_dump())
return _model_json_response(response)
if monitor_id is None and not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
@ -4983,7 +4997,7 @@ async def openai_chat_completions(
_monitor_context_length(),
)
api_monitor.finish(monitor_id)
return JSONResponse(content = response.model_dump())
return _model_json_response(response)
except Exception as e:
logger.error(f"Error during GGUF completion: {e}", exc_info = True)
@ -5383,7 +5397,7 @@ async def openai_chat_completions(
)
],
)
return JSONResponse(content = response.model_dump())
return _model_json_response(response)
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
@ -5593,7 +5607,7 @@ async def openai_chat_completions(
if _stats:
_monitor_usage(monitor_id, _stats.get("usage"))
api_monitor.finish(monitor_id)
return JSONResponse(content = response.model_dump())
return _model_json_response(response)
except Exception as e:
backend.reset_generation_state()
@ -5941,12 +5955,11 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
return StreamingResponse(_stream(), media_type = "text/event-stream")
else:
try:
async with httpx.AsyncClient() as client:
resp = await client.post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
@ -6011,12 +6024,11 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get
)
try:
async with httpx.AsyncClient() as client:
resp = await client.post(
target_url,
json = body,
timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
)
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
@ -6649,7 +6661,7 @@ async def _responses_non_streaming(
api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls))
_monitor_usage(monitor_id, usage_data, _monitor_context_length())
api_monitor.finish(monitor_id)
return JSONResponse(content = response.model_dump())
return _model_json_response(response)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
@ -8253,7 +8265,7 @@ async def _anthropic_tool_non_streaming(
output_tokens = usage.get("completion_tokens", 0),
),
)
return JSONResponse(content = resp.model_dump())
return _model_json_response(resp)
async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
@ -8295,7 +8307,7 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
output_tokens = usage.get("completion_tokens", 0),
),
)
return JSONResponse(content = resp.model_dump())
return _model_json_response(resp)
# =====================================================================
@ -8611,12 +8623,11 @@ async def _anthropic_passthrough_non_streaming(
backend_ctx = llama_backend.context_length,
)
async with httpx.AsyncClient() as client:
resp = await client.post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
if resp.status_code != 200:
raise HTTPException(
@ -8667,7 +8678,7 @@ async def _anthropic_passthrough_non_streaming(
output_tokens = usage.get("completion_tokens", 0),
),
)
return JSONResponse(content = resp_obj.model_dump())
return _model_json_response(resp_obj)
# =====================================================================
@ -9236,12 +9247,11 @@ async def _openai_passthrough_non_streaming(
)
while True:
try:
async with httpx.AsyncClient() as client:
resp = await client.post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise

84
studio/backend/tests/test_api_perf_serialization.py Normal file
View file

@ -0,0 +1,84 @@
# 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]

18
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -1771,9 +1771,9 @@ class TestApiMonitorProviderAndCompletionStreams:
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(
inf_mod.httpx,
"AsyncClient",
lambda *args, **kwargs: FailingAsyncClient(),
inf_mod,
"nonstreaming_client",
lambda: FailingAsyncClient(),
)
monkeypatch.setattr(
inf_mod,
@ -1899,9 +1899,9 @@ class TestApiMonitorProviderAndCompletionStreams:
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(
inf_mod.httpx,
"AsyncClient",
lambda *args, **kwargs: FakeAsyncClient(),
inf_mod,
"nonstreaming_client",
lambda: FakeAsyncClient(),
)
monkeypatch.setattr(
inf_mod,
@ -2014,9 +2014,9 @@ class TestApiMonitorProviderAndCompletionStreams:
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(
inf_mod.httpx,
"AsyncClient",
lambda *args, **kwargs: CancellingAsyncClient(),
inf_mod,
"nonstreaming_client",
lambda: CancellingAsyncClient(),
)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",