@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 224 lines 7.67 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import patch
5
6 import pytest
7
8 from vllm.v1.core.kv_cache_metrics import (
9 BlockMetricsState,
10 KVCacheMetricsCollector,
11 )
12 from vllm.v1.core.kv_cache_utils import KVCacheBlock
13
14
15 class TestBlockMetricsState:
16 def test_init(self):
17 with patch("time.monotonic_ns", return_value=1000000000):
18 state = BlockMetricsState()
19 assert state.birth_time_ns == 1000000000
20 assert state.last_access_ns == 1000000000
21 assert len(state.access_history) == 0
22
23 def test_access_tracking(self):
24 with patch("time.monotonic_ns", return_value=1000000000):
25 state = BlockMetricsState()
26
27 with patch("time.monotonic_ns", return_value=2000000000):
28 state.record_access()
29
30 assert state.last_access_ns == 2000000000
31 assert list(state.access_history) == [2000000000]
32
33 def test_ring_buffer_wraps_at_4(self):
34 with patch("time.monotonic_ns", return_value=1000000000):
35 state = BlockMetricsState()
36
37 for i in range(5):
38 t = 1000000000 + (i + 1) * 1000000000
39 with patch("time.monotonic_ns", return_value=t):
40 state.record_access()
41
42 assert len(state.access_history) == 4
43 assert list(state.access_history) == [
44 3000000000,
45 4000000000,
46 5000000000,
47 6000000000,
48 ]
49
50 def test_lifetime(self):
51 with patch("time.monotonic_ns", return_value=1000000000):
52 state = BlockMetricsState()
53 with patch("time.monotonic_ns", return_value=6500000000):
54 assert abs(state.get_lifetime_seconds() - 5.5) < 0.001
55
56 def test_idle_time(self):
57 with patch("time.monotonic_ns", return_value=1000000000):
58 state = BlockMetricsState()
59 state.last_access_ns = 2000000000
60 with patch("time.monotonic_ns", return_value=5200000000):
61 assert abs(state.get_idle_time_seconds() - 3.2) < 0.001
62
63 def test_reuse_gaps(self):
64 with patch("time.monotonic_ns", return_value=1000000000):
65 state = BlockMetricsState()
66
67 base = 1000000000
68 for offset in [0, 1.5, 3.0, 5.5]:
69 state.access_history.append(base + int(offset * 1e9))
70
71 gaps = state.get_reuse_gaps_seconds()
72 assert len(gaps) == 3
73 assert gaps[0] == 1.5 and gaps[1] == 1.5 and gaps[2] == 2.5
74
75 def test_ring_wrap_only_gives_3_gaps(self):
76 # 5 accesses in size-4 buffer = 3 gaps
77 with patch("time.monotonic_ns", return_value=1000000000):
78 state = BlockMetricsState()
79
80 for i in range(5):
81 state.access_history.append(1000000000 + i * 1000000000)
82
83 assert len(state.get_reuse_gaps_seconds()) == 3
84
85
86 class TestKVCacheMetricsCollector:
87 def test_sample_rate_validation(self):
88 with pytest.raises(AssertionError):
89 KVCacheMetricsCollector(sample_rate=-0.1)
90 with pytest.raises(AssertionError):
91 KVCacheMetricsCollector(sample_rate=1.5)
92 with pytest.raises(AssertionError):
93 KVCacheMetricsCollector(sample_rate=0.0)
94
95 def test_sampling(self):
96 c = KVCacheMetricsCollector(sample_rate=1.0)
97 assert sum(1 for _ in range(100) if c.should_sample_block()) == 100
98
99 c = KVCacheMetricsCollector(sample_rate=0.5)
100 samples = sum(1 for _ in range(1000) if c.should_sample_block())
101 assert 400 < samples < 600
102
103 def test_alloc(self):
104 c = KVCacheMetricsCollector(sample_rate=1.0)
105
106 blocks = [KVCacheBlock(block_id=i) for i in range(5)]
107 with patch("time.monotonic_ns", return_value=1000000000):
108 for block in blocks:
109 c.on_block_allocated(block)
110
111 assert len(c.block_metrics) == 5
112
113 def test_access(self):
114 c = KVCacheMetricsCollector(sample_rate=1.0)
115 block = KVCacheBlock(block_id=0)
116
117 with patch("time.monotonic_ns", return_value=1000000000):
118 c.on_block_allocated(block)
119
120 for i in range(3):
121 t = 1000000000 + (i + 1) * 1000000000
122 with patch("time.monotonic_ns", return_value=t):
123 c.on_block_accessed(block)
124
125 assert len(c.block_metrics[0].access_history) == 3
126
127 def test_evict_no_accesses(self):
128 # lifetime should equal idle if never accessed
129 c = KVCacheMetricsCollector(sample_rate=1.0)
130
131 block = KVCacheBlock(block_id=0)
132 with patch("time.monotonic_ns", return_value=1000000000):
133 c.on_block_allocated(block)
134
135 with patch("time.monotonic_ns", return_value=6000000000):
136 c.on_block_evicted(block)
137
138 events = c.drain_events()
139 assert len(events) == 1
140 assert abs(events[0].lifetime_seconds - 5.0) < 0.001
141 assert abs(events[0].idle_seconds - 5.0) < 0.001
142
143 def test_evict(self):
144 c = KVCacheMetricsCollector(sample_rate=1.0)
145
146 block = KVCacheBlock(block_id=0)
147 with patch("time.monotonic_ns", return_value=1000000000):
148 c.on_block_allocated(block)
149
150 with patch("time.monotonic_ns", return_value=2000000000):
151 c.on_block_accessed(block)
152 with patch("time.monotonic_ns", return_value=3000000000):
153 c.on_block_accessed(block)
154
155 with patch("time.monotonic_ns", return_value=4000000000):
156 c.on_block_evicted(block)
157
158 events = c.drain_events()
159 assert len(events) == 1
160 sample = events[0]
161 assert abs(sample.lifetime_seconds - 3.0) < 0.001
162 assert abs(sample.idle_seconds - 1.0) < 0.001
163 assert sample.reuse_gaps_seconds == (1.0,)
164 assert 0 not in c.block_metrics
165
166 def test_reset(self):
167 c = KVCacheMetricsCollector(sample_rate=1.0)
168
169 with patch("time.monotonic_ns", return_value=1000000000):
170 for i in range(5):
171 c.on_block_allocated(KVCacheBlock(block_id=i))
172
173 assert len(c.block_metrics) == 5
174 c.reset()
175 assert len(c.block_metrics) == 0
176
177 with patch("time.monotonic_ns", return_value=2000000000):
178 c.on_block_allocated(KVCacheBlock(block_id=10))
179 assert 10 in c.block_metrics
180
181 def test_huge_time_jump(self):
182 c = KVCacheMetricsCollector(sample_rate=1.0)
183
184 block = KVCacheBlock(block_id=0)
185 with patch("time.monotonic_ns", return_value=1000000000):
186 c.on_block_allocated(block)
187
188 with patch("time.monotonic_ns", return_value=9999999999999999):
189 c.on_block_evicted(block)
190
191 events = c.drain_events()
192 assert len(events) == 1
193 assert events[0].lifetime_seconds > 0
194
195
196 def test_kv_cache_metrics_collector_smoke() -> None:
197 """Simple smoke test for KVCacheMetricsCollector on CPU."""
198 collector = KVCacheMetricsCollector(sample_rate=1.0)
199 block = KVCacheBlock(block_id=123)
200
201 # Allocate at t = 1.0s.
202 with patch("time.monotonic_ns", return_value=1_000_000_000):
203 collector.on_block_allocated(block)
204
205 # Access at t = 2.0s and t = 3.0s.
206 with patch("time.monotonic_ns", return_value=2_000_000_000):
207 collector.on_block_accessed(block)
208 with patch("time.monotonic_ns", return_value=3_000_000_000):
209 collector.on_block_accessed(block)
210
211 # Evict at t = 4.0s.
212 with patch("time.monotonic_ns", return_value=4_000_000_000):
213 collector.on_block_evicted(block)
214
215 events = collector.drain_events()
216 assert len(events) == 1
217
218 event = events[0]
219 # Lifetime: 1.0s → 4.0s.
220 assert abs(event.lifetime_seconds - 3.0) < 1e-6
221 # Idle: last access at 3.0s, evicted at 4.0s.
222 assert abs(event.idle_seconds - 1.0) < 1e-6
223 # One reuse gap between the two accesses.
224 assert event.reuse_gaps_seconds == (1.0,)