| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """TieringOffloadingManager: Multi-tier KV cache offloading orchestrator. |
| 4 | |
| 5 | This manager coordinates between a CPU primary tier (with direct GPU access) |
| 6 | and zero or more secondary tiers (Storage, Network, etc.) to provide |
| 7 | hierarchical KV cache offloading. |
| 8 | |
| 9 | Key Design Principles: |
| 10 | 1. Always offload to all tiers — When a chunk is stored to the primary tier, |
| 11 | it is cascaded to ALL secondary tiers |
| 12 | 2. Primary tier is the gateway — Secondary tiers cannot access GPU memory |
| 13 | directly; all data flows through the CPU primary tier |
| 14 | 3. Staged promotion — Chunks in secondary tiers must be promoted to the |
| 15 | primary tier before GPU can access them |
| 16 | 4. Transparent retry mechanism — Return None from lookup() to signal |
| 17 | "data is being promoted, try later" |
| 18 | 5. ref_cnt as eviction protection — primary.prepare_read() increments ref_cnt, |
| 19 | protecting chunks from eviction until complete_read() is called |
| 20 | """ |
| 21 | |
| 22 | import time |
| 23 | from collections.abc import Collection, Iterable, Sequence |
| 24 | from dataclasses import dataclass, field |
| 25 | from typing import NamedTuple |
| 26 | |
| 27 | import numpy as np |
| 28 | from typing_extensions import override |
| 29 | |
| 30 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( |
| 31 | OffloadingConnectorStats, |
| 32 | ) |
| 33 | from vllm.logger import init_logger |
| 34 | from vllm.v1.kv_offload.base import ( |
| 35 | LoadStoreSpec, |
| 36 | LookupResult, |
| 37 | OffloadingEvent, |
| 38 | OffloadingManager, |
| 39 | OffloadKey, |
| 40 | OffloadPolicy, |
| 41 | PrepareStoreOutput, |
| 42 | ReqContext, |
| 43 | RequestOffloadingContext, |
| 44 | ScheduleEndContext, |
| 45 | ) |
| 46 | from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec |
| 47 | from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager |
| 48 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 49 | from vllm.v1.kv_offload.tiering.base import ( |
| 50 | JobId, |
| 51 | JobResult, |
| 52 | ParentManager, |
| 53 | SecondaryTierManager, |
| 54 | TransferJob, |
| 55 | ) |
| 56 | from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker |
| 57 | |
| 58 | logger = init_logger(__name__) |
| 59 | |
| 60 | |
| 61 | @dataclass |
| 62 | class PendingPromotion: |
| 63 | """Accumulator for chunks awaiting submit_load() for one (tier, request).""" |
| 64 | |
| 65 | req_context: ReqContext |
| 66 | keys: list[OffloadKey] = field(default_factory=list) |
| 67 | chunk_ids: list[int] = field(default_factory=list) |
| 68 | |
| 69 | |
| 70 | @dataclass(slots=True) |
| 71 | class RequestState: |
| 72 | req_context: ReqContext |
| 73 | pending_primary_stores: int = 0 |
| 74 | is_finished: bool = False |
| 75 | request_level_tiers: set[int] | None = None |
| 76 | pending_cascade_keys: list[OffloadKey] = field(default_factory=list) |
| 77 | |
| 78 | |
| 79 | class JobMetadata(NamedTuple): |
| 80 | transfer_job: TransferJob |
| 81 | tier_idx: int |
| 82 | |
| 83 | |
| 84 | class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): |
| 85 | """CPUOffloadingManager with a primary/secondary transfer interface. |
| 86 | |
| 87 | The inherited prepare_store/complete_store/prepare_load/complete_load are the |
| 88 | GPU-facing OffloadingManager interface. These aliases expose the same operations |
| 89 | from the secondary tier perspective, where read/write refers to secondary |
| 90 | accessing primary. This avoids confusion when reading TieringOffloadingManager |
| 91 | code (e.g. calling prepare_load inside a cascade/store path would be misleading). |
| 92 | """ |
| 93 | |
| 94 | def __init__( |
| 95 | self, |
| 96 | num_chunks: int, |
| 97 | mmap_region: SharedOffloadRegion, |
| 98 | cache_policy: str = "lru", |
| 99 | cache_policy_module_path: str | None = None, |
| 100 | enable_events: bool = False, |
| 101 | ): |
| 102 | super().__init__( |
| 103 | num_chunks=num_chunks, |
| 104 | cache_policy=cache_policy, |
| 105 | cache_policy_module_path=cache_policy_module_path, |
| 106 | enable_events=enable_events, |
| 107 | ) |
| 108 | self._mmap_region = mmap_region |
| 109 | # read/write is for CPU<->secondary transfers, |
| 110 | # load/store is for CPU<->GPU transfers. |
| 111 | # These aliases avoid calling prepare_load inside a store path. |
| 112 | self.complete_read = self.complete_load |
| 113 | self.prepare_write = self.prepare_store |
| 114 | self.complete_write = self.complete_store |
| 115 | |
| 116 | self._kv_memoryview = mmap_region.create_kv_memoryview() |
| 117 | |
| 118 | def prepare_read( |
| 119 | self, keys: Collection[OffloadKey], req_context: ReqContext |
| 120 | ) -> LoadStoreSpec: |
| 121 | """Pin chunks for a CPU-to-secondary transfer. |
| 122 | |
| 123 | Cascade reads are implementation details of tiering, not additional |
| 124 | request accesses, so they must not alter request-scoped recency. |
| 125 | """ |
| 126 | return self._prepare_load(keys, req_context, record_access=False) |
| 127 | |
| 128 | def get_kv_memoryview(self) -> memoryview: |
| 129 | """Return the memoryview over the primary tier's KV cache buffer. |
| 130 | |
| 131 | The view has shape (num_chunks, row_stride_bytes) and is backed by the |
| 132 | SharedOffloadRegion mmap. Secondary tiers address chunk *c* as |
| 133 | ``view[c]``. |
| 134 | """ |
| 135 | return self._kv_memoryview |
| 136 | |
| 137 | @override |
| 138 | def shutdown(self) -> None: |
| 139 | super().shutdown() |
| 140 | self._kv_memoryview.release() |
| 141 | self._mmap_region.cleanup() |
| 142 | |
| 143 | |
| 144 | class _SecondaryTierFacingParent(ParentManager): |
| 145 | """Wrapper that implements ParentManager by delegating to the |
| 146 | TieringOffloadingManager with exclude_tier_idx set to the origin tier.""" |
| 147 | |
| 148 | __slots__ = ("_m", "_origin_idx") |
| 149 | |
| 150 | def __init__( |
| 151 | self, |
| 152 | manager: "TieringOffloadingManager", |
| 153 | tier_idx: int, |
| 154 | ): |
| 155 | self._m = manager |
| 156 | self._origin_idx = tier_idx |
| 157 | |
| 158 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: |
| 159 | return self._m.on_new_request(req_context, exclude_tier_idx=self._origin_idx) |
| 160 | |
| 161 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: |
| 162 | return self._m.lookup(key, req_context, exclude_tier_idx=self._origin_idx) |
| 163 | |
| 164 | def create_store_job( |
| 165 | self, keys: Collection[OffloadKey], req_context: ReqContext |
| 166 | ) -> TransferJob: |
| 167 | return self._m.create_store_job(keys, req_context, self._origin_idx) |
| 168 | |
| 169 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 170 | return self._m.on_request_finished( |
| 171 | req_context, exclude_tier_idx=self._origin_idx |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | class TieringOffloadingManager(OffloadingManager): |
| 176 | """Orchestrates multi-tier KV cache offloading. |
| 177 | |
| 178 | This manager coordinates between a CPU primary tier (with direct GPU access) |
| 179 | and zero or more secondary tiers (Storage, Network, etc.) to provide |
| 180 | hierarchical KV cache offloading. |
| 181 | |
| 182 | Key internal state: |
| 183 | - Minimal state tracking; relies on secondary tiers to report completion |
| 184 | via get_finished_jobs() |
| 185 | - Secondary tiers return JobResult objects containing all necessary |
| 186 | information |
| 187 | - job_id_counter: monotonically increasing counter for job IDs |
| 188 | """ |
| 189 | |
| 190 | def __init__( |
| 191 | self, |
| 192 | primary_tier: CPUPrimaryTierOffloadingManager, |
| 193 | secondary_tiers: list[SecondaryTierManager] | None = None, |
| 194 | ): |
| 195 | """Initialize the TieringOffloadingManager. |
| 196 | |
| 197 | Args: |
| 198 | primary_tier: The primary tier manager (CPU-based). |
| 199 | secondary_tiers: List of secondary tier managers (e.g., Storage, |
| 200 | Network). Can be None or empty list. |
| 201 | |
| 202 | """ |
| 203 | self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier |
| 204 | self.secondary_tiers = secondary_tiers or [] |
| 205 | |
| 206 | self._job_id_counter: int = 0 |
| 207 | # Job tracking: maps job_id to metadata for all in-flight transfers. |
| 208 | # TransferJob.is_promotion distinguishes direction: |
| 209 | # True: secondary → primary (promotion) |
| 210 | # False: primary → secondary (cascade) |
| 211 | self._jobs: dict[JobId, JobMetadata] = {} |
| 212 | primary_view = self.primary_tier.get_kv_memoryview() |
| 213 | assert primary_view.strides is not None |
| 214 | self._metrics = TieringMetricsTracker( |
| 215 | tier_types=[tier.tier_type for tier in self.secondary_tiers], |
| 216 | num_primary_chunks=self.primary_tier._num_chunks, |
| 217 | primary_chunk_size=primary_view.strides[0], |
| 218 | ) |
| 219 | |
| 220 | # Pending promotion requests accumulated during lookup() calls; flushed |
| 221 | # as one batched submit_load() per (tier, request) in on_schedule_end(). |
| 222 | # Outer key: tier index. Inner key: req_context.req_id — the same ReqContext |
| 223 | # object is reused for all chunk lookups of a given request per engine step. |
| 224 | self._pending_load_submissions: dict[int, dict[str, PendingPromotion]] = {} |
| 225 | |
| 226 | # Gate for once-per-step execution of _maybe_process_finished_jobs(). |
| 227 | # Reset at the end of each step in on_schedule_end(). |
| 228 | self._processed_jobs_this_step: bool = False |
| 229 | |
| 230 | # Per-request state for prepared GPU->primary stores and finalization. |
| 231 | # Secondary tiers are finalized only after pending primary stores reach |
| 232 | # complete_store(), since complete_store() can still submit cascades. |
| 233 | self._req_state: dict[str, RequestState] = {} |
| 234 | |
| 235 | # Cached ParentManager wrappers for each secondary tier. |
| 236 | self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { |
| 237 | tier: _SecondaryTierFacingParent(self, tier_idx) |
| 238 | for tier_idx, tier in enumerate(self.secondary_tiers) |
| 239 | } |
| 240 | |
| 241 | self._tier_index: dict[SecondaryTierManager, int] = { |
| 242 | tier: i for i, tier in enumerate(self.secondary_tiers) |
| 243 | } |
| 244 | |
| 245 | @property |
| 246 | def _transfer_jobs(self) -> dict[JobId, JobMetadata]: |
| 247 | return self._jobs |
| 248 | |
| 249 | def _next_job_id(self) -> JobId: |
| 250 | """Generate a unique job ID for async transfer tracking.""" |
| 251 | job_id = self._job_id_counter |
| 252 | self._job_id_counter += 1 |
| 253 | return job_id |
| 254 | |
| 255 | def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None: |
| 256 | job_metadata = JobMetadata(transfer_job, tier_idx) |
| 257 | self._jobs[transfer_job.job_id] = job_metadata |
| 258 | self._metrics.on_job_registered(job_metadata) |
| 259 | |
| 260 | def _pop_job(self, job_id: JobId) -> JobMetadata | None: |
| 261 | return self._jobs.pop(job_id, None) |
| 262 | |
| 263 | def _maybe_process_finished_jobs(self): |
| 264 | """Poll secondary tiers for completed jobs (at most once per step). |
| 265 | |
| 266 | Guarded by _processed_jobs_this_step: the first call in an engine step |
| 267 | does the actual polling; subsequent calls are no-ops. The flag is reset |
| 268 | in on_schedule_end() at the end of each step. |
| 269 | """ |
| 270 | if self._processed_jobs_this_step: |
| 271 | return |
| 272 | self._processed_jobs_this_step = True |
| 273 | self._process_finished_jobs() |
| 274 | |
| 275 | def _complete_promotion( |
| 276 | self, job_metadata: JobMetadata, completed_job: JobResult |
| 277 | ) -> None: |
| 278 | transfer_job = job_metadata.transfer_job |
| 279 | successful_keys = completed_job.successful_keys |
| 280 | failed_keys: Collection[OffloadKey] |
| 281 | if completed_job.success: |
| 282 | successful_keys = transfer_job.keys |
| 283 | failed_keys = () |
| 284 | elif successful_keys: |
| 285 | failed_keys_set = set(transfer_job.keys) |
| 286 | assert failed_keys_set.issuperset(successful_keys), ( |
| 287 | f"Finished promotion job_id {completed_job.job_id} " |
| 288 | "reported unknown successful keys" |
| 289 | ) |
| 290 | failed_keys_set.difference_update(successful_keys) |
| 291 | failed_keys = failed_keys_set |
| 292 | else: |
| 293 | successful_keys = () |
| 294 | failed_keys = transfer_job.keys |
| 295 | |
| 296 | if successful_keys: |
| 297 | self.primary_tier.complete_write( |
| 298 | successful_keys, |
| 299 | transfer_job.req_context, |
| 300 | True, |
| 301 | ) |
| 302 | if failed_keys: |
| 303 | self.primary_tier.complete_write( |
| 304 | failed_keys, |
| 305 | transfer_job.req_context, |
| 306 | False, |
| 307 | ) |
| 308 | |
| 309 | def _process_finished_jobs(self): |
| 310 | """Unconditionally poll all secondary tiers for completed jobs. |
| 311 | |
| 312 | This method: |
| 313 | 1. Calls get_finished_jobs() on each secondary tier |
| 314 | 2. For completed stores (primary→secondary): calls primary.complete_read() |
| 315 | to decrement ref_cnt |
| 316 | 3. For completed loads (secondary→primary): calls primary.complete_write() |
| 317 | to make chunks available |
| 318 | """ |
| 319 | for i, tier in enumerate(self.secondary_tiers): |
| 320 | for completed_job in tier.get_finished_jobs(): |
| 321 | job_id = completed_job.job_id |
| 322 | job_metadata = self._pop_job(job_id) |
| 323 | assert job_metadata is not None, ( |
| 324 | f"Finished job_id {job_id} from tier #{i}" |
| 325 | f" ({tier.tier_type}) not in _jobs" |
| 326 | ) |
| 327 | assert job_metadata.tier_idx == i, ( |
| 328 | f"Finished job_id {job_id} reported by tier #{i}" |
| 329 | f" but belongs to tier #{job_metadata.tier_idx}" |
| 330 | ) |
| 331 | transfer_job = job_metadata.transfer_job |
| 332 | self._metrics.on_job_finished(job_metadata, completed_job) |
| 333 | |
| 334 | if transfer_job.is_promotion: |
| 335 | # secondary→primary transfer (promotion) completed. |
| 336 | # Make chunks available in primary tier. |
| 337 | self._complete_promotion(job_metadata, completed_job) |
| 338 | else: |
| 339 | # primary→secondary transfer completed. |
| 340 | # Decrement ref_cnt on primary chunks. |
| 341 | self.primary_tier.complete_read( |
| 342 | transfer_job.keys, transfer_job.req_context |
| 343 | ) |
| 344 | if completed_job.success: |
| 345 | self._update_backpressure(tier, job_metadata, completed_job) |
| 346 | |
| 347 | def _should_store_to_tier( |
| 348 | self, tier: SecondaryTierManager, num_blocks: int |
| 349 | ) -> bool: |
| 350 | detector = tier.bp_detector |
| 351 | if detector is None: |
| 352 | return True |
| 353 | return detector.should_store(num_blocks) |
| 354 | |
| 355 | def _update_backpressure( |
| 356 | self, |
| 357 | tier: SecondaryTierManager, |
| 358 | job_metadata: JobMetadata, |
| 359 | completed_job: JobResult, |
| 360 | ) -> None: |
| 361 | detector = tier.bp_detector |
| 362 | if detector is None: |
| 363 | return |
| 364 | was_under_pressure = detector.is_under_pressure() |
| 365 | tj = job_metadata.transfer_job |
| 366 | num_bytes = ( |
| 367 | completed_job.transfer_bytes |
| 368 | if completed_job.transfer_bytes is not None |
| 369 | else len(tj.keys) * tier.block_size_bytes |
| 370 | ) |
| 371 | detector.update(tj.submit_time, num_bytes) |
| 372 | if detector.is_under_pressure() != was_under_pressure: |
| 373 | tier_idx = self._tier_index[tier] |
| 374 | logger.info( |
| 375 | "Tier #%d (%s) back-pressure %s (stats=%s)", |
| 376 | tier_idx, |
| 377 | tier.tier_type, |
| 378 | "activated" if detector.is_under_pressure() else "cleared", |
| 379 | detector.stats, |
| 380 | ) |
| 381 | |
| 382 | @override |
| 383 | def lookup( |
| 384 | self, |
| 385 | key: OffloadKey, |
| 386 | req_context: ReqContext, |
| 387 | *, |
| 388 | exclude_tier_idx: int | None = None, |
| 389 | ) -> LookupResult: |
| 390 | """Check whether a single chunk is offloaded and ready. |
| 391 | |
| 392 | Algorithm: |
| 393 | 1. Process any completed async jobs first. |
| 394 | 2. Query primary tier — short-circuit on hit or in-flight. |
| 395 | 3. On primary miss, query secondary tiers — stop on first |
| 396 | hit and initiate promotion. |
| 397 | |
| 398 | Args: |
| 399 | key: Chunk hash to look up. |
| 400 | req_context: Per-request context. |
| 401 | exclude_tier_idx: Skip this tier index during the lookup. |
| 402 | |
| 403 | Returns: |
| 404 | HIT — chunk is ready in the primary tier. |
| 405 | HIT_PENDING — chunk found but not yet readable (write |
| 406 | in-flight on the primary tier). |
| 407 | RETRY — promotion started or a secondary tier is busy. |
| 408 | MISS — chunk not found in any tier, or primary is full |
| 409 | and cannot accept a promotion. |
| 410 | |
| 411 | """ |
| 412 | # Poll first so a promotion that finished since the last call is |
| 413 | # already reflected as HIT (not stale HIT_PENDING/MISS) below, and |
| 414 | # so chunks freed by cascade or promotion completions are evictable |
| 415 | # in time for a promotion this lookup may initiate. |
| 416 | self._maybe_process_finished_jobs() |
| 417 | |
| 418 | start_time = time.monotonic() |
| 419 | primary_hit = self.primary_tier.lookup(key, req_context) |
| 420 | lookup_duration = time.monotonic() - start_time |
| 421 | self._metrics.on_lookup( |
| 422 | req_context, |
| 423 | key, |
| 424 | self._metrics.primary_tier_label, |
| 425 | primary_hit, |
| 426 | lookup_duration, |
| 427 | ) |
| 428 | if primary_hit is LookupResult.HIT: |
| 429 | return LookupResult.HIT |
| 430 | if primary_hit is LookupResult.HIT_PENDING: |
| 431 | return LookupResult.HIT_PENDING |
| 432 | |
| 433 | any_retry = False |
| 434 | for i, tier in enumerate(self.secondary_tiers): |
| 435 | if i == exclude_tier_idx: |
| 436 | continue |
| 437 | if not req_context.load_tier_filter.allows(tier.medium, tier.locality): |
| 438 | continue |
| 439 | labelvalues = self._metrics.tier_label(i) |
| 440 | start_time = time.monotonic() |
| 441 | result = tier.lookup(key, req_context) |
| 442 | lookup_duration = time.monotonic() - start_time |
| 443 | if result is LookupResult.HIT: |
| 444 | self._metrics.on_lookup( |
| 445 | req_context, |
| 446 | key, |
| 447 | labelvalues, |
| 448 | result, |
| 449 | lookup_duration, |
| 450 | ) |
| 451 | promoted = self._initiate_promotion(i, key, req_context) |
| 452 | return LookupResult.MISS if not promoted else LookupResult.HIT_PENDING |
| 453 | if result is LookupResult.RETRY: |
| 454 | any_retry = True |
| 455 | self._metrics.on_lookup( |
| 456 | req_context, |
| 457 | key, |
| 458 | labelvalues, |
| 459 | result, |
| 460 | lookup_duration, |
| 461 | ) |
| 462 | |
| 463 | if any_retry: |
| 464 | return LookupResult.RETRY |
| 465 | return LookupResult.MISS |
| 466 | |
| 467 | def _initiate_promotion( |
| 468 | self, |
| 469 | tier_idx: int, |
| 470 | key: OffloadKey, |
| 471 | req_context: ReqContext, |
| 472 | ) -> bool: |
| 473 | """Queue a chunk for promotion from a secondary tier to the primary tier. |
| 474 | |
| 475 | Allocates space in the primary tier immediately (sets ref_cnt=-1 so |
| 476 | subsequent lookups within the same step see the slot as in-flight), |
| 477 | then defers the actual submit_load() call to _flush_pending_promotions() |
| 478 | so all chunks queued during one engine step are submitted as a single |
| 479 | batched job. |
| 480 | |
| 481 | Args: |
| 482 | tier_idx: The secondary tier index to promote from |
| 483 | key: Chunk to promote |
| 484 | req_context: Per-request context forwarded to primary.prepare_write(). |
| 485 | |
| 486 | Returns: |
| 487 | True if promotion was initiated, False if primary tier is full. |
| 488 | |
| 489 | """ |
| 490 | # Allocate space in primary tier for promoted chunk. |
| 491 | # Must happen immediately so primary.lookup() returns None (in-flight) |
| 492 | # for this key on any subsequent lookup() call within the same step, |
| 493 | # preventing duplicate promotion attempts. |
| 494 | primary_write_result = self.primary_tier.prepare_write([key], req_context) |
| 495 | |
| 496 | if primary_write_result is None: |
| 497 | # Primary tier is full; caller should treat the chunk as unavailable |
| 498 | # rather than retrying indefinitely. |
| 499 | self._metrics.on_promotion_allocation_failure() |
| 500 | return False |
| 501 | |
| 502 | store_spec = primary_write_result.store_spec |
| 503 | assert isinstance(store_spec, CPULoadStoreSpec) |
| 504 | # Defer submit_load to on_schedule_end(). Group by (tier, request) so |
| 505 | # each request's chunks are submitted as one batched job per tier. |
| 506 | tier_pending = self._pending_load_submissions.setdefault(tier_idx, {}) |
| 507 | ctx_id = req_context.req_id |
| 508 | if ctx_id not in tier_pending: |
| 509 | tier_pending[ctx_id] = PendingPromotion( |
| 510 | keys=[], chunk_ids=[], req_context=req_context |
| 511 | ) |
| 512 | entry = tier_pending[ctx_id] |
| 513 | entry.keys.extend(primary_write_result.keys_to_store) |
| 514 | entry.chunk_ids.extend(store_spec.chunk_ids) |
| 515 | return True |
| 516 | |
| 517 | def _flush_pending_promotions(self) -> None: |
| 518 | """Submit one batched submit_load() per (tier, request). |
| 519 | |
| 520 | Called from on_schedule_end() at the end of each scheduler step, |
| 521 | flushing all promotion requests deferred during lookup(). |
| 522 | """ |
| 523 | if not self._pending_load_submissions: |
| 524 | return |
| 525 | |
| 526 | for tier_idx, pending_by_ctx in self._pending_load_submissions.items(): |
| 527 | tier = self.secondary_tiers[tier_idx] |
| 528 | for entry in pending_by_ctx.values(): |
| 529 | job_id = self._next_job_id() |
| 530 | job_metadata = TransferJob( |
| 531 | job_id=job_id, |
| 532 | keys=entry.keys, |
| 533 | chunk_ids=np.array(entry.chunk_ids, dtype=np.int32), |
| 534 | is_promotion=True, |
| 535 | req_context=entry.req_context, |
| 536 | ) |
| 537 | self._register_job(job_metadata, tier_idx) |
| 538 | tier.submit_load(job_metadata) |
| 539 | |
| 540 | self._pending_load_submissions.clear() |
| 541 | |
| 542 | @override |
| 543 | def prepare_load( |
| 544 | self, keys: Collection[OffloadKey], req_context: ReqContext |
| 545 | ) -> LoadStoreSpec: |
| 546 | """Prepare chunks to be loaded from primary tier to GPU. |
| 547 | |
| 548 | Callers only pass keys already confirmed HIT by lookup() earlier this |
| 549 | step. |
| 550 | |
| 551 | This increments ref_cnt on the chunks in the primary tier, protecting |
| 552 | them from eviction during the transfer. |
| 553 | |
| 554 | Args: |
| 555 | keys: Chunks to prepare for loading. |
| 556 | req_context: Per-request context. |
| 557 | |
| 558 | Returns: |
| 559 | LoadStoreSpec for reading from primary tier. |
| 560 | |
| 561 | """ |
| 562 | return self.primary_tier.prepare_load(keys, req_context) |
| 563 | |
| 564 | @override |
| 565 | def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): |
| 566 | """Mark chunks as recently used in all tiers. |
| 567 | |
| 568 | Args: |
| 569 | keys: Chunks to mark as recently used. |
| 570 | req_context: Per-request context. |
| 571 | |
| 572 | """ |
| 573 | self.primary_tier.touch(keys, req_context) |
| 574 | for tier in self.secondary_tiers: |
| 575 | tier.touch(keys, req_context) |
| 576 | |
| 577 | @override |
| 578 | def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): |
| 579 | """Mark chunks as done loading from primary tier to GPU. |
| 580 | |
| 581 | This decrements ref_cnt on the chunks in the primary tier, allowing |
| 582 | them to be evicted again. |
| 583 | |
| 584 | Args: |
| 585 | keys: Chunks that finished loading. |
| 586 | req_context: Per-request context. |
| 587 | |
| 588 | """ |
| 589 | self.primary_tier.complete_load(keys, req_context) |
| 590 | |
| 591 | @override |
| 592 | def prepare_store( |
| 593 | self, keys: Collection[OffloadKey], req_context: ReqContext |
| 594 | ) -> PrepareStoreOutput | None: |
| 595 | """Prepare chunks to be stored from GPU to primary tier. |
| 596 | |
| 597 | CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure |
| 598 | that any completed async transfers have their ref_cnt decremented |
| 599 | before the primary tier makes eviction decisions. |
| 600 | |
| 601 | For request-level tiers, chunks already present in the primary tier |
| 602 | are immediately cascaded via submit_store(). |
| 603 | |
| 604 | Args: |
| 605 | keys: Chunks to prepare for storing. |
| 606 | req_context: Per-request context. |
| 607 | |
| 608 | Returns: |
| 609 | PrepareStoreOutput describing where to store chunks and what was |
| 610 | evicted, or None if store cannot proceed. |
| 611 | |
| 612 | """ |
| 613 | # Step 1: Poll for completed async jobs FIRST |
| 614 | # _process_finished_jobs() handles two kinds of completions here: |
| 615 | # - Cascade completions (store to a secondary tier, either a local |
| 616 | # cascade or a store job created for a remote requester via |
| 617 | # create_store_job()): decrements ref_cnt on the primary chunks |
| 618 | # that were read, making them evictable again once ref_cnt hits 0. |
| 619 | # - Promotion completions (secondary->primary loads): sets a |
| 620 | # not-yet-ready chunk's ref_cnt from -1 to 0 via complete_write(), |
| 621 | # making it evictable for the first time. |
| 622 | # Both must be accounted for before the eviction decision below. |
| 623 | self._maybe_process_finished_jobs() |
| 624 | |
| 625 | # Step 2: Store to primary tier (new chunks only). |
| 626 | # Cascading of these newly-stored chunks to ALL secondary tiers |
| 627 | # happens later in complete_store(), after the GPU→Primary transfer |
| 628 | # completes. |
| 629 | primary_result = self.primary_tier.prepare_store(keys, req_context) |
| 630 | |
| 631 | if primary_result is None: |
| 632 | return None |
| 633 | |
| 634 | if primary_result.keys_to_store: |
| 635 | state = self._req_state[req_context.req_id] |
| 636 | state.pending_primary_stores += 1 |
| 637 | |
| 638 | # Step 3: For request-level tiers, cascade chunks already in primary |
| 639 | request_level_tiers = self._req_state[req_context.req_id].request_level_tiers |
| 640 | if request_level_tiers: |
| 641 | keys_to_store_set = set(primary_result.keys_to_store) |
| 642 | keys_already_in_primary = tuple( |
| 643 | k for k in keys if k not in keys_to_store_set |
| 644 | ) |
| 645 | if keys_already_in_primary: |
| 646 | self._cascade_existing_chunks_to_request_level_tiers( |
| 647 | keys_already_in_primary, req_context, request_level_tiers |
| 648 | ) |
| 649 | |
| 650 | return primary_result |
| 651 | |
| 652 | def _cascade_existing_chunks_to_request_level_tiers( |
| 653 | self, |
| 654 | keys: Sequence[OffloadKey], |
| 655 | req_context: ReqContext, |
| 656 | request_level_tiers: set[int], |
| 657 | ) -> None: |
| 658 | """For tiers that requested request-level policy, submit_store() for |
| 659 | chunks that are already present in the primary tier. |
| 660 | |
| 661 | A key whose primary write is still in flight (HIT_PENDING) cannot be |
| 662 | dropped: prepare_store already excluded it as present, and the |
| 663 | scheduler advances past its chunk, so no path offers it again. Park it |
| 664 | instead. MISS keys are dropped, since nothing is there to read. |
| 665 | |
| 666 | The primary tier resolves every key it holds, so RETRY cannot reach |
| 667 | here. Parking on it would have no guarantee of ever draining, which is |
| 668 | what makes parking HIT_PENDING safe, so it is rejected rather than |
| 669 | guessed at. |
| 670 | """ |
| 671 | state = self._req_state[req_context.req_id] |
| 672 | ready_keys = [] |
| 673 | for key in keys: |
| 674 | result = self.primary_tier.lookup(key, req_context) |
| 675 | if result is LookupResult.HIT: |
| 676 | ready_keys.append(key) |
| 677 | elif result is LookupResult.HIT_PENDING: |
| 678 | state.pending_cascade_keys.append(key) |
| 679 | else: |
| 680 | assert result is LookupResult.MISS, ( |
| 681 | f"primary tier returned {result} for a cascade key" |
| 682 | ) |
| 683 | if not ready_keys: |
| 684 | return |
| 685 | |
| 686 | for tier_idx in request_level_tiers: |
| 687 | tier = self.secondary_tiers[tier_idx] |
| 688 | if not self._should_store_to_tier(tier, len(ready_keys)): |
| 689 | continue |
| 690 | job_metadata = self.create_store_job(ready_keys, req_context, tier_idx) |
| 691 | tier.submit_store(job_metadata) |
| 692 | |
| 693 | def _flush_pending_cascades(self) -> None: |
| 694 | """Retry request-level cascades parked on an in-flight primary write. |
| 695 | |
| 696 | A parked key always resolves, to HIT or to MISS, so the set drains and |
| 697 | a request cannot be held from finalization forever. |
| 698 | """ |
| 699 | for req_id, state in list(self._req_state.items()): |
| 700 | if not state.pending_cascade_keys: |
| 701 | continue |
| 702 | assert state.request_level_tiers |
| 703 | keys, state.pending_cascade_keys = state.pending_cascade_keys, [] |
| 704 | self._cascade_existing_chunks_to_request_level_tiers( |
| 705 | keys, state.req_context, state.request_level_tiers |
| 706 | ) |
| 707 | self._maybe_finalize_request(req_id) |
| 708 | |
| 709 | @override |
| 710 | def complete_store( |
| 711 | self, |
| 712 | keys: Collection[OffloadKey], |
| 713 | req_context: ReqContext, |
| 714 | success: bool = True, |
| 715 | ) -> None: |
| 716 | """Mark chunks as done storing from GPU to primary tier. |
| 717 | |
| 718 | This is where secondary tier cascading happens — after chunks are |
| 719 | confirmed to be in the primary tier, they are cascaded to ALL |
| 720 | secondary tiers. |
| 721 | |
| 722 | For each secondary tier: |
| 723 | 1. Call primary.prepare_read() to get LoadStoreSpec AND increment |
| 724 | ref_cnt (protecting chunks during async transfer) |
| 725 | 2. Call tier.submit_store() to start async transfer: primary→secondary |
| 726 | 3. Track the job in _store_jobs dictionary |
| 727 | |
| 728 | Args: |
| 729 | keys: Chunks that finished storing. |
| 730 | success: Whether the GPU→primary transfer succeeded. |
| 731 | req_context: Per-request context forwarded to primary.prepare_read(). |
| 732 | |
| 733 | """ |
| 734 | # Step 1: Complete store in primary tier (makes chunks loadable) |
| 735 | self.primary_tier.complete_store(keys, req_context, success) |
| 736 | |
| 737 | if success: |
| 738 | # Step 2: Cascade to ALL secondary tiers |
| 739 | # For each secondary tier, call primary.prepare_read() to get the |
| 740 | # LoadStoreSpec AND to increment ref_cnt (protecting chunks from |
| 741 | # eviction during the async transfer). One prepare_read() call per |
| 742 | # secondary tier. |
| 743 | for tier_idx, tier in enumerate(self.secondary_tiers): |
| 744 | if not self._should_store_to_tier(tier, len(keys)): |
| 745 | continue |
| 746 | job_metadata = self.create_store_job(keys, req_context, tier_idx) |
| 747 | tier.submit_store(job_metadata) |
| 748 | |
| 749 | # Note: The async transfers are now in flight. Their completion is |
| 750 | # tracked via get_finished_jobs() / _maybe_process_finished_jobs(). |
| 751 | req_id = req_context.req_id |
| 752 | state = self._req_state[req_id] |
| 753 | assert state.pending_primary_stores > 0 |
| 754 | state.pending_primary_stores -= 1 |
| 755 | self._maybe_finalize_request(req_id) |
| 756 | |
| 757 | def create_store_job( |
| 758 | self, |
| 759 | keys: Collection[OffloadKey], |
| 760 | req_context: ReqContext, |
| 761 | tier_idx: int = 0, |
| 762 | ) -> TransferJob: |
| 763 | """Pin chunks in the primary tier and create a tracked store job. |
| 764 | |
| 765 | Calls prepare_read() to increment ref_cnt (protecting chunks |
| 766 | from eviction during the async transfer), allocates a job ID, |
| 767 | and registers the job in _jobs. |
| 768 | |
| 769 | The caller is responsible for the actual data transfer and |
| 770 | reporting completion via get_finished_jobs(). |
| 771 | """ |
| 772 | primary_chunks_spec = self.primary_tier.prepare_read(keys, req_context) |
| 773 | assert isinstance(primary_chunks_spec, CPULoadStoreSpec) |
| 774 | job_id = self._next_job_id() |
| 775 | job_metadata = TransferJob( |
| 776 | job_id=job_id, |
| 777 | keys=keys, |
| 778 | chunk_ids=primary_chunks_spec.chunk_ids, |
| 779 | is_promotion=False, |
| 780 | req_context=req_context, |
| 781 | ) |
| 782 | self._register_job(job_metadata, tier_idx) |
| 783 | return job_metadata |
| 784 | |
| 785 | @override |
| 786 | def on_new_request( |
| 787 | self, |
| 788 | req_context: ReqContext, |
| 789 | *, |
| 790 | exclude_tier_idx: int | None = None, |
| 791 | ) -> RequestOffloadingContext: |
| 792 | """Query each secondary tier for its offload policy preference. |
| 793 | |
| 794 | Returns REQUEST_LEVEL if ANY secondary tier wants request-level. |
| 795 | Only stores REQUEST_LEVEL tier decisions for use in prepare_store. |
| 796 | """ |
| 797 | state = RequestState(req_context=req_context) |
| 798 | self._metrics.on_new_request(req_context) |
| 799 | for tier_idx, tier in enumerate(self.secondary_tiers): |
| 800 | if tier_idx == exclude_tier_idx: |
| 801 | continue |
| 802 | tier_ctx = tier.on_new_request(req_context) |
| 803 | if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL: |
| 804 | if state.request_level_tiers is None: |
| 805 | state.request_level_tiers = set() |
| 806 | state.request_level_tiers.add(tier_idx) |
| 807 | self._req_state[req_context.req_id] = state |
| 808 | |
| 809 | policy = ( |
| 810 | OffloadPolicy.REQUEST_LEVEL |
| 811 | if state.request_level_tiers |
| 812 | else OffloadPolicy.CHUNK_LEVEL |
| 813 | ) |
| 814 | return RequestOffloadingContext(policy=policy) |
| 815 | |
| 816 | @override |
| 817 | def on_request_finished( |
| 818 | self, |
| 819 | req_context: ReqContext, |
| 820 | *, |
| 821 | exclude_tier_idx: int | None = None, |
| 822 | ) -> None: |
| 823 | self.primary_tier.on_request_finished(req_context) |
| 824 | state = self._req_state[req_context.req_id] |
| 825 | state.is_finished = True |
| 826 | self._maybe_finalize_request(req_context.req_id, exclude_tier_idx) |
| 827 | |
| 828 | def _maybe_finalize_request( |
| 829 | self, |
| 830 | req_id: str, |
| 831 | exclude_tier_idx: int | None = None, |
| 832 | ) -> None: |
| 833 | """Finalize secondary tiers once no more cascades can be submitted. |
| 834 | |
| 835 | Their finalization is delayed until pending GPU->primary stores |
| 836 | finish, since those callbacks may still submit secondary stores. |
| 837 | """ |
| 838 | state = self._req_state[req_id] |
| 839 | if not state.is_finished: |
| 840 | return |
| 841 | if state.pending_primary_stores != 0: |
| 842 | return |
| 843 | if state.pending_cascade_keys: |
| 844 | return |
| 845 | |
| 846 | for tier_idx, tier in enumerate(self.secondary_tiers): |
| 847 | if tier_idx == exclude_tier_idx: |
| 848 | continue |
| 849 | tier.on_request_finished(state.req_context) |
| 850 | self._metrics.on_request_finished(state.req_context) |
| 851 | del self._req_state[req_id] |
| 852 | |
| 853 | @override |
| 854 | def on_schedule_end(self, context: ScheduleEndContext) -> None: |
| 855 | """End-of-schedule hook: process finished jobs, flush deferred |
| 856 | promotions, and reset the per-step gate. |
| 857 | |
| 858 | Called once per scheduler step from |
| 859 | OffloadingConnectorScheduler.build_connector_meta(). |
| 860 | """ |
| 861 | # Catch-all poll: guarantees jobs are processed even on steps where |
| 862 | # lookup()/prepare_store() were never called (e.g. no requests |
| 863 | # scheduled but a tier still has_pending_work()). |
| 864 | self._maybe_process_finished_jobs() |
| 865 | |
| 866 | for tier in self.secondary_tiers: |
| 867 | tier.serve_external_requests(self._tier_parents[tier]) |
| 868 | |
| 869 | # Reset the per-step gate AFTER serve_external_requests so that |
| 870 | # lookup() calls within it skip redundant _process_finished_jobs(). |
| 871 | self._processed_jobs_this_step = False |
| 872 | |
| 873 | self._flush_pending_promotions() |
| 874 | self._flush_pending_cascades() |
| 875 | for tier in self.secondary_tiers: |
| 876 | tier.on_schedule_end(context) |
| 877 | |
| 878 | for req_id in context.new_req_ids: |
| 879 | state = self._req_state.get(req_id) |
| 880 | if state is None: |
| 881 | continue |
| 882 | self._metrics.on_request_allocated(state.req_context) |
| 883 | |
| 884 | @override |
| 885 | def has_pending_work(self) -> bool: |
| 886 | # In-flight primary<->secondary transfers (pending promotions are |
| 887 | # translated to transfer jobs in on_schedule_end), plus any work the |
| 888 | # secondary tiers themselves still have outstanding. |
| 889 | return ( |
| 890 | bool(self._jobs) |
| 891 | or any(state.pending_cascade_keys for state in self._req_state.values()) |
| 892 | or any(tier.has_pending_work() for tier in self.secondary_tiers) |
| 893 | ) |
| 894 | |
| 895 | @override |
| 896 | def take_events(self) -> Iterable[OffloadingEvent]: |
| 897 | """Yield events owned by the primary and secondary tiers. |
| 898 | |
| 899 | Yields: |
| 900 | New OffloadingEvents collected by each tier since the last call. |
| 901 | |
| 902 | """ |
| 903 | yield from self.primary_tier.take_events() |
| 904 | for tier in self.secondary_tiers: |
| 905 | yield from tier.take_events() |
| 906 | |
| 907 | @override |
| 908 | def reset_cache(self) -> None: |
| 909 | """Reset transfer bookkeeping and primary-tier cache. |
| 910 | |
| 911 | Called during sleep, weight update, or resume. Each secondary tier |
| 912 | drains its in-flight transfers via drain_jobs() so no tier I/O is |
| 913 | touching primary memory before the primary tier is reset. A stuck |
| 914 | tier will block here visibly — preferable to silent corruption |
| 915 | from reusing primary slots while a transfer is mid-copy. |
| 916 | |
| 917 | Secondary tiers are intentionally not reset: persistent stores |
| 918 | (FS, network) keep their data across resets. Active request state is |
| 919 | retained so those requests can continue after the reset; finished |
| 920 | requests are finalized and removed. |
| 921 | """ |
| 922 | for tier in self.secondary_tiers: |
| 923 | tier.drain_jobs() |
| 924 | # All tier I/O has stopped; consume their completion notifications |
| 925 | # so manager bookkeeping is consistent before the primary reset. |
| 926 | self._process_finished_jobs() |
| 927 | assert not self._jobs |
| 928 | |
| 929 | # Deferred promotion submissions reserve primary slots that the |
| 930 | # reset below invalidates; their submit_load() has not yet been |
| 931 | # called so no tier I/O is touching that memory. |
| 932 | self._pending_load_submissions.clear() |
| 933 | self._metrics.assert_idle() |
| 934 | |
| 935 | finished_req_ids = [] |
| 936 | for req_id, state in self._req_state.items(): |
| 937 | state.pending_primary_stores = 0 |
| 938 | state.pending_cascade_keys.clear() |
| 939 | if not state.is_finished: |
| 940 | continue |
| 941 | for tier in self.secondary_tiers: |
| 942 | tier.on_request_finished(state.req_context) |
| 943 | self._metrics.on_request_finished(state.req_context) |
| 944 | finished_req_ids.append(req_id) |
| 945 | |
| 946 | self.primary_tier.reset_cache() |
| 947 | |
| 948 | for req_id in finished_req_ids: |
| 949 | del self._req_state[req_id] |
| 950 | self._processed_jobs_this_step = False |
| 951 | |
| 952 | for tier in self.secondary_tiers: |
| 953 | if tier.bp_detector is not None: |
| 954 | tier.bp_detector.reset() |
| 955 | |
| 956 | @override |
| 957 | def get_stats(self) -> OffloadingConnectorStats | None: |
| 958 | stats = self.primary_tier.get_stats() |
| 959 | |
| 960 | if stats is not None and stats.is_empty(): |
| 961 | stats = None |
| 962 | |
| 963 | self._metrics.record_backpressure(self.secondary_tiers) |
| 964 | metrics_stats = self._metrics.take_stats() |
| 965 | if metrics_stats is not None: |
| 966 | if stats is None: |
| 967 | stats = metrics_stats |
| 968 | else: |
| 969 | stats.aggregate(metrics_stats) |
| 970 | |
| 971 | for tier in self.secondary_tiers: |
| 972 | tier_stats = tier.get_stats() |
| 973 | if tier_stats is None or tier_stats.is_empty(): |
| 974 | continue |
| 975 | if stats is None: |
| 976 | stats = tier_stats |
| 977 | else: |
| 978 | stats.aggregate(tier_stats) |
| 979 | |
| 980 | return stats |
| 981 | |
| 982 | @override |
| 983 | def shutdown(self) -> None: |
| 984 | """Shut down secondary tiers before releasing primary resources. |
| 985 | |
| 986 | Every secondary tier is given a shutdown attempt. If any shutdown |
| 987 | fails, preserve the primary mmap because a failed tier may still use it. |
| 988 | """ |
| 989 | shutdown_error: Exception | None = None |
| 990 | for tier_idx, tier in enumerate(self.secondary_tiers): |
| 991 | try: |
| 992 | tier.shutdown() |
| 993 | except Exception as exc: |
| 994 | shutdown_error = exc |
| 995 | logger.exception( |
| 996 | "Failed to shut down secondary tier #%d " |
| 997 | "(tier_type=%s, impl_class=%s)", |
| 998 | tier_idx, |
| 999 | tier.tier_type, |
| 1000 | type(tier).__name__, |
| 1001 | ) |
| 1002 | |
| 1003 | if shutdown_error is not None: |
| 1004 | raise shutdown_error |
| 1005 | |
| 1006 | self.primary_tier.shutdown() |