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

Studio: admission control on /v1/messages, slot pool that tracks --parallel (#7436)

Browse source 
* Studio: admission control on /v1/messages, slot pool that tracks --parallel

/v1/chat/completions was gated by the llama admission queue but /v1/messages
was not, so an Anthropic client could oversubscribe llama-server's slots and
stall the backend. Wire the same queue into all six /v1/messages dispatch
sites, and rework the queue itself into an explicit slot pool.

- Queue keyed by base_url, so both API surfaces share one pool of slots.
- Waiting is unbounded by default instead of timing out; the wait line is
  sized at 16 x the serving slots so it follows --parallel.
- Neutral UNSLOTH_LLAMA_ADMISSION_* env names, legacy UNSLOTH_OPENAI_COMPAT_*
  spellings still honored.
- Passthrough retries once against a respawned llama-server, which comes back
  on a new ephemeral port.

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

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

* Fix over-admission on capacity shrink and restore the stream cancel contract

Review of the previous commit turned up two real regressions plus smaller gaps.

- Pool sizing looked only at free slot ids, so when capacity shrank while slots
  were held (an unload resets effective_parallel_slots to 1) a freed low id was
  handed out even though the holdovers already met the new ceiling. Count every
  held slot against capacity instead. A 1-slot backend could run 4 generations.
- The streaming wrapper closed the monitored body with aclose(), delivering
  GeneratorExit where _SameTaskStreamingResponse deliberately throws
  CancelledError. The monitor entry was never finalized, so it leaked as
  "running" for the process lifetime and cancel_event was never set. Close
  through the shared helper so cancellation reaches the handler.
- Finalize the monitor when a stream is abandoned before its body starts, and
  when a queued non-streaming request is cancelled (which also leaked the
  un-awaited generation coroutine).
- Floor the scaled wait line at 64, so a 1-slot backend keeps the depth it had
  before scaling existed instead of dropping from 64 to 16.
- Use the canonical Anthropic type map: a full queue is 429 rate_limit_error,
  which SDKs back off on; overloaded_error is 529.
- Treat non-positive max_queue/queue_per_slot as unbounded rather than "reject
  everything", and reclaim the slot if a waiter's event loop is gone.

Tests: regression tests for both defects, verified to fail without the fix.
Adds env coverage for QUEUE_PER_SLOT and the legacy fallbacks, a structural
check that all six dispatch sites stay admission-wrapped, and clears the new
env var in the isolation fixtures.

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

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

* Run the response pre-start cleanup when a queued stream is abandoned

From automated review of the earlier commits.

_anthropic_passthrough_stream enters its _TrackedCancel eagerly and relies on
the stream's finally to exit it, but aclose() on an async generator that never
started is a no-op, so that finally never runs. Admission made this reachable:
a client that disconnects while queued leaves the cancel id registered in
_CANCEL_REGISTRY forever.

- Give the passthrough response an unstarted_cleanup hook that exits the
  tracker, via a new optional arg on _sse_streaming_response.
- Chain to that hook from the admission wrapper rather than replacing it, and
  run it when the wrapper gives up before the body started.
- Defer to an in-progress MTP fallback instead of respawning underneath it;
  only the first caller gets True from _maybe_recover_from_mtp_crash.

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

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

* Restore Python 3.9 support, and stop the floor overriding an explicit setting

Second review round. The first item is a real break shipped by the earlier
commits, the rest are correctness and contract fixes.

- dataclass(slots = True) and int.bit_count() are both 3.10+, but the package
  declares requires-python >=3.9 and CI only runs 3.12, so nothing caught it.
  Importing the module raised TypeError on 3.9, taking down the whole backend,
  not just admission. Drop the dataclass slots and track the popcount in a
  counter. A test now asserts neither API comes back.
- The queue-depth floor applied even when an operator set QUEUE_PER_SLOT
  explicitly, so asking for a shallow line silently got 64 and, with no queue
  timeout, callers blocked instead of failing fast. The floor now only backs
  the default multiplier.
- Never let a failing close strand a slot: closing runs in its own try so the
  release always happens. A lost slot shrinks the pool permanently.
- Close the generation coroutine when reserving fails for any reason, not only
  on a full queue.
- Exit the passthrough cancel tracker if the client drops while the opening SSE
  lines are still being sent; those yields sit outside the teardown try.
- snapshot.free now reports what a caller could actually take, so the admission
  log cannot show free slots next to queued requests after a shrink.
- Correct the class docstring: the wait line is bounded by default, not
  unlimited. Document that abandoning wait() requires cancel(), and pin the
  thread assumption in _deliver_lease.

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

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

* Make the leak guards real, and cover the untested admission branches

Third review round, which attacked the previous round's tests by reverting each
fix. Two guards turned out to be hollow.

- The pre-start cleanup chain could be severed with the suite still green: the
  existing test drove the generator finally, never the response hook. A real
  pre-start disconnect leaked the passthrough cancel tracker permanently.
  Replaced with a test that runs the response hook and asserts _CANCEL_REGISTRY
  is empty; verified against both ways of reintroducing the leak.
- The structural check only asserted the unstarted_cleanup keyword was present,
  so passing a literal None passed it while leaking. It now asserts the hook is
  actually built.
- test_shares_queue_with_openai_by_base_url never touched the OpenAI helper; it
  was a duplicate under a misleading name. It now reserves through the same
  helper /v1/chat/completions uses, so it fails if either surface ever derives a
  different key. That is the PR's central shared-queue claim.
- Cover the passthrough dispatch site, 499 on disconnect-while-queued, and the
  streaming admission timeout. Four of six sites previously had only an AST
  node count behind them.
- Clear admission env in the autouse fixture rather than per test: an ambient
  canonical name silently beat the legacy name a test was exercising.
- Loosen the wall-clock assertion, which guarded against serialising on the
  uncontended path, not against a slow runner.
- The class docstring claimed a global concurrency cap; Studio's own chat
  endpoint does not reserve, so it is not one.

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

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

* Keep dataclass slots on 3.10+ via a version gate

Dropping slots = True for 3.9 gave it up everywhere, including the 3.12 CI runs
and every supported interpreter but one. Gate it instead: _SLOTS is
{"slots": True} on 3.10+ and empty below, unpacked into each dataclass.

The AST scan now requires the unpack rather than merely forbidding a literal
slots keyword, so a dataclass added later cannot quietly lose slots. Added a
test that the gate matches the running interpreter, since a gate that never
applies is worse than no gate. Verified the 3.9 branch by forcing _SLOTS empty
and reloading: the full admission suite passes either way.

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

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

* Fix two slot leaks, and cover the guards that had no test

Third review round, three reviewers working independently on the admission core,
the route wiring, and whether the PR regresses anything it is not about.

Leaks:

- cancel() made the same call_soon_threadsafe as _grant_waiters_locked but
  without its RuntimeError guard. Routes cancel from finally blocks, so a closed
  loop masked their exception and skipped the release, stranding the slot and
  pinning is_idle() false so the queue was never evicted either.
- The pre-start cleanup released the slot after an await that can raise
  BaseException, which is swallowed upstream. Nested it in a finally, as the
  streaming and OpenAI paths already do.

An unparseable QUEUE_PER_SLOT dropped the 64 floor while falling back to the
default multiplier, quietly giving a 1-slot backend a 16-deep line. Explicit now
means it parsed.

Guards that were correct but had no test. Each was reverted, confirmed the suite
stayed green, then covered and confirmed red:

- the slot released when stream setup raises, which is the reachable one:
  count_chat_tokens is a blocking call to llama-server, so a dead backend raises
  after the slot is taken and before a body exists to release it
- coro.close() on a cancelled queued request, the api_monitor.fail that
  distinguishes an admission timeout from a client hang-up, the MTP fallback
  short-circuit, and the BaseException guard around the opening stream lines
- the queue-full test asserted a type string OpenAI's 429 also uses, so it
  passed against an OpenAI envelope. It now pins the Anthropic shape.

Anthropic requests were invisible in the admission log while sharing the pool
with chat completions, so the same events are logged there with a mode. Renamed
the helper to match, since it is no longer OpenAI-only.

Corrected two comments that described behaviour the code does not have: the
slot is taken when the streaming response is built, not when the body starts
iterating, and the pool is not a cap on every generation, since /v1/completions,
Studio's chat endpoint and RAG captioning all reach llama-server directly.

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

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

* Cover the admission telemetry, and drop a dead helper

Fourth review round. No bugs found in the code this time; the finding was that
most of the previous commit's telemetry had no test. Only queue-full was
asserted, so removing any of the other four log calls left the suite green.

All five are covered now, each verified by removing its call and confirming only
its own test reds. Fixing the first attempt turned up a test bug of my own: the
log line carries a queued=N field, so asserting "queued" in the message matched
every admission log ever emitted. It asserts the event name now.

Also covered two guards that were correct but unguarded: waiters whose futures
die out of band stop counting against the queue depth, and a newcomer cannot
barge past a parked waiter. The second is pinned as behaviour rather than as the
`if not self._waiters` check, because that check cannot actually change the
outcome: _take_slot_locked consults _can_admit_locked anyway, so either alone
refuses the newcomer. The test fails only if both go.

_optional_positive_int_env lost its last caller when the env parsing was
rewritten last round. Removed.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-27 04:35:17 -07:00 • committed by GitHub
commit 2ab4b744ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 2357 additions and 113 deletions
Download patch file Download diff file Expand all files Collapse all files

274
studio/backend/core/inference/llama_admission.py
View file

@ -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)),
)

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

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

973
studio/backend/tests/test_anthropic_admission.py Normal file
View file

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

11
studio/backend/tests/test_anthropic_messages.py
View file

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

262
studio/backend/tests/test_anthropic_passthrough_respawn.py Normal file
View file

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

537
studio/backend/tests/test_llama_admission.py
View file

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