Skip to content

Make the embedded server's model cache memory-aware - #268

Draft
C-K-Loan wants to merge 9 commits into
Blaizzy:mainfrom
C-K-Loan:pr/memory-aware-model-cache
Draft

Make the embedded server's model cache memory-aware#268
C-K-Loan wants to merge 9 commits into
Blaizzy:mainfrom
C-K-Loan:pr/memory-aware-model-cache

Conversation

@C-K-Loan

Copy link
Copy Markdown
Contributor

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 number decide() 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.
  • Patches the function 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 runs decide(), evicts whatever it says to, performs the real cold load, then measures the actual resulting size for SizeCache.
  • A subtler race this also closes: which model's response-generator/cache a given request is actually talking to is now tracked via ContextVars scoped per-request, instead of shared mutable globals that two concurrent requests could stomp on.
  • New endpoint, GET /v1/models/active-generationsTenantTable.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.swiftswitch_model now 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 seam ChatSwitchModelToolExecutor is tested against.

Tests: unit tests for decide()/TenantTable covering 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.py updated so the freeze/bundle step actually ships the new nativ_model_memory.py file 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

  • Full test suite green (252/252 Swift, 18/18 pure-Python)
  • Unit tests for decide() covering reject/evict/coexist and multi-eviction paths
  • Manual verification: switching to an oversized model is rejected without disturbing the active one
  • Manual verification: switch_model no longer interrupts a sibling generation in flight

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.

1 participant