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",