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 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))

85
wsai/backends/melo.py
View File

@@ -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

117
wsai/backends/whisper.py
View File

@@ -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