@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 140 lines 3.45 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 from vllm.utils.cache import CacheInfo, LRUCache
6
7
8 class TestLRUCache(LRUCache):
9 def _on_remove(self, key, value):
10 if not hasattr(self, "_remove_counter"):
11 self._remove_counter = 0
12 self._remove_counter += 1
13
14
15 def test_lru_cache():
16 cache = TestLRUCache(3)
17 assert cache.stat() == CacheInfo(hits=0, total=0)
18 assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
19
20 cache.put(1, 1)
21 assert len(cache) == 1
22
23 cache.put(1, 1)
24 assert len(cache) == 1
25
26 cache.put(2, 2)
27 assert len(cache) == 2
28
29 cache.put(3, 3)
30 assert len(cache) == 3
31 assert set(cache.cache) == {1, 2, 3}
32
33 cache.put(4, 4)
34 assert len(cache) == 3
35 assert set(cache.cache) == {2, 3, 4}
36 assert cache._remove_counter == 1
37
38 assert cache.get(2) == 2
39 assert cache.stat() == CacheInfo(hits=1, total=1)
40 assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
41
42 assert cache[2] == 2
43 assert cache.stat() == CacheInfo(hits=2, total=2)
44 assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
45
46 cache.put(5, 5)
47 assert set(cache.cache) == {2, 4, 5}
48 assert cache._remove_counter == 2
49
50 assert cache.pop(5) == 5
51 assert len(cache) == 2
52 assert set(cache.cache) == {2, 4}
53 assert cache._remove_counter == 3
54
55 assert cache.get(-1) is None
56 assert cache.stat() == CacheInfo(hits=2, total=3)
57 assert cache.stat(delta=True) == CacheInfo(hits=0, total=1)
58
59 cache.pop(10)
60 assert len(cache) == 2
61 assert set(cache.cache) == {2, 4}
62 assert cache._remove_counter == 3
63
64 cache.get(10)
65 assert len(cache) == 2
66 assert set(cache.cache) == {2, 4}
67 assert cache._remove_counter == 3
68
69 cache.put(6, 6)
70 assert len(cache) == 3
71 assert set(cache.cache) == {2, 4, 6}
72 assert 2 in cache
73 assert 4 in cache
74 assert 6 in cache
75
76 cache.remove_oldest()
77 assert len(cache) == 2
78 assert set(cache.cache) == {2, 6}
79 assert cache._remove_counter == 4
80
81 cache.clear()
82 assert len(cache) == 0
83 assert cache._remove_counter == 6
84 assert cache.stat() == CacheInfo(hits=0, total=0)
85 assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
86
87 cache._remove_counter = 0
88
89 cache[1] = 1
90 assert len(cache) == 1
91
92 cache[1] = 1
93 assert len(cache) == 1
94
95 cache[2] = 2
96 assert len(cache) == 2
97
98 cache[3] = 3
99 assert len(cache) == 3
100 assert set(cache.cache) == {1, 2, 3}
101
102 cache[4] = 4
103 assert len(cache) == 3
104 assert set(cache.cache) == {2, 3, 4}
105 assert cache._remove_counter == 1
106 assert cache[2] == 2
107
108 cache[5] = 5
109 assert set(cache.cache) == {2, 4, 5}
110 assert cache._remove_counter == 2
111
112 del cache[5]
113 assert len(cache) == 2
114 assert set(cache.cache) == {2, 4}
115 assert cache._remove_counter == 3
116
117 cache.pop(10)
118 assert len(cache) == 2
119 assert set(cache.cache) == {2, 4}
120 assert cache._remove_counter == 3
121
122 cache[6] = 6
123 assert len(cache) == 3
124 assert set(cache.cache) == {2, 4, 6}
125 assert 2 in cache
126 assert 4 in cache
127 assert 6 in cache
128
129
130 def test_lru_cache_put_if_fits():
131 cache = LRUCache[str, int](10, getsizeof=lambda x: x)
132
133 assert cache.put_if_fits("ok", 4) is True
134 assert cache["ok"] == 4
135
136 assert cache.put_if_fits("big", 11) is False
137 assert "big" not in cache
138
139 with pytest.raises(ValueError, match="value too large"):
140 cache.put("big", 11)