From 1dcb9ea8cdc8030aaa1ad866e78987eb8341a5da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Mar 2026 10:28:34 +0000 Subject: [PATCH 1/4] 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, From b2da1e070b881941539f28bc3a55b32df5e823cb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 10:29:01 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/loggers/handlers.py | 30 ++++++++++++++++-------------- studio/backend/routes/inference.py | 3 ++- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index 9ffd0b0bd8..9cade6178a 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -22,12 +22,14 @@ import structlog 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_PATHS = frozenset( + { + "/api/train/status", + "/api/train/metrics", + "/api/train/hardware", + "/api/system", + } +) _EXCLUDED_SUFFIXES = (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf") @@ -64,10 +66,10 @@ class LoggingMiddleware: method = scope.get("method", "") logger.error( "request_failed", - path=path, - method=method, - error=str(e), - exc_info=True, + path = path, + method = method, + error = str(e), + exc_info = True, ) raise @@ -83,10 +85,10 @@ class LoggingMiddleware: method = scope.get("method", "") logger.info( "request_completed", - method=method, - path=path, - status_code=status_code, - process_time_ms=round(process_time, 2), + method = method, + path = path, + status_code = status_code, + process_time_ms = round(process_time, 2), ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9e8b359b7e..3d59a70b6f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -28,6 +28,7 @@ import re as _re _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() @@ -44,7 +45,7 @@ async def _run_gen_to_queue(gen, cancel_event): finally: loop.call_soon_threadsafe(q.put_nowait, _queue_sentinel) - thread = threading.Thread(target=_producer, daemon=True, name="sse-gen") + thread = threading.Thread(target = _producer, daemon = True, name = "sse-gen") thread.start() while True: From 3b1a416cd4393b0e69111bfe0f191360e26c3eb8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Mar 2026 10:34:12 +0000 Subject: [PATCH 3/4] Fix hand-rolled JSON to match Pydantic exclude_none=True output The old code used model_dump_json(exclude_none=True) which omits finish_reason entirely when it is None. The hand-rolled suffix was incorrectly including "finish_reason":null. Removed it so the SSE output is byte-for-byte identical to the old Pydantic path. --- studio/backend/routes/inference.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3d59a70b6f..a71fd1b8cc 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1134,9 +1134,10 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Pre-compute static JSON envelope for content tokens (hot path) + # Pre-compute static JSON envelope for content tokens (hot path). + # Must match model_dump_json(exclude_none=True): no finish_reason when null. _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}]}' + _chunk_suffix = '}}]}' gen = gguf_generate_with_tools() prev_text = "" @@ -1271,9 +1272,10 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Pre-compute static JSON envelope for content tokens (hot path) + # Pre-compute static JSON envelope for content tokens (hot path). + # Must match model_dump_json(exclude_none=True): no finish_reason when null. _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}]}' + _chunk_suffix = '}}]}' gen = gguf_generate() prev_text = "" @@ -1469,9 +1471,10 @@ async def openai_chat_completions( yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" prev_text = "" - # Pre-compute static JSON envelope for content tokens (hot path) + # Pre-compute static JSON envelope for content tokens (hot path). + # Must match model_dump_json(exclude_none=True): no finish_reason when null. _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}]}' + _chunk_suffix = '}}]}' gen = generate() _disconnect_check_interval = 20 From 996c0338adffa316c18de974b2405e8a10ca67d7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 26 Mar 2026 10:35:14 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a71fd1b8cc..8107e2d101 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1137,7 +1137,7 @@ async def openai_chat_completions( # Pre-compute static JSON envelope for content tokens (hot path). # Must match model_dump_json(exclude_none=True): no finish_reason when null. _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' - _chunk_suffix = '}}]}' + _chunk_suffix = "}}]}" gen = gguf_generate_with_tools() prev_text = "" @@ -1275,7 +1275,7 @@ async def openai_chat_completions( # Pre-compute static JSON envelope for content tokens (hot path). # Must match model_dump_json(exclude_none=True): no finish_reason when null. _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' - _chunk_suffix = '}}]}' + _chunk_suffix = "}}]}" gen = gguf_generate() prev_text = "" @@ -1474,7 +1474,7 @@ async def openai_chat_completions( # Pre-compute static JSON envelope for content tokens (hot path). # Must match model_dump_json(exclude_none=True): no finish_reason when null. _chunk_prefix = f'{{"id":"{completion_id}","object":"chat.completion.chunk","created":{created},"model":{json.dumps(model_name)},"choices":[{{"index":0,"delta":{{"content":' - _chunk_suffix = '}}]}' + _chunk_suffix = "}}]}" gen = generate() _disconnect_check_interval = 20