1"""CPU-only checks against the installed vLLM scheduler and shared-memory tier."""
2import mmap
3import os
4from pathlib import Path
5import tempfile
6import unittest
7import uuid
8
9from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext, make_offload_key
10from vllm.v1.kv_offload.config import OffloadingCacheConfig, OffloadingConfig, OffloadingGroupConfig, OffloadingModelConfig, OffloadingParallelConfig
11from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
12from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
13from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
14from vllm.v1.kv_offload.tiering.manager import CPUPrimaryTierOffloadingManager, TieringOffloadingManager
15
16from studio_ai_cache_vllm import register
17
18
19class VllmCacheTests(unittest.TestCase):
20 @classmethod
21 def setUpClass(cls):
22 register()
23
24 def setUp(self):
25 self.temporary = tempfile.TemporaryDirectory()
26 self.addCleanup(self.temporary.cleanup)
27 self.root = Path(self.temporary.name)
28 os.environ["STUDIO_AI_CACHE_KEY"] = "disposable-synthetic-cache-key"
29 config = OffloadingConfig(
30 groups=(OffloadingGroupConfig(16, ("attention",)), OffloadingGroupConfig(16, ("recurrent",))),
31 worker_kv_bytes_per_block=mmap.PAGESIZE,
32 enable_kv_cache_events=False, extra_config={"cpu_bytes_to_use": 4 * mmap.PAGESIZE},
33 engine_id="synthetic", model=OffloadingModelConfig("synthetic", "uint8"),
34 cache=OffloadingCacheConfig(16, 1),
35 parallel=OffloadingParallelConfig(0, 1, 1, 1, 1, 1, 0, False),
36 )
37 spec = CPUOffloadingSpec(config)
38 region = SharedOffloadRegion("synthetic-" + uuid.uuid4().hex, 4, None, mmap.PAGESIZE, mmap.PAGESIZE)
39 primary = CPUPrimaryTierOffloadingManager(4, region)
40 self.view = primary.get_kv_memoryview()
41 self.tier = SecondaryTierFactory.create_secondary_tier({
42 "type": "studio_disk", "ssd_dir": str(self.root / "ssd"), "hdd_dir": str(self.root / "hdd"),
43 "ssd_bytes": 4 * mmap.PAGESIZE, "hdd_bytes": 6 * mmap.PAGESIZE,
44 }, self.view, spec)
45 self.manager = TieringOffloadingManager(primary, [self.tier])
46 self.addCleanup(self.manager.shutdown)
47 self.keys = [make_offload_key(bytes([i + 1]) * 32, i % 2) for i in range(3)]
48 self.values = [bytes([i + 11]) * mmap.PAGESIZE for i in range(3)]
49 context = ReqContext("store")
50 self.manager.on_new_request(context)
51 stored = self.manager.prepare_store(self.keys, context)
52 with self.view.cast("B") as flat:
53 for block_id, value in zip(stored.store_spec.block_ids, self.values, strict=True):
54 offset = int(block_id) * mmap.PAGESIZE
55 flat[offset:offset + mmap.PAGESIZE] = value
56 self.manager.complete_store(self.keys, context)
57 self.tier.drain_jobs()
58 self.manager.on_schedule_end(ScheduleEndContext((), ()))
59 self.manager.on_request_finished(context)
60 self.manager.reset_cache()
61 with self.view.cast("B") as flat:
62 flat[:] = bytes(len(flat))
63
64 def test_disk_hit_restores_attention_and_recurrent_groups_into_ram(self):
65 context = ReqContext("restore")
66 self.manager.on_new_request(context)
67 for key, value in zip(self.keys, self.values, strict=True):
68 self.assertEqual(self.manager.lookup(key, context), LookupResult.RETRY)
69 self.manager.on_schedule_end(ScheduleEndContext((), ()))
70 self.tier.drain_jobs()
71 self.assertEqual(self.manager.lookup(key, context), LookupResult.HIT)
72 restored = self.manager.prepare_load([key], context)
73 offset = int(restored.block_ids[0]) * mmap.PAGESIZE
74 with self.view.cast("B") as flat:
75 self.assertEqual(bytes(flat[offset:offset + mmap.PAGESIZE]), value)
76 self.manager.complete_load([key], context)
77 self.manager.on_schedule_end(ScheduleEndContext((), ()))
78 self.manager.on_request_finished(context)
79
80 def test_corrupt_disk_promotion_returns_miss_to_scheduler(self):
81 key = self.keys[0]
82 name = self.tier.mapper.get_file_name(key)
83 for root in self.tier.cache.roots:
84 path = root / name
85 if path.exists():
86 path.write_bytes(b"invalid")
87 context = ReqContext("corrupt")
88 self.manager.on_new_request(context)
89 self.assertEqual(self.manager.lookup(key, context), LookupResult.RETRY)
90 self.manager.on_schedule_end(ScheduleEndContext((), ()))
91 self.tier.drain_jobs()
92 self.assertEqual(self.manager.lookup(key, context), LookupResult.MISS)
93 self.manager.on_request_finished(context)
94
95
96if __name__ == "__main__":
97 unittest.main()