Skip to content

Commit c0a7b89

Browse files
[Misc] Move LRUCache into its own file (#26342)
Signed-off-by: DarkLight1337 <[email protected]>
1 parent 6f59bea commit c0a7b89

File tree

6 files changed

+349
-378
lines changed

6 files changed

+349
-378
lines changed

tests/utils_/test_cache.py

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3+
from vllm.utils.cache import CacheInfo, LRUCache
4+
5+
6+
class TestLRUCache(LRUCache):
7+
def _on_remove(self, key, value):
8+
if not hasattr(self, "_remove_counter"):
9+
self._remove_counter = 0
10+
self._remove_counter += 1
11+
12+
13+
def test_lru_cache():
14+
cache = TestLRUCache(3)
15+
assert cache.stat() == CacheInfo(hits=0, total=0)
16+
assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
17+
18+
cache.put(1, 1)
19+
assert len(cache) == 1
20+
21+
cache.put(1, 1)
22+
assert len(cache) == 1
23+
24+
cache.put(2, 2)
25+
assert len(cache) == 2
26+
27+
cache.put(3, 3)
28+
assert len(cache) == 3
29+
assert set(cache.cache) == {1, 2, 3}
30+
31+
cache.put(4, 4)
32+
assert len(cache) == 3
33+
assert set(cache.cache) == {2, 3, 4}
34+
assert cache._remove_counter == 1
35+
36+
assert cache.get(2) == 2
37+
assert cache.stat() == CacheInfo(hits=1, total=1)
38+
assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
39+
40+
assert cache[2] == 2
41+
assert cache.stat() == CacheInfo(hits=2, total=2)
42+
assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
43+
44+
cache.put(5, 5)
45+
assert set(cache.cache) == {2, 4, 5}
46+
assert cache._remove_counter == 2
47+
48+
assert cache.pop(5) == 5
49+
assert len(cache) == 2
50+
assert set(cache.cache) == {2, 4}
51+
assert cache._remove_counter == 3
52+
53+
assert cache.get(-1) is None
54+
assert cache.stat() == CacheInfo(hits=2, total=3)
55+
assert cache.stat(delta=True) == CacheInfo(hits=0, total=1)
56+
57+
cache.pop(10)
58+
assert len(cache) == 2
59+
assert set(cache.cache) == {2, 4}
60+
assert cache._remove_counter == 3
61+
62+
cache.get(10)
63+
assert len(cache) == 2
64+
assert set(cache.cache) == {2, 4}
65+
assert cache._remove_counter == 3
66+
67+
cache.put(6, 6)
68+
assert len(cache) == 3
69+
assert set(cache.cache) == {2, 4, 6}
70+
assert 2 in cache
71+
assert 4 in cache
72+
assert 6 in cache
73+
74+
cache.remove_oldest()
75+
assert len(cache) == 2
76+
assert set(cache.cache) == {2, 6}
77+
assert cache._remove_counter == 4
78+
79+
cache.clear()
80+
assert len(cache) == 0
81+
assert cache._remove_counter == 6
82+
assert cache.stat() == CacheInfo(hits=0, total=0)
83+
assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
84+
85+
cache._remove_counter = 0
86+
87+
cache[1] = 1
88+
assert len(cache) == 1
89+
90+
cache[1] = 1
91+
assert len(cache) == 1
92+
93+
cache[2] = 2
94+
assert len(cache) == 2
95+
96+
cache[3] = 3
97+
assert len(cache) == 3
98+
assert set(cache.cache) == {1, 2, 3}
99+
100+
cache[4] = 4
101+
assert len(cache) == 3
102+
assert set(cache.cache) == {2, 3, 4}
103+
assert cache._remove_counter == 1
104+
assert cache[2] == 2
105+
106+
cache[5] = 5
107+
assert set(cache.cache) == {2, 4, 5}
108+
assert cache._remove_counter == 2
109+
110+
del cache[5]
111+
assert len(cache) == 2
112+
assert set(cache.cache) == {2, 4}
113+
assert cache._remove_counter == 3
114+
115+
cache.pop(10)
116+
assert len(cache) == 2
117+
assert set(cache.cache) == {2, 4}
118+
assert cache._remove_counter == 3
119+
120+
cache[6] = 6
121+
assert len(cache) == 3
122+
assert set(cache.cache) == {2, 4, 6}
123+
assert 2 in cache
124+
assert 4 in cache
125+
assert 6 in cache

tests/utils_/test_utils.py

Lines changed: 0 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,8 @@
2323
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
2424
from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
2525

26-
# isort: off
2726
from vllm.utils import (
28-
CacheInfo,
2927
FlexibleArgumentParser,
30-
LRUCache,
3128
MemorySnapshot,
3229
PlaceholderModule,
3330
bind_kv_cache,
@@ -50,7 +47,6 @@
5047
unique_filepath,
5148
)
5249

53-
# isort: on
5450
from ..utils import create_new_process_for_each_test, error_on_warning
5551

5652

@@ -557,128 +553,6 @@ def test_bind_kv_cache_pp():
557553
assert ctx["layers.0.self_attn"].kv_cache[1] is kv_cache[1][0]
558554

559555

560-
class TestLRUCache(LRUCache):
561-
def _on_remove(self, key, value):
562-
if not hasattr(self, "_remove_counter"):
563-
self._remove_counter = 0
564-
self._remove_counter += 1
565-
566-
567-
def test_lru_cache():
568-
cache = TestLRUCache(3)
569-
assert cache.stat() == CacheInfo(hits=0, total=0)
570-
assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
571-
572-
cache.put(1, 1)
573-
assert len(cache) == 1
574-
575-
cache.put(1, 1)
576-
assert len(cache) == 1
577-
578-
cache.put(2, 2)
579-
assert len(cache) == 2
580-
581-
cache.put(3, 3)
582-
assert len(cache) == 3
583-
assert set(cache.cache) == {1, 2, 3}
584-
585-
cache.put(4, 4)
586-
assert len(cache) == 3
587-
assert set(cache.cache) == {2, 3, 4}
588-
assert cache._remove_counter == 1
589-
590-
assert cache.get(2) == 2
591-
assert cache.stat() == CacheInfo(hits=1, total=1)
592-
assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
593-
594-
assert cache[2] == 2
595-
assert cache.stat() == CacheInfo(hits=2, total=2)
596-
assert cache.stat(delta=True) == CacheInfo(hits=1, total=1)
597-
598-
cache.put(5, 5)
599-
assert set(cache.cache) == {2, 4, 5}
600-
assert cache._remove_counter == 2
601-
602-
assert cache.pop(5) == 5
603-
assert len(cache) == 2
604-
assert set(cache.cache) == {2, 4}
605-
assert cache._remove_counter == 3
606-
607-
assert cache.get(-1) is None
608-
assert cache.stat() == CacheInfo(hits=2, total=3)
609-
assert cache.stat(delta=True) == CacheInfo(hits=0, total=1)
610-
611-
cache.pop(10)
612-
assert len(cache) == 2
613-
assert set(cache.cache) == {2, 4}
614-
assert cache._remove_counter == 3
615-
616-
cache.get(10)
617-
assert len(cache) == 2
618-
assert set(cache.cache) == {2, 4}
619-
assert cache._remove_counter == 3
620-
621-
cache.put(6, 6)
622-
assert len(cache) == 3
623-
assert set(cache.cache) == {2, 4, 6}
624-
assert 2 in cache
625-
assert 4 in cache
626-
assert 6 in cache
627-
628-
cache.remove_oldest()
629-
assert len(cache) == 2
630-
assert set(cache.cache) == {2, 6}
631-
assert cache._remove_counter == 4
632-
633-
cache.clear()
634-
assert len(cache) == 0
635-
assert cache._remove_counter == 6
636-
assert cache.stat() == CacheInfo(hits=0, total=0)
637-
assert cache.stat(delta=True) == CacheInfo(hits=0, total=0)
638-
639-
cache._remove_counter = 0
640-
641-
cache[1] = 1
642-
assert len(cache) == 1
643-
644-
cache[1] = 1
645-
assert len(cache) == 1
646-
647-
cache[2] = 2
648-
assert len(cache) == 2
649-
650-
cache[3] = 3
651-
assert len(cache) == 3
652-
assert set(cache.cache) == {1, 2, 3}
653-
654-
cache[4] = 4
655-
assert len(cache) == 3
656-
assert set(cache.cache) == {2, 3, 4}
657-
assert cache._remove_counter == 1
658-
assert cache[2] == 2
659-
660-
cache[5] = 5
661-
assert set(cache.cache) == {2, 4, 5}
662-
assert cache._remove_counter == 2
663-
664-
del cache[5]
665-
assert len(cache) == 2
666-
assert set(cache.cache) == {2, 4}
667-
assert cache._remove_counter == 3
668-
669-
cache.pop(10)
670-
assert len(cache) == 2
671-
assert set(cache.cache) == {2, 4}
672-
assert cache._remove_counter == 3
673-
674-
cache[6] = 6
675-
assert len(cache) == 3
676-
assert set(cache.cache) == {2, 4, 6}
677-
assert 2 in cache
678-
assert 4 in cache
679-
assert 6 in cache
680-
681-
682556
@pytest.mark.parametrize(
683557
("src_dtype", "tgt_dtype", "expected_result"),
684558
[

vllm/lora/models.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@
3232
from vllm.model_executor.models.module_mapping import MultiModelKeys
3333
from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper
3434
from vllm.model_executor.utils import get_packed_modules_mapping
35-
from vllm.utils import LRUCache, is_pin_memory_available
35+
from vllm.utils import is_pin_memory_available
36+
from vllm.utils.cache import LRUCache
3637

3738
logger = init_logger(__name__)
3839

vllm/multimodal/cache.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@
1717
)
1818
from vllm.envs import VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME
1919
from vllm.logger import init_logger
20-
from vllm.utils import GiB_bytes, LRUCache, MiB_bytes
20+
from vllm.utils import GiB_bytes, MiB_bytes
21+
from vllm.utils.cache import LRUCache
2122
from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves
2223

2324
from .inputs import (

0 commit comments

Comments
 (0)