From a35fbe22eacdfad8b19e6f1842b7ac74cd4ff855 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 17 Jun 2026 22:38:02 -0700 Subject: [PATCH] Studio: serialize non-streaming responses once and pool the proxy client (#6393) * 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> --- studio/backend/core/inference/llama_http.py | 56 +++++++++++++ studio/backend/main.py | 4 + studio/backend/routes/inference.py | 74 +++++++++------- .../tests/test_api_perf_serialization.py | 84 +++++++++++++++++++ .../tests/test_openai_tool_passthrough.py | 18 ++-- 5 files changed, 195 insertions(+), 41 deletions(-) create mode 100644 studio/backend/core/inference/llama_http.py create mode 100644 studio/backend/tests/test_api_perf_serialization.py diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py new file mode 100644 index 0000000000..b554949c3e --- /dev/null +++ b/studio/backend/core/inference/llama_http.py @@ -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 diff --git a/studio/backend/main.py b/studio/backend/main.py index aaeb7917d8..2313686a69 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 71cee33b24..0b95ad3987 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 diff --git a/studio/backend/tests/test_api_perf_serialization.py b/studio/backend/tests/test_api_perf_serialization.py new file mode 100644 index 0000000000..f5ad53306d --- /dev/null +++ b/studio/backend/tests/test_api_perf_serialization.py @@ -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] diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 8442d05cdc..2586076321 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -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",