@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 204 lines 6.89 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for SecondaryTierFactory.
4
5 These tests verify:
6 1. Pre-registration integrity — registered tier module paths can import
7 and yield correct SecondaryTierManager subclasses (CI sentinel).
8 2. Multi-tier creation via factory with correct tier_type propagation.
9 3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
10 """
11
12 from importlib.util import find_spec
13 from unittest.mock import MagicMock
14
15 import pytest
16
17 from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
18 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
19 from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
20
21 # ---------------------------------------------------------------------------
22 # Fixtures
23 # ---------------------------------------------------------------------------
24
25
26 @pytest.fixture(autouse=True)
27 def restore_registry():
28 """Save and restore SecondaryTierFactory._registry between tests."""
29 original = dict(SecondaryTierFactory._registry)
30 yield
31 SecondaryTierFactory._registry = original
32
33
34 def _make_mock_args():
35 """Build common mock args for create_secondary_tier."""
36 return MagicMock(), MagicMock() # primary_kv_view, offloading_spec
37
38
39 # ---------------------------------------------------------------------------
40 # Pre-registration integrity (CI sentinel)
41 # ---------------------------------------------------------------------------
42
43
44 def test_pre_registered_tiers_can_be_imported():
45 """CI sentinel: registered tiers import and yield SecondaryTierManager."""
46 for tier_type in SecondaryTierFactory._registry:
47 # KVCR is an optional external dependency and may not be installed.
48 if tier_type == "kvcr" and find_spec("kvcr") is None:
49 continue
50 cls = SecondaryTierFactory._registry[tier_type]()
51 assert issubclass(cls, SecondaryTierManager)
52
53
54 def test_example_tier_registered():
55 """Example tier is registered."""
56 cls = SecondaryTierFactory._registry["example"]()
57 assert cls is ExampleSecondaryTierManager
58
59
60 # ---------------------------------------------------------------------------
61 # Normal path — create_secondary_tier
62 # ---------------------------------------------------------------------------
63
64
65 def test_create_tier_from_registry():
66 """Registered tier_type creates instance with correct tier_type."""
67 primary_kv_view, offloading_spec = _make_mock_args()
68 tier_config = {"type": "example"}
69
70 tier = SecondaryTierFactory.create_secondary_tier(
71 tier_config, primary_kv_view, offloading_spec
72 )
73
74 assert isinstance(tier, SecondaryTierManager)
75 assert tier.tier_type == "example"
76
77
78 def test_create_multiple_tiers():
79 """Multiple tier configs can be created with correct tier_types."""
80 primary_kv_view, offloading_spec = _make_mock_args()
81 configs = [
82 {"type": "example", "custom_param": 1},
83 {"type": "example", "custom_param": 2},
84 ]
85
86 tiers = [
87 SecondaryTierFactory.create_secondary_tier(
88 cfg.copy(), primary_kv_view, offloading_spec
89 )
90 for cfg in configs
91 ]
92
93 assert len(tiers) == 2
94 assert all(tier.tier_type == "example" for tier in tiers)
95 assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers)
96
97
98 def test_register_new_tier_type():
99 """Verify that new tier types can be registered and created.
100
101 This is how external projects add custom secondary tiers
102 (e.g., llm-d FS backend was upstreamed as "fs" tier via this mechanism).
103 """
104 # Register a new tier type (reuse example manager for simplicity)
105 SecondaryTierFactory.register_tier(
106 "custom_tier",
107 "vllm.v1.kv_offload.tiering.example.manager",
108 "ExampleSecondaryTierManager",
109 )
110
111 primary_kv_view, offloading_spec = _make_mock_args()
112 tier = SecondaryTierFactory.create_secondary_tier(
113 {"type": "custom_tier", "custom_param": 99},
114 primary_kv_view,
115 offloading_spec,
116 )
117
118 assert tier.tier_type == "custom_tier"
119 assert isinstance(tier, ExampleSecondaryTierManager)
120
121
122 # ---------------------------------------------------------------------------
123 # Error paths
124 # ---------------------------------------------------------------------------
125
126
127 def test_missing_tier_type_raises():
128 """tier_config without 'type' → ValueError."""
129 primary_kv_view, offloading_spec = _make_mock_args()
130 tier_config: dict[str, str] = {}
131
132 with pytest.raises(ValueError, match="must include 'type'"):
133 SecondaryTierFactory.create_secondary_tier(
134 tier_config, primary_kv_view, offloading_spec
135 )
136
137
138 def test_unknown_tier_type_raises():
139 """Unrecognized tier_type without module_path → ValueError."""
140 primary_kv_view, offloading_spec = _make_mock_args()
141 tier_config = {"type": "nonexistent_tier"}
142
143 with pytest.raises(
144 ValueError,
145 match=r"Unknown secondary tier type.*also set 'module_path'",
146 ):
147 SecondaryTierFactory.create_secondary_tier(
148 tier_config, primary_kv_view, offloading_spec
149 )
150
151
152 def test_duplicate_registration_raises():
153 """register_tier with existing type → ValueError."""
154 with pytest.raises(ValueError, match="is already registered"):
155 SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")
156
157
158 # ---------------------------------------------------------------------------
159 # Out-of-tree dynamic loading via module_path
160 # ---------------------------------------------------------------------------
161
162
163 def test_create_tier_from_module_path():
164 """Full end-to-end: create_secondary_tier with module_path."""
165 assert "ExampleSecondaryTierManager" not in SecondaryTierFactory._registry
166
167 primary_kv_view, offloading_spec = _make_mock_args()
168 tier_config = {
169 "type": "ExampleSecondaryTierManager",
170 "module_path": "vllm.v1.kv_offload.tiering.example.manager",
171 "custom_param": 42,
172 }
173
174 tier = SecondaryTierFactory.create_secondary_tier(
175 tier_config, primary_kv_view, offloading_spec
176 )
177
178 assert isinstance(tier, ExampleSecondaryTierManager)
179 assert tier.tier_type == "ExampleSecondaryTierManager"
180
181
182 def test_module_path_invalid_module_raises():
183 """Non-existent module_path → ModuleNotFoundError."""
184 primary_kv_view, offloading_spec = _make_mock_args()
185 tier_config = {
186 "type": "SomeTier",
187 "module_path": "nonexistent.module.path",
188 }
189
190 with pytest.raises(ModuleNotFoundError):
191 SecondaryTierFactory.create_secondary_tier(
192 tier_config, primary_kv_view, offloading_spec
193 )
194
195
196 def test_module_path_not_subclass_raises():
197 """module_path resolving to a non-SecondaryTierManager → AssertionError."""
198 tier_config = {
199 "type": "MagicMock",
200 "module_path": "unittest.mock",
201 }
202
203 with pytest.raises(AssertionError):
204 SecondaryTierFactory.get_tier_class(tier_config)