thomas.kopp
6b0ee60d94
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.
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
from unittest.mock import MagicMock, patch
|
|
import pytest
|
|
|
|
|
|
def test_diarizer_returns_list_of_tuples(tmp_path):
|
|
"""Diarizer.diarize() returns [(start, end, speaker), ...]"""
|
|
wav = tmp_path / "test.wav"
|
|
wav.write_bytes(b"\x00" * 100)
|
|
|
|
mock_turn_1 = MagicMock()
|
|
mock_turn_1.start = 0.0
|
|
mock_turn_1.end = 2.5
|
|
|
|
mock_turn_2 = MagicMock()
|
|
mock_turn_2.start = 2.6
|
|
mock_turn_2.end = 5.0
|
|
|
|
mock_annotation = MagicMock()
|
|
mock_annotation.itertracks.return_value = [
|
|
(mock_turn_1, "A", "SPEAKER_00"),
|
|
(mock_turn_2, "B", "SPEAKER_01"),
|
|
]
|
|
|
|
mock_output = MagicMock()
|
|
mock_output.speaker_diarization = mock_annotation
|
|
mock_pipeline = MagicMock(return_value=mock_output)
|
|
|
|
import asyncio
|
|
from diarization import Diarizer
|
|
d = Diarizer.__new__(Diarizer)
|
|
d._pipeline = mock_pipeline
|
|
|
|
result = asyncio.run(d.diarize(str(wav)))
|
|
assert result == [(0.0, 2.5, "SPEAKER_00"), (2.6, 5.0, "SPEAKER_01")]
|
|
|
|
|
|
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
|