titan/atlas-iac
3
0
Fork 0
Code Issues Pull Requests 9 Actions Packages Projects Releases Wiki Activity

fix(hermes): warm STT before accepting speech

Browse Source
This commit is contained in:
jenkins 2026-08-23 23:52:30 -03:00
parent 70d797deb3
commit af5ee4697b
2 changed files with 68 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
dockerfiles/hermes-jetson-stt-server.py
View File

@ -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
# must not immediately satisfy the next rolling interval by itself.
self._rolling_last_audio_bytes = len(self._pcm)
# Only speech in the new epoch may satisfy the rolling threshold.
self._rolling_speech_bytes = 0
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 total_bytes - self._rolling_last_audio_bytes < required_new_bytes
or self._rolling_speech_bytes < minimum_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)

36
testing/tests/test_hermes_stt_streaming.py
View File

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