diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py index db9a5d8ce4..7bf0dd7429 100644 --- a/studio/backend/core/inference/llama_admission.py +++ b/studio/backend/core/inference/llama_admission.py @@ -58,6 +58,80 @@ DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16 DEFAULT_ADMISSION_MIN_QUEUE = 64 +def _executor_workers() -> int: + """Threads asyncio's default executor runs to_thread work on. + + Mirrors ThreadPoolExecutor's own default sizing, which is what + ``run_in_executor(None, ...)`` builds. 3.13 sizes it from + ``process_cpu_count()``, which honours CPU affinity and cgroup quotas; + ``cpu_count()`` would budget from the whole host inside a one-core container. + """ + cpus = getattr(os, "process_cpu_count", os.cpu_count)() or 1 + return min(32, cpus + 4) + + +def _executor_reserve(workers: int) -> int: + """Threads kept clear of parked approvals, for generation steps, stream + teardown and unrelated to_thread work. Scaled rather than flat: a flat count + would leave a 5-worker executor (one usable CPU) no budget at all. + """ + return max(2, workers // 8) + + +def _max_parked(capacity: int) -> int: + """How many holders may sit on an approval prompt with their slot given back. + + A pending prompt parks an executor thread (the loop blocks inside + to_thread(next, gen)) whether or not it parked its slot, the pool already + permits `capacity` of those, and every park admits one more, so budget only + what the executor has left over. Zero on a backend whose --parallel alone + fills it: the prompt then holds its slot, as it did before parking existed. + """ + workers = _executor_workers() + spare = workers - _executor_reserve(workers) - max(0, capacity) + # A quarter of the executor, floored at two while `spare` allows: a quarter of + # five is one, and one park cannot cover the two simultaneous prompts #7455 + # exists for. + return max(0, min(max(2, workers // 4), spare)) + + +# Process-wide, not per queue: there is one executor, and base_url takes a fresh +# port on every load, so a per-queue budget would hand the same allowance to each +# backend and to every reload, blind to the approvals parked on the old queue. +_PARK_LOCK = threading.Lock() +_parked_total = 0 + + +def _claim_park(limit: int) -> bool: + global _parked_total + with _PARK_LOCK: + if _parked_total >= limit: + return False + _parked_total += 1 + return True + + +def _drop_park() -> None: + global _parked_total + with _PARK_LOCK: + _parked_total = max(0, _parked_total - 1) + + +def _live_capacity(current: "LlamaAdmissionQueue") -> int: + """Slots across every backend still serving requests. + + One queue's capacity is the wrong denominator for a budget sized against the + one executor: a reload drains the old queue alongside the new one, and + prompts on both park threads. Idle queues hold nothing and are about to be + evicted. + """ + with _QUEUES_LOCK: + queues = list(_QUEUES.values()) + # is_idle takes each queue's own lock, so never while holding _QUEUES_LOCK. + total = sum(queue._capacity for queue in queues if queue is current or not queue.is_idle()) + return total if any(queue is current for queue in queues) else total + current._capacity + + @dataclass(frozen = True, **_SLOTS) class LlamaAdmissionConfig: enabled: bool = DEFAULT_ADMISSION_ENABLED @@ -214,7 +288,7 @@ class _Waiter: class LlamaAdmissionLease: - __slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked") + __slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked", "_budgeted") def __init__( self, @@ -226,27 +300,52 @@ class LlamaAdmissionLease: self._released = False self._release_lock = threading.Lock() self._parked = False + self._budgeted = False @property def slot(self) -> Optional[int]: """Pool slot this lease holds, or None when admission is disabled.""" return self._slot - def park(self) -> None: + def park(self) -> bool: """Hand the slot back while this holder waits on something off the GPU. A run stopped on a tool approval prompt is not decoding, so holding its slot would let unanswered prompts fill the pool while llama-server idles. The lease itself stays valid: releasing it after a park is still correct. + + False when the park budget is spent and nothing was given back: the + caller keeps its slot across the prompt, as it did before parking + existed. Slower for whoever is behind it, but each freed slot admits + another run that can park too, on the executor the generators run on. """ queue = self._queue - slot = None with self._release_lock: if queue is None or self._released or self._parked: - return + return False + # Under the lease lock so the decision and the handover cannot split. + # Nothing takes the queue lock then a lease lock, so this order is + # the only one in play. + if not queue.try_park(self._slot): + return False self._parked = True - slot, self._slot = self._slot, None - queue.park(slot) + self._budgeted = True + self._slot = None + return True + + def _drop_budget(self) -> None: + """Give the executor budget back now the prompt wait is over. + + Separate from the queue's parked count, which lasts until the slot is + back: the executor thread is free the moment the answer arrives. Holding + the budget until the resume lands would refuse someone else's park for a + finished wait, and that someone holds the slot the resumer wants. + """ + with self._release_lock: + if not self._budgeted: + return + self._budgeted = False + _drop_park() def unpark(self) -> None: """Drop the parked state without reclaiming a slot. @@ -259,6 +358,7 @@ class LlamaAdmissionLease: if not self._parked: return self._parked = False + self._drop_budget() if self._queue is not None: self._queue.unpark() @@ -278,6 +378,9 @@ class LlamaAdmissionLease: queue = self._queue if queue is None or not self._parked: return + # Before the wait, not after: the prompt is answered, so this holder is + # already off the executor and must not keep anyone else off it. + self._drop_budget() slot = await queue.acquire_parked_slot(cancel_event = cancel_event, poll_s = poll_s) stranded = None with self._release_lock: @@ -304,6 +407,7 @@ class LlamaAdmissionLease: self._released = True queue = self._queue parked, self._parked = self._parked, False + self._drop_budget() if queue is not None: if parked: queue.unpark() @@ -513,12 +617,20 @@ class LlamaAdmissionQueue: self._release_slot_locked(slot) self._grant_waiters_locked() - def park(self, slot: Optional[int]) -> None: - """Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``.""" + def try_park(self, slot: Optional[int]) -> bool: + """Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``. + + False leaves the slot with its holder, so a refused park costs nothing to + undo. The per-queue count is only what ``is_idle`` reads; the budget and + the capacity it is sized from are both process-wide. + """ + if not _claim_park(_max_parked(_live_capacity(self))): + return False with self._lock: self._parked += 1 self._release_slot_locked(slot) self._grant_waiters_locked() + return True def unpark(self) -> None: with self._lock: @@ -684,5 +796,10 @@ def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: def reset_llama_admission_queues() -> None: + global _parked_total with _QUEUES_LOCK: _QUEUES.clear() + # The budget outlives the queues it was claimed against, so dropping them + # without it leaks the count and shrinks the budget for good. + with _PARK_LOCK: + _parked_total = 0 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8b15779a50..53b4136e32 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -9413,7 +9413,10 @@ async def openai_chat_completions( if lease is None: return if on: - lease.park() + # Refused when the budget is spent: the slot stays here, + # so there is nothing to take back afterwards. + if not lease.park(): + return elif wait: # Resuming: park() may have handed our slot to a waiter, so wait for room instead # of putting two holders on one slot. diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py index 9ff19ec27d..1b1aeb1cc5 100644 --- a/studio/backend/tests/test_llama_admission.py +++ b/studio/backend/tests/test_llama_admission.py @@ -1066,3 +1066,229 @@ def test_an_immediate_arrival_cannot_take_an_approved_chats_slot(): assert queue.snapshot().active <= 1 asyncio.run(scenario()) + + +def test_parking_is_bounded_so_the_thread_pool_cannot_be_drained(monkeypatch): + # A pending prompt parks an executor thread (the loop blocks inside + # to_thread(next, gen)) and frees a slot that admits another run which can + # park too, so unbounded parking drains the pool the generators run on. + # Pinned because the real budget follows the runner's usable CPUs. + monkeypatch.setattr(llama_admission, "_executor_workers", lambda: 32) + + async def scenario(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + limit = llama_admission._max_parked(1) + assert limit >= 1 + + leases = [] + for _ in range(limit): + lease = queue.reserve(capacity = 1, config = config).lease_nowait() + assert lease is not None and lease.park() + leases.append(lease) + + refused = queue.reserve(capacity = 1, config = config).lease_nowait() + assert refused is not None + assert not refused.park(), "parking is unbounded" + # Refusing means keeping the slot, the old behaviour, not an error. + assert refused.slot is not None + assert queue.snapshot().active == 1 + + leases[0].unpark() + assert refused.park(), "budget was not returned" + for lease in leases[1:] + [refused]: + lease.release() + leases[0].release() + + asyncio.run(scenario()) + + +def test_the_park_budget_is_shared_by_every_queue(monkeypatch): + # One executor, so a per-queue budget would be handed out again to every + # backend and to every reload onto a fresh ephemeral port. + monkeypatch.setattr(llama_admission, "_executor_workers", lambda: 32) + + async def scenario(): + config = LlamaAdmissionConfig() + first = get_llama_admission_queue("http://llama.test:1") + second = get_llama_admission_queue("http://llama.test:2") + limit = llama_admission._max_parked(1) + + for index in range(limit): + queue = first if index % 2 == 0 else second + lease = queue.reserve(capacity = 1, config = config).lease_nowait() + assert lease.park() + + spare = second.reserve(capacity = 1, config = config).lease_nowait() + assert not spare.park(), "each queue got its own budget" + + # A reset drops the queues the count was claimed against, so it must drop + # the count too or the leak shrinks the budget process-wide. + reset_llama_admission_queues() + revived = get_llama_admission_queue("http://llama.test:1") + fresh = revived.reserve(capacity = 1, config = config).lease_nowait() + assert fresh.park(), "reset leaked the park count" + fresh.release() + + asyncio.run(scenario()) + + +def test_the_park_budget_leaves_the_executor_room_to_work(monkeypatch): + # The pool already permits `capacity` pending prompts and every park admits + # one more, so the budget must account for both. Swept across executor sizes + # rather than read off this host, since a container gets a small one. + for cpus in (1, 2, 4, 8, 16, 28, 64): + workers = min(32, cpus + 4) + monkeypatch.setattr(llama_admission, "_executor_workers", lambda w = workers: w) + reserve = llama_admission._executor_reserve(workers) + assert reserve >= 2, f"{workers} workers left no reserve" + + # Even the smallest executor fits the two simultaneous prompts #7455 needs. + assert llama_admission._max_parked(1) >= 2, f"no room for two on {workers} workers" + assert llama_admission._max_parked(1) <= workers // 2 + # A backend whose --parallel alone fills the executor gets no parks. + assert llama_admission._max_parked(workers) == 0 + for capacity in range(0, workers + 8): + budget = llama_admission._max_parked(capacity) + assert budget >= 0, f"negative budget at capacity {capacity}" + assert ( + budget == 0 or capacity + budget <= workers - reserve + ), f"{workers} workers: capacity {capacity} plus {budget} parks leaves no room" + + +def test_the_park_budget_follows_the_executors_own_cpu_count(monkeypatch): + # 3.13 sizes ThreadPoolExecutor from process_cpu_count(), which honours CPU + # affinity and cgroup quotas; cpu_count() would budget from the whole host + # inside a one-core container. Pulled apart here, since they usually match. + import concurrent.futures + + monkeypatch.setattr(os, "cpu_count", lambda: 64) + if hasattr(os, "process_cpu_count"): + monkeypatch.setattr(os, "process_cpu_count", lambda: 1) + # Against the real thing rather than the formula: the default executor is a + # plain ThreadPoolExecutor(), so its own sizing is the answer on any version. + with concurrent.futures.ThreadPoolExecutor() as pool: + assert llama_admission._executor_workers() == pool._max_workers + + +def test_the_stream_retries_a_park_that_was_refused(): + # _park_admission short-circuits on `on == _parked`, so recording a refused + # park as parked would skip every later approval in the run even once the + # budget frees up. Structural because that only shows on a second approval. + import ast + + # Read rather than import: routes.inference pulls in the whole app. + route = os.path.join(_backend, "routes", "inference.py") + with open(route, encoding = "utf-8") as handle: + tree = ast.parse(handle.read()) + helpers = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.AsyncFunctionDef) and node.name == "_park_admission" + ] + assert len(helpers) == 1, f"expected one _park_admission, found {len(helpers)}" + + guards = [ + node + for node in ast.walk(helpers[0]) + if isinstance(node, ast.If) + and isinstance(node.test, ast.UnaryOp) + and isinstance(node.test.op, ast.Not) + and isinstance(node.test.operand, ast.Call) + and getattr(node.test.operand.func, "attr", None) == "park" + and getattr(node.test.operand.func.value, "id", None) == "lease" + ] + assert len(guards) == 1, "lease.park()'s answer is ignored" + assert all( + isinstance(stmt, ast.Return) for stmt in guards[0].body + ), "a refused park must leave _parked alone, so a later approval retries it" + + +def test_the_park_budget_counts_every_live_backend(monkeypatch): + # base_url takes a fresh port on every load, so a reload mints a queue while + # the old one drains. Prompts on both park threads of the one executor, so a + # budget sized from either backend alone lets them add up past the reserve. + monkeypatch.setattr(llama_admission, "_executor_workers", lambda: 32) + + async def scenario(): + config = LlamaAdmissionConfig() + old = get_llama_admission_queue("http://llama.test:1") + draining = old.reserve(capacity = 16, config = config).lease_nowait() + assert draining is not None # in flight, so the registry keeps this queue + + new = get_llama_admission_queue("http://llama.test:2") + lease = new.reserve(capacity = 16, config = config).lease_nowait() + assert lease is not None + + # 16 slots each against 32 workers: their prompts alone can fill it. + assert llama_admission._max_parked(16) > 0, "this test needs a budget to remove" + assert not lease.park(), "budget sized from one backend of two" + + draining.release() # the old backend drains and is up for eviction + assert lease.park(), "an idle backend still counted against the budget" + lease.release() + + asyncio.run(scenario()) + + +def test_the_park_budget_is_freed_when_the_prompt_is_answered(monkeypatch): + # The executor thread comes back the moment the answer arrives, before the + # resume queues for a slot. Holding the budget until the slot lands refuses + # someone else's park, and that someone holds the slot the resumer wants. + monkeypatch.setattr(llama_admission, "_executor_workers", lambda: 32) + + async def scenario(): + config = LlamaAdmissionConfig() + queue = get_llama_admission_queue("http://llama.test") + + parked = [] + for _ in range(llama_admission._max_parked(1)): + lease = queue.reserve(capacity = 1, config = config).lease_nowait() + assert lease is not None and lease.park() + parked.append(lease) + + blocked = queue.reserve(capacity = 1, config = config).lease_nowait() + assert blocked is not None + assert not blocked.park(), "the budget was not full to begin with" + + # One prompt is answered. Its slot is taken, so the resume queues for one. + resumed = asyncio.ensure_future(parked[0].unpark_async(poll_s = 0.01)) + await asyncio.sleep(0.05) + assert not resumed.done(), "the resume needs to still be waiting for its slot" + + assert blocked.park(), "budget held for a prompt wait that is over" + # Which is what frees the slot the resumer was waiting for. + await asyncio.wait_for(resumed, timeout = 2) + for lease in parked[1:] + [blocked]: + lease.release() + parked[0].release() + + asyncio.run(scenario()) + + +def test_releasing_a_parked_holder_returns_its_budget(monkeypatch): + # A client that disconnects on the prompt releases straight out of parked, + # never unparking. Its executor thread went with it, so keeping the budget + # would lose one for the life of the process. + monkeypatch.setattr(llama_admission, "_executor_workers", lambda: 32) + + async def scenario(): + config = LlamaAdmissionConfig() + queue = get_llama_admission_queue("http://llama.test") + + parked = [] + for _ in range(llama_admission._max_parked(1)): + lease = queue.reserve(capacity = 1, config = config).lease_nowait() + assert lease is not None and lease.park() + parked.append(lease) + + blocked = queue.reserve(capacity = 1, config = config).lease_nowait() + assert blocked is not None + assert not blocked.park(), "the budget was not full to begin with" + + parked[0].release() + assert blocked.park(), "a released park never gave its budget back" + for lease in parked[1:] + [blocked]: + lease.release() + + asyncio.run(scenario())