From b2bfceb4fcff0f92c80fa68f4ef9848209a3794d Mon Sep 17 00:00:00 2001 From: Malcolm Habeeb Date: Mon, 27 Jul 2026 00:51:09 -0400 Subject: [PATCH 1/2] fix(stt): never block the caption path on the stream socket (XERK-128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the first cut, which only warmed the Nemotron stream and so did nothing when the stream itself is the problem — which is the actual case here. The deployment runs the hybrid backend (stream partials + Parakeet finals), but the Nemotron streaming server is unreachable/misconfigured (the host port that should serve it answers as an unrelated HTTP LLM). In that state the old code opened the stream *lazily on the first audio chunk*, inline on the caption path, and — because the api reads audio frames serially — stalled every frame behind the connect. Measured against a hung handshake, `open()` blocks for the whole `open_timeout`, which shipped at 15s. So each session stalled for up to 15s with no captions while audio piled up, then flushed the whole backlog at once: exactly "a delay, then all the words since the convo started". Re-decode partials were also suppressed the entire time because the stream still counted as "active". Meanwhile the offline Parakeet path is fast and healthy (~110-140ms per decode), so it can carry the live captions on its own with no perceptible startup delay. Fix — decouple captions from the stream entirely: - Open the stream only OFF the caption path (a background task, kicked off by warmup() and re-armed by push), never inline. `push` never awaits it. - Until the stream is confirmed open, and forever if it never opens, partials come from the fast re-decode path. A dead/slow stream no longer delays or suppresses captions. - Switch a turn to stream partials only if the stream was ready at the turn's first chunk (latched per turn), so a stream that finishes opening mid-turn can't rewrite a caption already on screen; a mid-turn stream failure falls straight back to re-decode for the rest of the turn. - Add API_STT_STREAM_OPEN_TIMEOUT_MS (default 4s, was a hard-coded 15s) bounding how long we keep trying the stream before settling on Parakeet-only partials. Verified end to end against the real Parakeet server with the stream pointed at a hung endpoint (full 4s timeout): streaming the JFK clip in real time, the first caption lands at ~790ms with real words and the full transcript follows; the worst single push blocks ~660ms (a Parakeet decode), never the stream open. Full api suite green (296 passed), coverage 96%, ruff clean. --- api/src/api/config.py | 7 ++ api/src/api/stt/__init__.py | 5 +- api/src/api/stt/streaming.py | 150 ++++++++++++++++++++++---------- api/tests/test_streaming_stt.py | 97 ++++++++++++++++++--- 4 files changed, 198 insertions(+), 61 deletions(-) diff --git a/api/src/api/config.py b/api/src/api/config.py index bf1c6f2..2af785b 100644 --- a/api/src/api/config.py +++ b/api/src/api/config.py @@ -72,6 +72,13 @@ class Settings(BaseSettings): # to ws(s)://. Empty with backend "hybrid" is a misconfiguration (no partials); # the factory raises rather than silently dropping the live caption band. stt_stream_endpoint: str = "" + # How long to wait for the streaming socket to connect before giving up on it for + # the session and staying on the re-decode partial path (XERK-128). The open runs + # off the caption path, so this only bounds how long a wedged/misconfigured stream + # server keeps us trying before we settle on Parakeet-only partials — captions + # themselves never wait on it. Kept short: on the single-host stack the stream + # server shares the api's network, so a healthy one connects in well under this. + stt_stream_open_timeout_ms: int = 4000 # ---- Cues (XERK-81) ---------------------------------------------------------- # Cues are private contextual info cards the api derives from the live diff --git a/api/src/api/stt/__init__.py b/api/src/api/stt/__init__.py index f128ccb..2be0d97 100644 --- a/api/src/api/stt/__init__.py +++ b/api/src/api/stt/__init__.py @@ -49,7 +49,10 @@ def make_transcriber(source_lang: Lang | None = None) -> Transcriber: "(API_STT_STREAM_ENDPOINT) — the WebSocket root of the " "nemotron-stt server, e.g. ws://nemotron:8000" ) - stream_engine = NemotronWsEngine(endpoint=settings.stt_stream_endpoint) + stream_engine = NemotronWsEngine( + endpoint=settings.stt_stream_endpoint, + timeout=settings.stt_stream_open_timeout_ms / 1000, + ) return StreamingTranscriber( offline, diff --git a/api/src/api/stt/streaming.py b/api/src/api/stt/streaming.py index 4217e67..0e66e3d 100644 --- a/api/src/api/stt/streaming.py +++ b/api/src/api/stt/streaming.py @@ -128,6 +128,19 @@ def __init__( self._stream_session = None self._stream_failed = False self._last_partial = "" + # The stream is opened OFF the caption path (a background task, kicked off by + # warmup() and/or the first push), never inline — connecting it, and the + # server building this connection's decoder, can take a moment or hang on an + # unhealthy/misconfigured stream server, and the api reads audio frames + # serially, so an inline open stalls every caption behind it (XERK-128). Until + # the open completes, partials come from the fast re-decode path; captions + # never wait on the stream. + self._open_task: asyncio.Task[bool] | None = None + # Latched at each turn boundary: does THIS turn's partials come from the stream? + # Fixing the source for a whole turn avoids a visible rewrite when the stream + # finishes opening mid-turn (its transcript starts empty and would briefly be + # shorter than the re-decode caption already on screen). + self._stream_turn = False self._buf = bytearray() self._since_partial = 0 @@ -139,47 +152,51 @@ def __init__( self._closed = False async def warmup(self) -> None: - """Open the live stream ahead of the first audio so its per-session setup - doesn't land on the first-caption path (XERK-128). - - On the hybrid path the first spoken chunk otherwise pays, inline, the cost of - connecting the streaming socket AND the server building this connection's - decoder — and because the api reads audio frames serially, that stall holds up - every frame behind it, so captions appear late and then in a burst carrying the - whole backlog. Opening the stream here (plus a scrap of silence to pay the - one-time cold-decode cost, then a reset so it never colours the first real - partial) moves all of that off the hot path. - - No-op when there's no streaming engine (the offline model is resident and - shared across sessions, so it has no per-session cold cost) or it already - failed. Best-effort: any failure just drops back to the lazy open on the first - chunk, exactly as before.""" - if not self._streaming_active() or self._stream_session is not None: - return - if not await self._ensure_stream(): - return - try: - # ~100 ms of silence: enough to trigger one decode step (paying the cold - # cost) without the reset that follows leaving any primed text behind. - await self._stream_session.feed(bytes(_ms_to_bytes(100))) - await self._stream_session.reset() - except Exception: # noqa: BLE001 — a failed warm just falls back to lazy open - log.warning("stream warmup failed; using lazy open", exc_info=True) - await self._drop_stream() - self._stream_failed = True + """Start opening the live stream ahead of the first audio, off the caption + path (XERK-128). + + On the hybrid path the stream's per-session setup — connecting the socket and + the server building this connection's decoder — used to be paid inline on the + first spoken chunk. Because the api reads audio frames serially, that stall + held up every frame behind it, so captions appeared late and then in a burst + carrying the whole backlog. Worse, while the stream was still "coming up" the + re-decode fallback was suppressed, so an unhealthy or misconfigured stream + server stalled captions for its whole connect timeout. + + Opening from here (in the background) means the first spoken words are captioned + from the fast re-decode path immediately and only switch to the stream once it + is confirmed up — so a slow or dead stream never delays the first caption. This + just kicks the open off and waits for it; a no-op without a streaming engine.""" + self._ensure_open_started() + if self._open_task is not None: + await self._open_task async def push(self, pcm: bytes) -> None: if not pcm: return + # Latch the partial source for this turn at its first chunk (buffer empty), + # so a stream that finishes opening mid-turn doesn't rewrite the caption. + if not self._buf: + self._stream_turn = self._stream_ready() self._buf.extend(pcm) self._since_partial += len(pcm) self._update_vad(pcm) + # Make sure the stream is opening (off the caption path); harmless once open or + # once it has failed. Never awaited here — captions must not wait on it. + self._ensure_open_started() + # Live partials from the cache-aware stream are driven by the audio itself - # (fed every chunk), not by the re-decode cadence timer below. - if self._streaming_active(): + # (fed every chunk), not by the re-decode cadence timer below — but only for a + # turn that started with the stream ready. + if self._stream_turn and self._stream_ready(): await self._feed_stream(pcm) + # Re-checked AFTER the feed: a feed that failed mid-turn drops the stream, and + # the rest of this turn must fall straight back to re-decode partials rather + # than go silent until the turn ends. + stream_driving = self._stream_turn and self._stream_ready() + if len(self._buf) >= self._max_segment_bytes: await self._finalize() elif ( @@ -188,13 +205,10 @@ async def push(self, pcm: bytes) -> None: and len(self._buf) >= self._min_segment_bytes ): await self._finalize() - elif ( - self._has_speech - and self._since_partial >= self._partial_bytes - and not self._streaming_active() - ): - # Re-decode cadence partial: only when there's no live stream (legacy - # single-model backend) or the stream failed and we fell back to it. + elif self._has_speech and self._since_partial >= self._partial_bytes and not stream_driving: + # Re-decode cadence partial: whenever the stream isn't driving — no stream + # configured, it hasn't finished opening, or it failed (before this turn or + # part-way through it) and we fell back. await self._emit_partial() def _speech_threshold(self) -> float: @@ -251,22 +265,51 @@ async def _run_engine( metrics.observe(f"stage.stt.{stage}_latency_ms", (time.perf_counter() - t0) * 1000) return result - def _streaming_active(self) -> bool: - """Whether live partials should come from the streaming engine right now.""" - return self._stream_engine is not None and not self._stream_failed + def _stream_ready(self) -> bool: + """Whether the stream is open and healthy, i.e. it can drive partials right + now. False keeps captions on the re-decode path — before the stream has + finished opening as well as after it has failed.""" + return ( + self._stream_engine is not None + and not self._stream_failed + and self._stream_session is not None + ) + + def _ensure_open_started(self) -> None: + """Kick off opening the stream in the background, at most once at a time. - async def _ensure_stream(self) -> bool: - """Open the streaming session lazily on first audio. False = unusable (then - the caller falls back to the re-decode partial path).""" + Never blocks the caller: the caption path keeps serving re-decode partials + until the open completes (see `push`). Called from both `warmup` and every + `push`, so the stream still opens even if `warmup` was never called.""" + if self._stream_engine is None or self._stream_failed: + return if self._stream_session is not None: - return True + return + if self._open_task is not None and not self._open_task.done(): + return + self._open_task = asyncio.create_task(self._open_and_prime()) + + async def _open_and_prime(self) -> bool: + """Open the streaming session and prime it, off the caption path. False = + unusable (the caption path then just stays on re-decode partials).""" try: self._stream_session = await self._stream_engine.open(language=self._language) - return True except Exception: # noqa: BLE001 — a dead stream must never stop captions log.warning("stream open failed; using re-decode partials", exc_info=True) self._stream_failed = True return False + try: + # A scrap of silence pays the stream's one-time cold-decode cost now, off + # the caption path, so the first spoken chunk gets a fast partial. No reset + # afterward: leading silence is benign to a cache-aware model (a real turn + # opens with some too) and the turn-boundary reset clears it regardless. + await self._stream_session.feed(bytes(_ms_to_bytes(100))) + except Exception: # noqa: BLE001 — a failed prime drops us to re-decode partials + log.warning("stream prime failed; using re-decode partials", exc_info=True) + await self._drop_stream() + self._stream_failed = True + return False + return True async def _drop_stream(self) -> None: session, self._stream_session = self._stream_session, None @@ -277,9 +320,9 @@ async def _drop_stream(self) -> None: pass async def _feed_stream(self, pcm: bytes) -> None: - """Feed one chunk to the live stream and emit a partial if the turn grew.""" - if not await self._ensure_stream(): - return + """Feed one chunk to the already-open live stream and emit a partial if the + turn grew. The caller only invokes this once the stream is confirmed open + (`_stream_ready`); opening happens off the caption path in `_open_and_prime`.""" try: text = await self._stream_session.feed(pcm) except Exception: # noqa: BLE001 — fall back to re-decode partials on any error @@ -395,6 +438,17 @@ async def flush(self) -> None: async def close(self) -> None: self._closed = True + # A still-running background open would race close()/drop below (both touch the + # stream session): cancel it and let it settle before dropping the session. + if self._open_task is not None and not self._open_task.done(): + self._open_task.cancel() + try: + await self._open_task + except asyncio.CancelledError: + pass + except Exception: # noqa: BLE001 — teardown swallows a late open error + pass + self._open_task = None await self._drop_stream() # Unblock a pending results() get with a skipped (empty) sentinel. await self._queue.put(CaptionPartial(type="caption.partial", text="", lang=None)) diff --git a/api/tests/test_streaming_stt.py b/api/tests/test_streaming_stt.py index 7d52085..108d55c 100644 --- a/api/tests/test_streaming_stt.py +++ b/api/tests/test_streaming_stt.py @@ -758,6 +758,7 @@ async def run() -> None: min_segment_ms=100, stream_engine=se, ) + await t.warmup() # open the stream off the caption path, as the session does for _ in range(3): # speech: three chunks -> three growing stream partials await t.push(_pcm(100, amplitude=8000)) for _ in range(3): # trailing silence closes the turn @@ -772,7 +773,7 @@ async def run() -> None: # The offline engine decoded ONLY the final — partials never touched it. assert eng.calls == 1 assert len(se.sessions) == 1 - assert se.sessions[0].feeds == 6 # every chunk, speech and silence, was fed + assert se.sessions[0].feeds == 7 # 1 warmup prime + 3 speech + 3 silence assert se.sessions[0].resets == 1 # cache cleared at the turn boundary asyncio.run(run()) @@ -794,6 +795,7 @@ async def run() -> None: t = StreamingTranscriber( FakeEngine(), language="en", stream_engine=StickyEngine(), silence_ms=100000 ) + await t.warmup() await t.push(_pcm(100, amplitude=8000)) await t.push(_pcm(100, amplitude=8000)) partials = [m for m in _drain(t) if isinstance(m, CaptionPartial)] @@ -814,6 +816,7 @@ async def run() -> None: local_agreement=False, stream_engine=se, ) + await t.warmup() # open fails here, off the caption path await t.push(_pcm(100, amplitude=8000)) await t.push(_pcm(100, amplitude=8000)) # 200ms cadence -> re-decode partial partials = [m for m in _drain(t) if isinstance(m, CaptionPartial)] @@ -835,13 +838,14 @@ async def run() -> None: local_agreement=False, stream_engine=se, ) - await t.push(_pcm(100, amplitude=8000)) # opens + one good stream partial + await t.warmup() # opens the stream off the caption path + await t.push(_pcm(100, amplitude=8000)) # one good stream partial se.sessions[0].fail_feed = True await t.push(_pcm(100, amplitude=8000)) # feed raises -> drop + fall back msgs = [m for m in _drain(t) if isinstance(m, CaptionPartial)] assert t._stream_failed is True assert se.sessions[0].closed is True # broken session was closed - assert any(p.text == "hello world" for p in msgs) # re-decode fallback kicked in + assert any(p.text == "hello world" for p in msgs) # re-decode fallback kicked in mid-turn asyncio.run(run()) @@ -866,6 +870,7 @@ async def run() -> None: min_segment_ms=100, stream_engine=se, ) + await t.warmup() # opens the stream (prime is a feed only, no reset) for _ in range(3): await t.push(_pcm(100, amplitude=8000)) for _ in range(3): # close the turn -> finalize -> reset() raises @@ -880,6 +885,7 @@ def test_hybrid_close_closes_stream_session() -> None: async def run() -> None: se = FakeStreamEngine() t = StreamingTranscriber(FakeEngine(), language="en", stream_engine=se) + await t.warmup() # opens the stream off the caption path await t.push(_pcm(100, amplitude=8000)) await t.close() assert se.sessions[0].closed is True @@ -896,21 +902,88 @@ async def run() -> None: t = StreamingTranscriber(FakeEngine(), language="en", stream_engine=se) await t.warmup() - # The socket is open and primed before a single audio frame arrived: one - # silent feed paid the cold-decode cost, one reset cleared it. + # The socket is open and primed before a single audio frame arrived: one silent + # feed paid the cold-decode cost (silence yields no partial, so nothing shows). assert len(se.sessions) == 1 assert se.sessions[0].language == "en" assert se.sessions[0].feeds == 1 - assert se.sessions[0].resets == 1 - # Priming emitted no caption — silence yields no partial, and the reset wipes it. assert _drain(t) == [] - # First real speech reuses the warmed session instead of opening a second one. + # First real speech reuses the warmed session instead of opening a second one, + # and drives partials straight from the stream (the turn started ready). await t.push(_pcm(100, amplitude=8000)) assert len(se.sessions) == 1 assert se.sessions[0].feeds == 2 partials = [m for m in _drain(t) if isinstance(m, CaptionPartial)] - assert [p.text for p in partials] == ["p0"] # numbering restarted after reset + assert [p.text for p in partials] == ["p0"] + + asyncio.run(run()) + + +def test_stream_open_never_blocks_the_caption_path() -> None: + """The core XERK-128 fix: even a slow-opening stream must not delay captions — + partials come from the re-decode path until the stream is confirmed up.""" + + class SlowOpenEngine(FakeStreamEngine): + def __init__(self) -> None: + super().__init__() + self.gate = asyncio.Event() + + async def open(self, *, language: str | None) -> FakeStreamSession: + await self.gate.wait() # stays "connecting" until the test releases it + return await super().open(language=language) + + async def run() -> None: + se = SlowOpenEngine() + t = StreamingTranscriber( + FakeEngine(), + language="en", + partial_interval_ms=200, + silence_ms=100000, + local_agreement=False, + stream_engine=se, + ) + # No warmup await: the open is still blocked on the gate. Captions must flow + # anyway, from the re-decode path — the first turn does not wait on the stream. + await t.push(_pcm(100, amplitude=8000)) + await t.push(_pcm(100, amplitude=8000)) # 200ms cadence -> re-decode partial + partials = [m for m in _drain(t) if isinstance(m, CaptionPartial)] + assert partials and partials[0].text == "hello world" + assert len(se.sessions) == 0 # the stream never opened, yet captions flowed + + await t.close() # cancels the still-pending open cleanly + + asyncio.run(run()) + + +def test_stream_upgrades_at_next_turn_once_open() -> None: + """A stream that finishes opening mid-turn drives partials from the NEXT turn, not + the current one — so it never rewrites a caption already on screen.""" + + async def run() -> None: + se = FakeStreamEngine() + t = StreamingTranscriber( + FakeEngine(), + language="en", + partial_interval_ms=100, + silence_ms=300, + min_segment_ms=100, + local_agreement=False, + stream_engine=se, + ) + # Turn 1: stream not open yet -> re-decode drives it. + await t.push(_pcm(100, amplitude=8000)) + turn1 = [m for m in _drain(t) if isinstance(m, CaptionPartial)] + assert turn1 and turn1[0].text == "hello world" # from the offline engine + # Let the background open finish, then close turn 1 on silence. + await t.warmup() + for _ in range(3): + await t.push(_pcm(100, amplitude=0)) + _drain(t) + # Turn 2: stream is ready at the turn's first chunk -> it drives partials now. + await t.push(_pcm(100, amplitude=8000)) + turn2 = [m for m in _drain(t) if isinstance(m, CaptionPartial)] + assert turn2 and turn2[0].text == "p0" # from the stream asyncio.run(run()) @@ -932,12 +1005,12 @@ async def run() -> None: await t.warmup() await t.warmup() # already open -> does nothing more assert len(se.sessions) == 1 - assert se.sessions[0].feeds == 1 + assert se.sessions[0].feeds == 1 # just the one prime feed asyncio.run(run()) -def test_warmup_open_failure_falls_back_to_lazy() -> None: +def test_warmup_open_failure_falls_back_to_redecode() -> None: async def run() -> None: se = FakeStreamEngine(fail_open=True) t = StreamingTranscriber( @@ -959,7 +1032,7 @@ async def run() -> None: asyncio.run(run()) -def test_warmup_feed_failure_falls_back_to_lazy() -> None: +def test_warmup_prime_failure_falls_back_to_redecode() -> None: class BadFeedSession(FakeStreamSession): async def feed(self, pcm: bytes) -> str | None: raise RuntimeError("warm feed exploded") From e14f8e07c504de1da789b37780c52f1126b62028 Mon Sep 17 00:00:00 2001 From: Malcolm Habeeb Date: Mon, 27 Jul 2026 08:36:25 -0400 Subject: [PATCH 2/2] fix(nemotron-stt): serve the WebSocket with wsproto so the handshake stops wedging (XERK-128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The live cause of the caption stall: the Nemotron streaming server on the GPU box answers HTTP /health in ~1ms but NEVER completes the /v1/audio/stream WebSocket handshake — the 101 is never sent, so the api's stream open hangs for its whole timeout every session. Reproduced against the live server (raw upgrade: no bytes in 8s; health: 1ms) and confirmed consistent (3/3 opens time out). This is the `websockets` library biting the handshake a second time. The server already switched off uvicorn's legacy `websockets` impl (which 400s a valid upgrade against websockets>=14) to `websockets-sansio`; that workaround now hangs instead against the version frozen in the image (deps are unpinned `>=`). Both failures are the same root: a hard dependency on uvicorn's fast-moving `websockets` adapter. Fix: serve with **wsproto**, a self-contained sans-I/O WebSocket implementation independent of the `websockets` library, via `ws=WS_IMPL`. Verified locally to complete the handshake AND the full partial/reset protocol driven by the api's real client. `websockets` stays installed (uvicorn[standard] pulls it; the api's own client uses it) but no longer serves this endpoint. Also add a real-uvicorn regression test: it starts an actual server on the configured WS impl and does a real handshake + protocol round-trip. The existing tests use FastAPI's TestClient, whose in-memory WS transport bypasses uvicorn's adapter entirely — so they could never have caught this. A wedged handshake now fails that test (open_timeout) instead of shipping. And align the repo compose's published Nemotron port 9402 -> 9403 to match the GPU deployment (DockerOps tenir-gpu.yaml: 9401 Parakeet, 9402 cue LLM, 9403 Nemotron), so the two topologies stop disagreeing on which port is which. NOTE: final proof needs a rebuild on the GPU host and a re-test of the live WS — which can't run from here; local repro + the api-side graceful fallback (#84) cover the rest. --- docker-compose.yml | 6 ++- nemotron-stt/Dockerfile | 11 +++-- nemotron-stt/requirements.txt | 11 ++++- nemotron-stt/server.py | 21 ++++++-- nemotron-stt/tests/test_server.py | 81 +++++++++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 12 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index c84c82d..17f1ec4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -182,7 +182,11 @@ services: dockerfile: nemotron-stt/Dockerfile image: tenir-nemotron-stt:latest ports: - - "9402:8000" + # 9403 to match the GPU deployment (DockerOps tenir-gpu.yaml): there 9401 is + # Parakeet and 9402 is the cue LLM, so Nemotron gets 9403. The api reaches this + # service by DNS (ws://nemotron:8000); this published port is only for ad-hoc + # external access, kept on the same number as prod to avoid confusion. + - "9403:8000" volumes: # Whole /root/.cache so the HF hub cache AND NeMo's torch cache persist. - "C:/docker/tenir/nemotron:/root/.cache" diff --git a/nemotron-stt/Dockerfile b/nemotron-stt/Dockerfile index 2cee637..c3793da 100644 --- a/nemotron-stt/Dockerfile +++ b/nemotron-stt/Dockerfile @@ -48,9 +48,10 @@ ENV NEMOTRON_MODEL=nvidia/nemotron-3.5-asr-streaming-0.6b \ EXPOSE 8000 # server.py's main() loads the model on the main thread BEFORE starting uvicorn, then -# serves with the websockets-sansio WS impl. Both are required for the WebSocket accept -# path to survive NeMo's cold CUDA init (loading in a background thread while the loop -# comes up wedges new handshakes) and modern websockets (the legacy uvicorn WS impl -# 400s a valid upgrade against websockets>=14). See _load_model / main() for the full -# rationale. One worker: a single resident model on one GPU. +# serves the WebSocket with wsproto (server.py WS_IMPL). Loading before the loop keeps +# the accept path healthy through NeMo's cold CUDA init; wsproto keeps the handshake +# off the `websockets` library, whose version churn broke it twice (a 400 on uvicorn's +# legacy impl, then a silent hang on websockets-sansio in production). See _load_model +# / main() / WS_IMPL for the full rationale. One worker: a single resident model on one +# GPU. CMD ["python", "server.py"] diff --git a/nemotron-stt/requirements.txt b/nemotron-stt/requirements.txt index 5dce3fe..e71a8c6 100644 --- a/nemotron-stt/requirements.txt +++ b/nemotron-stt/requirements.txt @@ -1,8 +1,15 @@ # Web layer on top of the NeMo base image (which already carries torch + NeMo ASR). # NeMo usually ships numpy/soundfile already; listing them is idempotent and makes -# the server's runtime deps explicit and self-contained. websockets backs the -# FastAPI/uvicorn WebSocket endpoint this server exposes. +# the server's runtime deps explicit and self-contained. +# +# The WebSocket endpoint is served with wsproto, NOT the `websockets` library +# (server.py WS_IMPL): riding `websockets` broke the handshake twice — a 400 on the +# legacy uvicorn impl, then a silent hang on `websockets-sansio` in production. +# wsproto is a self-contained sans-I/O WS implementation independent of that library's +# churn. `websockets` stays here only because uvicorn[standard] depends on it (and the +# api's own client uses it separately); it no longer serves this endpoint. fastapi>=0.115 uvicorn[standard]>=0.30 +wsproto>=1.2 websockets>=13 numpy diff --git a/nemotron-stt/server.py b/nemotron-stt/server.py index 737201e..a9a7355 100644 --- a/nemotron-stt/server.py +++ b/nemotron-stt/server.py @@ -64,6 +64,19 @@ app = FastAPI(title="tenir-nemotron-stt") +# WebSocket implementation uvicorn serves /v1/audio/stream with. `wsproto`, NOT the +# `websockets` library (XERK-128). The handshake has broken twice riding `websockets`: +# uvicorn's legacy `websockets` impl 400s a valid upgrade against websockets>=14, and +# its `websockets-sansio` impl — the previous workaround — was found in production to +# accept the TCP connection and answer /health but NEVER complete the /v1/audio/stream +# handshake (the 101 is never sent), so the api's stream open hangs for its whole +# timeout every session. wsproto is a self-contained sans-I/O WebSocket implementation +# that doesn't depend on the fast-moving `websockets` library at all, so it sidesteps +# that version churn; verified to complete the handshake and the full partial/reset +# protocol. (`websockets` stays installed — uvicorn[standard] pulls it — but is unused +# for serving; the api's own client still uses it independently.) +WS_IMPL = "wsproto" + # One resident model on one GPU. Each WebSocket owns its own encoder-cache state # (per-connection, cheap) but shares the model weights; NeMo decode steps aren't # safe to run concurrently on the same model, so every step runs under this lock. @@ -284,15 +297,15 @@ def main() -> None: # pragma: no cover - process entrypoint (needs NeMo + a GPU path healthy (see `_load_model`). The trade is that the port only binds after the model is resident (~tens of seconds) — during which health probes get connection-refused, covered by the compose healthcheck's start_period — instead - of binding immediately and answering 503. `--ws websockets-sansio` is required - because the base image ships websockets>=14, which uvicorn's legacy WS impl can't - handshake (it 400s a valid upgrade). + of binding immediately and answering 503. `ws=WS_IMPL` ("wsproto") serves the + WebSocket independently of the `websockets` library, whose churn broke this + handshake twice — see WS_IMPL for the full rationale. """ import uvicorn # noqa: PLC0415 port = int(os.environ.get("NEMOTRON_PORT", "8000")) _load_model() - uvicorn.run(app, host="0.0.0.0", port=port, ws="websockets-sansio") + uvicorn.run(app, host="0.0.0.0", port=port, ws=WS_IMPL) if __name__ == "__main__": diff --git a/nemotron-stt/tests/test_server.py b/nemotron-stt/tests/test_server.py index fe24e44..3798a97 100644 --- a/nemotron-stt/tests/test_server.py +++ b/nemotron-stt/tests/test_server.py @@ -10,9 +10,16 @@ from __future__ import annotations +import asyncio +import json +import socket +import threading +import time + import numpy as np import pytest import server +import uvicorn from fastapi.testclient import TestClient @@ -130,3 +137,77 @@ def test_stream_refused_while_loading() -> None: with client.websocket_connect("/v1/audio/stream") as ws: msg = ws.receive_json() assert msg["type"] == "error" + + +def _free_port() -> int: + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _wait_listening(port: int, timeout: float = 15.0) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.25): + return + except OSError: + time.sleep(0.05) + raise AssertionError("uvicorn did not start listening in time") + + +def test_real_uvicorn_ws_handshake_and_protocol(monkeypatch: pytest.MonkeyPatch) -> None: + """Drive a REAL uvicorn server over its configured WS impl (server.WS_IMPL). + + This is the regression guard for XERK-128: the handshake wedged in production — + uvicorn accepted the socket and served /health but never completed the + /v1/audio/stream upgrade — and TestClient's in-memory WebSocket transport can't + see that, because it bypasses uvicorn's WS adapter entirely. Only a real server + exercises the layer that broke. A wedged handshake fails this test (open_timeout) + rather than hanging it. + """ + import websockets # noqa: PLC0415 — only needed for this real-network test + + made: list[FakeDecoder] = [] + + def factory(*, target_lang: str) -> FakeDecoder: + d = FakeDecoder(target_lang=target_lang) + made.append(d) + return d + + monkeypatch.setattr(server, "make_decoder", factory) + server._ready.set() + + port = _free_port() + config = uvicorn.Config( + server.app, host="127.0.0.1", port=port, ws=server.WS_IMPL, log_level="error" + ) + srv = uvicorn.Server(config) + thread = threading.Thread(target=srv.run, daemon=True) + thread.start() + try: + _wait_listening(port) + + async def drive() -> tuple[dict, dict, dict]: + url = f"ws://127.0.0.1:{port}/v1/audio/stream" + async with websockets.connect(url, open_timeout=5) as ws: + await ws.send(json.dumps({"target_lang": "en-US"})) + await ws.send(_pcm()) + m1 = json.loads(await asyncio.wait_for(ws.recv(), 5)) + await ws.send(_pcm()) + m2 = json.loads(await asyncio.wait_for(ws.recv(), 5)) + await ws.send(json.dumps({"type": "reset"})) + m3 = json.loads(await asyncio.wait_for(ws.recv(), 5)) + return m1, m2, m3 + + m1, m2, m3 = asyncio.run(drive()) + assert m1 == {"type": "partial", "text": "w0"} # handshake completed; partials flow + assert m2 == {"type": "partial", "text": "w0 w1"} + assert m3 == {"type": "reset_ok"} + assert made and made[0].target_lang == "en-US" + finally: + srv.should_exit = True + thread.join(timeout=5) + server._ready.clear()