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

Studio: free the llama-server slot when a chat stream reaches [DONE] (#7564)

Browse source 
* Studio: free the llama-server slot when a chat stream reaches [DONE]

* Release the slot before yielding, only on a completed decode

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten the comments added by this PR

* Inline the done-sentinel check and use plain bools for the decode flags

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-28 18:34:00 -07:00 • committed by GitHub
commit 9e2fc49851
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 622 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

41
studio/backend/routes/inference.py
View file

@ -727,6 +727,7 @@ def _wants_stream_usage(payload) -> bool:
_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0
_SSE_DONE_LINE = "data: [DONE]"
_SSE_DONE_CHUNK = "data: [DONE]\n\n"
def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]:
@ -2440,10 +2441,16 @@ async def _await_cancel_or_disconnect_then_close_client(
return
async def _stop_local_disconnect_cancel_watcher(watcher) -> None:
async def _stop_local_disconnect_cancel_watcher(watcher, timeout_s: float = 5.0) -> None:
# Bounded: this runs in the stream's finally, so awaiting the watcher outright would let a
# wedged poll loop hold the response open forever. asyncio.wait neither cancels nor re-raises,
# and an abandoned watcher owns no resources.
watcher.cancel()
done, _pending = await asyncio.wait({watcher}, timeout = timeout_s)
if not done:
return
try:
await watcher
watcher.result()
except (asyncio.CancelledError, Exception):
pass
@ -9449,12 +9456,15 @@ async def openai_chat_completions(
raise _openai_admission_http_exception(exc, status_code = 429)
_tool_sentinel = object()
# True only once the sync generator returned on its own; see _gguf_decode_finished.
_tool_decode_finished = False
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel.for_payload(cancel_event, payload, *_cancel_keys)
_tracker.__enter__()
async def gguf_tool_stream():
nonlocal _tool_decode_finished
gen = None
next_task = None
stream_completed = False
@ -9542,6 +9552,7 @@ async def openai_chat_completions(
if next_task.done():
next_task = None
if event is _tool_sentinel:
_tool_decode_finished = True
break
# Anything after the gated tool_start means the user answered.
@ -9758,6 +9769,13 @@ async def openai_chat_completions(
stream_started = True
try:
async for chunk in iterator:
# Release before the yield; see gguf_stream_chunks.
if (
lease is not None
and _tool_decode_finished
and chunk == _SSE_DONE_CHUNK
):
lease.release()
yield chunk
except asyncio.CancelledError:
stream_cancelled = True
@ -10060,6 +10078,9 @@ async def openai_chat_completions(
)
_gguf_sentinel = object()
# True only once the sync generator returned on its own: only then has _open_stream's
# client exited. A cancel still emits [DONE] without it.
_gguf_decode_finished = False
if payload.stream:
if _wants_multiple_choices(payload):
@ -10086,6 +10107,7 @@ async def openai_chat_completions(
raise _openai_admission_http_exception(exc, status_code = 429)
async def gguf_stream_chunks():
nonlocal _gguf_decode_finished
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
@ -10130,6 +10152,7 @@ async def openai_chat_completions(
if next_task.done():
next_task = None
if cumulative is _gguf_sentinel:
_gguf_decode_finished = True
break
# Capture server metadata for the final usage chunk
if isinstance(cumulative, dict):
@ -10292,6 +10315,20 @@ async def openai_chat_completions(
stream_started = True
try:
async for chunk in iterator:
# The slot is idle once the sync generator returned and the stream ends
# with the plain sentinel. The finally only runs at ASGI teardown, so
# waiting for it starves the next request. Release before the yield: a
# stalled send() or a consumer that stops pulling parks us there, and
# Starlette never aclose()s a body iterator. Release is idempotent, so
# the finally stays the backstop. Exact equality, not endswith:
# _openai_stream_error_sse ends in the same sentinel before its
# cleanup runs, and that stream still owns the slot.
if (
lease is not None
and _gguf_decode_finished
and chunk == _SSE_DONE_CHUNK
):
lease.release()
yield chunk
except asyncio.CancelledError:
stream_cancelled = True

267
studio/backend/tests/test_gguf_stream_slot_release.py Normal file
View file

@ -0,0 +1,267 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""A finished GGUF chat stream must free its llama-server slot at [DONE].
llama-server has a fixed slot count, gated by an admission lease. Releasing that lease only in
the stream's outer finally, which runs at ASGI teardown, let a wedged teardown pin a slot
llama-server had already freed, so the next chat request queued behind a finished generation
with no timeout to bound the wait.
The wedge below stands in for the real one: the frontend never cancels its reader after [DONE]
(chat-api.ts), and uvicorn advertises ASGI spec_version 2.3, so Starlette's
OSError/ClientDisconnect path, the only disconnect detector _SameTaskStreamingResponse keeps,
cannot fire.
"""
import asyncio
import json
import pytest
from fastapi import FastAPI
from auth.authentication import get_current_subject
from core.inference import llama_admission
import routes.inference as inference_route
@pytest.fixture(autouse = True)
def _fresh_queues():
llama_admission.reset_llama_admission_queues()
yield
llama_admission.reset_llama_admission_queues()
def _active_slots() -> int:
with llama_admission._QUEUES_LOCK:
queues = list(llama_admission._QUEUES.values())
return sum(queue.snapshot().active for queue in queues)
_ONE_SLOT = llama_admission.LlamaAdmissionConfig(max_queue = 4)
def _reserve_one_slot():
"""Take the single slot of a 1-parallel backend. Needs a running loop."""
queue = llama_admission.get_llama_admission_queue("http://llama.test")
reservation = queue.reserve(capacity = 1, config = _ONE_SLOT)
return queue, reservation.lease_nowait()
def test_slot_is_freed_at_done_even_if_teardown_never_finishes():
"""Yield chunks, then wedge in the finally: without the release at [DONE] the slot stays
held for as long as the teardown is stuck, which is what starved the next request in CI.
"""
wedged = asyncio.Event()
async def _stream():
try:
yield 'data: {"choices": [{"delta": {"content": "hi"}}]}\n\n'
yield "data: [DONE]\n\n"
finally:
# Stand-in for a teardown that never completes.
await wedged.wait()
async def _admitted(held):
iterator = _stream()
try:
async for chunk in iterator:
yield chunk
if held is not None and chunk == inference_route._SSE_DONE_CHUNK:
held.release()
finally:
if held is not None:
held.release()
async def _drive():
queue, lease = _reserve_one_slot()
assert lease is not None
assert _active_slots() == 1
seen = []
saw_done = asyncio.Event()
async def _consume():
# Like Starlette's stream_response: it keeps pulling after the last chunk, so the
# generator resumes past [DONE] and only then runs into the wedged teardown.
async for chunk in _admitted(lease):
seen.append(chunk)
if chunk == inference_route._SSE_DONE_CHUNK:
saw_done.set()
task = asyncio.create_task(_consume())
try:
await asyncio.wait_for(saw_done.wait(), timeout = 5.0)
# Give the generator a turn to resume past the [DONE] yield and reach the wedge.
for _ in range(50):
if _active_slots() == 0:
break
await asyncio.sleep(0.01)
assert not task.done(), "teardown should still be wedged"
assert _active_slots() == 0, (
"slot still held after [DONE]; the next chat request would "
"queue behind a generation that already finished"
)
# A second caller must be admitted right away.
second = queue.reserve(capacity = 1, config = _ONE_SLOT).lease_nowait()
assert second is not None, "next request was refused a free slot"
second.release()
finally:
wedged.set()
task.cancel()
await asyncio.gather(task, return_exceptions = True)
return seen
seen = asyncio.run(_drive())
assert seen[-1] == "data: [DONE]\n\n"
def test_release_is_idempotent_so_the_finally_stays_a_backstop():
async def _drive():
_queue, lease = _reserve_one_slot()
assert _active_slots() == 1
lease.release()
lease.release()
assert _active_slots() == 0
asyncio.run(_drive())
def test_stopping_the_disconnect_watcher_cannot_hang():
"""The watcher stop runs in the stream's finally; it must be bounded."""
async def _drive():
started = asyncio.Event()
release = asyncio.Event()
async def _unstoppable():
started.set()
while not release.is_set():
try:
await asyncio.sleep(0.01)
except asyncio.CancelledError:
# Swallow cancellation, as the real watcher does on its way out.
if release.is_set():
raise
continue
watcher = asyncio.create_task(_unstoppable())
await started.wait()
# Would hang forever if the stop awaited the watcher outright.
await asyncio.wait_for(
inference_route._stop_local_disconnect_cancel_watcher(watcher, timeout_s = 0.2),
timeout = 5.0,
)
assert not watcher.done(), "watcher should have been abandoned, not awaited"
release.set()
watcher.cancel()
await asyncio.gather(watcher, return_exceptions = True)
asyncio.run(_drive())
class _OneSlotGgufBackend:
"""A loaded 1-parallel GGUF backend, the shape CI runs."""
is_loaded = True
model_identifier = "test/model.gguf"
base_url = "http://llama.test"
effective_parallel_slots = 1
_is_audio = False
is_vision = False
supports_tools = False
def generate_chat_completion(self, **kwargs):
yield "hi"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
"timings": {"prompt_n": 3, "predicted_n": 1},
"finish_reason": "stop",
}
def test_real_stream_frees_the_slot_at_done_with_a_wedged_teardown(monkeypatch):
"""Drive the real ASGI route, wedged exactly where CI wedged.
Hanging ``_stop_local_disconnect_cancel_watcher``, which runs in ``gguf_stream_chunks``'s
success-path finally, leaves a response that has sent [DONE] but cannot finish.
"""
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _OneSlotGgufBackend())
monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: False)
app = FastAPI()
app.include_router(inference_route.router)
app.dependency_overrides[get_current_subject] = lambda: "test-user"
async def _drive():
wedged = asyncio.Event()
async def _hang(watcher, *args, **kwargs):
watcher.cancel()
await wedged.wait()
monkeypatch.setattr(inference_route, "_stop_local_disconnect_cancel_watcher", _hang)
body = json.dumps(
{"messages": [{"role": "user", "content": "hi"}], "stream": True}
).encode()
scope = {
"type": "http",
"asgi": {"version": "3.0", "spec_version": "2.3"},
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/chat/completions",
"raw_path": b"/chat/completions",
"query_string": b"",
"root_path": "",
"headers": [
(b"host", b"testserver"),
(b"content-type", b"application/json"),
(b"content-length", str(len(body)).encode()),
],
"client": ("127.0.0.1", 12345),
"server": ("testserver", 80),
"app": app,
}
sent_body = asyncio.Event()
frames = []
async def receive():
if not frames:
return {"type": "http.request", "body": body, "more_body": False}
# Never disconnect: the browser keeps the socket open after [DONE].
await asyncio.Event().wait()
async def send(message):
frames.append(message)
if message.get("type") == "http.response.body":
chunk = message.get("body", b"").decode()
if chunk == inference_route._SSE_DONE_CHUNK:
sent_body.set()
task = asyncio.create_task(app(scope, receive, send))
try:
await asyncio.wait_for(sent_body.wait(), timeout = 20.0)
for _ in range(200):
if _active_slots() == 0:
break
await asyncio.sleep(0.01)
assert not task.done(), "response should still be wedged in teardown"
assert _active_slots() == 0, (
"slot still held after [DONE] on the real route; the next chat "
"request would queue behind a finished generation"
)
queue = llama_admission.get_llama_admission_queue("http://llama.test")
second = queue.reserve(capacity = 1, config = _ONE_SLOT).lease_nowait()
assert second is not None, "next request was refused a free slot"
second.release()
finally:
wedged.set()
task.cancel()
await asyncio.gather(task, return_exceptions = True)
asyncio.run(_drive())

316
studio/backend/tests/test_gguf_stream_slot_release_ordering.py Normal file
View file

@ -0,0 +1,316 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Ordering rules for the early admission release at ``data: [DONE]``.
Freeing the llama-server slot at the sentinel is only correct when two things hold, and on a
one-slot backend both are load-bearing:
1. The release happens *before* the sentinel reaches the ASGI ``send()``. Starlette's
``stream_response`` suspends the body iterator at its ``yield`` for the whole of
``await send(...)``, and uvicorn's ``send()`` awaits ``flow.drain()`` on a write-paused
transport, so a client that stops reading parks the generator there indefinitely. Starlette
never ``aclose()``s a body iterator either, so that generator's ``finally`` is left to GC.
2. The sentinel really means "llama-server is done with this request". Two other emitters end
in the same bytes: ``_openai_stream_error_sse``, yielded from inside the still-suspended
generator's ``except`` block, and the cancel path, which breaks the read loop while the sync
generator is still parked on a yield inside ``_open_stream``'s httpx client.
"""
import asyncio
import json
import threading
import pytest
from fastapi import FastAPI
from auth.authentication import get_current_subject
from core.inference import llama_admission
import routes.inference as inference_route
@pytest.fixture(autouse = True)
def _fresh_queues():
llama_admission.reset_llama_admission_queues()
yield
llama_admission.reset_llama_admission_queues()
def _active_slots() -> int:
with llama_admission._QUEUES_LOCK:
queues = list(llama_admission._QUEUES.values())
return sum(queue.snapshot().active for queue in queues)
class _OneSlotBackend:
"""A loaded 1-parallel GGUF backend, the shape CI runs."""
is_loaded = True
model_identifier = "test/model.gguf"
base_url = "http://llama.test"
effective_parallel_slots = 1
_is_audio = False
is_vision = False
supports_tools = False
def __init__(self):
self.closing = threading.Event()
self.finish_close = threading.Event()
self.closed = threading.Event()
self.cancel_event = None
def generate_chat_completion(self, **kwargs):
raise NotImplementedError
class _CompletingBackend(_OneSlotBackend):
def generate_chat_completion(self, **kwargs):
yield "hi"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
"timings": {"prompt_n": 3, "predicted_n": 1},
"finish_reason": "stop",
}
class _FailsMidStreamBackend(_OneSlotBackend):
"""Still decoding when the route's own chunk handling blows up.
``gen`` stays parked on its ``yield`` until the stream's ``finally`` closes it, and only
that close drops the httpx stream llama-server is writing to.
"""
def generate_chat_completion(self, **kwargs):
try:
yield "a"
yield "ab"
yield "abc"
except GeneratorExit:
self.closing.set()
# Stand in for the time llama-server needs to notice the drop and free its slot.
self.finish_close.wait(10.0)
self.closed.set()
raise
class _CancelledMidStreamBackend(_OneSlotBackend):
"""Cancelled by the user halfway through, the Stop-button path."""
def generate_chat_completion(
self,
cancel_event = None,
**kwargs,
):
self.cancel_event = cancel_event
try:
yield "a"
cancel_event.set()
yield "ab"
yield "abc"
except GeneratorExit:
self.closed.set()
raise
def _scope(app, body: bytes) -> dict:
return {
"type": "http",
"asgi": {"version": "3.0", "spec_version": "2.3"},
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/chat/completions",
"raw_path": b"/chat/completions",
"query_string": b"",
"root_path": "",
"headers": [
(b"host", b"testserver"),
(b"content-type", b"application/json"),
(b"content-length", str(len(body)).encode()),
],
"client": ("127.0.0.1", 12345),
"server": ("testserver", 80),
"app": app,
}
def _build_app(monkeypatch, backend):
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: False)
app = FastAPI()
app.include_router(inference_route.router)
app.dependency_overrides[get_current_subject] = lambda: "test-user"
return app
def _request_body() -> bytes:
return json.dumps({"messages": [{"role": "user", "content": "hi"}], "stream": True}).encode()
def test_slot_is_free_before_the_done_frame_reaches_send(monkeypatch):
"""The release must not sit behind ``await send(...)``.
uvicorn's ``send()`` awaits ``flow.drain()`` on a write-paused socket (h11_impl.py), so a
client that stops reading parks the body iterator on its ``yield`` indefinitely. Anything
after that ``yield`` is unreachable, and Starlette never ``aclose()``s the iterator, so the
outer ``finally`` is left to GC.
"""
backend = _CompletingBackend()
app = _build_app(monkeypatch, backend)
async def _drive():
body = _request_body()
frames = []
slots_at_done = []
finished = asyncio.Event()
async def receive():
if not frames:
return {"type": "http.request", "body": body, "more_body": False}
await asyncio.Event().wait()
async def send(message):
frames.append(message)
if message.get("type") != "http.response.body":
return
if message.get("body", b"").decode() == "data: [DONE]\n\n":
# Sampled exactly where a stalled client would wedge.
slots_at_done.append(_active_slots())
finished.set()
task = asyncio.create_task(app(_scope(app, body), receive, send))
try:
await asyncio.wait_for(finished.wait(), timeout = 20.0)
finally:
task.cancel()
await asyncio.gather(task, return_exceptions = True)
assert slots_at_done == [0], (
"the slot was still held while the [DONE] frame was being written; "
"a client that stops reading would pin it there indefinitely"
)
asyncio.run(_drive())
def test_error_sentinel_keeps_the_slot_until_the_generator_is_closed(monkeypatch):
"""``_openai_stream_error_sse`` ends in ``data: [DONE]`` but is not a finish.
It is yielded from inside ``gguf_stream_chunks``'s ``except`` block, so the generator has
not yet run its ``finally``: the worker is undrained and ``gen`` is still open with
llama-server streaming into it. Freeing the slot there puts two callers on a one-slot
backend.
"""
backend = _FailsMidStreamBackend()
app = _build_app(monkeypatch, backend)
calls = {"n": 0}
def _boom(monitor_id, text):
calls["n"] += 1
if calls["n"] >= 2:
raise RuntimeError("chunk handling failed")
monkeypatch.setattr(inference_route.api_monitor, "append_reply", _boom)
async def _drive():
body = _request_body()
frames = []
saw_error = asyncio.Event()
async def receive():
if not frames:
return {"type": "http.request", "body": body, "more_body": False}
await asyncio.Event().wait()
async def send(message):
frames.append(message)
if message.get("type") != "http.response.body":
return
chunk = message.get("body", b"").decode()
# The error form: a payload line plus the sentinel, in one chunk.
if chunk.endswith("data: [DONE]\n\n") and chunk != "data: [DONE]\n\n":
saw_error.set()
task = asyncio.create_task(app(_scope(app, body), receive, send))
try:
await asyncio.wait_for(saw_error.wait(), timeout = 20.0)
# Wait until cleanup reaches gen.close(), so llama-server still holds the slot.
for _ in range(500):
if backend.closing.is_set():
break
await asyncio.sleep(0.01)
assert backend.closing.is_set(), "cleanup never reached gen.close()"
assert _active_slots() == 1, (
"slot handed out while the failed request still owned "
"llama-server; the next request would exceed the configured "
"parallelism"
)
finally:
backend.finish_close.set()
task.cancel()
await asyncio.gather(task, return_exceptions = True)
asyncio.run(_drive())
def test_cancelled_stream_keeps_the_slot_until_the_generator_is_closed(monkeypatch):
"""A cancelled stream emits the plain sentinel with ``gen`` still open.
``cancel_event.is_set()`` breaks the read loop at the top, so the sync generator never
reaches StopIteration and stays parked on a ``yield`` inside ``_open_stream``'s httpx
client. ``stream_completed`` is set all the same, which also makes the ``finally`` skip
``gen.close()``, so ``data: [DONE]`` here does not mean llama-server is finished.
"""
backend = _CancelledMidStreamBackend()
app = _build_app(monkeypatch, backend)
wedged = asyncio.Event()
async def _hang(watcher, *args, **kwargs):
watcher.cancel()
await wedged.wait()
monkeypatch.setattr(inference_route, "_stop_local_disconnect_cancel_watcher", _hang)
async def _drive():
body = _request_body()
frames = []
saw_done = asyncio.Event()
async def receive():
if not frames:
return {"type": "http.request", "body": body, "more_body": False}
await asyncio.Event().wait()
async def send(message):
frames.append(message)
if message.get("type") != "http.response.body":
return
if message.get("body", b"").decode() == "data: [DONE]\n\n":
saw_done.set()
task = asyncio.create_task(app(_scope(app, body), receive, send))
try:
await asyncio.wait_for(saw_done.wait(), timeout = 20.0)
for _ in range(50):
if _active_slots() == 0:
break
await asyncio.sleep(0.01)
assert backend.cancel_event is not None and backend.cancel_event.is_set()
assert (
not backend.closed.is_set()
), "test setup: the generator should still be open here"
assert _active_slots() == 1, (
"slot freed on a cancelled stream whose llama-server request is "
"still open; the next request would exceed the configured "
"parallelism"
)
finally:
wedged.set()
task.cancel()
await asyncio.gather(task, return_exceptions = True)
asyncio.run(_drive())