From 6b0ee60d94bb9164f6dc6a08bf15349408bb902d Mon Sep 17 00:00:00 2001
From: "thomas.kopp" <thomas.kopp@tueit.de>
Date: Wed, 22 Jul 2026 09:03:42 +0200
Subject: [PATCH] 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.
---
 diarization.py            | 16 +++++++++++++---
 tests/test_diarization.py | 23 +++++++++++++++++++++++
 2 files changed, 36 insertions(+), 3 deletions(-)

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