diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ea9cc66..fe5b3f3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,7 @@ The bundled TTS model package (`pyproject.toml`) is versioned independently. ### Fixed +- Generation no longer crawls on CPU after a cancelled or failed dub: the TTS model is moved back to the GPU on every exit path, and each generation now verifies its own placement (#1191) - Subtitle parsing no longer stalls on a blank-line-heavy `.srt`: the timing-line regex backtracked across newlines, so a mis-saved export could pin an import for hours (#1203) - A broken ASR engine's fallback could silently auto-download multi-GB weights — every fallback now passes the same no-download preflight and shows the download CTA instead (#1189) - Dub transcription releases the ASR model from VRAM on every exit — crashes, early errors, and client disconnects included (#1175) diff --git a/backend/api/routers/dub_core.py b/backend/api/routers/dub_core.py index 3bc89ee1..ba358215 100644 --- a/backend/api/routers/dub_core.py +++ b/backend/api/routers/dub_core.py @@ -494,6 +494,12 @@ async def dub_transcribe_stream( # _gen_body parks the loaded backend here; the normal unload clears it; # gen()'s `finally` unloads whatever is still parked, on EVERY exit. _loaded_asr: dict = {"backend": None} + # Same shape, same reason, for the TTS offload (#1191): offload_tts_for_asr() + # moves the TTS model to CPU, and only _gen_body's success path moved it + # back — so an abort/error/disconnect stranded it there, silently making + # every subsequent /generate run on CPU. Set on a successful offload, + # cleared by the normal restore, honoured by gen()'s `finally` on EVERY exit. + _tts_offloaded: dict = {"v": False} async def _gen_body(): # ── Preflight — run INSIDE the stream, never before it (#1196) ── @@ -750,6 +756,8 @@ async def dub_transcribe_stream( # transcription can still proceed (it just has less headroom). try: await loop.run_in_executor(_cpu_pool, offload_tts_for_asr) + # Restore is now owed on every exit path, not just success (#1191). + _tts_offloaded["v"] = True except Exception as e: logger.warning("offload_tts_for_asr failed (continuing): %s", e) @@ -1337,6 +1345,8 @@ async def dub_transcribe_stream( _loaded_asr["backend"] = None await loop.run_in_executor(_cpu_pool, restore_tts_after_asr) + # Debt paid — don't make gen()'s finally repeat it. + _tts_offloaded["v"] = False if torch.backends.mps.is_available(): try: torch.mps.empty_cache() @@ -1380,6 +1390,32 @@ async def dub_transcribe_stream( # (GeneratorExit bypasses the except, never this finally). _b = _loaded_asr.get("backend") _loaded_asr["backend"] = None + # Pay the TTS-restore debt on every exit path (#1191). Leaving it + # unpaid is what stranded the TTS model on CPU after an abort or a + # disconnect, degrading every later generation by 10-50x. + _restore_tts = _tts_offloaded["v"] + _tts_offloaded["v"] = False + + def _log_bg(f, what): + if not f.cancelled() and f.exception(): + logger.warning("%s failed: %s", what, f.exception()) + + def _submit_tts_restore(_f=None): + if _f is not None: + _log_bg(_f, "Unloading ASR backend") + if not _restore_tts: + return + try: + _r = asyncio.get_running_loop().run_in_executor( + _cpu_pool, restore_tts_after_asr + ) + _r.add_done_callback(lambda f: _log_bg(f, "restore_tts_after_asr")) + except RuntimeError: + try: + restore_tts_after_asr() + except Exception as e: + logger.warning("restore_tts_after_asr failed: %s", e) + if _b is not None: # unload() blocks (gc.collect + CUDA cache drop can take # seconds) and this finally also runs under GeneratorExit, @@ -1390,17 +1426,19 @@ async def dub_transcribe_stream( _fut = asyncio.get_running_loop().run_in_executor( _gpu_pool, _b.unload ) - _fut.add_done_callback( - lambda f: f.cancelled() - or (f.exception() and logger.warning( - "Failed to unload ASR backend: %s", f.exception())) - ) + # Restore the TTS model only AFTER the ASR weights are + # freed — the same ordering the success path enforces, so + # the two never contend for VRAM. + _fut.add_done_callback(_submit_tts_restore) except RuntimeError: # No running loop (interpreter teardown) — best effort. try: _b.unload() except Exception as e: logger.warning("Failed to unload ASR backend: %s", e) + _submit_tts_restore() + else: + _submit_tts_restore() return StreamingResponse( gen(), diff --git a/backend/services/model_manager.py b/backend/services/model_manager.py index f90c15dc..53772ea0 100644 --- a/backend/services/model_manager.py +++ b/backend/services/model_manager.py @@ -1286,6 +1286,14 @@ async def get_model(): global model, _last_used _last_used = time.time() if model is not None: + # Placement self-heal (#1191). The ASR offload/restore pair below is a + # *balanced-call* contract, and any unbalanced path (abort, terminal + # error, client disconnect) used to leave the TTS model resident on CPU + # — where it stayed for EVERY later generation until the idle unload + # fired, at 10-50x the latency. Verifying placement here makes the + # contract unnecessary: a future unbalanced offload can no longer + # strand the model, because the next generation moves it back. + await _heal_tts_placement() return model async with _model_lock: @@ -1628,6 +1636,100 @@ def restore_tts_after_asr(): except Exception as e: logger.warning("TTS restore to %s failed: %s", get_best_device(), e) + +def _first_param_device(obj): + """Device the weights of ``obj`` actually live on, or None if undeterminable. + + The TTS runtime is a wrapper object, not necessarily an ``nn.Module``, so + fall back to the first sub-module that owns parameters. Never raises. + """ + try: + params = getattr(obj, "parameters", None) + if callable(params): + for p in params(): + return p.device + except Exception: # noqa: BLE001 — a probe must never break generation + pass + try: + torch = _lazy_torch() + for v in vars(obj).values(): + if isinstance(v, torch.nn.Module): + for p in v.parameters(): + return p.device + except Exception: # noqa: BLE001 + pass + return None + + +def _stranded_tts_target(): + """Target device string when the loaded TTS model is stranded off it, else None. + + Ordered cheapest-first so the hot path (model already on the accelerator) + costs a single parameter probe: anything not sitting on CPU is by + definition not stranded, because the only thing that moves the model is + ``offload_tts_for_asr()`` and it only ever moves it to CPU. + """ + m = model + if m is None: + return None + dev = _first_param_device(m) + if dev is None or getattr(dev, "type", None) != "cpu": + return None + if not _has_dedicated_vram(): + # Unified memory / CPU-only: the offload RELEASES the model rather than + # moving it, and CPU is the legitimate home here. Nothing to heal. + return None + try: + target = get_best_device() + except Exception: # noqa: BLE001 + return None + return target if target in ("cuda", "xpu") else None + + +def ensure_tts_on_device() -> bool: + """Move the TTS model back onto its target device if it was stranded on CPU. + + Returns True when a move actually happened. Never raises — a failed move + just leaves the model on CPU, which is exactly the pre-fix behaviour + (slow), never a failed generation. + """ + target = _stranded_tts_target() + m = model + if target is None or m is None: + return False + try: + logger.warning( + "TTS model found stranded on CPU (an ASR offload was never restored) — " + "moving it back to %s; generation would otherwise run 10-50x slower (#1191).", + target, + ) + m.to(target) + free_vram() + return True + except Exception as e: # noqa: BLE001 + logger.warning("TTS placement self-heal to %s failed (staying on CPU): %s", target, e) + return False + + +async def _heal_tts_placement() -> None: + """Async wrapper for :func:`ensure_tts_on_device` used by ``get_model()``. + + The cheap mismatch probe runs inline; only the rare actual move is handed + to the CPU pool (the same pool the offload/restore pair uses), because a + multi-GB host-to-device copy would otherwise stall the event loop. + """ + if _stranded_tts_target() is None: + return + async with _model_lock: + if _stranded_tts_target() is None: + return # another caller healed it while we waited + try: + await asyncio.get_running_loop().run_in_executor( + _cpu_pool, ensure_tts_on_device + ) + except Exception as e: # noqa: BLE001 + logger.warning("TTS placement self-heal could not run: %s", e) + _diar_pipeline = None # Sentinel error classes used by callers (dub_core) to decide whether to diff --git a/docs/performance.md b/docs/performance.md index 52407baa..315445ce 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -5,7 +5,7 @@ should leave alone. Everything here applies to the current release; numbers marked "measured" come from `scripts/bench_pipeline.py` on a 16 GB Apple Silicon M2 — your hardware will differ, but the *ratios* hold. -## First: the three classic causes of "it got slow" +## First: the classic causes of "it got slow" Before touching any knob, check these — they account for most slowness reports: @@ -43,6 +43,17 @@ Before touching any knob, check these — they account for most slowness reports the badge (full text on hover). Note: **GPU acceleration on Windows is NVIDIA/CUDA-only** — AMD and Intel GPUs run CPU-only there (see [Windows install notes](install/windows.md)). +5. **You aborted a dub earlier (fixed in v0.3.23).** Dubbing moves the TTS + model to CPU to free VRAM for the ASR model, then moves it back when the + transcription finishes. Before v0.3.23 that move-back only ran on the fully + successful path, so cancelling a dub, hitting a dub error, or closing the + tab mid-transcription left the TTS model stranded on CPU — and **every** + later generation ran there, 10-50x slower with the CPU pegged, until the + ~15-minute idle unload happened to fire. Restarting the backend cleared it, + which made it look random or time-of-day related (#1191). Since v0.3.23 the + move-back runs on every exit path, *and* each generation verifies the model + is on the expected device and moves it back itself — so no future code path + can strand it again. If you are on an older build, restart the backend. ## What a generation actually spends time on diff --git a/tests/test_tts_placement_self_heal.py b/tests/test_tts_placement_self_heal.py new file mode 100644 index 00000000..c1062200 --- /dev/null +++ b/tests/test_tts_placement_self_heal.py @@ -0,0 +1,297 @@ +"""The TTS model must never stay stranded on CPU after an ASR offload (#1191). + +Reported as "generation speed varies greatly depending on the time of day" — +which is a red herring. The bug is fully deterministic: + +``offload_tts_for_asr()`` moves the TTS model to CPU to make VRAM room for +WhisperX, and 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 because ``get_model()`` never re-checked placement, +EVERY subsequent /generate ran on CPU (10-50x slower, CPU pegged) until the +~15-minute idle unload happened to fire. Whether a user hit it came down to +whether they had aborted a dub earlier, which correlates with nothing but +feels like "time of day". + +Two independent guarantees are tested here: + +1. **Balanced pair** — gen()'s ``finally`` pays the restore debt on every exit + path, not just success. +2. **Self-heal (the class fix)** — ``get_model()`` verifies placement and moves + the model back itself, so a *future* unbalanced offload path cannot strand + it either. + +Fail-before: (1) restore was never called on abort/error; (2) ``get_model()`` +returned the CPU-resident model untouched. +""" +from __future__ import annotations + +import asyncio +import os +import struct +import threading +import wave + +os.environ.setdefault("OMNIVOICE_MODEL", "test") +os.environ.setdefault("OMNIVOICE_DISABLE_FILE_LOG", "1") + +import pytest + +import services.model_manager as mm + + +# ── Fakes: no real weights, no real torch ───────────────────────────────── +class _FakeDev: + def __init__(self, type_: str): + self.type = type_ + + def __repr__(self): # pragma: no cover — debugging aid + return f"device({self.type})" + + +class _FakeParam: + def __init__(self, device: str): + self.device = _FakeDev(device) + + +class _FakeTTS: + """Stand-in for the TTS runtime: records ``.to()`` moves, owns a parameter.""" + + def __init__(self, device: str = "cuda"): + self._param = _FakeParam(device) + self.moves: list[str] = [] + + def parameters(self): + yield self._param + + def to(self, device): + device = device if isinstance(device, str) else getattr(device, "type", str(device)) + self.moves.append(device) + self._param = _FakeParam(device) + return self + + @property + def device_type(self) -> str: + return self._param.device.type + + +@pytest.fixture +def gpu_host(monkeypatch): + """A dedicated-VRAM host (CUDA) with a loaded TTS model.""" + monkeypatch.setattr(mm, "_has_dedicated_vram", lambda: True) + monkeypatch.setattr(mm, "get_best_device", lambda: "cuda") + monkeypatch.setattr(mm, "free_vram", lambda: None) + + def _install(device: str = "cuda") -> _FakeTTS: + fake = _FakeTTS(device) + monkeypatch.setattr(mm, "model", fake, raising=False) + return fake + + return _install + + +# ── 1. The placement probe ──────────────────────────────────────────────── +def test_model_on_the_accelerator_is_not_stranded(gpu_host): + """The hot path must be a no-op: a healthy model costs one parameter probe.""" + gpu_host("cuda") + assert mm._stranded_tts_target() is None + + +def test_model_left_on_cpu_is_reported_stranded(gpu_host): + """CPU-resident weights on a CUDA box is the #1191 state.""" + gpu_host("cpu") + assert mm._stranded_tts_target() == "cuda" + + +def test_unified_memory_is_never_treated_as_stranded(gpu_host, monkeypatch): + """On Apple Silicon / CPU-only the offload RELEASES the model rather than + moving it, and CPU is the legitimate home — healing there would be wrong.""" + gpu_host("cpu") + monkeypatch.setattr(mm, "_has_dedicated_vram", lambda: False) + assert mm._stranded_tts_target() is None + + +def test_probe_tolerates_a_model_without_parameters(gpu_host, monkeypatch): + """An engine wrapper we can't introspect must not break generation.""" + monkeypatch.setattr(mm, "model", object(), raising=False) + monkeypatch.setattr(mm, "_lazy_torch", lambda: (_ for _ in ()).throw(RuntimeError("no torch"))) + assert mm._stranded_tts_target() is None + + +# ── 2. The self-heal ────────────────────────────────────────────────────── +def test_ensure_tts_on_device_moves_a_stranded_model_back(gpu_host): + fake = gpu_host("cpu") + assert mm.ensure_tts_on_device() is True + assert fake.moves == ["cuda"] + assert fake.device_type == "cuda" + + +def test_ensure_tts_on_device_is_a_noop_when_already_placed(gpu_host): + fake = gpu_host("cuda") + assert mm.ensure_tts_on_device() is False + assert fake.moves == [] + + +def test_self_heal_failure_degrades_to_cpu_never_raises(gpu_host, monkeypatch): + """An OOM while moving back must leave the pre-fix behaviour (slow), not a + failed generation.""" + fake = gpu_host("cpu") + + def _boom(_device): + raise RuntimeError("CUDA out of memory: simulated") + + monkeypatch.setattr(fake, "to", _boom) + assert mm.ensure_tts_on_device() is False # no exception escapes + + +def test_get_model_heals_placement_before_returning(gpu_host): + """THE CLASS FIX. Fail-before: get_model() returned the CPU-resident model + untouched, so every generation after a stranded offload ran on CPU.""" + fake = gpu_host("cpu") + + got = asyncio.run(mm.get_model()) + + assert got is fake + assert fake.moves == ["cuda"], "get_model() must move a stranded model back" + assert fake.device_type == "cuda" + + +def test_get_model_does_not_move_a_healthy_model(gpu_host): + """The self-heal must be free on the hot path — no spurious device moves.""" + fake = gpu_host("cuda") + assert asyncio.run(mm.get_model()) is fake + assert fake.moves == [] + + +# ── 3. The balanced pair, end to end through the SSE endpoint ───────────── +def _make_wav(path, seconds=1.0, rate=16000): + with wave.open(str(path), "wb") as w: + w.setnchannels(1) + w.setsampwidth(2) + w.setframerate(rate) + w.writeframes(struct.pack("<%dh" % int(rate * seconds), *([0] * int(rate * seconds)))) + + +def _run_transcribe_stream(job_id): + from api.routers import dub_core as dc + + async def _collect(): + resp = await dc.dub_transcribe_stream(job_id) + parts = [] + async for chunk in resp.body_iterator: + parts.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else str(chunk)) + # gen()'s finally hands the restore to the CPU pool fire-and-forget + # (it also runs under GeneratorExit, where awaiting is illegal), so + # give the loop a turn to run the done-callback chain. + await asyncio.sleep(0) + return "".join(parts) + + return asyncio.run(_collect()) + + +@pytest.fixture +def transcribe_job(tmp_path, monkeypatch): + """A dub job whose transcribe stream offloads the TTS model, with the + offload/restore pair instrumented.""" + from api.routers import dub_core as dc + + job_id = "t_1191" + audio = tmp_path / "a.wav" + _make_wav(audio) + dc._dub_jobs[job_id] = { + "audio_path": str(audio), "vocals_path": None, "scene_cuts": [], + } + + calls: dict = {"offload": 0, "restore": 0, "restored": threading.Event()} + + def _offload(): + calls["offload"] += 1 + + def _restore(): + calls["restore"] += 1 + calls["restored"].set() + + monkeypatch.setattr(dc, "offload_tts_for_asr", _offload) + monkeypatch.setattr(dc, "restore_tts_after_asr", _restore) + + fake_model = type("_FakeModel", (), {"_asr_pipe": None})() + + async def _ok_model(): + return fake_model + + monkeypatch.setattr(dc, "get_model", _ok_model) + + class _FakeASR: + id = "fake" + + def ensure_loaded(self): + pass + + def transcribe(self, path, *, word_timestamps=True): + return { + "chunks": [{"text": "hi", "timestamp": (0.0, 0.5)}], + "segments": [], + "language": "en", + } + + def unload(self): + pass + + monkeypatch.setattr( + "services.asr_backend.get_active_asr_backend", lambda *a, **k: _FakeASR() + ) + monkeypatch.setattr( + "services.asr_backend.asr_model_missing_error", lambda *a, **k: None + ) + + try: + yield job_id, dc._dub_jobs[job_id], calls + finally: + dc._dub_jobs.pop(job_id, None) + + +def test_aborted_transcribe_still_restores_the_tts_model(transcribe_job): + """THE REPORTED BUG. Fail-before: the restore lived only on the success + path, so aborting a dub left the TTS model on CPU for the rest of the + process — and every later generation crawled.""" + job_id, job, calls = transcribe_job + job["aborted"] = True + + body = _run_transcribe_stream(job_id) + + assert "event: aborted" in body, body + assert calls["offload"] == 1, "precondition: the stream must have offloaded" + assert calls["restored"].wait(timeout=10), ( + "an aborted transcribe left the TTS model stranded on CPU (#1191)" + ) + assert calls["restore"] == 1 + + +def test_crashed_transcribe_still_restores_the_tts_model(transcribe_job, monkeypatch): + """Same debt, different exit path: an unanticipated mid-stream exception.""" + from api.routers import dub_core as dc + + job_id, _job, calls = transcribe_job + + def _boom(*a, **k): + raise RuntimeError("segmentation exploded: simulated") + + monkeypatch.setattr(dc, "segment_transcript", _boom) + + body = _run_transcribe_stream(job_id) + + assert "event: error" in body, body + assert calls["offload"] == 1 + assert calls["restored"].wait(timeout=10), ( + "a crashed transcribe left the TTS model stranded on CPU (#1191)" + ) + + +def test_successful_transcribe_restores_exactly_once(transcribe_job): + """The finally must not double-restore what the success path already paid.""" + job_id, _job, calls = transcribe_job + + body = _run_transcribe_stream(job_id) + + assert "event: final" in body, body # the real success path, not an error exit + assert calls["offload"] == 1 + assert calls["restore"] == 1, "restore ran twice — the success path already paid it"