ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix inference stall during prefill (retry storm) (#4409)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-03-18 05:10:32 -07:00 • committed by GitHub
commit 95bfc50b35
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 84 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

103
studio/backend/core/inference/llama_cpp.py
View file

@ -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(