fix(tts): stop stranding the TTS model on CPU after a dub abort (#1191) - #1204
Merged
Conversation
`offload_tts_for_asr()` moves the TTS model to CPU to make VRAM room for WhisperX, but its partner `restore_tts_after_asr()` was only reachable on the dub-transcribe success path. Any abort, terminal error, or client disconnect skipped it, and `get_model()` never re-checked placement — so EVERY subsequent /generate ran on CPU (10-50x slower, CPU pegged) until the ~15-minute idle unload happened to fire. Reported as "speed varies by time of day"; it is fully deterministic. Two independent guarantees: - Balance the pair at the call site: gen()'s `finally` now pays the restore debt on every exit path, chained off the ASR unload so the two never contend for VRAM (and fire-and-forget, since the finally also runs under GeneratorExit where awaiting is illegal). - Self-heal placement (the class fix): `get_model()` verifies the model is on the resolved target device and moves it back if not, so a future unbalanced offload path cannot strand it either. Cheapest-first probe — one parameter check on the hot path; unified memory is exempt (its offload releases the model rather than moving it). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
|
| Filename | Overview |
|---|---|
| backend/api/routers/dub_core.py | Adds _tts_offloaded sentinel and _restore_tts_bg helper; gen()'s finally now pays the TTS-restore debt on every exit path (abort, error, disconnect) by chaining the restore off the ASR unload callback to preserve VRAM ordering. |
| backend/services/model_manager.py | Adds placement self-heal: _first_param_device, _stranded_tts_target, ensure_tts_on_device, and _heal_tts_placement are new; get_model() calls _heal_tts_placement() on every hot-path return, with GPU-pool-thread deadlock prevention and correct 1-worker mutual exclusion reasoning. |
| tests/test_tts_placement_self_heal.py | 12 new regression tests covering: probe hot/cold/unified-memory paths, self-heal success/noop/OOM-failure, get_model() heal-before-return, GPU-pool-thread deadlock prevention, and balanced-pair end-to-end through abort/crash/success SSE paths. |
| docs/performance.md | Adds item #5 to the "classic causes of slowness" list, explaining the v0.3.23 fix and advising users on older builds to restart the backend. |
| CHANGELOG.md | Single one-liner under ### Fixed referencing #1191, consistent with house-rule changelog style. |
Reviews (4): Last reviewed commit: "fix: strip leftover conflict markers fro..." | Re-trigger Greptile
CodeQL flagged the `return` inside the finally block ('break'/'return'
in finally swallows in-flight exceptions). The returns lived in a nested
def so nothing was actually swallowed, but the pattern is worth avoiding
outright: the dispatch helpers now live at endpoint scope and the finally
holds straight-line control flow only.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughThe change restores offloaded TTS models on every transcription-stream exit path, adds placement self-healing before subsequent model use, and adds regression tests for device detection, recovery, and successful or failed SSE flows. ChangesTTS placement recovery
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 7 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (7 passed)
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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
backend/services/model_manager.py (1)
1285-1297: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winbackend/services/model_manager.py:1285-1297 —
get_model()rereads the global afterawait _heal_tts_placement(), soidle_worker()can unload the model in between and this branch returnsNonedespite starting with a live instance. Capture the model before the await or re-check under_model_lockbefore returning.🤖 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/model_manager.py` around lines 1285 - 1297, The get_model() branch can return None if idle_worker() unloads the global model during the placement-healing await. Capture the existing model instance before awaiting _heal_tts_placement(), then return that captured instance, or re-check under _model_lock before returning while preserving the placement-healing behavior.
🤖 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/model_manager.py`:
- Around line 1689-1712: Update ensure_tts_on_device so the model transfer and
related cleanup run through the same lock or serialized execution path used by
model.generate(), preventing concurrent requests from accessing the shared model
during m.to(target). If the existing generation path cannot be reused directly,
defer the self-heal until the model is idle while preserving the current
no-raise and boolean return behavior.
In `@tests/test_tts_placement_self_heal.py`:
- Line 39: Move the module-level services.model_manager import into the fixture
or test function that uses mm, resolving it at runtime; preserve the existing
import alias and update any affected references without changing test behavior.
---
Outside diff comments:
In `@backend/services/model_manager.py`:
- Around line 1285-1297: The get_model() branch can return None if idle_worker()
unloads the global model during the placement-healing await. Capture the
existing model instance before awaiting _heal_tts_placement(), then return that
captured instance, or re-check under _model_lock before returning while
preserving the placement-healing 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: e27135e7-dc09-4b54-bedb-0849bd214c15
📒 Files selected for processing (5)
CHANGELOG.mdbackend/api/routers/dub_core.pybackend/services/model_manager.pydocs/performance.mdtests/test_tts_placement_self_heal.py
CodeRabbit (PR #1204): the self-heal moved the shared model on the CPU pool, so a concurrent generate() could hit the instance mid-transfer. Dispatch the move on the GPU pool instead — the hosts that can strand a model are always 1-worker (offload only fires below 8 GB free VRAM), so occupying a slot is real mutual exclusion. Deadlock guard: OmniVoiceBackend._ensure_loaded() reaches get_model() via asyncio.run() from inside generate(), already on a GPU-pool worker. Dispatching back into that pool (or blocking on the model lock held by the loop waiting on us) would deadlock, so that case heals inline — it already owns the GPU slot. Pinned by a regression test. Also resolve the test module's app import at run time (CodeRabbit). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…-on-cpu # Conflicts: # CHANGELOG.md
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
backend/services/model_manager.py (1)
1997-2052: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftSerialize all TTS placement moves
offload_tts_for_asr()/restore_tts_after_asr()can still raceensure_tts_on_device()on the same sharedmodelobject here;_model_lockonly covers the non-inline heal path, so concurrentmodel.to("cpu")andm.to(target)calls are still possible. Route all placement moves through one shared lock or executor so only one transfer runs at a time.🤖 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/model_manager.py` around lines 1997 - 2052, The TTS placement operations can race because ensure_tts_on_device() bypasses _model_lock on the GPU-pool path while offload_tts_for_asr() and restore_tts_after_asr() move the same model. Route every model.to placement operation, including those helpers and the inline heal path, through one shared synchronization mechanism so transfers are serialized without introducing a self-deadlock.Source: Path instructions
🧹 Nitpick comments (1)
backend/services/model_manager.py (1)
2035-2042: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDeadlock avoidance depends on a thread-name convention.
threading.current_thread().name.startswith("gpu-pool")is the only signal preventing the pool from being submitted-into from within itself. If_build_gpu_pool'sthread_name_prefixever changes, this silently reintroduces the deadlock it's meant to avoid. A thread-local sentinel set when entering the pool worker would be more robust than name-sniffing.🤖 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/model_manager.py` around lines 2035 - 2042, Replace the thread-name check in the inline move path with a thread-local sentinel that is set when entering and cleared when leaving each _build_gpu_pool worker. Use that sentinel in the surrounding ensure_tts_on_device flow to detect pool re-entry, preserving inline execution and return behavior while allowing thread_name_prefix to change safely.
🤖 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.
Outside diff comments:
In `@backend/services/model_manager.py`:
- Around line 1997-2052: The TTS placement operations can race because
ensure_tts_on_device() bypasses _model_lock on the GPU-pool path while
offload_tts_for_asr() and restore_tts_after_asr() move the same model. Route
every model.to placement operation, including those helpers and the inline heal
path, through one shared synchronization mechanism so transfers are serialized
without introducing a self-deadlock.
---
Nitpick comments:
In `@backend/services/model_manager.py`:
- Around line 2035-2042: Replace the thread-name check in the inline move path
with a thread-local sentinel that is set when entering and cleared when leaving
each _build_gpu_pool worker. Use that sentinel in the surrounding
ensure_tts_on_device flow to detect pool re-entry, preserving inline execution
and return behavior while allowing thread_name_prefix to change safely.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: f93b30df-af62-43f0-a65a-f678b53554bc
📒 Files selected for processing (4)
CHANGELOG.mdbackend/services/model_manager.pydocs/performance.mdtests/test_tts_placement_self_heal.py
🚧 Files skipped from review as they are similar to previous changes (2)
- CHANGELOG.md
- docs/performance.md
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
2 participants
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 #1191 (mis-framed as "speed varies by time of day" — the bug is deterministic).
offload_tts_for_asr()moves the TTS model to CPU for ASR headroom, butrestore_tts_after_asr()was only reachable on the dub-transcribe success path — so any abort/error/disconnect left it CPU-resident and every later/generateran on CPU until the ~15-min idle unload fired.gen()'sfinallypays the restore debt on every exit path, chained off the ASR unload so the two never contend for VRAM.get_model()self-heals placement (one parameter probe on the hot path; unified memory exempt), so a future unbalanced offload can't strand it either.🤖 Generated with Claude Code
Ensures the TTS model is returned to the correct accelerator after dub generation/transcribe streaming on every exit path (success, abort/error, or client disconnect) by paying “restore debt” in a
finallyblock, and further prevents reoccurrence by havingget_model()verify/repair TTS device placement before returning a loaded model (while keeping unified-memory exempt). This fixes the#1191slowdown where the TTS could be stranded on CPU after an interrupted workflow, causing persistent 10–50x slower generation until restart, and adds regression coverage (12 tests) plus documentation/changelog updates. Human review is most important around the async restore ordering vs ASR unload and ensuring the locking/device-move logic can’t race or deadlock under concurrent model/device mutations (including GPU-pool-thread scenarios).