feat(memory): one TTS engine resident at a time — stop stacking models on 16 GB (#1105)
Measured on a 16 GB M2: a generate on omnivoice (~2.8 GB core) followed by a generate on mlx-audio left BOTH resident (footprint 3.9 → 4.3 GB) because the OmniVoice core lives in model_manager.model while every other engine lives in engines._ENGINE_INSTANCES — two caches that never coordinated, and the latter was never unloaded. That accumulation is a direct contributor to the memory pressure behind the "Can't reach the local backend" OOM deaths. - services/engine_memory.py: evict_other_tts_engines(keep_id) unloads every OTHER resident TTS engine before the incoming one loads — spans both stores (the OmniVoice core under its async lock, and the instance cache). No-op when nothing else is resident, so steady-state single-engine use pays nothing; only a real switch evicts. Default on; OMNIVOICE_SINGLE_ENGINE_RESIDENT=0 to keep several warm. Wired into the /generate path right after the engine resolves. - TTSBackend.unload() (the ABC default) now actually frees the held model: it clears _MODEL_ATTRS (_model/_tts) and empties the device cache. Every in-process engine but OmniVoice previously inherited a NO-OP unload(), so an engine switch dropped the instance ref but left its model for GC with the GPU cache un-emptied. One change fixes all of them and is future-proof. - FasterWhisperBackend.unload() cleared self._asr — an attribute it never assigns — so its model in self._model was never freed. Fixed. Live-verified: omnivoice → mlx-audio now DROPS footprint 2300 → 1541 MB (core evicted) instead of climbing to 4305 with both resident. 7 new unit tests (eviction spans both stores / keeps the active engine / no-op when disabled / a failing unload doesn't abort the sweep / ABC unload frees + is idempotent), order-independent. Backend suite green. Refs the 16 GB OOM class (#1076 #1092 #1093 #1101) Co-authored-by: mergetest <nizam4103@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Palash Debnath
co-authored by
mergetest
Claude Fable 5
mergetest
Claude Fable 5
parent
7705343386
commit
8750b91522
@@ -14,6 +14,11 @@ The bundled TTS model package (`pyproject.toml`) is versioned independently.
|
||||
|
||||
- **The Storage panels got a design.** "Remove all data" and "Reset & remove" listed folders as a flat run of text, so a 7.5 GB model cache and a 391-byte config file carried exactly the same visual weight — the one thing you actually wanted to see (where the space went) was the one thing you couldn't. Every row now has an icon, a dimmed path, and a **proportional bar showing its share of what will be freed**, so the big one looks big. The shared Hugging Face cache is promoted out of the confirm dialog into its own "Optional" row with a checkbox, so ticking it moves the running total **in front of you** instead of springing a different number on you at the point of no return, and the dialog now lists exactly what is about to go.
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Switching TTS engines no longer stacks their models in memory.** Using a second engine in a session (or a per-request engine override) loaded its model *on top of* the first one's, because the OmniVoice core model and the other engines live in two separate caches that never coordinated — measured on a 16 GB M2, an `omnivoice` → `mlx-audio` switch left the machine holding both (footprint 3.9 GB → 4.3 GB, the ~2.8 GB core never freed). That accumulation is a direct contributor to the memory pressure behind the "Can't reach the local backend" OOM deaths. Now only one TTS engine's model stays resident: resolving an engine hands back every *other* resident engine first (the same `omnivoice → mlx-audio` switch now drops to ~1.5 GB). Steady-state single-engine use is unaffected; an A/B switch pays a re-load on the way back (~8 s for the OmniVoice core, ~1–2 s for the lighter engines). Opt out with `OMNIVOICE_SINGLE_ENGINE_RESIDENT=0` if you have RAM to keep several warm. Two underlying leaks are fixed as part of this: every in-process TTS engine's `unload()` now actually frees its model and empties the device cache (previously all but OmniVoice were silent no-ops), and `faster-whisper`'s `unload()` cleared the wrong attribute so its model was never released.
|
||||
|
||||
|
||||
### Fixed
|
||||
|
||||
- **The backend no longer sits on ~2 GB of idle dictation model — the real reason it was being killed on 16 GB Macs.** Four reports of *"Can't reach the local OmniVoice backend"* (#1076, #1092, #1093, #1101) all died at the same moment: during a generate, on a 16 GB machine. Measuring it showed the generate was never the problem — it costs about 116 MB. The problem was the **baseline**: the backend sat at **~6.2 GB even while idle**. The TTS model has always been unloaded after an idle timeout, but the speech-recognition model used for dictation never was — so once you dictated a single time, ~2 GB stayed resident for as long as the app ran. On a 16 GB Mac, that plus the app, macOS, and your other programs is enough for the system to run out of memory and kill the backend, which surfaced as the "can't reach the backend" error. Dictation's model now gets the same idle release the TTS model already had, handing that memory back. The only cost is a ~1.4-second re-warm on your next dictation after a long pause, and a live dictation session is pinned so nothing is ever unloaded mid-sentence.
|
||||
|
||||
@@ -783,6 +783,15 @@ async def generate_speech(
|
||||
),
|
||||
)
|
||||
|
||||
# Single-active-engine memory discipline: hand back any OTHER resident TTS
|
||||
# engine's model before loading this one, so switching engines (or a
|
||||
# per-request engine= override, which bypasses /engines/select entirely)
|
||||
# doesn't stack two multi-GB models in memory — the accumulation behind the
|
||||
# 16 GB-Mac OOM deaths. No-op when nothing else is resident, so steady-state
|
||||
# single-engine use pays nothing. Opt out: OMNIVOICE_SINGLE_ENGINE_RESIDENT=0.
|
||||
from services.engine_memory import evict_other_tts_engines
|
||||
await evict_other_tts_engines(engine_id)
|
||||
|
||||
_model = None
|
||||
_backend = None
|
||||
if backend_cls is OmniVoiceBackend:
|
||||
|
||||
@@ -992,7 +992,12 @@ class FasterWhisperBackend(ASRBackend):
|
||||
return out
|
||||
|
||||
def unload(self) -> None:
|
||||
self._asr = None
|
||||
# #memory: this cleared self._asr — an attribute FasterWhisperBackend
|
||||
# never assigns — so the actual model in self._model was never freed and
|
||||
# a warm faster-whisper stayed resident for the life of the process.
|
||||
# Clear the real handle so the model is released.
|
||||
self._model = None
|
||||
self._asr = None # harmless if a subclass ever used it; keeps idempotence
|
||||
import gc
|
||||
gc.collect()
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Single-active-TTS-engine memory discipline.
|
||||
|
||||
Only one TTS engine's model stays resident at a time. When the generate path
|
||||
resolves an engine, every *other* resident engine is unloaded first — so the
|
||||
previous engine's model is handed back instead of stacking in memory until GC.
|
||||
|
||||
Why this matters (measured on a 16 GB M2): a generate on ``omnivoice`` leaves
|
||||
its ~2.8 GB core model resident; a subsequent generate on ``mlx-audio`` loaded
|
||||
that engine's model **on top** (footprint 3.9 GB → 4.3 GB, both resident),
|
||||
because the two live in different caches with no coordination — the core in
|
||||
``model_manager.model``, the rest in ``engines._ENGINE_INSTANCES`` (which was
|
||||
never unloaded). That accumulation is the baseline that pushes a 16 GB machine
|
||||
into the memory pressure behind the "Can't reach the local backend" OOM deaths.
|
||||
|
||||
Default on. Opt out with ``OMNIVOICE_SINGLE_ENGINE_RESIDENT=0`` on machines with
|
||||
RAM to spare (keeping several engines warm avoids the reload latency on an A/B
|
||||
switch — ~8 s for the OmniVoice core, ~1–2 s for the lighter engines).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
logger = logging.getLogger("omnivoice.engine_memory")
|
||||
|
||||
_OFF = {"0", "false", "no", "off"}
|
||||
|
||||
|
||||
def single_engine_resident() -> bool:
|
||||
"""Whether the one-engine-at-a-time policy is active (default True)."""
|
||||
return (os.environ.get("OMNIVOICE_SINGLE_ENGINE_RESIDENT", "1").strip().lower()
|
||||
not in _OFF)
|
||||
|
||||
|
||||
def _evict_instance_cache(keep_cls) -> list[str]:
|
||||
"""Unload + drop every cached engine instance except ``keep_cls``.
|
||||
|
||||
Operates on the per-request instance cache the generate path shares with the
|
||||
engine health route (``engines._ENGINE_INSTANCES``). Each engine's
|
||||
``unload()`` frees its heavy model (the ABC default clears ``_MODEL_ATTRS``
|
||||
and empties the device cache; subprocess engines reap their sidecar). Never
|
||||
raises — a stuck unload must not block the generation that triggered it."""
|
||||
evicted: list[str] = []
|
||||
try:
|
||||
from api.routers.engines import _ENGINE_INSTANCES
|
||||
except Exception: # pragma: no cover — router import should always succeed
|
||||
return evicted
|
||||
for cls, inst in list(_ENGINE_INSTANCES.items()):
|
||||
if cls is keep_cls:
|
||||
continue
|
||||
try:
|
||||
inst.unload()
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning("evict: %s.unload() failed", getattr(cls, "id", cls.__name__),
|
||||
exc_info=True)
|
||||
_ENGINE_INSTANCES.pop(cls, None)
|
||||
evicted.append(getattr(cls, "id", cls.__name__))
|
||||
return evicted
|
||||
|
||||
|
||||
async def evict_other_tts_engines(keep_id: str) -> list[str]:
|
||||
"""Unload every resident TTS engine except ``keep_id`` and return their ids.
|
||||
|
||||
Spans both stores a TTS model can live in: the OmniVoice core singleton
|
||||
(``model_manager.model``, freed under its async lock when we're switching
|
||||
*away* from it) and the generic engine instance cache. A no-op when the
|
||||
policy is off or nothing else is resident, so steady-state single-engine use
|
||||
pays nothing — only an actual switch evicts. Never raises."""
|
||||
if not single_engine_resident():
|
||||
return []
|
||||
|
||||
evicted: list[str] = []
|
||||
|
||||
# The OmniVoice core singleton — only when the incoming engine isn't it.
|
||||
if keep_id != "omnivoice":
|
||||
try:
|
||||
import services.model_manager as mm
|
||||
|
||||
async with mm._model_lock:
|
||||
if mm.model is not None:
|
||||
mm.model = None
|
||||
mm.free_vram()
|
||||
evicted.append("omnivoice")
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning("evict: OmniVoice core unload failed", exc_info=True)
|
||||
|
||||
# Every other in-process / sidecar engine instance.
|
||||
keep_cls = None
|
||||
try:
|
||||
from services.tts_backend import get_backend_class
|
||||
|
||||
keep_cls = get_backend_class(keep_id)
|
||||
except Exception: # noqa: BLE001 — unknown id → evict all cached instances
|
||||
keep_cls = None
|
||||
evicted.extend(_evict_instance_cache(keep_cls))
|
||||
|
||||
if evicted:
|
||||
logger.info("single-engine eviction: freed %s (keeping %s)", evicted, keep_id)
|
||||
return evicted
|
||||
@@ -249,13 +249,37 @@ class TTSBackend(ABC):
|
||||
# `torch.cuda.empty_cache()` / `torch.mps.empty_cache()`).
|
||||
# • Safe to call before the first generate(): a backend that never
|
||||
# loaded has nothing to release.
|
||||
def unload(self) -> None:
|
||||
"""Release any GPU memory and file handles held by this backend.
|
||||
# Attribute(s) that hold this backend's heavy model, cleared by the default
|
||||
# unload(). Every in-process engine loads its weights lazily into one of
|
||||
# these in `_ensure_loaded()`; the next generate() re-runs that loader. An
|
||||
# engine that holds its model elsewhere (or nowhere — e.g. an external HTTP
|
||||
# server) overrides `unload()` or leaves these unset. OmniVoice overrides
|
||||
# entirely (it drives the shared model_manager singleton).
|
||||
_MODEL_ATTRS: tuple[str, ...] = ("_model", "_tts")
|
||||
|
||||
Called by the registry on engine switch and on app shutdown. Default
|
||||
is a no-op so engines that haven't migrated keep working; per-engine
|
||||
overrides arrive in Phase 2 (see ROADMAP.md). Must be idempotent.
|
||||
def unload(self) -> None:
|
||||
"""Release the heavy model this backend holds, and free device caches.
|
||||
|
||||
Called by the registry on engine switch, by the single-active-engine
|
||||
eviction (services.engine_memory), and on app shutdown. Clears each of
|
||||
``_MODEL_ATTRS`` that is set on this instance, then empties the device
|
||||
cache — so switching engines actually hands the memory back instead of
|
||||
leaving the old model resident until GC (the 16 GB-Mac OOM class). The
|
||||
next generate() lazily reloads. Idempotent and safe before first load:
|
||||
a backend that never loaded has every attr already None/absent.
|
||||
"""
|
||||
freed = False
|
||||
for attr in self._MODEL_ATTRS:
|
||||
if getattr(self, attr, None) is not None:
|
||||
setattr(self, attr, None)
|
||||
freed = True
|
||||
if freed:
|
||||
try:
|
||||
from services.model_manager import free_vram
|
||||
|
||||
free_vram()
|
||||
except Exception: # noqa: BLE001 — unload must never raise (idempotent contract)
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
"""Single-active-TTS-engine memory discipline (the 16 GB-Mac OOM class).
|
||||
|
||||
Measured before this: a generate on ``omnivoice`` (~2.8 GB core) followed by a
|
||||
generate on ``mlx-audio`` left BOTH resident (footprint 3.9 → 4.3 GB), because
|
||||
the OmniVoice core lives in ``model_manager.model`` and the other engines in
|
||||
``engines._ENGINE_INSTANCES`` — two caches with no coordination, and the latter
|
||||
was never unloaded. That accumulation is the baseline that OOM-kills a 16 GB Mac.
|
||||
|
||||
These tests pin the fix: resolving an engine evicts every OTHER resident engine
|
||||
first, across both stores, and the default ``unload()`` actually frees the held
|
||||
model.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("OMNIVOICE_MODEL", "test")
|
||||
os.environ.setdefault("OMNIVOICE_DISABLE_FILE_LOG", "1")
|
||||
|
||||
import pytest
|
||||
|
||||
from services import engine_memory as em
|
||||
|
||||
|
||||
class _FakeEngine:
|
||||
"""A backend holding a heavy model in `_model`, using the ABC unload()."""
|
||||
|
||||
def __init__(self, eid):
|
||||
self.id = eid
|
||||
self._model = object() # stand-in for multi-GB weights
|
||||
self.unloaded = 0
|
||||
|
||||
# Reuse the real ABC default unload via delegation so we test THAT logic.
|
||||
def unload(self):
|
||||
from services.tts_backend import TTSBackend
|
||||
|
||||
self.unloaded += 1
|
||||
TTSBackend.unload(self)
|
||||
|
||||
|
||||
# ── ABC default unload actually frees the model ─────────────────────────────
|
||||
|
||||
|
||||
def _concrete(tb):
|
||||
"""A minimal concrete TTSBackend subclass (satisfies the ABC) that inherits
|
||||
the real default unload() under test."""
|
||||
|
||||
class _Base(tb.TTSBackend):
|
||||
id = "fake"
|
||||
sample_rate = 24000
|
||||
supported_languages = ("en",)
|
||||
|
||||
@classmethod
|
||||
def is_available(cls):
|
||||
return True, "ready"
|
||||
|
||||
def generate(self, *a, **k): # never called in these tests
|
||||
raise NotImplementedError
|
||||
|
||||
return _Base
|
||||
|
||||
|
||||
def test_default_unload_clears_model_attrs_and_frees_vram(monkeypatch):
|
||||
from services import tts_backend as tb
|
||||
|
||||
freed = {"n": 0}
|
||||
monkeypatch.setattr("services.model_manager.free_vram", lambda: freed.__setitem__("n", freed["n"] + 1))
|
||||
|
||||
class Eng(_concrete(tb)):
|
||||
def __init__(self):
|
||||
self._model = object()
|
||||
|
||||
e = Eng()
|
||||
e.unload()
|
||||
assert e._model is None
|
||||
assert freed["n"] == 1
|
||||
# Idempotent + safe when nothing is loaded: a second call frees nothing more.
|
||||
e.unload()
|
||||
assert freed["n"] == 1
|
||||
|
||||
|
||||
def test_default_unload_handles_the_tts_attr_and_missing_attrs(monkeypatch):
|
||||
from services import tts_backend as tb
|
||||
|
||||
monkeypatch.setattr("services.model_manager.free_vram", lambda: None)
|
||||
|
||||
class Sherpa(_concrete(tb)):
|
||||
def __init__(self):
|
||||
self._tts = object() # sherpa holds its model here, not _model
|
||||
|
||||
s = Sherpa()
|
||||
s.unload()
|
||||
assert s._tts is None
|
||||
|
||||
class External(_concrete(tb)):
|
||||
def __init__(self):
|
||||
pass # no model attrs at all (e.g. an HTTP-server engine)
|
||||
|
||||
External().unload() # must not raise
|
||||
|
||||
|
||||
# ── evict_other_tts_engines ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_core_model(monkeypatch):
|
||||
"""Pin the OmniVoice core singleton to None so these tests are order-
|
||||
independent: evict_other_tts_engines() also frees model_manager.model when
|
||||
switching away from omnivoice, and an earlier full-suite test can leave a
|
||||
model loaded there. Tests that exercise the core eviction set it explicitly."""
|
||||
import services.model_manager as mm
|
||||
|
||||
monkeypatch.setattr(mm, "model", None, raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def instance_cache(monkeypatch):
|
||||
"""A stand-in for engines._ENGINE_INSTANCES keyed by class."""
|
||||
import api.routers.engines as eng
|
||||
|
||||
cache: dict = {}
|
||||
monkeypatch.setattr(eng, "_ENGINE_INSTANCES", cache, raising=False)
|
||||
return cache
|
||||
|
||||
|
||||
async def _evict(keep):
|
||||
return await em.evict_other_tts_engines(keep)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evicts_other_engine_instances_but_keeps_the_active_one(instance_cache, monkeypatch):
|
||||
class KittenTTSBackend:
|
||||
id = "kittentts"
|
||||
|
||||
class MLXAudioBackend:
|
||||
id = "mlx-audio"
|
||||
|
||||
keep = MLXAudioBackend()
|
||||
drop = KittenTTSBackend()
|
||||
drop.unloaded = 0
|
||||
drop.unload = lambda: setattr(drop, "unloaded", drop.unloaded + 1)
|
||||
instance_cache[MLXAudioBackend] = keep
|
||||
instance_cache[KittenTTSBackend] = drop
|
||||
|
||||
monkeypatch.setattr(em, "get_backend_class", None, raising=False)
|
||||
monkeypatch.setattr(
|
||||
"services.tts_backend.get_backend_class",
|
||||
lambda i: MLXAudioBackend if i == "mlx-audio" else KittenTTSBackend,
|
||||
)
|
||||
|
||||
evicted = await _evict("mlx-audio")
|
||||
|
||||
assert evicted == ["kittentts"]
|
||||
assert drop.unloaded == 1
|
||||
assert KittenTTSBackend not in instance_cache # dropped
|
||||
assert instance_cache[MLXAudioBackend] is keep # kept
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evicts_the_omnivoice_core_when_switching_away_from_it(instance_cache, monkeypatch):
|
||||
import services.model_manager as mm
|
||||
|
||||
monkeypatch.setattr(mm, "model", object(), raising=False)
|
||||
freed = {"n": 0}
|
||||
monkeypatch.setattr(mm, "free_vram", lambda: freed.__setitem__("n", freed["n"] + 1))
|
||||
monkeypatch.setattr("services.tts_backend.get_backend_class", lambda i: type("X", (), {"id": i}))
|
||||
|
||||
evicted = await _evict("mlx-audio")
|
||||
|
||||
assert "omnivoice" in evicted
|
||||
assert mm.model is None
|
||||
assert freed["n"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keeps_the_omnivoice_core_when_it_IS_the_active_engine(instance_cache, monkeypatch):
|
||||
import services.model_manager as mm
|
||||
|
||||
sentinel = object()
|
||||
monkeypatch.setattr(mm, "model", sentinel, raising=False)
|
||||
monkeypatch.setattr(mm, "free_vram", lambda: None)
|
||||
monkeypatch.setattr("services.tts_backend.get_backend_class", lambda i: type("X", (), {"id": i}))
|
||||
|
||||
evicted = await _evict("omnivoice")
|
||||
|
||||
assert "omnivoice" not in evicted
|
||||
assert mm.model is sentinel # the active engine's model is NOT evicted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_policy_can_be_disabled(instance_cache, monkeypatch):
|
||||
import services.model_manager as mm
|
||||
|
||||
monkeypatch.setenv("OMNIVOICE_SINGLE_ENGINE_RESIDENT", "0")
|
||||
monkeypatch.setattr(mm, "model", object(), raising=False)
|
||||
|
||||
class Other:
|
||||
id = "kittentts"
|
||||
|
||||
other = Other()
|
||||
other.unload = lambda: pytest.fail("must not unload when policy is off")
|
||||
instance_cache[Other] = other
|
||||
|
||||
assert await _evict("mlx-audio") == []
|
||||
assert mm.model is not None # untouched
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_unload_does_not_abort_the_eviction(instance_cache, monkeypatch):
|
||||
class A:
|
||||
id = "a"
|
||||
|
||||
class B:
|
||||
id = "b"
|
||||
|
||||
a, b = A(), B()
|
||||
a.unload = lambda: (_ for _ in ()).throw(RuntimeError("stuck"))
|
||||
b.unloaded = 0
|
||||
b.unload = lambda: setattr(b, "unloaded", b.unloaded + 1)
|
||||
instance_cache[A] = a
|
||||
instance_cache[B] = b
|
||||
monkeypatch.setattr("services.tts_backend.get_backend_class",
|
||||
lambda i: type("keep", (), {"id": i}))
|
||||
|
||||
evicted = await _evict("other") # keep nothing in the cache
|
||||
|
||||
# Both attempted; the raising one didn't stop the other from being freed.
|
||||
assert set(evicted) == {"a", "b"}
|
||||
assert b.unloaded == 1
|
||||
assert not instance_cache # both dropped despite the failure
|
||||
Reference in New Issue
Block a user