Handle CUDA context loss gracefully

Fall back to CPU for speaker identification and fail speech recognition with an actionable restart message when the NVIDIA driver resets.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

tests/test_gpu_lost.py

@@ -0,0 +1,152 @@
+"""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()