diff --git a/backend/api/routers/dub_core.py b/backend/api/routers/dub_core.py index 448a6044..87c483f8 100644 --- a/backend/api/routers/dub_core.py +++ b/backend/api/routers/dub_core.py @@ -886,6 +886,7 @@ class _ASRWorkLifetime: import threading self._lock = threading.Lock() self._closed = threading.Event() + self._cleaned = False def run(self, fn): with self._lock: @@ -900,6 +901,10 @@ class _ASRWorkLifetime: def cleanup(self, fn): self.stop() with self._lock: + if self._cleaned: + return + # Claim cleanup before invoking even a non-idempotent unload. + self._cleaned = True return fn() diff --git a/tests/test_dub_transcribe.py b/tests/test_dub_transcribe.py index e6a1175a..02092ad9 100644 --- a/tests/test_dub_transcribe.py +++ b/tests/test_dub_transcribe.py @@ -974,6 +974,7 @@ def test_stream_cleanup_waits_for_native_work_and_rejects_late_work(): from api.routers.dub_core import _ASRWorkLifetime lifetime = _ASRWorkLifetime() started, release, cleaned = threading.Event(), threading.Event(), threading.Event() + cleanup_started = threading.Event() events = [] def native(): started.set() @@ -986,9 +987,13 @@ def test_stream_cleanup_waits_for_native_work_and_rejects_late_work(): work = pool.submit(lifetime.run, native) assert started.wait(5) lifetime.stop() - removal = pool.submit(lifetime.cleanup, cleanup) + def remove(): + cleanup_started.set() + lifetime.cleanup(cleanup) + removal = pool.submit(remove) try: - assert not cleaned.wait(0.05) + assert cleanup_started.wait(5) + assert not cleaned.is_set() finally: release.set() work.result(timeout=5) @@ -996,3 +1001,26 @@ def test_stream_cleanup_waits_for_native_work_and_rejects_late_work(): assert events == ["native finished", "unloaded"] with pytest.raises(RuntimeError, match="stream has ended"): lifetime.run(lambda: pytest.fail("late work accessed unloaded model")) + + +def test_stream_unload_is_single_shot_during_disconnect(): + import threading + from concurrent.futures import ThreadPoolExecutor + from api.routers.dub_core import _ASRWorkLifetime + lifetime = _ASRWorkLifetime() + started, release = threading.Event(), threading.Event() + calls = [] + def unload(): + calls.append("unload") + assert len(calls) == 1 + started.set() + assert release.wait(5) + with ThreadPoolExecutor(max_workers=2) as pool: + normal = pool.submit(lifetime.cleanup, unload) + assert started.wait(5) + lifetime.stop() + disconnected = pool.submit(lifetime.cleanup, unload) + release.set() + normal.result(timeout=5) + disconnected.result(timeout=5) + assert calls == ["unload"]