fix(hermes): warm STT before accepting speech

dockerfiles/hermes-jetson-stt-server.py

@@ -41,7 +41,7 @@ VAD_END_SILENCE_MS = int(os.getenv("HERMES_STT_VAD_END_SILENCE_MS", "650"))
 VAD_TAIL_MS = int(os.getenv("HERMES_STT_VAD_TAIL_MS", "220"))
 ROLLING_MIN_AUDIO_MS = max(
     300,
-    min(2_000, int(os.getenv("HERMES_STT_ROLLING_MIN_AUDIO_MS", "650"))),
+    min(2_000, int(os.getenv("HERMES_STT_ROLLING_MIN_AUDIO_MS", "1800"))),
 )
 ROLLING_INTERVAL_MS = max(
     400,
@@ -127,6 +127,7 @@ class InferenceGate:
 
 
 INFERENCE_GATE = InferenceGate()
+STARTUP_WARMUP_MS: float | None = None
 
 
 def _repetitive_token(token: str) -> bool:
@@ -248,6 +249,21 @@ def _transcribe_pcm(model: object, pcm: bytes, language: str) -> dict:
                 pass
 
 
+def _warm_model(model: object) -> float:
+    """Pay Whisper's lazy CUDA initialization cost before readiness."""
+    temp_path = _write_pcm_wav(bytes(_pcm_bytes_for_ms(1_000)))
+    started = time.monotonic()
+    try:
+        _decode_path(model, temp_path, "auto")
+        torch.cuda.synchronize()
+    finally:
+        try:
+            os.unlink(temp_path)
+        except OSError:
+            pass
+    return (time.monotonic() - started) * 1000
+
+
 def _transcribe_pcm_rolling(
     model: object, pcm: bytes, language: str
 ) -> dict | None:
@@ -468,7 +484,8 @@ class StreamingTranscription:
         self._cached_payload: dict | None = None
         self._rolling_inflight = False
         self._rolling_last_started_at = 0.0
-        self._rolling_last_audio_bytes = 0
+        self._rolling_speech_bytes = 0
+        self._rolling_last_speech_bytes = 0
         self._rolling_previous = ""
         self._rolling_previous_window_start = 0
         self._rolling_revision = 0
@@ -530,9 +547,9 @@ class StreamingTranscription:
         """Discard provisional agreement after the client's speech epoch changes."""
         self._rolling_previous = ""
         self._rolling_previous_window_start = 0
-        # Audio before a resume belongs to an invalidated provisional epoch and
+        # Only speech in the new epoch may satisfy the rolling threshold.
-        # must not immediately satisfy the next rolling interval by itself.
+        self._rolling_speech_bytes = 0
-        self._rolling_last_audio_bytes = len(self._pcm)
+        self._rolling_last_speech_bytes = 0
         self._rolling_last_started_at = 0.0
 
     def _begin_rolling_locked(
@@ -543,12 +560,13 @@ class StreamingTranscription:
         minimum_bytes = _pcm_bytes_for_ms(ROLLING_MIN_AUDIO_MS)
         interval_bytes = _pcm_bytes_for_ms(ROLLING_INTERVAL_MS)
         required_new_bytes = (
-            interval_bytes if self._rolling_last_audio_bytes else minimum_bytes
+            interval_bytes if self._rolling_last_speech_bytes else minimum_bytes
         )
         if (
             self._rolling_inflight
-            or total_bytes < minimum_bytes
+            or self._rolling_speech_bytes < minimum_bytes
-            or total_bytes - self._rolling_last_audio_bytes < required_new_bytes
+            or self._rolling_speech_bytes - self._rolling_last_speech_bytes
+            < required_new_bytes
             or (
                 self._rolling_last_started_at
                 and (now - self._rolling_last_started_at) * 1000
@@ -559,7 +577,7 @@ class StreamingTranscription:
         pcm, window_start_ms = self._rolling_snapshot_locked()
         self._rolling_inflight = True
         self._rolling_last_started_at = now
-        self._rolling_last_audio_bytes = total_bytes
+        self._rolling_last_speech_bytes = self._rolling_speech_bytes
         return pcm, window_start_ms, total_bytes, self._epoch
 
     def _start_rolling(
@@ -657,6 +675,7 @@ class StreamingTranscription:
                 self._at_eos = False
                 self._silence_bytes = 0
                 self._last_speech_byte = len(self._pcm)
+                self._rolling_speech_bytes += len(pcm)
                 rolling = self._begin_rolling_locked(time.monotonic())
             elif self._heard_speech:
                 self._silence_bytes += len(pcm)
@@ -668,6 +687,7 @@ class StreamingTranscription:
                 )
                 if self._silence_bytes >= required and not self._at_eos:
                     self._at_eos = True
+                    self._reset_rolling_history_locked()
                     auto_speculate = True
         if auto_speculate:
             self.speculate()
@@ -828,6 +848,10 @@ class SpeechHandler(BaseHTTPRequestHandler):
                 "model": MODEL_NAME,
                 "device": "cuda" if torch.cuda.is_available() else "cpu",
                 "inference": INFERENCE_GATE.snapshot(),
+                "startup": {
+                    "warmed": STARTUP_WARMUP_MS is not None,
+                    "warmup_ms": round(STARTUP_WARMUP_MS or 0.0, 1),
+                },
                 "streaming": {
                     "enabled": True,
                     "path": "/v1/audio/transcriptions/stream",
@@ -838,6 +862,7 @@ class SpeechHandler(BaseHTTPRequestHandler):
                     "speculative": True,
                     "rolling": {
                         "enabled": True,
+                        "minimum_audio_ms": ROLLING_MIN_AUDIO_MS,
                         "interval_ms": ROLLING_INTERVAL_MS,
                         "window_ms": ROLLING_WINDOW_MS,
                         "stable_prefix": True,
@@ -993,11 +1018,14 @@ class SpeechHandler(BaseHTTPRequestHandler):
 
 def main() -> None:
     """Warm Whisper once, then serve concurrent clients through one GPU lock."""
+    global STARTUP_WARMUP_MS
     CACHE_DIR.mkdir(parents=True, exist_ok=True)
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for the Atlas Whisper service")
     print(f"[stt] loading Whisper {MODEL_NAME} into CUDA", flush=True)
     model = whisper.load_model(MODEL_NAME, device="cuda", download_root=str(CACHE_DIR))
+    STARTUP_WARMUP_MS = _warm_model(model)
+    print(f"[stt] CUDA warm-up completed in {STARTUP_WARMUP_MS:.1f}ms", flush=True)
     server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
     server.model = model  # type: ignore[attr-defined]
     print(f"[stt] ready on {HOST}:{PORT}", flush=True)

testing/tests/test_hermes_stt_streaming.py

@@ -28,7 +28,9 @@ def _load_server(monkeypatch):
     monkeypatch.setitem(
         sys.modules,
         "torch",
-        SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
+        SimpleNamespace(
+            cuda=SimpleNamespace(is_available=lambda: False, synchronize=lambda: None)
+        ),
     )
     monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace())
     spec.loader.exec_module(module)
@@ -214,7 +216,10 @@ def test_stream_speculation_is_reused_for_unchanged_commit(monkeypatch):
     session = module.StreamingTranscription(connection, model)
     _start(session)
 
-    session.append(_speech())
+    session.append(_silence(samples=32_000))
+    session.append(_speech(samples=20_800))
+    assert not session._rolling_inflight
+    assert model.calls == []
     session.append(_silence())
     partial = connection.wait_for("partial")
     assert partial == {
@@ -236,13 +241,12 @@ def test_stream_speculation_is_reused_for_unchanged_commit(monkeypatch):
 
 def test_active_speech_emits_rate_limited_stable_rolling_partials(monkeypatch):
     module = _load_server(monkeypatch)
-    _configure_fast_rolling(module, monkeypatch)
     connection = _Connection()
     model = _SequencedModel(["please open", "please open my calendar"])
     session = module.StreamingTranscription(connection, model)
     _start(session)
 
-    session.append(_speech())
+    session.append(_speech(samples=32_000))
     first = connection.wait_for_revision(1)
     assert first == {
         "type": "partial",
@@ -259,7 +263,7 @@ def test_active_speech_emits_rate_limited_stable_rolling_partials(monkeypatch):
 
     # Audio progress is required in addition to wall-clock rate limiting.
     session._rolling_last_started_at = 0.0
-    session.append(_speech())
+    session.append(_speech(samples=12_000))
     second = connection.wait_for_revision(2)
     assert second["transcript"] == "please open my calendar"
     assert second["stable_transcript"] == "please open"
@@ -470,4 +474,26 @@ def test_stream_health_contract_is_declared_in_source():
     assert '"server_vad": True' in source
     assert '"speculative": True' in source
     assert '"inference": INFERENCE_GATE.snapshot()' in source
+    assert '"minimum_audio_ms": ROLLING_MIN_AUDIO_MS' in source
+    assert '"warmed": STARTUP_WARMUP_MS is not None' in source
     assert "MAX_STREAM_AUDIO_BYTES" in source
+
+
+def test_startup_warms_canonical_decode_before_serving(monkeypatch):
+    module = _load_server(monkeypatch)
+    model = _Model()
+    synchronized: list[bool] = []
+    monkeypatch.setattr(module.torch.cuda, "synchronize", lambda: synchronized.append(True))
+
+    elapsed = module._warm_model(model)
+
+    assert elapsed >= 0
+    assert module.ROLLING_MIN_AUDIO_MS == 1_800
+    assert len(model.calls) == 1
+    assert model.calls[0][1]["language"] is None
+    assert synchronized == [True]
+    assert not Path(model.calls[0][0]).exists()
+    source = SERVER.read_text()
+    assert source.index("STARTUP_WARMUP_MS = _warm_model(model)") < source.index(
+        "server = ThreadingHTTPServer"
+    )