javis/tests/test_gpu_lost.py

"""Surviving a dead CUDA context (e.g. NVIDIA driver updated while running)."""
from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np

from javis.core.config import Config, PathsConfig
from javis.core.hardware import GpuUnavailableError, is_cuda_error
from javis.stt.transcriber import Transcriber
from javis.stt.voice_profiler import VoiceProfiler

try:
    import torch
except ImportError:  # pragma: no cover
    torch = None

SR = 16000
_AUDIO = np.full(SR, 0.1, dtype=np.float32)


class IsCudaErrorTests(unittest.TestCase):
    def test_classification(self) -> None:
        self.assertTrue(is_cuda_error(RuntimeError("CUDA error: unknown error")))
        self.assertTrue(is_cuda_error(RuntimeError("CUDA failed with error unknown error")))
        self.assertTrue(is_cuda_error(RuntimeError("cuBLAS failed with status X")))
        self.assertFalse(is_cuda_error(RuntimeError("CUDA failed with error out of memory")))
        self.assertFalse(is_cuda_error(ValueError("bad input")))


class _DeadModel:
    """Fake WhisperModel whose lazy segment generator hits a dead context."""

    def __init__(self, msg: str = "CUDA failed with error unknown error") -> None:
        self.msg = msg
        self.calls = 0

    def transcribe(self, audio, **kw):
        self.calls += 1

        def segments():
            raise RuntimeError(self.msg)
            yield  # pragma: no cover

        return segments(), SimpleNamespace(language="de", duration=1.0)


def _transcriber(model) -> Transcriber:
    stt = Transcriber(Config())
    stt._model = model
    stt._batched = None
    stt._device = "cuda"
    return stt


class TranscriberGpuLostTests(unittest.TestCase):
    def test_dead_context_raises_and_then_fails_fast(self) -> None:
        model = _DeadModel()
        stt = _transcriber(model)
        with patch.object(Transcriber, "_cuda_alive", return_value=False):
            with self.assertRaisesRegex(GpuUnavailableError, "Restart Javis"):
                stt._transcribe_sync(_AUDIO, SR)
            with self.assertRaises(GpuUnavailableError):
                stt._transcribe_sync(_AUDIO, SR)
        self.assertEqual(model.calls, 1)      # second call never touched the GPU
        self.assertIs(stt._model, model)      # freeing it would abort the process
        self.assertFalse(stt._force_cpu)

    def test_cuda_error_on_live_context_is_not_gpu_lost(self) -> None:
        stt = _transcriber(_DeadModel("CUDA failed with error invalid argument"))
        with patch.object(Transcriber, "_cuda_alive", return_value=True):
            with self.assertRaises(RuntimeError) as cm:
                stt._transcribe_sync(_AUDIO, SR)
        self.assertNotIsInstance(cm.exception, GpuUnavailableError)
        self.assertFalse(stt._gpu_lost)

    def test_cpu_model_errors_skip_the_probe(self) -> None:
        stt = _transcriber(_DeadModel())
        stt._device = "cpu"
        with patch.object(Transcriber, "_cuda_alive") as probe:
            with self.assertRaises(RuntimeError) as cm:
                stt._transcribe_sync(_AUDIO, SR)
        probe.assert_not_called()
        self.assertNotIsInstance(cm.exception, GpuUnavailableError)


class _BrokenCudaEncoder:
    def encode_batch(self, wav, lens):
        raise RuntimeError("CUDA error: unknown error")


class _CpuEncoder:
    def __init__(self) -> None:
        self.calls = 0

    def encode_batch(self, wav, lens):
        self.calls += 1
        return torch.ones(1, 1, 192)


@unittest.skipIf(torch is None, "torch not installed")
class ProfilerCpuFallbackTests(unittest.TestCase):
    def _profiler(self, base: Path) -> VoiceProfiler:
        prof = VoiceProfiler(Config(paths=PathsConfig(base_dir=base)))
        prof._encoder = _BrokenCudaEncoder()
        prof._device = "cuda:0"
        return prof

    def test_cuda_error_switches_encoder_to_cpu_once(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            prof = self._profiler(Path(d))
            cpu = _CpuEncoder()
            with patch.object(VoiceProfiler, "_build_encoder", return_value=cpu) as build:
                first = prof.embed(_AUDIO, SR)
                prof.embed(_AUDIO, SR)
            build.assert_called_once_with("cpu")
            self.assertIs(prof._encoder, cpu)
            self.assertEqual(prof._device, "cpu")
            self.assertEqual(cpu.calls, 2)
            self.assertAlmostEqual(float(np.linalg.norm(first)), 1.0, places=5)

    def test_errors_on_cpu_are_not_swallowed(self) -> None:
        with tempfile.TemporaryDirectory() as d:
            prof = self._profiler(Path(d))
            prof._device = "cpu"
            with patch.object(VoiceProfiler, "_build_encoder") as build:
                with self.assertRaises(RuntimeError):
                    prof.embed(_AUDIO, SR)
            build.assert_not_called()


class IdentifySpeakerTests(unittest.IsolatedAsyncioTestCase):
    async def test_identify_failure_returns_sentinel(self) -> None:
        from javis.agent.orchestrator import _ID_FAILED, Orchestrator

        orch = object.__new__(Orchestrator)

        def boom(audio, sr):
            raise RuntimeError("CUDA error: unknown error")

        orch.profiler = SimpleNamespace(identify=boom)
        self.assertIs(await orch._identify_speaker(_AUDIO, SR), _ID_FAILED)
        orch.profiler = SimpleNamespace(identify=lambda a, sr: ("Florian", 0.8))
        self.assertEqual(await orch._identify_speaker(_AUDIO, SR), ("Florian", 0.8))


if __name__ == "__main__":
    unittest.main()