fix(vosk): refuse to load a model that would exceed available memory (#850)
* fix(vosk): refuse to load a model that would exceed available memory vosk.Model() allocates the whole model in one spike during loading, including the RNNLM rescoring component some models ship, with no way to check the requirement up front or load incrementally. Walk the model directory, compare its on-disk size against psutil's available memory with a 1.5x margin, and raise before calling into the native loader instead of letting the kernel OOM killer take the process down mid-load. Fixes #676 * fix(vosk): match memory safety margin to the documented peak greptile flagged the guard's own comment: loading can spike to ~2x the model's on-disk size, but the margin only required 1.5x available memory, so a load between those two multiples could still trigger the OOM this guard exists to prevent. --------- Co-authored-by: Jatin K Malik <jatinkrmalik@gmail.com>
This commit is contained in:
Amir Fathi
co-authored by
Jatin K Malik
Jatin K Malik
parent
74d6f505d4
commit
a38cff59ea
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user