From 6e20f2cd790b4ed7252243b4465694c73035eabc Mon Sep 17 00:00:00 2001 From: EJClaw Date: Sat, 22 Aug 2026 22:43:43 +0900 Subject: [PATCH] fix: serialize worker warmup handshake to stop concurrent stdout reads The STT/TTS worker _ensure() treated a spawned-but-not-yet-handshaked subprocess as ready, so a voice turn arriving during warmup read the same stdout StreamReader concurrently with the warmup handshake and crashed with "readuntil() called while another coroutine is already waiting for incoming data". Add a _start_lock + _ready flag so (re)start and the ready handshake run atomically and callers wait for real readiness before reading stdout. Co-Authored-By: Claude Opus 4.7 --- tests/test_whisper_stt.py | 73 ++++++++++++++++++++++++ wsai/backends/melo.py | 85 ++++++++++++++++----------- wsai/backends/whisper.py | 117 ++++++++++++++++++++++---------------- 3 files changed, 191 insertions(+), 84 deletions(-) diff --git a/tests/test_whisper_stt.py b/tests/test_whisper_stt.py index 6620da3..68d7241 100644 --- a/tests/test_whisper_stt.py +++ b/tests/test_whisper_stt.py @@ -7,6 +7,7 @@ behaves with no source wired yet. """ import asyncio +import json from typing import AsyncIterator from wsai.backends.whisper import WhisperSTT @@ -58,3 +59,75 @@ def test_empty_transcript_is_skipped(monkeypatch): utts = _collect(stt) assert [u.text for u in utts] == ["안녕"] + + +def test_request_during_warmup_does_not_overlap_stdout(monkeypatch): + """Regression: a transcribe() arriving while warmup() is still awaiting the + worker's ready line must NOT read the same stdout StreamReader concurrently. + + Before the fix, _ensure()'s fast path returned as soon as the subprocess was + spawned (proc set, returncode None) even though the ready handshake was still + in flight, so the request's stdout.readline() overlapped warmup's and asyncio + raised "readuntil() called while another coroutine is already waiting for + incoming data" — the exact crash seen in the Discord voice server.""" + + async def run(): + stt = WhisperSTT() + stdout = asyncio.StreamReader() + stderr = asyncio.StreamReader() + stderr.feed_eof() # nothing on stderr; let the drain task finish cleanly + + class FakeStdin: + def write(self, _b): + pass + + async def drain(self): + pass + + class FakeProc: + returncode = None + + def __init__(self): + self.stdin = FakeStdin() + self.stdout = stdout + self.stderr = stderr + + def terminate(self): + self.returncode = 0 + + async def wait(self): + return 0 + + spawns = [] + + async def fake_create(*_a, **_k): + spawns.append(1) + return FakeProc() + + monkeypatch.setattr(asyncio, "create_subprocess_exec", fake_create) + + # warmup enters _ensure and blocks awaiting the ready line on stdout. + warm = asyncio.create_task(stt.warmup()) + await asyncio.sleep(0.05) + + # A concurrent request lands mid-warmup. It must wait for readiness, not + # crash and not read stdout yet. + tr = asyncio.create_task(stt.transcribe("x.wav")) + await asyncio.sleep(0.05) + assert not tr.done() # blocked on the start lock, no overlapping read + + # Complete the handshake -> warmup finishes and releases the request. + stdout.feed_data( + (json.dumps({"ready": True, "ms": 1, "device": "cpu"}) + "\n").encode() + ) + await asyncio.wait_for(warm, timeout=1) + await asyncio.sleep(0.02) + stdout.feed_data( + (json.dumps({"ok": True, "text": "안녕", "ms": 2}) + "\n").encode() + ) + assert await asyncio.wait_for(tr, timeout=1) == "안녕" + assert sum(spawns) == 1 # one worker, not one-per-concurrent-caller + + await stt.aclose() + + asyncio.run(asyncio.wait_for(run(), timeout=5)) diff --git a/wsai/backends/melo.py b/wsai/backends/melo.py index 2d39c9a..d1ef5f6 100644 --- a/wsai/backends/melo.py +++ b/wsai/backends/melo.py @@ -97,6 +97,11 @@ class MeloTTS: self.sink = sink or _log_sink self._proc: asyncio.subprocess.Process | None = None self._lock = asyncio.Lock() + # Serialises worker (re)start + the ready handshake so a caller that + # arrives mid-warmup waits for readiness instead of reading the same + # stdout StreamReader concurrently (asyncio forbids overlapping reads). + self._start_lock = asyncio.Lock() + self._ready = False # True only after the ready handshake completes self._n = 0 self.load_ms: int | None = None # Keep the worker's most recent stderr lines so a crash reports its real @@ -132,42 +137,53 @@ class MeloTTS: await self._ensure() async def _ensure(self) -> None: - if self._proc is not None and self._proc.returncode is None: + # Fast path: only skip when the worker is not just spawned but fully + # handshaked. Checking `_proc` alone would let a caller sail past while + # another coroutine (e.g. warmup) is still awaiting the ready line on + # this same stdout, causing overlapping StreamReader reads. + if self._proc is not None and self._proc.returncode is None and self._ready: return - self.out_dir.mkdir(parents=True, exist_ok=True) - env = {**os.environ, "WSAI_MELO_DEVICE": self.device} - # Run the worker module from the wsai source tree with the melo venv. - repo_root = str(Path(__file__).resolve().parents[2]) - self._proc = await asyncio.create_subprocess_exec( - self.python, "-m", "wsai.backends.melo_worker", - cwd=repo_root, env=env, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - self._stderr_tail.clear() - assert self._proc.stderr is not None - self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) - ready = await self._proc.stdout.readline() - if not ready: # worker died before signalling ready - await self._proc.wait() - raise RuntimeError( - f"melo worker exited before ready (code {self._proc.returncode})." - f"{self._stderr_hint()}" + async with self._start_lock: + # Re-check under the lock: another coroutine may have finished the + # (re)start + handshake while we waited. + if self._proc is not None and self._proc.returncode is None and self._ready: + return + self._ready = False + self.out_dir.mkdir(parents=True, exist_ok=True) + env = {**os.environ, "WSAI_MELO_DEVICE": self.device} + # Run the worker module from the wsai source tree with the melo venv. + repo_root = str(Path(__file__).resolve().parents[2]) + self._proc = await asyncio.create_subprocess_exec( + self.python, "-m", "wsai.backends.melo_worker", + cwd=repo_root, env=env, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) - try: - info = json.loads(ready.decode()) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"melo worker sent invalid ready line {ready!r}: {exc}." - f"{self._stderr_hint()}" - ) from exc - if not info.get("ready"): - raise RuntimeError( - f"melo worker failed to start: {info}.{self._stderr_hint()}" - ) - self.load_ms = info.get("ms") - log.info("melo worker ready in %s ms on %s", self.load_ms, info.get("device")) + self._stderr_tail.clear() + assert self._proc.stderr is not None + self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) + ready = await self._proc.stdout.readline() + if not ready: # worker died before signalling ready + await self._proc.wait() + raise RuntimeError( + f"melo worker exited before ready (code {self._proc.returncode})." + f"{self._stderr_hint()}" + ) + try: + info = json.loads(ready.decode()) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"melo worker sent invalid ready line {ready!r}: {exc}." + f"{self._stderr_hint()}" + ) from exc + if not info.get("ready"): + raise RuntimeError( + f"melo worker failed to start: {info}.{self._stderr_hint()}" + ) + self.load_ms = info.get("ms") + self._ready = True + log.info("melo worker ready in %s ms on %s", self.load_ms, info.get("device")) async def synth(self, text: str) -> str: """Synthesize `text` to a wav and return its path (no sink). Reusable by @@ -223,3 +239,4 @@ class MeloTTS: pass self._stderr_task = None self._proc = None + self._ready = False diff --git a/wsai/backends/whisper.py b/wsai/backends/whisper.py index ce03d28..265fbd4 100644 --- a/wsai/backends/whisper.py +++ b/wsai/backends/whisper.py @@ -75,6 +75,11 @@ class WhisperSTT: self.audio_source = audio_source self._proc: asyncio.subprocess.Process | None = None self._lock = asyncio.Lock() + # Serialises worker (re)start + the ready handshake so a caller that + # arrives mid-warmup waits for readiness instead of reading the same + # stdout StreamReader concurrently (asyncio forbids overlapping reads). + self._start_lock = asyncio.Lock() + self._ready = False # True only after the ready handshake completes self.load_ms: int | None = None self.resolved_device: str | None = None # "cuda" | "cpu", known after start # Keep the worker's most recent stderr so a crash reports its real cause @@ -106,59 +111,70 @@ class WhisperSTT: await self._ensure() async def _ensure(self) -> None: - if self._proc is not None and self._proc.returncode is None: + # Fast path: only skip when the worker is not just spawned but fully + # handshaked. Checking `_proc` alone would let a caller sail past while + # another coroutine (e.g. warmup) is still awaiting the ready line on + # this same stdout, causing overlapping StreamReader reads. + if self._proc is not None and self._proc.returncode is None and self._ready: return - env = { - **os.environ, - "WSAI_WHISPER_MODEL": self.model, - "WSAI_WHISPER_DEVICE": self.device, - } - # ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia - # pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured - # at exec, so inject those lib dirs into the child env here (harmless on - # CPU). Without this the CUDA model loads but transcribe() dies with - # "Library libcublas.so.12 is not found". - lib_dirs = _cuda_lib_dirs(self.python) - if lib_dirs: - prev = env.get("LD_LIBRARY_PATH", "") - env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else [])) - if self.language: - env["WSAI_WHISPER_LANGUAGE"] = self.language - repo_root = str(Path(__file__).resolve().parents[2]) - self._proc = await asyncio.create_subprocess_exec( - self.python, "-m", "wsai.backends.whisper_worker", - cwd=repo_root, env=env, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - self._stderr_tail.clear() - assert self._proc.stderr is not None - self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) - ready = await self._proc.stdout.readline() - if not ready: # worker died before signalling ready - await self._proc.wait() - raise RuntimeError( - f"whisper worker exited before ready (code {self._proc.returncode})." - f"{self._stderr_hint()}" + async with self._start_lock: + # Re-check under the lock: another coroutine may have finished the + # (re)start + handshake while we waited. + if self._proc is not None and self._proc.returncode is None and self._ready: + return + self._ready = False + env = { + **os.environ, + "WSAI_WHISPER_MODEL": self.model, + "WSAI_WHISPER_DEVICE": self.device, + } + # ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia + # pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured + # at exec, so inject those lib dirs into the child env here (harmless on + # CPU). Without this the CUDA model loads but transcribe() dies with + # "Library libcublas.so.12 is not found". + lib_dirs = _cuda_lib_dirs(self.python) + if lib_dirs: + prev = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else [])) + if self.language: + env["WSAI_WHISPER_LANGUAGE"] = self.language + repo_root = str(Path(__file__).resolve().parents[2]) + self._proc = await asyncio.create_subprocess_exec( + self.python, "-m", "wsai.backends.whisper_worker", + cwd=repo_root, env=env, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) - try: - info = json.loads(ready.decode()) - except json.JSONDecodeError as exc: - raise RuntimeError( - f"whisper worker sent invalid ready line {ready!r}: {exc}." - f"{self._stderr_hint()}" - ) from exc - if not info.get("ready"): - raise RuntimeError( - f"whisper worker failed to start: {info}.{self._stderr_hint()}" + self._stderr_tail.clear() + assert self._proc.stderr is not None + self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) + ready = await self._proc.stdout.readline() + if not ready: # worker died before signalling ready + await self._proc.wait() + raise RuntimeError( + f"whisper worker exited before ready (code {self._proc.returncode})." + f"{self._stderr_hint()}" + ) + try: + info = json.loads(ready.decode()) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"whisper worker sent invalid ready line {ready!r}: {exc}." + f"{self._stderr_hint()}" + ) from exc + if not info.get("ready"): + raise RuntimeError( + f"whisper worker failed to start: {info}.{self._stderr_hint()}" + ) + self.load_ms = info.get("ms") + self.resolved_device = info.get("device") + self._ready = True + log.info( + "whisper worker ready in %s ms on %s (model %s)", + self.load_ms, info.get("device"), info.get("model"), ) - self.load_ms = info.get("ms") - self.resolved_device = info.get("device") - log.info( - "whisper worker ready in %s ms on %s (model %s)", - self.load_ms, info.get("device"), info.get("model"), - ) async def transcribe(self, wav_path: str, *, language: str | None = None) -> str: """Transcribe one wav file to text using the warm worker.""" @@ -211,3 +227,4 @@ class WhisperSTT: pass self._stderr_task = None self._proc = None + self._ready = False