""" Unit tests for model management functionality. """ import pytest import os import tempfile import threading import time from unittest.mock import Mock, patch, MagicMock import torch import numpy as np from detector_worker.models.model_manager import ( ModelManager, ModelInfo, ModelConfig, ModelCache, ModelLoader, ModelError, ModelLoadError, ModelCacheError ) from detector_worker.core.exceptions import ConfigurationError class TestModelConfig: """Test model configuration.""" def test_creation(self): """Test model config creation.""" config = ModelConfig( model_id="yolo_v8_car", model_path="/models/yolo_v8_car.pt", model_type="detection", device="cuda:0" ) assert config.model_id == "yolo_v8_car" assert config.model_path == "/models/yolo_v8_car.pt" assert config.model_type == "detection" assert config.device == "cuda:0" assert config.confidence_threshold == 0.5 assert config.max_memory_mb == 1024 def test_creation_with_optional_params(self): """Test config creation with optional parameters.""" config = ModelConfig( model_id="classifier_v1", model_path="/models/classifier.pt", model_type="classification", device="cpu", confidence_threshold=0.8, max_memory_mb=512, class_names={0: "car", 1: "truck", 2: "bus"}, preprocessing_config={"resize": (224, 224), "normalize": True} ) assert config.confidence_threshold == 0.8 assert config.max_memory_mb == 512 assert config.class_names[0] == "car" assert config.preprocessing_config["resize"] == (224, 224) def test_from_dict(self): """Test creating config from dictionary.""" config_dict = { "model_id": "detection_model", "model_path": "/path/to/model.pt", "model_type": "detection", "device": "cuda:0", "confidence_threshold": 0.75, "class_names": {0: "person", 1: "vehicle"}, "unknown_field": "ignored" } config = ModelConfig.from_dict(config_dict) assert config.model_id == "detection_model" assert config.confidence_threshold == 0.75 assert config.class_names[1] == "vehicle" def test_validation(self): """Test config validation.""" # Valid config valid_config = ModelConfig( model_id="test_model", model_path="/valid/path/model.pt", model_type="detection", device="cpu" ) assert valid_config.is_valid() is True # Invalid config (empty model_id) invalid_config = ModelConfig( model_id="", model_path="/path/model.pt", model_type="detection", device="cpu" ) assert invalid_config.is_valid() is False def test_get_memory_limit_bytes(self): """Test getting memory limit in bytes.""" config = ModelConfig( model_id="test", model_path="/path", model_type="detection", device="cpu", max_memory_mb=256 ) assert config.get_memory_limit_bytes() == 256 * 1024 * 1024 class TestModelInfo: """Test model information.""" def test_creation(self): """Test model info creation.""" config = ModelConfig( model_id="test_model", model_path="/path/model.pt", model_type="detection", device="cuda:0" ) mock_model = Mock() info = ModelInfo( config=config, model_instance=mock_model, load_time=1.5 ) assert info.config == config assert info.model_instance == mock_model assert info.load_time == 1.5 assert info.reference_count == 0 assert info.last_used <= time.time() assert info.memory_usage == 0 def test_increment_reference(self): """Test incrementing reference count.""" config = ModelConfig("test", "/path", "detection", "cpu") info = ModelInfo(config, Mock(), 1.0) assert info.reference_count == 0 info.increment_reference() assert info.reference_count == 1 info.increment_reference() assert info.reference_count == 2 def test_decrement_reference(self): """Test decrementing reference count.""" config = ModelConfig("test", "/path", "detection", "cpu") info = ModelInfo(config, Mock(), 1.0) info.reference_count = 3 assert info.decrement_reference() == 2 assert info.reference_count == 2 assert info.decrement_reference() == 1 assert info.decrement_reference() == 0 # Should not go below 0 assert info.decrement_reference() == 0 def test_update_usage(self): """Test updating usage statistics.""" config = ModelConfig("test", "/path", "detection", "cpu") info = ModelInfo(config, Mock(), 1.0) original_time = info.last_used original_count = info.usage_count time.sleep(0.01) # Small delay info.update_usage(memory_usage=512*1024*1024) # 512MB assert info.last_used > original_time assert info.usage_count == original_count + 1 assert info.memory_usage == 512*1024*1024 def test_age_calculation(self): """Test age calculation.""" config = ModelConfig("test", "/path", "detection", "cpu") info = ModelInfo(config, Mock(), 1.0) time.sleep(0.01) age = info.age() assert age > 0 assert age < 1 # Should be less than 1 second def test_get_stats(self): """Test getting model statistics.""" config = ModelConfig("test_model", "/path", "detection", "cuda:0") info = ModelInfo(config, Mock(), 2.5) info.reference_count = 3 info.usage_count = 100 info.memory_usage = 1024*1024*1024 # 1GB stats = info.get_stats() assert stats["model_id"] == "test_model" assert stats["device"] == "cuda:0" assert stats["load_time"] == 2.5 assert stats["reference_count"] == 3 assert stats["usage_count"] == 100 assert stats["memory_usage_mb"] == 1024 assert "age_seconds" in stats class TestModelLoader: """Test model loading functionality.""" def test_creation(self): """Test model loader creation.""" loader = ModelLoader() assert loader.supported_formats == [".pt", ".pth", ".onnx", ".trt"] assert loader.default_device == "cpu" def test_detect_device_cuda_available(self): """Test device detection when CUDA is available.""" loader = ModelLoader() with patch('torch.cuda.is_available', return_value=True): with patch('torch.cuda.device_count', return_value=2): device = loader.detect_optimal_device() assert device == "cuda:0" def test_detect_device_cuda_unavailable(self): """Test device detection when CUDA is not available.""" loader = ModelLoader() with patch('torch.cuda.is_available', return_value=False): device = loader.detect_optimal_device() assert device == "cpu" def test_load_pytorch_model(self): """Test loading PyTorch model.""" loader = ModelLoader() with patch('torch.load') as mock_torch_load: with patch('os.path.exists', return_value=True): mock_model = Mock() mock_torch_load.return_value = mock_model config = ModelConfig( model_id="test_model", model_path="/path/to/model.pt", model_type="detection", device="cpu" ) loaded_model = loader.load_model(config) assert loaded_model == mock_model mock_torch_load.assert_called_once_with("/path/to/model.pt", map_location="cpu") def test_load_model_file_not_exists(self): """Test loading model when file doesn't exist.""" loader = ModelLoader() with patch('os.path.exists', return_value=False): config = ModelConfig( model_id="missing_model", model_path="/nonexistent/model.pt", model_type="detection", device="cpu" ) with pytest.raises(ModelLoadError) as exc_info: loader.load_model(config) assert "does not exist" in str(exc_info.value) def test_load_model_invalid_format(self): """Test loading model with invalid format.""" loader = ModelLoader() with patch('os.path.exists', return_value=True): config = ModelConfig( model_id="invalid_model", model_path="/path/to/model.invalid", model_type="detection", device="cpu" ) with pytest.raises(ModelLoadError) as exc_info: loader.load_model(config) assert "unsupported format" in str(exc_info.value).lower() def test_load_model_torch_error(self): """Test loading model with torch loading error.""" loader = ModelLoader() with patch('os.path.exists', return_value=True): with patch('torch.load', side_effect=RuntimeError("CUDA out of memory")): config = ModelConfig( model_id="error_model", model_path="/path/to/model.pt", model_type="detection", device="cuda:0" ) with pytest.raises(ModelLoadError) as exc_info: loader.load_model(config) assert "CUDA out of memory" in str(exc_info.value) def test_validate_model_pytorch(self): """Test validating PyTorch model.""" loader = ModelLoader() mock_model = Mock() mock_model.__class__.__module__ = "torch.nn" config = ModelConfig("test", "/path", "detection", "cpu") is_valid = loader.validate_model(mock_model, config) assert is_valid is True def test_validate_model_invalid(self): """Test validating invalid model.""" loader = ModelLoader() invalid_model = "not_a_model" config = ModelConfig("test", "/path", "detection", "cpu") is_valid = loader.validate_model(invalid_model, config) assert is_valid is False def test_estimate_model_memory(self): """Test estimating model memory usage.""" loader = ModelLoader() mock_model = Mock() mock_param1 = Mock() mock_param1.numel.return_value = 1000000 # 1M parameters mock_param1.element_size.return_value = 4 # 4 bytes per parameter mock_param2 = Mock() mock_param2.numel.return_value = 500000 # 0.5M parameters mock_param2.element_size.return_value = 4 mock_model.parameters.return_value = [mock_param1, mock_param2] memory_bytes = loader.estimate_memory_usage(mock_model) expected_bytes = (1000000 + 500000) * 4 # 6MB assert memory_bytes == expected_bytes class TestModelCache: """Test model caching functionality.""" def test_creation(self): """Test model cache creation.""" cache = ModelCache(max_size=5, max_memory_mb=2048) assert cache.max_size == 5 assert cache.max_memory_mb == 2048 assert len(cache.models) == 0 assert len(cache.access_order) == 0 def test_put_and_get_model(self): """Test putting and getting model from cache.""" cache = ModelCache(max_size=3) config = ModelConfig("test_model", "/path", "detection", "cpu") mock_model = Mock() model_info = ModelInfo(config, mock_model, 1.5) cache.put("test_model", model_info) retrieved_info = cache.get("test_model") assert retrieved_info == model_info assert retrieved_info.reference_count == 1 # Should be incremented on get def test_get_nonexistent_model(self): """Test getting non-existent model.""" cache = ModelCache(max_size=3) result = cache.get("nonexistent_model") assert result is None def test_contains_check(self): """Test checking if model exists in cache.""" cache = ModelCache(max_size=3) config = ModelConfig("test_model", "/path", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) cache.put("test_model", model_info) assert cache.contains("test_model") is True assert cache.contains("nonexistent_model") is False def test_remove_model(self): """Test removing model from cache.""" cache = ModelCache(max_size=3) config = ModelConfig("test_model", "/path", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) cache.put("test_model", model_info) assert cache.contains("test_model") is True removed_info = cache.remove("test_model") assert removed_info == model_info assert cache.contains("test_model") is False def test_lru_eviction(self): """Test LRU eviction policy.""" cache = ModelCache(max_size=2) # Add models to fill cache for i in range(2): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) cache.put(f"model_{i}", model_info) # Access model_0 to make it recently used cache.get("model_0") # Add another model (should evict model_1, the least recently used) config = ModelConfig("model_2", "/path_2", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) cache.put("model_2", model_info) assert cache.size() == 2 assert cache.contains("model_0") is True # Recently accessed assert cache.contains("model_1") is False # Evicted assert cache.contains("model_2") is True # Newly added def test_memory_based_eviction(self): """Test memory-based eviction.""" cache = ModelCache(max_size=10, max_memory_mb=1) # 1MB limit # Add model that uses 0.8MB config1 = ModelConfig("model_1", "/path_1", "detection", "cpu") model1 = Mock() info1 = ModelInfo(config1, model1, 1.0) info1.memory_usage = 0.8 * 1024 * 1024 # 0.8MB cache.put("model_1", info1) # Add model that would exceed memory limit config2 = ModelConfig("model_2", "/path_2", "detection", "cpu") model2 = Mock() info2 = ModelInfo(config2, model2, 1.0) info2.memory_usage = 0.5 * 1024 * 1024 # 0.5MB cache.put("model_2", info2) # First model should be evicted due to memory constraint assert cache.contains("model_1") is False assert cache.contains("model_2") is True def test_get_stats(self): """Test getting cache statistics.""" cache = ModelCache(max_size=5) # Add some models for i in range(3): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) model_info.memory_usage = 100 * 1024 * 1024 # 100MB each cache.put(f"model_{i}", model_info) # Access some models cache.get("model_0") cache.get("model_1") cache.get("nonexistent") # Miss stats = cache.get_stats() assert stats["size"] == 3 assert stats["max_size"] == 5 assert stats["hits"] == 2 assert stats["misses"] == 1 assert stats["hit_rate"] == 2/3 assert stats["memory_usage_mb"] == 300 def test_clear_cache(self): """Test clearing entire cache.""" cache = ModelCache(max_size=5) # Add models for i in range(3): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) cache.put(f"model_{i}", model_info) assert cache.size() == 3 cache.clear() assert cache.size() == 0 assert len(cache.models) == 0 assert len(cache.access_order) == 0 class TestModelManager: """Test main model manager functionality.""" def test_initialization(self): """Test model manager initialization.""" manager = ModelManager() assert isinstance(manager.cache, ModelCache) assert isinstance(manager.loader, ModelLoader) assert manager.models_directory == "models" assert manager.default_device == "cpu" def test_initialization_with_config(self): """Test initialization with custom configuration.""" config = { "models_directory": "/custom/models", "default_device": "cuda:0", "cache_max_size": 20, "cache_max_memory_mb": 4096 } manager = ModelManager(config) assert manager.models_directory == "/custom/models" assert manager.default_device == "cuda:0" assert manager.cache.max_size == 20 assert manager.cache.max_memory_mb == 4096 def test_load_model_new(self): """Test loading new model.""" manager = ModelManager() config = ModelConfig( model_id="test_model", model_path="/path/to/model.pt", model_type="detection", device="cpu" ) with patch.object(manager.loader, 'load_model') as mock_load: with patch.object(manager.loader, 'estimate_memory_usage', return_value=512*1024*1024): mock_model = Mock() mock_load.return_value = mock_model loaded_model = manager.load_model(config) assert loaded_model == mock_model assert manager.cache.contains("test_model") is True mock_load.assert_called_once_with(config) def test_load_model_from_cache(self): """Test loading model from cache.""" manager = ModelManager() # Pre-populate cache config = ModelConfig("cached_model", "/path", "detection", "cpu") mock_model = Mock() model_info = ModelInfo(config, mock_model, 1.0) manager.cache.put("cached_model", model_info) with patch.object(manager.loader, 'load_model') as mock_load: loaded_model = manager.load_model(config) assert loaded_model == mock_model mock_load.assert_not_called() # Should not load from disk def test_get_model_by_id(self): """Test getting model by ID.""" manager = ModelManager() config = ModelConfig("test_model", "/path", "detection", "cpu") mock_model = Mock() model_info = ModelInfo(config, mock_model, 1.0) manager.cache.put("test_model", model_info) retrieved_model = manager.get_model("test_model") assert retrieved_model == mock_model def test_get_nonexistent_model(self): """Test getting non-existent model.""" manager = ModelManager() model = manager.get_model("nonexistent_model") assert model is None def test_unload_model_with_references(self): """Test unloading model with active references.""" manager = ModelManager() config = ModelConfig("ref_model", "/path", "detection", "cpu") mock_model = Mock() model_info = ModelInfo(config, mock_model, 1.0) model_info.reference_count = 2 # Active references manager.cache.put("ref_model", model_info) result = manager.unload_model("ref_model") assert result is False # Should not unload with active references assert manager.cache.contains("ref_model") is True def test_unload_model_no_references(self): """Test unloading model without references.""" manager = ModelManager() config = ModelConfig("no_ref_model", "/path", "detection", "cpu") mock_model = Mock() model_info = ModelInfo(config, mock_model, 1.0) model_info.reference_count = 0 # No references manager.cache.put("no_ref_model", model_info) result = manager.unload_model("no_ref_model") assert result is True assert manager.cache.contains("no_ref_model") is False def test_list_loaded_models(self): """Test listing loaded models.""" manager = ModelManager() # Add models to cache for i in range(3): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) manager.cache.put(f"model_{i}", model_info) loaded_models = manager.list_loaded_models() assert len(loaded_models) == 3 assert all(info["model_id"].startswith("model_") for info in loaded_models) def test_get_model_info(self): """Test getting model information.""" manager = ModelManager() config = ModelConfig("info_model", "/path", "detection", "cuda:0") mock_model = Mock() model_info = ModelInfo(config, mock_model, 2.5) model_info.usage_count = 10 manager.cache.put("info_model", model_info) info = manager.get_model_info("info_model") assert info is not None assert info["model_id"] == "info_model" assert info["device"] == "cuda:0" assert info["load_time"] == 2.5 assert info["usage_count"] == 10 def test_cleanup_unused_models(self): """Test cleaning up unused models.""" manager = ModelManager() # Add models with different reference counts models_data = [ ("used_model", 2), # Has references ("unused_model_1", 0), # No references ("unused_model_2", 0) # No references ] for model_id, ref_count in models_data: config = ModelConfig(model_id, f"/path/{model_id}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) model_info.reference_count = ref_count manager.cache.put(model_id, model_info) cleaned_count = manager.cleanup_unused_models() assert cleaned_count == 2 # Two unused models cleaned assert manager.cache.contains("used_model") is True assert manager.cache.contains("unused_model_1") is False assert manager.cache.contains("unused_model_2") is False def test_get_memory_usage(self): """Test getting total memory usage.""" manager = ModelManager() # Add models with different memory usage memory_sizes = [256, 512, 1024] # MB for i, memory_mb in enumerate(memory_sizes): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) model_info.memory_usage = memory_mb * 1024 * 1024 # Convert to bytes manager.cache.put(f"model_{i}", model_info) total_usage = manager.get_memory_usage() expected_bytes = sum(memory_sizes) * 1024 * 1024 assert total_usage == expected_bytes def test_health_check(self): """Test model manager health check.""" manager = ModelManager() # Add models for i in range(3): config = ModelConfig(f"model_{i}", f"/path_{i}", "detection", "cpu") model_info = ModelInfo(config, Mock(), 1.0) model_info.memory_usage = 100 * 1024 * 1024 # 100MB each manager.cache.put(f"model_{i}", model_info) health_report = manager.health_check() assert health_report["status"] == "healthy" assert health_report["loaded_models"] == 3 assert health_report["total_memory_mb"] == 300 assert health_report["cache_hit_rate"] >= 0 class TestModelManagerIntegration: """Integration tests for model manager.""" def test_concurrent_model_loading(self): """Test concurrent model loading.""" manager = ModelManager() # Mock loader to simulate loading time def slow_load(config): time.sleep(0.1) # Simulate loading time mock_model = Mock() mock_model.model_id = config.model_id return mock_model with patch.object(manager.loader, 'load_model', side_effect=slow_load): with patch.object(manager.loader, 'estimate_memory_usage', return_value=100*1024*1024): # Create multiple threads loading different models results = {} errors = [] def load_model_thread(model_id): try: config = ModelConfig( model_id=model_id, model_path=f"/path/{model_id}.pt", model_type="detection", device="cpu" ) model = manager.load_model(config) results[model_id] = model except Exception as e: errors.append((model_id, str(e))) threads = [] for i in range(5): thread = threading.Thread(target=load_model_thread, args=(f"model_{i}",)) threads.append(thread) thread.start() for thread in threads: thread.join() # All models should be loaded successfully assert len(errors) == 0 assert len(results) == 5 assert len(manager.cache.models) == 5 def test_memory_pressure_handling(self): """Test handling memory pressure.""" # Create manager with small memory limit manager = ModelManager({ "cache_max_memory_mb": 200 # 200MB limit }) with patch.object(manager.loader, 'load_model') as mock_load: with patch.object(manager.loader, 'estimate_memory_usage', return_value=100*1024*1024): # 100MB per model def create_mock_model(config): mock_model = Mock() mock_model.model_id = config.model_id return mock_model mock_load.side_effect = create_mock_model # Load models that exceed memory limit for i in range(4): # 4 * 100MB = 400MB > 200MB limit config = ModelConfig( model_id=f"large_model_{i}", model_path=f"/path/large_model_{i}.pt", model_type="detection", device="cpu" ) manager.load_model(config) # Should not exceed memory limit due to eviction total_memory = manager.get_memory_usage() memory_limit = 200 * 1024 * 1024 assert total_memory <= memory_limit def test_model_lifecycle_management(self): """Test complete model lifecycle.""" manager = ModelManager() with patch.object(manager.loader, 'load_model') as mock_load: with patch.object(manager.loader, 'estimate_memory_usage', return_value=50*1024*1024): mock_model = Mock() mock_load.return_value = mock_model config = ModelConfig( model_id="lifecycle_model", model_path="/path/lifecycle_model.pt", model_type="detection", device="cpu" ) # 1. Load model loaded_model = manager.load_model(config) assert loaded_model == mock_model assert manager.cache.contains("lifecycle_model") is True # 2. Get model multiple times (increase usage) for _ in range(5): model = manager.get_model("lifecycle_model") assert model == mock_model # 3. Check model info info = manager.get_model_info("lifecycle_model") assert info["usage_count"] >= 5 # 4. Simulate model still in use model_info = manager.cache.get("lifecycle_model") model_info.reference_count = 1 # Should not unload while in use unloaded = manager.unload_model("lifecycle_model") assert unloaded is False assert manager.cache.contains("lifecycle_model") is True # 5. Release reference and unload model_info.reference_count = 0 unloaded = manager.unload_model("lifecycle_model") assert unloaded is True assert manager.cache.contains("lifecycle_model") is False def test_error_recovery(self): """Test error recovery scenarios.""" manager = ModelManager() # Test loading model that fails initially then succeeds call_count = 0 def failing_then_success_load(config): nonlocal call_count call_count += 1 if call_count == 1: raise ModelLoadError("First attempt failed") return Mock() with patch.object(manager.loader, 'load_model', side_effect=failing_then_success_load): with patch.object(manager.loader, 'estimate_memory_usage', return_value=50*1024*1024): config = ModelConfig( model_id="retry_model", model_path="/path/retry_model.pt", model_type="detection", device="cpu" ) # First attempt should fail with pytest.raises(ModelLoadError): manager.load_model(config) # Model should not be in cache assert manager.cache.contains("retry_model") is False # Second attempt should succeed model = manager.load_model(config) assert model is not None assert manager.cache.contains("retry_model") is True