wehub-resource-sync
583 lines
19 KiB
Python
583 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for cache factory, recovery manager, and hybrid cache configuration.
|
|
|
|
This module tests the factory pattern for creating cache instances,
|
|
recovery from cache corruption, and hybrid model cache configurations.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
from unittest.mock import MagicMock, patch, PropertyMock
|
|
|
|
import pytest
|
|
|
|
from omlx.cache.factory import CacheConfig, CacheFactory
|
|
from omlx.cache.recovery import CacheRecoveryManager
|
|
from omlx.cache.hybrid_cache import (
|
|
LayerCacheConfig,
|
|
ModelCacheConfig,
|
|
create_default_kvcache_config,
|
|
)
|
|
from omlx.cache.type_handlers import CacheType
|
|
|
|
|
|
class TestCacheConfig:
|
|
"""Tests for CacheConfig dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default values."""
|
|
config = CacheConfig()
|
|
|
|
assert config.block_size == 64
|
|
assert config.max_num_blocks == 1024
|
|
assert config.initial_blocks == 256
|
|
assert config.paged_ssd_cache_dir is None
|
|
assert config.max_paged_ssd_cache_size == 100 * 1024 * 1024 * 1024 # 100GB
|
|
assert config.max_kv_cache_memory is None
|
|
assert config.model_name == ""
|
|
|
|
def test_with_custom_values(self):
|
|
"""Test with custom values."""
|
|
config = CacheConfig(
|
|
block_size=128,
|
|
max_num_blocks=2048,
|
|
initial_blocks=512,
|
|
paged_ssd_cache_dir=Path("/tmp/cache"),
|
|
max_paged_ssd_cache_size=50 * 1024**3,
|
|
max_kv_cache_memory=8 * 1024**3,
|
|
model_name="test-model",
|
|
)
|
|
|
|
assert config.block_size == 128
|
|
assert config.max_num_blocks == 2048
|
|
assert config.initial_blocks == 512
|
|
assert config.paged_ssd_cache_dir == Path("/tmp/cache")
|
|
assert config.max_paged_ssd_cache_size == 50 * 1024**3
|
|
assert config.max_kv_cache_memory == 8 * 1024**3
|
|
assert config.model_name == "test-model"
|
|
|
|
|
|
class TestCacheFactory:
|
|
"""Tests for CacheFactory."""
|
|
|
|
def test_create_paged_cache_disabled(self):
|
|
"""Test create_paged_cache returns None when SSD cache disabled."""
|
|
config = CacheConfig(paged_ssd_cache_dir=None)
|
|
|
|
result = CacheFactory.create_paged_cache(config)
|
|
|
|
assert result is None
|
|
|
|
def test_create_paged_cache_enabled(self, tmp_path: Path):
|
|
"""Test create_paged_cache creates manager when enabled."""
|
|
config = CacheConfig(
|
|
paged_ssd_cache_dir=tmp_path / "cache",
|
|
block_size=64,
|
|
max_num_blocks=100,
|
|
initial_blocks=50,
|
|
model_name="test-model",
|
|
)
|
|
|
|
result = CacheFactory.create_paged_cache(config)
|
|
|
|
assert result is not None
|
|
assert result.block_size == 64
|
|
assert result.max_blocks == 100
|
|
assert result.model_name == "test-model"
|
|
|
|
def test_create_paged_ssd_cache_disabled(self):
|
|
"""Test create_paged_ssd_cache returns None when disabled."""
|
|
config = CacheConfig(paged_ssd_cache_dir=None)
|
|
|
|
result = CacheFactory.create_paged_ssd_cache(config)
|
|
|
|
assert result is None
|
|
|
|
def test_create_paged_ssd_cache_enabled(self, tmp_path: Path):
|
|
"""Test create_paged_ssd_cache creates manager."""
|
|
config = CacheConfig(
|
|
paged_ssd_cache_dir=tmp_path / "cache",
|
|
max_paged_ssd_cache_size=1024**3,
|
|
)
|
|
|
|
result = CacheFactory.create_paged_ssd_cache(config, model_name="test-model")
|
|
|
|
assert result is not None
|
|
# Check that model_name is used in path
|
|
assert "test-model" in str(result._cache_dir)
|
|
|
|
def test_create_prefix_cache_disabled(self):
|
|
"""Test create_prefix_cache returns None when disabled."""
|
|
config = CacheConfig(paged_ssd_cache_dir=None)
|
|
|
|
result = CacheFactory.create_prefix_cache(config)
|
|
|
|
assert result is None
|
|
|
|
def test_create_prefix_cache_no_paged_cache(self, tmp_path: Path):
|
|
"""Test create_prefix_cache returns None without paged cache."""
|
|
config = CacheConfig(paged_ssd_cache_dir=tmp_path / "cache")
|
|
|
|
result = CacheFactory.create_prefix_cache(config, paged_cache=None)
|
|
|
|
assert result is None
|
|
|
|
def test_create_prefix_cache_enabled(self, tmp_path: Path):
|
|
"""Test create_prefix_cache creates cache."""
|
|
config = CacheConfig(paged_ssd_cache_dir=tmp_path / "cache")
|
|
|
|
paged_cache = CacheFactory.create_paged_cache(config)
|
|
mock_model = MagicMock()
|
|
mock_model.layers = [MagicMock() for _ in range(4)]
|
|
|
|
result = CacheFactory.create_prefix_cache(
|
|
config,
|
|
model=mock_model,
|
|
paged_cache=paged_cache,
|
|
)
|
|
|
|
assert result is not None
|
|
|
|
def test_create_memory_monitor(self, tmp_path: Path):
|
|
"""Test create_memory_monitor creates monitor."""
|
|
config = CacheConfig(
|
|
paged_ssd_cache_dir=tmp_path / "cache",
|
|
max_kv_cache_memory=8 * 1024**3,
|
|
)
|
|
|
|
paged_cache = CacheFactory.create_paged_cache(config)
|
|
|
|
result = CacheFactory.create_memory_monitor(config, paged_cache)
|
|
|
|
assert result is not None
|
|
|
|
def test_create_memory_monitor_default_memory(self, tmp_path: Path):
|
|
"""Test create_memory_monitor with default memory limit."""
|
|
config = CacheConfig(
|
|
paged_ssd_cache_dir=tmp_path / "cache",
|
|
max_kv_cache_memory=None, # Should use default
|
|
)
|
|
|
|
result = CacheFactory.create_memory_monitor(config)
|
|
|
|
assert result is not None
|
|
|
|
def test_create_full_cache_stack_disabled(self):
|
|
"""Test create_full_cache_stack returns Nones when disabled."""
|
|
config = CacheConfig(paged_ssd_cache_dir=None)
|
|
|
|
result = CacheFactory.create_full_cache_stack(config)
|
|
|
|
assert result["paged_cache"] is None
|
|
assert result["paged_ssd_cache"] is None
|
|
assert result["prefix_cache"] is None
|
|
assert result["memory_monitor"] is None
|
|
|
|
def test_create_full_cache_stack_enabled(self, tmp_path: Path):
|
|
"""Test create_full_cache_stack creates all components."""
|
|
config = CacheConfig(
|
|
paged_ssd_cache_dir=tmp_path / "cache",
|
|
model_name="test-model",
|
|
)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.layers = [MagicMock() for _ in range(4)]
|
|
|
|
result = CacheFactory.create_full_cache_stack(
|
|
config,
|
|
model=mock_model,
|
|
)
|
|
|
|
assert result["paged_cache"] is not None
|
|
assert result["paged_ssd_cache"] is not None
|
|
assert result["prefix_cache"] is not None
|
|
assert result["memory_monitor"] is not None
|
|
|
|
|
|
class TestCacheRecoveryManager:
|
|
"""Tests for CacheRecoveryManager."""
|
|
|
|
@pytest.fixture
|
|
def mock_prefix_cache(self):
|
|
"""Create a mock prefix cache."""
|
|
mock = MagicMock()
|
|
mock.clear.return_value = 10
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
def recovery_manager(self, mock_prefix_cache):
|
|
"""Create a recovery manager."""
|
|
return CacheRecoveryManager(block_aware_cache=mock_prefix_cache)
|
|
|
|
def test_is_cache_corruption_shape_mismatch(self, recovery_manager):
|
|
"""Test detecting shape mismatch errors."""
|
|
error = ValueError("Shapes (1, 32, 64) and (1, 32, 128) do not match")
|
|
|
|
# This depends on is_cache_corruption_error implementation
|
|
# Just verify method exists and returns bool
|
|
result = recovery_manager.is_cache_corruption(error)
|
|
assert isinstance(result, bool)
|
|
|
|
def test_is_cache_corruption_dimension_error(self, recovery_manager):
|
|
"""Test detecting dimension errors."""
|
|
error = ValueError("Dimension mismatch in concatenate")
|
|
|
|
result = recovery_manager.is_cache_corruption(error)
|
|
assert isinstance(result, bool)
|
|
|
|
def test_recover(self, recovery_manager, mock_prefix_cache):
|
|
"""Test recovery process."""
|
|
# Setup mock objects
|
|
mock_holder = MagicMock()
|
|
mock_holder.batch_generator = MagicMock()
|
|
mock_holder._current_sampler_params = MagicMock()
|
|
|
|
request_id_to_uid = {"req-001": 1, "req-002": 2}
|
|
uid_to_request_id = {1: "req-001", 2: "req-002"}
|
|
request_detokenizers = {"req-001": MagicMock()}
|
|
|
|
# Perform recovery
|
|
recovery_manager.recover(
|
|
mock_holder,
|
|
request_id_to_uid,
|
|
uid_to_request_id,
|
|
request_detokenizers,
|
|
)
|
|
|
|
# Verify batch_generator cleared
|
|
assert mock_holder.batch_generator is None
|
|
assert mock_holder._current_sampler_params is None
|
|
|
|
# Verify cache cleared
|
|
mock_prefix_cache.clear.assert_called_once()
|
|
|
|
# Verify mappings cleared
|
|
assert len(request_id_to_uid) == 0
|
|
assert len(uid_to_request_id) == 0
|
|
assert len(request_detokenizers) == 0
|
|
|
|
def test_recover_without_cache(self):
|
|
"""Test recovery without prefix cache."""
|
|
manager = CacheRecoveryManager(block_aware_cache=None)
|
|
|
|
mock_holder = MagicMock()
|
|
request_id_to_uid = {"req-001": 1}
|
|
uid_to_request_id = {1: "req-001"}
|
|
request_detokenizers = {}
|
|
|
|
# Should not raise
|
|
manager.recover(
|
|
mock_holder,
|
|
request_id_to_uid,
|
|
uid_to_request_id,
|
|
request_detokenizers,
|
|
)
|
|
|
|
assert mock_holder.batch_generator is None
|
|
|
|
def test_reschedule_running_requests(self, recovery_manager):
|
|
"""Test rescheduling running requests."""
|
|
from collections import deque
|
|
|
|
# Create mock requests
|
|
mock_request1 = MagicMock()
|
|
mock_request1.prompt_token_ids = [1, 2, 3, 4]
|
|
mock_request2 = MagicMock()
|
|
mock_request2.prompt_token_ids = [5, 6, 7, 8]
|
|
|
|
running = {
|
|
"req-001": mock_request1,
|
|
"req-002": mock_request2,
|
|
}
|
|
waiting = deque()
|
|
mock_waiting_status = MagicMock()
|
|
|
|
count = recovery_manager.reschedule_running_requests(
|
|
running,
|
|
waiting,
|
|
mock_waiting_status,
|
|
)
|
|
|
|
assert count == 2
|
|
assert len(running) == 0
|
|
assert len(waiting) == 2
|
|
|
|
# Verify requests were reset
|
|
assert mock_request1.status is mock_waiting_status
|
|
assert mock_request1.batch_uid is None
|
|
assert mock_request1.prompt_cache is None
|
|
|
|
def test_reschedule_empty_running(self, recovery_manager):
|
|
"""Test rescheduling with no running requests."""
|
|
from collections import deque
|
|
|
|
running = {}
|
|
waiting = deque()
|
|
mock_waiting_status = MagicMock()
|
|
|
|
count = recovery_manager.reschedule_running_requests(
|
|
running,
|
|
waiting,
|
|
mock_waiting_status,
|
|
)
|
|
|
|
assert count == 0
|
|
|
|
|
|
class TestLayerCacheConfig:
|
|
"""Tests for LayerCacheConfig."""
|
|
|
|
def test_creation(self):
|
|
"""Test creating layer config."""
|
|
config = LayerCacheConfig(
|
|
layer_idx=5,
|
|
cache_type=CacheType.KVCACHE,
|
|
supports_block_slicing=True,
|
|
class_name="KVCache",
|
|
)
|
|
|
|
assert config.layer_idx == 5
|
|
assert config.cache_type == CacheType.KVCACHE
|
|
assert config.supports_block_slicing is True
|
|
assert config.class_name == "KVCache"
|
|
|
|
def test_handler_property(self):
|
|
"""Test handler property returns correct handler."""
|
|
config = LayerCacheConfig(
|
|
layer_idx=0,
|
|
cache_type=CacheType.KVCACHE,
|
|
supports_block_slicing=True,
|
|
class_name="KVCache",
|
|
)
|
|
|
|
handler = config.handler
|
|
|
|
assert handler is not None
|
|
assert handler.cache_type == CacheType.KVCACHE
|
|
|
|
def test_arrays_cache_layer(self):
|
|
"""Test config for ArraysCache layer."""
|
|
config = LayerCacheConfig(
|
|
layer_idx=1,
|
|
cache_type=CacheType.ARRAYS_CACHE,
|
|
supports_block_slicing=False,
|
|
class_name="ArraysCache",
|
|
)
|
|
|
|
assert config.supports_block_slicing is False
|
|
assert config.handler.cache_type == CacheType.ARRAYS_CACHE
|
|
|
|
|
|
class TestModelCacheConfig:
|
|
"""Tests for ModelCacheConfig."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default values."""
|
|
config = ModelCacheConfig()
|
|
|
|
assert config.model_name == ""
|
|
assert config.num_layers == 0
|
|
assert config.layer_configs == []
|
|
assert config.is_hybrid is False
|
|
assert config.sliceable_layer_count == 0
|
|
|
|
def test_from_cache_list_empty(self):
|
|
"""Test from_cache_list with empty list."""
|
|
config = ModelCacheConfig.from_cache_list([], model_name="test")
|
|
|
|
assert config.model_name == "test"
|
|
assert config.num_layers == 0
|
|
|
|
def test_from_cache_list_kvcache_only(self):
|
|
"""Test from_cache_list with KVCache only."""
|
|
mock_caches = []
|
|
for i in range(4):
|
|
mock = MagicMock()
|
|
mock.__class__.__name__ = "KVCache"
|
|
mock_caches.append(mock)
|
|
|
|
config = ModelCacheConfig.from_cache_list(mock_caches, model_name="test-model")
|
|
|
|
assert config.num_layers == 4
|
|
assert config.is_hybrid is False
|
|
assert config.sliceable_layer_count == 4
|
|
|
|
def test_from_cache_list_hybrid(self):
|
|
"""Test from_cache_list with mixed cache types."""
|
|
mock_caches = []
|
|
|
|
# KVCache layer
|
|
mock_kv = MagicMock()
|
|
mock_kv.__class__.__name__ = "KVCache"
|
|
mock_caches.append(mock_kv)
|
|
|
|
# ArraysCache layer
|
|
mock_arrays = MagicMock()
|
|
mock_arrays.__class__.__name__ = "ArraysCache"
|
|
mock_arrays.cache = [MagicMock(), MagicMock()]
|
|
mock_caches.append(mock_arrays)
|
|
|
|
config = ModelCacheConfig.from_cache_list(mock_caches)
|
|
|
|
assert config.num_layers == 2
|
|
assert config.is_hybrid is True
|
|
assert config.sliceable_layer_count == 1
|
|
|
|
def test_from_type_list(self):
|
|
"""Test from_type_list."""
|
|
type_names = ["KVCache", "KVCache", "ArraysCache", "KVCache"]
|
|
|
|
config = ModelCacheConfig.from_type_list(type_names, model_name="test")
|
|
|
|
assert config.num_layers == 4
|
|
assert config.is_hybrid is True
|
|
assert config.sliceable_layer_count == 3
|
|
|
|
def test_from_type_list_empty(self):
|
|
"""Test from_type_list with empty list."""
|
|
config = ModelCacheConfig.from_type_list([])
|
|
|
|
assert config.num_layers == 0
|
|
|
|
def test_get_sliceable_layers(self):
|
|
"""Test get_sliceable_layers."""
|
|
type_names = ["KVCache", "ArraysCache", "KVCache", "RotatingKVCache"]
|
|
config = ModelCacheConfig.from_type_list(type_names)
|
|
|
|
sliceable = config.get_sliceable_layers()
|
|
|
|
# KVCache layers only (indices 0 and 2)
|
|
assert 0 in sliceable
|
|
assert 2 in sliceable
|
|
assert 1 not in sliceable # ArraysCache
|
|
assert 3 not in sliceable # RotatingKVCache
|
|
|
|
def test_get_non_sliceable_layers(self):
|
|
"""Test get_non_sliceable_layers."""
|
|
type_names = ["KVCache", "ArraysCache", "KVCache", "ArraysCache"]
|
|
config = ModelCacheConfig.from_type_list(type_names)
|
|
|
|
non_sliceable = config.get_non_sliceable_layers()
|
|
|
|
assert 1 in non_sliceable
|
|
assert 3 in non_sliceable
|
|
assert len(non_sliceable) == 2
|
|
|
|
def test_get_layer_type(self):
|
|
"""Test get_layer_type."""
|
|
type_names = ["KVCache", "ArraysCache"]
|
|
config = ModelCacheConfig.from_type_list(type_names)
|
|
|
|
assert config.get_layer_type(0) == CacheType.KVCACHE
|
|
assert config.get_layer_type(1) == CacheType.ARRAYS_CACHE
|
|
# Out of bounds returns default
|
|
assert config.get_layer_type(99) == CacheType.KVCACHE
|
|
|
|
def test_get_layer_handler(self):
|
|
"""Test get_layer_handler."""
|
|
type_names = ["KVCache", "ArraysCache"]
|
|
config = ModelCacheConfig.from_type_list(type_names)
|
|
|
|
handler0 = config.get_layer_handler(0)
|
|
handler1 = config.get_layer_handler(1)
|
|
|
|
assert handler0.cache_type == CacheType.KVCACHE
|
|
assert handler1.cache_type == CacheType.ARRAYS_CACHE
|
|
|
|
def test_get_type_names(self):
|
|
"""Test get_type_names."""
|
|
type_names = ["KVCache", "ArraysCache", "RotatingKVCache"]
|
|
config = ModelCacheConfig.from_type_list(type_names)
|
|
|
|
names = config.get_type_names()
|
|
|
|
assert names == ["KVCache", "ArraysCache", "RotatingKVCache"]
|
|
|
|
def test_get_meta_states(self):
|
|
"""Test get_meta_states."""
|
|
mock_caches = []
|
|
for i in range(2):
|
|
mock = MagicMock()
|
|
mock.__class__.__name__ = "KVCache"
|
|
mock_keys = MagicMock()
|
|
mock_keys.shape = (1, 8, 64, 64)
|
|
mock.state = (mock_keys, MagicMock())
|
|
mock.offset = 64 + i * 10
|
|
mock_caches.append(mock)
|
|
|
|
config = ModelCacheConfig.from_cache_list(mock_caches)
|
|
meta_states = config.get_meta_states(mock_caches)
|
|
|
|
assert len(meta_states) == 2
|
|
# Each should have at least offset
|
|
for meta in meta_states:
|
|
assert len(meta) > 0
|
|
|
|
def test_supports_full_block_slicing(self):
|
|
"""Test supports_full_block_slicing."""
|
|
# All KVCache
|
|
config1 = ModelCacheConfig.from_type_list(["KVCache", "KVCache"])
|
|
assert config1.supports_full_block_slicing() is True
|
|
|
|
# Mixed
|
|
config2 = ModelCacheConfig.from_type_list(["KVCache", "ArraysCache"])
|
|
assert config2.supports_full_block_slicing() is False
|
|
|
|
def test_has_rotating_layers(self):
|
|
"""Test has_rotating_layers."""
|
|
config1 = ModelCacheConfig.from_type_list(["KVCache", "KVCache"])
|
|
assert config1.has_rotating_layers() is False
|
|
|
|
config2 = ModelCacheConfig.from_type_list(["KVCache", "RotatingKVCache"])
|
|
assert config2.has_rotating_layers() is True
|
|
|
|
def test_repr(self):
|
|
"""Test string representation."""
|
|
config = ModelCacheConfig.from_type_list(
|
|
["KVCache", "ArraysCache", "KVCache"],
|
|
model_name="test-model",
|
|
)
|
|
|
|
repr_str = repr(config)
|
|
|
|
assert "test-model" in repr_str
|
|
assert "layers=3" in repr_str
|
|
assert "hybrid=True" in repr_str
|
|
|
|
|
|
class TestCreateDefaultKVCacheConfig:
|
|
"""Tests for create_default_kvcache_config helper."""
|
|
|
|
def test_creates_all_kvcache(self):
|
|
"""Test creating all-KVCache config."""
|
|
config = create_default_kvcache_config(num_layers=32, model_name="llama")
|
|
|
|
assert config.num_layers == 32
|
|
assert config.model_name == "llama"
|
|
assert config.is_hybrid is False
|
|
assert config.sliceable_layer_count == 32
|
|
|
|
def test_all_layers_kvcache(self):
|
|
"""Test all layers are KVCache type."""
|
|
config = create_default_kvcache_config(num_layers=4)
|
|
|
|
for layer_config in config.layer_configs:
|
|
assert layer_config.cache_type == CacheType.KVCACHE
|
|
assert layer_config.supports_block_slicing is True
|
|
assert layer_config.class_name == "KVCache"
|
|
|
|
def test_layer_indices_correct(self):
|
|
"""Test layer indices are sequential."""
|
|
config = create_default_kvcache_config(num_layers=4)
|
|
|
|
for i, layer_config in enumerate(config.layer_configs):
|
|
assert layer_config.layer_idx == i
|
|
|
|
def test_zero_layers(self):
|
|
"""Test with zero layers."""
|
|
config = create_default_kvcache_config(num_layers=0)
|
|
|
|
assert config.num_layers == 0
|
|
assert len(config.layer_configs) == 0
|