# SPDX-License-Identifier: Apache-2.0 """Tests for EnginePool functionality.""" import asyncio import json import logging import shutil from unittest.mock import AsyncMock, MagicMock, patch import pytest from omlx.engine_pool import EngineEntry, EnginePool from omlx.exceptions import ( InsufficientMemoryError, ModelBusyError, ModelLoadingError, ModelNotFoundError, ModelTooLargeError, ModelUnavailableError, ) from omlx.scheduler import PrefillEvictionRequest def _make_pool(ceiling: int | None = None, **kwargs) -> EnginePool: """Helper: create an EnginePool with a stubbed final-ceiling callback. Tests historically constructed EnginePool with `max_model_memory=X`. The pre-load admission ceiling now comes from the ProcessMemoryEnforcer via a callback, so tests inject a fake callback that returns the desired ceiling. `ceiling=None` (or 0) disables the limit (callback returns 0). """ pool = EnginePool(**kwargs) if ceiling is None or ceiling <= 0: pool._get_final_ceiling = lambda: 0 else: pool._get_final_ceiling = lambda c=int(ceiling): c return pool @pytest.fixture def mock_model_dir(tmp_path): """Create a mock model directory with multiple models.""" # Create model-a (1GB) model_a = tmp_path / "model-a" model_a.mkdir() (model_a / "config.json").write_text(json.dumps({"model_type": "llama"})) (model_a / "model.safetensors").write_bytes(b"0" * (1024 * 1024 * 1024)) # 1GB # Create model-b (2GB) model_b = tmp_path / "model-b" model_b.mkdir() (model_b / "config.json").write_text(json.dumps({"model_type": "qwen"})) (model_b / "model.safetensors").write_bytes(b"0" * (2 * 1024 * 1024 * 1024)) # 2GB # Create model-c (500MB MLLM) model_c = tmp_path / "model-c" model_c.mkdir() (model_c / "config.json").write_text(json.dumps({"vision_config": {}})) (model_c / "model.safetensors").write_bytes(b"0" * (512 * 1024 * 1024)) # 500MB return tmp_path @pytest.fixture def small_mock_model_dir(tmp_path): """Create a mock model directory with small models for fast tests.""" # Create model-a (1KB) model_a = tmp_path / "model-a" model_a.mkdir() (model_a / "config.json").write_text(json.dumps({"model_type": "llama"})) (model_a / "model.safetensors").write_bytes(b"0" * 1024) # Create model-b (2KB) model_b = tmp_path / "model-b" model_b.mkdir() (model_b / "config.json").write_text(json.dumps({"model_type": "qwen"})) (model_b / "model.safetensors").write_bytes(b"0" * 2048) return tmp_path class TestEnginePoolInit: """Tests for EnginePool initialization.""" def test_init_default_config(self): """Test initialization with default config.""" pool = _make_pool(ceiling=32 * 1024**3) assert pool._current_ceiling() == 32 * 1024**3 assert pool.current_model_memory == 0 assert pool.model_count == 0 assert pool.loaded_model_count == 0 def test_init_disabled_memory(self): """Test initialization with disabled (None) memory limit.""" pool = _make_pool(ceiling=None) assert pool._current_ceiling() == 0 assert pool.current_model_memory == 0 def test_model_too_large_skipped_when_disabled(self, small_mock_model_dir): """Test that ModelTooLargeError is NOT raised when memory is disabled.""" pool = _make_pool(ceiling=None) pool.discover_models(str(small_mock_model_dir)) # With None (disabled), the size check should be skipped entirely # Model loading will fail for other reasons (mock), but ModelTooLargeError # should not be raised entry = pool.get_entry("model-a") assert entry is not None # Verify the entry has a nonzero estimated size assert entry.estimated_size > 0 def test_discover_models(self, small_mock_model_dir): """Test model discovery.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) assert pool.model_count == 2 assert "model-a" in pool.get_model_ids() assert "model-b" in pool.get_model_ids() def test_discover_models_with_pinned(self, small_mock_model_dir): """Test model discovery with pinned models.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir), pinned_models=["model-a"]) entry = pool.get_entry("model-a") assert entry is not None assert entry.is_pinned is True entry_b = pool.get_entry("model-b") assert entry_b is not None assert entry_b.is_pinned is False class TestExposedProfileModelResolution: """Tests for exposed profile model IDs that share a physical engine.""" def _manager_with_exposed_profile(self, tmp_path): from omlx.model_settings import ModelSettingsManager manager = ModelSettingsManager(tmp_path) manager.save_profile( model_id="model-b", name="thinking", display_name="Thinking", description=None, settings={"enable_thinking": True}, expose_as_model=True, ) return manager def test_resolve_model_id_maps_exposed_profile_to_source_model( self, small_mock_model_dir, tmp_path ): """Exposed profile model IDs resolve to their base model for loading.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) manager = self._manager_with_exposed_profile(tmp_path) resolved = pool.resolve_model_id("model-b:thinking", manager) assert resolved == "model-b" def test_resolve_model_id_maps_provider_prefixed_exposed_profile_to_source( self, small_mock_model_dir, tmp_path ): """Provider prefixes do not prevent exposed profile resolution.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) manager = self._manager_with_exposed_profile(tmp_path) resolved = pool.resolve_model_id("omlx/model-b:thinking", manager) assert resolved == "model-b" class TestDiscoverModelsMerge: """Tests for discover_models merge behavior (issue #89).""" def test_rediscover_preserves_loaded_engine(self, small_mock_model_dir): """Test that re-discovery preserves loaded engine state.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Simulate loaded model-a entry_a = pool.get_entry("model-a") mock_engine = MagicMock() entry_a.engine = mock_engine entry_a.last_access = 42.0 pool._current_model_memory = entry_a.estimated_size original_size = entry_a.estimated_size # Re-discover (simulates model deletion or download completion) pool.discover_models(str(small_mock_model_dir)) # Loaded model should be preserved entry_a_after = pool.get_entry("model-a") assert entry_a_after is entry_a # Same object assert entry_a_after.engine is mock_engine assert entry_a_after.last_access == 42.0 assert entry_a_after.estimated_size == original_size assert pool._current_model_memory == original_size assert pool.model_count == 2 def test_rediscover_removes_stale_unloaded(self, small_mock_model_dir): """Test that unloaded models missing from disk are removed.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) assert "model-a" in pool.get_model_ids() assert "model-b" in pool.get_model_ids() # Delete model-b from disk import shutil shutil.rmtree(small_mock_model_dir / "model-b") # Re-discover pool.discover_models(str(small_mock_model_dir)) assert "model-a" in pool.get_model_ids() assert "model-b" not in pool.get_model_ids() assert pool.model_count == 1 def test_rediscover_keeps_loaded_model_missing_from_disk( self, small_mock_model_dir ): """Test that loaded models are kept even if missing from disk.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Simulate loaded model-b entry_b = pool.get_entry("model-b") entry_b.engine = MagicMock() entry_b.last_access = 99.0 # Delete model-b from disk import shutil shutil.rmtree(small_mock_model_dir / "model-b") # Re-discover pool.discover_models(str(small_mock_model_dir)) # model-b should still be in entries (loaded in memory) assert "model-b" in pool.get_model_ids() assert pool.get_entry("model-b").engine is not None assert pool.get_entry("model-b").last_access == 99.0 def test_rediscover_updates_pinned_flag_on_loaded(self, small_mock_model_dir): """Test that pinned flag is updated on loaded models during re-discovery.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Simulate loaded model-a (not pinned) entry_a = pool.get_entry("model-a") entry_a.engine = MagicMock() assert entry_a.is_pinned is False # Re-discover with model-a pinned pool.discover_models(str(small_mock_model_dir), pinned_models=["model-a"]) # Pinned flag should be updated, engine preserved assert pool.get_entry("model-a").is_pinned is True assert pool.get_entry("model-a").engine is not None # Re-discover without pinning pool.discover_models(str(small_mock_model_dir), pinned_models=[]) assert pool.get_entry("model-a").is_pinned is False assert pool.get_entry("model-a").engine is not None class TestEnginePoolErrors: """Tests for EnginePool error handling.""" def test_model_not_found_error(self, small_mock_model_dir): """Test ModelNotFoundError when model doesn't exist.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) with pytest.raises(ModelNotFoundError) as exc_info: asyncio.run(pool.get_engine("nonexistent")) assert exc_info.value.model_id == "nonexistent" assert "model-a" in exc_info.value.available_models def test_model_too_large_error(self, small_mock_model_dir): """Test ModelTooLargeError when model exceeds memory limit.""" # Set very small memory limit pool = _make_pool(ceiling=100) # 100 bytes pool.discover_models(str(small_mock_model_dir)) with pytest.raises(ModelTooLargeError) as exc_info: asyncio.run(pool.get_engine("model-a")) assert exc_info.value.model_id == "model-a" assert exc_info.value.ceiling == 100 def test_missing_model_path_removes_unloaded_entry(self, small_mock_model_dir): """A deleted model directory is removed and reported as not found.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) shutil.rmtree(small_mock_model_dir / "model-a") with pytest.raises(ModelNotFoundError) as exc_info: asyncio.run(pool.get_engine("model-a")) assert exc_info.value.model_id == "model-a" assert pool.get_entry("model-a") is None class TestEnginePoolStatus: """Tests for EnginePool status reporting.""" def test_get_status(self, small_mock_model_dir): """Test get_status returns correct information.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir), pinned_models=["model-a"]) status = pool.get_status() assert status["final_ceiling"] == 10 * 1024**3 assert status["current_model_memory"] == 0 assert status["model_count"] == 2 assert status["loaded_count"] == 0 assert len(status["models"]) == 2 # Check model details model_ids = {m["id"] for m in status["models"]} assert model_ids == {"model-a", "model-b"} # Check pinned status model_a_status = next(m for m in status["models"] if m["id"] == "model-a") assert model_a_status["pinned"] is True assert model_a_status["loaded"] is False def test_get_model_ids(self, small_mock_model_dir): """Test get_model_ids returns all model IDs.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) ids = pool.get_model_ids() assert set(ids) == {"model-a", "model-b"} def test_get_loaded_model_ids_empty(self, small_mock_model_dir): """Test get_loaded_model_ids when no models loaded.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) assert pool.get_loaded_model_ids() == [] class TestEngineEntry: """Tests for EngineEntry dataclass.""" def test_entry_defaults(self): """Test EngineEntry default values.""" entry = EngineEntry( model_id="test", model_path="/path/to/model", model_type="llm", engine_type="batched", estimated_size=1000, ) assert entry.engine is None assert entry.last_access == 0.0 assert entry.is_loading is False assert entry.is_pinned is False def test_entry_with_values(self): """Test EngineEntry with custom values.""" entry = EngineEntry( model_id="test", model_path="/path/to/model", model_type="embedding", engine_type="embedding", estimated_size=2000, is_pinned=True, ) assert entry.model_type == "embedding" assert entry.engine_type == "embedding" assert entry.is_pinned is True class TestApplySettingsOverrides: """Tests for apply_settings_overrides method.""" def test_override_changes_model_type(self, small_mock_model_dir): """Test that model_type_override changes entry model_type and engine_type.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Verify auto-detected types assert pool.get_entry("model-a").model_type == "llm" assert pool.get_entry("model-a").engine_type == "batched" # Mock settings manager from omlx.model_settings import ModelSettings settings_manager = MagicMock() settings_manager.get_settings.side_effect = lambda mid: ( ModelSettings(model_type_override="vlm") if mid == "model-a" else ModelSettings() ) pool.apply_settings_overrides(settings_manager) # model-a should be overridden to vlm assert pool.get_entry("model-a").model_type == "vlm" assert pool.get_entry("model-a").engine_type == "vlm" # model-b should remain unchanged (no override) assert pool.get_entry("model-b").model_type == "llm" assert pool.get_entry("model-b").engine_type == "batched" def test_no_override_leaves_entry_unchanged(self, small_mock_model_dir): """Test that None override doesn't change entry types.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) from omlx.model_settings import ModelSettings settings_manager = MagicMock() settings_manager.get_settings.return_value = ModelSettings() pool.apply_settings_overrides(settings_manager) assert pool.get_entry("model-a").model_type == "llm" assert pool.get_entry("model-a").engine_type == "batched" class TestVLMFallback: """Tests for VLM-to-LLM fallback during engine loading.""" @pytest.mark.asyncio async def test_vlm_fallback_to_llm_on_start_failure(self, small_mock_model_dir): """Test that VLM loading failure falls back to LLM BatchedEngine.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Force model-a to be VLM type entry = pool.get_entry("model-a") entry.model_type = "vlm" entry.engine_type = "vlm" # VLM engine that fails on start mock_vlm_engine = MagicMock() mock_vlm_engine.start = AsyncMock( side_effect=Exception("Missing vision_tower parameters") ) mock_vlm_engine.stop = AsyncMock() # Batched engine that succeeds mock_batched_engine = MagicMock() mock_batched_engine.start = AsyncMock() with ( patch("omlx.engine_pool.VLMBatchedEngine", return_value=mock_vlm_engine), patch("omlx.engine_pool.BatchedEngine", return_value=mock_batched_engine), ): await pool._load_engine("model-a") # Should have fallen back to LLM assert entry.model_type == "llm" assert entry.engine_type == "batched" assert entry.engine is mock_batched_engine @pytest.mark.asyncio async def test_non_vlm_failure_still_raises(self, small_mock_model_dir): """Test that non-VLM engine failures propagate normally.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool.get_entry("model-a") assert entry.engine_type == "batched" # LLM, not VLM mock_engine = MagicMock() mock_engine.start = AsyncMock(side_effect=Exception("Load failed")) with ( patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine), pytest.raises(Exception, match="Load failed"), ): await pool._load_engine("model-a") @pytest.mark.asyncio async def test_force_lm_fallback_to_vlm_on_start_failure( self, small_mock_model_dir ): """Test that force_lm failure for VLM model falls back to VLMBatchedEngine.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Force model-a to be VLM type entry = pool.get_entry("model-a") entry.model_type = "vlm" entry.engine_type = "vlm" # BatchedEngine (force_lm) fails on start mock_batched_engine = MagicMock() mock_batched_engine.start = AsyncMock( side_effect=TypeError( "ModelArgs.__init__() missing 1 required positional argument: " "'tie_word_embeddings'" ) ) mock_batched_engine.stop = AsyncMock() # VLMBatchedEngine succeeds mock_vlm_engine = MagicMock() mock_vlm_engine.start = AsyncMock() with ( patch("omlx.engine_pool.BatchedEngine", return_value=mock_batched_engine), patch("omlx.engine_pool.VLMBatchedEngine", return_value=mock_vlm_engine), ): await pool._load_engine("model-a", force_lm=True) # Should have fallen back to VLM assert entry.model_type == "vlm" assert entry.engine_type == "vlm" assert entry.engine is mock_vlm_engine @pytest.mark.asyncio async def test_force_lm_no_fallback_for_non_vlm(self, small_mock_model_dir): """Test that force_lm failure for non-VLM model still raises.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool.get_entry("model-a") assert entry.engine_type == "batched" # LLM, not VLM mock_engine = MagicMock() mock_engine.start = AsyncMock(side_effect=Exception("Load failed")) with ( patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine), pytest.raises(Exception, match="Load failed"), ): await pool._load_engine("model-a", force_lm=True) @pytest.mark.asyncio async def test_vlm_fallback_to_llm_both_fail_surfaces_both_errors( self, small_mock_model_dir ): """When VLM start fails AND the LLM fallback also fails, the raised unavailable error should embed both messages and preserve the original fallback RuntimeError in the cause chain (PR #1283).""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool.get_entry("model-a") entry.model_type = "vlm" entry.engine_type = "vlm" mock_vlm_engine = MagicMock() mock_vlm_engine.start = AsyncMock( side_effect=Exception("Missing vision_tower parameters") ) mock_vlm_engine.stop = AsyncMock() mock_batched_engine = MagicMock() mock_batched_engine.start = AsyncMock( side_effect=Exception("Model type lfm2_vl not supported") ) with ( patch("omlx.engine_pool.VLMBatchedEngine", return_value=mock_vlm_engine), patch("omlx.engine_pool.BatchedEngine", return_value=mock_batched_engine), pytest.raises(ModelUnavailableError) as excinfo, ): await pool._load_engine("model-a") msg = str(excinfo.value) assert "VLM load failed" in msg assert "Missing vision_tower parameters" in msg assert "LLM fallback also failed" in msg assert "Model type lfm2_vl not supported" in msg assert excinfo.value.__cause__ is not None assert isinstance(excinfo.value.__cause__, RuntimeError) assert excinfo.value.__cause__.__cause__ is not None assert "Missing vision_tower parameters" in str( excinfo.value.__cause__.__cause__ ) @pytest.mark.asyncio async def test_force_lm_fallback_to_vlm_both_fail_surfaces_both_errors( self, small_mock_model_dir ): """force_lm path: LM start fails AND VLM fallback also fails. Both error messages should land in the unavailable error, with the LM error preserved in the cause chain (PR #1283).""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool.get_entry("model-a") entry.model_type = "vlm" entry.engine_type = "vlm" mock_batched_engine = MagicMock() mock_batched_engine.start = AsyncMock( side_effect=TypeError( "ModelArgs.__init__() missing 1 required positional argument: " "'tie_word_embeddings'" ) ) mock_batched_engine.stop = AsyncMock() mock_vlm_engine = MagicMock() mock_vlm_engine.start = AsyncMock( side_effect=Exception("vision encoder weights missing") ) with ( patch("omlx.engine_pool.BatchedEngine", return_value=mock_batched_engine), patch("omlx.engine_pool.VLMBatchedEngine", return_value=mock_vlm_engine), pytest.raises(ModelUnavailableError) as excinfo, ): await pool._load_engine("model-a", force_lm=True) msg = str(excinfo.value) assert "LM load failed" in msg assert "force_lm=True" in msg assert "tie_word_embeddings" in msg assert "VLM fallback also failed" in msg assert "vision encoder weights missing" in msg assert isinstance(excinfo.value.__cause__, RuntimeError) assert isinstance(excinfo.value.__cause__.__cause__, TypeError) class TestEnginePoolLRU: """Tests for LRU eviction logic.""" @pytest.fixture def pool_with_entries(self, small_mock_model_dir): """Create pool with mock entries for LRU testing.""" pool = _make_pool(ceiling=5000) # 5KB limit pool.discover_models(str(small_mock_model_dir)) return pool def test_find_lru_victim_no_loaded(self, pool_with_entries): """Test finding LRU victim when no models loaded.""" victim = pool_with_entries._find_lru_victim() assert victim is None def test_find_lru_victim_all_pinned(self, pool_with_entries): """Test finding LRU victim when all loaded models are pinned.""" # Mark both as pinned pool_with_entries._entries["model-a"].is_pinned = True pool_with_entries._entries["model-b"].is_pinned = True # Simulate loaded state pool_with_entries._entries["model-a"].engine = MagicMock() pool_with_entries._entries["model-a"].last_access = 100 victim = pool_with_entries._find_lru_victim() assert victim is None def test_find_lru_victim_oldest_first(self, pool_with_entries): """Test that oldest (lowest last_access) is selected.""" # Simulate loaded state with different access times mock_a = MagicMock() mock_a.has_active_requests.return_value = False pool_with_entries._entries["model-a"].engine = mock_a pool_with_entries._entries["model-a"].last_access = 100 # Older mock_b = MagicMock() mock_b.has_active_requests.return_value = False pool_with_entries._entries["model-b"].engine = mock_b pool_with_entries._entries["model-b"].last_access = 200 # Newer victim = pool_with_entries._find_lru_victim() assert victim == "model-a" def test_pinned_model_skipped_for_eviction(self, pool_with_entries): """Test that pinned models are skipped during eviction.""" # model-a is pinned and older pool_with_entries._entries["model-a"].is_pinned = True mock_a = MagicMock() mock_a.has_active_requests.return_value = False pool_with_entries._entries["model-a"].engine = mock_a pool_with_entries._entries["model-a"].last_access = 50 # model-b is not pinned and newer mock_b = MagicMock() mock_b.has_active_requests.return_value = False pool_with_entries._entries["model-b"].engine = mock_b pool_with_entries._entries["model-b"].last_access = 200 victim = pool_with_entries._find_lru_victim() # model-a is skipped (pinned), model-b is selected assert victim == "model-b" def test_find_lru_victim_skips_active_requests(self, pool_with_entries): """Test that models with active requests are skipped during eviction.""" # model-a has active requests mock_engine_a = MagicMock() mock_engine_a.has_active_requests.return_value = True pool_with_entries._entries["model-a"].engine = mock_engine_a pool_with_entries._entries["model-a"].last_access = 50 # Older # model-b has no active requests mock_engine_b = MagicMock() mock_engine_b.has_active_requests.return_value = False pool_with_entries._entries["model-b"].engine = mock_engine_b pool_with_entries._entries["model-b"].last_access = 200 # Newer victim = pool_with_entries._find_lru_victim() # model-a skipped (active requests), model-b selected assert victim == "model-b" def test_find_lru_victim_all_active(self, pool_with_entries): """Test that None is returned when all models have active requests.""" for mid in ("model-a", "model-b"): mock_engine = MagicMock() mock_engine.has_active_requests.return_value = True pool_with_entries._entries[mid].engine = mock_engine pool_with_entries._entries[mid].last_access = 100 victim = pool_with_entries._find_lru_victim() assert victim is None def test_find_lru_victim_no_has_active_requests(self, pool_with_entries): """Test graceful handling when engine lacks has_active_requests.""" mock_engine = MagicMock(spec=[]) # No has_active_requests pool_with_entries._entries["model-a"].engine = mock_engine pool_with_entries._entries["model-a"].last_access = 100 victim = pool_with_entries._find_lru_victim() assert victim == "model-a" class TestEnginePoolDFlashIsolation: """Tests for DFlash process-global runtime isolation.""" class DFlashEngine: def __init__(self, *, active: bool = False): self.active = active def has_active_requests(self): return self.active class OtherEngine: def has_active_requests(self): return False @staticmethod def _entry(model_id: str, engine) -> EngineEntry: return EngineEntry( model_id=model_id, model_path=f"/models/{model_id}", model_type="llm", engine_type="batched", estimated_size=1024, engine=engine, ) @pytest.mark.asyncio async def test_unload_other_dflash_engines_unloads_idle_dflash_only(self): pool = EnginePool() pool._entries["old-dflash"] = self._entry("old-dflash", self.DFlashEngine()) pool._entries["other"] = self._entry("other", self.OtherEngine()) pool._entries["new-dflash"] = self._entry("new-dflash", None) unloaded = [] async def fake_unload(model_id): unloaded.append(model_id) pool._entries[model_id].engine = None pool._unload_engine = fake_unload await pool._unload_other_dflash_engines("new-dflash") assert unloaded == ["old-dflash"] assert pool._entries["other"].engine is not None @pytest.mark.asyncio async def test_unload_other_dflash_engines_blocks_active_dflash(self): pool = EnginePool() pool._entries["active-dflash"] = self._entry( "active-dflash", self.DFlashEngine(active=True) ) pool._entries["new-dflash"] = self._entry("new-dflash", None) pool._unload_engine = AsyncMock() with pytest.raises(RuntimeError, match="active-dflash"): await pool._unload_other_dflash_engines("new-dflash") pool._unload_engine.assert_not_awaited() class TestEnginePoolAsync: """Async tests for EnginePool (mocked).""" @pytest.fixture def pool_with_mock_engines(self, small_mock_model_dir): """Create pool with mocked engine loading.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) return pool @pytest.mark.asyncio async def test_get_engine_loads_model(self, pool_with_mock_engines): """Test that get_engine loads the model.""" pool = pool_with_mock_engines # Mock the engine loading mock_engine = MagicMock() mock_engine.start = AsyncMock() mock_engine.stop = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine): engine = await pool.get_engine("model-a") assert engine == mock_engine mock_engine.start.assert_called_once() assert pool.loaded_model_count == 1 assert pool.current_model_memory > 0 @pytest.mark.asyncio async def test_load_failure_is_cached_until_discovery_refresh( self, pool_with_mock_engines, small_mock_model_dir ): """A failed model load is not retried until the model list is refreshed.""" pool = pool_with_mock_engines mock_engine = MagicMock() mock_engine.start = AsyncMock(side_effect=RuntimeError("broken weights")) mock_engine.stop = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine): with pytest.raises(ModelUnavailableError): await pool.get_engine("model-a") with pytest.raises(ModelUnavailableError): await pool.get_engine("model-a") assert mock_engine.start.await_count == 1 entry = pool.get_entry("model-a") assert entry is not None assert entry.load_failed is True assert "broken weights" in (entry.load_failure_message or "") pool.discover_models(str(small_mock_model_dir)) refreshed = pool.get_entry("model-a") assert refreshed is not None assert refreshed.load_failed is False @pytest.mark.asyncio async def test_runtime_settings_signature_reload(self, pool_with_mock_engines): """A profile runtime variant with engine fields reloads the base engine.""" from omlx.model_settings import ModelSettings pool = pool_with_mock_engines pool._settings_manager = MagicMock() pool._settings_manager.get_settings.return_value = ModelSettings( mtp_enabled=False ) base_engine = MagicMock() base_engine.start = AsyncMock() base_engine.stop = AsyncMock() profile_engine = MagicMock() profile_engine.start = AsyncMock() profile_engine.stop = AsyncMock() with patch( "omlx.engine_pool.BatchedEngine", side_effect=[base_engine, profile_engine], ): first = await pool.get_engine("model-a") second = await pool.get_engine( "model-a", runtime_settings=ModelSettings(mtp_enabled=True), ) assert first is base_engine assert second is profile_engine base_engine.stop.assert_awaited_once() assert pool.get_entry("model-a").runtime_settings_signature == ( pool._engine_runtime_signature( "model-a", ModelSettings(mtp_enabled=True), ) ) @pytest.mark.asyncio async def test_runtime_settings_reload_rejected_while_leased( self, pool_with_mock_engines ): """A profile variant switch must not unload an engine held by a request.""" from omlx.model_settings import ModelSettings pool = pool_with_mock_engines pool._settings_manager = MagicMock() pool._settings_manager.get_settings.return_value = ModelSettings( mtp_enabled=False ) base_engine = MagicMock() base_engine.start = AsyncMock() base_engine.stop = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=base_engine): first = await pool.get_engine("model-a") pool.get_entry("model-a").in_use = 1 with pytest.raises(ModelBusyError, match="runtime settings variant"): await pool.get_engine( "model-a", runtime_settings=ModelSettings(mtp_enabled=True), ) assert first is base_engine assert pool.get_entry("model-a").engine is base_engine base_engine.stop.assert_not_awaited() @pytest.mark.asyncio async def test_runtime_sampling_only_profile_reuses_loaded_engine( self, pool_with_mock_engines ): from omlx.model_settings import ModelSettings pool = pool_with_mock_engines pool._settings_manager = MagicMock() pool._settings_manager.get_settings.return_value = ModelSettings( mtp_enabled=False ) base_engine = MagicMock() base_engine.start = AsyncMock() base_engine.stop = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=base_engine): first = await pool.get_engine("model-a") second = await pool.get_engine( "model-a", runtime_settings=ModelSettings(temperature=0.9, mtp_enabled=False), ) assert first is base_engine assert second is base_engine base_engine.stop.assert_not_awaited() def test_runtime_signature_ignores_request_only_profile_fields( self, pool_with_mock_engines ): from omlx.model_settings import ModelSettings pool = pool_with_mock_engines pure = ModelSettings( enable_thinking=False, dflash_enabled=False, dflash_draft_model="/stale/draft", vlm_mtp_enabled=False, vlm_mtp_draft_model="/stale/assistant", ) pure_think = ModelSettings( enable_thinking=True, dflash_enabled=False, dflash_draft_model=None, vlm_mtp_enabled=False, vlm_mtp_draft_model=None, ) dflash = ModelSettings( dflash_enabled=True, dflash_draft_model="/draft", ) vlm_mtp = ModelSettings( vlm_mtp_enabled=True, vlm_mtp_draft_model="/assistant", ) pure_signature = pool._engine_runtime_signature("model-a", pure) pure_think_signature = pool._engine_runtime_signature("model-a", pure_think) assert pure_signature == pure_think_signature assert pool._engine_runtime_signature("model-a", dflash) != pure_signature assert pool._engine_runtime_signature("model-a", vlm_mtp) != pure_signature @pytest.mark.asyncio async def test_base_request_reloads_after_profile_variant( self, pool_with_mock_engines ): from omlx.model_settings import ModelSettings pool = pool_with_mock_engines pool._settings_manager = MagicMock() pool._settings_manager.get_settings.return_value = ModelSettings( mtp_enabled=False ) profile_engine = MagicMock() profile_engine.start = AsyncMock() profile_engine.stop = AsyncMock() base_engine = MagicMock() base_engine.start = AsyncMock() base_engine.stop = AsyncMock() with patch( "omlx.engine_pool.BatchedEngine", side_effect=[profile_engine, base_engine], ): first = await pool.get_engine( "model-a", runtime_settings=ModelSettings(mtp_enabled=True), ) second = await pool.get_engine("model-a") assert first is profile_engine assert second is base_engine profile_engine.stop.assert_awaited_once() @pytest.mark.asyncio async def test_embedding_engine_receives_scheduler_config(self, tmp_path): """Embedding chunk sizing should come from the shared scheduler config.""" from omlx.scheduler import SchedulerConfig model_path = tmp_path / "embed-model" model_path.mkdir() (model_path / "config.json").write_text(json.dumps({"model_type": "bert"})) scheduler_config = SchedulerConfig( completion_batch_size=6, embedding_batch_size=4, ) pool = _make_pool( ceiling=10 * 1024**3, scheduler_config=scheduler_config, ) pool._entries["embed-model"] = EngineEntry( model_id="embed-model", model_path=str(model_path), model_type="embedding", engine_type="embedding", estimated_size=1024, ) mock_engine = MagicMock() mock_engine.start = AsyncMock() with patch( "omlx.engine_pool.EmbeddingEngine", return_value=mock_engine, ) as MockEmbeddingEngine: engine = await pool.get_engine("embed-model") assert engine is mock_engine MockEmbeddingEngine.assert_called_once_with( model_name=str(model_path), trust_remote_code=False, scheduler_config=scheduler_config, ) @pytest.mark.asyncio async def test_embedding_engine_receives_fallback_scheduler_config(self, tmp_path): """A bare EnginePool should pass its fallback scheduler config consistently.""" model_path = tmp_path / "embed-model" model_path.mkdir() (model_path / "config.json").write_text(json.dumps({"model_type": "bert"})) pool = _make_pool(ceiling=10 * 1024**3) pool._entries["embed-model"] = EngineEntry( model_id="embed-model", model_path=str(model_path), model_type="embedding", engine_type="embedding", estimated_size=1024, ) mock_engine = MagicMock() mock_engine.start = AsyncMock() with patch( "omlx.engine_pool.EmbeddingEngine", return_value=mock_engine, ) as MockEmbeddingEngine: engine = await pool.get_engine("embed-model") assert engine is mock_engine MockEmbeddingEngine.assert_called_once_with( model_name=str(model_path), trust_remote_code=False, scheduler_config=pool._scheduler_config, ) @pytest.mark.asyncio async def test_apply_embedding_batch_size_updates_loaded_embedding_engines(self): """Runtime setting changes should update pool config and loaded embedding engines.""" from omlx.engine.embedding import EmbeddingEngine from omlx.scheduler import SchedulerConfig engine = EmbeddingEngine("embed-model", batch_size=8) pool = _make_pool( ceiling=10 * 1024**3, scheduler_config=SchedulerConfig(embedding_batch_size=8), ) pool._entries["embed-model"] = EngineEntry( model_id="embed-model", model_path="/tmp/embed-model", model_type="embedding", engine_type="embedding", estimated_size=1024, engine=engine, ) await pool.apply_embedding_batch_size(5) assert pool._scheduler_config.embedding_batch_size == 5 assert engine.get_stats()["batch_size"] == 5 @pytest.mark.asyncio async def test_get_engine_returns_cached(self, pool_with_mock_engines): """Test that get_engine returns cached engine on second call.""" pool = pool_with_mock_engines mock_engine = MagicMock() mock_engine.start = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine): engine1 = await pool.get_engine("model-a") engine2 = await pool.get_engine("model-a") assert engine1 is engine2 # start() should only be called once assert mock_engine.start.call_count == 1 @pytest.mark.asyncio async def test_unload_engine(self, pool_with_mock_engines): """Test unloading an engine.""" pool = pool_with_mock_engines mock_engine = MagicMock() mock_engine.start = AsyncMock() mock_engine.stop = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine): await pool.get_engine("model-a") initial_memory = pool.current_model_memory await pool._unload_engine("model-a") mock_engine.stop.assert_called_once() assert pool.current_model_memory < initial_memory assert pool._entries["model-a"].engine is None @pytest.mark.asyncio async def test_shutdown_unloads_all(self, pool_with_mock_engines): """Test that shutdown unloads all engines.""" pool = pool_with_mock_engines mock_engine_a = MagicMock() mock_engine_a.start = AsyncMock() mock_engine_a.stop = AsyncMock() mock_engine_b = MagicMock() mock_engine_b.start = AsyncMock() mock_engine_b.stop = AsyncMock() engines = [mock_engine_a, mock_engine_b] engine_idx = [0] def create_engine(*args, **kwargs): engine = engines[engine_idx[0]] engine_idx[0] += 1 return engine with patch("omlx.engine_pool.BatchedEngine", side_effect=create_engine): await pool.get_engine("model-a") await pool.get_engine("model-b") await pool.shutdown() mock_engine_a.stop.assert_called_once() mock_engine_b.stop.assert_called_once() assert pool.loaded_model_count == 0 class TestEnginePoolEviction: """Tests for memory-based eviction.""" @pytest.fixture def tight_memory_pool(self, small_mock_model_dir, monkeypatch): """Create pool with tight memory limit. Admission compares `phys_footprint + model_size` against the ceiling. For these byte-sized synthetic ceilings we proxy phys_footprint to the pool's tracked model weight sum so loading a second model triggers eviction the same way the original max_model_memory logic did. """ # Model-a is ~1KB (1024 bytes + overhead ~1.1KB) # Model-b is ~2KB (2048 bytes + overhead ~2.1KB) # Set limit to allow each model individually but not both together pool = _make_pool(ceiling=2500) # Allows each but not both pool.discover_models(str(small_mock_model_dir)) monkeypatch.setattr( "omlx.engine_pool.get_phys_footprint", lambda: pool._current_model_memory, ) monkeypatch.setattr("omlx.engine_pool.mx.get_active_memory", lambda: 0) return pool @pytest.mark.asyncio async def test_eviction_before_load(self, tight_memory_pool): """Test that eviction happens before loading new model.""" pool = tight_memory_pool mock_engine_a = MagicMock() mock_engine_a.start = AsyncMock() mock_engine_a.stop = AsyncMock() mock_engine_a.has_active_requests.return_value = False mock_engine_b = MagicMock() mock_engine_b.start = AsyncMock() mock_engine_b.has_active_requests.return_value = False call_count = [0] def create_engine(*args, **kwargs): call_count[0] += 1 if "model-a" in str(kwargs.get("model_name", args[0] if args else "")): return mock_engine_a return mock_engine_b with patch("omlx.engine_pool.BatchedEngine", side_effect=create_engine): # Load model-a first await pool.get_engine("model-a") assert pool.loaded_model_count == 1 # Load model-b - should evict model-a first await pool.get_engine("model-b") # model-a should have been unloaded mock_engine_a.stop.assert_called_once() assert pool._entries["model-a"].engine is None assert pool._entries["model-b"].engine is not None @pytest.fixture def undercounting_memory_pool(self, small_mock_model_dir, monkeypatch): """Pool where live Metal memory under-reports the committed footprint. The #1623 condition: after a model settles/idles, both `mx.get_active_memory()` and `get_phys_footprint()` can read well below the model's true resident size, while the tracked accumulator (`_current_model_memory`) still reflects it. Unlike `tight_memory_pool`, this fixture does NOT proxy phys_footprint to the accumulator — it leaves live memory at 0 so admission has to consult the accumulator itself. """ pool = _make_pool(ceiling=2500) # each model fits alone, not both pool.discover_models(str(small_mock_model_dir)) monkeypatch.setattr("omlx.engine_pool.get_phys_footprint", lambda: 0) monkeypatch.setattr("omlx.engine_pool.mx.get_active_memory", lambda: 0) return pool @pytest.mark.asyncio async def test_eviction_when_live_memory_undercounts( self, undercounting_memory_pool ): """#1623: a second model must still evict the first when live memory under-reports but the tracked accumulator shows the pair over-commits. Without consulting `_current_model_memory`, admission sees `current = max(0, 0) = 0`, projects `0 + model-b` under the ceiling, and loads model-b alongside model-a — over-committing past the ceiling. """ pool = undercounting_memory_pool mock_engine_a = MagicMock() mock_engine_a.start = AsyncMock() mock_engine_a.stop = AsyncMock() mock_engine_a.has_active_requests.return_value = False mock_engine_b = MagicMock() mock_engine_b.start = AsyncMock() mock_engine_b.has_active_requests.return_value = False def create_engine(*args, **kwargs): if "model-a" in str(kwargs.get("model_name", args[0] if args else "")): return mock_engine_a return mock_engine_b with patch("omlx.engine_pool.BatchedEngine", side_effect=create_engine): await pool.get_engine("model-a") assert pool.loaded_model_count == 1 # Load model-b: the pair exceeds the ceiling, so model-a must be # evicted first even though live memory reads 0. await pool.get_engine("model-b") mock_engine_a.stop.assert_called_once() assert pool._entries["model-a"].engine is None assert pool._entries["model-b"].engine is not None assert pool.loaded_model_count == 1 @pytest.fixture def residue_memory_pool(self, small_mock_model_dir, monkeypatch): """Pool where phys_footprint lags high after an eviction. Reproduces the false-507-on-swap bug: the macOS phys_footprint ledger still counts reclaimable residue (dirty file-backed / IOAccelerator pages from a just-evicted model) that the kernel has not yet reaped, while both `mx.get_active_memory()` and the tracked accumulator (`_current_model_memory`) have already dropped. We model that as a high-water mark that never falls on its own. """ pool = _make_pool(ceiling=2500) # each model fits alone, not both pool.discover_models(str(small_mock_model_dir)) hwm = {"v": 0} def lagging_footprint(): hwm["v"] = max(hwm["v"], pool._current_model_memory) return hwm["v"] monkeypatch.setattr( "omlx.engine_pool.get_phys_footprint", lagging_footprint ) monkeypatch.setattr("omlx.engine_pool.mx.get_active_memory", lambda: 0) return pool @pytest.mark.asyncio async def test_swap_admits_when_footprint_lags_after_eviction( self, residue_memory_pool ): """A model that fits cleanly must load after eviction even when phys_footprint still reflects the evicted model's reclaimable residue. Without the committed-baseline recheck, admission sees `current = max(0, lagging_footprint, 0)` still pinned at the evicted model's size, projects past the ceiling with nothing left to evict, and wrongly raises InsufficientMemoryError (the 507 on swap). """ pool = residue_memory_pool mock_engine_a = MagicMock() mock_engine_a.start = AsyncMock() mock_engine_a.stop = AsyncMock() mock_engine_a.has_active_requests.return_value = False mock_engine_b = MagicMock() mock_engine_b.start = AsyncMock() mock_engine_b.has_active_requests.return_value = False def create_engine(*args, **kwargs): if "model-a" in str(kwargs.get("model_name", args[0] if args else "")): return mock_engine_a return mock_engine_b with patch("omlx.engine_pool.BatchedEngine", side_effect=create_engine): await pool.get_engine("model-a") assert pool.loaded_model_count == 1 # Swap to model-b: model-a is evicted, but phys_footprint still # reports model-a's residue. model-b fits alone, so it must load. await pool.get_engine("model-b") mock_engine_a.stop.assert_called_once() assert pool._entries["model-a"].engine is None assert pool._entries["model-b"].engine is not None assert pool.loaded_model_count == 1 @pytest.mark.asyncio async def test_high_footprint_without_eviction_still_blocks_first_load( self, small_mock_model_dir ): """Do not discount phys_footprint unless this admission evicted a model. A high footprint with no evicted victim may be unrelated process pressure, not reclaimable model residue. The committed-baseline retry must not let the first model load past the process memory ceiling. """ pool = _make_pool(ceiling=2500) pool.discover_models(str(small_mock_model_dir)) mock_engine = MagicMock() mock_engine.start = AsyncMock() with ( patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine) as cls, patch("omlx.engine_pool.get_phys_footprint", return_value=2000), patch("omlx.engine_pool.mx.get_active_memory", return_value=0), ): with pytest.raises(InsufficientMemoryError) as exc_info: await pool.get_engine("model-a") assert exc_info.value.current == 2000 cls.assert_not_called() assert pool.loaded_model_count == 0 @pytest.mark.asyncio async def test_insufficient_memory_all_pinned(self, tight_memory_pool): """Test InsufficientMemoryError when all models are pinned.""" pool = tight_memory_pool # Pin model-a pool._entries["model-a"].is_pinned = True mock_engine = MagicMock() mock_engine.start = AsyncMock() with patch("omlx.engine_pool.BatchedEngine", return_value=mock_engine): # Load pinned model-a await pool.get_engine("model-a") # Try to load model-b - should fail (can't evict pinned model-a) with pytest.raises(InsufficientMemoryError): await pool.get_engine("model-b") class TestEnginePoolPrefillEviction: """Tests for request-time idle LRU eviction before prefill throttling.""" @staticmethod def _entry(model_id: str, size: int, *, active: bool = False) -> EngineEntry: engine = MagicMock() engine.has_active_requests.return_value = active engine.scheduler = None engine._engine = None return EngineEntry( model_id=model_id, model_path=f"/models/{model_id}", model_type="llm", engine_type="batched", estimated_size=size, engine=engine, last_access=0.0, ) @pytest.mark.asyncio async def test_prefill_eviction_evicts_idle_lru_until_target(self): gb = 1024**3 pool = _make_pool(ceiling=0) pool._entries = { "idle-a": self._entry("idle-a", 20 * gb), "idle-b": self._entry("idle-b", 15 * gb), "target": self._entry("target", 25 * gb), } pool._entries["idle-a"].last_access = 1.0 pool._entries["idle-b"].last_access = 2.0 pool._entries["target"].last_access = 3.0 pool._current_model_memory = 60 * gb unloaded = [] async def fake_unload(model_id): unloaded.append(model_id) entry = pool._entries[model_id] entry.engine = None pool._current_model_memory -= entry.estimated_size pool._unload_engine = fake_unload req = PrefillEvictionRequest( request_id="req-1", model_id="target", current_bytes=60 * gb, target_cap_bytes=40 * gb, predicted_transient_bytes=10 * gb, requested_tokens=2048, reason="adaptive_prefill_throttle", ) with ( patch("omlx.engine_pool.mx.get_active_memory", return_value=0), patch("omlx.engine_pool.get_phys_footprint", return_value=0), ): evicted = await pool._evict_idle_lru_for_prefill("target", req) assert evicted is True assert unloaded == ["idle-a", "idle-b"] assert pool._entries["target"].engine is not None @pytest.mark.asyncio async def test_prefill_eviction_skips_active_pinned_loading_and_current(self): gb = 1024**3 pool = _make_pool(ceiling=0) pool._entries = { "active": self._entry("active", 20 * gb, active=True), "pinned": self._entry("pinned", 20 * gb), "loading": self._entry("loading", 20 * gb), "target": self._entry("target", 25 * gb), } pool._entries["pinned"].is_pinned = True pool._entries["loading"].is_loading = True pool._current_model_memory = 85 * gb pool._unload_engine = AsyncMock() req = PrefillEvictionRequest( request_id="req-1", model_id="target", current_bytes=85 * gb, target_cap_bytes=40 * gb, predicted_transient_bytes=10 * gb, requested_tokens=2048, reason="adaptive_prefill_throttle", ) with ( patch("omlx.engine_pool.mx.get_active_memory", return_value=0), patch("omlx.engine_pool.get_phys_footprint", return_value=0), ): evicted = await pool._evict_idle_lru_for_prefill("target", req) assert evicted is False pool._unload_engine.assert_not_awaited() class TestEnginePoolStatus: """Tests for get_status is_loading field.""" def test_get_status_includes_is_loading(self, small_mock_model_dir): """Test get_status includes is_loading field.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) status = pool.get_status() for model in status["models"]: assert "is_loading" in model assert model["is_loading"] is False class TestEnginePoolTTL: """Tests for TTL expiration checking.""" @pytest.fixture def pool_with_loaded_model(self, small_mock_model_dir): """Create pool with a mock-loaded model.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Mock-load model-a entry = pool._entries["model-a"] entry.engine = MagicMock(spec=[]) entry.engine.stop = AsyncMock() entry.engine.has_active_requests = MagicMock(return_value=False) entry.last_access = 100.0 # Old access time pool._current_model_memory = entry.estimated_size return pool @pytest.mark.asyncio async def test_ttl_expires_idle_model(self, pool_with_loaded_model): """Test that TTL unloads idle model.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): # 100s idle > 60s TTL expired = await pool.check_ttl_expirations(settings_manager) assert "model-a" in expired assert pool._entries["model-a"].engine is None @pytest.mark.asyncio async def test_ttl_skips_model_within_ttl(self, pool_with_loaded_model): """Test that TTL does not unload model within TTL.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 300 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): # 100s idle < 300s TTL expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["model-a"].engine is not None @pytest.mark.asyncio async def test_ttl_skips_no_ttl_model(self, pool_with_loaded_model): """Test that models without TTL are not unloaded.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = None settings_manager.get_settings.return_value = settings with patch("time.time", return_value=99999.0): expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["model-a"].engine is not None @pytest.mark.asyncio async def test_ttl_skips_pinned_model(self, pool_with_loaded_model): """Test that pinned models are not unloaded by TTL.""" pool = pool_with_loaded_model pool._entries["model-a"].is_pinned = True settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["model-a"].engine is not None @pytest.mark.asyncio async def test_ttl_skips_model_with_active_requests(self, pool_with_loaded_model): """Test that TTL does not unload model with active requests.""" pool = pool_with_loaded_model # Mock an engine that reports active requests mock_engine = MagicMock() mock_engine.has_active_requests.return_value = True pool._entries["model-a"].engine = mock_engine settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] # last_access should be refreshed assert pool._entries["model-a"].last_access == 200.0 @pytest.mark.asyncio async def test_ttl_skips_vlm_with_active_requests(self, pool_with_loaded_model): """Test that TTL does not unload VLM engine with active requests.""" pool = pool_with_loaded_model mock_engine = MagicMock() mock_engine.has_active_requests.return_value = True pool._entries["model-a"].engine = mock_engine settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["model-a"].last_access == 200.0 @pytest.mark.asyncio async def test_ttl_skips_non_streaming_with_active_requests( self, pool_with_loaded_model ): """Test that TTL does not unload non-streaming engine with active requests.""" pool = pool_with_loaded_model mock_engine = MagicMock() mock_engine.has_active_requests.return_value = True pool._entries["model-a"].engine = mock_engine settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["model-a"].last_access == 200.0 @pytest.mark.asyncio async def test_ttl_falls_back_to_global_idle_timeout(self, pool_with_loaded_model): """Per-model TTL None falls back to global idle timeout.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = None settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): # 100s idle > 60s global expired = await pool.check_ttl_expirations( settings_manager, global_idle_timeout_seconds=60 ) assert "model-a" in expired assert pool._entries["model-a"].engine is None @pytest.mark.asyncio async def test_ttl_global_disabled_when_none(self, pool_with_loaded_model): """Per-model None + global None keeps model loaded.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = None settings_manager.get_settings.return_value = settings with patch("time.time", return_value=99999.0): expired = await pool.check_ttl_expirations( settings_manager, global_idle_timeout_seconds=None ) assert expired == [] assert pool._entries["model-a"].engine is not None @pytest.mark.asyncio async def test_per_model_ttl_overrides_global(self, pool_with_loaded_model): """Per-model TTL wins over global idle timeout.""" pool = pool_with_loaded_model settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 300 # per-model wider than global settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): # 100s idle < 300s per-model expired = await pool.check_ttl_expirations( settings_manager, global_idle_timeout_seconds=60 ) assert expired == [] assert pool._entries["model-a"].engine is not None class TestHasActiveRequests: """Tests for has_active_requests() on engine types.""" def test_base_non_streaming_engine_active_count(self): """Test BaseNonStreamingEngine active request tracking.""" from omlx.engine.base import BaseNonStreamingEngine class DummyEngine(BaseNonStreamingEngine): @property def model_name(self): return "dummy" async def start(self): pass async def stop(self): pass def get_stats(self): return {} engine = DummyEngine() assert engine.has_active_requests() is False with engine._active_lock: engine._active_count += 1 assert engine.has_active_requests() is True with engine._active_lock: engine._active_count -= 1 assert engine.has_active_requests() is False def test_batched_engine_has_active_requests(self): """Test BatchedEngine.has_active_requests() via _output_collectors.""" from omlx.engine.batched import BatchedEngine engine = BatchedEngine.__new__(BatchedEngine) engine._engine = None assert engine.has_active_requests() is False # Simulate engine with active collectors mock_engine_core = MagicMock() mock_inner = MagicMock() mock_inner._output_collectors = {"req1": MagicMock()} mock_engine_core.engine = mock_inner engine._engine = mock_engine_core assert engine.has_active_requests() is True # Empty collectors mock_inner._output_collectors = {} assert engine.has_active_requests() is False def test_vlm_engine_has_active_requests(self): """Test VLMBatchedEngine.has_active_requests() via _output_collectors.""" from omlx.engine.vlm import VLMBatchedEngine engine = VLMBatchedEngine.__new__(VLMBatchedEngine) engine._engine = None assert engine.has_active_requests() is False mock_engine_core = MagicMock() mock_inner = MagicMock() mock_inner._output_collectors = {"req1": MagicMock()} mock_engine_core.engine = mock_inner engine._engine = mock_engine_core assert engine.has_active_requests() is True class TestResolveModelId: """Tests for resolve_model_id method.""" def test_direct_match(self, small_mock_model_dir): """Test direct model_id match returns immediately.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) result = pool.resolve_model_id("model-a", settings_manager=None) assert result == "model-a" def test_alias_match(self, small_mock_model_dir): """Test alias resolution returns real model_id.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) settings_manager = MagicMock() settings_manager.get_exposed_profile_source_model_id.return_value = None from omlx.model_settings import ModelSettings settings_manager.get_all_settings.return_value = { "model-a": ModelSettings(model_alias="gpt-4"), "model-b": ModelSettings(), } result = pool.resolve_model_id("gpt-4", settings_manager) assert result == "model-a" def test_no_match_returns_original(self, small_mock_model_dir): """Test unresolved name returns original string.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) settings_manager = MagicMock() settings_manager.get_exposed_profile_source_model_id.return_value = None from omlx.model_settings import ModelSettings settings_manager.get_all_settings.return_value = { "model-a": ModelSettings(), } result = pool.resolve_model_id("nonexistent", settings_manager) assert result == "nonexistent" def test_alias_match_no_settings_manager(self, small_mock_model_dir): """Test with None settings_manager falls back to direct match only.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) result = pool.resolve_model_id("some-alias", settings_manager=None) assert result == "some-alias" def test_provider_prefix_alias_match(self, small_mock_model_dir): """Test alias resolution with provider prefix (e.g. omlx/alias).""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) settings_manager = MagicMock() settings_manager.get_exposed_profile_source_model_id.return_value = None from omlx.model_settings import ModelSettings settings_manager.get_all_settings.return_value = { "model-a": ModelSettings(model_alias="gpt-4"), "model-b": ModelSettings(), } result = pool.resolve_model_id("omlx/gpt-4", settings_manager) assert result == "model-a" def test_provider_prefix_direct_match(self, small_mock_model_dir): """Test direct match with provider prefix (e.g. provider/model-a).""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) result = pool.resolve_model_id("provider/model-a", settings_manager=None) assert result == "model-a" def test_provider_prefix_no_match(self, small_mock_model_dir): """Test prefix strip still returns original when no match found.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) settings_manager = MagicMock() settings_manager.get_exposed_profile_source_model_id.return_value = None from omlx.model_settings import ModelSettings settings_manager.get_all_settings.return_value = { "model-a": ModelSettings(), } result = pool.resolve_model_id("omlx/nonexistent", settings_manager) assert result == "omlx/nonexistent" def test_case_insensitive_match(self, small_mock_model_dir): """Test case-insensitive fallback when exact match fails.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) result = pool.resolve_model_id("MODEL-A", settings_manager=None) assert result == "model-a" def test_case_insensitive_with_provider_prefix(self, small_mock_model_dir): """Test case-insensitive match after stripping provider prefix.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) result = pool.resolve_model_id("omlx/MODEL-B", settings_manager=None) assert result == "model-b" def test_exact_match_preferred_over_case_insensitive(self, small_mock_model_dir): """Test exact match takes priority over case-insensitive.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) # Exact match should be returned directly result = pool.resolve_model_id("model-a", settings_manager=None) assert result == "model-a" class TestMemorySettleBarrier: """Tests for memory settle barrier in _unload_engine().""" @pytest.fixture def pool_with_loaded_model(self, small_mock_model_dir): """Create pool with a mock-loaded model for settle barrier testing. Sets estimated_size to 5GB. With scaled tolerance (max(2GB, 5% of 5GB) = max(2GB, 0.25GB) = 2GB), the barrier requires at least 3GB freed. """ pool = _make_pool(ceiling=100 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool._entries["model-a"] entry.estimated_size = 5 * 1024**3 # 5GB (> 2GB tolerance) mock_engine = MagicMock() mock_engine.stop = AsyncMock() mock_engine.has_active_requests = MagicMock(return_value=False) entry.engine = mock_engine entry.last_access = 100.0 pool._current_model_memory = entry.estimated_size return pool @pytest.mark.asyncio async def test_settle_succeeds_first_round(self, pool_with_loaded_model): """Test that settle barrier passes on first round when memory is freed.""" pool = pool_with_loaded_model est_size = pool._entries["model-a"].estimated_size # 5GB initial_memory = pool._current_model_memory # Pre-unload: 10GB active. After GC: drops to 5GB (5GB freed >= 3GB needed). active_memory_values = [10 * 1024**3, 5 * 1024**3] call_idx = [0] def mock_get_active(): val = active_memory_values[min(call_idx[0], len(active_memory_values) - 1)] call_idx[0] += 1 return val with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") assert pool._entries["model-a"].engine is None assert pool._current_model_memory == initial_memory - est_size @pytest.mark.asyncio async def test_settle_takes_multiple_rounds(self, pool_with_loaded_model): """Test settle barrier succeeds after multiple rounds of GC.""" pool = pool_with_loaded_model # est_size = 5GB, tolerance = 2GB, so need >= 3GB freed # Pre-unload: 10GB. Need active <= 7GB to settle. # Round 1: 9GB (freed=1GB < 3GB), Round 2: 8GB (freed=2GB < 3GB), # Round 3: 5GB (freed=5GB >= 3GB) → settled active_memory_values = [10 * 1024**3, 9 * 1024**3, 8 * 1024**3, 5 * 1024**3] call_idx = [0] def mock_get_active(): val = active_memory_values[min(call_idx[0], len(active_memory_values) - 1)] call_idx[0] += 1 return val sleep_calls = [] async def mock_sleep(duration): sleep_calls.append(duration) with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", side_effect=mock_sleep), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") # Should have slept at least once (0.5s between rounds) assert any(d == 0.5 for d in sleep_calls) assert pool._entries["model-a"].engine is None assert pool._current_model_memory == 0 @pytest.mark.asyncio async def test_settle_timeout_triggers_emergency(self, pool_with_loaded_model): """Test emergency reclaim is triggered when settle barrier times out.""" pool = pool_with_loaded_model # est_size = 5GB, tolerance = 2GB, need >= 3GB freed # Memory stays at 10GB during all 10 settle rounds (0GB freed < 3GB) # After emergency reclaim: drops to safe level settle_calls = [0] def mock_get_active(): settle_calls[0] += 1 # 1 pre-unload + 10 settle rounds (each calls once) = 11 # After emergency: return safe level if settle_calls[0] <= 11: return 10 * 1024**3 return 0 sleep_calls = [] async def mock_sleep(duration): sleep_calls.append(duration) with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", side_effect=mock_sleep), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") # Emergency reclaim uses 1.0s sleeps (3 rounds) assert sleep_calls.count(1.0) == 3 assert pool._entries["model-a"].engine is None @pytest.mark.asyncio async def test_emergency_reclaim_failure_logs_error(self, pool_with_loaded_model): """Test error is logged when emergency reclaim fails to free enough memory.""" pool = pool_with_loaded_model # Memory never drops — stays at 10GB throughout (well above 5GB threshold) with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = MagicMock(return_value=10 * 1024**3) mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() with patch("omlx.engine_pool.logger") as mock_logger: await pool._unload_engine("model-a") # Should have logged an error about emergency reclaim failure error_calls = [str(c) for c in mock_logger.error.call_args_list] assert any("Emergency reclaim failed" in s for s in error_calls) @pytest.mark.asyncio async def test_memory_counter_decremented_after_barrier( self, pool_with_loaded_model ): """Regression test: _current_model_memory must not be decremented before the settle barrier completes.""" pool = pool_with_loaded_model est_size = pool._entries["model-a"].estimated_size # 5GB original_memory = pool._current_model_memory memory_during_settle = [] def mock_get_active(): # Record the pool's memory counter state during settle polling memory_during_settle.append(pool._current_model_memory) # Return high value for 3 rounds, then settle # Need >= 3GB freed (5GB - 2GB tolerance) if len(memory_during_settle) <= 3: return 10 * 1024**3 # 0GB freed return 5 * 1024**3 # 5GB freed >= 3GB needed with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") # During settle barrier, memory counter should NOT have been decremented for mem in memory_during_settle[:-1]: assert mem == original_memory, ( f"Memory counter was {mem} during settle, expected {original_memory}. " "Counter must not be decremented before barrier completes." ) # After barrier, it should be decremented assert pool._current_model_memory == original_memory - est_size @pytest.mark.asyncio async def test_settle_large_model_proportional_tolerance( self, small_mock_model_dir ): """Test that settle tolerance scales with model size for large models. For a 60GB model, 5% = 3GB > 2GB floor, so tolerance = 3GB. min_expected_freed = 60GB - 3GB = 57GB. """ pool = _make_pool(ceiling=200 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool._entries["model-a"] entry.estimated_size = 60 * 1024**3 # 60GB mock_engine = MagicMock() mock_engine.stop = AsyncMock() mock_engine.has_active_requests = MagicMock(return_value=False) entry.engine = mock_engine entry.last_access = 100.0 pool._current_model_memory = entry.estimated_size # Freed = 80 - 23 = 57GB. With proportional tolerance (3GB) this # settles, but would fail with the old fixed 2GB tolerance (needed 58GB). active_memory_values = [80 * 1024**3, 23 * 1024**3] call_idx = [0] def mock_get_active(): val = active_memory_values[min(call_idx[0], len(active_memory_values) - 1)] call_idx[0] += 1 return val with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") assert pool._entries["model-a"].engine is None assert pool._current_model_memory == 0 @pytest.mark.asyncio async def test_settle_small_model_uses_floor_tolerance(self, small_mock_model_dir): """Test that 2GB floor tolerance applies for small models. For a 1GB model, 5% = 0.05GB << 2GB, so tolerance = 2GB floor. min_expected_freed = max(0, 1GB - 2GB) = 0, settle is trivially true. """ pool = _make_pool(ceiling=100 * 1024**3) pool.discover_models(str(small_mock_model_dir)) entry = pool._entries["model-a"] entry.estimated_size = 1 * 1024**3 # 1GB mock_engine = MagicMock() mock_engine.stop = AsyncMock() mock_engine.has_active_requests = MagicMock(return_value=False) entry.engine = mock_engine entry.last_access = 100.0 pool._current_model_memory = entry.estimated_size # Even 0 freed should settle (min_expected_freed = 0) active_memory_values = [10 * 1024**3, 10 * 1024**3] call_idx = [0] def mock_get_active(): val = active_memory_values[min(call_idx[0], len(active_memory_values) - 1)] call_idx[0] += 1 return val with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = mock_get_active mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") assert pool._entries["model-a"].engine is None assert pool._current_model_memory == 0 @pytest.mark.asyncio async def test_settle_bails_out_under_concurrent_activity( self, pool_with_loaded_model, caplog ): """1774 regression: with another engine serving, the global freed delta is unmeasurable (it can read negative as the other engine allocates). The barrier must bail after one sample instead of burning 10 settle rounds + emergency reclaim — ~8s of gc/synchronize/ clear_cache serialized against live decode, under the pool lock. """ pool = pool_with_loaded_model # Second entry actively serving. other_engine = MagicMock() other_engine.has_active_requests = MagicMock(return_value=True) pool._entries["model-b"].engine = other_engine # Global gauge RISES during settle (concurrent prefill/KV growth), # so freed = pre_unload - active_now is negative every round. call_idx = [0] def rising_gauge(): val = (10 + call_idx[0]) * 1024**3 call_idx[0] += 1 return val sleep_calls: list[float] = [] async def record_sleep(duration, *args, **kwargs): sleep_calls.append(duration) with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", side_effect=record_sleep), caplog.at_level(logging.DEBUG, logger="omlx.engine_pool"), ): mock_mx.get_active_memory = rising_gauge mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") assert "indeterminate under concurrent activity" in caplog.text # No settle-round burn, no timeout warning, no emergency reclaim. assert sleep_calls.count(0.5) == 0 assert "Settle barrier timed out" not in caplog.text assert "Emergency reclaim" not in caplog.text # Only the initial pre-barrier release cycle touched the executor. assert mock_mx.synchronize.call_count == 1 assert mock_mx.clear_cache.call_count == 1 # The unload itself still completes and is accounted. assert pool._entries["model-a"].engine is None assert pool._current_model_memory == 0 @pytest.mark.asyncio async def test_settle_still_waits_when_pool_otherwise_idle( self, pool_with_loaded_model, caplog ): """Idle-pool behavior is unchanged (#768 protection): with no other entry serving, an unsatisfied barrier still burns its settle rounds and escalates to emergency reclaim. """ pool = pool_with_loaded_model call_idx = [0] def rising_gauge(): val = (10 + call_idx[0]) * 1024**3 call_idx[0] += 1 return val with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), caplog.at_level(logging.DEBUG, logger="omlx.engine_pool"), ): mock_mx.get_active_memory = rising_gauge mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") assert "indeterminate under concurrent activity" not in caplog.text assert "Settle barrier timed out" in caplog.text # Full barrier behavior preserved: 1 initial release cycle + 10 settle # rounds + 3 emergency-reclaim rounds on the executor. assert mock_mx.synchronize.call_count == 14 assert mock_mx.clear_cache.call_count == 14 def test_other_entries_serving_in_use_lease_counts(self, pool_with_loaded_model): """The in-use lease (acquired but not yet active) also marks the pool as serving — eviction paths already treat it as activity (#1667). """ pool = pool_with_loaded_model entry_b = pool._entries["model-b"] assert pool._other_entries_serving("model-a") is False entry_b.engine = MagicMock() entry_b.engine.has_active_requests = MagicMock(return_value=False) assert pool._other_entries_serving("model-a") is False entry_b.in_use = 1 assert pool._other_entries_serving("model-a") is True class TestEnginePoolInUseLease: """Tests for the acquire-vs-use in-use lease (#1667). The memory enforcer must never evict an engine that a request has acquired but not yet registered activity on. The lease is taken under the pool lock at acquire time and skipped by every eviction path, in addition to the existing has_active_requests() gate. """ @staticmethod def _loaded_entry(model_id: str, last_access: float = 0.0) -> EngineEntry: engine = MagicMock() engine.has_active_requests.return_value = False # looks idle engine.stop = AsyncMock() # Mirror TestEnginePoolPrefillEviction: no scheduler → prefill-idle. engine.scheduler = None engine._engine = None return EngineEntry( model_id=model_id, model_path=f"/models/{model_id}", model_type="embedding", engine_type="embedding", estimated_size=1024, engine=engine, last_access=last_access, ) def test_engine_entry_in_use_default_zero(self): """EngineEntry exposes an in_use lease counter defaulting to 0.""" entry = EngineEntry( model_id="x", model_path="/p", model_type="embedding", engine_type="embedding", estimated_size=1, ) assert entry.in_use == 0 def test_find_lru_victim_skips_in_use_engine(self): """An idle-looking engine with in_use > 0 must not be an LRU victim.""" pool = _make_pool(ceiling=0) pool._entries = { "leased": self._loaded_entry("leased", last_access=1.0), # oldest "free": self._loaded_entry("free", last_access=2.0), } pool._entries["leased"].in_use = 1 victim = pool._find_lru_victim() # Without the lease the older "leased" would be picked; lease forces "free". assert victim == "free" def test_find_lru_victim_none_when_only_in_use(self): """If the only loaded engine is leased, there is no LRU victim.""" pool = _make_pool(ceiling=0) pool._entries = {"leased": self._loaded_entry("leased", last_access=1.0)} pool._entries["leased"].in_use = 2 assert pool._find_lru_victim() is None def test_is_idle_for_prefill_eviction_skips_in_use(self): """Prefill-pressure eviction must skip a leased (in_use > 0) engine.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("leased") # Idle-looking → would be evictable, but the lease blocks it. assert pool._is_idle_for_prefill_eviction(entry) is True entry.in_use = 1 assert pool._is_idle_for_prefill_eviction(entry) is False @pytest.mark.asyncio async def test_check_ttl_expirations_skips_in_use(self): """TTL eviction must not unload a leased engine even when idle past TTL.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("leased", last_access=100.0) entry.in_use = 1 pool._entries = {"leased": entry} settings_manager = MagicMock() settings = MagicMock() settings.ttl_seconds = 60 settings_manager.get_settings.return_value = settings with patch("time.time", return_value=200.0): # 100s idle > 60s TTL expired = await pool.check_ttl_expirations(settings_manager) assert expired == [] assert pool._entries["leased"].engine is not None # in_use counts as activity → last_access refreshed assert pool._entries["leased"].last_access == 200.0 @pytest.mark.asyncio async def test_release_engine_floors_at_zero(self): """release_engine decrements under lock and never goes below 0.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("leased") entry.in_use = 1 pool._entries = {"leased": entry} await pool.release_engine("leased") assert entry.in_use == 0 # Extra release is a no-op (floor at 0), not a negative count. await pool.release_engine("leased") assert entry.in_use == 0 # Unknown model id is a harmless no-op. await pool.release_engine("nope") @pytest.mark.asyncio async def test_release_engine_unloads_pending_after_lease_drains(self): """Pending hard-pressure unload runs as soon as the lease drains.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("leased") entry.in_use = 1 entry.pending_unload_reason = "hard memory pressure" entry.abort_requested = True pool._entries = {"leased": entry} pool._unload_engine = AsyncMock() await pool.release_engine("leased") assert entry.in_use == 0 assert entry.pending_unload_reason is None assert entry.abort_requested is False pool._unload_engine.assert_awaited_once_with("leased") @pytest.mark.asyncio async def test_release_engine_keeps_pending_while_scheduler_active(self): """A drained lease is not enough if scheduler requests are still active.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("leased") entry.in_use = 1 entry.engine.has_active_requests.return_value = True entry.pending_unload_reason = "hard memory pressure" entry.abort_requested = True pool._entries = {"leased": entry} pool._unload_engine = AsyncMock() await pool.release_engine("leased") assert entry.in_use == 0 assert entry.pending_unload_reason == "hard memory pressure" assert entry.abort_requested is True pool._unload_engine.assert_not_awaited() @pytest.mark.asyncio async def test_acquire_leases_then_releases_on_success(self): """acquire() leases on enter and releases in finally on normal exit.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("model-a") pool._entries = {"model-a": entry} async with pool.acquire("model-a") as engine: assert engine is entry.engine assert entry.in_use == 1 # leased while in use assert entry.in_use == 0 # released on exit @pytest.mark.asyncio async def test_acquire_releases_lease_on_exception(self): """acquire() releases the lease even if the body raises.""" pool = _make_pool(ceiling=0) entry = self._loaded_entry("model-a") pool._entries = {"model-a": entry} with pytest.raises(RuntimeError, match="boom"): async with pool.acquire("model-a"): assert entry.in_use == 1 raise RuntimeError("boom") assert entry.in_use == 0 # no leaked lease @pytest.mark.asyncio async def test_get_engine_rejects_loaded_llm_without_tokenizer(self): """A stale/half-loaded LLM engine must not reach chat tokenization.""" from omlx.engine.batched import BatchedEngine pool = _make_pool(ceiling=0) engine = BatchedEngine.__new__(BatchedEngine) engine._tokenizer = None entry = self._loaded_entry("model-a") entry.engine = engine entry.model_type = "llm" entry.engine_type = "batched" pool._entries = {"model-a": entry} with pytest.raises(ModelLoadingError, match="usable tokenizer"): await pool.get_engine("model-a") class TestResetActivityTracking: """Tests for BaseNonStreamingEngine._reset_activity_tracking (#1595).""" @staticmethod def _engine(): from omlx.engine.base import BaseNonStreamingEngine class DummyEngine(BaseNonStreamingEngine): @property def model_name(self): return "dummy" async def start(self): pass async def stop(self): pass def get_stats(self): return {} return DummyEngine() def test_reset_clears_leaked_counter_and_activities(self): """Reset zeros a leaked _active_count and clears _activities. Mirrors the immediate-abort eviction case where stop() runs without the per-request completion callbacks, leaving a phantom busy count. """ engine = self._engine() with engine._active_lock: engine._active_count = 2 engine._activities = {"a": {}, "b": {}} assert engine.has_active_requests() is True # phantom busy engine._reset_activity_tracking() assert engine._active_count == 0 assert engine._activities == {} assert engine.has_active_requests() is False @pytest.mark.asyncio async def test_unload_engine_calls_reset_on_teardown(self): """EnginePool._unload_engine resets activity tracking after stop().""" pool = _make_pool(ceiling=0) engine = self._engine() with engine._active_lock: engine._active_count = 1 # leaked phantom busy entry = EngineEntry( model_id="model-a", model_path="/models/model-a", model_type="embedding", engine_type="embedding", estimated_size=1024, engine=engine, last_access=0.0, ) pool._entries = {"model-a": entry} pool._current_model_memory = entry.estimated_size with ( patch("omlx.engine_pool.mx") as mock_mx, patch("omlx.engine_pool.get_mlx_executor", return_value=None), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_mx.get_active_memory = MagicMock(return_value=0) mock_mx.synchronize = MagicMock() mock_mx.clear_cache = MagicMock() await pool._unload_engine("model-a") # Teardown reset the leaked counter (getattr-guarded reset called). assert engine._active_count == 0 assert engine.has_active_requests() is False class TestFailedLoadReclaim: """Failed model loads schedule a deferred memory reclaim. Weights loaded before a load failure can remain reachable only through the propagating exception's traceback frames, so synchronous gc/clear_cache at raise time cannot free them. The pool schedules a background task (_schedule_failed_load_reclaim) that retries reclaim after the exception has been handled and dropped. """ async def test_reclaim_settles_when_memory_returns(self): pool = _make_pool() gauge = {"active": 80 * 1024**3} def _fake_active(): return gauge["active"] def _drop_on_clear(): # Simulate the leaked buffers becoming collectable after the # first gc + clear_cache round. gauge["active"] = 1 * 1024**3 import concurrent.futures executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) try: with ( patch( "omlx.engine_pool.mx.get_active_memory", side_effect=_fake_active, ), patch("omlx.engine_pool.mx.synchronize"), patch( "omlx.engine_pool.mx.clear_cache", side_effect=_drop_on_clear, ), patch("omlx.engine_pool.get_phys_footprint", return_value=0), patch( "omlx.engine_pool.get_mlx_executor", return_value=executor ), patch( "omlx.engine_pool.asyncio.sleep", new=AsyncMock(return_value=None), ), ): pool._schedule_failed_load_reclaim( "model-x", pre_load_memory=1 * 1024**3 ) await asyncio.wait_for( pool._failed_load_reclaim_task, timeout=10 ) finally: executor.shutdown(wait=False) assert gauge["active"] == 1 * 1024**3 async def test_reclaim_gives_up_after_max_rounds(self): """A permanently-inflated gauge must not hang the reclaim task.""" pool = _make_pool() import concurrent.futures executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) try: with ( patch( "omlx.engine_pool.mx.get_active_memory", return_value=80 * 1024**3, ), patch("omlx.engine_pool.mx.synchronize"), patch("omlx.engine_pool.mx.clear_cache"), patch("omlx.engine_pool.get_phys_footprint", return_value=0), patch( "omlx.engine_pool.get_mlx_executor", return_value=executor ), patch( "omlx.engine_pool.asyncio.sleep", new=AsyncMock(return_value=None), ), ): pool._schedule_failed_load_reclaim( "model-x", pre_load_memory=1 * 1024**3 ) await asyncio.wait_for( pool._failed_load_reclaim_task, timeout=10 ) finally: executor.shutdown(wait=False) async def test_failed_load_schedules_reclaim(self, small_mock_model_dir): pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) failing_engine = MagicMock() failing_engine.start = AsyncMock(side_effect=RuntimeError("boom")) failing_engine.stop = AsyncMock() with ( patch("omlx.engine_pool.BatchedEngine", return_value=failing_engine), patch.object(pool, "_schedule_failed_load_reclaim") as scheduled, pytest.raises(ModelUnavailableError), ): await pool._load_engine("model-a") scheduled.assert_called_once() args = scheduled.call_args[0] assert args[0] == "model-a" class TestSchedulerConfigModelId: """engine_pool sets model_name=model_id and model_path=entry.model_path on the shared scheduler config so the scheduler uses the engine pool's canonical model_id (not a filesystem snapshot hash) for tracker lookups.""" async def test_wires_model_name_and_model_path(self, small_mock_model_dir): """engine_pool sets model_name=model_id and model_path on the scheduler config when loading an engine.""" pool = _make_pool(ceiling=10 * 1024**3) pool.discover_models(str(small_mock_model_dir)) mock_engine = MagicMock() mock_engine.start = AsyncMock() with patch( "omlx.engine_pool.BatchedEngine", return_value=mock_engine, ): await pool.get_engine("model-a") assert pool._scheduler_config.model_name == "model-a" assert pool._scheduler_config.model_path == str(small_mock_model_dir / "model-a")