Merge pull request #3684 from ai-coustics/goedev/aic-model-caching

AIC model caching
This commit is contained in:
Mark Backman
2026-02-16 10:43:14 -05:00
committed by GitHub
3 changed files with 398 additions and 30 deletions

View File

@@ -5,6 +5,7 @@
#
import asyncio
import time
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
@@ -23,6 +24,13 @@ except ImportError:
AIC_FILTER_MODULE = "pipecat.audio.filters.aic_filter"
def _model_manager_ref_count(manager, key: str) -> int:
"""Test helper: return reference count for a cache key (reads internal cache)."""
with manager._lock:
entry = manager._cache.get(key)
return entry[1] if entry else 0
class MockProcessor:
"""A lightweight mock for AIC ProcessorAsync that mimics real behavior."""
@@ -99,10 +107,11 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
@classmethod
def setUpClass(cls):
"""Import AICFilter after confirming aic_sdk is available."""
from pipecat.audio.filters.aic_filter import AICFilter
from pipecat.audio.filters.aic_filter import AICFilter, AICModelManager
from pipecat.frames.frames import FilterEnableFrame
cls.AICFilter = AICFilter
cls.AICModelManager = AICModelManager
cls.FilterEnableFrame = FilterEnableFrame
def setUp(self):
@@ -122,13 +131,13 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
async def _start_filter_with_mocks(self, filter_instance, sample_rate=16000):
"""Start a filter with mocked SDK components."""
cache_key = "test-cache-key"
with (
patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls,
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor),
):
mock_model_cls.from_file.return_value = self.mock_model
mock_model_cls.download_async = AsyncMock(return_value="/tmp/model")
mock_manager_cls.acquire = AsyncMock(return_value=(self.mock_model, cache_key))
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(sample_rate)
@@ -171,37 +180,44 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
filter_instance = self._create_filter_with_mocks(model_id=None, model_path=model_path)
with (
patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls,
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor),
):
mock_model_cls.from_file.return_value = self.mock_model
mock_manager_cls.acquire = AsyncMock(
return_value=(self.mock_model, "path:/tmp/test.aicmodel")
)
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000)
mock_model_cls.from_file.assert_called_once_with(str(model_path))
mock_manager_cls.acquire.assert_called_once()
call_kw = mock_manager_cls.acquire.call_args[1]
self.assertEqual(call_kw["model_path"], model_path)
self.assertIsNone(call_kw["model_id"])
self.assertTrue(filter_instance._aic_ready)
self.assertEqual(filter_instance._sample_rate, 16000)
self.assertEqual(filter_instance._frames_per_block, 160)
async def test_start_with_model_id_downloads(self):
"""Test starting filter with model_id triggers download."""
"""Test starting filter with model_id uses manager (download happens in manager)."""
filter_instance = self._create_filter_with_mocks()
with (
patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls,
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor),
):
mock_model_cls.from_file.return_value = self.mock_model
mock_model_cls.download_async = AsyncMock(return_value="/tmp/model")
mock_manager_cls.acquire = AsyncMock(
return_value=(self.mock_model, "id:test-model:/custom/cache")
)
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000)
mock_model_cls.download_async.assert_called_once()
mock_model_cls.from_file.assert_called_once()
mock_manager_cls.acquire.assert_called_once()
call_kw = mock_manager_cls.acquire.call_args[1]
self.assertEqual(call_kw["model_id"], "test-model")
self.assertTrue(filter_instance._aic_ready)
async def test_start_creates_processor(self):
@@ -209,14 +225,13 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
filter_instance = self._create_filter_with_mocks()
with (
patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls,
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(
f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor
) as mock_processor_cls,
):
mock_model_cls.from_file.return_value = self.mock_model
mock_model_cls.download_async = AsyncMock(return_value="/tmp/model")
mock_manager_cls.acquire = AsyncMock(return_value=(self.mock_model, "test-cache-key"))
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000)
@@ -241,17 +256,21 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
self.assertEqual(bypass_params[-1][1], 0.0)
async def test_stop_cleans_up_resources(self):
"""Test that stop properly cleans up resources."""
"""Test that stop properly cleans up resources and releases model reference."""
filter_instance = self._create_filter_with_mocks()
await self._start_filter_with_mocks(filter_instance)
cache_key = filter_instance._model_cache_key
await filter_instance.stop()
with patch(f"{AIC_FILTER_MODULE}.AICModelManager.release") as mock_release:
await filter_instance.stop()
mock_release.assert_called_once_with(cache_key)
self.assertTrue(self.mock_processor.processor_ctx.reset_called)
self.assertIsNone(filter_instance._processor)
self.assertIsNone(filter_instance._processor_ctx)
self.assertIsNone(filter_instance._vad_ctx)
self.assertIsNone(filter_instance._model)
self.assertIsNone(filter_instance._model_cache_key)
self.assertFalse(filter_instance._aic_ready)
async def test_stop_without_start(self):
@@ -261,6 +280,177 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
# Should not raise
await filter_instance.stop()
async def test_model_manager_reference_count(self):
"""Test that AICModelManager reference count increments and decrements correctly."""
model_path = Path("/tmp/refcount-test.aicmodel")
mock_model = MockModel()
manager = self.AICModelManager
with patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls:
mock_model_cls.from_file.return_value = mock_model
# Acquire first reference
model1, key = await manager.acquire(model_path=model_path)
self.assertEqual(model1, mock_model)
self.assertEqual(_model_manager_ref_count(manager, key), 1)
# Acquire second reference (same key, cached)
model2, key2 = await manager.acquire(model_path=model_path)
self.assertIs(model2, model1)
self.assertEqual(key2, key)
self.assertEqual(_model_manager_ref_count(manager, key), 2)
# Release one reference
manager.release(key)
self.assertEqual(_model_manager_ref_count(manager, key), 1)
# Release last reference (model evicted from cache)
manager.release(key)
self.assertEqual(_model_manager_ref_count(manager, key), 0)
async def test_model_manager_concurrent_load_deduplication(self):
"""Test that concurrent acquire calls for the same key share a single load task."""
model_path = Path("/tmp/concurrent-load-test.aicmodel")
mock_model = MockModel()
manager = self.AICModelManager
load_count = 0
def from_file_once(path):
nonlocal load_count
load_count += 1
time.sleep(0.02) # yield so other acquire callers can hit _loading and await same task
return mock_model
with patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls:
mock_model_cls.from_file.side_effect = from_file_once
# Start several acquire calls concurrently before any completes
results = await asyncio.gather(
manager.acquire(model_path=model_path),
manager.acquire(model_path=model_path),
manager.acquire(model_path=model_path),
)
self.assertEqual(
load_count, 1, "Model.from_file should be called once for concurrent callers"
)
model1, key1 = results[0]
model2, key2 = results[1]
model3, key3 = results[2]
self.assertIs(model1, mock_model)
self.assertIs(model2, mock_model)
self.assertIs(model3, mock_model)
self.assertEqual(key1, key2)
self.assertEqual(key2, key3)
self.assertEqual(_model_manager_ref_count(manager, key1), 3)
# Release all references
manager.release(key1)
manager.release(key1)
manager.release(key1)
self.assertEqual(_model_manager_ref_count(manager, key1), 0)
async def test_load_model_from_file_invalid_args_raises(self):
"""Test _load_model_from_file defensive else: raises ValueError."""
manager = self.AICModelManager
with self.assertRaises(ValueError) as ctx:
await manager._load_model_from_file(
"key",
model_path=None,
model_id=None,
model_download_dir=None,
)
self.assertIn("Unexpected", str(ctx.exception))
async def test_model_manager_acquire_by_model_id_hits_download_path(self):
"""Test acquire with model_id runs download path in _load_model_from_file."""
model_id = "test-model-id"
model_download_dir = Path("/tmp/aic-downloads")
mock_model = MockModel()
manager = self.AICModelManager
with patch(f"{AIC_FILTER_MODULE}.Model") as mock_model_cls:
mock_model_cls.download_async = AsyncMock(
return_value="/tmp/aic-downloads/model.aicmodel"
)
mock_model_cls.from_file.return_value = mock_model
model, key = await manager.acquire(
model_id=model_id,
model_download_dir=model_download_dir,
)
mock_model_cls.download_async.assert_called_once()
mock_model_cls.from_file.assert_called_once_with("/tmp/aic-downloads/model.aicmodel")
self.assertIs(model, mock_model)
self.assertEqual(_model_manager_ref_count(manager, key), 1)
manager.release(key)
def test_get_cache_key_invalid_raises(self):
"""Test _get_cache_key raises ValueError for invalid args."""
with self.assertRaises(ValueError) as ctx:
self.AICModelManager._get_cache_key(model_path=None, model_id=None)
self.assertIn("model_path", str(ctx.exception))
with self.assertRaises(ValueError) as ctx2:
self.AICModelManager._get_cache_key(
model_path=None,
model_id="x",
model_download_dir=None,
)
self.assertIn("model_download_dir", str(ctx2.exception))
async def test_start_processor_init_failure(self):
"""Test start() when ProcessorAsync raises: exception logged, _aic_ready False."""
filter_instance = self._create_filter_with_mocks()
with (
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(
f"{AIC_FILTER_MODULE}.ProcessorAsync",
side_effect=RuntimeError("SDK init failed"),
),
):
mock_manager_cls.acquire = AsyncMock(return_value=(self.mock_model, "test-key"))
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000)
self.assertIsNone(filter_instance._processor)
self.assertFalse(filter_instance._aic_ready)
async def test_start_parameter_fixed_error_logged(self):
"""Test start() when set_parameter raises ParameterFixedError: logged, no raise."""
filter_instance = self._create_filter_with_mocks()
self.mock_processor.processor_ctx.set_parameter = MagicMock(
side_effect=aic_sdk.ParameterFixedError("fixed")
)
with (
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls,
patch(f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor),
):
mock_manager_cls.acquire = AsyncMock(return_value=(self.mock_model, "test-key"))
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000)
self.assertTrue(filter_instance._aic_ready)
async def test_process_frame_set_parameter_exception_logged(self):
"""Test process_frame when set_parameter raises: exception logged, no raise."""
filter_instance = self._create_filter_with_mocks()
await self._start_filter_with_mocks(filter_instance)
filter_instance._processor_ctx.set_parameter = MagicMock(
side_effect=ValueError("param error")
)
await filter_instance.process_frame(self.FilterEnableFrame(enable=True))
self.assertFalse(filter_instance._bypass)
async def test_process_frame_enable(self):
"""Test processing FilterEnableFrame to enable filtering."""
filter_instance = self._create_filter_with_mocks()