From 9d133870e5a722fa91380b4804491856deaf7fec Mon Sep 17 00:00:00 2001 From: debpalash <4178343+debpalash@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:46:29 +0530 Subject: [PATCH 1/3] fix(watermark): isolate lifecycle state between tests --- CHANGELOG.md | 4 ++-- backend/services/model_manager.py | 9 ++++++++- tests/conftest.py | 15 +++++++++++++++ tests/test_watermark_prefetch_coldstart.py | 4 ---- 4 files changed, 25 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 591c66f2..502e1647 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +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 @paoloantinoro! +- 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) @@ -32,7 +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 @paoloantinoro! +- 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/services/model_manager.py b/backend/services/model_manager.py index e6910939..17dd9176 100644 --- a/backend/services/model_manager.py +++ b/backend/services/model_manager.py @@ -1112,6 +1112,10 @@ def is_stopped(self) -> bool: 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, @@ -1150,7 +1154,10 @@ def begin_watermark_pool_lifecycle() -> None: and _watermark_pool_singleton.is_stopped() ): _watermark_pool_singleton = None - _watermark_pool_accepting = _watermark_pool_singleton is None + _watermark_pool_accepting = ( + _watermark_pool_singleton is None + or not _watermark_pool_singleton.is_shutdown() + ) def get_watermark_pool() -> _WatermarkExecutor: diff --git a/tests/conftest.py b/tests/conftest.py index e059c1bf..ae3110cc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,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_watermark_prefetch_coldstart.py b/tests/test_watermark_prefetch_coldstart.py index 1696c5de..e4ef156e 100644 --- a/tests/test_watermark_prefetch_coldstart.py +++ b/tests/test_watermark_prefetch_coldstart.py @@ -210,7 +210,6 @@ def test_watermark_pool_rebuilds_after_shutdown_drain(): 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() - begin_watermark_pool_lifecycle() def test_watermark_pool_shutdown_waits_for_active_worker(): @@ -241,7 +240,6 @@ def _blocking_load(): release.set() shutdown_thread.join(timeout=2) assert shutdown_done.is_set() - begin_watermark_pool_lifecycle() def test_watermark_pool_shutdown_deadline_bounds_stuck_worker(): @@ -274,7 +272,6 @@ def _stuck_load(): assert pool._thread is not None and pool._thread.daemon release.set() pool._thread.join(timeout=1) - begin_watermark_pool_lifecycle() def test_watermark_pool_cannot_be_replaced_while_timed_out_worker_is_alive(): @@ -316,7 +313,6 @@ def _stuck_load(): assert replacement is not pool assert replacement.submit(lambda: "ok").result(timeout=1) == "ok" shutdown_watermark_pool() - begin_watermark_pool_lifecycle() def test_prefetched_model_gets_one_extra_idle_window(monkeypatch, watermark): From afe013a6bccd418e382b677f009129d0f78543ce Mon Sep 17 00:00:00 2001 From: debpalash <4178343+debpalash@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:47:45 +0530 Subject: [PATCH 2/3] fix(events): dispatch foreign loops through serving loop --- backend/core/event_bus.py | 26 ++++++--------- tests/test_event_bus_thread_emit.py | 52 +++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 15 deletions(-) diff --git a/backend/core/event_bus.py b/backend/core/event_bus.py index 74fd662c..a9242ed3 100644 --- a/backend/core/event_bus.py +++ b/backend/core/event_bus.py @@ -64,31 +64,27 @@ def emit(kind: str, payload: dict[str, Any] | None = None) -> None: **(payload or {}), } event_str = json.dumps(event) - loop = asyncio.get_running_loop() if _on_event_loop() else _serving_loop - if loop is None: + try: + caller_loop = asyncio.get_running_loop() + except RuntimeError: + caller_loop = None + target_loop = _serving_loop or caller_loop + if target_loop is None: # No serving loop yet — nobody to notify; dropping is correct. logger.debug("No event loop — event dropped: %s", kind) return try: - if _on_event_loop(): - loop.create_task(_broadcast(event_str)) + if caller_loop is target_loop: + target_loop.create_task(_broadcast(event_str)) else: - # Threadpool worker (sync endpoint body): the only thread-safe way in. - loop.call_soon_threadsafe(_schedule_broadcast, event_str) + # Sync endpoints and async producers on a foreign loop must both + # hand off: the lock and listener queues belong to serving_loop. + target_loop.call_soon_threadsafe(_schedule_broadcast, event_str) except RuntimeError: # The serving loop closed between capture and use (app shutdown). logger.debug("Event loop closed — event dropped: %s", kind) -def _on_event_loop() -> bool: - """True when called from the running event loop (async-context emit).""" - try: - asyncio.get_running_loop() - except RuntimeError: - return False - return True - - def _schedule_broadcast(event_str: str) -> None: """Run `_broadcast` on the serving loop; called via call_soon_threadsafe.""" asyncio.get_running_loop().create_task(_broadcast(event_str)) diff --git a/tests/test_event_bus_thread_emit.py b/tests/test_event_bus_thread_emit.py index 49b9960c..16002fc3 100644 --- a/tests/test_event_bus_thread_emit.py +++ b/tests/test_event_bus_thread_emit.py @@ -66,6 +66,58 @@ def sync_endpoint_body(): loop.close() +def test_emit_from_foreign_running_loop_reaches_serving_loop(bus): + """An async producer may run on a worker loop, but listener state belongs + to the WebSocket serving loop and must only be touched there.""" + serving_loop = asyncio.new_event_loop() + serving_loop.set_debug(True) + received: list[str] = [] + started = threading.Event() + done = threading.Event() + + async def serve_with_waiter_ready(): + q = await bus.subscribe() + waiter = asyncio.create_task(q.get()) + await asyncio.sleep(0) # q.get() has installed its serving-loop Future + started.set() + try: + received.append(await asyncio.wait_for(waiter, 0.5)) + except asyncio.TimeoutError: + pass + finally: + done.set() + await bus.unsubscribe(q) + + def run_serving_loop(): + asyncio.set_event_loop(serving_loop) + serving_loop.run_until_complete(serve_with_waiter_ready()) + + serving_thread = threading.Thread( + target=run_serving_loop, name="test-serving-loop" + ) + serving_thread.start() + try: + assert started.wait(2.0), "serving loop never subscribed" + + async def foreign_async_caller(): + assert asyncio.get_running_loop() is not serving_loop + bus.emit("profiles", {"action": "updated", "id": "foreign-loop"}) + await asyncio.sleep(0.05) + + asyncio.run(foreign_async_caller()) + + assert done.wait(2.0), ( + "emit() ran listener delivery on the caller's foreign loop" + ) + assert received, "foreign-loop event never reached the serving loop" + payload = json.loads(received[0]) + assert payload["id"] == "foreign-loop" + finally: + serving_loop.call_soon_threadsafe(done.set) + serving_thread.join(2.0) + serving_loop.close() + + async def _serve(bus, received: list[str], started: threading.Event, done: threading.Event): q = await bus.subscribe() started.set() From 93616a9c2aabb10586d2d015700ef131bdc793de Mon Sep 17 00:00:00 2001 From: debpalash <4178343+debpalash@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:53:42 +0530 Subject: [PATCH 3/3] fix(watermark): fail open while pool drains --- backend/api/routers/archetypes.py | 12 +++------ backend/api/routers/batch.py | 14 ++++------ backend/api/routers/generation.py | 30 ++++++++-------------- backend/services/watermark.py | 30 ++++++++++++++++++++++ tests/test_watermark_prefetch_coldstart.py | 14 +++++++++- 5 files changed, 63 insertions(+), 37 deletions(-) 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/services/watermark.py b/backend/services/watermark.py index 47b7c070..98130ad0 100644 --- a/backend/services/watermark.py +++ b/backend/services/watermark.py @@ -286,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, diff --git a/tests/test_watermark_prefetch_coldstart.py b/tests/test_watermark_prefetch_coldstart.py index e4ef156e..dc212639 100644 --- a/tests/test_watermark_prefetch_coldstart.py +++ b/tests/test_watermark_prefetch_coldstart.py @@ -274,7 +274,7 @@ def _stuck_load(): pool._thread.join(timeout=1) -def test_watermark_pool_cannot_be_replaced_while_timed_out_worker_is_alive(): +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, @@ -303,6 +303,18 @@ def _stuck_load(): 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)