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:
+2
-1
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user