| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Object store secondary tier implementation.""" |
| 4 | |
| 5 | import ctypes |
| 6 | import time |
| 7 | from collections.abc import Iterable |
| 8 | from typing import TYPE_CHECKING, ClassVar, NamedTuple |
| 9 | |
| 10 | import numpy as np |
| 11 | |
| 12 | from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent |
| 13 | from vllm.distributed.nixl_utils import nixl_agent_config |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.v1.kv_offload.base import ( |
| 16 | Locality, |
| 17 | LookupResult, |
| 18 | Medium, |
| 19 | OffloadingEvent, |
| 20 | OffloadKey, |
| 21 | ReqContext, |
| 22 | ) |
| 23 | from vllm.v1.kv_offload.file_mapper import FileMapper |
| 24 | from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager |
| 25 | from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector |
| 26 | from vllm.v1.kv_offload.tiering.base import ( |
| 27 | JobId, |
| 28 | JobResult, |
| 29 | RequestOffloadingContext, |
| 30 | ScheduleEndContext, |
| 31 | SecondaryTierManager, |
| 32 | TransferJob, |
| 33 | ) |
| 34 | from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig |
| 35 | |
| 36 | if TYPE_CHECKING: |
| 37 | from nixl._api import nixl_prepped_dlist_handle, nixl_xfer_handle |
| 38 | |
| 39 | from vllm.v1.kv_offload.base import OffloadingSpec |
| 40 | |
| 41 | logger = init_logger(__name__) |
| 42 | |
| 43 | NIXL_WRITE = "WRITE" |
| 44 | NIXL_READ = "READ" |
| 45 | NIXL_PROC = "PROC" |
| 46 | NIXL_DONE = "DONE" |
| 47 | |
| 48 | # Device ID for CPU DRAM descriptors. DRAM is not a multi-device resource so |
| 49 | # the device ID is always 0. |
| 50 | NIXL_DEV_ID: int = 0 |
| 51 | |
| 52 | # Fields for NIXL OBJ descriptors: (addr, len, dev_id, obj_key). |
| 53 | # For existence probes addr and len are placeholders — no data is read. |
| 54 | # dev_id=0 is reserved for probes; transfers start from 1. |
| 55 | _PROBE_ADDR: int = 0 |
| 56 | _PROBE_LEN: int = 1 |
| 57 | _PROBE_DEV_ID: int = 0 |
| 58 | |
| 59 | |
| 60 | class TransferEntry(NamedTuple): |
| 61 | xfer_handle: "nixl_xfer_handle" |
| 62 | files_desc: object |
| 63 | obj_handle: "nixl_prepped_dlist_handle" |
| 64 | |
| 65 | |
| 66 | class ObjAsyncLookupManager(AsyncLookupManager): |
| 67 | """Async lookup manager for ObjectStoreSecondaryTierManager. |
| 68 | |
| 69 | Batches existence probes into a single query_memory() call so the |
| 70 | background thread issues one round-trip per step instead of one per key. |
| 71 | """ |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | tier: "ObjectStoreSecondaryTierManager", |
| 76 | tier_type: str, |
| 77 | ) -> None: |
| 78 | super().__init__(tier_type=tier_type) |
| 79 | self._tier = tier |
| 80 | |
| 81 | def batch_lookup( |
| 82 | self, keys: list[OffloadKey], req_context: ReqContext |
| 83 | ) -> Iterable[bool]: |
| 84 | descriptors = [ |
| 85 | ( |
| 86 | _PROBE_ADDR, |
| 87 | _PROBE_LEN, |
| 88 | _PROBE_DEV_ID, |
| 89 | self._tier._file_mapper.get_file_name(k), |
| 90 | ) |
| 91 | for k in keys |
| 92 | ] |
| 93 | results = self._tier._agent.query_memory(descriptors, "OBJ", "OBJ") |
| 94 | return (r is not None for r in results) |
| 95 | |
| 96 | |
| 97 | class ObjectStoreSecondaryTierManager(SecondaryTierManager): |
| 98 | """Secondary tier that offloads KV cache blocks to an S3-compatible store. |
| 99 | |
| 100 | Handles CPU DRAM <-> S3 transfers only. GPU <-> CPU is managed by the |
| 101 | primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``. |
| 102 | """ |
| 103 | |
| 104 | medium: ClassVar[Medium] = Medium.STORAGE |
| 105 | |
| 106 | def __init__( |
| 107 | self, |
| 108 | offloading_spec: "OffloadingSpec", |
| 109 | primary_kv_view: memoryview, |
| 110 | tier_type: str, |
| 111 | store_config: dict, |
| 112 | prefix: str = "", |
| 113 | io_threads: int = 4, |
| 114 | enable_kv_events: bool = False, |
| 115 | locality: str | None = None, |
| 116 | backpressure_detector: BackpressureDetector | None = None, |
| 117 | ): |
| 118 | """Args: |
| 119 | offloading_spec: Offloading configuration. |
| 120 | primary_kv_view: Memoryview of the primary tier's CPU KV cache. |
| 121 | tier_type: Tier type identifier, set by SecondaryTierFactory. |
| 122 | store_config: Object store connection parameters (see ObjStoreConfig). |
| 123 | prefix: Key prefix prepended to all object keys. |
| 124 | io_threads: Number of NIXL I/O threads. |
| 125 | enable_kv_events: Emit BlockStored KV events for blocks |
| 126 | successfully stored to this tier. Effective only when KV |
| 127 | cache events are enabled globally (kv_events_config). |
| 128 | locality: Whether this tier's storage is LOCAL or REMOTE relative |
| 129 | to the publishing vLLM instance. |
| 130 | backpressure_detector: Optional backpressure detector. |
| 131 | |
| 132 | """ |
| 133 | super().__init__( |
| 134 | offloading_spec, |
| 135 | primary_kv_view, |
| 136 | tier_type, |
| 137 | backpressure_detector, |
| 138 | ) |
| 139 | self.locality = Locality(locality) if locality is not None else None |
| 140 | |
| 141 | self.events: list[OffloadingEvent] | None = None |
| 142 | if enable_kv_events: |
| 143 | if offloading_spec.kv_events_config.enable_kv_cache_events: |
| 144 | self.events = [] |
| 145 | else: |
| 146 | logger.warning( |
| 147 | "enable_kv_events is set on secondary tier '%s' but KV " |
| 148 | "cache events are disabled globally; the tier will not " |
| 149 | "emit events.", |
| 150 | tier_type, |
| 151 | ) |
| 152 | # Keys of in-flight store jobs, tracked only when events are enabled. |
| 153 | self._store_job_keys: dict[JobId, list[OffloadKey]] = {} |
| 154 | # Keys of in-flight load (promotion) jobs, so a failed download can |
| 155 | # mark its own cached lookup verdicts False (see get_finished_jobs). |
| 156 | self._load_job_keys: dict[JobId, list[OffloadKey]] = {} |
| 157 | |
| 158 | agent_config = nixl_agent_config(backends=[]) |
| 159 | self._agent = nixl_agent("ObjAgent", agent_config) |
| 160 | obj_config = ObjStoreConfig(**store_config) |
| 161 | params = {**obj_config.to_nixl_params(), "num_threads": str(io_threads)} |
| 162 | self._agent.create_backend("OBJ", params) |
| 163 | self._transfers: dict[int, TransferEntry] = {} |
| 164 | # Buffered results awaiting the next get_finished_jobs() call: |
| 165 | # submission-time failures + poll-time completions accumulated |
| 166 | # during drain_jobs(). |
| 167 | self._pending_results: list[JobResult] = [] |
| 168 | self._primary_reg = None |
| 169 | self._block_size_bytes: int = 0 |
| 170 | root_dir = f"{prefix}/" if prefix else "" |
| 171 | # Opt in; FileMapper enables it only for a parallelism-invariant block. |
| 172 | self._file_mapper = FileMapper.from_offloading_spec( |
| 173 | root_dir, offloading_spec, parallel_agnostic=True |
| 174 | ) |
| 175 | self._next_obj_dev_id: int = 1 # dev_id=0 is reserved for _exists() probes |
| 176 | |
| 177 | self._probe_connectivity() |
| 178 | |
| 179 | base_addr = ctypes.addressof(ctypes.c_char.from_buffer(primary_kv_view)) |
| 180 | assert primary_kv_view.strides is not None |
| 181 | stride = primary_kv_view.strides[0] |
| 182 | self._primary_reg = self._agent.register_memory( |
| 183 | [(base_addr, primary_kv_view.nbytes, NIXL_DEV_ID, "")], "DRAM" |
| 184 | ) |
| 185 | self._block_size_bytes = stride |
| 186 | all_blocks = [ |
| 187 | (base_addr + i * stride, stride, NIXL_DEV_ID) |
| 188 | for i in range(len(primary_kv_view)) |
| 189 | ] |
| 190 | # NIXL_INIT_AGENT marks this as the local side; make_prepped_xfer requires |
| 191 | # local_xfer_side tagged with NIXL_INIT_AGENT and remote_xfer_side tagged |
| 192 | # with the peer agent name ("ObjAgent"). |
| 193 | self._dram_prepped_handle: nixl_prepped_dlist_handle = ( |
| 194 | self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", all_blocks, "DRAM") |
| 195 | ) |
| 196 | |
| 197 | self._lookup_manager = ObjAsyncLookupManager( |
| 198 | tier=self, tier_type=self.tier_type |
| 199 | ) |
| 200 | |
| 201 | def _probe_connectivity(self) -> None: |
| 202 | """Verify object store connectivity at startup via a NIXL lookup probe. |
| 203 | |
| 204 | Performs a single exists() check against a synthetic key that will |
| 205 | never exist. A True/False result confirms the bucket is reachable; |
| 206 | an exception indicates misconfigured obj store params and raises RuntimeError. |
| 207 | """ |
| 208 | probe_key = "__nixl_probe__/connectivity_test" |
| 209 | try: |
| 210 | self._exists(probe_key) |
| 211 | logger.info("Object store tier connectivity probe succeeded") |
| 212 | except Exception as e: |
| 213 | raise RuntimeError( |
| 214 | f"Object store tier connectivity probe failed — check bucket, " |
| 215 | f"endpoint_override, and scheme. If using explicit credentials " |
| 216 | f"verify access_key and secret_key; otherwise ensure the AWS " |
| 217 | f"SDK default credential chain is configured (IAM role, env " |
| 218 | f"vars, credential file). Error: {e}" |
| 219 | ) from e |
| 220 | |
| 221 | def _exists(self, obj_key: str) -> bool: |
| 222 | results = self._agent.query_memory( |
| 223 | [(_PROBE_ADDR, _PROBE_LEN, _PROBE_DEV_ID, obj_key)], "OBJ", "OBJ" |
| 224 | ) |
| 225 | return results[0] is not None |
| 226 | |
| 227 | def _submit_transfer( |
| 228 | self, |
| 229 | job_id: int, |
| 230 | block_ids: np.ndarray, |
| 231 | obj_keys: Iterable[str], |
| 232 | op: str, |
| 233 | ) -> None: |
| 234 | """Submit an async transfer. op is 'WRITE' (store) or 'READ' (load).""" |
| 235 | # NIXL takes indices as int32; a wider dtype silently costs a conversion. |
| 236 | assert block_ids.dtype == np.int32 |
| 237 | # The OBJ backend maps devId -> obj_key. All descriptors must have |
| 238 | # unique devIds or later registrations overwrite earlier ones. |
| 239 | nixl_files = [ |
| 240 | (0, self._block_size_bytes, dev_id, key) |
| 241 | for dev_id, key in enumerate(obj_keys, self._next_obj_dev_id) |
| 242 | ] |
| 243 | self._next_obj_dev_id += len(nixl_files) |
| 244 | |
| 245 | files_desc = self._agent.register_memory(nixl_files, "OBJ") |
| 246 | if files_desc is None: |
| 247 | logger.warning("register_memory (OBJ) failed for job %d", job_id) |
| 248 | self._pending_results.append(JobResult(job_id=job_id, success=False)) |
| 249 | return |
| 250 | |
| 251 | obj_handle = self._agent.prep_xfer_dlist("ObjAgent", files_desc.trim()) |
| 252 | if not obj_handle: |
| 253 | logger.warning("prep_xfer_dlist (OBJ) failed for job %d", job_id) |
| 254 | self._agent.deregister_memory(files_desc) |
| 255 | self._pending_results.append(JobResult(job_id=job_id, success=False)) |
| 256 | return |
| 257 | |
| 258 | xfer_handle = self._agent.make_prepped_xfer( |
| 259 | op, |
| 260 | self._dram_prepped_handle, |
| 261 | block_ids, |
| 262 | obj_handle, |
| 263 | np.arange(len(nixl_files), dtype=np.int32), |
| 264 | ) |
| 265 | if not xfer_handle: |
| 266 | logger.warning("make_prepped_xfer failed for job %d", job_id) |
| 267 | self._agent.release_dlist_handle(obj_handle) |
| 268 | self._agent.deregister_memory(files_desc) |
| 269 | self._pending_results.append(JobResult(job_id=job_id, success=False)) |
| 270 | return |
| 271 | |
| 272 | state = self._agent.transfer(xfer_handle) |
| 273 | if state == "ERR": |
| 274 | logger.warning("agent.transfer failed for job %d", job_id) |
| 275 | self._agent.release_dlist_handle(obj_handle) |
| 276 | self._agent.deregister_memory(files_desc) |
| 277 | self._agent.release_xfer_handle(xfer_handle) |
| 278 | self._pending_results.append(JobResult(job_id=job_id, success=False)) |
| 279 | return |
| 280 | |
| 281 | self._transfers[job_id] = TransferEntry( |
| 282 | xfer_handle=xfer_handle, |
| 283 | files_desc=files_desc, |
| 284 | obj_handle=obj_handle, |
| 285 | ) |
| 286 | |
| 287 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: |
| 288 | result = self._lookup_manager.lookup(key, req_context) |
| 289 | if result is None: |
| 290 | return LookupResult.RETRY |
| 291 | return LookupResult.HIT if result else LookupResult.MISS |
| 292 | |
| 293 | def submit_store(self, job_metadata: TransferJob) -> None: |
| 294 | if self.events is not None: |
| 295 | self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) |
| 296 | obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) |
| 297 | self._submit_transfer( |
| 298 | job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_WRITE |
| 299 | ) |
| 300 | |
| 301 | def submit_load(self, job_metadata: TransferJob) -> None: |
| 302 | self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys) |
| 303 | obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) |
| 304 | self._submit_transfer( |
| 305 | job_metadata.job_id, job_metadata.chunk_ids, obj_keys, NIXL_READ |
| 306 | ) |
| 307 | |
| 308 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 309 | self._lookup_manager.cleanup(req_context.req_id) |
| 310 | |
| 311 | def on_schedule_end(self, context: ScheduleEndContext) -> None: |
| 312 | self._lookup_manager.flush() |
| 313 | |
| 314 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: |
| 315 | return RequestOffloadingContext() |
| 316 | |
| 317 | def _poll_active_transfers(self) -> None: |
| 318 | """Poll all in-flight transfers once; move newly-completed (success or |
| 319 | failure) into ``_pending_results`` and release their NIXL handles.""" |
| 320 | for job_id, entry in list(self._transfers.items()): |
| 321 | try: |
| 322 | state = self._agent.check_xfer_state(entry.xfer_handle) |
| 323 | except Exception as exc: |
| 324 | success = False |
| 325 | logger.warning("check_xfer_state raised for job %d: %s", job_id, exc) |
| 326 | else: |
| 327 | if state == NIXL_PROC: |
| 328 | continue |
| 329 | if state == NIXL_DONE: |
| 330 | success = True |
| 331 | else: |
| 332 | success = False |
| 333 | logger.warning("transfer failed job=%d state=%s", job_id, state) |
| 334 | |
| 335 | transfer_time = None |
| 336 | if success: |
| 337 | telemetry = self._agent.get_xfer_telemetry(entry.xfer_handle) |
| 338 | transfer_time = telemetry.xferDuration / 1e6 |
| 339 | |
| 340 | try: |
| 341 | self._agent.release_xfer_handle(entry.xfer_handle) |
| 342 | except Exception as exc: |
| 343 | # Keep the entry until NIXL confirms that the transfer handle |
| 344 | # can be released. The transfer may still access primary-tier |
| 345 | # memory, so publishing its result would allow unsafe reuse. |
| 346 | logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc) |
| 347 | continue |
| 348 | |
| 349 | # Once the transfer handle is released, these remaining cleanup |
| 350 | # failures must not suppress the job completion. They can leak |
| 351 | # NIXL metadata, but cannot leave an active data transfer behind. |
| 352 | try: |
| 353 | self._agent.release_dlist_handle(entry.obj_handle) |
| 354 | except Exception as exc: |
| 355 | logger.warning( |
| 356 | "release_dlist_handle failed for job %d: %s", job_id, exc |
| 357 | ) |
| 358 | try: |
| 359 | self._agent.deregister_memory(entry.files_desc) |
| 360 | except Exception as exc: |
| 361 | logger.warning("deregister_memory failed for job %d: %s", job_id, exc) |
| 362 | |
| 363 | del self._transfers[job_id] |
| 364 | self._pending_results.append( |
| 365 | JobResult( |
| 366 | job_id=job_id, |
| 367 | success=success, |
| 368 | transfer_time=transfer_time, |
| 369 | ) |
| 370 | ) |
| 371 | |
| 372 | def get_finished_jobs(self) -> Iterable[JobResult]: |
| 373 | """Poll transfers; a failed promotion marks its cached verdicts False |
| 374 | here (scheduler thread).""" |
| 375 | self._poll_active_transfers() |
| 376 | results = self._pending_results |
| 377 | self._pending_results = [] |
| 378 | for result in results: |
| 379 | if self.events is not None: |
| 380 | keys = self._store_job_keys.pop(result.job_id, None) |
| 381 | if result.success and keys: |
| 382 | self.events.append( |
| 383 | OffloadingEvent( |
| 384 | keys=keys, |
| 385 | medium=self.medium, |
| 386 | removed=False, |
| 387 | locality=self.locality, |
| 388 | ) |
| 389 | ) |
| 390 | # Mark only the keys that did not load as a miss; the request |
| 391 | # recomputes them. The miss is per-request (cleared when the request |
| 392 | # finishes), so other requests still HIT the blocks that loaded |
| 393 | # fine. nixl reports the batch as a whole (successful_keys is None), |
| 394 | # so today this marks all keys; the subtraction keeps it correct if |
| 395 | # partial results are ever reported. |
| 396 | load_keys = self._load_job_keys.pop(result.job_id, None) |
| 397 | if load_keys is not None and not result.success: |
| 398 | successful = set(result.successful_keys or ()) |
| 399 | failed = [k for k in load_keys if k not in successful] |
| 400 | self._lookup_manager.mark_miss(failed) |
| 401 | return results |
| 402 | |
| 403 | def take_events(self) -> Iterable[OffloadingEvent]: |
| 404 | if self.events is not None: |
| 405 | yield from self.events |
| 406 | self.events.clear() |
| 407 | |
| 408 | def drain_jobs(self) -> None: |
| 409 | """Block until every submitted transfer has completed or failed. |
| 410 | |
| 411 | nixl exposes only ``check_xfer_state`` (poll-based), so this loops |
| 412 | until ``_transfers`` is empty. Results accumulate in |
| 413 | ``_pending_results`` and are surfaced by the next |
| 414 | ``get_finished_jobs()`` call. |
| 415 | """ |
| 416 | start = time.monotonic() |
| 417 | warned = False |
| 418 | while self._transfers: |
| 419 | self._poll_active_transfers() |
| 420 | if not self._transfers: |
| 421 | break |
| 422 | if not warned and time.monotonic() - start > 5.0: |
| 423 | logger.warning( |
| 424 | "ObjectStoreSecondaryTierManager.drain_jobs: still " |
| 425 | "draining after 5s (%d transfers in flight); a stuck " |
| 426 | "transfer will block the engine.", |
| 427 | len(self._transfers), |
| 428 | ) |
| 429 | warned = True |
| 430 | time.sleep(0.001) |
| 431 | |
| 432 | def shutdown(self) -> None: |
| 433 | self._lookup_manager.shutdown() |
| 434 | for job_id, entry in self._transfers.items(): |
| 435 | try: |
| 436 | self._agent.release_xfer_handle(entry.xfer_handle) |
| 437 | except Exception as exc: |
| 438 | logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc) |
| 439 | try: |
| 440 | self._agent.release_dlist_handle(entry.obj_handle) |
| 441 | except Exception as exc: |
| 442 | logger.warning( |
| 443 | "release_dlist_handle failed for job %d: %s", job_id, exc |
| 444 | ) |
| 445 | try: |
| 446 | self._agent.deregister_memory(entry.files_desc) |
| 447 | except Exception as exc: |
| 448 | logger.warning("deregister_memory failed for job %d: %s", job_id, exc) |
| 449 | self._transfers.clear() |
| 450 | if self._dram_prepped_handle is not None: |
| 451 | try: |
| 452 | self._agent.release_dlist_handle(self._dram_prepped_handle) |
| 453 | except Exception as exc: |
| 454 | logger.warning("failed to release DRAM prepped handle: %s", exc) |
| 455 | self._dram_prepped_handle = None |
| 456 | if self._primary_reg is not None: |
| 457 | try: |
| 458 | self._agent.deregister_memory(self._primary_reg) |
| 459 | except Exception as exc: |
| 460 | logger.warning("failed to deregister primary buffer: %s", exc) |
| 461 | self._primary_reg = None |