diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 12547277f5..d0a2d97f74 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 diff --git a/studio/backend/tests/test_gguf_stream_slot_release.py b/studio/backend/tests/test_gguf_stream_slot_release.py new file mode 100644 index 0000000000..4390f364c8 --- /dev/null +++ b/studio/backend/tests/test_gguf_stream_slot_release.py @@ -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()) diff --git a/studio/backend/tests/test_gguf_stream_slot_release_ordering.py b/studio/backend/tests/test_gguf_stream_slot_release_ordering.py new file mode 100644 index 0000000000..7a8ceb4f53 --- /dev/null +++ b/studio/backend/tests/test_gguf_stream_slot_release_ordering.py @@ -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())