atlas-iac/testing/tests/test_hermes_stt_rolling_model.py

"""Release gates for Hermes' disposable tiny rolling transcription model."""

from __future__ import annotations

import importlib.util
import struct
import sys
import threading
import time
import wave
from pathlib import Path
from types import SimpleNamespace


ROOT = Path(__file__).resolve().parents[2]
SERVER = ROOT / "dockerfiles" / "hermes-jetson-stt-server.py"
DOCKERFILE = ROOT / "dockerfiles" / "Dockerfile.hermes-jetson-stt"
PIPELINE = ROOT / "ci" / "Jenkinsfile.hermes-voice-image"


def _load_server(monkeypatch):
    """Import the server without requiring image-only CUDA dependencies."""
    spec = importlib.util.spec_from_file_location("hermes_stt_dual_model", SERVER)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(sys.modules, "cgi", SimpleNamespace())
    monkeypatch.setitem(
        sys.modules,
        "torch",
        SimpleNamespace(
            cuda=SimpleNamespace(is_available=lambda: False, synchronize=lambda: None)
        ),
    )
    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace())
    spec.loader.exec_module(module)
    return module


class _Connection:
    """Capture asynchronous stream responses."""

    def __init__(self) -> None:
        self.messages: list[dict] = []
        self.condition = threading.Condition()

    def send_json(self, payload: dict) -> None:
        with self.condition:
            self.messages.append(payload)
            self.condition.notify_all()

    def wait_for(self, message_type: str, timeout: float = 1.0) -> dict:
        deadline = time.monotonic() + timeout
        with self.condition:
            while time.monotonic() < deadline:
                for message in self.messages:
                    if message.get("type") == message_type:
                        return message
                self.condition.wait(deadline - time.monotonic())
        raise AssertionError(f"missing {message_type}: {self.messages}")


class _Model:
    """Return deterministic Whisper output while recording model ownership."""

    def __init__(self, text: str) -> None:
        self.text = text
        self.calls: list[dict] = []

    def transcribe(self, path: str, **options: object) -> dict:
        self.calls.append({"path": path, "options": options})
        return {
            "text": self.text,
            "language": "en",
            "segments": [
                {
                    "text": self.text,
                    "no_speech_prob": 0.01,
                    "avg_logprob": -0.1,
                }
            ],
        }


class _BlockingModel(_Model):
    """Hold one authoritative decode to exercise commit/speculation races."""

    def __init__(self, text: str = "hello") -> None:
        super().__init__(text)
        self.entered = threading.Event()
        self.release = threading.Event()

    def transcribe(self, path: str, **options: object) -> dict:
        self.entered.set()
        assert self.release.wait(2.0)
        return super().transcribe(path, **options)


class _SequencedModel(_Model):
    """Return evolving text and retain the audio size decoded each time."""

    def __init__(self, texts: list[str]) -> None:
        super().__init__(texts[0])
        self.texts = texts
        self.frame_counts: list[int] = []

    def transcribe(self, path: str, **options: object) -> dict:
        with wave.open(path, "rb") as wav_file:
            self.frame_counts.append(wav_file.getnframes())
        self.text = self.texts[min(len(self.calls), len(self.texts) - 1)]
        return super().transcribe(path, **options)


def _start(session) -> None:
    session.start(
        {
            "type": "start",
            "turn_id": "dual-model-turn",
            "format": "pcm_s16le",
            "sample_rate": 16_000,
            "language": "auto",
        }
    )


def _speech(samples: int = 1_600) -> bytes:
    return struct.pack(f"<{samples}h", *([12_000] * samples))


def _silence(samples: int = 12_000) -> bytes:
    return bytes(samples * 2)


def _wait_idle(session) -> None:
    deadline = time.monotonic() + 1.0
    with session._lock:
        while session._rolling_inflight and time.monotonic() < deadline:
            session._lock.wait(deadline - time.monotonic())
        assert not session._rolling_inflight


def test_rolling_uses_tiny_once_while_final_remains_small(monkeypatch):
    module = _load_server(monkeypatch)
    monkeypatch.setattr(module, "ROLLING_MIN_AUDIO_MS", 100)
    monkeypatch.setattr(module, "ROLLING_MAX_RESULT_LAG_MS", 5_000)
    final_model = _Model("authoritative small")
    rolling_model = _Model("provisional tiny")
    connection = _Connection()
    session = module.StreamingTranscription(connection, final_model, rolling_model)
    _start(session)

    session.append(_speech())
    partial = connection.wait_for("partial")
    assert partial["transcript"] == "provisional tiny"
    assert partial["stable_transcript"] == ""
    assert len(rolling_model.calls) == 1
    assert final_model.calls == []

    session.append(_speech())
    _wait_idle(session)
    assert len(rolling_model.calls) == 1
    session.commit()
    assert connection.wait_for("final")["transcript"] == "authoritative small"
    assert len(final_model.calls) == 1


def test_trailing_silence_after_speculation_keeps_exact_cache_key(monkeypatch):
    """Recorder shutdown silence must not force a second Whisper-small decode."""
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _Model("hello")
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech(samples=20_800))
    session.speculate()
    assert connection.wait_for("partial")["transcript"] == "hello"
    session.append(_silence(samples=8_000))
    session.commit()

    assert connection.wait_for("final")["transcript"] == "hello"
    assert len(model.calls) == 1
    telemetry = module.STREAM_TELEMETRY.snapshot()
    assert telemetry["speculative_cache_hits"] == 1
    assert telemetry["direct_final_decodes"] == 0
    assert telemetry["last_commit_cache_hit"] is True
    assert telemetry["last_commit_audio_ms"] == 1_300.0


def test_commit_waits_for_one_inflight_exact_speculation_without_redecoding(
    monkeypatch,
):
    """Exercise the browser timing race while Whisper-small is still running."""
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _BlockingModel()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech(samples=20_800))
    session.speculate()
    assert model.entered.wait(1.0)
    session.append(_silence(samples=8_000))
    committed = threading.Event()
    worker = threading.Thread(target=lambda: (session.commit(), committed.set()))
    worker.start()
    time.sleep(0.02)
    assert not committed.is_set()
    model.release.set()
    worker.join(timeout=1.0)

    assert committed.is_set()
    assert connection.wait_for("final")["transcript"] == "hello"
    assert len(model.calls) == 1
    telemetry = module.STREAM_TELEMETRY.snapshot()
    assert telemetry["last_commit_cache_hit"] is True
    assert telemetry["last_commit_pending_wait_ms"] >= 10


def test_resumed_speech_invalidates_frozen_eos_and_decodes_full_turn(monkeypatch):
    """Freezing EOS cannot discard an intentional barge-in continuation."""
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _SequencedModel(["first thought", "complete thought"])
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech(samples=8_000))
    session.speculate()
    connection.wait_for("partial")
    session.append(_silence(samples=4_000))
    session.resume()
    session.append(_speech(samples=4_000))
    session.commit()

    assert connection.wait_for("final")["transcript"] == "complete thought"
    assert model.frame_counts == [8_000, 16_000]
    telemetry = module.STREAM_TELEMETRY.snapshot()
    assert telemetry["speculative_cache_hits"] == 0
    assert telemetry["direct_final_decodes"] == 1
    assert telemetry["last_commit_cache_hit"] is False


def test_eos_does_not_rearm_but_resumed_speech_does(monkeypatch):
    module = _load_server(monkeypatch)
    monkeypatch.setattr(module, "ROLLING_MIN_AUDIO_MS", 100)
    monkeypatch.setattr(module, "ROLLING_MAX_RESULT_LAG_MS", 5_000)
    final_model = _Model("final")
    rolling_model = _Model("rolling")
    session = module.StreamingTranscription(_Connection(), final_model, rolling_model)
    _start(session)

    session.append(_speech())
    _wait_idle(session)
    assert session._rolling_attempted is True
    session.append(bytes(24_000))
    assert session._at_eos is True
    assert session._rolling_attempted is True

    session.resume()
    assert session._rolling_attempted is False
    session.append(_speech())
    _wait_idle(session)
    assert len(rolling_model.calls) == 2


def test_dual_warmup_and_separate_timing_are_declared(monkeypatch):
    module = _load_server(monkeypatch)
    model = _Model("warm")
    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]["options"]["language"] is None
    assert synchronized == [True]
    assert not Path(model.calls[0]["path"]).exists()
    source = SERVER.read_text()
    rolling_warm = "STARTUP_ROLLING_WARMUP_MS = _warm_model(rolling_model)"
    final_warm = "STARTUP_WARMUP_MS = _warm_model(model, _decode_final)"
    assert source.index(rolling_warm) < source.index(final_warm)
    assert source.index(final_warm) < source.index("server = ThreadingHTTPServer")
    assert 'kind="rolling"' in source
    assert 'kind="speculative"' in source
    assert '"last_rolling_decode_ms"' in source
    assert '"last_speculative_decode_ms"' in source


def test_tiny_model_is_checksum_pinned_read_only_in_the_runtime_image():
    dockerfile = DOCKERFILE.read_text()
    digest = "65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9"
    assert f"--checksum=sha256:{digest} --chmod=0444" in dockerfile
    assert f"/models/{digest}/tiny.pt" in dockerfile
    assert "HERMES_STT_ROLLING_MODEL=tiny" in dockerfile
    assert "('large-v3-turbo.pt','small.pt','tiny.pt')" in dockerfile


def test_dual_model_contract_is_an_exact_voice_release_gate():
    pipeline = PIPELINE.read_text()
    assert "testing/tests/test_hermes_stt_rolling_model.py" in pipeline
    source = SERVER.read_text()
    assert '"path": "/v1/audio/transcriptions/stream"' in source
    assert '"inference": INFERENCE_GATE.snapshot()' in source
    assert '"rolling_model": ROLLING_MODEL_NAME' in source
    assert "STARTUP_ROLLING_WARMUP_MS is not None" in source