Make the embedded server's model cache memory-aware - #268
Draft
C-K-Loan wants to merge 9 commits into
Draft
Conversation
C-K-Loan
force-pushed
the
pr/memory-aware-model-cache
branch
from
August 11, 2026 05:03
f897b60 to
d615a29
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Depends on #267 and #263 (MCP consent adapter, spawn_agent) — this branch is stacked on top of them, so the diff includes their commits until those merge.
The embedded server previously held exactly one text-generation model at a time with no memory accounting, so a model swap could silently evict whatever was currently generating against it — a real risk once multiple concurrent requests (coordinator plus sub-agents) can be in flight against different models at once. This adds a memory-aware cache that rejects a load cleanly with a clear error when there genuinely isn't room, prefers keeping models resident over evicting when there's space for both, and evicts the least-recently-used model only when actually necessary — never one that's mid-generation.
Eviction is a fallback, not the default: a second model only evicts anything if it genuinely doesn't fit in the current budget; otherwise both stay resident.
How it works
The core decision —
nativ_model_memory.py(new file, pure Python, no I/O):decide()— given a requested model's size, the currently resident models, and the available memory budget: if it fits, load with no eviction. If not, evict the minimum number of idle (zero active generations) resident models, least-recently-released first, until it fits. If it still doesn't fit after evicting everything evictable, reject cleanly with a specific error instead of failing generically.TenantTable— refcounts active generations per resident model (acquire/release/active_count/any_active). This is what makes "never evict a model mid-generation" an enforced invariant rather than a hope.SizeCache— a model's real measured size (actual memory delta after its first load) is cached and reused afterward; before that first load, falls back to a pessimistic on-disk size estimate.Wiring into the real server —
nativ_server.py:_text_generation_memory_budget_bytes()— the available-memory numberdecide()checks against, measured fresh each time: Apple's own recommended GPU working-set ceiling for the current device, minus whatever MLX actually has active right now. Live per-device measurement, not a static config value.mlx_vlm's server calls to fetch/load a model: a pool hit (the common case) acquires a tenant and returns immediately, no change in cost. A pool miss runsdecide(), evicts whatever it says to, performs the real cold load, then measures the actual resulting size forSizeCache.ContextVars scoped per-request, instead of shared mutable globals that two concurrent requests could stomp on.GET /v1/models/active-generations→TenantTable.any_active().switch_model's server restart (its only mechanism) now checks this and refuses to restart while any generation is active anywhere — a restart kills every resident model's generations at once, so it has to check globally, not just the model being switched away from.Swift side:
ChatSwitchModelTool.swift—switch_modelnow does a pre-flight check against that new endpoint before attempting anything, and surfaces the new insufficient-memory rejection as a specific, clean error instead of a generic failure.NativModel.swift/NativChatClient.swift— the client-side plumbing for the pre-flight check (hasActiveTextGenerations(),preloadMemoryWarning), which exist purely to satisfy the protocol seamChatSwitchModelToolExecutoris tested against.Tests: unit tests for
decide()/TenantTablecovering reject/evict/coexist and multi-eviction paths, integration tests against the real server wiring, and Swift-side tests for the new switch_model edge cases (memory rejection, active-generation guard).Build script:
build_mlx_vlm_server.pyupdated so the freeze/bundle step actually ships the newnativ_model_memory.pyfile inside the app's embedded Python distribution.Known limitation
The fit check is based on a model's measured/estimated weight size only — it doesn't reserve headroom for that model's own KV-cache growth during generation (which scales with context length and concurrent in-flight requests). A model can pass the fits-check and load, then still pressure memory once generation is underway, since the eviction decision only runs at load time and isn't re-consulted afterward. Not fixed here — flagging as a real, currently open gap rather than a solved case.
Test plan
decide()covering reject/evict/coexist and multi-eviction pathsswitch_modelno longer interrupts a sibling generation in flight