diff --git a/src/vocalinux/speech_recognition/recognition_manager.py b/src/vocalinux/speech_recognition/recognition_manager.py index a7763ea..4fc6a90 100644 --- a/src/vocalinux/speech_recognition/recognition_manager.py +++ b/src/vocalinux/speech_recognition/recognition_manager.py @@ -1036,6 +1036,12 @@ def _get_system_model_paths() -> list: # Alternative locations for pre-installed models (now dynamic) SYSTEM_MODELS_DIRS = _get_system_model_paths() +# vosk.Model()'s peak RSS during loading runs well above the model's on-disk +# size (issue #676: an RNNLM-bearing model's load-time spike was ~1.5-2x its +# file size). Require this multiple of available memory before attempting +# the load. +_VOSK_MODEL_MEMORY_SAFETY_MARGIN = 2.0 + def detect_pywhispercpp_gpu_backend() -> str: """Detect whether pywhispercpp's native library actually has GPU support.""" @@ -1296,6 +1302,8 @@ class SpeechRecognitionManager: else: logger.info(f"Using existing VOSK model from {self.vosk_model_path}") + self._check_vosk_model_memory(self.vosk_model_path) + logger.info(f"Loading VOSK model from {self.vosk_model_path}") # Ensure previous model/recognizer are released if re-initializing self.model = None @@ -1310,6 +1318,56 @@ class SpeechRecognitionManager: self.state = RecognitionState.ERROR raise + @staticmethod + def _vosk_model_size_bytes(model_path: str) -> int: + """Sum on-disk file sizes under a VOSK model directory. + + Walks the whole directory rather than reading a catalog size, since + the RNNLM component (rnnlm/final.raw) that drives the memory spike + this guards against is only present for some languages/models and + is not reflected in VOSK_MODEL_INFO's nominal download size. + """ + total = 0 + for root, _dirs, files in os.walk(model_path): + for name in files: + try: + total += os.path.getsize(os.path.join(root, name)) + except OSError: + continue + return total + + def _check_vosk_model_memory(self, model_path: str) -> None: + """Refuse to load a VOSK model that would exhaust available memory. + + vosk.Model() allocates the whole model (acoustic model plus, for + models that ship one, the RNNLM rescoring component) in one spike + during loading. There is no way to load it incrementally or query + vosk for how much it will need up front, so the only guard + available here is comparing the model's on-disk footprint against + currently available memory before calling into the native loader + (issue #676: this spike killed the process via the kernel OOM + killer on a 16GB machine, with no exception Vocalinux could catch). + """ + import psutil + + model_bytes = self._vosk_model_size_bytes(model_path) + if model_bytes <= 0: + return + + available_bytes = psutil.virtual_memory().available + required_bytes = model_bytes * _VOSK_MODEL_MEMORY_SAFETY_MARGIN + + if available_bytes < required_bytes: + model_mb = model_bytes / (1024 * 1024) + available_mb = available_bytes / (1024 * 1024) + raise RuntimeError( + f"Not enough memory to load the VOSK model at {model_path} " + f"(model is ~{model_mb:.0f}MB on disk, {available_mb:.0f}MB " + "available). Loading it would likely be killed by the " + "system's out-of-memory killer partway through. Close other " + "applications or switch to a smaller model size in Settings." + ) + def _init_whisper(self): """Initialize the Whisper speech recognition engine.""" import warnings diff --git a/tests/test_vosk_memory_guard.py b/tests/test_vosk_memory_guard.py new file mode 100644 index 0000000..17c80aa --- /dev/null +++ b/tests/test_vosk_memory_guard.py @@ -0,0 +1,201 @@ +""" +Tests for the VOSK model-load memory guard (issue #676). + +vosk.Model() allocates a model's full footprint (plus, for models that ship +an RNNLM rescoring component, a further multi-GB spike) in one shot with no +way to query the requirement up front or load incrementally. These tests +exercise the guard that refuses the load rather than letting the kernel OOM +killer take the process down mid-load. +""" + +import sys +from unittest.mock import MagicMock, patch + +import pytest + +if "gi" not in sys.modules: + sys.modules["gi"] = MagicMock() +if "gi.repository" not in sys.modules: + sys.modules["gi.repository"] = MagicMock() + + +@pytest.fixture(autouse=True) +def _restore_sys_modules(): + saved = dict(sys.modules) + yield + added = set(sys.modules.keys()) - set(saved.keys()) + for k in added: + del sys.modules[k] + for k, v in saved.items(): + if k not in sys.modules or sys.modules[k] is not v: + sys.modules[k] = v + + +def _make_manager(**kw): + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + with ( + patch.object(SpeechRecognitionManager, "_init_vosk"), + patch.object(SpeechRecognitionManager, "_init_whisper"), + patch.object(SpeechRecognitionManager, "_init_whispercpp"), + patch.object(SpeechRecognitionManager, "_init_parakeet"), + patch.object(SpeechRecognitionManager, "_init_faster_whisper"), + ): + return SpeechRecognitionManager( + engine="whisper_cpp", + model_size="small", + language="en-us", + defer_download=True, + **kw, + ) + + +def _write_file_of_size(path: str, size_bytes: int) -> None: + with open(path, "wb") as f: + f.write(b"\0" * size_bytes) + + +class TestVoskModelSizeBytes: + """_vosk_model_size_bytes walks the model directory on disk.""" + + def test_sums_all_files_recursively(self, tmp_path): + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + model_dir = tmp_path / "vosk-model-ru-0.22" + (model_dir / "am").mkdir(parents=True) + (model_dir / "rnnlm").mkdir() + _write_file_of_size(str(model_dir / "am" / "final.mdl"), 1000) + _write_file_of_size(str(model_dir / "rnnlm" / "final.raw"), 2000) + + total = SpeechRecognitionManager._vosk_model_size_bytes(str(model_dir)) + assert total == 3000 + + def test_nonexistent_directory_is_zero(self, tmp_path): + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + total = SpeechRecognitionManager._vosk_model_size_bytes(str(tmp_path / "does-not-exist")) + assert total == 0 + + def test_unreadable_entry_is_skipped_not_raised(self, tmp_path): + # os.walk lists a directory's entries up front; a broken symlink (or a + # file removed by a concurrent process) still shows up in that list + # but os.path.getsize() on it raises OSError. The walk must keep + # summing the rest of the model rather than dying on one bad entry. + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + model_dir = tmp_path / "vosk-model-ru-0.22" + model_dir.mkdir() + _write_file_of_size(str(model_dir / "final.mdl"), 1000) + (model_dir / "broken-link").symlink_to(tmp_path / "does-not-exist-target") + + total = SpeechRecognitionManager._vosk_model_size_bytes(str(model_dir)) + assert total == 1000 + + +class TestCheckVoskModelMemory: + """_check_vosk_model_memory raises before the native loader is ever called.""" + + def test_raises_when_available_memory_is_below_the_safety_margin(self, tmp_path): + model_dir = tmp_path / "vosk-model-ru-0.22" + model_dir.mkdir() + # 100MB on disk; guard requires 2.0x = 200MB available. + _write_file_of_size(str(model_dir / "final.raw"), 100 * 1024 * 1024) + + manager = _make_manager() + with patch("psutil.virtual_memory") as mock_mem: + mock_mem.return_value.available = 120 * 1024 * 1024 # below the 200MB requirement + with pytest.raises(RuntimeError, match="Not enough memory"): + manager._check_vosk_model_memory(str(model_dir)) + + def test_passes_when_available_memory_clears_the_safety_margin(self, tmp_path): + model_dir = tmp_path / "vosk-model-small-en-us-0.15" + model_dir.mkdir() + _write_file_of_size(str(model_dir / "final.mdl"), 40 * 1024 * 1024) + + manager = _make_manager() + with patch("psutil.virtual_memory") as mock_mem: + mock_mem.return_value.available = 16 * 1024 * 1024 * 1024 # 16GB, plenty + manager._check_vosk_model_memory(str(model_dir)) # must not raise + + def test_empty_model_directory_is_never_guarded(self, tmp_path): + # A directory with no files (or one os.walk cannot see) sizes to 0. + # There is nothing on disk to compare, so the guard is a no-op here + # by design, not a claim that zero bytes always fits in memory. + model_dir = tmp_path / "empty-model" + model_dir.mkdir() + + manager = _make_manager() + with patch("psutil.virtual_memory") as mock_mem: + mock_mem.return_value.available = 0 + manager._check_vosk_model_memory(str(model_dir)) # must not raise + + +class TestInitVoskMemoryGuardIntegration: + """The guard sits in front of vosk.Model(), inside the real _init_vosk.""" + + def test_insufficient_memory_blocks_the_native_load(self, tmp_path): + model_dir = tmp_path / "vosk-model-ru-0.22" + model_dir.mkdir() + _write_file_of_size(str(model_dir / "final.raw"), 100 * 1024 * 1024) + + vosk_mock = MagicMock() + sys.modules["vosk"] = vosk_mock + + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + with ( + patch.object( + SpeechRecognitionManager, "_get_vosk_model_path", return_value=str(model_dir) + ), + patch("os.path.exists", return_value=True), + patch.object(SpeechRecognitionManager, "_init_whisper"), + patch.object(SpeechRecognitionManager, "_init_whispercpp"), + patch.object(SpeechRecognitionManager, "_init_parakeet"), + patch.object(SpeechRecognitionManager, "_init_faster_whisper"), + patch("psutil.virtual_memory") as mock_mem, + ): + mock_mem.return_value.available = 1 * 1024 * 1024 # far below the requirement + with pytest.raises(RuntimeError, match="Not enough memory"): + SpeechRecognitionManager( + engine="vosk", + model_size="large", + language="ru", + defer_download=True, + ) + + # The control: the same fixture with ample memory reaches vosk.Model(). + # If this assertion were dropped, a guard that (by bug) always raises + # would pass the test above for the wrong reason. + vosk_mock.Model.assert_not_called() + + def test_sufficient_memory_reaches_the_native_load(self, tmp_path): + model_dir = tmp_path / "vosk-model-small-en-us-0.15" + model_dir.mkdir() + _write_file_of_size(str(model_dir / "final.mdl"), 40 * 1024 * 1024) + + vosk_mock = MagicMock() + sys.modules["vosk"] = vosk_mock + + from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager + + with ( + patch.object( + SpeechRecognitionManager, "_get_vosk_model_path", return_value=str(model_dir) + ), + patch("os.path.exists", return_value=True), + patch.object(SpeechRecognitionManager, "_init_whisper"), + patch.object(SpeechRecognitionManager, "_init_whispercpp"), + patch.object(SpeechRecognitionManager, "_init_parakeet"), + patch.object(SpeechRecognitionManager, "_init_faster_whisper"), + patch("psutil.virtual_memory") as mock_mem, + ): + mock_mem.return_value.available = 16 * 1024 * 1024 * 1024 + manager = SpeechRecognitionManager( + engine="vosk", + model_size="small", + language="en-us", + defer_download=True, + ) + + vosk_mock.Model.assert_called_once_with(str(model_dir)) + assert manager._model_initialized is True