Skip to content

server: bound ModelProvider's model cache with LRU eviction - #884

Open
mstkyvz wants to merge 1 commit into
Blaizzy:mainfrom
mstkyvz:fix/model-provider-lru
Open

server: bound ModelProvider's model cache with LRU eviction#884
mstkyvz wants to merge 1 commit into
Blaizzy:mainfrom
mstkyvz:fix/model-provider-lru

Conversation

@mstkyvz

@mstkyvz mstkyvz commented Aug 11, 2026

Copy link
Copy Markdown

Fixes #835.

ModelProvider cached every model it ever loaded in an unbounded dict and never evicted anything on its own. remove_model() exists but only the DELETE /v1/models endpoint calls it, so a server used with several voices kept each one resident until the process was restarted.

Change

The cache is now an OrderedDict used as an LRU:

  • a cache hit moves the entry to the most-recently-used end
  • after a miss loads a new model, the least recently used entries are dropped until the cache is within its limit
  • the limit comes from MLX_AUDIO_MAX_LOADED_MODELS, parsed the same way as the existing MLX_AUDIO_TTS_MAX_BATCH_SIZE (invalid values fall back to the default)
  • eviction and remove_model() both call mx.clear_cache(), matching how the generation loops already release memory

The default

I set DEFAULT_MAX_LOADED_MODELS = 2 so a mixed server can keep one TTS and one STT model hot, which seems like the common case for this server. That is the one judgement call in the patch — say the word and I will change it to 1, to a larger number, or to unbounded-unless-configured if you would rather not alter current behaviour by default.

Not included

load_model() is synchronous and does not take self.lock, so two concurrent requests for the same uncached model can still both load it. That predates this change and fixing it means making the method async and touching all six call sites, so I left it alone rather than widening the diff. Happy to do it separately if you want it.

Tests

Six tests in mlx_audio/tests/test_server.py cover cache hits not reloading, LRU eviction order, use refreshing recency, the env limit, invalid env falling back to the default, and the default staying bounded.

mlx_audio/tests/test_server.py     master: 11 failed, 17 passed
                                this branch: 11 failed, 23 passed

The 11 failures are identical on both and are unrelated to this change — they reproduce on a clean checkout in my environment (Python 3.14, macOS 26.5, M4 Pro).

ModelProvider kept every model it ever loaded in an unbounded dict. Nothing
called remove_model() automatically, so a server used with several voices
grew until the process was restarted.

Back the cache with an OrderedDict and evict the least recently used entry
once it exceeds MLX_AUDIO_MAX_LOADED_MODELS (default 2, which keeps a TTS
and an STT model resident at the same time). Eviction and the existing
DELETE /v1/models path both call mx.clear_cache() to release the freed
memory.
@tahazarif10

Copy link
Copy Markdown

One edge case here: this bounds cache entries but not necessarily resident models

The realtime STT websocket keeps stt_model alive for the lifetime of the connection. If that model gets evicted from the LRU while the socket is still open it is still referenced, so loading another model can exceed the intended resident-model limit

Might be worth pinning in-use models or adding a test for this case

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.

ModelProvider (mlx_audio.server) never evicts loaded models — unbounded memory growth with multiple voices

2 participants