tkrmagid/watch_sceen_ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: serialize worker warmup handshake to stop concurrent stdout reads

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
EJClaw
2026-08-22 22:43:43 +09:00
parent 2ea7d04289
commit 6e20f2cd79
3 changed files with 191 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
tests/test_whisper_stt.py
View File

@@ -7,6 +7,7 @@ behaves with no source wired yet.
""" """
import asyncio import asyncio
import json
from typing import AsyncIterator from typing import AsyncIterator
from wsai.backends.whisper import WhisperSTT from wsai.backends.whisper import WhisperSTT
@@ -58,3 +59,75 @@ def test_empty_transcript_is_skipped(monkeypatch):
utts = _collect(stt) utts = _collect(stt)
assert [u.text for u in utts] == ["안녕"] 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))

85
wsai/backends/melo.py
View File

@@ -97,6 +97,11 @@ class MeloTTS:
self.sink = sink or _log_sink self.sink = sink or _log_sink
self._proc: asyncio.subprocess.Process | None = None self._proc: asyncio.subprocess.Process | None = None
self._lock = asyncio.Lock() 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._n = 0
self.load_ms: int | None = None self.load_ms: int | None = None
# Keep the worker's most recent stderr lines so a crash reports its real # Keep the worker's most recent stderr lines so a crash reports its real
@@ -132,42 +137,53 @@ class MeloTTS:
await self._ensure() await self._ensure()
async def _ensure(self) -> None: 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 return
self.out_dir.mkdir(parents=True, exist_ok=True) async with self._start_lock:
env = {**os.environ, "WSAI_MELO_DEVICE": self.device} # Re-check under the lock: another coroutine may have finished the
# Run the worker module from the wsai source tree with the melo venv. # (re)start + handshake while we waited.
repo_root = str(Path(__file__).resolve().parents[2]) if self._proc is not None and self._proc.returncode is None and self._ready:
self._proc = await asyncio.create_subprocess_exec( return
self.python, "-m", "wsai.backends.melo_worker", self._ready = False
cwd=repo_root, env=env, self.out_dir.mkdir(parents=True, exist_ok=True)
stdin=asyncio.subprocess.PIPE, env = {**os.environ, "WSAI_MELO_DEVICE": self.device}
stdout=asyncio.subprocess.PIPE, # Run the worker module from the wsai source tree with the melo venv.
stderr=asyncio.subprocess.PIPE, repo_root = str(Path(__file__).resolve().parents[2])
) self._proc = await asyncio.create_subprocess_exec(
self._stderr_tail.clear() self.python, "-m", "wsai.backends.melo_worker",
assert self._proc.stderr is not None cwd=repo_root, env=env,
self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) stdin=asyncio.subprocess.PIPE,
ready = await self._proc.stdout.readline() stdout=asyncio.subprocess.PIPE,
if not ready: # worker died before signalling ready stderr=asyncio.subprocess.PIPE,
await self._proc.wait()
raise RuntimeError(
f"melo worker exited before ready (code {self._proc.returncode})."
f"{self._stderr_hint()}"
) )
try: self._stderr_tail.clear()
info = json.loads(ready.decode()) assert self._proc.stderr is not None
except json.JSONDecodeError as exc: self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr))
raise RuntimeError( ready = await self._proc.stdout.readline()
f"melo worker sent invalid ready line {ready!r}: {exc}." if not ready: # worker died before signalling ready
f"{self._stderr_hint()}" await self._proc.wait()
) from exc raise RuntimeError(
if not info.get("ready"): f"melo worker exited before ready (code {self._proc.returncode})."
raise RuntimeError( f"{self._stderr_hint()}"
f"melo worker failed to start: {info}.{self._stderr_hint()}" )
) try:
self.load_ms = info.get("ms") info = json.loads(ready.decode())
log.info("melo worker ready in %s ms on %s", self.load_ms, info.get("device")) 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: async def synth(self, text: str) -> str:
"""Synthesize `text` to a wav and return its path (no sink). Reusable by """Synthesize `text` to a wav and return its path (no sink). Reusable by
@@ -223,3 +239,4 @@ class MeloTTS:
pass pass
self._stderr_task = None self._stderr_task = None
self._proc = None self._proc = None
self._ready = False

117
wsai/backends/whisper.py
View File

@@ -75,6 +75,11 @@ class WhisperSTT:
self.audio_source = audio_source self.audio_source = audio_source
self._proc: asyncio.subprocess.Process | None = None self._proc: asyncio.subprocess.Process | None = None
self._lock = asyncio.Lock() 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.load_ms: int | None = None
self.resolved_device: str | None = None # "cuda" | "cpu", known after start 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 # Keep the worker's most recent stderr so a crash reports its real cause
@@ -106,59 +111,70 @@ class WhisperSTT:
await self._ensure() await self._ensure()
async def _ensure(self) -> None: 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 return
env = { async with self._start_lock:
**os.environ, # Re-check under the lock: another coroutine may have finished the
"WSAI_WHISPER_MODEL": self.model, # (re)start + handshake while we waited.
"WSAI_WHISPER_DEVICE": self.device, if self._proc is not None and self._proc.returncode is None and self._ready:
} return
# ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia self._ready = False
# pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured env = {
# at exec, so inject those lib dirs into the child env here (harmless on **os.environ,
# CPU). Without this the CUDA model loads but transcribe() dies with "WSAI_WHISPER_MODEL": self.model,
# "Library libcublas.so.12 is not found". "WSAI_WHISPER_DEVICE": self.device,
lib_dirs = _cuda_lib_dirs(self.python) }
if lib_dirs: # ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia
prev = env.get("LD_LIBRARY_PATH", "") # pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured
env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else [])) # at exec, so inject those lib dirs into the child env here (harmless on
if self.language: # CPU). Without this the CUDA model loads but transcribe() dies with
env["WSAI_WHISPER_LANGUAGE"] = self.language # "Library libcublas.so.12 is not found".
repo_root = str(Path(__file__).resolve().parents[2]) lib_dirs = _cuda_lib_dirs(self.python)
self._proc = await asyncio.create_subprocess_exec( if lib_dirs:
self.python, "-m", "wsai.backends.whisper_worker", prev = env.get("LD_LIBRARY_PATH", "")
cwd=repo_root, env=env, env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else []))
stdin=asyncio.subprocess.PIPE, if self.language:
stdout=asyncio.subprocess.PIPE, env["WSAI_WHISPER_LANGUAGE"] = self.language
stderr=asyncio.subprocess.PIPE, repo_root = str(Path(__file__).resolve().parents[2])
) self._proc = await asyncio.create_subprocess_exec(
self._stderr_tail.clear() self.python, "-m", "wsai.backends.whisper_worker",
assert self._proc.stderr is not None cwd=repo_root, env=env,
self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr)) stdin=asyncio.subprocess.PIPE,
ready = await self._proc.stdout.readline() stdout=asyncio.subprocess.PIPE,
if not ready: # worker died before signalling ready stderr=asyncio.subprocess.PIPE,
await self._proc.wait()
raise RuntimeError(
f"whisper worker exited before ready (code {self._proc.returncode})."
f"{self._stderr_hint()}"
) )
try: self._stderr_tail.clear()
info = json.loads(ready.decode()) assert self._proc.stderr is not None
except json.JSONDecodeError as exc: self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr))
raise RuntimeError( ready = await self._proc.stdout.readline()
f"whisper worker sent invalid ready line {ready!r}: {exc}." if not ready: # worker died before signalling ready
f"{self._stderr_hint()}" await self._proc.wait()
) from exc raise RuntimeError(
if not info.get("ready"): f"whisper worker exited before ready (code {self._proc.returncode})."
raise RuntimeError( f"{self._stderr_hint()}"
f"whisper worker failed to start: {info}.{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: async def transcribe(self, wav_path: str, *, language: str | None = None) -> str:
"""Transcribe one wav file to text using the warm worker.""" """Transcribe one wav file to text using the warm worker."""
@@ -211,3 +227,4 @@ class WhisperSTT:
pass pass
self._stderr_task = None self._stderr_task = None
self._proc = None self._proc = None
self._ready = False