# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. import asyncio import os import sys import threading import pytest _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) from core.inference import llama_admission from core.inference.llama_admission import ( ADMISSION_CONTROL_ENV, ADMISSION_KEEPALIVE_INTERVAL_ENV, ADMISSION_MAX_QUEUE_ENV, ADMISSION_QUEUE_PER_SLOT_ENV, ADMISSION_QUEUE_TIMEOUT_ENV, DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, DEFAULT_ADMISSION_MAX_QUEUE, DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, LlamaAdmissionConfig, LlamaAdmissionQueueFull, get_llama_admission_queue, llama_admission_config_from_env, reset_llama_admission_queues, ) _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(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() def test_admission_config_defaults(monkeypatch): for name in ( ADMISSION_CONTROL_ENV, 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 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): monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") config = llama_admission_config_from_env() assert config.enabled is False assert config.queue_timeout_s is None assert config.keepalive_interval_s == 0.25 assert config.max_queue is None def test_admission_config_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") config = llama_admission_config_from_env() assert config.queue_timeout_s == 600.0 def test_fifo_capacity_one_grants_next_waiter_on_release(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) third = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() assert first_lease is not None assert second.lease_nowait() is None assert third.lease_nowait() is None assert queue.snapshot().queued == 2 first_lease.release() second_lease = await second.wait(0.1) assert second_lease is not None assert third.lease_nowait() is None second_lease.release() third_lease = await third.wait(0.1) assert third_lease is not None third_lease.release() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_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") config = LlamaAdmissionConfig(max_queue = 1) first = queue.reserve(capacity = 1, config = config) queued = queue.reserve(capacity = 1, config = config) assert first.lease_nowait() is not None assert queued.lease_nowait() is None with pytest.raises(LlamaAdmissionQueueFull): queue.reserve(capacity = 1, config = config) asyncio.run(_run()) def test_disabled_admission_bypasses_active_slot_limit(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig(enabled = False) first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) assert first.lease_nowait() is not None assert second.lease_nowait() is not None assert queue.snapshot().active == 0 assert queue.snapshot().queued == 0 asyncio.run(_run()) def test_cancelling_promoted_waiter_releases_slot(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() first_lease.release() await asyncio.sleep(0) second.cancel() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_cancelling_promoted_waiter_before_delivery_releases_slot(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() first_lease.release() second.cancel() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_external_waiter_future_cancel_invalidates_reservation(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() assert first_lease is not None assert second._waiter is not None second._waiter.future.cancel() assert second.lease_nowait() is None assert second.is_cancelled is True assert await second.wait(0.01) is None first_lease.release() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() assert first_lease is not None assert second._waiter is not None wait_task = asyncio.create_task(second.wait(1.0)) await asyncio.sleep(0) second._waiter.future.cancel() assert await asyncio.wait_for(wait_task, timeout = 0.1) is None assert second.is_cancelled is True first_lease.release() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_capacity_increase_promotes_existing_waiter_fifo(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() first = queue.reserve(capacity = 1, config = config) second = queue.reserve(capacity = 1, config = config) first_lease = first.lease_nowait() assert first_lease is not None assert second.lease_nowait() is None assert queue.snapshot().active == 1 assert queue.snapshot().queued == 1 third = queue.reserve(capacity = 2, config = config) second_lease = await second.wait(0.1) assert second_lease is not None assert third.lease_nowait() is None snapshot = queue.snapshot() assert snapshot.capacity == 2 assert snapshot.active == 2 assert snapshot.queued == 1 first_lease.release() third_lease = await third.wait(0.1) assert third_lease is not None second_lease.release() third_lease.release() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_lease_release_is_idempotent_under_concurrent_calls(): async def _run(): queue = get_llama_admission_queue("http://llama.test") config = LlamaAdmissionConfig() reservation = queue.reserve(capacity = 1, config = config) lease = reservation.lease_nowait() assert lease is not None threads = [threading.Thread(target = lease.release) for _ in range(16)] for thread in threads: thread.start() for thread in threads: thread.join() snapshot = queue.snapshot() assert snapshot.active == 0 assert snapshot.queued == 0 asyncio.run(_run()) def test_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. get_llama_admission_queue("http://127.0.0.1:1001") get_llama_admission_queue("http://127.0.0.1:1002") assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} get_llama_admission_queue("http://127.0.0.1:1003") assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} def test_new_key_retains_in_flight_prior_load_queue(): config = LlamaAdmissionConfig() busy = get_llama_admission_queue("http://127.0.0.1:2001") async def _run(): reservation = busy.reserve(capacity = 1, config = config) lease = reservation.lease_nowait() assert lease is not None # A new load must not drop a queue that still has an in-flight request. get_llama_admission_queue("http://127.0.0.1:2002") assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} # Once it drains, the next load reclaims it. lease.release() get_llama_admission_queue("http://127.0.0.1:2003") assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} asyncio.run(_run()) 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())