| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FileSystemTierManager: pure-Python filesystem tier for KV cache offloading. |
| 4 | |
| 5 | Store path: |
| 6 | Data is written to a temp file (<dest_path.tmp>) via os.write, |
| 7 | then os.replace'd to the final path (without .tmp). |
| 8 | |
| 9 | Load path: |
| 10 | Data is read from the block file directly via os.readv into the |
| 11 | provided memoryview slice. |
| 12 | |
| 13 | File naming: <base_path>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash_hex>.bin |
| 14 | (hash-based subdirectories to limit directory fan-out) |
| 15 | """ |
| 16 | |
| 17 | import functools |
| 18 | import json |
| 19 | import os |
| 20 | from collections.abc import Iterable |
| 21 | from typing import TYPE_CHECKING, ClassVar |
| 22 | |
| 23 | try: |
| 24 | from vllm.fs_io_C import batch_lookup as batch_lookup_C |
| 25 | |
| 26 | _HAS_BATCH_LOOKUP_C = True |
| 27 | except ImportError: |
| 28 | _HAS_BATCH_LOOKUP_C = False |
| 29 | |
| 30 | from typing_extensions import override |
| 31 | |
| 32 | from vllm.logger import init_logger |
| 33 | from vllm.v1.kv_offload.base import ( |
| 34 | Locality, |
| 35 | LookupResult, |
| 36 | Medium, |
| 37 | OffloadingEvent, |
| 38 | OffloadKey, |
| 39 | ReqContext, |
| 40 | ) |
| 41 | from vllm.v1.kv_offload.file_mapper import FileMapper |
| 42 | from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager |
| 43 | from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector |
| 44 | from vllm.v1.kv_offload.tiering.base import ( |
| 45 | JobId, |
| 46 | JobResult, |
| 47 | RequestOffloadingContext, |
| 48 | ScheduleEndContext, |
| 49 | SecondaryTierManager, |
| 50 | TransferJob, |
| 51 | ) |
| 52 | from vllm.v1.kv_offload.tiering.fs.io import ( |
| 53 | batch_load_block, |
| 54 | batch_store_block, |
| 55 | probe_o_direct, |
| 56 | ) |
| 57 | from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool |
| 58 | |
| 59 | if TYPE_CHECKING: |
| 60 | from vllm.v1.kv_offload.base import OffloadingSpec |
| 61 | |
| 62 | logger = init_logger(__name__) |
| 63 | |
| 64 | |
| 65 | class FsAsyncLookupManager(AsyncLookupManager): |
| 66 | """Async lookup manager for FileSystemTierManager.""" |
| 67 | |
| 68 | def __init__( |
| 69 | self, |
| 70 | tier: "FileSystemTierManager", |
| 71 | tier_type: str, |
| 72 | ) -> None: |
| 73 | super().__init__(tier_type=tier_type) |
| 74 | self._tier = tier |
| 75 | |
| 76 | def batch_lookup( |
| 77 | self, keys: list[OffloadKey], req_context: ReqContext |
| 78 | ) -> Iterable[bool]: |
| 79 | paths = [self._tier.file_mapper.get_file_name(k) for k in keys] |
| 80 | if _HAS_BATCH_LOOKUP_C: |
| 81 | # C extension: GIL released for the entire faccessat() batch. |
| 82 | return batch_lookup_C(paths) |
| 83 | return (os.path.exists(p) for p in paths) |
| 84 | |
| 85 | |
| 86 | class FileSystemTierManager(SecondaryTierManager): |
| 87 | """Pure-Python disk-backed secondary tier. |
| 88 | |
| 89 | Read-priority threads service load jobs preferentially; write-priority |
| 90 | threads service store jobs preferentially. Both groups can drain either |
| 91 | queue, so neither starves. |
| 92 | |
| 93 | submit_store / submit_load are non-blocking: they enqueue tasks and return. |
| 94 | get_finished_jobs() polls job completion and returns completed JobResults. |
| 95 | |
| 96 | Cross-process sharing: |
| 97 | KV cache sharing between multiple vLLM instances using the same |
| 98 | ``root_dir`` (e.g., via a shared PVC) works by default: ``NONE_HASH`` |
| 99 | (the chain-hash seed for block content hashes) is derived from a fixed |
| 100 | default seed, so identical token content produces identical block |
| 101 | filenames across instances. Setting the ``PYTHONHASHSEED`` environment |
| 102 | variable to the same value on all instances overrides the default seed, |
| 103 | and is required to share a cache when using a non-cryptographic |
| 104 | prefix-caching hash algorithm, which seeds ``NONE_HASH`` randomly. |
| 105 | """ |
| 106 | |
| 107 | medium: ClassVar[Medium] = Medium.STORAGE |
| 108 | |
| 109 | def __init__( |
| 110 | self, |
| 111 | offloading_spec: "OffloadingSpec", |
| 112 | primary_kv_view: memoryview, |
| 113 | tier_type: str, |
| 114 | root_dir: str, |
| 115 | n_read_threads: int = 16, |
| 116 | n_write_threads: int = 16, |
| 117 | enable_kv_events: bool = False, |
| 118 | locality: str | None = None, |
| 119 | backpressure_detector: BackpressureDetector | None = None, |
| 120 | ): |
| 121 | """Args: |
| 122 | offloading_spec: Contains normalized offloading configuration and |
| 123 | blocks_per_chunk. |
| 124 | primary_kv_view: Memoryview of the primary tier's CPU KV cache. |
| 125 | tier_type: Tier type identifier, set by SecondaryTierFactory. |
| 126 | root_dir: Root directory for block files. |
| 127 | n_read_threads: Number of read-priority I/O threads. |
| 128 | n_write_threads: Number of write-priority I/O threads. |
| 129 | enable_kv_events: Emit BlockStored KV events for blocks |
| 130 | successfully stored to this tier. Effective only when KV |
| 131 | cache events are enabled globally (kv_events_config). |
| 132 | locality: Whether this tier's storage is LOCAL or REMOTE relative |
| 133 | to the publishing vLLM instance. |
| 134 | backpressure_detector: Optional backpressure detector. |
| 135 | |
| 136 | """ |
| 137 | super().__init__( |
| 138 | offloading_spec, primary_kv_view, tier_type, backpressure_detector |
| 139 | ) |
| 140 | self.locality = Locality(locality) if locality is not None else None |
| 141 | |
| 142 | self.events: list[OffloadingEvent] | None = None |
| 143 | if enable_kv_events: |
| 144 | if offloading_spec.kv_events_config.enable_kv_cache_events: |
| 145 | self.events = [] |
| 146 | else: |
| 147 | logger.warning( |
| 148 | "enable_kv_events is set on secondary tier '%s' but KV " |
| 149 | "cache events are disabled globally; the tier will not " |
| 150 | "emit events.", |
| 151 | tier_type, |
| 152 | ) |
| 153 | # Keys of in-flight store jobs, tracked only when events are enabled. |
| 154 | self._store_job_keys: dict[JobId, list[OffloadKey]] = {} |
| 155 | # Keys of in-flight load (promotion) jobs, so a failed load can mark |
| 156 | # its own cached lookup verdicts False (see get_finished_jobs). |
| 157 | self._load_job_keys: dict[JobId, list[OffloadKey]] = {} |
| 158 | # Block count per in-flight job, used to report transfer_bytes. |
| 159 | self._job_block_counts: dict[JobId, int] = {} |
| 160 | # Per load job: how many blocks loaded before a failure (partial keep). |
| 161 | # Written by the pool worker inside the load task before it raises (so |
| 162 | # before task_done publishes the job); read on the scheduler thread in |
| 163 | # get_finished_jobs only for job ids the finished queue returned. Under |
| 164 | # the GIL that read cannot observe the finished job without the prior |
| 165 | # write, so no extra lock is needed (get_finished is itself lock-free). |
| 166 | self._load_progress: dict[JobId, int] = {} |
| 167 | |
| 168 | # Extract block size from primary view |
| 169 | assert primary_kv_view.strides is not None, ( |
| 170 | "primary_kv_view.strides cannot be None" |
| 171 | ) |
| 172 | self._block_size: int = primary_kv_view.strides[0] |
| 173 | |
| 174 | # Opt in; FileMapper enables it only for a parallelism-invariant block. |
| 175 | self.file_mapper = FileMapper.from_offloading_spec( |
| 176 | root_dir=root_dir, |
| 177 | offloading_spec=offloading_spec, |
| 178 | blocks_per_file=offloading_spec.blocks_per_chunk, |
| 179 | parallel_agnostic=True, |
| 180 | ) |
| 181 | |
| 182 | # Write config file |
| 183 | config_path = self.file_mapper.get_config_file_path() |
| 184 | os.makedirs(os.path.dirname(config_path), exist_ok=True) |
| 185 | if not os.path.exists(config_path): |
| 186 | with open(config_path, "w") as f: |
| 187 | json.dump( |
| 188 | self.file_mapper.get_run_config(), f, indent=2, sort_keys=True |
| 189 | ) |
| 190 | |
| 191 | # Prefer O_DIRECT to bypass the page cache, but fall back to buffered |
| 192 | # I/O on filesystems that reject it (e.g. overlayfs, some NFS mounts) |
| 193 | # rather than failing every block. |
| 194 | self._use_o_direct = probe_o_direct(os.path.dirname(config_path)) |
| 195 | if not self._use_o_direct: |
| 196 | logger.warning( |
| 197 | "O_DIRECT is not supported at '%s'; falling back to buffered " |
| 198 | "I/O for the '%s' KV offload tier.", |
| 199 | root_dir, |
| 200 | tier_type, |
| 201 | ) |
| 202 | |
| 203 | self._pool = DualQueueThreadPool( |
| 204 | n_read_threads, |
| 205 | n_write_threads, |
| 206 | thread_name_prefix="vllm_kv_py_fs", |
| 207 | ) |
| 208 | |
| 209 | self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type) |
| 210 | |
| 211 | @override |
| 212 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: |
| 213 | return RequestOffloadingContext() |
| 214 | |
| 215 | @override |
| 216 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: |
| 217 | result = self._lookup_manager.lookup(key, req_context) |
| 218 | if result is None: |
| 219 | return LookupResult.RETRY |
| 220 | return LookupResult.HIT if result else LookupResult.MISS |
| 221 | |
| 222 | @override |
| 223 | def submit_store(self, job_metadata: TransferJob) -> None: |
| 224 | keys = list(job_metadata.keys) |
| 225 | if self.events is not None: |
| 226 | self._store_job_keys[job_metadata.job_id] = keys |
| 227 | task = functools.partial( |
| 228 | batch_store_block, |
| 229 | [self.file_mapper.get_file_name(key) for key in keys], |
| 230 | self._primary_kv_view, |
| 231 | [int(cid) * self._block_size for cid in job_metadata.chunk_ids], |
| 232 | self._block_size, |
| 233 | self._use_o_direct, |
| 234 | ) |
| 235 | self._job_block_counts[job_metadata.job_id] = len(keys) |
| 236 | self._pool.enqueue_store(job_metadata.job_id, 1, [task]) |
| 237 | |
| 238 | @override |
| 239 | def submit_load(self, job_metadata: TransferJob) -> None: |
| 240 | job_id = job_metadata.job_id |
| 241 | # Track this load's keys so a failed promotion can mark only its failed |
| 242 | # keys as a miss (see get_finished_jobs). |
| 243 | keys = list(job_metadata.keys) |
| 244 | self._load_job_keys[job_id] = keys |
| 245 | self._job_block_counts[job_id] = len(keys) |
| 246 | paths = [self.file_mapper.get_file_name(key) for key in keys] |
| 247 | offsets = [int(cid) * self._block_size for cid in job_metadata.chunk_ids] |
| 248 | |
| 249 | def load_task() -> None: |
| 250 | try: |
| 251 | batch_load_block( |
| 252 | paths, |
| 253 | self._primary_kv_view, |
| 254 | offsets, |
| 255 | self._block_size, |
| 256 | self._use_o_direct, |
| 257 | ) |
| 258 | except OSError as exc: |
| 259 | # Runs on the pool worker thread. Record how many blocks loaded |
| 260 | # before the failure so get_finished_jobs can keep them; this |
| 261 | # write precedes task_done, so the scheduler reads it safely |
| 262 | # under the GIL once the finished queue hands back this job. |
| 263 | num_succeeded = getattr(exc, "num_succeeded", 0) |
| 264 | self._load_progress[job_id] = num_succeeded |
| 265 | # Surfaces errno (e.g. EMFILE "Too many open files") for both |
| 266 | # the C and Python load paths. |
| 267 | logger.debug( |
| 268 | "Load of %d blocks for job %s failed at block %d: %s", |
| 269 | len(paths), |
| 270 | job_id, |
| 271 | num_succeeded, |
| 272 | exc, |
| 273 | ) |
| 274 | raise |
| 275 | |
| 276 | self._pool.enqueue_load(job_id, 1, [load_task]) |
| 277 | |
| 278 | @override |
| 279 | def get_finished_jobs(self) -> Iterable[JobResult]: |
| 280 | """Collect finished jobs; a failed promotion marks only its failed keys |
| 281 | as a miss here (scheduler thread).""" |
| 282 | results = [] |
| 283 | for job_id, success, transfer_time in self._pool.get_finished(): |
| 284 | block_count = self._job_block_counts.pop(job_id, 0) |
| 285 | transfer_bytes = block_count * self._block_size if block_count else None |
| 286 | if self.events is not None: |
| 287 | keys = self._store_job_keys.pop(job_id, None) |
| 288 | if success and keys: |
| 289 | self.events.append( |
| 290 | OffloadingEvent( |
| 291 | keys=keys, |
| 292 | medium=self.medium, |
| 293 | removed=False, |
| 294 | locality=self.locality, |
| 295 | ) |
| 296 | ) |
| 297 | load_keys = self._load_job_keys.pop(job_id, None) |
| 298 | num_succeeded = self._load_progress.pop(job_id, 0) |
| 299 | if load_keys is not None and not success: |
| 300 | # A batched load stops at the first bad block and reports how |
| 301 | # many loaded before it. Those earlier blocks are kept in the |
| 302 | # primary tier (reported via successful_keys); only this block |
| 303 | # and the ones after it are marked a miss and recomputed. |
| 304 | successful = load_keys[:num_succeeded] |
| 305 | failed = load_keys[num_succeeded:] |
| 306 | self._lookup_manager.mark_miss(failed) |
| 307 | results.append( |
| 308 | JobResult( |
| 309 | job_id=job_id, |
| 310 | success=False, |
| 311 | successful_keys=tuple(successful) if successful else None, |
| 312 | transfer_time=transfer_time, |
| 313 | transfer_bytes=transfer_bytes, |
| 314 | ) |
| 315 | ) |
| 316 | continue |
| 317 | results.append( |
| 318 | JobResult( |
| 319 | job_id=job_id, |
| 320 | success=success, |
| 321 | transfer_time=transfer_time, |
| 322 | transfer_bytes=transfer_bytes, |
| 323 | ) |
| 324 | ) |
| 325 | return results |
| 326 | |
| 327 | @override |
| 328 | def take_events(self) -> Iterable[OffloadingEvent]: |
| 329 | if self.events is not None: |
| 330 | yield from self.events |
| 331 | self.events.clear() |
| 332 | |
| 333 | @override |
| 334 | def drain_jobs(self) -> None: |
| 335 | """Block until all in-flight transfers in the threadpool finish.""" |
| 336 | self._pool.wait_idle() |
| 337 | |
| 338 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 339 | self._lookup_manager.cleanup(req_context.req_id) |
| 340 | |
| 341 | @override |
| 342 | def on_schedule_end(self, context: ScheduleEndContext) -> None: |
| 343 | self._lookup_manager.flush() |
| 344 | |
| 345 | @override |
| 346 | def shutdown(self) -> None: |
| 347 | """Release resources held by this tier. |
| 348 | |
| 349 | Shuts down the lookup manager and the thread pool, |
| 350 | clearing pending tasks and waiting for active threads to complete. |
| 351 | """ |
| 352 | self._lookup_manager.shutdown() |
| 353 | self._pool.shutdown(wait=True) |