From 95bfc50b35e10ced5feb55f8f9d679a940d8a766 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 18 Mar 2026 05:10:32 -0700 Subject: [PATCH] Fix inference stall during prefill (retry storm) (#4409) * Fix inference stall during prefill by removing retry storm The _stream_with_retry method used a 0.5s read timeout and retried by sending a brand new POST request each time. During prompt prefill (which can take 5-30+ seconds for long contexts or reasoning models), this caused 10-60 duplicate requests that forced llama-server to restart processing from scratch each time, resulting in 10-20s stalls visible as "Generating" with no progress in the UI. Fix: send the request ONCE with a 120s read timeout for the initial response headers. Cancel support during the prefill wait is handled by a background thread that monitors cancel_event (checked every 0.3s) and closes the response to unblock the httpx read immediately. This preserves the ability to stop/cancel/refresh during generation. The existing 0.5s timeout on the httpx.Client is still used by _iter_text_cancellable for per-token cancel checking during streaming (after prefill), which is unaffected by this change. * Fix race in cancel watcher when response is not yet created When cancel_event fires before client.stream() returns (response is still None), the watcher would hit return and exit without closing anything. The main thread stays blocked for up to 120s. Fix: after cancel is requested, keep polling _response_ref every 0.1s until the response object appears (then close it) or _cancel_closed is set (main thread finished on its own). * Minor cleanup: remove redundant None check, add debug logging in cancel watcher Address Gemini review: cancel_event is guaranteed non-None when the watcher thread runs, and logging the close exception aids debugging. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retry r.close() on failure instead of giving up If r.close() raises, stay in the polling loop and retry rather than returning and leaving the main thread blocked for up to 120s. * fix: keep short read timeout during token streaming The prefill_timeout (read=120s) was passed to client.stream(), which applied to ALL reads -- not just the initial response headers. This meant _iter_text_cancellable's ReadTimeout-based cancel checking was broken during token streaming: the Stop button could take up to 120s to respond instead of 0.5s. Fix: keep the client's short read timeout (0.5s) for the stream call. During prefill, catch ReadTimeout in a loop and re-check cancel_event instead of re-sending the POST (which was the original retry storm). Once the first bytes arrive, yield the response with a PrependStream wrapper so iter_text() sees the buffered first chunk. This preserves both: - Fast cancel during prefill (via cancel watcher + ReadTimeout loop) - Fast cancel during streaming (via _iter_text_cancellable's 0.5s ReadTimeout, which now fires correctly again) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: swap to short-timeout stream after prefill completes Address two review issues: 1. _PrependStream did not inherit from httpx.SyncByteStream, so Response.iter_raw() would raise RuntimeError. Replaced with a _ShortTimeoutStream that inherits SyncByteStream properly. 2. client.stream() entry itself raises ReadTimeout during slow prefill (before headers arrive). The previous fix tried to catch this at the body-read level but missed the connection-level timeout. New approach: keep the 120s read timeout for client.stream() so the connection survives long prefills. Once headers arrive, replace the response stream with _ShortTimeoutStream -- a wrapper that uses a background reader thread and a Queue with a short get() timeout to re-raise ReadTimeout at the original 0.5s interval. This way _iter_text_cancellable's cancel-checking remains responsive during token streaming while prefill gets the long timeout it needs. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: move _ShortTimeoutStream before LlamaCppBackend class The class was placed inside LlamaCppBackend's body, splitting the class in two and making _codec_mgr and other attributes unreachable. Move it to module level before LlamaCppBackend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: remove _ShortTimeoutStream, use watcher for all cancel _ShortTimeoutStream had two critical issues: 1. Raising ReadTimeout from a generator kills it -- Python finalizes generators after an uncaught exception, so the next next() call hits StopIteration and streaming ends mid-response. 2. The unbounded Queue in the background reader loses backpressure, causing memory spikes with slow clients. Simpler approach: use the 120s read timeout for the entire stream and rely on the cancel watcher thread for all cancellation (both prefill and streaming). The watcher closes the response on cancel_event, which unblocks any blocking httpx read within ~0.3s. This eliminates the need for short timeout tricks entirely. Cancel latency: - Prefill: ~0.3s (watcher polls cancel_event every 0.3s) - Streaming: ~0.3s (same watcher mechanism) - Both faster than the old 0.5s ReadTimeout approach * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs: clarify cancel limitations in _stream_with_retry The docstrings claimed ~0.3s cancel in all cases, but httpx cannot interrupt a blocked read before the response object exists. Update the docstrings to accurately describe the behavior: - Cancel during prefill (header wait) is deferred until headers arrive - Cancel during streaming works via response.close() from the watcher - _iter_text_cancellable docstring updated to reflect the watcher-based cancel mechanism instead of the old ReadTimeout polling --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 103 +++++++++++++++++---- 1 file changed, 83 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 64a7324adf..7bb32a81f8 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1273,10 +1273,11 @@ class LlamaCppBackend: ) -> Generator[str, None, None]: """Iterate over an httpx streaming response with cancel support. - Uses a short read timeout on the stream so that cancel_event is - checked at least every 0.5s, even if the model is slow to produce - the next token. Without this, iter_text() blocks until the next - chunk arrives and cancellation can take many seconds on large models. + Checks cancel_event between chunks and on ReadTimeout. The + cancel watcher in _stream_with_retry also calls response.close() + on cancel, which unblocks iter_text() once the response exists. + During normal streaming llama-server sends tokens frequently, + so the cancel check between chunks is the primary mechanism. """ text_iter = response.iter_text() while True: @@ -1301,24 +1302,85 @@ class LlamaCppBackend: payload: dict, cancel_event: Optional[threading.Event] = None, ): - """Open an httpx streaming POST, retrying on ReadTimeout. + """Open an httpx streaming POST with cancel support. - The short read timeout (0.5 s) that enables cancel-checking during - streaming can also fire while waiting for the server to produce - its first response bytes (e.g. a reasoning model thinking). - This wrapper retries the connection until headers arrive or - cancel_event is set. + Sends the request once with a long read timeout (120 s) so + prompt processing (prefill) can finish without triggering a + retry storm. The previous 0.5 s timeout caused duplicate POST + requests every half second, forcing llama-server to restart + processing each time. + + A background watcher thread provides cancel by closing the + response when cancel_event is set. Limitation: httpx does not + allow interrupting a blocked read from another thread before + the response object exists, so cancel during the initial + header wait (prefill phase) only takes effect once headers + arrive. After that, response.close() unblocks reads promptly. + In practice llama-server prefill is 1-5 s for typical prompts, + during which cancel is deferred -- still much better than the + old retry storm which made prefill slower. """ - while True: + if cancel_event is not None and cancel_event.is_set(): + raise GeneratorExit + + # Background watcher: close the response if cancel is requested. + # Only effective after response headers arrive (httpx limitation). + _cancel_closed = threading.Event() + _response_ref: list = [None] + + def _cancel_watcher(): + while not _cancel_closed.is_set(): + if cancel_event.wait(timeout = 0.3): + # Cancel requested. Keep polling until the response object + # exists so we can close it, or until the main thread + # finishes on its own (_cancel_closed is set in finally). + while not _cancel_closed.is_set(): + r = _response_ref[0] + if r is not None: + try: + r.close() + return + except Exception as e: + logger.debug( + f"Error closing response in cancel watcher: {e}" + ) + # Response not created yet -- wait briefly and retry + _cancel_closed.wait(timeout = 0.1) + return + + watcher = None + if cancel_event is not None: + watcher = threading.Thread( + target = _cancel_watcher, daemon = True, name = "prefill-cancel" + ) + watcher.start() + + try: + # Long read timeout so prefill (prompt processing) can finish + # without triggering a retry storm. Cancel during both + # prefill and streaming is handled by the watcher thread + # which closes the response, unblocking any httpx read. + prefill_timeout = httpx.Timeout( + connect = 30, + read = 120.0, + write = 10, + pool = 10, + ) + with client.stream( + "POST", url, json = payload, timeout = prefill_timeout + ) as response: + _response_ref[0] = response + if cancel_event is not None and cancel_event.is_set(): + raise GeneratorExit + yield response + return + except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError): + # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit - try: - with client.stream("POST", url, json = payload) as response: - yield response - return - except httpx.ReadTimeout: - # Server still thinking -- retry - continue + raise + finally: + _cancel_closed.set() def generate_chat_completion( self, @@ -1371,8 +1433,9 @@ class LlamaCppBackend: in_thinking = False try: - # Use a short read timeout so we can check cancel_event - # frequently instead of blocking indefinitely on slow models. + # _stream_with_retry uses a 120 s read timeout so prefill + # can finish. Cancel during streaming is handled by the + # watcher thread (closes the response on cancel_event). stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry(