wehub-resource-sync
345 lines
13 KiB
Python
345 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for audio engine memory tracking in EnginePool (INV-06).
|
|
|
|
Verifies that audio (STT/TTS) engines participate in the same LRU memory
|
|
management lifecycle as LLM/VLM/embedding engines:
|
|
- Loading updates _current_model_memory
|
|
- Unloading decrements _current_model_memory
|
|
- last_access is updated on get_engine()
|
|
- Audio engines are eligible for LRU eviction unless pinned
|
|
- _find_lru_victim() can select an audio model
|
|
- Pre-load eviction evicts audio when memory is tight
|
|
|
|
All tests run with mocked engines — mlx-audio is not required.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from omlx.engine_pool import EngineEntry, EnginePool
|
|
|
|
|
|
def _engine_pool_with_ceiling(ceiling=None):
|
|
"""Helper: EnginePool with a stubbed pre-load ceiling callback."""
|
|
pool = EnginePool()
|
|
if ceiling and ceiling > 0:
|
|
pool._get_final_ceiling = lambda c=ceiling: c
|
|
else:
|
|
pool._get_final_ceiling = lambda: 0
|
|
return pool
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def audio_model_dir(tmp_path):
|
|
"""Model directory with one LLM and one STT model (small sizes for fast tests)."""
|
|
llm_dir = tmp_path / "llama-3b"
|
|
llm_dir.mkdir()
|
|
(llm_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(llm_dir / "model.safetensors").write_bytes(b"0" * 1024) # ~1KB
|
|
|
|
stt_dir = tmp_path / "whisper-tiny"
|
|
stt_dir.mkdir()
|
|
(stt_dir / "config.json").write_text(json.dumps({
|
|
"model_type": "whisper",
|
|
"architectures": ["WhisperForConditionalGeneration"],
|
|
}))
|
|
(stt_dir / "model.safetensors").write_bytes(b"0" * 2048) # ~2KB
|
|
|
|
tts_dir = tmp_path / "kokoro-tts"
|
|
tts_dir.mkdir()
|
|
(tts_dir / "config.json").write_text(json.dumps({"model_type": "qwen3_tts"}))
|
|
(tts_dir / "model.safetensors").write_bytes(b"0" * 1536) # ~1.5KB
|
|
|
|
return tmp_path
|
|
|
|
|
|
@pytest.fixture
|
|
def pool_with_audio(audio_model_dir):
|
|
"""EnginePool with audio + LLM models, generous memory limit."""
|
|
pool = _engine_pool_with_ceiling(10 * 1024**3)
|
|
pool.discover_models(str(audio_model_dir))
|
|
return pool
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestAudioMemoryTracking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAudioMemoryTracking:
|
|
"""Loading and unloading audio engines updates _current_model_memory."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_loading_stt_updates_memory(self, pool_with_audio):
|
|
"""Loading an STT engine increments _current_model_memory."""
|
|
pool = pool_with_audio
|
|
assert pool.current_model_memory == 0
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
mock_engine.stop = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_engine, create=True):
|
|
await pool.get_engine("whisper-tiny")
|
|
|
|
assert pool.current_model_memory > 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_loading_tts_updates_memory(self, pool_with_audio):
|
|
"""Loading a TTS engine increments _current_model_memory."""
|
|
pool = pool_with_audio
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
mock_engine.stop = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.TTSEngine", return_value=mock_engine, create=True):
|
|
await pool.get_engine("kokoro-tts")
|
|
|
|
assert pool.current_model_memory > 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unloading_stt_decrements_memory(self, pool_with_audio):
|
|
"""Unloading an STT engine decrements _current_model_memory."""
|
|
pool = pool_with_audio
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
mock_engine.stop = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_engine, create=True):
|
|
await pool.get_engine("whisper-tiny")
|
|
memory_after_load = pool.current_model_memory
|
|
assert memory_after_load > 0
|
|
|
|
await pool._unload_engine("whisper-tiny")
|
|
|
|
assert pool.current_model_memory < memory_after_load
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unload_clears_engine_reference(self, pool_with_audio):
|
|
"""After unload, EngineEntry.engine is None."""
|
|
pool = pool_with_audio
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
mock_engine.stop = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_engine, create=True):
|
|
await pool.get_engine("whisper-tiny")
|
|
await pool._unload_engine("whisper-tiny")
|
|
|
|
assert pool._entries["whisper-tiny"].engine is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestAudioLastAccess
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAudioLastAccess:
|
|
"""last_access is updated when an audio engine is retrieved."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_engine_updates_last_access(self, pool_with_audio):
|
|
"""get_engine() updates last_access timestamp for audio entry."""
|
|
pool = pool_with_audio
|
|
entry = pool._entries["whisper-tiny"]
|
|
assert entry.last_access == 0.0
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_engine, create=True):
|
|
with patch("time.time", return_value=1234.0):
|
|
await pool.get_engine("whisper-tiny")
|
|
|
|
assert entry.last_access == 1234.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_second_get_engine_refreshes_last_access(self, pool_with_audio):
|
|
"""Second call to get_engine() refreshes last_access."""
|
|
pool = pool_with_audio
|
|
|
|
mock_engine = MagicMock()
|
|
mock_engine.start = AsyncMock()
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_engine, create=True):
|
|
with patch("time.time", return_value=1000.0):
|
|
await pool.get_engine("whisper-tiny")
|
|
|
|
with patch("time.time", return_value=2000.0):
|
|
await pool.get_engine("whisper-tiny")
|
|
|
|
assert pool._entries["whisper-tiny"].last_access == 2000.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestAudioLRUEviction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAudioLRUEviction:
|
|
"""Audio engines are eligible for LRU eviction by default."""
|
|
|
|
def test_audio_entry_not_pinned_by_default(self, pool_with_audio):
|
|
"""STT and TTS entries are not pinned by default."""
|
|
pool = pool_with_audio
|
|
assert pool._entries["whisper-tiny"].is_pinned is False
|
|
assert pool._entries["kokoro-tts"].is_pinned is False
|
|
|
|
def test_find_lru_victim_can_select_stt(self, pool_with_audio):
|
|
"""_find_lru_victim() returns STT model when it is the oldest loaded entry."""
|
|
pool = pool_with_audio
|
|
|
|
# Mark whisper-tiny as loaded and older than all others
|
|
mock_engine = MagicMock()
|
|
mock_engine.has_active_requests.return_value = False
|
|
pool._entries["whisper-tiny"].engine = mock_engine
|
|
pool._entries["whisper-tiny"].last_access = 10.0
|
|
|
|
victim = pool._find_lru_victim()
|
|
assert victim == "whisper-tiny"
|
|
|
|
def test_find_lru_victim_selects_oldest_audio_over_newer_llm(self, pool_with_audio):
|
|
"""_find_lru_victim() picks the oldest entry regardless of model type."""
|
|
pool = pool_with_audio
|
|
|
|
mock_stt = MagicMock()
|
|
mock_stt.has_active_requests.return_value = False
|
|
pool._entries["whisper-tiny"].engine = mock_stt
|
|
pool._entries["whisper-tiny"].last_access = 50.0 # Older
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.has_active_requests.return_value = False
|
|
pool._entries["llama-3b"].engine = mock_llm
|
|
pool._entries["llama-3b"].last_access = 100.0 # Newer
|
|
|
|
victim = pool._find_lru_victim()
|
|
assert victim == "whisper-tiny"
|
|
|
|
def test_pinned_audio_not_evicted(self, pool_with_audio):
|
|
"""Pinned audio engine is skipped by _find_lru_victim()."""
|
|
pool = pool_with_audio
|
|
|
|
mock_stt = MagicMock()
|
|
mock_stt.has_active_requests.return_value = False
|
|
pool._entries["whisper-tiny"].engine = mock_stt
|
|
pool._entries["whisper-tiny"].last_access = 50.0
|
|
pool._entries["whisper-tiny"].is_pinned = True # pinned
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.has_active_requests.return_value = False
|
|
pool._entries["llama-3b"].engine = mock_llm
|
|
pool._entries["llama-3b"].last_access = 100.0 # Newer but not pinned
|
|
|
|
victim = pool._find_lru_victim()
|
|
# whisper-tiny is pinned — llama-3b must be the victim
|
|
assert victim == "llama-3b"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestAudioPinning
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAudioPinning:
|
|
"""Audio engines can be pinned to prevent eviction."""
|
|
|
|
def test_discover_with_pinned_audio(self, audio_model_dir):
|
|
"""discover_models() with pinned_models pins the audio entry."""
|
|
pool = _engine_pool_with_ceiling(10 * 1024**3)
|
|
pool.discover_models(str(audio_model_dir), pinned_models=["whisper-tiny"])
|
|
|
|
assert pool._entries["whisper-tiny"].is_pinned is True
|
|
assert pool._entries["llama-3b"].is_pinned is False
|
|
|
|
def test_pinned_audio_not_selected_as_lru_victim(self, audio_model_dir):
|
|
"""Pinned audio model is excluded from LRU eviction candidates."""
|
|
pool = _engine_pool_with_ceiling(10 * 1024**3)
|
|
pool.discover_models(str(audio_model_dir), pinned_models=["whisper-tiny"])
|
|
|
|
pool._entries["whisper-tiny"].engine = MagicMock()
|
|
pool._entries["whisper-tiny"].last_access = 1.0 # Oldest
|
|
|
|
pool._entries["llama-3b"].engine = MagicMock()
|
|
pool._entries["llama-3b"].last_access = 99.0
|
|
|
|
victim = pool._find_lru_victim()
|
|
assert victim != "whisper-tiny"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestAudioPreLoadEviction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAudioPreLoadEviction:
|
|
"""Pre-load eviction works when loading an audio model requires freeing memory."""
|
|
|
|
@pytest.fixture
|
|
def tight_audio_pool(self, tmp_path):
|
|
"""Pool tight enough that only one model fits at a time."""
|
|
llm_dir = tmp_path / "llama-3b"
|
|
llm_dir.mkdir()
|
|
(llm_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(llm_dir / "model.safetensors").write_bytes(b"0" * 1024)
|
|
|
|
stt_dir = tmp_path / "whisper-tiny"
|
|
stt_dir.mkdir()
|
|
(stt_dir / "config.json").write_text(json.dumps({
|
|
"model_type": "whisper",
|
|
"architectures": ["WhisperForConditionalGeneration"],
|
|
}))
|
|
(stt_dir / "model.safetensors").write_bytes(b"0" * 2048)
|
|
|
|
# Limit: allow one model but not both simultaneously
|
|
pool = _engine_pool_with_ceiling(2500)
|
|
pool.discover_models(str(tmp_path))
|
|
return pool
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_loading_stt_evicts_llm(self, tight_audio_pool, monkeypatch):
|
|
"""When memory is tight, loading STT evicts the loaded LLM."""
|
|
pool = tight_audio_pool
|
|
# Proxy phys_footprint to the pool's tracked weight sum so admission
|
|
# against the byte-sized synthetic ceiling matches the test's intent
|
|
# (real phys_footprint is ~100 MB and would dominate).
|
|
monkeypatch.setattr(
|
|
"omlx.engine_pool.get_phys_footprint",
|
|
lambda: pool._current_model_memory,
|
|
)
|
|
monkeypatch.setattr(
|
|
"omlx.engine_pool.mx.get_active_memory", lambda: 0
|
|
)
|
|
|
|
mock_llm = MagicMock()
|
|
mock_llm.start = AsyncMock()
|
|
mock_llm.stop = AsyncMock()
|
|
mock_llm.has_active_requests.return_value = False
|
|
|
|
mock_stt = MagicMock()
|
|
mock_stt.start = AsyncMock()
|
|
mock_stt.stop = AsyncMock()
|
|
mock_stt.has_active_requests.return_value = False
|
|
|
|
with patch("omlx.engine_pool.BatchedEngine", return_value=mock_llm):
|
|
await pool.get_engine("llama-3b")
|
|
|
|
with patch("omlx.engine_pool.STTEngine", return_value=mock_stt, create=True):
|
|
await pool.get_engine("whisper-tiny")
|
|
|
|
# llama-3b should have been evicted
|
|
mock_llm.stop.assert_called_once()
|
|
assert pool._entries["llama-3b"].engine is None
|
|
assert pool._entries["whisper-tiny"].engine is not None
|