diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py index b6a939c87b..1a9ae04b0e 100644 --- a/studio/backend/core/inference/llama_admission.py +++ b/studio/backend/core/inference/llama_admission.py @@ -13,37 +13,85 @@ from __future__ import annotations import asyncio import os +import sys 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" +# dataclass(slots = True) halves per-instance overhead. Measured as perf-neutral +# here, not a speed win: it costs a little on construction and gains it back on +# access. It is 3.10+ and this package declares >=3.9, so gate it rather than +# dropping it outright. Empty on 3.9 means a plain dataclass. +_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} + + +ADMISSION_CONTROL_ENV = "UNSLOTH_LLAMA_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_LLAMA_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_LLAMA_ADMISSION_MAX_QUEUE" +ADMISSION_QUEUE_PER_SLOT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_PER_SLOT" + +# The UNSLOTH_OPENAI_COMPAT_* spellings predate this queue being shared with the +# Anthropic /v1/messages route (same llama-server slots). Still honored; the +# neutral name above wins when both are set. +_LEGACY_ENV = { + 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 +# None: a queued request waits for its slot indefinitely rather than timing out. DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 -DEFAULT_ADMISSION_MAX_QUEUE = 64 +# None: no absolute cap, the wait line is sized from the pool instead. +DEFAULT_ADMISSION_MAX_QUEUE = None +# Wait line = 16 x the serving slots, so it tracks --parallel (4 slots -> 64 +# waiters, 8 -> 128). Purely a memory guard; waiting itself is never timed out. +DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16 +# Floor for the scaled line, so a 1-slot backend (plain `unsloth studio`, or any +# load downshifted to fit VRAM) keeps the depth it had before scaling existed +# rather than dropping to 16 and rejecting callers that used to queue. +DEFAULT_ADMISSION_MIN_QUEUE = 64 -@dataclass(frozen = True) +@dataclass(frozen = True, **_SLOTS) 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 + queue_per_slot: Optional[int] = DEFAULT_ADMISSION_QUEUE_PER_SLOT + # Unconditional floor on the scaled line. The env path clears it when the + # operator sets QUEUE_PER_SLOT, so only the default multiplier is floored. + min_queue: Optional[int] = DEFAULT_ADMISSION_MIN_QUEUE + + def queue_limit(self, capacity: int) -> Optional[int]: + """How many callers may line up for a pool of ``capacity`` slots. + + An explicit ``max_queue`` wins; otherwise the line scales with the slots + so it follows ``--parallel``. The default multiplier is floored, so a + 1-slot backend does not end up shallower than it was before scaling. None + (or any non-positive setting) means an unbounded line. + """ + if self.max_queue is not None: + return self.max_queue if self.max_queue > 0 else None + if not self.queue_per_slot or self.queue_per_slot <= 0: + return None + scaled = self.queue_per_slot * max(1, capacity) + return max(self.min_queue, scaled) if self.min_queue else scaled -@dataclass(frozen = True) +@dataclass(frozen = True, **_SLOTS) class LlamaAdmissionSnapshot: key: str capacity: int active: int queued: int + free: int = 0 class LlamaAdmissionError(Exception): @@ -69,8 +117,17 @@ class LlamaAdmissionCancelled(LlamaAdmissionError): pass -def _bool_env(name: str, default: bool) -> bool: +def _raw_env(name: str) -> Optional[str]: + """Value for a canonical name, falling back to its legacy spelling.""" value = os.environ.get(name) + if value is None or not value.strip(): + legacy = _LEGACY_ENV.get(name) + value = os.environ.get(legacy) if legacy else None + return value + + +def _bool_env(name: str, default: bool) -> bool: + value = _raw_env(name) if value is None or not value.strip(): return default value = value.strip().lower() @@ -82,7 +139,7 @@ def _bool_env(name: str, default: bool) -> bool: def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: - value = os.environ.get(name) + value = _raw_env(name) if value is None or not value.strip(): return default try: @@ -93,7 +150,7 @@ def _optional_positive_float_env(name: str, default: Optional[float]) -> Optiona def _positive_float_env(name: str, default: float) -> float: - value = os.environ.get(name) + value = _raw_env(name) if value is None or not value.strip(): return default try: @@ -103,19 +160,38 @@ def _positive_float_env(name: str, default: float) -> float: 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 +def _queue_limits_from_env() -> tuple[Optional[int], Optional[int], Optional[int]]: + """(max_queue, queue_per_slot, min_queue) from the environment. + + An absolute MAX_QUEUE wins outright; MAX_QUEUE=0 asks for an unbounded line. + Unset leaves the per-slot multiplier in charge (itself 0 for unbounded). The + floor applies only to the default multiplier: setting QUEUE_PER_SLOT means + the operator wants that exact depth, however shallow. + """ + # Explicit means it parsed, not just that something was set: a typo falls back + # to the default multiplier, so it has to keep the default's floor too. + raw_per_slot = _raw_env(ADMISSION_QUEUE_PER_SLOT_ENV) try: - parsed = int(value.strip()) + per_slot = int((raw_per_slot or "").strip()) except ValueError: - return default - return parsed if parsed > 0 else None + per_slot, min_queue = DEFAULT_ADMISSION_QUEUE_PER_SLOT, DEFAULT_ADMISSION_MIN_QUEUE + else: + per_slot, min_queue = (per_slot if per_slot > 0 else None), None + raw = _raw_env(ADMISSION_MAX_QUEUE_ENV) + if raw is None or not raw.strip(): + return None, per_slot, min_queue + try: + parsed = int(raw.strip()) + except ValueError: + return None, per_slot, min_queue + return (parsed, None, None) if parsed > 0 else (None, None, None) def llama_admission_config_from_env() -> LlamaAdmissionConfig: + max_queue, queue_per_slot, min_queue = _queue_limits_from_env() return LlamaAdmissionConfig( + queue_per_slot = queue_per_slot, + min_queue = min_queue, enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), queue_timeout_s = _optional_positive_float_env( ADMISSION_QUEUE_TIMEOUT_ENV, @@ -125,14 +201,11 @@ def llama_admission_config_from_env() -> LlamaAdmissionConfig: ADMISSION_KEEPALIVE_INTERVAL_ENV, DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, ), - max_queue = _optional_positive_int_env( - ADMISSION_MAX_QUEUE_ENV, - DEFAULT_ADMISSION_MAX_QUEUE, - ), + max_queue = max_queue, ) -@dataclass +@dataclass(**_SLOTS) class _Waiter: loop: asyncio.AbstractEventLoop future: asyncio.Future @@ -141,11 +214,23 @@ class _Waiter: class LlamaAdmissionLease: - def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + __slots__ = ("_queue", "_slot", "_released", "_release_lock") + + def __init__( + self, + queue: Optional["LlamaAdmissionQueue"], + slot: Optional[int] = None, + ): self._queue = queue + self._slot = slot self._released = False self._release_lock = threading.Lock() + @property + def slot(self) -> Optional[int]: + """Pool slot this lease holds, or None when admission is disabled.""" + return self._slot + def release(self) -> None: queue = None with self._release_lock: @@ -154,7 +239,7 @@ class LlamaAdmissionLease: self._released = True queue = self._queue if queue is not None: - queue.release() + queue.release(self._slot) async def __aenter__(self) -> "LlamaAdmissionLease": return self @@ -164,6 +249,8 @@ class LlamaAdmissionLease: class LlamaAdmissionReservation: + __slots__ = ("_queue", "_lease", "_waiter", "snapshot") + def __init__( self, *, @@ -195,6 +282,13 @@ class LlamaAdmissionReservation: return self._lease async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + """Wait up to ``timeout_s`` for a slot. + + A timeout leaves this reservation queued so the caller can poll again. + Any exit that abandons the wait for good must call ``cancel()``, or the + slot granted later is delivered to a future nobody reads and is never + released. + """ lease = self.lease_nowait() if lease is not None: return lease @@ -229,35 +323,80 @@ class LlamaAdmissionReservation: class LlamaAdmissionQueue: + """A fixed pool of generation slots for one llama-server, plus a FIFO wait line. + + The pool mirrors llama-server's own ``--parallel`` slots: ``capacity`` slot ids + are each either free or held by exactly one caller. A caller that finds every + slot busy waits in arrival order and is handed the next slot to free, so no + caller is starved. This bounds only the callers that reserve: chat completions + and messages do, while /v1/completions, Studio's own chat endpoint and RAG + captioning all reach llama-server directly, so it is not a global cap. + Waiting is unbounded in time by default (``queue_timeout_s`` + None); the wait line itself is bounded, and only how many may line up before + new arrivals are rejected. By default that is ``16 x slots`` floored at 64, + not unlimited: an unbounded line takes ``max_queue`` or ``queue_per_slot`` + set to 0. See ``LlamaAdmissionConfig.queue_limit``. + """ + + __slots__ = ("key", "_lock", "_capacity", "_free", "_in_use", "_held", "_waiters") + def __init__(self, key: str): self.key = key self._lock = threading.Lock() - self._active = 0 self._capacity = 1 + self._free: list[int] = [0] + # Held slots as a bitmask: one int instead of a set, so the pool costs the + # same whether it is idle or saturated. _held is its popcount, kept as a + # counter because int.bit_count() is 3.10+ and this package targets 3.9. + self._in_use = 0 + self._held = 0 self._waiters: Deque[_Waiter] = deque() + def _resize_pool_locked(self, capacity: int) -> None: + # Slots past a shrunk capacity retire when their holder releases them. + if capacity == self._capacity: + return + self._capacity = capacity + self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1] + + def _can_admit_locked(self) -> bool: + # Slots still held above a shrunk capacity keep occupying the backend, so + # count every held slot against the ceiling, not just the ids below it. + return bool(self._free) and self._held < self._capacity + + def _take_slot_locked(self) -> Optional[int]: + if not self._can_admit_locked(): + return None + slot = self._free.pop() + self._in_use |= 1 << slot + self._held += 1 + return slot + 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), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0, capacity), ) loop = asyncio.get_running_loop() with self._lock: - self._capacity = capacity - self._prune_waiters_locked() + self._resize_pool_locked(capacity) 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: + if not self._waiters: + slot = self._take_slot_locked() + if slot is not None: + # No snapshot here: callers read it through snapshot_now(), + # which re-reads the queue, so building one per admitted + # request would be pure allocation on the hot path. + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self, slot), + ) + limit = config.queue_limit(self._capacity) + if limit is not None and self._live_waiters_locked() >= limit: raise LlamaAdmissionQueueFull( "llama-server generation queue is full", snapshot = self._snapshot_locked(), @@ -270,13 +409,20 @@ class LlamaAdmissionQueue: return LlamaAdmissionReservation( queue = self, waiter = waiter, - snapshot = self._snapshot_locked(), ) - def release(self) -> None: + def _release_slot_locked(self, slot: Optional[int]) -> None: + # A slot id at or past a shrunk capacity retires instead of returning. + if slot is None or not self._in_use >> slot & 1: + return + self._in_use &= ~(1 << slot) + self._held -= 1 + if slot < self._capacity: + self._free.append(slot) + + def release(self, slot: Optional[int]) -> None: with self._lock: - if self._active > 0: - self._active -= 1 + self._release_slot_locked(slot) self._grant_waiters_locked() def cancel(self, waiter: _Waiter) -> None: @@ -291,7 +437,13 @@ class LlamaAdmissionQueue: lease_to_release = waiter.granted_lease waiter.granted_lease = None if not waiter.future.done(): - waiter.loop.call_soon_threadsafe(waiter.future.cancel) + try: + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + except RuntimeError: + # Loop gone. Routes call cancel() from finally blocks, so + # raising here would both mask their exception and skip the + # release below, stranding the slot for the process lifetime. + pass if lease_to_release is not None: lease_to_release.release() @@ -303,20 +455,30 @@ class LlamaAdmissionQueue: def is_idle(self) -> bool: with self._lock: self._prune_waiters_locked() - return self._active == 0 and not self._waiters + return self._in_use == 0 and not self._waiters def _grant_waiters_locked(self) -> None: - self._prune_waiters_locked() - while self._waiters and self._active < self._capacity: + # Dead waiters are skipped as they are popped, so no prune is needed here. + while self._waiters and self._can_admit_locked(): waiter = self._waiters.popleft() if waiter.cancelled or waiter.future.done(): continue - self._active += 1 - lease = LlamaAdmissionLease(self) + slot = self._take_slot_locked() + lease = LlamaAdmissionLease(self, slot) waiter.granted_lease = lease - waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + try: + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + except RuntimeError: + # Waiter's loop is gone. Reclaim the slot; leaving the bit set + # would strand it, since _free is rebuilt from the bitmask. + waiter.granted_lease = None + self._release_slot_locked(slot) def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + # Runs on the waiter's own loop thread, which is also the only thread that + # cancels that reservation, so waiter state is safe to touch unlocked here. + # release() may be called from any thread, but only reaches this via + # call_soon_threadsafe. Cancelling off-loop would need this under _lock. if waiter.cancelled or waiter.future.done(): waiter.granted_lease = None if not waiter.future.done(): @@ -331,16 +493,32 @@ class LlamaAdmissionQueue: lease.release() def _prune_waiters_locked(self) -> None: + # Rebuilding the deque on every reserve/release dominated the hot path, so + # only pay it when a waiter actually died out of band (an externally + # cancelled future); cancel() already drops its own waiter eagerly. + for waiter in self._waiters: + if waiter.cancelled or waiter.future.done(): + break + else: + return self._waiters = deque( waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() ) + def _live_waiters_locked(self) -> int: + self._prune_waiters_locked() + return len(self._waiters) + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: return LlamaAdmissionSnapshot( key = self.key, capacity = self._capacity, - active = self._active, + active = self._held, queued = len(self._waiters), + # What another caller could actually take, so the admission log never + # shows free slots next to queued requests: after a shrink, ids below + # the new capacity can be free while holdovers still fill the ceiling. + free = min(len(self._free), max(0, self._capacity - self._held)), ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 06911fd866..cf95e743bf 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -387,7 +387,7 @@ def _raise_unsupported_n(path_label: str) -> None: _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") -def _sse_streaming_response(content) -> StreamingResponse: +def _sse_streaming_response(content, *, unstarted_cleanup = None) -> StreamingResponse: """A ``text/event-stream`` response with the standard SSE headers used by every streaming path here: no client/proxy caching, no proxy buffering, and a one-shot connection. Two callers build their response inline instead: the @@ -409,6 +409,7 @@ def _sse_streaming_response(content) -> StreamingResponse: "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = unstarted_cleanup, ) @@ -1141,7 +1142,7 @@ def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: return None -def _openai_admission_log( +def _llama_admission_log( event: str, reservation: Optional[LlamaAdmissionReservation] = None, *, @@ -1159,13 +1160,15 @@ def _openai_admission_log( 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", + "llama admission %s: mode=%s path=%s completion_id=%s " + "pool=%s/%s free=%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, "capacity", None), + getattr(snapshot, "free", None), getattr(snapshot, "queued", None), wait_ms, ) @@ -1189,6 +1192,23 @@ def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTT ) +def _anthropic_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + """Anthropic-shaped error for an admission reject/timeout/cancel (429/503/499).""" + 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})" + ) + # Types come from ANTHROPIC_TYPE_BY_STATUS (429 -> rate_limit_error, which is + # what Anthropic SDKs back off on); overloaded_error is reserved for 529. + return HTTPException( + status_code = status_code, + detail = anthropic_error_body(message, status = status_code), + ) + + def _openai_admission_timeout_error( reservation: LlamaAdmissionReservation, ) -> LlamaAdmissionTimeout: @@ -1494,6 +1514,24 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +async def _release_unstarted_anthropic_stream(iterator, prior_cleanup) -> None: + """Close a stream whose body never started, running the response's own + pre-start hook. aclose() on an unstarted async generator is a no-op, so its + finally never runs and anything the builder acquired eagerly (the passthrough + cancel tracker) would leak without the hook.""" + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + try: + await aclose() + except Exception: + pass + if prior_cleanup is not None: + try: + await prior_cleanup() + except Exception: + pass + + def _tracked_cancel_unstarted_cleanup(tracker): """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when the generator's finally (which normally exits it) never runs.""" @@ -8286,7 +8324,7 @@ async def openai_chat_completions( llama_backend = llama_backend, ) except LlamaAdmissionQueueFull as exc: - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -8524,7 +8562,7 @@ async def openai_chat_completions( admission_wait_started_at = None if stream_lease is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -8549,7 +8587,7 @@ async def openai_chat_completions( yield wait_item continue lease = wait_item - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -8580,7 +8618,7 @@ async def openai_chat_completions( cancelled = stream_cancelled, ) except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -8594,7 +8632,7 @@ async def openai_chat_completions( _openai_admission_error_body(exc, status_code = 503) ) except LlamaAdmissionCancelled: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -8705,7 +8743,7 @@ async def openai_chat_completions( try: if reservation.lease_nowait() is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -8720,7 +8758,7 @@ async def openai_chat_completions( cancel_event = cancel_event, ) if admission_wait_started_at is not None: - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -8791,7 +8829,7 @@ async def openai_chat_completions( _tracker.__exit__(None, None, None) raise except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -8806,7 +8844,7 @@ async def openai_chat_completions( _tracker.__exit__(None, None, None) raise _openai_admission_http_exception(exc, status_code = 503) except LlamaAdmissionCancelled as exc: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -8886,7 +8924,7 @@ async def openai_chat_completions( ) except LlamaAdmissionQueueFull as exc: _tracker.__exit__(None, None, None) - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -9058,7 +9096,7 @@ async def openai_chat_completions( admission_wait_started_at = None if stream_lease is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -9083,7 +9121,7 @@ async def openai_chat_completions( yield wait_item continue lease = wait_item - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -9114,7 +9152,7 @@ async def openai_chat_completions( cancelled = stream_cancelled, ) except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -9128,7 +9166,7 @@ async def openai_chat_completions( _openai_admission_error_body(exc, status_code = 503) ) except LlamaAdmissionCancelled: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -9184,7 +9222,7 @@ async def openai_chat_completions( llama_backend = llama_backend, ) except LlamaAdmissionQueueFull as exc: - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -9203,7 +9241,7 @@ async def openai_chat_completions( try: if reservation.lease_nowait() is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -9218,7 +9256,7 @@ async def openai_chat_completions( cancel_event = cancel_event, ) if admission_wait_started_at is not None: - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -9240,7 +9278,7 @@ async def openai_chat_completions( _tracker.__exit__(None, None, None) raise except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -9255,7 +9293,7 @@ async def openai_chat_completions( _tracker.__exit__(None, None, None) raise _openai_admission_http_exception(exc, status_code = 503) except LlamaAdmissionCancelled as exc: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -11713,7 +11751,7 @@ async def _responses_stream( llama_backend = llama_backend, ) except LlamaAdmissionQueueFull as exc: - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -12571,7 +12609,7 @@ async def _responses_stream( try: if lease is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -12589,7 +12627,7 @@ async def _responses_stream( yield wait_item continue lease = wait_item - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -12621,7 +12659,7 @@ async def _responses_stream( cancelled = stream_cancelled, ) except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -12633,7 +12671,7 @@ async def _responses_stream( api_monitor.fail(monitor_id, str(exc)) yield _responses_admission_failed_sse(exc, status_code = 503) except LlamaAdmissionCancelled: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -13233,12 +13271,206 @@ async def anthropic_messages( cancel_event, ) + # ── Admission control ───────────────────────────────────── + # Bound concurrent llama-server generations to the backend's serving slots via a + # FIFO queue keyed by base_url (shared with /v1/chat/completions, same slots). + # Excess requests queue; a streaming waiter gets SSE keep-alives, the queue 429s + # once full. Mirrors the OpenAI passthrough admission wiring. Streaming takes the + # slot when the response is built and drops it when the body finishes or is + # abandoned; the non-stream path holds it across the single awaited generation. + _anthropic_admission_mode = "anthropic_stream" if payload.stream else "anthropic_nonstream" + + async def _admitted_anthropic_stream( + orig_body, + reservation, + admission_config, + stream_lease, + prior_cleanup = None, + ): + lease = stream_lease + stream_cancelled = False + body_started = False + wait_started_at = None + try: + if lease is None: + wait_started_at = time.monotonic() + _llama_admission_log( + "queued", + reservation, + request = request, + mode = _anthropic_admission_mode, + ) + 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 + break + _llama_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _anthropic_admission_mode, + wait_started_at = wait_started_at, + ) + if lease is None: + return + body_started = True + async for chunk in orig_body: + yield chunk + except asyncio.CancelledError: + # Must reach the monitored generator as CancelledError, not aclose's + # GeneratorExit, or its handler never finalizes the monitor entry. + stream_cancelled = True + raise + except LlamaAdmissionTimeout as exc: + api_monitor.fail(monitor_id, str(exc)) + _llama_admission_log( + "timeout", + reservation, + request = request, + mode = _anthropic_admission_mode, + wait_started_at = wait_started_at, + level = "warning", + ) + yield build_anthropic_sse_event( + "error", + anthropic_error_body(str(exc), status = 503), + ) + except LlamaAdmissionCancelled: + _llama_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _anthropic_admission_mode, + wait_started_at = wait_started_at, + ) + return + finally: + # Closing can raise (a raw body re-raises CancelledError after + # teardown), and a slot lost that way never comes back: with no queue + # timeout the pool just shrinks and later callers wait forever. Keep + # the release in its own finally, as the /responses wiring does. + try: + if body_started: + await _close_openai_admitted_stream_iterator( + orig_body, + cancelled = stream_cancelled, + ) + else: + # Gave up while queued: the monitored body never ran, so nothing + # downstream finalizes the entry or exits the response's tracker. + api_monitor.finish(monitor_id, "cancelled") + await _release_unstarted_anthropic_stream(orig_body, prior_cleanup) + finally: + if lease is not None: + lease.release() + else: + reservation.cancel() + + async def _admitted_anthropic(coro): + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, llama_backend = llama_backend + ) + except LlamaAdmissionQueueFull as exc: + coro.close() + api_monitor.fail(monitor_id, str(exc)) + _llama_admission_log( + "queue-full", + snapshot = getattr(exc, "snapshot", None), + request = request, + mode = _anthropic_admission_mode, + level = "warning", + ) + raise _anthropic_admission_http_exception(exc, status_code = 429) + except BaseException: + # Reserving never awaited the generation, so close it rather than + # leave an un-awaited coroutine behind. + coro.close() + raise + + if payload.stream: + stream_lease = reservation.lease_nowait() + # Set up the stream (token count + tracker enter) and surface a pre-response + # cancel now, exactly as the un-admitted path did; the upstream generation is + # deferred to body iteration, so the slot is only held while tokens flow. + try: + # Token counting calls llama-server, so a dead backend raises here + # with the slot already taken. cancel() covers both cases: it + # releases the lease if one was granted, else drops the waiter. + monitored = await _monitored_anthropic(coro) + except BaseException: + reservation.cancel() + raise + orig_body = getattr(monitored, "body_iterator", None) + if orig_body is None: + reservation.cancel() + return monitored + + # Replacing body_iterator would strand the response's own pre-start + # hook (the passthrough uses one to exit its cancel tracker), so chain + # to it instead of clobbering it. + prior_cleanup = getattr(monitored, "_unstarted_cleanup", None) + + async def _unstarted_cleanup() -> None: + # The body never ran, so nothing else closes out the monitor entry. + api_monitor.finish(monitor_id, "cancelled") + try: + await _release_unstarted_anthropic_stream(orig_body, prior_cleanup) + finally: + # A BaseException here is swallowed upstream, so releasing + # outside the finally would shrink the pool silently. + reservation.cancel() + + monitored.body_iterator = _admitted_anthropic_stream( + orig_body, reservation, admission_config, stream_lease, prior_cleanup + ) + monitored._unstarted_cleanup = _unstarted_cleanup + return monitored + + lease = None + try: + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + monitored = await _monitored_anthropic(coro) + return monitored + except LlamaAdmissionTimeout as exc: + coro.close() + api_monitor.fail(monitor_id, str(exc)) + raise _anthropic_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + coro.close() + api_monitor.finish(monitor_id, "cancelled") + raise _anthropic_admission_http_exception(exc, status_code = 499) + except BaseException: + # Cancelled while queued (shutdown, outer task cancel): the generation + # coroutine was never awaited, so close it rather than leak it. + if lease is None: + coro.close() + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + else: + reservation.cancel() + # ── Client-side pass-through path ───────────────────────── if client_tools: openai_tools = openai_client_tools if payload.stream: - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_passthrough_stream( request, cancel_event, @@ -13262,7 +13494,7 @@ async def anthropic_messages( auto_heal_tool_calls = payload.auto_heal_tool_calls, ) ) - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_passthrough_non_streaming( llama_backend, openai_messages, @@ -13367,7 +13599,7 @@ async def anthropic_messages( ) if payload.stream: - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_tool_stream( request, cancel_event, @@ -13380,7 +13612,7 @@ async def anthropic_messages( disable_parallel_tool_use = _disable_parallel, ) ) - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_tool_non_streaming( _run_tool_gen, message_id, @@ -13407,7 +13639,7 @@ async def anthropic_messages( ) if payload.stream: - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_plain_stream( request, cancel_event, @@ -13418,7 +13650,7 @@ async def anthropic_messages( openai_messages = openai_messages, ) ) - return await _monitored_anthropic( + return await _admitted_anthropic( _anthropic_plain_non_streaming( _run_plain_gen, message_id, @@ -13930,6 +14162,28 @@ def _build_passthrough_payload( return body +async def _anthropic_passthrough_retry_url(llama_backend, exc): + """Fresh upstream URL after respawning a dead llama-server, else None. + + A crashed server relaunches on a NEW ephemeral port, so a passthrough still + holding the old base_url keeps failing until the next load. Mirrors the + respawn-and-retry in generate_chat_completion. None when an MTP+tensor crash + already scheduled its own recovery, or when nothing needed respawning. + """ + recover = getattr(llama_backend, "_maybe_recover_from_mtp_crash", None) + if recover is not None and recover(exc): + return None + # Only the first caller gets True above; the rest must not respawn the same + # MTP config underneath the fallback that is already reloading without it. + if getattr(llama_backend, "_mtp_runtime_fallback_in_progress", False): + return None + respawn = getattr(llama_backend, "_respawn_if_dead", None) + if respawn is None or not await asyncio.to_thread(respawn): + return None + logger.warning("llama-server was unreachable; respawned it and retrying the passthrough") + return f"{llama_backend.base_url}/v1/chat/completions" + + async def _anthropic_passthrough_stream( request, cancel_event, @@ -13997,8 +14251,15 @@ async def _anthropic_passthrough_stream( openai_tools, disable_parallel_tool_use = disable_parallel_tool_use, ) - for line in emitter.start(message_id, model_name, input_tokens = input_tokens): - yield line + # These yields sit outside the teardown try below, so a disconnect while + # the opening lines are being sent would strand the tracker. __exit__ is + # idempotent, so the normal path still exits once, down there. + try: + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): + yield line + except BaseException: + _tracker.__exit__(None, None, None) + raise # Manage the httpx client, response, AND the aiter_lines() async # generator MANUALLY -- no `async with`, no anonymous iterator. @@ -14033,13 +14294,24 @@ async def _anthropic_passthrough_stream( cancel_watcher = None disconnect_watcher = None try: - req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} - ) - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request - ) + url = target_url + try: + req = client.build_request("POST", url, json = body, headers = {"Connection": "close"}) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + except httpx.ConnectError as exc: + # Nothing has streamed yet, so a respawned server can be retried once + # on its new port without duplicating output. + url = await _anthropic_passthrough_retry_url(llama_backend, exc) + if url is None: + raise + req = client.build_request("POST", url, json = body, headers = {"Connection": "close"}) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) if resp is None: return @@ -14118,7 +14390,13 @@ async def _anthropic_passthrough_stream( for line in emitter.finish(): yield line - return _sse_streaming_response(_stream()) + # The tracker is entered eagerly above, but _stream()'s finally is what exits + # it. Closing an async generator that never started is a no-op, so hand the + # response a cleanup hook or a pre-start give-up leaks the registry entry. + return _sse_streaming_response( + _stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + ) async def _anthropic_passthrough_non_streaming( @@ -14158,11 +14436,24 @@ async def _anthropic_passthrough_non_streaming( backend_ctx = llama_backend.context_length, ) - resp = await nonstreaming_client().post( - target_url, - json = body, - timeout = _llama_non_streaming_generation_timeout(), - ) + try: + resp = await nonstreaming_client().post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) + except httpx.ConnectError as exc: + # Nothing was returned yet, so retry once against the respawned server's + # new port; the nudge retry below then reuses the same fresh URL. + retry_url = await _anthropic_passthrough_retry_url(llama_backend, exc) + if retry_url is None: + raise + target_url = retry_url + resp = await nonstreaming_client().post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise HTTPException( @@ -14667,7 +14958,7 @@ async def _openai_passthrough_stream( ) except LlamaAdmissionQueueFull as exc: _tracker.__exit__(None, None, None) - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -14712,7 +15003,7 @@ async def _openai_passthrough_stream( ) admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -14736,7 +15027,7 @@ async def _openai_passthrough_stream( if isinstance(wait_item, str): yield wait_item continue - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -14781,7 +15072,7 @@ async def _openai_passthrough_stream( await cleanup() return except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -14793,7 +15084,7 @@ async def _openai_passthrough_stream( 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( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, @@ -15570,7 +15861,7 @@ async def _openai_passthrough_non_streaming( llama_backend = llama_backend, ) except LlamaAdmissionQueueFull as exc: - _openai_admission_log( + _llama_admission_log( "queue-full", snapshot = exc.snapshot, request = request, @@ -15585,7 +15876,7 @@ async def _openai_passthrough_non_streaming( try: if reservation.lease_nowait() is None: admission_wait_started_at = time.monotonic() - _openai_admission_log( + _llama_admission_log( "queued", reservation, request = request, @@ -15599,7 +15890,7 @@ async def _openai_passthrough_non_streaming( cancel_event = cancel_event, ) if admission_wait_started_at is not None: - _openai_admission_log( + _llama_admission_log( "granted-after-wait", reservation, request = request, @@ -15621,7 +15912,7 @@ async def _openai_passthrough_non_streaming( cancel_event = cancel_event, ) except LlamaAdmissionTimeout as exc: - _openai_admission_log( + _llama_admission_log( "timeout", reservation, request = request, @@ -15632,7 +15923,7 @@ async def _openai_passthrough_non_streaming( api_monitor.fail(monitor_id, str(exc)) raise _openai_admission_http_exception(exc, status_code = 503) except LlamaAdmissionCancelled as exc: - _openai_admission_log( + _llama_admission_log( "cancelled-before-upstream", reservation, request = request, diff --git a/studio/backend/tests/test_anthropic_admission.py b/studio/backend/tests/test_anthropic_admission.py new file mode 100644 index 0000000000..de01accd08 --- /dev/null +++ b/studio/backend/tests/test_anthropic_admission.py @@ -0,0 +1,973 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission-control wiring for the Anthropic /v1/messages endpoint. + +The FIFO queue itself is unit-tested in test_llama_admission.py; here we exercise +how anthropic_messages reserves a slot, queues when the backend is saturated, +streams keep-alives while waiting, releases on completion, and maps rejects to +429/503. Slot occupancy is driven directly through the shared queue (keyed by the +backend base_url) so generation stays fast and no thread has to block. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import gc +import os +import re +import sys +import threading +import time +import warnings +from types import SimpleNamespace + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import routes.inference as inf_mod +from routes.inference import ( + _anthropic_passthrough_retry_url, + _anthropic_passthrough_stream, + anthropic_messages, +) +from models.inference import AnthropicMessagesRequest +from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_PER_SLOT_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) +from fastapi import HTTPException + +_KEY = "http://llama.admission.test:9999" + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + reset_llama_admission_queues() + monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 64)) + monkeypatch.setattr(inf_mod, "_CANCEL_REGISTRY", {}) + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_PER_SLOT_ENV, + # Legacy spellings resolve too, so clear both for isolation. + "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", + ): + monkeypatch.delenv(name, raising = False) + yield + reset_llama_admission_queues() + + +class _Request: + def __init__(self, disconnected = False): + self.state = SimpleNamespace() + self.url = SimpleNamespace(path = "/v1/messages") + self.method = "POST" + self._disconnected = disconnected + + async def is_disconnected(self): + return self._disconnected + + +def _install_backend( + monkeypatch, + *, + slots = 1, + base_url = _KEY, + count_tokens = None, +): + def _gen_plain(**_kwargs): + yield "ok" + + def _gen_tools(**_kwargs): + yield {"type": "content", "text": "ok"} + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_tool_passthrough = False, + model_identifier = "test-model", + context_length = 2048, + count_chat_tokens = count_tokens or (lambda *a, **k: 2), + generate_chat_completion = _gen_plain, + generate_chat_completion_with_tools = _gen_tools, + effective_parallel_slots = slots, + base_url = base_url, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + return backend + + +def _payload(**fields) -> AnthropicMessagesRequest: + base = {"max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]} + base.update(fields) + return AnthropicMessagesRequest(**base) + + +def _record_admission_logs(monkeypatch): + """Capture _llama_admission_log output. + + Through the logger rather than caplog: this one is a structlog bound logger, + so it never reaches the stdlib handlers caplog installs. + """ + records = [] + + def _record(level): + return lambda fmt, *args: records.append((level, fmt % args)) + + monkeypatch.setattr( + inf_mod, + "logger", + SimpleNamespace( + debug = _record("debug"), + info = _record("info"), + warning = _record("warning"), + ), + ) + return records + + +def _snapshot(key = _KEY): + return get_llama_admission_queue(key).snapshot() + + +def _occupy(key, capacity, n): + """Hold ``n`` slots on the queue so the next reserve must wait; returns leases.""" + leases = [] + for _ in range(n): + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, config = LlamaAdmissionConfig() + ) + lease = reservation.lease_nowait() + assert lease is not None + leases.append(lease) + return leases + + +async def _consume(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk) + return "".join(chunks) + + +# ── Non-streaming ───────────────────────────────────────────── + + +def test_non_streaming_completes_and_releases_slot(monkeypatch): + _install_backend(monkeypatch, slots = 2) + + async def _run(): + response = await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + assert response.status_code == 200 + snap = _snapshot() + assert snap.active == 0 and snap.queued == 0 + + asyncio.run(_run()) + + +def test_non_streaming_queue_full_returns_429(monkeypatch): + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # slot busy + # One waiter fills the max_queue=1; the next reserve rejects. + get_llama_admission_queue(_KEY).reserve( + capacity = 1, config = LlamaAdmissionConfig(max_queue = 1) + ) + with pytest.raises(HTTPException) as exc: + await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + assert exc.value.status_code == 429 + # rate_limit_error is what Anthropic SDKs back off on; overloaded_error is 529. + # The type string alone does not pin the envelope, since OpenAI's 429 uses the + # same word. Assert the shape too, or emitting an OpenAI body still passes. + detail = exc.value.detail + assert detail["type"] == "error" + assert "request_id" in detail + assert set(detail["error"]) == {"type", "message"} + assert detail["error"]["type"] == "rate_limit_error" + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_admission_events_are_logged_on_the_anthropic_surface(monkeypatch): + # The OpenAI passthrough logs these with a mode; without the same on /v1/messages + # an operator debugging a slow Anthropic client has nothing to look at, and the + # pool is shared, so it is the same triage. + records = _record_admission_logs(monkeypatch) + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + get_llama_admission_queue(_KEY).reserve( + capacity = 1, config = LlamaAdmissionConfig(max_queue = 1) + ) + with pytest.raises(HTTPException): + await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + for lease in held: + lease.release() + + asyncio.run(_run()) + full = [msg for _level, msg in records if "queue-full" in msg] + assert full, records + assert "llama admission queue-full" in full[0] + assert "mode=anthropic_nonstream" in full[0] + + +def test_streaming_admission_waiting_is_logged(monkeypatch): + # queued and granted-after-wait were both emitted with nothing asserting them. + records = _record_admission_logs(monkeypatch) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + task = asyncio.create_task(_consume(response)) + await asyncio.sleep(0.15) + for lease in held: + lease.release() + await asyncio.wait_for(task, timeout = 5) + + asyncio.run(_run()) + events = [msg for _level, msg in records if "llama admission" in msg] + # "llama admission queued", not "queued": every line carries a queued=N field, + # so the bare substring matches any admission log at all. + assert any( + "llama admission queued" in m and "mode=anthropic_stream" in m for m in events + ), events + granted = [m for m in events if "granted-after-wait" in m] + assert granted, events + # wait_ms is the point of the event: a grant that reports nothing is useless. + assert re.search(r"wait_ms=\d+", granted[0]), granted + + +def test_streaming_admission_timeout_is_logged(monkeypatch): + records = _record_admission_logs(monkeypatch) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # never released, so the waiter times out + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + await _consume(response) + for lease in held: + lease.release() + + asyncio.run(_run()) + timeouts = [msg for level, msg in records if "timeout" in msg and level == "warning"] + assert timeouts, records + assert "mode=anthropic_stream" in timeouts[0] + + +def test_streaming_give_up_while_queued_is_logged(monkeypatch): + # cancelled-before-upstream is the one that tells an operator a client walked + # away rather than the backend being slow. + records = _record_admission_logs(monkeypatch) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), + request = _Request(disconnected = True), + current_subject = "t", + ) + await _consume(response) + for lease in held: + lease.release() + + asyncio.run(_run()) + events = [msg for _level, msg in records if "llama admission" in msg] + assert any("llama admission cancelled-before-upstream" in m for m in events), events + + +def test_non_streaming_times_out_returns_503(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # never released -> waiter times out + with pytest.raises(HTTPException) as exc: + await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + assert exc.value.status_code == 503 + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_non_streaming_queued_then_admitted(monkeypatch): + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + task = asyncio.create_task( + anthropic_messages(_payload(), request = _Request(), current_subject = "t") + ) + await asyncio.sleep(0.1) + assert _snapshot().queued == 1 # waiting on the busy slot + held[0].release() # free it + response = await asyncio.wait_for(task, timeout = 2) + assert response.status_code == 200 + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_capacity_enforced_from_effective_parallel_slots(monkeypatch): + _install_backend(monkeypatch, slots = 3) + + async def _run(): + held = _occupy(_KEY, 3, 3) # all 3 slots busy + task = asyncio.create_task( + anthropic_messages(_payload(), request = _Request(), current_subject = "t") + ) + await asyncio.sleep(0.1) + snap = _snapshot() + assert snap.capacity == 3 and snap.active == 3 and snap.queued == 1 + for lease in held: + lease.release() + response = await asyncio.wait_for(task, timeout = 2) + assert response.status_code == 200 + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_limit(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # would block if admission were on + response = await asyncio.wait_for( + anthropic_messages(_payload(), request = _Request(), current_subject = "t"), + timeout = 2, + ) + assert response.status_code == 200 + for lease in held: + lease.release() + + asyncio.run(_run()) + + +# ── Streaming ───────────────────────────────────────────────── + + +def test_streaming_completes_and_releases_slot(monkeypatch): + _install_backend(monkeypatch, slots = 1) + + async def _run(): + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + blob = await _consume(response) + assert "event: message_start" in blob + assert "event: message_stop" in blob + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_streaming_emits_keepalives_while_queued_then_streams(monkeypatch): + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + # First chunk must be a keep-alive comment (slot still busy). + first = await asyncio.wait_for(body.__anext__(), timeout = 2) + first = first.decode() if isinstance(first, (bytes, bytearray)) else first + assert first.startswith(":") # SSE comment keep-alive + held[0].release() # free the slot -> real stream follows + rest = await asyncio.wait_for(_drain(body), timeout = 2) + assert "event: message_start" in rest + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_streaming_queue_full_returns_429(monkeypatch): + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + get_llama_admission_queue(_KEY).reserve( + capacity = 1, config = LlamaAdmissionConfig(max_queue = 1) + ) + with pytest.raises(HTTPException) as exc: + await anthropic_messages(_payload(stream = True), request = _Request(), current_subject = "t") + assert exc.value.status_code == 429 + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_streaming_disconnect_while_queued_frees_slot(monkeypatch): + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # one keep-alive + assert _snapshot().queued == 1 + await body.aclose() # client goes away mid-wait + held[0].release() + await asyncio.sleep(0.05) + snap = _snapshot() + assert snap.queued == 0 and snap.active == 0 + + asyncio.run(_run()) + + +# ── Shared queue + fairness + speed ─────────────────────────── + + +def test_shares_queue_with_openai_by_base_url(monkeypatch): + """The two API surfaces must land on one pool of the same llama-server slots. + + Reserves through the OpenAI helper the /v1/chat/completions path uses, rather + than poking the queue directly, so this fails if either side ever derives a + different key. + """ + _install_backend(monkeypatch, slots = 1) + + async def _run(): + openai_reservation, _ = inf_mod._openai_llama_admission_reserve( + request = _Request(), llama_backend = inf_mod.get_llama_cpp_backend() + ) + openai_lease = openai_reservation.lease_nowait() + assert openai_lease is not None + assert _snapshot().active == 1 # same key the Anthropic side will use + + task = asyncio.create_task( + anthropic_messages(_payload(), request = _Request(), current_subject = "t") + ) + await asyncio.sleep(0.1) + assert _snapshot().queued == 1 # queued behind the OpenAI generation + openai_lease.release() + assert (await asyncio.wait_for(task, timeout = 2)).status_code == 200 + + asyncio.run(_run()) + + +def test_non_streaming_client_gone_while_queued_returns_499(monkeypatch): + # The disconnect-while-queued branch; nothing else exercised 499. + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + with pytest.raises(HTTPException) as exc: + await anthropic_messages( + _payload(), request = _Request(disconnected = True), current_subject = "t" + ) + assert exc.value.status_code == 499 + assert _snapshot().queued == 0 # waiter cleaned up, not left parked + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_streaming_timeout_emits_an_error_event_and_frees_the_slot(monkeypatch): + # Only the non-streaming 503 was covered; streaming reports in-band instead. + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # never released, so the waiter times out + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = await _consume(response) + assert "event: error" in body + assert "message_start" not in body # never reached the model + for lease in held: + lease.release() + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_fifo_fairness_across_many_waiters(monkeypatch): + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + order = [] + + async def _one(i): + resp = await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + order.append(i) + return resp + + tasks = [asyncio.create_task(_one(i)) for i in range(8)] + await asyncio.sleep(0.2) + assert _snapshot().queued == 8 + held[0].release() + await asyncio.wait_for(asyncio.gather(*tasks), timeout = 5) + assert order == list(range(8)) # granted in arrival order + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_uncontended_hot_path_is_fast(monkeypatch): + _install_backend(monkeypatch, slots = 4) + + async def _run(): + start = time.perf_counter() + for _ in range(50): + resp = await anthropic_messages(_payload(), request = _Request(), current_subject = "t") + assert resp.status_code == 200 + elapsed = time.perf_counter() - start + # Generous ceiling on purpose: this guards against admission accidentally + # serialising or sleeping on the uncontended path, not against a slow + # runner, so it must not flake on a loaded CI box. + assert elapsed < 10.0, f"50 uncontended round-trips took {elapsed:.2f}s" + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +async def _drain(body): + chunks = [] + async for chunk in body: + chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk) + return "".join(chunks) + + +def test_streaming_midstream_cancel_finalizes_the_monitor(monkeypatch): + # A mid-stream disconnect is delivered as CancelledError so the monitored body + # can finalize its entry. Closing the inner iterator with aclose() instead + # delivers GeneratorExit, and the entry stays "running" for the process life. + _install_backend(monkeypatch, slots = 1) + + async def _run(): + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # stream started + assert inf_mod.api_monitor.active_count() == 1 + + # Propagates back out, as the un-admitted path did; what matters is that + # the monitored body saw it on the way through. + with pytest.raises(asyncio.CancelledError): + await body.athrow(asyncio.CancelledError()) # client vanished + + assert inf_mod.api_monitor.active_count() == 0 + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_streaming_give_up_while_queued_finalizes_the_monitor(monkeypatch): + # Cancelled before the body ever ran, so nothing downstream can close the + # entry out; the wrapper has to do it. + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # keep-alive, still queued + assert inf_mod.api_monitor.active_count() == 1 + + await body.aclose() # give up while waiting + + assert inf_mod.api_monitor.active_count() == 0 + for lease in held: + lease.release() + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_every_dispatch_site_goes_through_admission(): + """All six generation returns in anthropic_messages are admission-wrapped. + + The tool paths need a passthrough-capable backend and a tools payload to reach + at runtime, so guard them structurally instead: a new dispatch site added + without admission (or one reverted to _monitored_anthropic) fails here. + """ + import ast + import inspect + + tree = ast.parse(inspect.getsource(inf_mod).replace("\t", " ")) + handler = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.AsyncFunctionDef) and node.name == "anthropic_messages" + ) + # The wrappers themselves call _monitored_anthropic; only the dispatch sites count. + nested = { + node + for node in ast.walk(handler) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name.startswith("_admitted_anthropic") + } + inner = {id(n) for wrapper in nested for n in ast.walk(wrapper)} + + called = [] + for node in ast.walk(handler): + if id(node) in inner or not isinstance(node, ast.Call): + continue + if isinstance(node.func, ast.Name): + called.append(node.func.id) + + assert called.count("_admitted_anthropic") == 6 + assert called.count("_monitored_anthropic") == 0 + + +def test_queued_give_up_runs_the_response_pre_start_cleanup(monkeypatch): + """A stream abandoned while queued must run the builder's eager cleanup. + + The passthrough enters a _TrackedCancel before returning its response and + relies on the stream's finally to exit it. That finally never runs for a + generator that never started, so the response carries a pre-start hook and + the admission wrapper has to chain to it instead of replacing it. + """ + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + ran = [] + + async def _hook(): + ran.append(True) + + real = inf_mod._sse_streaming_response + + def _tagged(content, *, unstarted_cleanup = None): + return real(content, unstarted_cleanup = _hook) + + monkeypatch.setattr(inf_mod, "_sse_streaming_response", _tagged) + + async def _run(): + held = _occupy(_KEY, 1, 1) + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # keep-alive, still queued + await body.aclose() # give up before the body ran + + assert ran == [True] + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_passthrough_stream_registers_a_pre_start_cleanup(): + # Structural guard: the tracker is entered eagerly, so the response must + # carry the hook that exits it when the body never starts. + import ast + import inspect + + src = inspect.getsource(inf_mod._anthropic_passthrough_stream) + tree = ast.parse(src.replace("\t", " ").lstrip()) + returns = [n for n in ast.walk(tree) if isinstance(n, ast.Return) and n.value is not None] + call = next( + n.value + for n in returns + if isinstance(n.value, ast.Call) + and getattr(n.value.func, "id", "") == "_sse_streaming_response" + ) + hook = next(kw.value for kw in call.keywords if kw.arg == "unstarted_cleanup") + # Not just present: a literal None passes the keyword check and still leaks. + assert isinstance(hook, ast.Call) + assert getattr(hook.func, "id", None) == "_tracked_cancel_unstarted_cleanup" + + +def test_slot_is_released_even_if_closing_the_body_raises(monkeypatch): + # A slot lost here never comes back: with no queue timeout the pool silently + # shrinks and later callers wait forever, so the release must not sit behind + # anything that can throw. + _install_backend(monkeypatch, slots = 1) + + async def _boom(iterator, *, cancelled): + raise RuntimeError("close failed") + + monkeypatch.setattr(inf_mod, "_close_openai_admitted_stream_iterator", _boom) + + async def _run(): + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # stream started + assert _snapshot().active == 1 + + with pytest.raises(RuntimeError): + await body.aclose() + + assert _snapshot().active == 0 # slot returned despite the failure + # And the pool still serves the next caller. + again = get_llama_admission_queue(_KEY).reserve(capacity = 1, config = LlamaAdmissionConfig()) + lease = again.lease_nowait() + assert lease is not None + lease.release() + + asyncio.run(_run()) + + +_CLIENT_TOOLS = [ + {"name": "get_time", "description": "t", "input_schema": {"type": "object", "properties": {}}} +] + + +def _passthrough_payload(**fields): + # server_tools off + declared tools + a passthrough-capable backend routes + # anthropic_messages down the client-tool passthrough dispatch site. + return _payload(tools = _CLIENT_TOOLS, enable_tools = False, **fields) + + +def test_response_pre_start_cleanup_exits_the_passthrough_tracker(monkeypatch): + """A disconnect before the body starts must still exit the cancel tracker. + + The wrapper replaces the response's own pre-start hook, so it has to chain to + it. Asserting through _CANCEL_REGISTRY rather than the wiring, because the + hook can be present and still be a no-op. + """ + backend = _install_backend(monkeypatch, slots = 1) + backend.supports_tool_passthrough = True + monkeypatch.setattr(inf_mod, "_CANCEL_REGISTRY", {}) + + async def _run(): + response = await anthropic_messages( + _passthrough_payload(stream = True), request = _Request(), current_subject = "t" + ) + assert inf_mod._CANCEL_REGISTRY, "passthrough should have registered a tracker" + + cleanup = getattr(response, "_unstarted_cleanup", None) + assert cleanup is not None + await cleanup() # what _SameTaskStreamingResponse runs on a pre-start disconnect + + assert inf_mod._CANCEL_REGISTRY == {} + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_passthrough_dispatch_site_reserves_and_releases(monkeypatch): + # Behavioural cover for a dispatch site the other tests never reach. + backend = _install_backend(monkeypatch, slots = 1) + backend.supports_tool_passthrough = True + + async def _run(): + held = _occupy(_KEY, 1, 1) + task = asyncio.create_task( + anthropic_messages(_passthrough_payload(), request = _Request(), current_subject = "t") + ) + await asyncio.sleep(0.1) + assert _snapshot().queued == 1 # queued behind the busy slot, not bypassing + for lease in held: + lease.release() + with contextlib.suppress(Exception): + await asyncio.wait_for(task, timeout = 2) # upstream is not mocked + assert _snapshot().active == 0 and _snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_stream_setup_failure_returns_the_slot(monkeypatch): + # count_chat_tokens makes a blocking HTTP call to llama-server, so a dead + # server raises here: after lease_nowait() took the slot, before a body + # exists to release it. Nothing else can hand the slot back. + def _boom(*_a, **_k): + raise RuntimeError("tokenizer unreachable") + + _install_backend(monkeypatch, slots = 1, count_tokens = _boom) + + async def _run(): + with pytest.raises(RuntimeError): + await anthropic_messages(_payload(stream = True), request = _Request(), current_subject = "t") + snap = _snapshot() + assert snap.active == 0, f"slot leaked after stream setup failed: {snap}" + # And the pool still serves the next caller. + again = get_llama_admission_queue(_KEY).reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert again.lease_nowait() is not None + + asyncio.run(_run()) + + +def test_queued_non_stream_cancel_does_not_leak_a_coroutine(monkeypatch): + # The non-stream path builds the generation coroutine before reserving and + # only awaits it once admitted. Giving up while queued must close it. + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) + task = asyncio.create_task( + anthropic_messages(_payload(), request = _Request(), current_subject = "t") + ) + await asyncio.sleep(0.1) + assert _snapshot().queued == 1 + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + for lease in held: + lease.release() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + asyncio.run(_run()) + gc.collect() + leaked = [w for w in caught if "never awaited" in str(w.message)] + assert not leaked, [str(w.message) for w in leaked] + + +def test_stream_timeout_marks_the_monitor_entry_as_error(monkeypatch): + # The finally finishes the entry as "cancelled"; without the fail() first, a + # timed-out request is indistinguishable from a client hang-up in the + # monitor. api_monitor.finish is a no-op on an already terminal entry. + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05") + _install_backend(monkeypatch, slots = 1) + + async def _run(): + held = _occupy(_KEY, 1, 1) # never released, so the waiter times out + response = await anthropic_messages( + _payload(stream = True), request = _Request(), current_subject = "t" + ) + async for _ in response.body_iterator: + pass + entries = inf_mod.api_monitor.snapshot() + assert entries and entries[0]["status"] == "error", entries + for lease in held: + lease.release() + + asyncio.run(_run()) + + +class _RespawnBackend: + """Backend whose base_url moves to a new port once respawned.""" + + def __init__( + self, + *, + mtp_handled = False, + fallback_in_progress = False, + ): + self.base_url = "http://127.0.0.1:57953" + self.context_length = 4096 + self.respawn_calls = 0 + self._mtp_handled = mtp_handled + self._mtp_runtime_fallback_in_progress = fallback_in_progress + + def count_chat_tokens(self, *_a, **_k): + return 2 + + def _maybe_recover_from_mtp_crash(self, _exc): + return self._mtp_handled + + def _respawn_if_dead(self): + self.respawn_calls += 1 + self.base_url = "http://127.0.0.1:62933" + return True + + +def test_retry_url_stands_down_while_an_mtp_fallback_is_reloading(): + # Only the first caller gets True from _maybe_recover_from_mtp_crash; the rest + # see False and must still stand down, or they respawn the same MTP config + # underneath the fallback already reloading without it. + backend = _RespawnBackend(mtp_handled = False, fallback_in_progress = True) + + url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x"))) + + assert url is None + assert backend.respawn_calls == 0 + + +class _PtRequest: + async def is_disconnected(self): + return False + + +async def _passthrough_response(backend): + return await _anthropic_passthrough_stream( + _PtRequest(), + threading.Event(), + backend, + [{"role": "user", "content": "hi"}], + [], + 0.7, + 0.95, + 20, + 16, + "msg_tracker_probe", + "test-model", + ) + + +def test_disconnect_during_the_opening_lines_exits_the_tracker(): + # Suspended inside emitter.start()'s yields the generator has not reached the + # try/finally that exits the tracker, so those yields need their own. + backend = _RespawnBackend() + + async def _run(): + response = await _passthrough_response(backend) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) # first start line + assert inf_mod._CANCEL_REGISTRY, "tracker should be registered" + await body.aclose() + assert inf_mod._CANCEL_REGISTRY == {}, "tracker leaked" + + asyncio.run(_run()) + + +def test_cancel_during_the_opening_lines_exits_the_tracker(): + # Same window, delivered the way _SameTaskStreamingResponse delivers it. + backend = _RespawnBackend() + + async def _run(): + response = await _passthrough_response(backend) + body = response.body_iterator + await asyncio.wait_for(body.__anext__(), timeout = 2) + assert inf_mod._CANCEL_REGISTRY, "tracker should be registered" + with pytest.raises(asyncio.CancelledError): + await body.athrow(asyncio.CancelledError()) + assert inf_mod._CANCEL_REGISTRY == {}, "tracker leaked" + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 621ac9aaca..296cb80911 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1523,6 +1523,17 @@ def _reset_policy(): reset_tool_policy() +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + # The admission queue is process-global; isolate the shared "llama-server" key + # so one test's leftover reservation can't stall the next. + from core.inference.llama_admission import reset_llama_admission_queues + + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + class TestAnthropicMessagesToolRouting: class _Request: state = SimpleNamespace() diff --git a/studio/backend/tests/test_anthropic_passthrough_respawn.py b/studio/backend/tests/test_anthropic_passthrough_respawn.py new file mode 100644 index 0000000000..a9f31208ed --- /dev/null +++ b/studio/backend/tests/test_anthropic_passthrough_respawn.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Restart survival for the Anthropic /v1/messages passthrough. + +A crashed llama-server relaunches on a NEW ephemeral port. Before the retry the +passthrough kept posting to the dead port, so a Claude Code session stayed broken +until the next explicit load. These cover the respawn-and-retry on both the +streaming and non-streaming passthroughs. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import threading +from types import SimpleNamespace + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import routes.inference as inf_mod +from routes.inference import ( + _anthropic_passthrough_non_streaming, + _anthropic_passthrough_retry_url, + _anthropic_passthrough_stream, +) + +_DEAD = "http://127.0.0.1:57953" +_FRESH = "http://127.0.0.1:62933" + + +class _Backend: + """Stub llama backend whose base_url moves to a new port once respawned.""" + + def __init__( + self, + *, + respawn_ok = True, + mtp_handled = False, + ): + self.base_url = _DEAD + self.context_length = 4096 + self.respawn_calls = 0 + self.mtp_calls = 0 + self._respawn_ok = respawn_ok + self._mtp_handled = mtp_handled + + def count_chat_tokens(self, *_args, **_kwargs): + return 2 + + def _maybe_recover_from_mtp_crash(self, _exc): + self.mtp_calls += 1 + return self._mtp_handled + + def _respawn_if_dead(self): + self.respawn_calls += 1 + if not self._respawn_ok: + return False + self.base_url = _FRESH + return True + + +class _Request: + async def is_disconnected(self): + return False + + +class _FakeNonStreamingClient: + def __init__(self): + self.urls = [] + + async def post(self, url, **_kwargs): + self.urls.append(url) + if url.startswith(_DEAD): + raise httpx.ConnectError("connection refused") + return httpx.Response( + 200, + json = { + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 2, "completion_tokens": 1}, + }, + ) + + +def _install_stream_transport(monkeypatch, calls): + def handler(request: httpx.Request) -> httpx.Response: + calls.append(str(request.url)) + if str(request.url).startswith(_DEAD): + raise httpx.ConnectError("connection refused") + content = ( + f"data: {json.dumps({'choices': [{'delta': {'content': 'hi'}}]})}\n\n" + "data: [DONE]\n\n" + ) + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def _client(*_args, **kwargs): + return real_client(transport = transport, timeout = kwargs.get("timeout", 600)) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + + +async def _run_stream(backend): + response = await _anthropic_passthrough_stream( + _Request(), + threading.Event(), + backend, + [{"role": "user", "content": "hi"}], + [], + 0.7, + 0.95, + 20, + 16, + "msg_1", + "test-model", + ) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk) + return "".join(chunks) + + +async def _run_non_streaming(backend): + return await _anthropic_passthrough_non_streaming( + backend, + [{"role": "user", "content": "hi"}], + [], + 0.7, + 0.95, + 20, + 16, + "msg_1", + "test-model", + ) + + +# ── Helper ──────────────────────────────────────────────────── + + +def test_retry_url_rebuilds_from_the_respawned_base_url(): + backend = _Backend() + + url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x"))) + + assert url == f"{_FRESH}/v1/chat/completions" + assert backend.respawn_calls == 1 + + +def test_retry_url_is_none_when_nothing_respawned(): + backend = _Backend(respawn_ok = False) + + url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x"))) + + assert url is None + + +def test_retry_url_defers_to_the_mtp_crash_recovery(): + # An MTP+tensor crash schedules its own reload; retrying would race it. + backend = _Backend(mtp_handled = True) + + url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x"))) + + assert url is None + assert backend.respawn_calls == 0 + + +def test_retry_url_tolerates_a_backend_without_respawn_hooks(): + backend = SimpleNamespace(base_url = _DEAD) + + url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x"))) + + assert url is None + + +# ── Non-streaming ───────────────────────────────────────────── + + +def test_non_streaming_retries_against_the_new_port(monkeypatch): + client = _FakeNonStreamingClient() + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + backend = _Backend() + + response = asyncio.run(_run_non_streaming(backend)) + + assert response.status_code == 200 + assert backend.respawn_calls == 1 + assert client.urls == [f"{_DEAD}/v1/chat/completions", f"{_FRESH}/v1/chat/completions"] + + +def test_non_streaming_raises_when_the_server_stays_dead(monkeypatch): + client = _FakeNonStreamingClient() + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + backend = _Backend(respawn_ok = False) + + with pytest.raises(httpx.ConnectError): + asyncio.run(_run_non_streaming(backend)) + + assert client.urls == [f"{_DEAD}/v1/chat/completions"] # no blind retry + + +def test_non_streaming_does_not_retry_an_mtp_crash(monkeypatch): + client = _FakeNonStreamingClient() + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + backend = _Backend(mtp_handled = True) + + with pytest.raises(httpx.ConnectError): + asyncio.run(_run_non_streaming(backend)) + + assert backend.respawn_calls == 0 + + +# ── Streaming ───────────────────────────────────────────────── + + +def test_streaming_retries_against_the_new_port(monkeypatch): + calls = [] + _install_stream_transport(monkeypatch, calls) + backend = _Backend() + + blob = asyncio.run(_run_stream(backend)) + + assert backend.respawn_calls == 1 + assert calls == [f"{_DEAD}/v1/chat/completions", f"{_FRESH}/v1/chat/completions"] + # The retried stream really produced the turn, not just a clean-looking stop. + assert "event: message_start" in blob + assert "event: message_stop" in blob + assert "hi" in blob + + +def test_streaming_emits_an_error_event_when_the_server_stays_dead(monkeypatch): + calls = [] + _install_stream_transport(monkeypatch, calls) + backend = _Backend(respawn_ok = False) + + blob = asyncio.run(_run_stream(backend)) + + assert calls == [f"{_DEAD}/v1/chat/completions"] # no blind retry + assert "event: error" in blob + + +def test_streaming_does_not_retry_an_mtp_crash(monkeypatch): + calls = [] + _install_stream_transport(monkeypatch, calls) + backend = _Backend(mtp_handled = True) + + blob = asyncio.run(_run_stream(backend)) + + assert backend.respawn_calls == 0 + assert calls == [f"{_DEAD}/v1/chat/completions"] + assert "event: error" in blob diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py index 2f04e81926..f69eb7c5c9 100644 --- a/studio/backend/tests/test_llama_admission.py +++ b/studio/backend/tests/test_llama_admission.py @@ -16,6 +16,7 @@ from core.inference.llama_admission import ( ADMISSION_CONTROL_ENV, ADMISSION_KEEPALIVE_INTERVAL_ENV, ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_PER_SLOT_ENV, ADMISSION_QUEUE_TIMEOUT_ENV, DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, DEFAULT_ADMISSION_MAX_QUEUE, @@ -28,8 +29,23 @@ from core.inference.llama_admission import ( ) +_ADMISSION_ENV = ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_PER_SLOT_ENV, + *llama_admission._LEGACY_ENV.values(), +) + + @pytest.fixture(autouse = True) -def _reset_queues(): +def _reset_queues(monkeypatch): + # Clear ambient settings for every test, not just the ones that remember to: + # a canonical name set on the machine silently beats the legacy name a test + # is exercising, and the queue registry is process-global. + for name in _ADMISSION_ENV: + monkeypatch.delenv(name, raising = False) reset_llama_admission_queues() yield reset_llama_admission_queues() @@ -41,15 +57,25 @@ def test_admission_config_defaults(monkeypatch): ADMISSION_QUEUE_TIMEOUT_ENV, ADMISSION_KEEPALIVE_INTERVAL_ENV, ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_PER_SLOT_ENV, + "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL", + "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", ): monkeypatch.delenv(name, raising = False) config = llama_admission_config_from_env() + # Literals, not the module constants: comparing a default to itself would let + # any future value change through silently. 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 + assert config.queue_timeout_s is None # wait forever + assert config.keepalive_interval_s == 5.0 + assert config.max_queue is None # no absolute cap + assert config.queue_per_slot == 16 + assert (DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, DEFAULT_ADMISSION_MAX_QUEUE) == (None, None) + assert DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S == 5.0 def test_admission_config_env_overrides(monkeypatch): @@ -66,6 +92,25 @@ def test_admission_config_env_overrides(monkeypatch): assert config.max_queue is None +def test_admission_config_honors_legacy_openai_compat_env(monkeypatch): + # The queue is shared with /v1/messages now, but existing OPENAI_COMPAT + # settings must keep working. + monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7") + monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", "off") + + config = llama_admission_config_from_env() + + assert config.max_queue == 7 + assert config.enabled is False + + +def test_admission_config_prefers_neutral_env_over_legacy(monkeypatch): + monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "3") + + assert llama_admission_config_from_env().max_queue == 3 + + def test_admission_config_positive_queue_timeout_env(monkeypatch): monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") @@ -106,6 +151,160 @@ def test_fifo_capacity_one_grants_next_waiter_on_release(): asyncio.run(_run()) +def test_pool_hands_out_distinct_slots_and_reuses_them(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + leases = [queue.reserve(capacity = 3, config = config).lease_nowait() for _ in range(3)] + assert sorted(lease.slot for lease in leases) == [0, 1, 2] # one slot each + snapshot = queue.snapshot() + assert (snapshot.active, snapshot.free, snapshot.capacity) == (3, 0, 3) + + # A freed slot returns to the pool and is handed to the next caller. + freed = leases[1].slot + leases[1].release() + assert queue.snapshot().free == 1 + reused = queue.reserve(capacity = 3, config = config).lease_nowait() + assert reused.slot == freed + + reused.release() + leases[0].release() + leases[2].release() + snapshot = queue.snapshot() + assert (snapshot.active, snapshot.free) == (0, 3) + + asyncio.run(_run()) + + +def test_pool_waiter_is_handed_a_real_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + held = queue.reserve(capacity = 1, config = config).lease_nowait() + waiting = queue.reserve(capacity = 1, config = config) + assert waiting.lease_nowait() is None + assert queue.snapshot().free == 0 + + held.release() + granted = await waiting.wait(0.1) + assert granted is not None and granted.slot == 0 # the slot just freed + granted.release() + + asyncio.run(_run()) + + +def test_shrinking_capacity_retires_slots_beyond_the_new_pool(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)] + assert queue.snapshot().capacity == 4 + + # llama-server reloaded with fewer --parallel slots; in-flight holders keep + # running and their slots retire instead of returning to the smaller pool. + shrunk = queue.reserve(capacity = 2, config = config) + assert shrunk.lease_nowait() is None # all 4 still held, nothing free + for lease in leases: + lease.release() + + granted = await shrunk.wait(0.1) + assert granted is not None and granted.slot < 2 + granted.release() + snapshot = queue.snapshot() + assert (snapshot.capacity, snapshot.active, snapshot.free) == (2, 0, 2) + + asyncio.run(_run()) + + +def test_queue_limit_scales_with_the_serving_slots(): + # The wait line follows --parallel: 16 per slot, floored at 64 so a 1-slot + # backend keeps the depth it had before scaling existed. + config = LlamaAdmissionConfig() + assert config.queue_limit(4) == 64 # --parallel 4 (the default) + assert config.queue_limit(8) == 128 # --parallel 8 + assert config.queue_limit(16) == 256 + assert config.queue_limit(1) == 64 # floor, not 16 + assert config.queue_limit(2) == 64 # floor, not 32 + # An explicit cap wins, and a None multiplier means an unbounded line. + assert LlamaAdmissionConfig(max_queue = 5).queue_limit(8) == 5 + assert LlamaAdmissionConfig(queue_per_slot = None).queue_limit(8) is None + # Non-positive settings mean unbounded, never "reject everything". + assert LlamaAdmissionConfig(max_queue = 0).queue_limit(4) is None + assert LlamaAdmissionConfig(max_queue = -1).queue_limit(4) is None + assert LlamaAdmissionConfig(queue_per_slot = 0).queue_limit(4) is None + assert LlamaAdmissionConfig(queue_per_slot = -3).queue_limit(4) is None + + +def test_queue_limit_rejects_only_once_the_line_is_full(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + # Explicit cap, so the test drives rejection without standing up the 64 + # waiters the scaled floor would otherwise require. + config = LlamaAdmissionConfig(max_queue = 4) + + held = [queue.reserve(capacity = 2, config = config).lease_nowait() for _ in range(2)] + parked = [queue.reserve(capacity = 2, config = config) for _ in range(4)] + assert queue.snapshot().queued == 4 + + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 2, config = config) + + for reservation in parked: + reservation.cancel() + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_waiting_is_never_timed_out_by_default(): + # "Wait forever": the default config sets no queue timeout at all. + assert llama_admission_config_from_env().queue_timeout_s is None + assert LlamaAdmissionConfig().queue_timeout_s is None + + +def test_single_request_at_a_time_never_queues_or_allocates_waiters(): + # The common serving case: one request in flight at a time must take a slot + # straight away and never touch the wait line. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + for _ in range(50): + reservation = queue.reserve(capacity = 4, config = config) + lease = reservation.lease_nowait() + assert lease is not None # admitted immediately + assert queue.snapshot().queued == 0 # nobody ever lined up + lease.release() + snapshot = queue.snapshot() + assert (snapshot.active, snapshot.free, snapshot.queued) == (0, 4, 0) + + asyncio.run(_run()) + + +def test_unbounded_queue_keeps_waiting_instead_of_rejecting(): + # queue_per_slot None is the "pool + unbounded wait line" mode: nothing is + # ever rejected, callers just line up for the next free slot. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = None, queue_per_slot = None) + + held = queue.reserve(capacity = 1, config = config).lease_nowait() + waiters = [queue.reserve(capacity = 1, config = config) for _ in range(200)] + assert queue.snapshot().queued == 200 # no LlamaAdmissionQueueFull + + held.release() + first = await waiters[0].wait(0.1) + assert first is not None + first.release() + for waiter in waiters[1:]: + waiter.cancel() + + asyncio.run(_run()) + + def test_queue_full_rejects_excess_waiter(): async def _run(): queue = get_llama_admission_queue("http://llama.test") @@ -288,6 +487,105 @@ def test_lease_release_is_idempotent_under_concurrent_calls(): asyncio.run(_run()) +def test_releasing_a_stale_lease_does_not_free_someone_elses_slot(): + # The concurrent test above passes without the _released guard: the racing + # calls all target a still-live slot, which the bitmask already absorbs. The + # case the guard exists for is a slot released twice with a reuse in between. + # It is live: _wait_for_openai_admission_non_streaming releases and re-raises, + # then the caller's finally cancels the reservation and releases the same + # lease again, by which point the slot can belong to another request. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + stale = queue.reserve(capacity = 1, config = config).lease_nowait() + stale.release() + other = queue.reserve(capacity = 1, config = config).lease_nowait() + assert other.slot == stale.slot # the slot got reused + + stale.release() + assert queue.snapshot().active == 1, "stale release handed back a live slot" + other.release() + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + +def test_grant_reclaims_the_slot_when_the_waiters_loop_is_gone(): + # _grant_waiters_locked takes the slot before scheduling delivery, so if the + # schedule fails the bit is already set. Leaving it set strands the slot for + # good, because _free is rebuilt from the bitmask. + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + held = None + + dead = asyncio.new_event_loop() + try: + + async def _fill_and_queue(): + nonlocal held + held = queue.reserve(capacity = 1, config = config).lease_nowait() + assert queue.reserve(capacity = 1, config = config).lease_nowait() is None + + dead.run_until_complete(_fill_and_queue()) + finally: + dead.close() + + held.release() # grant path now hits the closed loop + assert queue.snapshot().active == 0 + assert queue.is_idle() + + +def test_cancel_returns_the_granted_slot_when_the_waiters_loop_is_gone(): + # Routes cancel() from finally blocks, so a raise here would mask their + # exception and skip the release that hands the granted slot back. + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + held = reservation = None + + dead = asyncio.new_event_loop() + try: + + async def _fill_and_queue(): + nonlocal held, reservation + held = queue.reserve(capacity = 1, config = config).lease_nowait() + reservation = queue.reserve(capacity = 1, config = config) + + dead.run_until_complete(_fill_and_queue()) + held.release() # promotes the waiter, so cancel() has a lease to return + finally: + dead.close() + + reservation.cancel() + assert queue.snapshot().active == 0 + assert queue.is_idle() + + +def test_delivery_to_an_already_finished_waiter_releases_the_slot(): + # A slot is taken before delivery is scheduled, so if the waiter finishes in + # that window someone has to hand it back. _deliver_lease does it twice over, + # in the dead-waiter branch and in the InvalidStateError backstop; this pins + # the outcome, not which one. Reaches into the waiter because no public call + # leaves that window open: queue.cancel() reclaims granted_lease itself. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + held = queue.reserve(capacity = 1, config = config).lease_nowait() + reservation = queue.reserve(capacity = 1, config = config) + waiter = reservation._waiter + + held.release() # schedules _deliver_lease, sets granted_lease + waiter.future.cancel() # finishes the future before the callback runs + assert waiter.granted_lease is not None + await asyncio.sleep(0) # let the callback run + + assert queue.snapshot().active == 0 + assert queue.is_idle() + + 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. @@ -318,3 +616,234 @@ def test_new_key_retains_in_flight_prior_load_queue(): assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} asyncio.run(_run()) + + +def test_capacity_shrink_never_admits_past_the_new_ceiling(): + # A load that downshifts --parallel (or an unload resetting it to 1) shrinks the + # pool while slots are still held. Those holdovers keep occupying the backend, so + # they must count against the ceiling; sizing on free ids alone over-admits. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)] + assert all(lease is not None for lease in held) + waiter = queue.reserve(capacity = 4, config = config) + + queue.reserve(capacity = 1, config = config) # capacity collapses to 1 + # Release the one id that still falls inside the shrunk pool, so it goes + # back on the free list; ids at or above capacity retire instead. + low = min(held, key = lambda lease: lease.slot) + assert low.slot == 0 + low.release() + + # The other 3 holdovers are still generating, which already meets the new + # ceiling, so the freed id must not be handed on. Gating on "is an id free" + # alone grants it here and puts 4 generations on a 1-slot backend. + with pytest.raises(asyncio.TimeoutError): + await waiter.wait(0.2) + assert queue.snapshot().active == 3 + + waiter.cancel() + for lease in held: + if lease is not low: + lease.release() + + asyncio.run(_run()) + + +def test_queue_per_slot_env_is_parsed(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "4") + assert llama_admission_config_from_env().queue_limit(32) == 128 + # Non-positive asks for an unbounded line rather than rejecting everything. + monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "0") + assert llama_admission_config_from_env().queue_limit(32) is None + + +def test_max_queue_zero_from_env_is_unbounded_end_to_end(monkeypatch): + # Guards the whole env path, not just the parsed field: a regression that let + # queue_per_slot survive MAX_QUEUE=0 would silently re-bound the line. + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + config = llama_admission_config_from_env() + assert config.max_queue is None and config.queue_per_slot is None + assert config.queue_limit(1) is None and config.queue_limit(64) is None + + +def test_legacy_env_fallback_covers_every_setting(monkeypatch): + for canonical, legacy in llama_admission._LEGACY_ENV.items(): + monkeypatch.delenv(canonical, raising = False) + monkeypatch.setenv(legacy, "0" if "CONTROL" in canonical else "7") + config = llama_admission_config_from_env() + assert config.enabled is False + assert config.queue_timeout_s == 7.0 + assert config.keepalive_interval_s == 7.0 + assert config.max_queue == 7 + + +def test_empty_canonical_env_falls_through_to_legacy(monkeypatch): + # The branch _raw_env exists for: set but blank must not mask the legacy name. + monkeypatch.setenv(ADMISSION_CONTROL_ENV, " ") + monkeypatch.setenv(llama_admission._LEGACY_ENV[ADMISSION_CONTROL_ENV], "0") + assert llama_admission_config_from_env().enabled is False + + +def test_explicit_queue_per_slot_is_not_floored(monkeypatch): + # The floor exists so a 1-slot backend keeps its old depth by default, not to + # override an operator who asked for a shallow line. + monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "2") + config = llama_admission_config_from_env() + assert config.queue_limit(1) == 2 + assert config.queue_limit(8) == 16 + + # Unset, the default multiplier is floored instead. + monkeypatch.delenv(ADMISSION_QUEUE_PER_SLOT_ENV, raising = False) + assert llama_admission_config_from_env().queue_limit(1) == 64 + + # A value that does not parse falls back to the default multiplier, so it has + # to keep the default's floor. Otherwise a typo quietly shrinks the line 4x. + for garbage in ("abc", "1e3", "16.0"): + monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, garbage) + assert llama_admission_config_from_env().queue_limit(1) == 64, garbage + + +def test_module_imports_on_python_39(monkeypatch): + """No 3.10+ API on an import path. The package declares >=3.9 but CI only + runs 3.12, so a regression here would ship broken.""" + import ast + import pathlib + + src = pathlib.Path(llama_admission.__file__).read_text(encoding = "utf-8") + tree = ast.parse(src) + + # int.bit_count() (3.10+) + assert not [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "bit_count" + ] + # dataclass(slots = ...) is 3.10+, so every dataclass must take it through + # the version gate instead of naming it. A new one that forgets the gate + # loses slots silently, so require the **_SLOTS unpack rather than allow it. + seen = 0 + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = getattr(node.func, "id", None) or getattr(node.func, "attr", None) + if name != "dataclass": + continue + seen += 1 + assert "slots" not in {kw.arg for kw in node.keywords} + assert [ + kw + for kw in node.keywords + if kw.arg is None and getattr(kw.value, "id", None) == "_SLOTS" + ], ast.dump(node) + assert seen + + +def test_slots_gate_matches_the_running_interpreter(): + """The gate is only worth having if it actually applies where it can.""" + import sys + + gated = (LlamaAdmissionConfig, llama_admission.LlamaAdmissionSnapshot, llama_admission._Waiter) + if sys.version_info >= (3, 10): + assert llama_admission._SLOTS == {"slots": True} + for cls in gated: + assert getattr(cls, "__slots__", None), cls + else: + assert llama_admission._SLOTS == {} + + # Construct through the gate either way: slots=True rebuilds the class, so a + # field it cannot carry over would only show up on instantiation. + config = LlamaAdmissionConfig(max_queue = 7) + assert config.max_queue == 7 and config.queue_limit(4) == 7 + assert llama_admission.LlamaAdmissionSnapshot("k", 1, 1, 0).capacity == 1 + + +def test_held_count_tracks_the_bitmask(): + # _held replaces int.bit_count(); the two must never drift apart. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + popcount = lambda: bin(queue._in_use).count("1") + + leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)] + assert queue._held == popcount() == 4 + leases[1].release() + assert queue._held == popcount() == 3 + shrunk = queue.reserve(capacity = 2, config = config) # shrink with slots held + assert queue._held == popcount() == 3 + shrunk.cancel() # else it is granted a slot as the others drain + for lease in leases: + lease.release() + assert queue._held == popcount() == 0 + + asyncio.run(_run()) + + +def test_snapshot_free_never_exceeds_what_can_be_admitted(): + # After a shrink, low ids can sit in _free while holdovers fill the ceiling. + # Reporting them as free made the admission log contradict itself. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)] + queue.reserve(capacity = 1, config = config) # capacity collapses to 1 + min(held, key = lambda lease: lease.slot).release() + + snapshot = queue.snapshot() + assert snapshot.free == 0, snapshot # nothing is actually takeable + assert snapshot.active == 3 + for lease in held: + lease.release() + + asyncio.run(_run()) + + +def test_a_newcomer_does_not_barge_past_a_parked_waiter(): + # Anti-starvation, pinned as behaviour rather than as the `if not self._waiters` + # check: _take_slot_locked consults _can_admit_locked anyway, so either alone + # refuses the newcomer. This fails if both ever go. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + held = queue.reserve(capacity = 1, config = config).lease_nowait() + parked = queue.reserve(capacity = 1, config = config) + assert parked.lease_nowait() is None + + held.release() + newcomer = queue.reserve(capacity = 1, config = config) + assert newcomer.lease_nowait() is None, "newcomer barged past the parked waiter" + assert (await parked.wait(0.1)) is not None + + asyncio.run(_run()) + + +def test_dead_waiters_stop_counting_against_the_queue_limit(): + # A future cancelled out of band leaves the entry in the deque: cancel() is not + # called, so only the prune drops it. Without that, depth, is_idle() and the + # queue-full limit all drift for the life of the queue. + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 2) + + held = queue.reserve(capacity = 1, config = config).lease_nowait() + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + assert queue.snapshot().queued == 2 + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + first._waiter.future.cancel() + second._waiter.future.cancel() + assert queue.snapshot().queued == 0, "dead waiters still occupy the line" + # The freed depth is usable again, and an idle queue is evictable. + queue.reserve(capacity = 1, config = config).cancel() + held.release() + assert queue.is_idle() + + asyncio.run(_run())