STT-Assistant/tests/test_accel.py

"""Tests for the ROCm accelerator helpers."""

import logging
import os
import sys
import types

import pytest

from stt_daemon import accel
from stt_daemon import transcriber as transcriber_mod
from stt_daemon.transcriber import Transcriber


def _fake_torch(*, hip=None, cuda_available=False):
    return types.SimpleNamespace(
        version=types.SimpleNamespace(hip=hip),
        cuda=types.SimpleNamespace(is_available=lambda: cuda_available),
    )


class _RecordingWhisperModel:
    """Records the allocator and arguments visible at construction time."""

    last: "_RecordingWhisperModel | None" = None

    def __init__(self, model, device, compute_type, cpu_threads=0):
        self.model = None  # no CTranslate2 inner model in the fake
        self.allocator_at_creation = os.environ.get(accel.CT2_ALLOCATOR_ENV)
        self.device_arg = device
        self.compute_type_arg = compute_type
        _RecordingWhisperModel.last = self


def test_is_rocm_host_true_when_torch_reports_hip(monkeypatch):
    monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip="7.2.0"))
    assert accel.is_rocm_host() is True


def test_is_rocm_host_false_when_torch_reports_no_hip(monkeypatch):
    monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip=None))
    assert accel.is_rocm_host() is False


def test_is_rocm_host_false_without_torch_env_or_marker(monkeypatch):
    monkeypatch.delitem(sys.modules, "torch", raising=False)
    monkeypatch.delenv("ROCM_PATH", raising=False)
    monkeypatch.delenv("ROCM_HOME", raising=False)
    monkeypatch.setattr(accel, "_ROCM_MARKER_PATHS", ())
    assert accel.is_rocm_host() is False


def test_is_rocm_host_true_when_marker_path_is_a_directory(monkeypatch, tmp_path):
    monkeypatch.delitem(sys.modules, "torch", raising=False)
    monkeypatch.delenv("ROCM_PATH", raising=False)
    monkeypatch.delenv("ROCM_HOME", raising=False)
    monkeypatch.setattr(accel, "_ROCM_MARKER_PATHS", (str(tmp_path),))
    assert accel.is_rocm_host() is True


def test_configure_ct2_allocator_is_noop_on_non_rocm_host(monkeypatch):
    monkeypatch.setattr(accel, "is_rocm_host", lambda: False)
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)

    assert accel.configure_ct2_allocator() is None
    assert accel.CT2_ALLOCATOR_ENV not in os.environ


def test_configure_ct2_allocator_sets_cub_caching_on_rocm_host(monkeypatch):
    monkeypatch.setattr(accel, "is_rocm_host", lambda: True)
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)

    assert accel.configure_ct2_allocator() == accel.ROCM_CT2_ALLOCATOR
    assert os.environ[accel.CT2_ALLOCATOR_ENV] == "cub_caching"


def test_configure_ct2_allocator_preserves_explicit_value(monkeypatch):
    monkeypatch.setattr(accel, "is_rocm_host", lambda: True)
    monkeypatch.setenv(accel.CT2_ALLOCATOR_ENV, "default")

    assert accel.configure_ct2_allocator() == "default"
    assert os.environ[accel.CT2_ALLOCATOR_ENV] == "default"


@pytest.mark.parametrize(
    ("configured", "resolved", "count"),
    [
        ("cpu", "cpu", 0),
        ("auto", "cuda", 1),
    ],
)
def test_acceleration_warning_returns_none_when_healthy(configured, resolved, count):
    assert accel.acceleration_warning(configured, resolved, count) is None


def test_acceleration_warning_cuda_states_request_not_honored(monkeypatch):
    message = accel.acceleration_warning("cuda", "cpu", 0)

    assert message is not None
    assert "device=cuda requested" in message
    assert "no CUDA/HIP device" in message
    assert "install-rocm.sh" in message
    assert "\n" not in message


def test_acceleration_warning_auto_names_missing_device(monkeypatch):
    message = accel.acceleration_warning("auto", "cpu", 0)

    assert message is not None
    assert "device=auto resolved to CPU" in message
    assert "no CUDA/HIP" in message
    assert "install-rocm.sh" in message
    assert "\n" not in message


def test_transcriber_sets_allocator_before_model_creation(monkeypatch, caplog):
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)
    monkeypatch.setattr(accel, "is_rocm_host", lambda: True)
    monkeypatch.setattr(transcriber_mod, "ct2_cuda_device_count", lambda: 1)
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _RecordingWhisperModel)

    with caplog.at_level(logging.INFO, logger="stt_daemon.transcriber"):
        transcriber = Transcriber(
            "small", device="cpu", compute_type="int8", language=None
        )

    assert _RecordingWhisperModel.last is not None
    assert _RecordingWhisperModel.last.allocator_at_creation == "cub_caching"
    assert transcriber.device == "cpu"
    assert "device=cpu" in caplog.text
    assert "ct2_allocator=cub_caching" in caplog.text


def test_transcriber_cuda_fallback_logs_without_stdout(monkeypatch, caplog, capsys):
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)
    monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip=None))
    monkeypatch.setattr(transcriber_mod, "ct2_cuda_device_count", lambda: 0)
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _RecordingWhisperModel)

    with caplog.at_level(logging.WARNING, logger="stt_daemon.transcriber"):
        transcriber = Transcriber(
            "small", device="cuda", compute_type="float16", language=None
        )

    assert capsys.readouterr().out == ""
    assert transcriber.device == "cpu"
    assert transcriber.compute_type == "int8"
    assert "CUDA requested but not available" in caplog.text


@pytest.mark.parametrize(
    ("compute_type", "expected"),
    [
        ("float16", "int8"),
        ("int8_float16", "int8"),
        ("bfloat16", "int8"),
        ("int8_bfloat16", "int8"),
        ("int8", "int8"),
        ("int8_float32", "int8_float32"),
        ("float32", "float32"),
    ],
)
def test_cpu_safe_compute_type(compute_type, expected):
    assert accel.cpu_safe_compute_type(compute_type) == expected


def test_transcriber_downgrades_gpu_compute_type_when_no_device(monkeypatch, caplog):
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)
    monkeypatch.setattr(accel, "is_rocm_host", lambda: True)
    monkeypatch.setattr(transcriber_mod, "ct2_cuda_device_count", lambda: 0)
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _RecordingWhisperModel)

    with caplog.at_level(logging.WARNING, logger="stt_daemon.transcriber"):
        transcriber = Transcriber(
            "small", device="auto", compute_type="float16", language=None
        )

    assert _RecordingWhisperModel.last is not None
    assert _RecordingWhisperModel.last.device_arg == "cpu"
    assert _RecordingWhisperModel.last.compute_type_arg == "int8"
    assert transcriber.device == "cpu"
    assert transcriber.compute_type == "int8"
    assert "device=auto resolved to CPU" in caplog.text
    assert "needs a GPU" in caplog.text


def test_transcriber_downgrades_gpu_compute_type_on_explicit_cpu(monkeypatch, caplog):
    monkeypatch.delenv(accel.CT2_ALLOCATOR_ENV, raising=False)
    monkeypatch.setattr(accel, "is_rocm_host", lambda: False)
    monkeypatch.setattr(transcriber_mod, "ct2_cuda_device_count", lambda: 1)
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _RecordingWhisperModel)

    with caplog.at_level(logging.WARNING, logger="stt_daemon.transcriber"):
        Transcriber("small", device="cpu", compute_type="float16", language=None)

    assert _RecordingWhisperModel.last is not None
    assert _RecordingWhisperModel.last.compute_type_arg == "int8"
    assert "needs a GPU" in caplog.text