From fef37cb25b35e36b12c1d9c7d020aa0fb402ce89 Mon Sep 17 00:00:00 2001 From: Apoze <158856608+Apoze@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:05:48 +0200 Subject: [PATCH] Studio: queue local GGUF OpenAI-compatible requests before llama-server (#7047) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../backend/core/inference/llama_admission.py | 368 +++++ studio/backend/core/inference/llama_cpp.py | 23 + studio/backend/routes/inference.py | 1257 ++++++++++++++- studio/backend/tests/test_llama_admission.py | 320 ++++ ...test_llama_cpp_effective_parallel_slots.py | 53 + .../tests/test_openai_tool_passthrough.py | 1396 +++++++++++++++++ .../test_stream_cancel_registration_timing.py | 24 +- 7 files changed, 3408 insertions(+), 33 deletions(-) create mode 100644 studio/backend/core/inference/llama_admission.py create mode 100644 studio/backend/tests/test_llama_admission.py create mode 100644 studio/backend/tests/test_llama_cpp_effective_parallel_slots.py diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py new file mode 100644 index 0000000000..b6a939c87b --- /dev/null +++ b/studio/backend/core/inference/llama_admission.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission control for local llama-server generation requests. + +The helpers in this module deliberately know nothing about FastAPI, SSE, or the +OpenAI-compatible route shape. They only coordinate how many upstream generation +requests may be active for one llama-server backend and provide a cancellable +FIFO queue for excess requests. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections import deque +from dataclasses import dataclass +from typing import Deque, Optional + + +ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" + +DEFAULT_ADMISSION_ENABLED = True +DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None +DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 +DEFAULT_ADMISSION_MAX_QUEUE = 64 + + +@dataclass(frozen = True) +class LlamaAdmissionConfig: + enabled: bool = DEFAULT_ADMISSION_ENABLED + queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + + +@dataclass(frozen = True) +class LlamaAdmissionSnapshot: + key: str + capacity: int + active: int + queued: int + + +class LlamaAdmissionError(Exception): + def __init__( + self, + message: str, + *, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + super().__init__(message) + self.snapshot = snapshot + + +class LlamaAdmissionQueueFull(LlamaAdmissionError): + pass + + +class LlamaAdmissionTimeout(LlamaAdmissionError): + pass + + +class LlamaAdmissionCancelled(LlamaAdmissionError): + pass + + +def _bool_env(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + value = value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def _positive_float_env(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else default + + +def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = int(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def llama_admission_config_from_env() -> LlamaAdmissionConfig: + return LlamaAdmissionConfig( + enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), + queue_timeout_s = _optional_positive_float_env( + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + ), + keepalive_interval_s = _positive_float_env( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + ), + max_queue = _optional_positive_int_env( + ADMISSION_MAX_QUEUE_ENV, + DEFAULT_ADMISSION_MAX_QUEUE, + ), + ) + + +@dataclass +class _Waiter: + loop: asyncio.AbstractEventLoop + future: asyncio.Future + cancelled: bool = False + granted_lease: Optional["LlamaAdmissionLease"] = None + + +class LlamaAdmissionLease: + def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + self._queue = queue + self._released = False + self._release_lock = threading.Lock() + + def release(self) -> None: + queue = None + with self._release_lock: + if self._released: + return + self._released = True + queue = self._queue + if queue is not None: + queue.release() + + async def __aenter__(self) -> "LlamaAdmissionLease": + return self + + async def __aexit__(self, *_args) -> None: + self.release() + + +class LlamaAdmissionReservation: + def __init__( + self, + *, + queue: Optional["LlamaAdmissionQueue"], + lease: Optional[LlamaAdmissionLease] = None, + waiter: Optional[_Waiter] = None, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + self._queue = queue + self._lease = lease + self._waiter = waiter + self.snapshot = snapshot + + @property + def is_cancelled(self) -> bool: + return self._lease is None and self._waiter is None + + def lease_nowait(self) -> Optional[LlamaAdmissionLease]: + if self._lease is not None: + return self._lease + if self._waiter is None or not self._waiter.future.done(): + return None + if self._waiter.future.cancelled(): + self._waiter.cancelled = True + self._waiter = None + return None + self._lease = self._waiter.future.result() + self._waiter = None + return self._lease + + async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + lease = self.lease_nowait() + if lease is not None: + return lease + if self._waiter is None: + return None + waiter = self._waiter + try: + await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s) + except asyncio.CancelledError: + if waiter.future.cancelled(): + waiter.cancelled = True + if self._waiter is waiter: + self._waiter = None + return None + raise + return self.lease_nowait() + + def cancel(self) -> None: + lease = self.lease_nowait() + if lease is not None: + lease.release() + self._lease = None + return + if self._queue is not None and self._waiter is not None: + self._queue.cancel(self._waiter) + self._waiter = None + + def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]: + if self._queue is None: + return self.snapshot + return self._queue.snapshot() + + +class LlamaAdmissionQueue: + def __init__(self, key: str): + self.key = key + self._lock = threading.Lock() + self._active = 0 + self._capacity = 1 + self._waiters: Deque[_Waiter] = deque() + + def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + capacity = max(1, int(capacity or 1)) + if not config.enabled: + return LlamaAdmissionReservation( + queue = None, + lease = LlamaAdmissionLease(None), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + ) + + loop = asyncio.get_running_loop() + with self._lock: + self._capacity = capacity + self._prune_waiters_locked() + self._grant_waiters_locked() + if self._active < self._capacity and not self._waiters: + self._active += 1 + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self), + snapshot = self._snapshot_locked(), + ) + if config.max_queue is not None and len(self._waiters) >= config.max_queue: + raise LlamaAdmissionQueueFull( + "llama-server generation queue is full", + snapshot = self._snapshot_locked(), + ) + waiter = _Waiter( + loop = loop, + future = loop.create_future(), + ) + self._waiters.append(waiter) + return LlamaAdmissionReservation( + queue = self, + waiter = waiter, + snapshot = self._snapshot_locked(), + ) + + def release(self) -> None: + with self._lock: + if self._active > 0: + self._active -= 1 + self._grant_waiters_locked() + + def cancel(self, waiter: _Waiter) -> None: + lease_to_release = None + with self._lock: + waiter.cancelled = True + try: + self._waiters.remove(waiter) + except ValueError: + pass + if waiter.granted_lease is not None: + lease_to_release = waiter.granted_lease + waiter.granted_lease = None + if not waiter.future.done(): + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + if lease_to_release is not None: + lease_to_release.release() + + def snapshot(self) -> LlamaAdmissionSnapshot: + with self._lock: + self._prune_waiters_locked() + return self._snapshot_locked() + + def is_idle(self) -> bool: + with self._lock: + self._prune_waiters_locked() + return self._active == 0 and not self._waiters + + def _grant_waiters_locked(self) -> None: + self._prune_waiters_locked() + while self._waiters and self._active < self._capacity: + waiter = self._waiters.popleft() + if waiter.cancelled or waiter.future.done(): + continue + self._active += 1 + lease = LlamaAdmissionLease(self) + waiter.granted_lease = lease + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + + def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + if waiter.cancelled or waiter.future.done(): + waiter.granted_lease = None + if not waiter.future.done(): + waiter.future.cancel() + lease.release() + return + try: + waiter.future.set_result(lease) + waiter.granted_lease = None + except asyncio.InvalidStateError: + waiter.granted_lease = None + lease.release() + + def _prune_waiters_locked(self) -> None: + self._waiters = deque( + waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + ) + + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: + return LlamaAdmissionSnapshot( + key = self.key, + capacity = self._capacity, + active = self._active, + queued = len(self._waiters), + ) + + +_QUEUES_LOCK = threading.Lock() +_QUEUES: dict[str, LlamaAdmissionQueue] = {} + + +def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: + with _QUEUES_LOCK: + queue = _QUEUES.get(key) + if queue is None: + queue = LlamaAdmissionQueue(key) + _QUEUES[key] = queue + # base_url carries a fresh ephemeral port on every model load, so + # each load registers a new key. Drop the now-idle queues from prior + # loads so the registry can't grow without bound on a long-running + # server. Queues with in-flight requests are kept until they drain. + for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]: + del _QUEUES[stale_key] + return queue + + +def reset_llama_admission_queues() -> None: + with _QUEUES_LOCK: + _QUEUES.clear() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c2e3adf815..b06c6eb5cb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1541,6 +1541,7 @@ class LlamaCppBackend: self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None + self._effective_parallel_slots: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1726,6 +1727,15 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def effective_parallel_slots(self) -> int: + """Return the serving-slot count the active llama-server actually uses.""" + try: + slots = int(getattr(self, "_effective_parallel_slots", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1742,6 +1752,16 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def _commit_effective_parallel_slots(self, n_parallel: int) -> None: + try: + slots = int(n_parallel) + except (TypeError, ValueError): + slots = 1 + self._effective_parallel_slots = max(1, slots) + + def _reset_effective_parallel_slots(self) -> None: + self._effective_parallel_slots = 1 + @staticmethod def _read_rss_bytes(pid: int) -> Optional[int]: """Resident set size of ``pid`` in bytes, from /proc//status (Linux). @@ -7187,6 +7207,7 @@ class LlamaCppBackend: ) self._healthy = True + self._commit_effective_parallel_slots(n_parallel) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -7724,6 +7745,7 @@ class LlamaCppBackend: self._context_length = None self._effective_context_length = None self._max_context_length = None + self._reset_effective_parallel_slots() self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -7779,6 +7801,7 @@ class LlamaCppBackend: # Stop the watchdog before a deliberate kill so a planned reload/unload # isn't seen as a crash; a real crash never routes through here. self._stop_mtp_crash_watchdog() + self._reset_effective_parallel_slots() if self._process is None: return try: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 826669d922..db350240ac 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -28,6 +28,16 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) def _positive_int_or_none(value: Any) -> Optional[int]: @@ -1072,6 +1082,270 @@ _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 _OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() def _openai_compat_stream_stall_timeout(): @@ -6649,6 +6923,24 @@ async def openai_chat_completions( bypass_permissions = bool(payload.bypass_permissions), ) + _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + _tool_sentinel = object() _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) @@ -6657,6 +6949,8 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + next_task = None + stream_completed = False disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -6694,7 +6988,14 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") return - event = await asyncio.to_thread(next, gen, _tool_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _tool_sentinel) + ) + try: + event = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if event is _tool_sentinel: break @@ -6793,6 +7094,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6807,18 +7109,155 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - if gen is not None: - try: - gen.close() - except (RuntimeError, ValueError): - pass - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF tool stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) if payload.stream: + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_tool_stream(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_tool_stream() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_tool_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_tool_stream(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_tool_stream(), + unstarted_cleanup = _gguf_tool_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6864,10 +7303,61 @@ async def openai_chat_completions( except (RuntimeError, ValueError): pass + drain_task = None + + async def _drain_cancelled_gguf_tool_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass + + admission_lease = None + admission_wait_started_at = None try: - full_text, completion_usage, completion_finish = await asyncio.to_thread( - _drain_gguf_tool_loop + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) + full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) reasoning_text, visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _responses_should_parse_think_markers( @@ -6913,6 +7403,48 @@ async def openai_chat_completions( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_tool_task() + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) except Exception as e: logger.error(f"Error during GGUF tool completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6933,6 +7465,8 @@ async def openai_chat_completions( ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) finally: + if admission_lease is not None: + admission_lease.release() _tracker.__exit__(None, None, None) # ── Standard GGUF path (no tools) ───────────────────── @@ -6967,6 +7501,23 @@ async def openai_chat_completions( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) async def gguf_stream_chunks(): disconnect_watcher = asyncio.create_task( @@ -7114,9 +7665,122 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_stream_chunks(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_stream_chunks() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_stream_chunks(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_stream_chunks(), + unstarted_cleanup = _gguf_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -7125,6 +7789,101 @@ async def openai_chat_completions( }, ) else: + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. @@ -7268,6 +8027,10 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -9399,6 +10162,52 @@ async def _responses_stream( ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + def _responses_admission_failed_sse(exc: Exception, *, status_code: int) -> str: + return ( + "event: response.failed\n" + "data: " + + json.dumps( + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": _llama_public_model_id(llama_backend, payload.model) + or payload.model, + "output": [], + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + "error": { + "code": status_code, + "message": str(exc), + }, + }, + } + ) + + "\n\n" + ) async def event_generator(): # Clean public id for every response envelope. Prefer the loaded model's @@ -10208,14 +11017,111 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) + async def admitted_event_generator(): + lease = reservation.lease_nowait() + admission_wait_started_at = None + stream_started = False + stream_cancelled = False + iterator = None + try: + if lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "debug", + ) + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = None, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = None, + ) + iterator = event_generator() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _responses_admission_failed_sse(exc, status_code = 503) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + async def _responses_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + return _SameTaskStreamingResponse( - event_generator(), + admitted_event_generator(), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = _responses_admission_unstarted_cleanup, ) @@ -12059,7 +12965,202 @@ async def _openai_passthrough_stream( completion_id, monitor_id: Optional[str] = None, ): - """Streaming client-side pass-through for /v1/chat/completions. + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionCancelled as exc: + lease.release() + _tracker.__exit__(None, None, None) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + return await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = lease, + tracker = _tracker, + ) + + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "debug", + ) + + async def _queued_stream(): + admitted_started = False + admitted_body_owns_cleanup = False + admitted_response = None + admitted_body_cancelled = False + try: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + admitted_response = await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = wait_item, + tracker = _tracker, + ) + admitted_started = True + iterator = admitted_response.body_iterator + admitted_body_owns_cleanup = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + admitted_body_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = admitted_body_cancelled, + ) + if not admitted_body_owns_cleanup: + cleanup = getattr(admitted_response, "_unstarted_cleanup", None) + if cleanup is not None: + await cleanup() + return + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if not admitted_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _queued_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _queued_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _queued_unstarted_cleanup, + ) + + +async def _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, + *, + admission_lease: LlamaAdmissionLease, + tracker, +): + """Streaming client-side pass-through after Studio granted an upstream slot. Forwards the client's OpenAI function-calling request to llama-server and relays the SSE stream back with minimal normalization (reasoning-only @@ -12073,9 +13174,7 @@ async def _openai_passthrough_stream( --reasoning-format auto``), so ``delta.content`` carries no raw markup and is deliberately not re-parsed locally, unlike the ``/completion`` paths. """ - _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) - _tracker = _TrackedCancel(cancel_event, *_cancel_keys) - _tracker.__enter__() + _tracker = tracker target_url = f"{llama_backend.base_url}/v1/chat/completions" upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) @@ -12166,7 +13265,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", @@ -12712,7 +13814,10 @@ async def _openai_passthrough_stream( client = client, ) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) async def _unstarted_cleanup() -> None: # Client disconnected before the body stream started, so _stream()'s @@ -12723,7 +13828,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( _stream(), @@ -12747,7 +13855,10 @@ async def _openai_passthrough_stream( await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) finally: - _tracker.__exit__(None, None, None) + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) raise @@ -12759,6 +13870,106 @@ async def _openai_passthrough_non_streaming( *, request: Optional[Request] = None, cancel_event = None, +): + """Non-streaming pass-through guarded by local llama-server admission.""" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_nonstream", + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + level = "debug", + ) + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + return await _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + raise + finally: + if lease is not None: + lease.release() + + +async def _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py new file mode 100644 index 0000000000..2f04e81926 --- /dev/null +++ b/studio/backend/tests/test_llama_admission.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import threading + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_admission +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + DEFAULT_ADMISSION_MAX_QUEUE, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + LlamaAdmissionConfig, + LlamaAdmissionQueueFull, + get_llama_admission_queue, + llama_admission_config_from_env, + reset_llama_admission_queues, +) + + +@pytest.fixture(autouse = True) +def _reset_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_admission_config_defaults(monkeypatch): + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ): + monkeypatch.delenv(name, raising = False) + + config = llama_admission_config_from_env() + + assert config.enabled is True + assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE + + +def test_admission_config_env_overrides(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + + config = llama_admission_config_from_env() + + assert config.enabled is False + assert config.queue_timeout_s is None + assert config.keepalive_interval_s == 0.25 + assert config.max_queue is None + + +def test_admission_config_positive_queue_timeout_env(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") + + config = llama_admission_config_from_env() + + assert config.queue_timeout_s == 600.0 + + +def test_fifo_capacity_one_grants_next_waiter_on_release(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + third = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert third.lease_nowait() is None + assert queue.snapshot().queued == 2 + + first_lease.release() + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + second_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + third_lease.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_queue_full_rejects_excess_waiter(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 1) + + first = queue.reserve(capacity = 1, config = config) + queued = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert queued.lease_nowait() is None + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_active_slot_limit(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(enabled = False) + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert second.lease_nowait() is not None + assert queue.snapshot().active == 0 + assert queue.snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + await asyncio.sleep(0) + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_before_delivery_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_external_waiter_future_cancel_invalidates_reservation(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + second._waiter.future.cancel() + + assert second.lease_nowait() is None + assert second.is_cancelled is True + assert await second.wait(0.01) is None + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + wait_task = asyncio.create_task(second.wait(1.0)) + await asyncio.sleep(0) + second._waiter.future.cancel() + + assert await asyncio.wait_for(wait_task, timeout = 0.1) is None + assert second.is_cancelled is True + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_capacity_increase_promotes_existing_waiter_fifo(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert queue.snapshot().active == 1 + assert queue.snapshot().queued == 1 + + third = queue.reserve(capacity = 2, config = config) + + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + snapshot = queue.snapshot() + assert snapshot.capacity == 2 + assert snapshot.active == 2 + assert snapshot.queued == 1 + + first_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + + second_lease.release() + third_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_lease_release_is_idempotent_under_concurrent_calls(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + reservation = queue.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + threads = [threading.Thread(target = lease.release) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_new_key_evicts_idle_prior_load_queues(): + # Each model load carries a fresh ephemeral port, so a new base_url key must + # not leave the drained queues from earlier loads accumulating forever. + get_llama_admission_queue("http://127.0.0.1:1001") + get_llama_admission_queue("http://127.0.0.1:1002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} + + get_llama_admission_queue("http://127.0.0.1:1003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} + + +def test_new_key_retains_in_flight_prior_load_queue(): + config = LlamaAdmissionConfig() + busy = get_llama_admission_queue("http://127.0.0.1:2001") + + async def _run(): + reservation = busy.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + # A new load must not drop a queue that still has an in-flight request. + get_llama_admission_queue("http://127.0.0.1:2002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + + # Once it drains, the next load reclaims it. + lease.release() + get_llama_admission_queue("http://127.0.0.1:2003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py new file mode 100644 index 0000000000..5525bc3ea9 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import os +import sys + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) + monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + return LlamaCppBackend() + + +def test_effective_parallel_slots_initial_value_is_one(backend): + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): + backend._commit_effective_parallel_slots(3) + + assert backend.effective_parallel_slots == 3 + + +@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): + backend._commit_effective_parallel_slots(value) + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_reset_returns_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend._reset_effective_parallel_slots() + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_unload_resets_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend.unload_model() + + assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 10cbd2aa01..910818d7d8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -30,6 +30,15 @@ from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) from routes.inference import ( _aclose_stream_resources, _build_chat_request, @@ -50,13 +59,17 @@ from routes.inference import ( _monitor_openai_sse_event, _normalize_openai_passthrough_sse_line, _openai_compat_stream_stall_timeout, + _openai_llama_admission_capacity, _openai_messages_for_gguf_chat, _openai_passthrough_sse_line_terminal_state, _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _responses_stream, _openai_stream_error_sse, _openai_stream_usage_chunk, + _openai_admission_wait_stream_chunks, + _wait_for_openai_admission_non_streaming, _proxy_to_external_provider, _SameTaskStreamingResponse, _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, @@ -68,6 +81,13 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy, set_tool_policy +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): class Closeable: def __init__(self, *, cancel = False): @@ -1338,6 +1358,80 @@ class TestOpenAICompatibilityHelpers: assert headers["Authorization"] == "Bearer secret" assert headers["Connection"] == "close" + def test_openai_admission_capacity_prefers_backend_effective_slots(self): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + ) + backend = SimpleNamespace(effective_parallel_slots = 3) + + assert _openai_llama_admission_capacity(request, backend) == 3 + + @pytest.mark.parametrize("backend_value", [None, 0, -1, "not-an-int"]) + def test_openai_admission_capacity_falls_back_to_app_state(self, backend_value): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 2)) + ) + backend = SimpleNamespace(effective_parallel_slots = backend_value) + + assert _openai_llama_admission_capacity(request, backend) == 2 + + def test_openai_admission_capacity_falls_back_to_one_without_request(self): + assert _openai_llama_admission_capacity(None, SimpleNamespace()) == 1 + + def test_openai_admission_non_streaming_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for( + _wait_for_openai_admission_non_streaming( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ), + timeout = 0.1, + ) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_admission_stream_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.stream.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + chunks = _openai_admission_wait_stream_chunks( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ) + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for(chunks.__anext__(), timeout = 0.1) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) assert _openai_compat_stream_stall_timeout() == 120.0 @@ -1993,6 +2087,335 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-stream-close-cleanup" + + def _generate(**_kwargs): + yield "visible" + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + cancel_id = cancel_id, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert cancel_id in inf_mod._CANCEL_REGISTRY + await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + aclose = getattr(iterator, "aclose", None) + assert aclose is not None + await aclose() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + + asyncio.run(_run()) + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): import routes.inference as inf_mod @@ -2434,6 +2857,313 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_cancel_id_stops_queued_request_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after cancel_id") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "standard-nonstream-admission-cancel" + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + try: + for _ in range(50): + if cancel_id in inf_mod._CANCEL_REGISTRY: + break + await asyncio.sleep(0.01) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(task, timeout = 0.5) + assert exc.value.status_code == 499 + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_admission_task_cancel_cleans_tracker_and_slot( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-nonstream-task-cancel" + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after task cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod @@ -4086,6 +4816,270 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_immediate_task_cancel_releases_admission_and_tracker( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + async def fake_cancel_check(*_args, **_kwargs): + raise asyncio.CancelledError() + + cancel_id = "passthrough-stream-immediate-task-cancel" + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_raise_if_openai_admission_cancelled", + fake_cancel_check, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_stream( + self._Request(), + threading.Event(), + backend, + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_before_inner_first_chunk_runs_cleanup( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + body_holder = {} + cleanup_called = threading.Event() + + async def fake_admitted(*_args, admission_lease, tracker, **_kwargs): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + cleanup_called.set() + + class BlockingBody: + def __init__(self): + self.started = threading.Event() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started.set() + await asyncio.sleep(3600) + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + await cleanup() + + body = BlockingBody() + body_holder["body"] = body + return _SameTaskStreamingResponse( + body, + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-unstarted-cleanup" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + assert cancel_id in inf_mod._CANCEL_REGISTRY + + blocker.release() + pending = asyncio.create_task(iterator.__anext__()) + for _ in range(100): + if "body" in body_holder: + break + await asyncio.sleep(0.01) + body = body_holder["body"] + assert await asyncio.to_thread(body.started.wait, 1.0) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert body_holder["body"].closed + assert cleanup_called.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_after_inner_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_admitted( + *_args, + monitor_id = None, + admission_lease, + tracker, + **_kwargs, + ): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + + async def body(): + try: + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + await asyncio.sleep(3600) + except asyncio.CancelledError: + inf_mod.api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await cleanup() + + return _SameTaskStreamingResponse( + body(), + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-cancel-monitor" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + + blocker.release() + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "hello" in first + + pending = asyncio.create_task(iterator.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_stream_synthesizes_missing_finish_reason(self, monkeypatch): async def _run(): result = await self._run_passthrough_stream( @@ -4189,6 +5183,292 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_admitted(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start when admission queue is full") + + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve( + capacity = 1, + config = LlamaAdmissionConfig(max_queue = 1), + ).lease_nowait() + queued = queue.reserve(capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)) + assert blocker is not None + assert queued.lease_nowait() is None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 429 + finally: + queued.cancel() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after client cancellation") + + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + cancel_event = threading.Event() + cancel_event.set() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = cancel_event, + ) + + assert exc.value.status_code == 499 + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after admission task cancel") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = threading.Event(), + ) + + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -5381,6 +6661,15 @@ class TestApiMonitorAudioInput: class TestResponsesChatTemplateKwargs: _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + class _Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/responses") + method = "POST" + + async def is_disconnected(self): + return False + def test_enable_thinking_lifted_from_extra_body(self): payload = ResponsesRequest( model = "qwen-local", @@ -5413,6 +6702,113 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None + def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + + queue = get_llama_admission_queue("http://llama.responses.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start after created cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "event: response.created" in first + + with pytest.raises(asyncio.CancelledError): + await iterator.athrow(asyncio.CancelledError()) + + assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + # ===================================================================== # GGUF chat-template role alternation: coalesce orphaned user turns left diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 7cfe9bf32c..a833006873 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -17,7 +17,7 @@ from pathlib import Path SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py" -SRC = SOURCE_PATH.read_text() +SRC = SOURCE_PATH.read_text(encoding = "utf-8") _TREE = ast.parse(SRC) @@ -166,16 +166,20 @@ def test_chat_completions_streams_avoid_starlette_task_group(): def test_openai_passthrough_stream_avoids_starlette_task_group(): - top = _async_function("_openai_passthrough_stream") + functions = [ + _async_function("_openai_passthrough_stream"), + _async_function("_openai_passthrough_stream_admitted"), + ] legacy_calls = [] same_task_calls = 0 - for sub in ast.walk(top): - if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): - continue - if sub.func.id == "StreamingResponse": - legacy_calls.append(sub.lineno) - if sub.func.id == "_SameTaskStreamingResponse": - same_task_calls += 1 + for fn in functions: + for sub in ast.walk(fn): + if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): + continue + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 assert not legacy_calls, ( "OpenAI passthrough streams must use _SameTaskStreamingResponse, " "not Starlette's legacy task-group StreamingResponse. Lines: " @@ -197,7 +201,7 @@ def test_direct_llama_server_streams_install_disconnect_watcher(): "openai_completions", "_responses_stream", "_anthropic_passthrough_stream", - "_openai_passthrough_stream", + "_openai_passthrough_stream_admitted", } missing = [ name