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(