fix: cache pyannote diarization pipeline to stop memory leak / OOM
The meeting pipeline created a fresh Diarizer per recording, each loading the multi-GB pyannote speaker-diarization model anew (api/pipeline.py). Whisper and Ollama run remotely in this deployment, so pyannote was the only heavy in-process consumer. Reloading it per recording leaked CPU memory (torch reference cycles + glibc arena fragmentation) that was never returned to the OS, climbing to a 37 GB peak over a multi-day run until the kernel OOM-killed the service. Cache the loaded pipeline on the class and reuse it across Diarizer instances, mirroring TranscriptionEngine._model. RSS now stays flat.
This commit is contained in:
+13
-3
@@ -2,19 +2,29 @@ import asyncio
|
|||||||
|
|
||||||
|
|
||||||
class Diarizer:
|
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):
|
def __init__(self, hf_token: str):
|
||||||
if not hf_token:
|
if not hf_token:
|
||||||
raise ValueError("hf_token is required for diarization")
|
raise ValueError("hf_token is required for diarization")
|
||||||
self._hf_token = hf_token
|
self._hf_token = hf_token
|
||||||
self._pipeline = None
|
self._pipeline = None # per-instance override (used by tests)
|
||||||
|
|
||||||
def _load_pipeline(self):
|
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
|
from pyannote.audio import Pipeline
|
||||||
self._pipeline = Pipeline.from_pretrained(
|
Diarizer._shared_pipeline = Pipeline.from_pretrained(
|
||||||
"pyannote/speaker-diarization-3.1",
|
"pyannote/speaker-diarization-3.1",
|
||||||
token=self._hf_token,
|
token=self._hf_token,
|
||||||
)
|
)
|
||||||
|
self._pipeline = Diarizer._shared_pipeline
|
||||||
return self._pipeline
|
return self._pipeline
|
||||||
|
|
||||||
async def diarize(self, wav_path: str) -> list[tuple[float, float, str]]:
|
async def diarize(self, wav_path: str) -> list[tuple[float, float, str]]:
|
||||||
|
|||||||
@@ -38,3 +38,26 @@ def test_diarizer_requires_hf_token():
|
|||||||
from diarization import Diarizer
|
from diarization import Diarizer
|
||||||
with pytest.raises(ValueError, match="hf_token"):
|
with pytest.raises(ValueError, match="hf_token"):
|
||||||
Diarizer(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
|
||||||
|
|||||||
Reference in New Issue
Block a user