server: bound ModelProvider's model cache with LRU eviction - #884
Open
mstkyvz wants to merge 1 commit into
Open
Conversation
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.
|
One edge case here: this bounds cache entries but not necessarily resident models The realtime STT websocket keeps Might be worth pinning in-use models or adding a test for this case |
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.
Fixes #835.
ModelProvidercached every model it ever loaded in an unbounded dict and never evicted anything on its own.remove_model()exists but only theDELETE /v1/modelsendpoint calls it, so a server used with several voices kept each one resident until the process was restarted.Change
The cache is now an
OrderedDictused as an LRU:MLX_AUDIO_MAX_LOADED_MODELS, parsed the same way as the existingMLX_AUDIO_TTS_MAX_BATCH_SIZE(invalid values fall back to the default)remove_model()both callmx.clear_cache(), matching how the generation loops already release memoryThe default
I set
DEFAULT_MAX_LOADED_MODELS = 2so 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 takeself.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.pycover 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.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).