882 lines
No EOL
31 KiB
Python
882 lines
No EOL
31 KiB
Python
"""
|
|
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 |