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>
56 lines
1.9 KiB
Python
56 lines
1.9 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
|
|
|
|
"""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
|