| 1 | """vLLM secondary tier for the bounded SSD/HDD cache.""" |
| 2 | import functools |
| 3 | import hashlib |
| 4 | import os |
| 5 | from pathlib import Path |
| 6 | |
| 7 | import vllm |
| 8 | from vllm.v1.kv_offload.base import LookupResult, Medium |
| 9 | from vllm.v1.kv_offload.file_mapper import FileMapper |
| 10 | from vllm.v1.kv_offload.tiering.base import JobResult, RequestOffloadingContext, SecondaryTierManager |
| 11 | from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory |
| 12 | from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool |
| 13 | |
| 14 | from studio_ai_cache import DiskCache |
| 15 | |
| 16 | |
| 17 | class DiskTier(SecondaryTierManager): |
| 18 | medium = Medium.STORAGE |
| 19 | |
| 20 | def __init__(self, offloading_spec, primary_kv_view, tier_type, ssd_dir, hdd_dir, ssd_bytes, hdd_bytes): |
| 21 | super().__init__(offloading_spec, primary_kv_view, tier_type) |
| 22 | key = hashlib.sha256(os.environ["STUDIO_AI_CACHE_KEY"].encode()).digest() |
| 23 | namespace = hashlib.sha256( |
| 24 | Path("/models/manifest.json").read_bytes() + vllm.__version__.encode() |
| 25 | + Path(__file__).read_bytes() + Path(__file__).with_name("studio_ai_cache.py").read_bytes() + key |
| 26 | ).hexdigest() |
| 27 | self.mapper = FileMapper.from_offloading_spec(namespace, offloading_spec, offloading_spec.blocks_per_chunk) |
| 28 | self.cache = DiskCache(ssd_dir, hdd_dir, ssd_bytes, hdd_bytes, key) |
| 29 | self.block_size = primary_kv_view.strides[0] |
| 30 | self.pool = DualQueueThreadPool(16, 2, thread_name_prefix="studio_kv_disk") |
| 31 | |
| 32 | def on_new_request(self, req_context): |
| 33 | return RequestOffloadingContext() |
| 34 | |
| 35 | def lookup(self, key, req_context): |
| 36 | return LookupResult.HIT if self.cache.acquire(self.mapper.get_file_name(key), req_context.req_id) else LookupResult.MISS |
| 37 | |
| 38 | def submit_store(self, job_metadata): |
| 39 | tasks = [functools.partial(self._transfer, self.cache.store, key, block_id) |
| 40 | for key, block_id in zip(job_metadata.keys, job_metadata.block_ids, strict=True)] |
| 41 | self.pool.enqueue_store(job_metadata.job_id, len(tasks), tasks) |
| 42 | |
| 43 | def submit_load(self, job_metadata): |
| 44 | tasks = [functools.partial(self._transfer, self.cache.load, key, block_id) |
| 45 | for key, block_id in zip(job_metadata.keys, job_metadata.block_ids, strict=True)] |
| 46 | self.pool.enqueue_load(job_metadata.job_id, len(tasks), tasks) |
| 47 | |
| 48 | def _transfer(self, operation, key, block_id): |
| 49 | offset = int(block_id) * self.block_size |
| 50 | with self._primary_kv_view.cast("B")[offset:offset + self.block_size] as view: |
| 51 | operation(self.mapper.get_file_name(key), view) |
| 52 | |
| 53 | def get_finished_jobs(self): |
| 54 | return [JobResult(job_id, success) for job_id, success in self.pool.get_finished()] |
| 55 | |
| 56 | def touch(self, keys, req_context): |
| 57 | for key in keys: |
| 58 | self.cache.touch(self.mapper.get_file_name(key)) |
| 59 | |
| 60 | def on_request_finished(self, req_context): |
| 61 | self.cache.release(req_context.req_id) |
| 62 | |
| 63 | def drain_jobs(self): |
| 64 | self.pool.wait_idle() |
| 65 | |
| 66 | def shutdown(self): |
| 67 | self.pool.wait_idle() |
| 68 | self.pool.shutdown(wait=True) |
| 69 | self.cache.close() |
| 70 | |
| 71 | |
| 72 | def register(): |
| 73 | SecondaryTierFactory.register_tier("studio_disk", __name__, "DiskTier") |