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

Studio: bound how many tool approvals may park their slot (#7496)

Browse source 
* Bound how many approvals may park, against the executor

#7455 landed parking, which is the right shape and supersedes what this branch
was carrying. It is unbounded, though, and the thing it is unbounded against is
not the GPU.

A run stopped on an approval prompt is blocked inside the to_thread(next, gen)
call that drives it, so it holds one of asyncio's default min(32, cpu + 4)
executor threads until the user answers. The slot cap used to bound that.
Parking hands the slot back, which admits another run that can park too, so the
ceiling became the wait line: 64 deep on a 1-slot backend. Long before that, the
executor is full and nothing else in the backend runs, including generation
steps for chats that already hold slots and the stream teardown that would clean
up after a disconnect.

The pool already permits `capacity` pending prompts, and each park adds one
more, so the budget is what the executor has left after the cap and a reserve of
4. On this machine (32 workers) --parallel 4 gets 8 parks and 20 free threads,
--parallel 24 gets 4 and 4, and --parallel 28 or higher gets none: there the
prompt keeps its slot and behaves exactly as it did before parking existed.

Counted process-wide rather than per queue. There is one executor, but a
per-queue budget is the same allowance again for every backend, and base_url
carries a fresh port on every model load, so a reload would mint a queue that
knows nothing about the approvals still parked on the old one. A reset clears it
too, or a leaked claim shrinks the budget for the life of the process.

park() reports whether it took the budget, and a refusal costs nothing to undo
because the slot never left its holder. The stream reads that answer rather than
recording a refused park as parked, which would make it skip the park for every
later approval in the same run even once the budget freed up.

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

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

* Size the park budget from the executor's own CPU count

Two review findings, both real.

The budget read os.cpu_count(). 3.13 sizes ThreadPoolExecutor from
os.process_cpu_count(), which honours CPU affinity and cgroup quotas, and
asyncio's default executor is a plain ThreadPoolExecutor(), so a container
pinned to one core on a 64-core host got a 5-thread executor and a budget
computed from 64. The bound was then looser than no bound at all in exactly the
environment that can least afford it. It asks the same source the executor does,
and the test compares against a real ThreadPoolExecutor rather than restating
the formula, so it stays right on 3.12 as well.

The reserve was a flat 4, which on that same 5-thread executor left nothing to
budget and turned parking off entirely. Small hosts are where a chat most needs
to keep moving while another sits on a prompt. It scales now, and the ceiling
has a floor of two: a quarter of five is one, and one park cannot cover two
chats on prompts at once, which is what #7455's own two-approvals test needs.
Without that floor, that test fails on a one or two CPU runner. `spare` still
takes the budget to zero when the pool already fills the executor, so nothing
about a 32-worker machine changes: --parallel 4 still gets 8 parks, 24 gets 4,
28 gets none.

The two behavioural budget tests pin the worker count rather than reading it off
the runner, and the property test sweeps executor sizes from one CPU to 64
instead of asserting against whatever the host happens to have. The whole suite
passes with the CPU count faked to 1, 2 and 4, which is how both of these were
reproduced.

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

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

* Size the park budget from every live backend, and free it on the answer

Two review findings, both real.

The budget was global but sized from one queue's capacity. A reload mints a
queue on a new port while the old one drains, so both are live, and prompts on
both park executor threads. Eight parks on an old 1-slot queue plus a new
24-slot backend is 32 threads on a 32-thread executor, with the new backend's
prompts refused and holding their slots, which is the state the reserve exists
to prevent. It sums the capacity of every backend still serving instead. Idle
queues are skipped: those are the ones the registry is about to evict, and they
are holding nothing.

The budget also outlived the wait it was paying for. unpark_async only dropped
it after reacquiring a slot, but the generator yields its post-approval event
first, so the executor thread is already back in the pool while the resume
queues. An approved chat waiting on a slot would refuse a different chat's park,
and that chat then keeps the slot the resumer is waiting for, so an unanswered
prompt strands chats that were already approved. The budget is released when the
prompt wait ends now, and the queue's parked count still runs until the slot is
back, which is what guards idle eviction and the resume ordering.

Both are separate counters on the lease as a result, and every exit from a park
drops the budget: unpark, unpark_async and release. That last one was the mutant
that came back missed, since a client disconnecting on a prompt releases
straight out of parked and would otherwise lose a budget slot for good.

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

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

* Tighten the comments on the park budget

---------

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-28 14:49:21 -07:00 • committed by GitHub
commit 5fe457ad01
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 355 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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