| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """vLLM secondary-tier adapter for KVCR.""" |
| 4 | |
| 5 | import ctypes |
| 6 | import mmap |
| 7 | import socket |
| 8 | import time |
| 9 | import uuid |
| 10 | from collections.abc import Collection, Iterable |
| 11 | from pathlib import Path |
| 12 | from typing import TYPE_CHECKING, Any, cast |
| 13 | |
| 14 | import msgspec |
| 15 | from kvcr import ( |
| 16 | DURATION_METRIC, |
| 17 | KVCR, |
| 18 | ROUTER_HINT_CAPABILITIES, |
| 19 | STATE_METRIC, |
| 20 | TRANSFER_BLOCKS_METRIC, |
| 21 | TRANSFER_BYTES_METRIC, |
| 22 | KVCRBindings, |
| 23 | ) |
| 24 | from kvcr.config import ( |
| 25 | FrameworkDramInput, |
| 26 | G3Options, |
| 27 | KeyAdapter, |
| 28 | KVCRBackendConfigs, |
| 29 | KVCRConfig, |
| 30 | KVCRGuardConfig, |
| 31 | LocalDramOptions, |
| 32 | RemoteFWDramOptions, |
| 33 | ) |
| 34 | from kvcr.control_channels import ZmqPeerControlChannel |
| 35 | from kvcr.policy import ( |
| 36 | FIFOPolicy, |
| 37 | G3FIFOPolicy, |
| 38 | G3LRUPolicy, |
| 39 | KVCachePolicy, |
| 40 | LRUPolicy, |
| 41 | ) |
| 42 | from kvcr.types import ( |
| 43 | BlockKey, |
| 44 | CacheTier, |
| 45 | InventoryEvent, |
| 46 | MemDescriptor, |
| 47 | OpHandle, |
| 48 | PinRequestId, |
| 49 | PinResult, |
| 50 | QueryStatus, |
| 51 | ) |
| 52 | from typing_extensions import override |
| 53 | |
| 54 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( |
| 55 | OffloadingConnectorStats, |
| 56 | ) |
| 57 | from vllm.logger import init_logger |
| 58 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 59 | from vllm.v1.core.kv_cache_utils import ( |
| 60 | BlockHash, |
| 61 | maybe_convert_block_hash, |
| 62 | ) |
| 63 | from vllm.v1.kv_offload.base import ( |
| 64 | LookupResult, |
| 65 | Medium, |
| 66 | OffloadingCounterMetadata, |
| 67 | OffloadingEvent, |
| 68 | OffloadingGaugeMetadata, |
| 69 | OffloadingHistogramMetadata, |
| 70 | OffloadingMetricMetadata, |
| 71 | OffloadKey, |
| 72 | ReqContext, |
| 73 | RequestOffloadingContext, |
| 74 | get_offload_block_hash, |
| 75 | ) |
| 76 | from vllm.v1.kv_offload.tiering.base import ( |
| 77 | JobResult, |
| 78 | ParentManager, |
| 79 | SecondaryTierManager, |
| 80 | TransferJob, |
| 81 | ) |
| 82 | |
| 83 | if TYPE_CHECKING: |
| 84 | from vllm.v1.kv_offload.base import OffloadingSpec |
| 85 | |
| 86 | |
| 87 | _REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES |
| 88 | |
| 89 | logger = init_logger(__name__) |
| 90 | |
| 91 | _BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = { |
| 92 | "fifo": FIFOPolicy, |
| 93 | "lru": LRUPolicy, |
| 94 | "g3_fifo": G3FIFOPolicy, |
| 95 | "g3_lru": G3LRUPolicy, |
| 96 | } |
| 97 | |
| 98 | |
| 99 | def _resolve_policy(name: str | None) -> KVCachePolicy | None: |
| 100 | if name is None: |
| 101 | return None |
| 102 | policy_type = _BUILTIN_POLICIES.get(name) |
| 103 | if policy_type is None: |
| 104 | if "." not in name: |
| 105 | raise ValueError( |
| 106 | f"Unknown KVCR policy: {name!r}. " |
| 107 | f"Supported: {list(_BUILTIN_POLICIES)}; external policies " |
| 108 | "must use a fully qualified module.Class path." |
| 109 | ) |
| 110 | policy_type = resolve_obj_by_qualname(name) |
| 111 | if not isinstance(policy_type, type) or not issubclass( |
| 112 | policy_type, KVCachePolicy |
| 113 | ): |
| 114 | raise TypeError(f"{name} is not a KVCachePolicy") |
| 115 | return policy_type() |
| 116 | |
| 117 | |
| 118 | def _vllm_metric_name(name: str) -> str: |
| 119 | return f"vllm:{name}" |
| 120 | |
| 121 | |
| 122 | def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]: |
| 123 | return { |
| 124 | _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata( |
| 125 | documentation="KVCR operation and stage duration in seconds.", |
| 126 | labelnames=("scope", "result"), |
| 127 | ), |
| 128 | _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata( |
| 129 | documentation="Bytes transferred by KVCR operations.", |
| 130 | labelnames=("operation",), |
| 131 | ), |
| 132 | _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata( |
| 133 | documentation="KV blocks transferred by successful KVCR operations.", |
| 134 | labelnames=("operation",), |
| 135 | ), |
| 136 | _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata( |
| 137 | documentation="Current KVCR metadata and operation counts.", |
| 138 | labelnames=("resource",), |
| 139 | ), |
| 140 | } |
| 141 | |
| 142 | |
| 143 | class _FrameworkPinAdapter: |
| 144 | """Adapt KVCR framework pin callbacks to primary-tier store jobs.""" |
| 145 | |
| 146 | def __init__(self, tier: "KVCRSecondaryTierManager") -> None: |
| 147 | self._tier = tier |
| 148 | self._next_request_id = 0 |
| 149 | self._next_pin_request_id = 0 |
| 150 | self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {} |
| 151 | self._completed_pins: list[tuple[PinRequestId, PinResult]] = [] |
| 152 | self._pin_jobs: dict[str, int] = {} |
| 153 | self._pin_job_results: list[JobResult] = [] |
| 154 | |
| 155 | def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId: |
| 156 | # The parent callback is valid only from serve_external_requests(). |
| 157 | request = PinRequestId(self._next_pin_request_id) |
| 158 | self._next_pin_request_id += 1 |
| 159 | self._pending_pins[request] = tuple(keys) |
| 160 | return request |
| 161 | |
| 162 | def process_pending(self, parent: ParentManager) -> None: |
| 163 | pending = self._pending_pins |
| 164 | self._pending_pins = {} |
| 165 | self._completed_pins.extend( |
| 166 | (request, self._resolve_pin(parent, keys)) |
| 167 | for request, keys in pending.items() |
| 168 | ) |
| 169 | |
| 170 | def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]: |
| 171 | completed = self._completed_pins |
| 172 | self._completed_pins = [] |
| 173 | return completed |
| 174 | |
| 175 | def cancel_pin_request(self, request: PinRequestId) -> None: |
| 176 | self._pending_pins.pop(request, None) |
| 177 | |
| 178 | def _resolve_pin( |
| 179 | self, parent: ParentManager, keys: Collection[BlockKey] |
| 180 | ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: |
| 181 | offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) |
| 182 | request_id = f"kvcr-source:{self._next_request_id}" |
| 183 | self._next_request_id += 1 |
| 184 | req_context = ReqContext(req_id=request_id) |
| 185 | started = False |
| 186 | job: TransferJob | None = None |
| 187 | pin_handle: str | None = None |
| 188 | try: |
| 189 | parent.on_new_request(req_context) |
| 190 | started = True |
| 191 | hit_keys = tuple( |
| 192 | key |
| 193 | for key in offload_keys |
| 194 | if parent.lookup(key, req_context) is LookupResult.HIT |
| 195 | ) |
| 196 | if not hit_keys: |
| 197 | return None |
| 198 | |
| 199 | job = parent.create_store_job(hit_keys, req_context) |
| 200 | job_keys = tuple(job.keys) |
| 201 | chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) |
| 202 | if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): |
| 203 | return None |
| 204 | |
| 205 | descriptors: dict[BlockKey, list[MemDescriptor] | None] = { |
| 206 | BlockKey(bytes(key)): None for key in offload_keys |
| 207 | } |
| 208 | descriptors.update( |
| 209 | ( |
| 210 | BlockKey(bytes(key)), |
| 211 | [self._tier._make_descriptor(chunk_id)], |
| 212 | ) |
| 213 | for key, chunk_id in zip(job_keys, chunk_ids) |
| 214 | ) |
| 215 | pin_handle = f"kvcr-job-{job.job_id}" |
| 216 | self._pin_jobs[pin_handle] = job.job_id |
| 217 | return pin_handle, descriptors |
| 218 | except Exception: |
| 219 | logger.warning("KVCR source acquisition failed", exc_info=True) |
| 220 | return None |
| 221 | finally: |
| 222 | if started: |
| 223 | try: |
| 224 | parent.on_request_finished(req_context) |
| 225 | except Exception: |
| 226 | logger.warning("KVCR source request cleanup failed", exc_info=True) |
| 227 | if job is not None and pin_handle is None: |
| 228 | self._pin_job_results.append( |
| 229 | JobResult(job_id=job.job_id, success=False) |
| 230 | ) |
| 231 | |
| 232 | def release_pin(self, pin_handle: str) -> bool: |
| 233 | job_id = self._pin_jobs.pop(pin_handle, None) |
| 234 | if job_id is not None: |
| 235 | self._pin_job_results.append(JobResult(job_id=job_id, success=True)) |
| 236 | return True |
| 237 | |
| 238 | def has_active_pins(self) -> bool: |
| 239 | return bool(self._pin_jobs) |
| 240 | |
| 241 | def take_pin_job_results(self) -> list[JobResult]: |
| 242 | results = self._pin_job_results |
| 243 | self._pin_job_results = [] |
| 244 | return results |
| 245 | |
| 246 | |
| 247 | class _VllmKeyAdapter: |
| 248 | """Adapt vLLM's local key format to KVCR.""" |
| 249 | |
| 250 | def encode(self, framework_key: object) -> BlockKey: |
| 251 | if not isinstance(framework_key, bytes): |
| 252 | raise TypeError("vLLM offload keys must be bytes") |
| 253 | return BlockKey(framework_key) |
| 254 | |
| 255 | def decode(self, key: BlockKey) -> int | bytes: |
| 256 | block_hash = BlockHash(get_offload_block_hash(OffloadKey(key))) |
| 257 | return maybe_convert_block_hash(block_hash) |
| 258 | |
| 259 | |
| 260 | # Job ID, remaining blocks, aggregate success, and successful load keys. |
| 261 | _JobState = tuple[int, int, bool, set[OffloadKey] | None] |
| 262 | |
| 263 | |
| 264 | class KVCRSecondaryTierManager(SecondaryTierManager): |
| 265 | """Secondary tier wrapper around the KVCR KV P2P API.""" |
| 266 | |
| 267 | @classmethod |
| 268 | @override |
| 269 | def build_metric_definitions( |
| 270 | cls, extra_config: dict[str, Any] |
| 271 | ) -> dict[str, OffloadingMetricMetadata]: |
| 272 | if not extra_config.get("enable_telemetry", False): |
| 273 | return {} |
| 274 | return _kvcr_metric_definitions() |
| 275 | |
| 276 | def __init__( |
| 277 | self, |
| 278 | offloading_spec: "OffloadingSpec", |
| 279 | primary_kv_view: memoryview, |
| 280 | tier_type: str, |
| 281 | router_capabilities: Iterable[str] | None = None, |
| 282 | control_host: str = "0.0.0.0", |
| 283 | control_ports: list[int] | None = None, |
| 284 | control_advertise_host: str | None = None, |
| 285 | eager_ctrl_connect: bool = True, |
| 286 | opportunistic_query: bool = True, |
| 287 | enable_telemetry: bool = False, |
| 288 | operation_timeout_ms: int = 1000, |
| 289 | abandon_timeout_ms: int = 5000, |
| 290 | metadata_retry_interval_ms: int = 100, |
| 291 | secondary_g2_slots: int = 0, |
| 292 | kvcr_service_socket_path: str | None = None, |
| 293 | compatibility_digest: str | None = None, |
| 294 | policy: str | None = None, |
| 295 | g3: dict[str, Any] | None = None, |
| 296 | local_dram_backend: str = "UCX", |
| 297 | remote_fw_dram_backend: str = "UCX", |
| 298 | ) -> None: |
| 299 | super().__init__(offloading_spec, primary_kv_view, tier_type) |
| 300 | selected_policy = _resolve_policy(policy) |
| 301 | if (kvcr_service_socket_path is None) != (compatibility_digest is None): |
| 302 | raise ValueError( |
| 303 | "kvcr_service_socket_path and compatibility_digest must be " |
| 304 | "configured together" |
| 305 | ) |
| 306 | missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ()) |
| 307 | if missing: |
| 308 | raise ValueError( |
| 309 | f"KVCR secondary tier requires router capabilities: {sorted(missing)}" |
| 310 | ) |
| 311 | events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events |
| 312 | self_describing_events = ( |
| 313 | offloading_spec.kv_events_config.self_describing_kv_events |
| 314 | ) |
| 315 | if ( |
| 316 | events_enabled |
| 317 | and ( |
| 318 | secondary_g2_slots > 0 |
| 319 | or kvcr_service_socket_path is not None |
| 320 | or g3 is not None |
| 321 | ) |
| 322 | and not self_describing_events |
| 323 | ): |
| 324 | raise ValueError( |
| 325 | "KVCR local G2/G3 inventory requires self_describing_kv_events" |
| 326 | ) |
| 327 | self._key_adapter: KeyAdapter = _VllmKeyAdapter() |
| 328 | advertise_host = control_advertise_host or socket.gethostname() |
| 329 | dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local |
| 330 | if dp_local_rank is None: |
| 331 | if control_ports is None or len(control_ports) != 1: |
| 332 | raise ValueError( |
| 333 | "control_ports must contain exactly one port when " |
| 334 | "data_parallel_rank_local is unset" |
| 335 | ) |
| 336 | dp_local_rank = 0 |
| 337 | if control_ports is None or not 0 <= dp_local_rank < len(control_ports): |
| 338 | raise ValueError( |
| 339 | "control_ports must contain a port for " |
| 340 | f"local data-parallel rank {dp_local_rank}" |
| 341 | ) |
| 342 | if any(not 1 <= port <= 65535 for port in control_ports): |
| 343 | raise ValueError("control_ports must be between 1 and 65535") |
| 344 | control = ZmqPeerControlChannel( |
| 345 | control_host, |
| 346 | control_ports[dp_local_rank], |
| 347 | advertise_host, |
| 348 | ) |
| 349 | # Give colocated workers distinct transport listen ports. |
| 350 | with socket.socket() as _s: |
| 351 | _s.bind(("", 0)) |
| 352 | _nixl_listen_port = _s.getsockname()[1] |
| 353 | self._primary_base_addr = ctypes.addressof( |
| 354 | ctypes.c_char.from_buffer(primary_kv_view) |
| 355 | ) |
| 356 | if primary_kv_view.strides is None: |
| 357 | raise ValueError("primary KV memoryview must expose strides") |
| 358 | self._primary_row_stride = int(primary_kv_view.strides[0]) |
| 359 | if secondary_g2_slots < 0: |
| 360 | raise ValueError("secondary_g2_slots must be non-negative") |
| 361 | if ( |
| 362 | secondary_g2_slots or kvcr_service_socket_path is not None |
| 363 | ) and not events_enabled: |
| 364 | logger.warning( |
| 365 | "KVCR local DRAM is enabled, but KV cache events are disabled; " |
| 366 | "local DRAM inventory will not be published" |
| 367 | ) |
| 368 | local_mapping: mmap.mmap | None = None |
| 369 | local_dram: LocalDramOptions | None = None |
| 370 | if secondary_g2_slots: |
| 371 | if kvcr_service_socket_path is not None: |
| 372 | logger.warning( |
| 373 | "secondary_g2_slots is ignored when " |
| 374 | "kvcr_service_socket_path is configured" |
| 375 | ) |
| 376 | else: |
| 377 | local_mapping = mmap.mmap( |
| 378 | -1, secondary_g2_slots * self._primary_row_stride |
| 379 | ) |
| 380 | local_dram = LocalDramOptions( |
| 381 | pools=[ |
| 382 | ( |
| 383 | "", |
| 384 | ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)), |
| 385 | len(local_mapping), |
| 386 | ) |
| 387 | ], |
| 388 | backend=local_dram_backend, |
| 389 | ) |
| 390 | guard_config = ( |
| 391 | KVCRGuardConfig( |
| 392 | kvcr_service_socket_path=kvcr_service_socket_path, |
| 393 | guard_index=dp_local_rank, |
| 394 | compatibility_digest=compatibility_digest, |
| 395 | ) |
| 396 | if kvcr_service_socket_path is not None and compatibility_digest is not None |
| 397 | else None |
| 398 | ) |
| 399 | g3_config = None |
| 400 | if g3 is not None: |
| 401 | g3_values = dict(g3) |
| 402 | g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"]) |
| 403 | g3_config = G3Options(**g3_values) |
| 404 | nixl_agent_name = f"KVCR-{uuid.uuid4()}" |
| 405 | self._framework_pin_adapter = _FrameworkPinAdapter(self) |
| 406 | self._inventory_events: list[OffloadingEvent] = [] |
| 407 | |
| 408 | try: |
| 409 | self._kvcr = KVCR( |
| 410 | KVCRConfig( |
| 411 | pool_layouts=[("", self._primary_row_stride)], |
| 412 | enable_telemetry=enable_telemetry, |
| 413 | operation_timeout_ms=operation_timeout_ms, |
| 414 | abandon_timeout_ms=abandon_timeout_ms, |
| 415 | nixl_agent_name=nixl_agent_name, |
| 416 | nixl_listen_port=_nixl_listen_port, |
| 417 | ), |
| 418 | KVCRBindings( |
| 419 | request_pin=self._framework_pin_adapter.request_pin, |
| 420 | poll_pin_results=self._framework_pin_adapter.poll_pin_results, |
| 421 | release_pin=self._framework_pin_adapter.release_pin, |
| 422 | cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request), |
| 423 | framework_control=control, |
| 424 | key_adapter=self._key_adapter, |
| 425 | inventory_sink=(self._record_inventory if events_enabled else None), |
| 426 | stats_factory=( |
| 427 | OffloadingConnectorStats if enable_telemetry else None |
| 428 | ), |
| 429 | policy=selected_policy, |
| 430 | ), |
| 431 | KVCRBackendConfigs( |
| 432 | framework_dram=FrameworkDramInput( |
| 433 | self._primary_base_addr, primary_kv_view.nbytes |
| 434 | ), |
| 435 | local_dram=local_dram, |
| 436 | g3=g3_config, |
| 437 | remote_fw_dram=RemoteFWDramOptions( |
| 438 | eager_ctrl_connect=eager_ctrl_connect, |
| 439 | opportunistic_query=opportunistic_query, |
| 440 | metadata_retry_interval_ms=metadata_retry_interval_ms, |
| 441 | backend=remote_fw_dram_backend, |
| 442 | ), |
| 443 | ), |
| 444 | guard_config, |
| 445 | ) |
| 446 | except BaseException: |
| 447 | control.close() |
| 448 | if local_mapping is not None: |
| 449 | local_mapping.close() |
| 450 | raise |
| 451 | self._local_dram_mmap = local_mapping |
| 452 | self._finished_jobs: list[JobResult] = [] |
| 453 | self._jobs_by_op: dict[OpHandle, _JobState] = {} |
| 454 | |
| 455 | @override |
| 456 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: |
| 457 | block_key = self._key_adapter.encode(key) |
| 458 | status, _ = self._kvcr.query((block_key,), req_context.req_id)[0] |
| 459 | if status is QueryStatus.FETCHING: |
| 460 | return LookupResult.RETRY |
| 461 | if status in (QueryStatus.HIT, QueryStatus.FETCHABLE): |
| 462 | return LookupResult.HIT |
| 463 | return LookupResult.MISS |
| 464 | |
| 465 | @override |
| 466 | def submit_load(self, job_metadata: TransferJob) -> None: |
| 467 | blocks = { |
| 468 | self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] |
| 469 | for key, chunk_id in zip( |
| 470 | job_metadata.keys, job_metadata.chunk_ids, strict=True |
| 471 | ) |
| 472 | } |
| 473 | if not blocks: |
| 474 | self._finished_jobs.append( |
| 475 | JobResult(job_id=job_metadata.job_id, success=True) |
| 476 | ) |
| 477 | return |
| 478 | |
| 479 | op_handle = self._kvcr.deliver( |
| 480 | blocks, request_id=job_metadata.req_context.req_id |
| 481 | ) |
| 482 | self._jobs_by_op[op_handle] = ( |
| 483 | job_metadata.job_id, |
| 484 | len(blocks), |
| 485 | True, |
| 486 | set(), |
| 487 | ) |
| 488 | |
| 489 | @override |
| 490 | def submit_store(self, job_metadata: TransferJob) -> None: |
| 491 | blocks = { |
| 492 | self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] |
| 493 | for key, chunk_id in zip( |
| 494 | job_metadata.keys, job_metadata.chunk_ids, strict=True |
| 495 | ) |
| 496 | } |
| 497 | if not blocks: |
| 498 | self._finished_jobs.append( |
| 499 | JobResult(job_id=job_metadata.job_id, success=True) |
| 500 | ) |
| 501 | return |
| 502 | op_handle = self._kvcr.deposit(blocks) |
| 503 | self._jobs_by_op[op_handle] = ( |
| 504 | job_metadata.job_id, |
| 505 | len(blocks), |
| 506 | True, |
| 507 | None, |
| 508 | ) |
| 509 | |
| 510 | @override |
| 511 | def get_finished_jobs(self) -> Iterable[JobResult]: |
| 512 | results = self._finished_jobs |
| 513 | self._finished_jobs = [] |
| 514 | results.extend(self._poll_finished_jobs()) |
| 515 | return results |
| 516 | |
| 517 | @override |
| 518 | def get_stats(self) -> OffloadingConnectorStats | None: |
| 519 | stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats()) |
| 520 | if stats is not None: |
| 521 | # Add the framework namespace only at the vLLM boundary. |
| 522 | for metrics in stats.data.values(): |
| 523 | for name in tuple(metrics): |
| 524 | metrics[_vllm_metric_name(name)] = metrics.pop(name) |
| 525 | return stats |
| 526 | |
| 527 | @override |
| 528 | def serve_external_requests(self, parent: ParentManager) -> None: |
| 529 | self._framework_pin_adapter.process_pending(parent) |
| 530 | |
| 531 | def _poll_finished_jobs(self) -> list[JobResult]: |
| 532 | results: list[JobResult] = [] |
| 533 | for op_handle, entries in self._kvcr.poll_completed(): |
| 534 | job_state = self._jobs_by_op.get(op_handle) |
| 535 | if job_state is None: |
| 536 | continue |
| 537 | job_id, remaining, success, successful_keys = job_state |
| 538 | remaining -= len(entries) |
| 539 | success = success and all(entry.success for entry in entries.values()) |
| 540 | if successful_keys is not None: |
| 541 | successful_keys.update( |
| 542 | OffloadKey(bytes(key)) |
| 543 | for key, entry in entries.items() |
| 544 | if entry.success |
| 545 | ) |
| 546 | if remaining > 0: |
| 547 | self._jobs_by_op[op_handle] = ( |
| 548 | job_id, |
| 549 | remaining, |
| 550 | success, |
| 551 | successful_keys, |
| 552 | ) |
| 553 | continue |
| 554 | self._jobs_by_op.pop(op_handle, None) |
| 555 | results.append( |
| 556 | JobResult( |
| 557 | job_id=job_id, |
| 558 | success=success, |
| 559 | successful_keys=successful_keys if not success else None, |
| 560 | ) |
| 561 | ) |
| 562 | results.extend(self._framework_pin_adapter.take_pin_job_results()) |
| 563 | return results |
| 564 | |
| 565 | @override |
| 566 | def has_pending_work(self) -> bool: |
| 567 | # Keep scheduler polling: control messages may arrive without a load job. |
| 568 | return True |
| 569 | |
| 570 | @override |
| 571 | def take_events(self) -> Iterable[OffloadingEvent]: |
| 572 | events = self._inventory_events |
| 573 | self._inventory_events = [] |
| 574 | return events |
| 575 | |
| 576 | @override |
| 577 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: |
| 578 | if (hint := req_context.kv_hints) is not None: |
| 579 | self._kvcr.submit_hint( |
| 580 | request_id=req_context.req_id, |
| 581 | hints=msgspec.to_builtins(hint), |
| 582 | ) |
| 583 | return RequestOffloadingContext() |
| 584 | |
| 585 | @override |
| 586 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 587 | self._kvcr.discard_hint(req_context.req_id) |
| 588 | |
| 589 | @override |
| 590 | def drain_jobs(self) -> None: |
| 591 | while self._jobs_by_op or self._framework_pin_adapter.has_active_pins(): |
| 592 | drained = self._poll_finished_jobs() |
| 593 | if drained: |
| 594 | self._finished_jobs.extend(drained) |
| 595 | else: |
| 596 | time.sleep(0.001) |
| 597 | |
| 598 | @override |
| 599 | def shutdown(self) -> None: |
| 600 | self.drain_jobs() |
| 601 | # TODO: Keep registered buffers alive if KVCR.close() fails. |
| 602 | self._kvcr.close() |
| 603 | if self._local_dram_mmap is not None: |
| 604 | self._local_dram_mmap.close() |
| 605 | self._local_dram_mmap = None |
| 606 | |
| 607 | def _record_inventory(self, event: InventoryEvent) -> None: |
| 608 | if event.tier is CacheTier.LOCAL_G2: |
| 609 | medium = Medium.CPU |
| 610 | elif event.tier is CacheTier.G3: |
| 611 | medium = Medium.STORAGE |
| 612 | else: |
| 613 | raise ValueError(f"unsupported KVCR inventory tier: {event.tier}") |
| 614 | self._inventory_events.append( |
| 615 | OffloadingEvent( |
| 616 | keys=[OffloadKey(bytes(key)) for key in event.keys], |
| 617 | medium=medium, |
| 618 | removed=event.removed, |
| 619 | ownership="kvcr", |
| 620 | removal_expected=True, |
| 621 | ) |
| 622 | ) |
| 623 | |
| 624 | def _make_descriptor(self, chunk_id: int) -> MemDescriptor: |
| 625 | return MemDescriptor( |
| 626 | end_point_name=self._kvcr.config.nixl_agent_name, |
| 627 | mem_type="DRAM", |
| 628 | addr=self._primary_base_addr + chunk_id * self._primary_row_stride, |
| 629 | size=self._primary_row_stride, |
| 630 | device_Id=0, |
| 631 | info="", |
| 632 | ) |