diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ee2e6c4..502e1647 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ the frozen-backend fallback mirror it for their toolchains. - The backend binds its port immediately and reports startup progress live — `/health` answers 503-with-step and a new `/startup/progress` endpoint lists every step while PyTorch, API routes, and database migrations load in the background, so "starting at step X" is never mistakable for "dead"; the desktop splash narrates each step (#1550) ### Added +- The locally cached AudioSeal watermark generator warms on a background thread ~35s after boot (`OMNIVOICE_PRELOAD_WATERMARK=0` opts out; explicitly setting `=1` may download it), so the first synthesis no longer serializes the audioseal import + model load inline — measured at ~42s on a cold filesystem, 3s short of a 90s client timeout (#1576) — thanks @paoloantinori! - Voices you've cloned stay "warm" across restarts — encoded references now persist to disk (~10 KB each), so the first generation of a session skips the re-encode and any transcription pass; `OMNIVOICE_PROMPT_DISK_CACHE=0` opts out (#1565) - Optional FlashInfer acceleration for the default engine on CUDA (`OMNIVOICE_FLASHINFER=1`, ~2.2x measured) — needs the optional `flashinfer-python` package; missing package or kernel failure logs why and falls back to the standard path (#1565) - The bug reporter notices when you're on an outdated build and offers the latest release before filing — with a "File anyway" escape hatch — and stamps a `Build status` line into every report so up-to-date reports are tellable from stale ones (#1547) @@ -31,6 +32,7 @@ the frozen-backend fallback mirror it for their toolchains. - The OmniVoice guide now covers combining style attributes with a reference clip (consistent instruct stabilizes cloning; the reference wins conflicts), inline pronunciation control (pinyin / CMU phonemes), and corrects the claim that the default engine can't do voice design — it can, from attributes (#1565) ### Fixed +- Watermark embedding failures now log the full traceback instead of just the exception message, so a silently-unmarked-audio incident (audio passes through unmarked by design) is diagnosable from the log alone (#1576) — thanks @paoloantinori! - Stored artifact subpaths now resolve after moving a data directory between Windows, macOS, Linux, and Docker, while traversal and symlink escapes remain blocked (#1559) — thanks @Eman-Yousaf! - A remote browser hitting an API-key-configured server's admin 403 now gets the API-key login form instead of endless console 403s, while desktop and PIN-only/no-key servers keep the plain loopback error so guests are never offered a login no key can satisfy (#1568) — thanks @paoloantinori! - The crash-isolated ASR sidecar and its download preflight now agree on which model to load — setting the shared faster-whisper model variable applies to both variants instead of the sidecar quietly using a different one (#1556) diff --git a/backend/api/routers/archetypes.py b/backend/api/routers/archetypes.py index 32142559..067be538 100644 --- a/backend/api/routers/archetypes.py +++ b/backend/api/routers/archetypes.py @@ -357,15 +357,11 @@ def _infer(seed: int): # Runs on the dedicated watermark pool (#1190): AudioSeal embedding is CPU # work that holds no VRAM, so it must not occupy a GPU worker ahead of the # next generate on 1-worker hosts. - from services.watermark import mark_synthetic - from services.model_manager import get_watermark_pool - import functools - audio_tensor = await run_on_gpu_pool_guarded( - functools.partial(mark_synthetic, audio_tensor, model.sampling_rate, - context="archetypes.render"), - what="Archetype watermark", + from services.watermark import mark_synthetic_async + audio_tensor = await mark_synthetic_async( + audio_tensor, model.sampling_rate, + context="archetypes.render", timeout=generate_timeout_s(""), - executor=get_watermark_pool(), ) out_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/backend/api/routers/batch.py b/backend/api/routers/batch.py index 20f05baa..debb1029 100644 --- a/backend/api/routers/batch.py +++ b/backend/api/routers/batch.py @@ -413,19 +413,15 @@ def _gen(text=seg_text, lang=target_lang, dur=seg_duration): # unmarked while the interactive dub pipeline marked every segment. # One whole-track embed (chunked internally, #1045) is equivalent to # dub_generate's per-segment marks: the 16-bit message repeats - # throughout. Runs in the GPU pool like generate's finalize; never - # raises (degrades to unmarked on failure, same as every producer). + # throughout. Never raises (degrades to unmarked on failure, same as + # every producer). # Dispatched to the dedicated watermark pool, not the GPU pool (#1190): # AudioSeal embedding is CPU work that holds no VRAM, and a whole-track # embed is long enough that occupying a GPU worker with it stalled the # next language's segments on 1-worker hosts. - from services.watermark import mark_synthetic - from services.model_manager import get_watermark_pool - import functools - full_audio = await loop.run_in_executor( - get_watermark_pool(), - functools.partial(mark_synthetic, full_audio, sr, - context="batch.dub_track"), + from services.watermark import mark_synthetic_async + full_audio = await mark_synthetic_async( + full_audio, sr, context="batch.dub_track", ) # Same assembly pattern as dub_generate.py:390 — `full_audio` is a diff --git a/backend/api/routers/generation.py b/backend/api/routers/generation.py index 51bb1db3..b95d4429 100644 --- a/backend/api/routers/generation.py +++ b/backend/api/routers/generation.py @@ -870,7 +870,6 @@ async def _finalize_generation( Returns ``(watermarked_tensor, meta)`` where ``meta`` carries ``id`` / ``filename`` / ``duration`` / ``gen_time``. """ - loop = asyncio.get_running_loop() # Invisible AudioSeal provenance watermark on the final audio. Embedding # was previously only wired into the dub pipeline (dub_generate.py), so # plain TTS came out unmarked despite the setting being on — and the same @@ -882,12 +881,9 @@ async def _finalize_generation( # AudioSeal embedding is CPU work that holds no VRAM, so occupying a GPU # worker with it only delays the next generate on 1-worker hosts. if not already_marked: - from services.watermark import mark_synthetic - from services.model_manager import get_watermark_pool - audio_tensor = await loop.run_in_executor( - get_watermark_pool(), - functools.partial(mark_synthetic, audio_tensor, sample_rate, - context="generate.finalize"), + from services.watermark import mark_synthetic_async + audio_tensor = await mark_synthetic_async( + audio_tensor, sample_rate, context="generate.finalize", ) gen_time = round(time.time() - start_time, 2) @@ -1822,12 +1818,10 @@ def _line(obj) -> bytes: # (#1190): AudioSeal embedding is CPU work that owns no # VRAM, and on a 1-worker host it used to serialize # directly ahead of the next generate. - from services.watermark import mark_synthetic - from services.model_manager import get_watermark_pool - _preview = await asyncio.get_running_loop().run_in_executor( - get_watermark_pool(), - functools.partial(mark_synthetic, audio_tensor, sample_rate, - context="generate.stream_preview"), + from services.watermark import mark_synthetic_async + _preview = await mark_synthetic_async( + audio_tensor, sample_rate, + context="generate.stream_preview", ) yield _line({"type": "chunk", "seq": 0, "pcm": _pcm16_b64(_preview)}) else: @@ -1849,12 +1843,10 @@ def _line(obj) -> bytes: # Provenance-mark the streamed copy off the GPU pool # (#1169 mark, #1190 placement): CPU-only AudioSeal # work must not occupy a GPU worker between chunks. - from services.watermark import mark_synthetic - from services.model_manager import get_watermark_pool - preview = await asyncio.get_running_loop().run_in_executor( - get_watermark_pool(), - functools.partial(mark_synthetic, preview, sample_rate, - context="generate.stream_preview"), + from services.watermark import mark_synthetic_async + preview = await mark_synthetic_async( + preview, sample_rate, + context="generate.stream_preview", ) if i == 0: # After the first render so lazy-loading engines diff --git a/backend/main.py b/backend/main.py index f2082700..3e2dce02 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,3 +1,4 @@ +import math import os import sys @@ -369,19 +370,36 @@ def _env_flag(name: str, default: bool = False) -> bool: _EAGER = _env_flag("OMNIVOICE_EAGER_INIT", default=("pytest" in sys.modules)) +def _env_float(name: str, default: float) -> float: + """Parse a float env override, rejecting negative and non-finite values. + + Shared by the preload-delay / timeout knobs: NaN would silently never + fire, a negative would fire during startup I/O, so both fall back to the + default instead (the bug class CodeRabbit flagged on the watermark knob + in PR #1577 — latent in the older copies too, closed here for all).""" + raw = os.environ.get(name, "") + try: + value = float(raw) if raw.strip() else default + except ValueError: + return default + return value if math.isfinite(value) and value >= 0 else default + + def _capture_preload_delay_s() -> float: """Seconds after boot before the dictation (capture ASR) model warms. Late enough that it never competes with startup I/O or the TTS preload; overridable via OMNIVOICE_CAPTURE_PRELOAD_DELAY (mostly for tests).""" - raw = os.environ.get("OMNIVOICE_CAPTURE_PRELOAD_DELAY", "") - try: - v = float(raw) - if v >= 0: - return v - except (TypeError, ValueError): - pass - return 30.0 + return _env_float("OMNIVOICE_CAPTURE_PRELOAD_DELAY", 30.0) + +def _watermark_preload_delay_s() -> float: + """Seconds after boot before the AudioSeal generator warm-up fires. + + Own knob, NOT ``_capture_preload_delay_s`` + offset: a capture-specific + env override must not retime the watermark warm too, and the two cold + imports shouldn't fire on the same tick (CodeRabbit, PR #1577). Default + 35s sits ~5s past the capture-ASR warm for the same reason.""" + return _env_float("OMNIVOICE_PRELOAD_WATERMARK_DELAY", 35.0) def _capture_preload_ram_ok(min_free_bytes: int = 4 * 1024**3) -> bool: @@ -398,14 +416,7 @@ def _capture_preload_ram_ok(min_free_bytes: int = 4 * 1024**3) -> bool: def _mcp_start_timeout_s() -> float: """Seconds to wait for the MCP session manager to start before giving up and serving without it (#632). Overridable via OMNIVOICE_MCP_START_TIMEOUT_S.""" - raw = os.environ.get("OMNIVOICE_MCP_START_TIMEOUT_S", "") - try: - v = float(raw) - if v > 0: - return v - except (TypeError, ValueError): - pass - return 30.0 + return max(_env_float("OMNIVOICE_MCP_START_TIMEOUT_S", 30.0), 0.001) async def _serve_mcp(session_manager, ready: "asyncio.Event", stop: "asyncio.Event") -> None: @@ -852,6 +863,8 @@ async def _phase_b(app: FastAPI) -> None: # #1174: arm model loads for THIS run — an in-process relaunch may carry a # stale shutting-down flag from a previous lifespan. model_loads_reset_shutdown() + from services.model_manager import begin_watermark_pool_lifecycle + begin_watermark_pool_lifecycle() app.state.idle_task = asyncio.create_task(idle_worker()) app.state.worker_task = asyncio.create_task(task_manager.worker()) # Warm the TTS model in the background so first /generate is instant. @@ -905,6 +918,50 @@ def _warm(): else: logger.info("Capture ASR preload disabled; dictation ASR will load on first use.") + # Watermark: warm the AudioSeal generator in the background so the first + # mark_synthetic doesn't serialize the audioseal import + model load + # inside the first synthesis (measured ~42 s inline on a cold filesystem, + # 2026-08-17 macOS report — 3 s short of the client's 90 s timeout). + # Small model on CPU; deferred a few seconds past the capture-ASR warm so + # the two cold imports don't contend for the same disk, and no RAM guard + # is needed. Runs on the watermark pool — where the model is used — not + # the shared default executor. + if _env_flag("OMNIVOICE_PRELOAD_WATERMARK", default=True): + async def _preload_watermark(): + await asyncio.sleep(_watermark_preload_delay_s()) + loop = asyncio.get_running_loop() + from services import watermark as _watermark + + # Gate BEFORE touching get_watermark_pool(): the pool is lazy so + # hosts with watermarking disabled never spawn its thread, and + # creating it unconditionally would break that invariant. The + # race with a first embed is benign — pool creation is itself + # lock-guarded. + if not _watermark.will_mark(): + logger.debug("Watermark preload skipped (disabled or audioseal absent)") + return + from services.model_manager import get_watermark_pool + + # Default startup may warm an existing local checkpoint but may + # not fetch one. Only an explicit user opt-in permits a download. + raw_preload = os.environ.get("OMNIVOICE_PRELOAD_WATERMARK", "") + allow_download = raw_preload.strip().lower() in {"1", "true", "yes", "on"} + + try: + await loop.run_in_executor( + get_watermark_pool(), + lambda: _watermark.prefetch_generator( + allow_download=allow_download + ), + ) + except Exception: + # prefetch_generator swallows its own errors; this guards the + # setup half (imports, pool construction) so a broken warm-up + # is visible now, not as an unretrieved exception at shutdown. + logger.warning("Watermark preload task failed", exc_info=True) + + app.state.watermark_preload_task = asyncio.create_task(_preload_watermark()) + # ── MCP session manager (Wave 2.2) ──────────────────────────────────── # Run it in its OWN task owning the full enter→exit lifecycle (anyio # task-affinity, see _serve_mcp); only wait, with a timeout, for ready — @@ -1080,8 +1137,20 @@ async def lifespan(app: FastAPI): getattr(app.state, "worker_task", None), getattr(app.state, "preload_task", None), getattr(app.state, "capture_preload_task", None), + getattr(app.state, "watermark_preload_task", None), timeout=20.0, ) + # The watermark warm-up runs on its dedicated 1-worker pool. Cancellation + # detaches the asyncio future but cannot kill a thread inside AudioSeal, + # so drain it fully before lifespan teardown reports completion. + try: + from services.model_manager import shutdown_watermark_pool as _wm_drain + + _wm_drain() + except Exception: + # Best-effort drain: a failure here must not abort the remaining + # shutdown steps (model unload, MCP teardown) below. + logger.warning("Watermark pool drain failed at shutdown", exc_info=True) # Unload the model and free GPU memory try: import services.model_manager as mm diff --git a/backend/services/model_manager.py b/backend/services/model_manager.py index 52781c8d..17dd9176 100644 --- a/backend/services/model_manager.py +++ b/backend/services/model_manager.py @@ -4,8 +4,9 @@ import time import asyncio import logging +import queue import threading -from concurrent.futures import ThreadPoolExecutor, Executor +from concurrent.futures import Executor, Future, ThreadPoolExecutor from utils.containment import contain_system_exit @@ -1057,21 +1058,166 @@ def _timeout_guidance( # doubling the effective queue depth of a streamed multi-chunk render. # Giving it its own tiny pool removes that head-of-line blocking with no VRAM # risk, because the work was never on the device to begin with. -_watermark_pool_singleton: "ThreadPoolExecutor | None" = None +_WATERMARK_STOP = object() + + +class _WatermarkExecutor(Executor): + """Single daemon worker with a bounded shutdown contract. + + ``ThreadPoolExecutor`` uses non-daemon workers that Python joins at exit, + so ``wait=False`` still delays process exit while ``wait=True`` can hang + lifespan teardown forever. AudioSeal loading is not cooperatively + cancellable; a daemon worker plus a bounded join is the only thread-based + contract that both preserves in-process model warm-up and guarantees exit. + """ + + def __init__(self) -> None: + self._items: queue.Queue = queue.Queue() + self._lock = threading.Lock() + self._shutdown = False + self._thread: threading.Thread | None = None + + def submit(self, fn, /, *args, **kwargs) -> Future: + future: Future = Future() + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + if self._thread is None: + self._thread = threading.Thread( + target=self._run, + name="watermark_0", + daemon=True, + ) + self._thread.start() + self._items.put((future, fn, args, kwargs)) + return future + + def _run(self) -> None: + while True: + item = self._items.get() + if item is _WATERMARK_STOP: + return + future, fn, args, kwargs = item + if not future.set_running_or_notify_cancel(): + continue + try: + future.set_result(fn(*args, **kwargs)) + except (Exception, SystemExit, KeyboardInterrupt) as exc: + future.set_exception(exc) + + def is_stopped(self) -> bool: + """Whether shutdown has completed and this executor can be replaced.""" + with self._lock: + return self._shutdown and ( + self._thread is None or not self._thread.is_alive() + ) + + def is_shutdown(self) -> bool: + with self._lock: + return self._shutdown + + def shutdown( + self, + wait: bool = True, + *, + cancel_futures: bool = False, + timeout: float | None = None, + ) -> bool: + with self._lock: + self._shutdown = True + thread = self._thread + if cancel_futures: + while True: + try: + item = self._items.get_nowait() + except queue.Empty: + break + if item is not _WATERMARK_STOP: + item[0].cancel() + self._items.put(_WATERMARK_STOP) + if wait and thread is not None: + thread.join(timeout=timeout) + return thread is None or not thread.is_alive() + + +_watermark_pool_singleton: "_WatermarkExecutor | None" = None _watermark_pool_lock = threading.Lock() +_watermark_pool_accepting = True + + +def begin_watermark_pool_lifecycle() -> None: + """Open watermark submissions for a newly-started app lifespan.""" + global _watermark_pool_accepting, _watermark_pool_singleton + with _watermark_pool_lock: + if ( + _watermark_pool_singleton is not None + and _watermark_pool_singleton.is_stopped() + ): + _watermark_pool_singleton = None + _watermark_pool_accepting = ( + _watermark_pool_singleton is None + or not _watermark_pool_singleton.is_shutdown() + ) -def get_watermark_pool() -> ThreadPoolExecutor: +def get_watermark_pool() -> _WatermarkExecutor: """Dedicated 1-worker pool for provenance marking. Built lazily so hosts - with watermarking disabled never spawn the thread.""" - global _watermark_pool_singleton - if _watermark_pool_singleton is None: - with _watermark_pool_lock: - if _watermark_pool_singleton is None: - _watermark_pool_singleton = ThreadPoolExecutor( - max_workers=1, thread_name_prefix="watermark", - ) - return _watermark_pool_singleton + with watermarking disabled never spawn the thread. + + The executor is captured and returned UNDER the lock: reading the global + again after an unlocked null-check could race shutdown_watermark_pool's + reset and hand out None (CodeRabbit, PR #1577).""" + global _watermark_pool_accepting, _watermark_pool_singleton + with _watermark_pool_lock: + if not _watermark_pool_accepting: + if ( + _watermark_pool_singleton is not None + and _watermark_pool_singleton.is_stopped() + ): + _watermark_pool_singleton = None + _watermark_pool_accepting = True + else: + raise RuntimeError("watermark executor is shutting down") + if ( + _watermark_pool_singleton is not None + and _watermark_pool_singleton.is_stopped() + ): + _watermark_pool_singleton = None + if _watermark_pool_singleton is None: + _watermark_pool_singleton = _WatermarkExecutor() + return _watermark_pool_singleton + + +def shutdown_watermark_pool(*, timeout: float = 20.0) -> None: + """Drain the watermark pool at app shutdown (PR #1577). + + Refuse queued work and wait for the active operation: Python cannot kill + a thread inside AudioSeal loading, so returning early would let model + initialization continue during interpreter teardown. The draining pool + remains published until its worker stops, preventing concurrent producers + from creating a replacement that escapes this shutdown. A process that + keeps running after lifespan shutdown (the test suite does exactly this) + gets a fresh pool once the old worker has actually stopped.""" + global _watermark_pool_accepting, _watermark_pool_singleton + with _watermark_pool_lock: + _watermark_pool_accepting = False + pool = _watermark_pool_singleton + if pool is not None: + stopped = pool.shutdown( + wait=True, + cancel_futures=True, + timeout=max(0.0, float(timeout)), + ) + if stopped: + with _watermark_pool_lock: + if _watermark_pool_singleton is pool: + _watermark_pool_singleton = None + else: + logger.warning( + "Watermark worker exceeded the %.1fs shutdown deadline; " + "abandoning its daemon thread", + timeout, + ) model = None # type: ignore diff --git a/backend/services/watermark.py b/backend/services/watermark.py index aef54655..98130ad0 100644 --- a/backend/services/watermark.py +++ b/backend/services/watermark.py @@ -22,10 +22,14 @@ import logging import math +import os +import threading import time -import torch +from pathlib import Path from typing import Optional +import torch + from core.prefs import resolve logger = logging.getLogger("omnivoice.watermark") @@ -37,6 +41,20 @@ _audioseal_available: Optional[bool] = None # Monotonic stamp of the last embed/detect, for the idle release below. _last_used = 0.0 +# Per-model locks for the lazy builds below: the startup prefetch thread +# races the first embed, and both must share ONE build (a double load doubles +# the cold-start cost the prefetch exists to hide). One lock PER MODEL — a +# single shared lock made the ~42s generator prefetch block unrelated detector +# loads and the idle reaper behind it. release_idle_models acquires both, in +# this fixed order (nothing else nests them, so no cycle is possible). +_generator_lock = threading.Lock() +_detector_lock = threading.Lock() + +# True when the generator exists ONLY because the startup prefetch built it +# and no embed/detect has used it since. The idle reaper grants one extra +# idle window before dropping such a model, so a first synthesis at minute +# 20 still finds it warm (code-review finding 2 on the prefetch PR). +_prefetched_unused = False # 16-bit message: "OM" in ASCII = 0x4F 0x4D = 0100_1111 0100_1101 # This is our signature — every VoiceStudio-generated audio carries it. @@ -77,28 +95,82 @@ def _check_available() -> bool: return _audioseal_available -def _get_generator(): - """Lazy-load the AudioSeal generator model.""" - global _generator, _last_used - _last_used = time.monotonic() - if _generator is None: - from audioseal import AudioSeal - _generator = AudioSeal.load_generator("audioseal_wm_16bits") - _generator.eval() - logger.info("AudioSeal generator loaded (16-bit message mode)") - return _generator +def _get_generator(mark_prefetched: bool = False): + """Lazy-load the AudioSeal generator model. + + Owns the idle-reaper grace in ONE critical section: the startup prefetch + claims it (``mark_prefetched=True``) only when THIS call builds the model, + and every other call (a real embed) consumes it — no call-site blocks, no + window between two lock scopes where the claim could land on an + already-used model. + """ + global _generator, _last_used, _prefetched_unused + with _generator_lock: + _last_used = time.monotonic() + if _generator is None: + from audioseal import AudioSeal + _generator = AudioSeal.load_generator("audioseal_wm_16bits") + _generator.eval() + logger.info("AudioSeal generator loaded (16-bit message mode)") + _prefetched_unused = mark_prefetched + elif not mark_prefetched: + _prefetched_unused = False + return _generator def _get_detector(): """Lazy-load the AudioSeal detector model.""" global _detector, _last_used - _last_used = time.monotonic() - if _detector is None: - from audioseal import AudioSeal - _detector = AudioSeal.load_detector("audioseal_detector_16bits") - _detector.eval() - logger.info("AudioSeal detector loaded (16-bit message mode)") - return _detector + with _detector_lock: + _last_used = time.monotonic() + if _detector is None: + from audioseal import AudioSeal + _detector = AudioSeal.load_detector("audioseal_detector_16bits") + _detector.eval() + logger.info("AudioSeal detector loaded (16-bit message mode)") + return _detector + + +def _generator_checkpoint_cached() -> bool: + """Return whether AudioSeal can warm without contacting Hugging Face. + + AudioSeal 0.2 stores the checkpoint in ``/audioseal`` even though + it uses huggingface_hub to fetch it. Keep startup local-first: an ordinary + boot may consume that file, but must never turn prefetch into a download. + """ + cache_root = os.environ.get("AUDIOSEAL_CACHE_DIR") or os.environ.get( + "XDG_CACHE_HOME" + ) + root = Path(cache_root).expanduser() if cache_root else Path.home() / ".cache" + return (root / "audioseal" / "generator_base.pth").is_file() + + +def prefetch_generator(*, allow_download: bool = False) -> None: + """Warm the AudioSeal generator eagerly (startup background thread). + + The first ``mark_synthetic`` otherwise pays the audioseal import plus the + generator load inline — measured at ~42 s on a cold filesystem (2026-08-17 + macOS deployment), serialized inside the first synthesis and 3 s short of + a 90 s client timeout. Warming here overlaps that span with the TTS model + load. No-op when watermarking is off or audioseal is absent; a failure + logs and leaves the lazy path to retry on first embed. Default startup is + also cache-only; a download is allowed only when the user explicitly set + ``OMNIVOICE_PRELOAD_WATERMARK=1``. + """ + try: + if not will_mark(): + logger.debug("Watermark prefetch skipped (disabled or audioseal absent)") + return + if not allow_download and not _generator_checkpoint_cached(): + logger.info("Watermark prefetch skipped: AudioSeal checkpoint is not cached") + return + _get_generator(mark_prefetched=True) + logger.info("AudioSeal generator prefetched in the background") + except Exception: + logger.warning( + "Watermark prefetch failed; the first embed will retry inline", + exc_info=True, + ) def release_idle_models(idle_seconds: float, *, now: Optional[float] = None) -> bool: @@ -114,14 +186,28 @@ def release_idle_models(idle_seconds: float, *, now: Optional[float] = None) -> 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 + global _generator, _detector, _prefetched_unused + with _generator_lock, _detector_lock: + 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 + if _prefetched_unused: + # The startup prefetch built the generator and nothing has used + # it yet. Drop the grace (one extra idle window only) instead of + # the model, so a first synthesis shortly after boot still finds + # it warm — the exact scenario the prefetch exists for. + _prefetched_unused = False + logger.info( + "Idle watermark models are prefetch-warmed but unused; " + "granting one more idle window before releasing." + ) + return False + # Under the locks so a release racing the prefetch or a first embed + # can't wipe a model the lazy path just built. + _generator = None + _detector = None logger.info("Idle timeout reached. Released the AudioSeal watermark models.") return True @@ -200,6 +286,36 @@ def mark_synthetic( return marked +async def mark_synthetic_async( + waveform: torch.Tensor, + sample_rate: int, + *, + context: str, + force: bool = False, + timeout: float | None = None, +) -> torch.Tensor: + """Dispatch marking without letting a draining pool lose finished audio.""" + import asyncio + import functools + + from services.model_manager import get_watermark_pool, run_on_gpu_pool_guarded + + try: + pool = get_watermark_pool() + except RuntimeError: + logger.warning("Watermark skipped while the prior worker is shutting down") + return waveform + + job = functools.partial( + mark_synthetic, waveform, sample_rate, context=context, force=force + ) + if timeout is not None: + return await run_on_gpu_pool_guarded( + job, what="Audio watermark", timeout=timeout, executor=pool + ) + return await asyncio.get_running_loop().run_in_executor(pool, job) + + @torch.no_grad() def embed_watermark( waveform: torch.Tensor, @@ -260,7 +376,7 @@ def embed_watermark( return watermarked except Exception as e: - logger.warning("Watermark embedding failed (passing through original): %s", e) + logger.warning("Watermark embedding failed (passing through original): %s", e, exc_info=True) return waveform @@ -337,7 +453,7 @@ def detect_watermark( } except Exception as e: - logger.warning("Watermark detection failed: %s", e) + logger.warning("Watermark detection failed: %s", e, exc_info=True) return { "is_watermarked": False, "confidence": 0.0, diff --git a/tests/conftest.py b/tests/conftest.py index 6e4c2b3c..ae3110cc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -45,6 +45,15 @@ # need a different value monkeypatch it explicitly. os.environ["OMNIVOICE_MODEL"] = "test" +# Background warm-ups must not fire mid-suite: many tests boot the app +# lifespan via TestClient, and any that exits without a lifespan shutdown +# leaves the deferred preload task pending — 35s later (mid-suite, in +# another thread) it loads real models and mutates watermark module state +# under whatever test happens to be running (seen as a CI-only flake in the +# prefetch cold-start tests). Unconditional: a stray export from the runner +# shell must not re-enable it; a test that wants the warm-up monkeypatches. +os.environ["OMNIVOICE_PRELOAD_WATERMARK"] = "0" + # ── Test fixtures ────────────────────────────────────────────────────────── @@ -53,6 +62,21 @@ import warnings as _warnings +# App lifespans deliberately close watermark admission during shutdown. The +# test process then keeps running and many route tests call handlers without +# starting another lifespan, so restore the equivalent of a fresh lifespan at +# every test boundary. ``begin_watermark_pool_lifecycle`` still refuses to +# reopen while a timed-out worker is genuinely alive, preserving that race +# signal instead of hiding it. +@pytest.fixture(autouse=True) +def _watermark_pool_lifecycle_baseline(): + model_manager = sys.modules.get("services.model_manager") + begin = getattr(model_manager, "begin_watermark_pool_lifecycle", None) + if callable(begin): + begin() + yield + + # ── torch default-dtype isolation (CI flaky trio) ─────────────────────────── # Three tests (test_effects_chain / test_generation_audio_guard / # test_persona_bundle) fail intermittently on CI — never locally — with diff --git a/tests/test_shutdown_preload_race_1000.py b/tests/test_shutdown_preload_race_1000.py index f4ff6a47..8b4d89fa 100644 --- a/tests/test_shutdown_preload_race_1000.py +++ b/tests/test_shutdown_preload_race_1000.py @@ -140,6 +140,7 @@ def test_production_shutdown_wait_is_generous_enough_for_a_cold_import(): r"getattr\(app\.state, \"worker_task\", None\),\s*" r"getattr\(app\.state, \"preload_task\", None\),\s*" r"getattr\(app\.state, \"capture_preload_task\", None\),\s*" + r"getattr\(app\.state, \"watermark_preload_task\", None\),\s*" r"timeout=([\d.]+),?\s*\)", src, ) diff --git a/tests/test_watermark_prefetch_coldstart.py b/tests/test_watermark_prefetch_coldstart.py new file mode 100644 index 00000000..dc212639 --- /dev/null +++ b/tests/test_watermark_prefetch_coldstart.py @@ -0,0 +1,381 @@ +"""Background warm-up + thread safety for the AudioSeal watermark models. + +The 2026-08-17 cold-start report on a macOS deployment: the first +``mark_synthetic`` serialized the audioseal import + generator load (~42 s on +a cold filesystem) INSIDE the first synthesis, and a 90 s client timeout +missed the audio by 3 s. The generator now warms on a background thread +during startup; because that thread races the first embed, the lazy getters +must be thread-safe — exactly one load, no torn state. +""" +from __future__ import annotations + +import importlib +import sys +import threading +import types +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def watermark(): + """Resolve app state after per-test setup, never at collection time.""" + return importlib.import_module("services.watermark") + + +@pytest.fixture(autouse=True) +def _reset_models(monkeypatch, watermark): + # Reset ALL lifecycle globals (CodeRabbit, PR #1577): a stale warm-up + # stamp or availability cache from a prior test changes this test's + # conditions. + watermark._generator = None + watermark._detector = None + watermark._last_used = 0.0 + watermark._prefetched_unused = False + monkeypatch.setattr(watermark, "_audioseal_available", None, raising=False) + yield + watermark._generator = None + watermark._detector = None + watermark._last_used = 0.0 + watermark._prefetched_unused = False + + +def _fake_audioseal(monkeypatch, load_s: float) -> list[int]: + """Install a fake ``audioseal`` module whose load blocks ``load_s`` and + records every invocation. The block is what makes a missing lock fail + reliably instead of winning the interleaving lottery.""" + calls: list[int] = [] + import time + + def _slow_load(name): + calls.append(1) + time.sleep(load_s) + return SimpleNamespace(eval=lambda: None) + + fake = types.ModuleType("audioseal") + fake.AudioSeal = SimpleNamespace(load_generator=_slow_load) + monkeypatch.setitem(sys.modules, "audioseal", fake) + return calls + + +def test_get_generator_loads_exactly_once_under_concurrency(monkeypatch, watermark): + """A background prefetch thread + the first embed race the lazy load; + both must share ONE generator build, not one each.""" + calls = _fake_audioseal(monkeypatch, load_s=0.05) + + results = [] + errors = [] + + def _hit(): + try: + results.append(watermark._get_generator()) + except Exception as exc: # pragma: no cover - surfaced by assertion + errors.append(exc) + + threads = [threading.Thread(target=_hit) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors + assert calls == [1], f"load_generator ran {len(calls)}x; the lazy load races" + assert all(g is results[0] for g in results) + + +def test_prefetch_generator_loads_when_watermarking_is_on(monkeypatch, watermark): + """prefetch_generator() must build the generator eagerly (the startup + warm-up path) while the pref is enabled and audioseal is importable.""" + calls = _fake_audioseal(monkeypatch, load_s=0) + monkeypatch.setattr(watermark, "is_enabled", lambda: True) + + watermark.prefetch_generator(allow_download=True) + + assert calls == [1] + assert watermark._generator is not None + + +def test_prefetch_generator_no_ops_when_disabled_or_absent(monkeypatch, watermark): + """Pref disabled, or audioseal not installed: the warm-up must touch + nothing — no import attempts, no model, no exception.""" + calls = _fake_audioseal(monkeypatch, load_s=0) + monkeypatch.setattr(watermark, "is_enabled", lambda: False) + watermark.prefetch_generator() + assert calls == [] + assert watermark._generator is None + + monkeypatch.setattr(watermark, "is_enabled", lambda: True) + monkeypatch.setattr(watermark, "_check_available", lambda: False) + watermark.prefetch_generator() + assert calls == [] + assert watermark._generator is None + + +def test_prefetch_generator_degrades_silently_on_failure(monkeypatch, watermark): + """A failed warm-up must never take the backend down or wedge the lazy + path: log, leave _generator None; the first embed retries inline.""" + monkeypatch.setattr(watermark, "is_enabled", lambda: True) + monkeypatch.setattr(watermark, "_check_available", lambda: True) + + def _boom(): + raise RuntimeError("hub exploded (test)") + + monkeypatch.setattr(watermark, "_get_generator", _boom) + watermark.prefetch_generator(allow_download=True) # must not raise + assert watermark._generator is None + + +def test_prefetch_generator_does_not_download_with_empty_offline_cache( + monkeypatch, tmp_path, watermark +): + """Default startup stays local-first even when watermarking is enabled.""" + monkeypatch.setenv("AUDIOSEAL_CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(watermark, "will_mark", lambda: True) + calls = [] + monkeypatch.setattr(watermark, "_get_generator", lambda **kw: calls.append(kw)) + + watermark.prefetch_generator() + + assert calls == [] + + +def test_prefetch_generator_loads_cached_checkpoint_offline( + monkeypatch, tmp_path, watermark +): + cache_dir = tmp_path / "audioseal" + cache_dir.mkdir() + (cache_dir / "generator_base.pth").touch() + monkeypatch.setenv("AUDIOSEAL_CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(watermark, "will_mark", lambda: True) + calls = [] + monkeypatch.setattr(watermark, "_get_generator", lambda **kw: calls.append(kw)) + + watermark.prefetch_generator() + + assert calls == [{"mark_prefetched": True}] + + +def test_detector_load_is_not_blocked_by_a_generator_prefetch(monkeypatch, watermark): + """Per-model locks (review finding): a ~42s generator build in the + prefetch thread must not stall an unrelated detector load — with the old + single shared lock, _get_detector queued behind the whole build.""" + gen_started = threading.Event() + gen_release = threading.Event() + det_done = threading.Event() + + fake = types.ModuleType("audioseal") + + def _slow_gen(name): + gen_started.set() + gen_release.wait(10) + return SimpleNamespace(eval=lambda: None) + + def _fast_det(name): + det_done.set() + return SimpleNamespace(eval=lambda: None) + + fake.AudioSeal = SimpleNamespace(load_generator=_slow_gen, load_detector=_fast_det) + monkeypatch.setitem(sys.modules, "audioseal", fake) + + t = threading.Thread(target=watermark._get_generator) + t.start() + assert gen_started.wait(5) + det = watermark._get_detector() # must NOT queue behind the generator build + assert det_done.wait(1), "detector load blocked behind the generator load" + gen_release.set() + t.join(timeout=10) + + +def test_watermark_pool_rebuilds_after_shutdown_drain(): + """The lifespan shutdown drains the watermark pool; a process that keeps + running afterwards (the test suite) must get a FRESH pool on next use, + not "cannot schedule new futures after shutdown" (CI, PR #1577).""" + from services.model_manager import ( + begin_watermark_pool_lifecycle, + get_watermark_pool, + shutdown_watermark_pool, + ) + + begin_watermark_pool_lifecycle() + pool_before = get_watermark_pool() + shutdown_watermark_pool() + with pytest.raises(RuntimeError): + # The drained pool refuses new work… + pool_before.submit(lambda: None).result(timeout=5) + # …but the next app lifespan hands out a live replacement. + begin_watermark_pool_lifecycle() + assert get_watermark_pool().submit(lambda: "ok").result(timeout=5) == "ok" + # Clean up the replacement so later tests start from a fresh pool too. + shutdown_watermark_pool() + + +def test_watermark_pool_shutdown_waits_for_active_worker(): + """Lifespan teardown cannot finish while AudioSeal is still loading.""" + from services.model_manager import ( + begin_watermark_pool_lifecycle, + get_watermark_pool, + shutdown_watermark_pool, + ) + + begin_watermark_pool_lifecycle() + started = threading.Event() + release = threading.Event() + shutdown_done = threading.Event() + + def _blocking_load(): + started.set() + release.wait(5) + + get_watermark_pool().submit(_blocking_load) + assert started.wait(1) + + shutdown_thread = threading.Thread( + target=lambda: (shutdown_watermark_pool(), shutdown_done.set()) + ) + shutdown_thread.start() + assert not shutdown_done.wait(0.1), "shutdown returned while worker was active" + release.set() + shutdown_thread.join(timeout=2) + assert shutdown_done.is_set() + + +def test_watermark_pool_shutdown_deadline_bounds_stuck_worker(): + """A stuck AudioSeal import cannot hang backend shutdown forever.""" + import time + + from services.model_manager import ( + begin_watermark_pool_lifecycle, + get_watermark_pool, + shutdown_watermark_pool, + ) + + started = threading.Event() + release = threading.Event() + + def _stuck_load(): + started.set() + release.wait(5) + + begin_watermark_pool_lifecycle() + pool = get_watermark_pool() + pool.submit(_stuck_load) + assert started.wait(1) + + before = time.monotonic() + shutdown_watermark_pool(timeout=0.05) + elapsed = time.monotonic() - before + + assert elapsed < 0.5 + assert pool._thread is not None and pool._thread.daemon + release.set() + pool._thread.join(timeout=1) + + +def test_watermark_pool_cannot_be_replaced_while_timed_out_worker_is_alive(watermark): + """A producer racing bounded shutdown cannot create an undrained pool.""" + from services.model_manager import ( + begin_watermark_pool_lifecycle, + get_watermark_pool, + shutdown_watermark_pool, + ) + + started = threading.Event() + release = threading.Event() + + def _stuck_load(): + started.set() + release.wait(5) + + begin_watermark_pool_lifecycle() + pool = get_watermark_pool() + pool.submit(_stuck_load) + assert started.wait(1) + + try: + shutdown_watermark_pool(timeout=0.01) + with pytest.raises(RuntimeError, match="shutting down"): + get_watermark_pool() + # A deliberate in-process relaunch must not overlap the retired + # worker: both touch the process-global AudioSeal model state. + begin_watermark_pool_lifecycle() + with pytest.raises(RuntimeError, match="shutting down"): + get_watermark_pool() + # Finished synthesis must still be returned unchanged: pool admission + # is outside mark_synthetic's synchronous fail-open boundary. + import asyncio + import torch + + audio = torch.zeros(1, 240) + marked = asyncio.run( + watermark.mark_synthetic_async( + audio, 24000, context="test.restart_during_shutdown" + ) + ) + assert marked is audio + finally: + release.set() + pool._thread.join(timeout=1) + + # Once the retired worker really exits, the reopened lifecycle becomes + # usable without requiring another startup signal. + replacement = get_watermark_pool() + assert replacement is not pool + assert replacement.submit(lambda: "ok").result(timeout=1) == "ok" + shutdown_watermark_pool() + + +def test_prefetched_model_gets_one_extra_idle_window(monkeypatch, watermark): + """Review finding: the reaper freed the prefetch-warmed, never-used + generator at the first idle tick, re-imposing the cold start the prefetch + exists to hide. It now survives ONE extra window; real use clears the + grace entirely. + + Each phase re-establishes its module state IMMEDIATELY before its + release_idle_models call and passes an explicit far-future ``now=``: a + leaked idle reaper (a test lifespan that exits without shutdown keeps + idle_worker running) mutates these same globals from another thread, and + re-stamping _last_used mid-test made the real-assertions flake on CI. + With preconditions set adjacent to each call and now pinned, an + interleaved tick cannot change the outcome. + """ + import time + + # Full isolation from a leaked idle reaper (idle_worker resolves + # watermark.release_idle_models per call): divert it to a no-op for the + # duration of this test, and call the real function via the saved ref. + real_release = watermark.release_idle_models + # The test owns the reaper's decisions for its duration. + monkeypatch.setattr(watermark, "release_idle_models", lambda *a, **k: False) + + far_future = time.monotonic() + 1_000_000 + + def _given(generator_set: bool, grace: bool): + watermark._generator = SimpleNamespace(eval=lambda: None) if generator_set else None + watermark._prefetched_unused = grace + watermark._last_used = 0.0 + + # First reaper pass on a prefetched-never-used model: grace, model kept. + _given(generator_set=True, grace=True) + assert real_release(900, now=far_future) is False + # Second pass: grace consumed, model released. + _given(generator_set=True, grace=False) + assert real_release(900, now=far_future) is True + # After grace was consumed, an idle model with no models at all is a no-op. + _given(generator_set=False, grace=False) + assert real_release(900, now=far_future) is False + + # Real use clears the grace: embed (even a failing one) resets the flag, + # so the next reaper pass releases without a second window. + import torch as _torch + monkeypatch.setattr(watermark, "is_enabled", lambda: True) + monkeypatch.setattr(watermark, "_check_available", lambda: True) + watermark._generator = SimpleNamespace(eval=lambda: None) + watermark._prefetched_unused = True + watermark.embed_watermark(_torch.zeros(1, 2400), 24000) + # The embed call itself must have cleared the grace — assert it, don't + # re-establish it, or a failing embed would pass unnoticed. + assert watermark._prefetched_unused is False + assert real_release(900, now=far_future) is True