diff --git a/backend/api/routers/dub_core.py b/backend/api/routers/dub_core.py index 11049b83..fcff6c2f 100644 --- a/backend/api/routers/dub_core.py +++ b/backend/api/routers/dub_core.py @@ -1920,7 +1920,7 @@ async def dub_transcribe_stream( "heuristic", ) - fut_diar = loop.run_in_executor(_gpu_pool, _diarize) + fut_diar = loop.run_in_executor(_gpu_pool, lambda: _asr_work.run(_diarize)) async for _ping in _ping_while(fut_diar): yield _ping final_segs, diar_warning, labels_source = fut_diar.result() diff --git a/backend/api/routers/dub_export.py b/backend/api/routers/dub_export.py index 559ad56a..99c2e967 100644 --- a/backend/api/routers/dub_export.py +++ b/backend/api/routers/dub_export.py @@ -291,7 +291,10 @@ async def stream_task(task_id: str, after_seq: int = 0): finally: await task_manager.remove_listener(task_id, q) - return StreamingResponse(_reader(), media_type="text/event-stream") + return StreamingResponse( + _reader(), media_type="text/event-stream", + headers={"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"}, + ) @router.get("/jobs") diff --git a/docs/dubbing/translation-engines.md b/docs/dubbing/translation-engines.md index 8d7b0e19..a93b3806 100644 --- a/docs/dubbing/translation-engines.md +++ b/docs/dubbing/translation-engines.md @@ -235,4 +235,4 @@ panels instead so neither editor becomes unusably small. - **Installed it but still "needs install"** — restart the backend so Python picks up the newly-installed module. -Dubbing transcription emits keepalives during quiet diarization, reference-refinement, and cleanup steps. Disconnecting stops queued model work; native calls already running retain their model until they finish, then cleanup restores TTS. +Dubbing transcription emits keepalives during quiet diarization, reference-refinement, and cleanup steps. Disconnecting stops queued model work; native calls already running retain their model until they finish, then cleanup restores TTS. Task streams also request that proxies disable buffering so keepalives reach the client promptly. diff --git a/tests/test_dub_transcribe.py b/tests/test_dub_transcribe.py index 02092ad9..2df9064d 100644 --- a/tests/test_dub_transcribe.py +++ b/tests/test_dub_transcribe.py @@ -1024,3 +1024,69 @@ def test_stream_unload_is_single_shot_during_disconnect(): normal.result(timeout=5) disconnected.result(timeout=5) assert calls == ["unload"] + + +def test_disconnect_during_diarization_waits_before_unload_and_restore(tmp_path, monkeypatch): + import asyncio + import threading + from concurrent.futures import ThreadPoolExecutor + from api.routers import dub_core as dc + from services import asr_backend + + started, release, cleanup_started, restored = [threading.Event() for _ in range(4)] + events, guarded = [], [] + original_cleanup = dc._ASRWorkLifetime.cleanup + def cleanup(lifetime, fn): + guarded.append(lifetime._lock.locked()) + cleanup_started.set() + return original_cleanup(lifetime, fn) + monkeypatch.setattr(dc._ASRWorkLifetime, "cleanup", cleanup) + def diarize(**kwargs): + started.set() + assert release.wait(5) + events.append("diarization finished") + return None, None + class Backend: + id = "fake" + def ensure_loaded(self): pass + def transcribe(self, *a, **kw): + return {"chunks": [{"text": "hi", "timestamp": (0., .5)}], "segments": [], "language": "en"} + def unload(self): events.append("unloaded") + async def model(): + result = MagicMock() + result._asr_pipe = MagicMock() + return result + def restore(): + events.append("restored") + restored.set() + monkeypatch.setattr(dc, "get_model", model) + monkeypatch.setattr(asr_backend, "get_active_asr_backend", lambda *a, **kw: Backend()) + monkeypatch.setattr(dc, "get_diarization_pipeline", diarize) + monkeypatch.setattr(dc, "offload_tts_for_asr", lambda: None) + monkeypatch.setattr(dc, "restore_tts_after_asr", restore) + monkeypatch.setattr(dc, "POST_ASR_PING_S", .01) + audio = tmp_path / "diar.wav" + _make_wav(audio, seconds=1.) + job_id = "diar_disconnect" + dc._dub_jobs[job_id] = {"audio_path": str(audio), "vocals_path": None, "scene_cuts": []} + async def scenario(): + response = await dc.dub_transcribe_stream(job_id) + try: + async for _ in response.body_iterator: + if started.is_set(): + break + await response.body_iterator.aclose() + assert await asyncio.to_thread(cleanup_started.wait, 5) + assert guarded == [True] + assert not restored.is_set() + finally: + release.set() + await response.body_iterator.aclose() + assert await asyncio.to_thread(restored.wait, 5) + try: + with ThreadPoolExecutor(max_workers=2) as pool: + monkeypatch.setattr(dc, "_gpu_pool", pool) + asyncio.run(asyncio.wait_for(scenario(), timeout=10)) + finally: + dc._dub_jobs.pop(job_id, None) + assert events == ["diarization finished", "unloaded", "restored"] diff --git a/tests/test_tasks_stream_keepalive.py b/tests/test_tasks_stream_keepalive.py index d7d35135..24008f12 100644 --- a/tests/test_tasks_stream_keepalive.py +++ b/tests/test_tasks_stream_keepalive.py @@ -28,6 +28,8 @@ def test_quiet_task_stream_emits_keepalive_comments(monkeypatch): "error": None, "cancelled": False, } resp = await de.stream_task(task_id) + assert resp.headers["cache-control"] == "no-cache, no-transform" + assert resp.headers["x-accel-buffering"] == "no" frames = [] async def _read():