wehub-resource-sync
197 lines
6.6 KiB
Python
197 lines
6.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for IndexCache monkey-patch."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
class TestBuildLayerPattern:
|
|
"""Test _build_layer_pattern function."""
|
|
|
|
def test_freq_2(self):
|
|
from omlx.patches.index_cache import _build_layer_pattern
|
|
|
|
pattern = _build_layer_pattern(8, 2)
|
|
assert pattern == [True, False, True, False, True, False, True, False]
|
|
|
|
def test_freq_4(self):
|
|
from omlx.patches.index_cache import _build_layer_pattern
|
|
|
|
pattern = _build_layer_pattern(8, 4)
|
|
assert pattern == [True, False, False, False, True, False, False, False]
|
|
|
|
def test_layer_0_always_full(self):
|
|
from omlx.patches.index_cache import _build_layer_pattern
|
|
|
|
for freq in range(2, 10):
|
|
pattern = _build_layer_pattern(60, freq)
|
|
assert pattern[0] is True
|
|
|
|
def test_freq_4_counts(self):
|
|
from omlx.patches.index_cache import _build_layer_pattern
|
|
|
|
pattern = _build_layer_pattern(60, 4)
|
|
full_count = sum(pattern)
|
|
shared_count = len(pattern) - full_count
|
|
assert full_count == 15
|
|
assert shared_count == 45
|
|
|
|
def test_freq_larger_than_layers(self):
|
|
from omlx.patches.index_cache import _build_layer_pattern
|
|
|
|
pattern = _build_layer_pattern(3, 8)
|
|
# Only layer 0 is Full
|
|
assert pattern == [True, False, False]
|
|
|
|
|
|
class TestGetModelType:
|
|
"""Test _get_model_type function."""
|
|
|
|
def test_model_type_attribute(self):
|
|
from omlx.patches.index_cache import _get_model_type
|
|
|
|
model = MagicMock(spec=[])
|
|
model.model_type = "deepseek_v32"
|
|
assert _get_model_type(model) == "deepseek_v32"
|
|
|
|
def test_args_model_type(self):
|
|
from omlx.patches.index_cache import _get_model_type
|
|
|
|
model = MagicMock(spec=[])
|
|
model.args = MagicMock(spec=[])
|
|
model.args.model_type = "glm_moe_dsa"
|
|
assert _get_model_type(model) == "glm_moe_dsa"
|
|
|
|
def test_no_model_type(self):
|
|
from omlx.patches.index_cache import _get_model_type
|
|
|
|
model = MagicMock(spec=[])
|
|
assert _get_model_type(model) is None
|
|
|
|
|
|
class TestApplyIndexCache:
|
|
"""Test apply_index_cache function."""
|
|
|
|
def test_unsupported_model_returns_false(self):
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
model = MagicMock(spec=[])
|
|
model.model_type = "llama"
|
|
assert apply_index_cache(model, 4) is False
|
|
|
|
def test_no_model_type_returns_false(self):
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
model = MagicMock(spec=[])
|
|
assert apply_index_cache(model, 4) is False
|
|
|
|
def test_glm_moe_dsa_returns_false(self):
|
|
"""GLM-5.2 uses its native indexer_types schedule, not IndexCache."""
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
model = MagicMock(spec=[])
|
|
model.args = MagicMock(spec=[])
|
|
model.args.model_type = "glm_moe_dsa"
|
|
assert apply_index_cache(model, 4) is False
|
|
|
|
def test_freq_less_than_2_returns_false(self):
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
model = MagicMock(spec=[])
|
|
model.model_type = "deepseek_v32"
|
|
assert apply_index_cache(model, 1) is False
|
|
|
|
@patch("omlx.patches.index_cache._class_patch_applied", True)
|
|
def test_applies_flags_to_layers(self):
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
# Create a mock model with 4 layers
|
|
model = MagicMock(spec=[])
|
|
model.model_type = "deepseek_v32"
|
|
model.args = MagicMock(spec=[])
|
|
model.args.model_type = "deepseek_v32"
|
|
|
|
layers = []
|
|
for _ in range(4):
|
|
layer = MagicMock(spec=[])
|
|
layer.self_attn = MagicMock(spec=[])
|
|
layers.append(layer)
|
|
|
|
model.model = MagicMock(spec=[])
|
|
model.model.layers = layers
|
|
|
|
result = apply_index_cache(model, 2)
|
|
assert result is True
|
|
|
|
# Check flags on each layer
|
|
assert layers[0].self_attn._ic_is_full is True # layer 0: Full
|
|
assert layers[1].self_attn._ic_is_full is False # layer 1: Shared
|
|
assert layers[2].self_attn._ic_is_full is True # layer 2: Full
|
|
assert layers[3].self_attn._ic_is_full is False # layer 3: Shared
|
|
|
|
# Check shared state
|
|
assert hasattr(model.model, "_index_cache_state")
|
|
assert model.model._index_cache_state["last_topk_indices"] is None
|
|
|
|
@patch("omlx.patches.index_cache._class_patch_applied", True)
|
|
def test_skips_none_layers(self):
|
|
"""None layers (pipeline parallel placeholders) should be skipped."""
|
|
from omlx.patches.index_cache import apply_index_cache
|
|
|
|
model = MagicMock(spec=[])
|
|
model.model_type = "deepseek_v32"
|
|
model.args = MagicMock(spec=[])
|
|
model.args.model_type = "deepseek_v32"
|
|
|
|
layer0 = MagicMock(spec=[])
|
|
layer0.self_attn = MagicMock(spec=[])
|
|
|
|
model.model = MagicMock(spec=[])
|
|
model.model.layers = [layer0, None, None]
|
|
|
|
result = apply_index_cache(model, 2)
|
|
assert result is True
|
|
assert layer0.self_attn._ic_is_full is True
|
|
|
|
|
|
class TestApplyPostLoadTransforms:
|
|
"""Test the centralized transform entry point."""
|
|
|
|
def test_none_settings_returns_model(self):
|
|
from omlx.utils.model_loading import apply_post_load_transforms
|
|
|
|
model = MagicMock()
|
|
result = apply_post_load_transforms(model, None)
|
|
assert result is model
|
|
|
|
def test_no_index_cache_freq_returns_model(self):
|
|
from omlx.utils.model_loading import apply_post_load_transforms
|
|
|
|
model = MagicMock()
|
|
settings = MagicMock(spec=[])
|
|
settings.index_cache_freq = None
|
|
result = apply_post_load_transforms(model, settings)
|
|
assert result is model
|
|
|
|
@patch("omlx.patches.index_cache.apply_index_cache")
|
|
def test_calls_apply_index_cache(self, mock_apply):
|
|
from omlx.utils.model_loading import apply_post_load_transforms
|
|
|
|
mock_apply.return_value = True
|
|
model = MagicMock()
|
|
settings = MagicMock(spec=[])
|
|
settings.index_cache_freq = 4
|
|
result = apply_post_load_transforms(model, settings)
|
|
mock_apply.assert_called_once_with(model, 4)
|
|
assert result is model
|
|
|
|
@patch("omlx.patches.index_cache.apply_index_cache")
|
|
def test_freq_1_skipped(self, mock_apply):
|
|
from omlx.utils.model_loading import apply_post_load_transforms
|
|
|
|
model = MagicMock()
|
|
settings = MagicMock(spec=[])
|
|
settings.index_cache_freq = 1
|
|
result = apply_post_load_transforms(model, settings)
|
|
mock_apply.assert_not_called()
|
|
assert result is model
|