Repository navigation
feat(server): put the reranker and the embedder in the pool's budget - #144
Conversation
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.
|
Local pipeline on d9b269b, the change rebased onto ea04cf2 after #140 and #141 merged (
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 Gated, same head: the five residency smokes in Cursor Bugbot has reported "usage limit reached" since 2026-10-06 and did not review this PR. |
Closes #130, the half #133 left open.
/v1/rerankand/v1/embeddingseach keep one engine on the server, and neither went throughModelPool. 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:
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 orderload(_:)does: sweep idle entries, evict to fit, reserve the cost, run and drop each victim's unload, and only then the owner'sload; on success the reservation becomes an entry whose unload isrelease, on failure it is rolled back. An admission can start in use;beginUsingAuxiliaryis 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) andendUsingAuxiliaryends it;forgetAuxiliarydrops 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.AuxiliaryResidency. The app hands inPoolAuxiliaryResidencyover 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 whenstop()is called still runs, admits and installs, and the stop then forgets what that request installed, not the other way round.unloadof 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 insidebeginUse, 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.POST /x/models/unloadtakes{"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/statusreportsrerankerandembedder.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;unloadruns the release andforgetAuxiliarydoes not;unloadleaves 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 andnullpaths, the two slots in/x/status, and a stop that forgets nothing it does not hold.PoolAuxiliaryResidencyTestshas 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).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 realModelPoolwhose 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.beginUsea 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 astop()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 makeembedfail 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.