Add opt-in idle model unloading to free VRAM#889
Open
michael-borck wants to merge 1 commit into
Open
Conversation
Models load on demand at generate time but are only unloaded explicitly (POST /models/unload) or at shutdown, so an idle instance pins its weights indefinitely -- ~5.6GB for qwen-tts-1.7B, most of an 8GB card. Adds an asyncio reaper that unloads after a period of inactivity, gated on VOICEBOX_IDLE_UNLOAD_SECONDS and off by default. Idleness is read from the serial generation queue's own state so a model is never pulled out from under queued or in-flight work; the next request reloads it transparently. Refs jamiepine#595
Contributor
📝 WalkthroughWalkthroughThe backend adds an optional idle reaper that tracks generation activity and unloads TTS, Whisper, and LLM models after a configured idle period. It starts during application startup and resets its timer after each generation cycle. ChangesIdle model unloading
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant AppStartup
participant GenerationQueue
participant IdleReaper
participant ModelBackends
AppStartup->>GenerationQueue: init_queue()
AppStartup->>IdleReaper: start()
GenerationQueue->>IdleReaper: mark_activity() after each work cycle
IdleReaper->>GenerationQueue: check queued and running generations
IdleReaper->>ModelBackends: unload TTS, Whisper, and LLM after timeout
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/services/idle_unload.py`:
- Around line 76-82: Update the backend-loaded guard in _reaper so it considers
all backends cleared by _unload_all, including TTS, Whisper/STT, and LLM, before
skipping an unload cycle. Either remove the guard or reuse the corresponding
backend loaded-state checks, while preserving the existing exception-safe reaper
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6f4b4425-a983-4b08-8fd9-998861fa4425
📒 Files selected for processing (3)
backend/app.pybackend/services/idle_unload.pybackend/services/task_queue.py
Comment on lines
+76
to
+82
| try: | ||
| from ..backends import get_tts_backend | ||
|
|
||
| if not get_tts_backend().is_loaded(): | ||
| continue # nothing loaded — nothing to do | ||
| except Exception: | ||
| continue |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find backend getter functions and is_loaded methods to confirm the fix is feasible.
rg -nP 'def get_.*_backend' backend/backends/ --type py -C2
rg -nP 'def is_loaded' backend/backends/ --type py -C2
rg -nP 'def unload_whisper_model|def unload_llm_model|def unload_tts_model' backend/services/ --type py -C3Repository: jamiepine/voicebox
Length of output: 7680
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== backend/services/idle_unload.py =="
nl -ba backend/services/idle_unload.py | sed -n '1,140p'
echo
echo "== backend/backends/__init__.py (relevant getters) =="
nl -ba backend/backends/__init__.py | sed -n '710,760p'
echo
echo "== backend/services/transcribe.py =="
nl -ba backend/services/transcribe.py | sed -n '1,80p'
echo
echo "== backend/services/llm.py =="
nl -ba backend/services/llm.py | sed -n '1,60p'
echo
echo "== backend/services/tts.py =="
nl -ba backend/services/tts.py | sed -n '1,60p'Repository: jamiepine/voicebox
Length of output: 234
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== backend/services/idle_unload.py =="
sed -n '1,140p' backend/services/idle_unload.py | cat -n
echo
echo "== backend/backends/__init__.py (relevant getters) =="
sed -n '710,760p' backend/backends/__init__.py | cat -n
echo
echo "== backend/services/transcribe.py =="
sed -n '1,80p' backend/services/transcribe.py | cat -n
echo
echo "== backend/services/llm.py =="
sed -n '1,60p' backend/services/llm.py | cat -n
echo
echo "== backend/services/tts.py =="
sed -n '1,60p' backend/services/tts.py | cat -nRepository: jamiepine/voicebox
Length of output: 8026
Check all loaded backends before skipping unload
_reaper() only looks at get_tts_backend().is_loaded(), but _unload_all() also clears Whisper and LLM. If TTS is idle while STT or LLM is loaded, the reaper never runs and those models stay resident. Remove this guard or make it cover all three backends.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/services/idle_unload.py` around lines 76 - 82, Update the
backend-loaded guard in _reaper so it considers all backends cleared by
_unload_all, including TTS, Whisper/STT, and LLM, before skipping an unload
cycle. Either remove the guard or reuse the corresponding backend loaded-state
checks, while preserving the existing exception-safe reaper behavior.
michael-borck
added a commit
to michael-borck/loco-puente
that referenced
this pull request
Jul 13, 2026
Upstream voicebox never unloads models once loaded, so an idle instance pins ~4.3GB of GPU 1 indefinitely (most of the 2060 Super's 8GB). Adds VoiceboxConfig.build_ref (owner/repo@ref) so the service can build from a fork carrying the idle-unload patch, and defaults the timeout to 600s. Pointed at michael-borck/voicebox@feat/idle-model-unload pending jamiepine/voicebox#889; drop build_ref from puente.yml once that merges. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
1 participant
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.
Closes #595 (the VRAM half of it — this adds the automatic trigger, not the manual load/unload UI button).
Problem
On a headless/server deployment there's no way for models to unload automatically. They load on demand at generate time, but the only things that ever unload them are an explicit
POST /models/unloador process shutdown (app.py). There's no idle path.So an instance that served one request an hour ago is still pinning its weights — ~5.6GB for qwen-tts-1.7B, which is most of an 8GB card. Anything else sharing that GPU is starved until you restart the container or remember to call the endpoint by hand.
Change
An asyncio reaper that unloads models after a period of inactivity.
VOICEBOX_IDLE_UNLOAD_SECONDS;0/unset disables it entirely, so existing behaviour is unchanged for anyone who doesn't opt in.tts/transcribe/llmunload calls already used on shutdown, with per-model failure isolation so one failure can't block the others._queued_generation_ids/_running_generation_tasks), not inferred externally, so a model is never pulled out from under a queued or running generation. A long generation continually re-arms the countdown.Three files: a new
backend/services/idle_unload.py, two lines inapp.pyto start it, and amark_activity()call in the generation worker'sfinallyblock (the one place every generation exits through, success/failure/cancel alike).Testing
Manual, on an RTX 2060 Super (8GB), Docker/CUDA, qwen-tts-1.7B:
Reclaims ~4.3GB of resident model weights. Verified separately that:
status=completed, no errors);I did not run Black across the touched files —
app.pyhas pre-existing formatting drift (theall_originsCORS block) and reformatting it would add unrelated churn to this diff. The new file and mytask_queue.pyedit are both Black-clean on their own.Disclosure: this patch was written with AI assistance (Claude), as was this description. It's tested end-to-end on my own build and the numbers above are real measurements from my hardware, not generated. If AI-assisted contributions aren't welcome here I understand — happy for this to be closed, and it stays a local patch on my build.
Summary by CodeRabbit