From 1dcb9ea8cdc8030aaa1ad866e78987eb8341a5da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Mar 2026 10:28:34 +0000 Subject: [PATCH] Fix ~20% SSE streaming overhead in Studio inference Benchmarking showed llama-server through Studio's web layer runs at 284 TPS vs 357 TPS direct (~20% overhead, +0.72ms per token). The generator code itself has zero overhead -- the bottleneck is entirely in how the route handler and middleware pipeline yield SSE chunks. Three structural changes: 1. Convert LoggingMiddleware from BaseHTTPMiddleware to pure ASGI middleware. BaseHTTPMiddleware routes every response body chunk through an anyio memory channel (two coroutine context switches per chunk). The pure ASGI version wraps the send callable directly. 2. Replace per-token asyncio.to_thread / run_in_executor with a single background thread feeding an asyncio.Queue via call_soon_threadsafe (~2us vs ~50us per token). All three streaming loops (gguf_stream_chunks, gguf_tool_stream, stream_chunks) now use the shared _run_gen_to_queue helper. 3. Pre-format JSON for content-token chunks (the 99% hot path) instead of constructing 3 Pydantic objects and calling model_dump_json per token. Pydantic is still used for the first chunk (role), final chunk (finish_reason), and usage chunk where correctness matters more than speed. Additionally, request.is_disconnected() is now checked every 20 tokens instead of every token. The cancel_event mechanism still provides prompt cancellation since the watcher in llama_cpp.py closes the httpx response on cancel. --- studio/backend/loggers/handlers.py | 97 ++++++++++-------- studio/backend/routes/inference.py | 153 ++++++++++++++--------------- 2 files changed, 132 insertions(+), 118 deletions(-) diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index 3add92ea1e..9ffd0b0bd8 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -16,60 +16,79 @@ Key Components: """ import time -from typing import Callable import structlog -from fastapi import Request, Response -from starlette.middleware.base import BaseHTTPMiddleware logger = structlog.get_logger(__name__) +# Paths and suffixes excluded from request logging. +_EXCLUDED_PATHS = frozenset({ + "/api/train/status", + "/api/train/metrics", + "/api/train/hardware", + "/api/system", +}) +_EXCLUDED_SUFFIXES = (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf") + + +class LoggingMiddleware: + """Pure ASGI middleware for request logging. + + Unlike Starlette's BaseHTTPMiddleware, this wraps the ``send`` + callable directly -- no anyio memory channel, no extra coroutine + context switches per response body chunk. This matters for SSE + streaming where hundreds of small chunks are sent per second. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return -class LoggingMiddleware(BaseHTTPMiddleware): - async def dispatch(self, request: Request, call_next: Callable) -> Response: start_time = time.time() + status_code = None + + async def send_wrapper(message): + nonlocal status_code + if message["type"] == "http.response.start": + status_code = message.get("status", 0) + await send(message) try: - response = await call_next(request) - - # Log response - process_time = (time.time() - start_time) * 1000 - - EXCLUDED_PATHS = { - "/api/train/status", - "/api/train/metrics", - "/api/train/hardware", - "/api/system", - } - is_excluded = ( - request.url.path in EXCLUDED_PATHS - or request.url.path.startswith("/assets/") - or request.url.path.endswith( - (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf") - ) - ) - - if not is_excluded: - logger.info( - "request_completed", - method = request.method, - path = request.url.path, - status_code = response.status_code, - process_time_ms = round(process_time, 2), - ) - - return response - + await self.app(scope, receive, send_wrapper) except Exception as e: + path = scope.get("path", "") + method = scope.get("method", "") logger.error( "request_failed", - path = request.url.path, - method = request.method, - error = str(e), - exc_info = True, + path=path, + method=method, + error=str(e), + exc_info=True, ) raise + # Log after response completes (same exclusion logic as before) + process_time = (time.time() - start_time) * 1000 + path = scope.get("path", "") + is_excluded = ( + path in _EXCLUDED_PATHS + or path.startswith("/assets/") + or path.endswith(_EXCLUDED_SUFFIXES) + ) + if not is_excluded: + method = scope.get("method", "") + logger.info( + "request_completed", + method=method, + path=path, + status_code=status_code, + process_time_ms=round(process_time, 2), + ) + def filter_sensitive_data(logger, method_name, event_dict): """Structlog processor to filter out base64 data from logs.""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 78d95fedbd..9e8b359b7e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -22,6 +22,40 @@ import threading import re as _re +# ── Shared helper: run sync generator via background thread + asyncio.Queue ── +# Replaces per-token asyncio.to_thread/run_in_executor calls with a single +# thread that pushes items through call_soon_threadsafe (~2us vs ~50us). +_queue_sentinel = object() +_queue_error = object() + +async def _run_gen_to_queue(gen, cancel_event): + """Run a synchronous generator in a background thread, yield items via asyncio.Queue.""" + loop = asyncio.get_running_loop() + q = asyncio.Queue() + + def _producer(): + try: + for item in gen: + if cancel_event is not None and cancel_event.is_set(): + break + loop.call_soon_threadsafe(q.put_nowait, item) + except Exception as exc: + loop.call_soon_threadsafe(q.put_nowait, (_queue_error, exc)) + finally: + loop.call_soon_threadsafe(q.put_nowait, _queue_sentinel) + + thread = threading.Thread(target=_producer, daemon=True, name="sse-gen") + thread.start() + + while True: + item = await q.get() + if item is _queue_sentinel: + break + if isinstance(item, tuple) and len(item) == 2 and item[0] is _queue_error: + raise item[1] + yield item + + def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" msg = str(exc) @@ -1084,8 +1118,6 @@ async def openai_chat_completions( session_id = payload.session_id, ) - _tool_sentinel = object() - async def gguf_tool_stream(): try: first_chunk = ChatCompletionChunk( @@ -1101,21 +1133,17 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Iterate the synchronous generator in a thread so - # the event loop stays free for disconnect detection. + # Pre-compute static JSON envelope for content tokens (hot path) + _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' + _chunk_suffix = '},"finish_reason":null}]}' + gen = gguf_generate_with_tools() prev_text = "" _stream_usage = None _stream_timings = None - while True: - if await request.is_disconnected(): - cancel_event.set() - return - - event = await asyncio.to_thread(next, gen, _tool_sentinel) - if event is _tool_sentinel: - break - + _disconnect_check_interval = 20 + _token_count = 0 + async for event in _run_gen_to_queue(gen, cancel_event): if event["type"] == "status": # Emit tool status as a custom SSE event status_data = json.dumps( @@ -1142,18 +1170,12 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [ - ChunkChoice( - delta = ChoiceDelta(content = new_text), - finish_reason = None, - ) - ], - ) - yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + _token_count += 1 + if _token_count % _disconnect_check_interval == 0: + if await request.is_disconnected(): + cancel_event.set() + return + yield f"data: {_chunk_prefix}{json.dumps(new_text)}{_chunk_suffix}\n\n" final_chunk = ChatCompletionChunk( id = completion_id, @@ -1230,8 +1252,6 @@ async def openai_chat_completions( enable_thinking = payload.enable_thinking, ) - _gguf_sentinel = object() - if payload.stream: async def gguf_stream_chunks(): @@ -1250,19 +1270,17 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Iterate the synchronous generator in a thread so - # the event loop stays free for disconnect detection. + # Pre-compute static JSON envelope for content tokens (hot path) + _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' + _chunk_suffix = '},"finish_reason":null}]}' + gen = gguf_generate() prev_text = "" _stream_usage = None _stream_timings = None - while True: - if await request.is_disconnected(): - cancel_event.set() - return - cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) - if cumulative is _gguf_sentinel: - break + _disconnect_check_interval = 20 + _token_count = 0 + async for cumulative in _run_gen_to_queue(gen, cancel_event): # Capture server metadata for final usage chunk if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": @@ -1282,18 +1300,12 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [ - ChunkChoice( - delta = ChoiceDelta(content = new_text), - finish_reason = None, - ) - ], - ) - yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + _token_count += 1 + if _token_count % _disconnect_check_interval == 0: + if await request.is_disconnected(): + cancel_event.set() + return + yield f"data: {_chunk_prefix}{json.dumps(new_text)}{_chunk_suffix}\n\n" # Final chunk final_chunk = ChatCompletionChunk( @@ -1456,42 +1468,25 @@ async def openai_chat_completions( yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" prev_text = "" - # Run sync generator in thread pool to avoid blocking - # the event loop. Critical for compare mode: two SSE - # requests arrive concurrently but the orchestrator - # serializes them via _gen_lock. Without run_in_executor - # the second request's blocking lock acquisition would - # freeze the entire event loop, stalling both streams. - _DONE = object() # sentinel for generator exhaustion - loop = asyncio.get_event_loop() + # Pre-compute static JSON envelope for content tokens (hot path) + _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' + _chunk_suffix = '},"finish_reason":null}]}' + gen = generate() - while True: - # next(gen, _DONE) returns _DONE instead of raising - # StopIteration — StopIteration cannot propagate - # through asyncio futures (Python limitation). - cumulative = await loop.run_in_executor(None, next, gen, _DONE) - if cumulative is _DONE: - break - if await request.is_disconnected(): - cancel_event.set() - backend.reset_generation_state() - return + _disconnect_check_interval = 20 + _token_count = 0 + async for cumulative in _run_gen_to_queue(gen, cancel_event): new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue - chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [ - ChunkChoice( - delta = ChoiceDelta(content = new_text), - finish_reason = None, - ) - ], - ) - yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + _token_count += 1 + if _token_count % _disconnect_check_interval == 0: + if await request.is_disconnected(): + cancel_event.set() + backend.reset_generation_state() + return + yield f"data: {_chunk_prefix}{json.dumps(new_text)}{_chunk_suffix}\n\n" final_chunk = ChatCompletionChunk( id = completion_id,