Skip to content

feat(server): put the reranker and the embedder in the pool's budget - #144

Merged
magicnight merged 1 commit into
mainfrom
feat/auxiliary-residency
Oct 11, 2026
Merged

magicnight merged 1 commit into
mainfrom
feat/auxiliary-residency

Conversation

@magicnight

Copy link
Copy Markdown
Owner

Closes #130, the half #133 left open.

/v1/rerank and /v1/embeddings each keep one engine on the server, and neither went through ModelPool. The pool budgets and evicts chat models, but a loaded reranker or embedder was not in that account: the pool could fill its budget while a Qwen3-Reranker-8B held 16 GB next to it, nothing ever unloaded them (no idle sweep, no unload path, only replacement by the next one), and a GUI cold swap evicting under budget pressure never knew they existed.

What changes:

  • The pool gains auxiliary residents. An entry the pool budgets, evicts, sweeps and protects like an engine's, without holding the engine, under its own key (reranker:<id>, embedder:<id>) so a reranker and a chat model with the same id, which the library will happily offer both ways for a Qwen3 reranker, never share an entry. admitAuxiliary(kind, model, ttlSeconds:, inUse:, load:, release:) runs in the order load(_:) does: sweep idle entries, evict to fit, reserve the cost, run and drop each victim's unload, and only then the owner's load; on success the reservation becomes an entry whose unload is release, on failure it is rolled back. An admission can start in use; beginUsingAuxiliary is POOL-3 for them (it refreshes the LRU clock in the same step, so a hit needs no separate touch, and answers false when the pool no longer holds the entry) and endUsingAuxiliary ends it; forgetAuxiliary drops an entry without its release for an owner that already let go. unload(_:) now detaches before it awaits, so an admission that lands during a release is kept, and the load task consumes its victims' unloads one by one instead of holding them across the new allocation.
  • The server reports through AuxiliaryResidency. The app hands in PoolAuxiliaryResidency over its pool, which reads the model's TTL from the per-model parameters the way chat models' is read and talks to the pool directly, off the main actor; the CLI has no pool, passes nothing, and loads them directly as before. A cold request admits before it loads (the load runs inside the admission, after the pool evicted and reserved) and the admission returns with the model in use; a hit marks it in use in one step with the pool, or learns that the pool has already let the model go (an eviction or sweep, whether or not its release has reached the server) and loads it afresh rather than computing on an engine the pool no longer counts; the request ends its use when done; a swap forgets the old id before it drains and admits the new one; an admission the slot moved past (the pool's own release, or the stop) is forgotten rather than installed; and the server's stop lets both go once it has stopped serving: a request parked on the generation lock when stop() is called still runs, admits and installs, and the stop then forgets what that request installed, not the other way round.
  • A release that cannot keep the engine alive. The release the pool keeps carries only the admission's token, never the engine: the slot holds the last reference, so emptying it frees the weights and the drain that follows returns them. A release for a model the server has since swapped out does nothing (the token moved on), and a release that reaches the server before it has installed the engine (an unload of its key through the pool, or a stop that did not wait for the request, since an admission in use is neither evicted nor swept) is honoured: the request runs on the engine it loaded and nothing is installed, so the pool holds no entry for a model that is not there. A request that finds the pool has let its model go (an eviction or sweep, whether or not its release has reached the server) drains what it held and admits afresh instead of computing on a model the pool no longer counts; a swap drains the engine it replaces even when the pool's release emptied the slot, and drained to no effect, while the request was suspended checking that engine's model, since the drain keys on an engine having been resident when the call began, not on the slot (no test drives that ordering: the stand-in pool can run a release only inside beginUse, which a swap never reaches, so it holds by construction); a request that confirmed its engine before a stop that did not wait finishes on it and drains after itself.
  • The release never takes the generation lock. The pool runs a release while a handler holding the lock awaits its admission (admitting one slot can evict the other), so a release that waited for the lock would deadlock the server; the doc says so, and a gated smoke over a real pool whose budget holds one model drives exactly that path.
  • The Models tab no longer offers to pin a reranker or an embedder: the pin is the pool's, for chat models.
  • An unload path and a view. POST /x/models/unload takes {"model": id}: the resident reranker or embedder is dropped under the generation lock, the pool is told, and MLX's cache drained; the chat model as before when the id is the active model's or there is no body, and the check is made again under the lock; a 404 for anything else, where the door used to unload the chat model regardless, and a 400 for a body that is not JSON, where it used to be ignored. GET /x/status reports reranker and embedder.

Not changed: whether the two should share one slot when memory is tight, which #130 called a separate decision; the Models tab's loaded indicator, which still shows only the chat model.

Tests:

  • ModelPoolTests, fourteen new cases: the victim's unload runs before the auxiliary load; a release's captured object is freed before the next engine is built; admission sweeps idle entries; the idle sweep reclaims past the TTL; an in-use entry is spared until balanced; an admission in use is spared until ended; beginning the use of an entry the pool let go is refused; unload runs the release and forgetAuxiliary does not; unload leaves the resident set before its release runs; a failed load leaves no reservation; re-admission replaces without releasing even under a budget that would evict; a chat model and a reranker with the same id are separate entries; marking in use moves one off the LRU end; key scoping.
  • HummingbirdServerAuxiliaryResidencyTests (5) with a recording pool stand-in that also records whether the load ran inside the admission: a cold rerank and a cold embeddings request admit and load inside, a refusal is a 500 with the lock free and no load run, the door's 404, 400, bodiless and null paths, the two slots in /x/status, and a stop that forgets nothing it does not hold.
  • PoolAuxiliaryResidencyTests has six cases (the TTL from the parameters store, an admission in use until ended, a hit's use refreshing the clock, a hit's use sparing the entry, a hit's use refused once the pool let the model go, forget without the release).
  • Gated, real checkpoints (MACMLX_RUN_RERANK_SMOKE=1; ms-marco-MiniLM-L-6-v2, Qwen3-Reranker-0.6B-4bit, BAAI/bge-small-zh-v1.5 and BAAI/bge-base-zh-v1.5 through the real server): admit, the use begun and ended, a hit's use, a hit the pool refuses (it let the model go) loading the model afresh with the stale release then doing nothing, the same with the release arriving first and the readmission failing (the engine the request held is drained before it, so nothing stays in MLX's cache), release-then-readmit with MLX's cache under 64 MB after the release, a swap and the swap back with the same bound (the fix(server): release the resident reranker or embedder before loading its replacement, then drain MLX's cache #133 measurement, now with the pool in the loop), a stale release ignored, a swap into a model that fails to load forgetting the old one and leaving the slot empty, the door by id with the bound and a 404 on the second try, a release that beats the install served uninstalled and drained by the request that held it, the embedder's admit, hit, refused hit in both orderings (the release arriving first fails the readmission and the bound shows the drain), release-then-readmit, swap to the base model and back, stale release, door and release-at-once, an overlong input (truncated by the model) leaving the in-use marks balanced, the stop forgetting the embedder it still held and draining, a stop triggered while a request is parked on the generation lock (polled until the stop has begun): the request is served, and the stop forgets what it installed; and the server over a real ModelPool whose budget holds one of the two models, where admitting the reranker evicts the embedder through a release that re-enters the server while the handler holds the lock, bounded so a release waiting for the lock shows as a failure within the client's timeout (the test then leaves the stuck server behind rather than hang in its stop). All passed on this machine. The HF cache here holds bge-base-zh-v1.5's config and weights in two snapshots (fetched at different revisions), so the test's resolver links the two into one directory when that happens. Each residency smoke that asserts the 64 MB cache bound first checks that the checkpoint it would leave behind outweighs it (fp32 ms-marco is 91 MB, bge-small 96 MB), so a half-precision copy could not pass the bound vacuously.
  • Mutation run, thirty-eight mutants: no eviction on admission, detach ignoring auxiliaries, unload skipping the action, a failed load keeping its reservation, an admission in use not marked, no forget at the top of admit, no sweep on admit, victims unloaded after the owner's load, the load keeping the victims' actions alive, the entry removed only after the release; the residency not reading the TTL, its beginUse a no-op, admitting not in use; the reranker loading outside the pool, a hit not beginning the reranker's use, the door not telling the pool, a swap not forgetting, the reranker's use never ended, a release ignoring the token, a release keeping the reranker alive, an admission the pool already released being installed, a stop that lets nothing go, a hit not beginning the embedder's use, the orphan drain removed from the cross-encoder rerank handler, the orphan drain removed from the embeddings handler, the embedder's use ended only on success, an admission the slot moved past not forgotten, the embedder door branch gone, a swap not forgetting the old embedder, a release dropping whatever embedder is resident, the stop letting go before the shutdown, the pool reporting a missing entry as held, the residency reporting a model the pool let go as held, the reranker hit path ignoring a refused use, the embedder hit path ignoring one, every drain on the reranker's path keyed on an engine having been resident when the call began removed (the swap's, the refused hit's and the failed load's), the same on the embedder's, and, run on its own, the embedder's release taking the generation lock. Thirty-seven caught, the lock-taking release among them by the real-pool smoke (the request never completes; the first version of that smoke rethrew the client's timeout into a stop() that then waited for the stuck handler forever, which is how the mutant was first seen). Each of those two runs records one failed bound (the reranker's release-first, the embedder's swap-back: the first to fail); a probe printing MLX's active and cached memory at every bound under the same mutants showed the dropped engine freed each time (active memory back to its baseline of a few bytes) and its weights, 91 to 690 MB, in the cache at every later bound as well, so each bound discriminates on its own. The one that ends the embedder's use only when the work succeeded survives: no request can make embed fail on a loaded engine (it throws only without a model), so that branch is reached only by a fault inside MLX, and the end being unconditional is covered by review.

The server's two auxiliary engines were loaded beside the model pool and
never counted (#130): the pool could fill its budget while a 16 GB
reranker sat next to it, and nothing ever unloaded them.

The pool gains auxiliary residents: entries it budgets, evicts, sweeps and
protects like engines without holding the engine, keyed by kind
("reranker:<id>", "embedder:<id>") so a reranker and a chat model with the
same id stay apart, admitted through a load closure (evict to fit and
reserve first, then load) and reclaimed through the release closure their
owner registered. The server reports its reranker and embedder through an
AuxiliaryResidency the app builds over the pool: an admission returns with
the model in use, a hit marks it in use in one step with the pool (the
refresh of its idle clock included) or learns that the pool has already
let the model go and loads it afresh, and the request ends its use when
done; the server forgets what it swaps out or did not install, and lets
both go once it has stopped serving. A release carries the admission's token,
never the engine, so the pool holding it keeps no weights alive, and it
never takes the generation lock, since the pool runs it while a handler
holding the lock awaits an admission; a release for a model already
swapped out does nothing, and one that reaches the server before the
engine is installed empties the slot so the engine is not installed; a
request that finds its engine orphaned that way, or by a release landing
mid-request, drains MLX's cache after itself; a swap, or a hit the pool
refused, drains the engine it lets go of once it has, whether this call
or the pool's release emptied the slot meanwhile. The pool drops each
victim's unload as soon as it has run, for the same reason.
`/x/models/unload` takes a model id and drops a resident reranker or
embedder; `/x/status` shows both slots; the Models tab no longer offers
to pin a reranker or an embedder. The CLI has no pool and keeps loading
them directly.
@magicnight

Copy link
Copy Markdown
Owner Author

Local pipeline on d9b269b, the change rebased onto ea04cf2 after #140 and #141 merged (scripts/ci-local.sh, main's copy, CI_LOCAL_REF=d9b269b), macOS 27.0.1, Xcode 27.0 (ci.yml pins 26.4.1; Actions minutes are still out, so this is the gate):

job result evidence
website passed
spm passed MacMLXCore 1,050 Swift Testing + 335 XCTest (42 skipped); macmlx-cli 28
metal passed 1,343 passed, 0 failed, 42 skipped of 1,385
app passed 39 passed; app resolved the MLX fork at 1026d23

An earlier run, before the rebase, failed the app job on a protocol conformance in the app target (it builds with approachable concurrency, where an async closure type defaults to nonisolated(nonsending) and the protocol in MacMLXCore means @concurrent); the residency now lives in MacMLXCore as PoolAuxiliaryResidency, and the run above is on the final head.

Gated, same head: the five residency smokes in RerankEngineSmokeTests under xcodebuild with MACMLX_RUN_RERANK_SMOKE=1 (ms-marco-MiniLM-L-6-v2, Qwen3-Reranker-0.6B-4bit, BAAI/bge-small-zh-v1.5 and BAAI/bge-base-zh-v1.5 through the real server) passed, with MLX's cache under 64 MB after every release, swap-back, door, orphaned request, eviction through a real pool and stop, and the stop triggered while a request was parked on the lock serving that request first. Mutation run, same head: thirty-eight mutants, thirty-seven caught (the release that takes the generation lock by the real-pool smoke, which then records the request that never completes); the one that ends the embedder's use only on success survives because no request can make a loaded embedder fail (review covers it).

Cursor Bugbot has reported "usage limit reached" since 2026-10-06 and did not review this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Reranker and embedder engines are outside the model pool budget, and a swap keeps both models loaded

1 participant