fix(memory): release the model before emptying the cache, not after

The shared voice model's unload emptied the allocator caches and *then*
dropped the reference. That frees nothing: the weights are still reachable
when gc.collect() runs, empty_cache() only returns blocks the allocator
already considered free, and the reference drops a moment later into a cache
nothing will flush again. The unload logs success, the engine leaves the
registry, and nvidia-smi does not move.

Six modules open-coded the same two lines. Exactly one had them inverted --
OmniVoiceBackend.unload, which is the path the engine-registry idle sweep
reaches, which is the sweep a headless worker node runs. So every unload a
user could trigger from the UI worked, and the one that runs unattended on a
machine lending its GPU held 3.6 GB indefinitely. Found on hardware: the
sweep fired on schedule, logged "Released 1 idle engine(s)", and VRAM stayed
flat at 3656 MiB for the next two minutes.

Replace all six with model_manager.unload_shared_model(), which clears the
reference, drops the clone-prompt side cache, then frees -- in that order,
in one place. Two callers gain the side-cache drop they were missing
(/system/flush-memory and the shutdown path), which is the same defect one
step down: an unload that kept the encoded reference tensors belonging to the
model it had just released.

A source guard asserts nothing outside model_manager assigns the shared
reference, so the next caller cannot reintroduce the ordering. It caught the
sixth site while being written.

Also give the AudioSeal watermark models the bargain every other model in the
app already makes: they loaded on the first embed and stayed resident for the
life of the process. CPU-resident, so this is system RAM rather than VRAM,
and the machines that notice are the ones running batches.

The error text on a failing unload changes with the ordering. "Could not be
unloaded, retry after the current generation finishes" was accurate when the
cache flush ran first and aborted before the release; now the release has
already happened and only the flush can fail, so it says that instead of
sending the user to repeat work that is done.
This commit is contained in:
velixio
2026-08-12 02:33:04 +05:30
parent 6a6f3fbc29
commit 090cc37144
12 changed files with 279 additions and 36 deletions
+2 -1
View File
@@ -61,7 +61,8 @@ the frozen-backend fallback mirror it for their toolchains.
### Fixed
- Remote workers that lack required task inputs, progress leases, or model-download commands are now refused visibly instead of returning wrong audio or hanging. (#1478)
- An idle voice model now actually hands its memory back. The unload emptied the GPU cache a moment before releasing the model, so it freed nothing while reporting success — a GPU machine lending its card sat on 3.6 GB indefinitely. (#1495)
- The AudioSeal watermark models are released after the same idle period as everything else, instead of staying in memory for the life of the app once anything was watermarked. (#1495)
- Remote GPU workers now synthesize a dub's fresh segments as one coarse job with live progress and cancellation; fitting, assembly and RVC remain local. (#1478)
- Gallery voice previews now fall back to a local render when a downloaded clip cannot be decoded, instead of failing silently. (#1478)
- A second VoiceStudio instance can no longer silently share the remote-worker port; it keeps running locally and explains how to resolve the conflict. (#1478)
+4 -3
View File
@@ -544,9 +544,10 @@ async def flush_memory(unload_model: bool = False):
if unload_model:
import services.model_manager as mm
async with mm._model_lock:
if mm.model is not None:
mm.model = None
freed_model = True
# Also drops the clone-prompt side cache, which this path used to
# leave resident — an "unload" that kept the encoded reference
# tensors belonging to the model it just released (#1495).
freed_model = mm.unload_shared_model()
# Multi-pass GC to break reference cycles
gc.collect(generation=2)
+3 -2
View File
@@ -839,9 +839,10 @@ async def lifespan(app: FastAPI):
# Unload the model and free GPU memory
try:
import services.model_manager as mm
if mm.model is not None:
mm.model = None
if mm.unload_shared_model():
logger.info("Shutdown: model unloaded.")
# Still unconditional: there are allocator caches to hand back even when
# no model was resident.
mm.free_vram()
# Abandon a still-running preload's GPU-pool thread (Python can't kill
# a thread mid blocking call) so it can't outlive this shutdown block
+1 -3
View File
@@ -78,9 +78,7 @@ async def evict_other_tts_engines(keep_id: str) -> list[str]:
import services.model_manager as mm
async with mm._model_lock:
if mm.model is not None:
mm.model = None
mm.free_vram()
if mm.unload_shared_model():
evicted.append("omnivoice")
except Exception: # noqa: BLE001
logger.warning("evict: VoiceStudio core unload failed", exc_info=True)
+1 -3
View File
@@ -220,9 +220,7 @@ async def unload(model_id: str) -> dict:
if model_id == "tts":
async with mm._model_lock:
if mm.model is not None:
mm.model = None
mm.free_vram()
if mm.unload_shared_model():
return {"unloaded": "tts", "success": True}
return {"unloaded": "tts", "success": False, "reason": "not loaded"}
+45 -7
View File
@@ -2595,7 +2595,6 @@ def _resolve_idle_timeout() -> float:
async def idle_worker():
global model
torch = _lazy_torch()
while True:
await asyncio.sleep(30)
@@ -2603,9 +2602,7 @@ async def idle_worker():
async with _model_lock:
if model is not None and time.time() - _last_used > idle_timeout:
logger.info("Idle timeout reached. Unloading VoiceStudio model to free VRAM.")
model = None
release_tts_side_caches()
free_vram()
unload_shared_model()
# The capture/dictation ASR was never idle-released — so once a user
# dictated, its model stayed resident for the life of the process while
# the TTS model dutifully freed its 3.8 GB. On a 16 GB Mac that left the
@@ -2621,6 +2618,16 @@ async def idle_worker():
free_vram()
except Exception: # noqa: BLE001 — the reaper must never kill idle_worker
logger.warning("idle capture-ASR release failed", exc_info=True)
# Same bargain for the AudioSeal watermark models, which loaded on the
# first embed and were never released. Deliberately only here and not
# in the make-room paths: watermarking runs immediately *after* a
# generate, so evicting it just before one would only buy a reload.
try:
from services.watermark import release_idle_models
release_idle_models(idle_timeout)
except Exception: # noqa: BLE001 — the reaper must never kill idle_worker
logger.warning("idle watermark-model release failed", exc_info=True)
def release_tts_side_caches():
"""Drop caches keyed to the TTS model, for when the model itself is released.
@@ -2667,6 +2674,39 @@ def free_vram():
torch.xpu.empty_cache()
def unload_shared_model() -> bool:
"""Drop the shared VoiceStudio model and actually give the memory back.
The order is the entire point of this function. Clearing the reference has
to come FIRST, then the allocator caches. ``free_vram()`` run while
``model`` is still bound releases nothing: the weights are still reachable,
so ``gc.collect()`` keeps them and ``empty_cache()`` only returns blocks
the allocator already considered free. The reference drops a moment later,
the weights go back into torch's cache, and nobody ever hands that cache to
the driver so the unload is logged, the engine is dropped from the
registry, and ``nvidia-smi`` does not move.
Six call sites open-coded this pair and one of them had it inverted the
one the engine-registry sweep reaches, which is the sweep a headless worker
node runs. A worker therefore sat on 3.6 GB indefinitely while reporting
the engine released, and every other path looked fine (#1495). One helper,
so there is one ordering and nowhere left to get it wrong.
Takes no lock of its own: the sync engine-registry path
(``OmniVoiceBackend.unload``) cannot await one, and callers that do hold
``_model_lock`` simply keep holding it across the call. Assignment is
GIL-atomic, so the worst a race costs is a redundant reload. Idempotent
returns False when nothing was resident.
"""
global model
if model is None:
return False
model = None
release_tts_side_caches()
free_vram()
return True
def _has_dedicated_vram():
"""Check if the current device has limited dedicated VRAM that needs offloading."""
torch = _lazy_torch()
@@ -2704,9 +2744,7 @@ def _offload_unified_memory() -> bool:
"(it reloads on the next generation).",
"unknown" if free_gb is None else f"{free_gb:.1f}",
)
model = None
release_tts_side_caches()
free_vram()
unload_shared_model()
return True
except Exception as e: # noqa: BLE001
logger.warning("unified-memory TTS offload failed (continuing): %s", e)
+13 -6
View File
@@ -608,18 +608,25 @@ class OmniVoiceBackend(TTSBackend):
clear the shared one and free GPU memory too. Idempotent and safe before
the first generate(). Best-effort: assignment is GIL-atomic, so we don't
take the async ``_model_lock`` from this sync path; the registry wraps
this call in try/except so a race can never block an engine switch."""
this call in try/except so a race can never block an engine switch.
Delegates to ``model_manager.unload_shared_model`` rather than clearing
the singleton here: this path used to free the device caches *before*
dropping the shared reference, which frees nothing, and it is the path
the idle sweep on a headless worker node runs (#1495)."""
self._model = None
clear_clone_prompt_cache() # #427: drop cached prompts so VRAM is freed
try:
import services.model_manager as mm
if mm.model is not None:
mm.free_vram()
mm.model = None
mm.unload_shared_model()
except Exception as exc:
logger.warning("Shared voice model unload did not complete")
# The reference is already gone by the time anything in here can
# raise — only the device-cache flush is left, and that failing is
# a driver problem, not a stuck model. Saying "retry" would send
# the user to repeat an unload that already happened.
logger.warning("Shared voice model released, but the device cache flush failed")
raise RuntimeError(
"The shared voice model could not be unloaded. Retry after the current generation finishes."
"The voice model was released, but the GPU memory cache could not be flushed."
) from exc
+32 -2
View File
@@ -22,6 +22,7 @@ from __future__ import annotations
import logging
import math
import time
import torch
from typing import Optional
@@ -34,6 +35,8 @@ logger = logging.getLogger("omnivoice.watermark")
_generator = None
_detector = None
_audioseal_available: Optional[bool] = None
# Monotonic stamp of the last embed/detect, for the idle release below.
_last_used = 0.0
# 16-bit message: "OM" in ASCII = 0x4F 0x4D = 0100_1111 0100_1101
# This is our signature — every VoiceStudio-generated audio carries it.
@@ -76,7 +79,8 @@ def _check_available() -> bool:
def _get_generator():
"""Lazy-load the AudioSeal generator model."""
global _generator
global _generator, _last_used
_last_used = time.monotonic()
if _generator is None:
from audioseal import AudioSeal
_generator = AudioSeal.load_generator("audioseal_wm_16bits")
@@ -87,7 +91,8 @@ def _get_generator():
def _get_detector():
"""Lazy-load the AudioSeal detector model."""
global _detector
global _detector, _last_used
_last_used = time.monotonic()
if _detector is None:
from audioseal import AudioSeal
_detector = AudioSeal.load_detector("audioseal_detector_16bits")
@@ -96,6 +101,31 @@ def _get_detector():
return _detector
def release_idle_models(idle_seconds: float, *, now: Optional[float] = None) -> bool:
"""Drop the AudioSeal models if nothing has watermarked for ``idle_seconds``.
These load on the first embed or detect and then stayed resident for the
life of the process the same bargain the TTS model and the capture ASR
both stopped making. Modest next to those (they run on CPU, so this is
system RAM rather than VRAM), but a batch job that watermarks once leaves
them held forever afterwards, and the machines that hit memory pressure are
the ones running batches.
Returns True if anything was released. Never raises: this runs from the
idle reaper, which must survive it.
"""
global _generator, _detector
if _generator is None and _detector is None:
return False
stamp = time.monotonic() if now is None else float(now)
if stamp - _last_used < idle_seconds:
return False
_generator = None
_detector = None
logger.info("Idle timeout reached. Released the AudioSeal watermark models.")
return True
def is_enabled() -> bool:
"""Check if invisible watermarking is enabled in user preferences."""
return resolve("watermark.invisible", default=True) is not False
+21 -7
View File
@@ -249,16 +249,30 @@ kept, so turning it back on does not mean setting everything up again.
| `OMNIVOICE_INBOUND_PORT` | Port to accept them on (default `7444`) |
| `OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` | How long a model may sit unused before its VRAM is handed back (default `600`, minimum `5`) |
| `OMNIVOICE_IDLE_SWEEP_SECONDS` | How often that check runs (default `60`, minimum `1`) |
The last two exist so the ten-minute unload can be watched in a minute while
testing — set them together, since shortening only the threshold still means
waiting a full sweep interval to see it fire. Values that are unparseable or
below the floor are ignored with a warning rather than honoured: a zero
threshold would unload a model the instant it went idle and reload it for the
next request.
| `OMNIVOICE_WORKER_MODE` | `1` on the worker machine |
| `OMNIVOICE_WORKER_TOKEN` | Enrollment token, first run only |
`OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` and `OMNIVOICE_IDLE_SWEEP_SECONDS` exist
so the ten-minute unload can be watched in a minute while testing — set them
together, since shortening only the threshold still means waiting a full sweep
interval to see it fire. Values that are unparseable or below the floor are
ignored with a warning rather than honoured: a zero threshold would unload a
model the instant it went idle and reload it for the next request.
### Two idle timers, not one
A worker node runs the full app, so two independent reapers can release the
same model and they are configured separately:
| Timer | Default | Set with |
|---|---|---|
| Engine registry — drops the cached engine instance and, for VoiceStudio, the shared model with it | 600 s | `OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` |
| In-process model reaper — the backstop, also releases the dictation ASR and the watermark models | 900 s | `OMNIVOICE_IDLE_TIMEOUT` (or Settings) |
In practice the first one gets there first and the second finds nothing to do.
Shortening only `OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` is the right move when
testing; the backstop is not worth touching.
Only one VoiceStudio instance can accept remote workers on a given port. If
another instance already owns the configured port, the app continues running
with remote workers unavailable and shows the conflict in Settings. Close the
+103
View File
@@ -9,6 +9,7 @@ from __future__ import annotations
import asyncio
import os
import sys
import pytest
@@ -275,3 +276,105 @@ def test_list_models_downgrades_truncated_cache(tmp_path, monkeypatch):
assert row["installed"] is False
assert row["incomplete"] is True
models.invalidate_cache()
# ── The unload ordering (#1495) ─────────────────────────────────────────────
#
# Dropping the shared reference has to happen BEFORE the allocator caches are
# emptied. Inverted, the unload frees nothing and says it worked: the weights
# are still reachable so gc keeps them, empty_cache() only returns blocks the
# allocator already considered free, and the reference drops a moment later
# into a cache nobody will flush again. That is how a headless worker node held
# 3.6 GB across an idle sweep whose log line read "Released 1 idle engine(s)".
#
# Four call sites open-coded the pair; the one the engine-registry sweep reaches
# was the inverted one, which is why every UI-driven unload looked fine. These
# pin the ordering at the helper, at the sweep path, and at the facade — and the
# last test keeps new callers from open-coding it again.
class _Weights:
"""Stands in for the model. Identity is all these tests need."""
def _watch_free_vram(monkeypatch, manager=mm):
"""Record what the shared ref held at each free_vram() call.
``manager`` is explicit because other suites reimport
services.model_manager, so more than one module object can be alive at
once. Each caller here patches the exact object the code under test will
reach ``OmniVoiceBackend.unload`` imports at call time and gets whatever
sys.modules holds now, the facade uses the alias it bound at its own import.
Patching this file's alias for all of them passes alone and fails in a full
run, which is how this test first went red.
"""
seen: list = []
monkeypatch.setattr(manager, "free_vram", lambda: seen.append(manager.model))
monkeypatch.setattr(manager, "release_tts_side_caches", lambda: None)
monkeypatch.setattr(manager, "model", _Weights())
return seen
def test_unload_shared_model_clears_the_ref_before_freeing(monkeypatch):
seen = _watch_free_vram(monkeypatch)
assert mm.unload_shared_model() is True
assert seen == [None], "free_vram() ran while the model was still referenced"
assert mm.model is None
def test_unload_shared_model_is_idempotent(monkeypatch):
seen = _watch_free_vram(monkeypatch)
monkeypatch.setattr(mm, "model", None)
assert mm.unload_shared_model() is False
assert seen == [], "nothing was resident, so the allocator was left alone"
def test_engine_unload_releases_the_shared_model(monkeypatch):
"""The idle-sweep path — the one that was inverted."""
manager = sys.modules["services.model_manager"]
seen = _watch_free_vram(monkeypatch, manager)
monkeypatch.setattr(tb, "clear_clone_prompt_cache", lambda: None)
tb.OmniVoiceBackend().unload()
assert seen == [None], "the engine sweep emptied the cache before releasing"
assert manager.model is None
def test_facade_unload_tts_releases_the_shared_model(monkeypatch):
seen = _watch_free_vram(monkeypatch, ml.mm)
assert _run(ml.unload("tts")) == {"unloaded": "tts", "success": True}
assert seen == [None]
assert ml.mm.model is None
def test_no_caller_open_codes_the_shared_unload():
"""One ordering, in one place.
The bug was not that someone wrote the two lines wrongly it was that five
modules each wrote them at all, so getting one wrong stayed invisible next
to four that were right. Assigning ``model_manager.model`` from outside the
module is the shape that made that possible; ``unload_shared_model()`` is
the replacement.
"""
import pathlib
import re
root = pathlib.Path(mm.__file__).resolve().parents[1]
skip = {".venv", "venv", "site-packages", "node_modules", "__pycache__", "build", "dist"}
pattern = re.compile(r"^\w+\.model\s*=\s*(?!=)")
offenders = []
for path in root.rglob("*.py"):
if path.name == "model_manager.py" or skip & set(path.parts):
continue
for number, line in enumerate(path.read_text(encoding="utf-8", errors="ignore").splitlines(), 1):
stripped = line.strip()
if pattern.match(stripped) and ("mm." in stripped or "model_manager." in stripped):
offenders.append(f"{path.relative_to(root)}:{number}: {stripped}")
assert not offenders, (
"assign model_manager.model only inside model_manager; callers use "
"unload_shared_model(), which frees in the right order:\n "
+ "\n ".join(offenders)
)
+7 -2
View File
@@ -214,8 +214,13 @@ def test_shared_voice_unload_failure_is_stable(monkeypatch):
monkeypatch.setattr(manager, "free_vram", lambda: (_ for _ in ()).throw(RuntimeError("/secret/gpu")))
with pytest.raises(RuntimeError) as caught:
backend.unload()
assert str(caught.value) == "The shared voice model could not be unloaded. Retry after the current generation finishes."
assert manager.model is not None
assert str(caught.value) == "The voice model was released, but the GPU memory cache could not be flushed."
# The reference drop is the durable half and must survive a failing cache
# flush (#1495). This used to assert the opposite — that a raising
# free_vram() left the model bound for a retry — which was not a decision
# but a consequence of freeing before clearing, the very ordering that let
# a worker node hold 3.6 GB through an idle sweep.
assert manager.model is None
monkeypatch.setattr(manager, "free_vram", lambda: None)
backend.unload()
assert manager.model is None
+47
View File
@@ -121,3 +121,50 @@ def test_iter_chunks_covers_everything_exactly_once():
def test_iter_chunks_empty_audio_yields_nothing():
assert list(_iter_chunks(torch.zeros(1, 1, 0), SR)) == []
# ── Idle release of the AudioSeal models (#1495) ───────────────────────────
#
# They loaded on the first embed and then stayed resident for the life of the
# process — the one model in the app still making that bargain after the TTS
# model and the capture ASR both stopped. CPU-resident, so this is system RAM,
# and the machines that notice are the ones running batches.
def test_idle_release_drops_both_models(monkeypatch):
monkeypatch.setattr(watermark, "_generator", object())
monkeypatch.setattr(watermark, "_detector", object())
monkeypatch.setattr(watermark, "_last_used", 100.0)
assert watermark.release_idle_models(60.0, now=200.0) is True
assert watermark._generator is None
assert watermark._detector is None
def test_idle_release_keeps_a_recently_used_model(monkeypatch):
generator = object()
monkeypatch.setattr(watermark, "_generator", generator)
monkeypatch.setattr(watermark, "_detector", None)
monkeypatch.setattr(watermark, "_last_used", 100.0)
assert watermark.release_idle_models(60.0, now=130.0) is False
assert watermark._generator is generator
def test_idle_release_is_a_noop_when_nothing_loaded(monkeypatch):
monkeypatch.setattr(watermark, "_generator", None)
monkeypatch.setattr(watermark, "_detector", None)
monkeypatch.setattr(watermark, "_last_used", 0.0)
assert watermark.release_idle_models(0.0) is False
def test_embedding_restarts_the_idle_clock(monkeypatch):
"""Without this the models are released mid-batch: `_last_used` would sit
at whatever the first embed set it to while the batch kept watermarking."""
monkeypatch.setattr(watermark, "_generator", object())
monkeypatch.setattr(watermark, "_last_used", 0.0)
monkeypatch.setattr(watermark.time, "monotonic", lambda: 500.0)
watermark._get_generator()
assert watermark._last_used == 500.0