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
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,
                }
            ],
        }


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