diff --git a/diarization.py b/diarization.py index ade0e39..6dae63a 100644 --- a/diarization.py +++ b/diarization.py @@ -2,19 +2,29 @@ import asyncio class Diarizer: + # The pyannote pipeline holds multi-GB torch models. A fresh instance is + # created per recording (see api/pipeline.py), so cache the loaded pipeline + # on the class and reuse it — reloading per recording leaks CPU memory + # (torch reference cycles + glibc arena fragmentation) and eventually OOMs + # the long-running service. Mirrors TranscriptionEngine._model. + _shared_pipeline = None + def __init__(self, hf_token: str): if not hf_token: raise ValueError("hf_token is required for diarization") self._hf_token = hf_token - self._pipeline = None + self._pipeline = None # per-instance override (used by tests) def _load_pipeline(self): - if self._pipeline is None: + if self._pipeline is not None: + return self._pipeline + if Diarizer._shared_pipeline is None: from pyannote.audio import Pipeline - self._pipeline = Pipeline.from_pretrained( + Diarizer._shared_pipeline = Pipeline.from_pretrained( "pyannote/speaker-diarization-3.1", token=self._hf_token, ) + self._pipeline = Diarizer._shared_pipeline return self._pipeline async def diarize(self, wav_path: str) -> list[tuple[float, float, str]]: diff --git a/tests/test_diarization.py b/tests/test_diarization.py index 1becbf6..5dba2d4 100644 --- a/tests/test_diarization.py +++ b/tests/test_diarization.py @@ -38,3 +38,26 @@ def test_diarizer_requires_hf_token(): from diarization import Diarizer with pytest.raises(ValueError, match="hf_token"): Diarizer(hf_token="") + + +def test_pipeline_loaded_once_across_instances(): + """The heavy pyannote pipeline must be loaded once and shared, not reloaded + per recording — reloading leaks torch/CPU memory and OOM-kills the service.""" + import sys, types + from diarization import Diarizer + + Diarizer._shared_pipeline = None # reset shared cache for the test + + fake_module = types.ModuleType("pyannote.audio") + fake_pipeline_cls = MagicMock() + fake_pipeline_cls.from_pretrained.return_value = MagicMock(name="loaded_pipeline") + fake_module.Pipeline = fake_pipeline_cls + + with patch.dict(sys.modules, {"pyannote.audio": fake_module}): + first = Diarizer(hf_token="tok")._load_pipeline() + second = Diarizer(hf_token="tok")._load_pipeline() + + assert first is second + fake_pipeline_cls.from_pretrained.assert_called_once() + + Diarizer._shared_pipeline = None # avoid leaking mock into other tests