| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Collection, Iterable, Mapping |
| 5 | from types import SimpleNamespace |
| 6 | from typing import Any |
| 7 | |
| 8 | import msgspec |
| 9 | import numpy as np |
| 10 | import pytest |
| 11 | |
| 12 | pytest.importorskip("kvcr") |
| 13 | |
| 14 | from kvcr import KVCRBindings |
| 15 | from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig |
| 16 | from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy |
| 17 | from kvcr.types import ( |
| 18 | BlockKey, |
| 19 | CacheTier, |
| 20 | InventoryEvent, |
| 21 | MemDescriptor, |
| 22 | OpEntryResult, |
| 23 | OpEntryStatus, |
| 24 | OpHandle, |
| 25 | QueryStatus, |
| 26 | ) |
| 27 | |
| 28 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( |
| 29 | OffloadingConnectorStats, |
| 30 | ) |
| 31 | from vllm.v1.kv_hints import KvHintsEnvelope |
| 32 | from vllm.v1.kv_offload.base import ( |
| 33 | LookupResult, |
| 34 | Medium, |
| 35 | OffloadKey, |
| 36 | ReqContext, |
| 37 | make_offload_key, |
| 38 | ) |
| 39 | from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob |
| 40 | from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager |
| 41 | from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager |
| 42 | |
| 43 | |
| 44 | def _op_entries( |
| 45 | entries: Mapping[BlockKey, bool], |
| 46 | ) -> dict[BlockKey, OpEntryResult]: |
| 47 | return { |
| 48 | key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED) |
| 49 | for key, success in entries.items() |
| 50 | } |
| 51 | |
| 52 | |
| 53 | class _StubControlChannel: |
| 54 | def __init__(self, endpoint: str) -> None: |
| 55 | self.endpoint = endpoint |
| 56 | |
| 57 | def send(self, endpoint: str, message: bytes) -> bool: |
| 58 | return True |
| 59 | |
| 60 | def recv(self) -> list[bytes]: |
| 61 | return [] |
| 62 | |
| 63 | def close(self) -> None: |
| 64 | pass |
| 65 | |
| 66 | |
| 67 | class RecordingKVCR: |
| 68 | def __init__(self) -> None: |
| 69 | self.config: KVCRConfig | None = None |
| 70 | self.guard_config: KVCRGuardConfig | None = None |
| 71 | self.backend_configs = KVCRBackendConfigs() |
| 72 | self.constructor_bindings: KVCRBindings | None = None |
| 73 | self.nixl_agent_name = "recording" |
| 74 | self.framework_control: _StubControlChannel | None = None |
| 75 | self.inventory_sink = None |
| 76 | self.query_status = QueryStatus.MISS |
| 77 | self.stats: OffloadingConnectorStats | None = None |
| 78 | self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] |
| 79 | self.discard_hint_calls: list[str] = [] |
| 80 | self.deliver_calls: list[ |
| 81 | tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] |
| 82 | ] = [] |
| 83 | self.deposit_calls: list[ |
| 84 | tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] |
| 85 | ] = [] |
| 86 | self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] |
| 87 | self._next_op_handle = 1 |
| 88 | |
| 89 | def submit_hint(self, *args: Any, **kwargs: Any) -> None: |
| 90 | self.submit_hint_calls.append((args, kwargs)) |
| 91 | |
| 92 | def discard_hint(self, request_id: str) -> None: |
| 93 | self.discard_hint_calls.append(request_id) |
| 94 | |
| 95 | def query( |
| 96 | self, |
| 97 | keys: Collection[BlockKey], |
| 98 | request_id: str | None = None, |
| 99 | ) -> list[tuple[QueryStatus, CacheTier | None]]: |
| 100 | tier = { |
| 101 | QueryStatus.HIT: CacheTier.LOCAL_G2, |
| 102 | QueryStatus.FETCHING: CacheTier.LOCAL_G2, |
| 103 | QueryStatus.FETCHABLE: CacheTier.REMOTE_G2, |
| 104 | QueryStatus.MISS: None, |
| 105 | }[self.query_status] |
| 106 | return [(self.query_status, tier) for _ in keys] |
| 107 | |
| 108 | def deliver( |
| 109 | self, |
| 110 | blocks: Mapping[BlockKey, list[MemDescriptor]], |
| 111 | request_id: str | None = None, |
| 112 | ) -> OpHandle: |
| 113 | op_handle = self._next_op_handle |
| 114 | self._next_op_handle += 1 |
| 115 | self.deliver_calls.append((op_handle, dict(blocks), request_id)) |
| 116 | self.complete(op_handle, {key: True for key in blocks}) |
| 117 | return op_handle |
| 118 | |
| 119 | def deposit( |
| 120 | self, |
| 121 | blocks: Mapping[BlockKey, list[MemDescriptor]], |
| 122 | ) -> OpHandle: |
| 123 | op_handle = self._next_op_handle |
| 124 | self._next_op_handle += 1 |
| 125 | self.deposit_calls.append((op_handle, dict(blocks))) |
| 126 | self.complete(op_handle, {key: True for key in blocks}) |
| 127 | return op_handle |
| 128 | |
| 129 | def complete( |
| 130 | self, |
| 131 | op_handle: OpHandle, |
| 132 | entries: Mapping[BlockKey, bool], |
| 133 | ) -> None: |
| 134 | self.completed.append((op_handle, _op_entries(entries))) |
| 135 | |
| 136 | def poll_completed( |
| 137 | self, |
| 138 | ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]: |
| 139 | completed = self.completed |
| 140 | self.completed = [] |
| 141 | return completed |
| 142 | |
| 143 | def get_stats(self) -> OffloadingConnectorStats | None: |
| 144 | stats = self.stats |
| 145 | self.stats = None |
| 146 | return stats |
| 147 | |
| 148 | def close(self) -> None: |
| 149 | pass |
| 150 | |
| 151 | |
| 152 | class _ExternalPolicy(FIFOPolicy): |
| 153 | pass |
| 154 | |
| 155 | |
| 156 | def _job( |
| 157 | job_id: int, |
| 158 | req_context: ReqContext, |
| 159 | key: OffloadKey | None = None, |
| 160 | chunk_id: int = 0, |
| 161 | ) -> TransferJob: |
| 162 | return TransferJob( |
| 163 | job_id=job_id, |
| 164 | keys=[key if key is not None else OffloadKey(b"k0")], |
| 165 | chunk_ids=np.array([chunk_id], dtype=np.int64), |
| 166 | is_promotion=True, |
| 167 | req_context=req_context, |
| 168 | ) |
| 169 | |
| 170 | |
| 171 | def _make_tier( |
| 172 | monkeypatch, |
| 173 | kvcr: RecordingKVCR, |
| 174 | *, |
| 175 | enable_telemetry: bool = False, |
| 176 | secondary_g2_slots: int = 0, |
| 177 | kvcr_service_socket_path: str | None = None, |
| 178 | compatibility_digest: str | None = None, |
| 179 | enable_kv_cache_events: bool = False, |
| 180 | self_describing_kv_events: bool = False, |
| 181 | policy: str | None = None, |
| 182 | g3: dict[str, object] | None = None, |
| 183 | control_ports: list[int] | None = None, |
| 184 | data_parallel_rank_local: int | None = None, |
| 185 | ) -> KVCRSecondaryTierManager: |
| 186 | def make_control(_bind_host, bind_port, advertise_host): |
| 187 | return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") |
| 188 | |
| 189 | def make_kvcr(config, bindings, backend_configs, guard_config): |
| 190 | kvcr.config = config |
| 191 | kvcr.guard_config = guard_config |
| 192 | kvcr.nixl_agent_name = config.nixl_agent_name |
| 193 | kvcr.backend_configs = backend_configs |
| 194 | kvcr.constructor_bindings = bindings |
| 195 | kvcr.framework_control = bindings.framework_control |
| 196 | kvcr.inventory_sink = bindings.inventory_sink |
| 197 | return kvcr |
| 198 | |
| 199 | monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) |
| 200 | monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) |
| 201 | return KVCRSecondaryTierManager( |
| 202 | offloading_spec=SimpleNamespace( |
| 203 | config=SimpleNamespace( |
| 204 | parallel=SimpleNamespace( |
| 205 | data_parallel_rank_local=data_parallel_rank_local, |
| 206 | ) |
| 207 | ), |
| 208 | kv_events_config=SimpleNamespace( |
| 209 | enable_kv_cache_events=enable_kv_cache_events, |
| 210 | self_describing_kv_events=self_describing_kv_events, |
| 211 | ), |
| 212 | ), |
| 213 | primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), |
| 214 | tier_type="kvcr", |
| 215 | router_capabilities=["router_hint"], |
| 216 | control_host="127.0.0.1", |
| 217 | control_ports=control_ports if control_ports is not None else [7777], |
| 218 | control_advertise_host="127.0.0.1", |
| 219 | enable_telemetry=enable_telemetry, |
| 220 | secondary_g2_slots=secondary_g2_slots, |
| 221 | kvcr_service_socket_path=kvcr_service_socket_path, |
| 222 | compatibility_digest=compatibility_digest, |
| 223 | policy=policy, |
| 224 | g3=g3, |
| 225 | local_dram_backend="UCX", |
| 226 | remote_fw_dram_backend="UCX", |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): |
| 231 | """Keep the control endpoint and guard pool aligned to the local DP rank.""" |
| 232 | kvcr = RecordingKVCR() |
| 233 | tier = _make_tier( |
| 234 | monkeypatch, |
| 235 | kvcr, |
| 236 | kvcr_service_socket_path="/tmp/kvcr.sock", |
| 237 | compatibility_digest="Opaque-Digest", |
| 238 | secondary_g2_slots=1, |
| 239 | control_ports=[7001, 7002], |
| 240 | data_parallel_rank_local=1, |
| 241 | ) |
| 242 | |
| 243 | assert kvcr.framework_control is not None |
| 244 | assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002" |
| 245 | assert kvcr.guard_config == KVCRGuardConfig( |
| 246 | kvcr_service_socket_path="/tmp/kvcr.sock", |
| 247 | guard_index=1, |
| 248 | compatibility_digest="Opaque-Digest", |
| 249 | ) |
| 250 | assert kvcr.config is not None |
| 251 | assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] |
| 252 | assert kvcr.backend_configs.local_dram is None |
| 253 | |
| 254 | |
| 255 | def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): |
| 256 | """Convert user-provided G3 paths to KVCR's typed configuration.""" |
| 257 | kvcr = RecordingKVCR() |
| 258 | path = tmp_path / "g3.data" |
| 259 | |
| 260 | _make_tier( |
| 261 | monkeypatch, |
| 262 | kvcr, |
| 263 | g3={"paths": [str(path)], "capacity_bytes_per_file": 64}, |
| 264 | ) |
| 265 | |
| 266 | assert kvcr.backend_configs.g3 == G3Options( |
| 267 | paths=(path,), |
| 268 | capacity_bytes_per_file=64, |
| 269 | ) |
| 270 | |
| 271 | |
| 272 | def test_kvcr_tier_adapts_request_and_load(monkeypatch): |
| 273 | """Check hint forwarding, key conversion, load descriptors, and cleanup.""" |
| 274 | kvcr = RecordingKVCR() |
| 275 | tier = _make_tier(monkeypatch, kvcr) |
| 276 | kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[]) |
| 277 | ctx = ReqContext( |
| 278 | req_id="req", |
| 279 | kv_hints=kv_hint, |
| 280 | ) |
| 281 | key = make_offload_key((123).to_bytes(8, "big"), 0) |
| 282 | same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) |
| 283 | other_key = make_offload_key((124).to_bytes(8, "big"), 0) |
| 284 | |
| 285 | tier.on_new_request(ctx) |
| 286 | |
| 287 | assert kvcr.submit_hint_calls == [ |
| 288 | ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)}) |
| 289 | ] |
| 290 | |
| 291 | bindings = kvcr.constructor_bindings |
| 292 | assert bindings is not None |
| 293 | assert bindings.key_adapter is not None |
| 294 | monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0") |
| 295 | decode = bindings.key_adapter.decode |
| 296 | hash_123 = (123).to_bytes(8, "big") |
| 297 | assert decode(BlockKey(bytes(key))) == hash_123 |
| 298 | assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123 |
| 299 | assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big") |
| 300 | |
| 301 | tier.submit_load(_job(7, ctx, key=key, chunk_id=2)) |
| 302 | |
| 303 | # Check load destinations and job completion independently of hint parsing. |
| 304 | assert len(kvcr.submit_hint_calls) == 1 |
| 305 | _, blocks, request_id = kvcr.deliver_calls[0] |
| 306 | assert request_id == "req" |
| 307 | assert list(blocks) == [key] |
| 308 | (descriptor,) = blocks[key] |
| 309 | assert descriptor.end_point_name == kvcr.nixl_agent_name |
| 310 | assert descriptor.addr == tier._primary_base_addr + 2 * 16 |
| 311 | assert descriptor.size == 16 |
| 312 | assert list(tier.get_finished_jobs()) == [JobResult(7, True)] |
| 313 | |
| 314 | # Here we verify that request cleanup discards the request-scoped hint in |
| 315 | # KVCR. |
| 316 | tier.on_request_finished(ctx) |
| 317 | assert kvcr.discard_hint_calls == ["req"] |
| 318 | |
| 319 | |
| 320 | def test_kvcr_tier_allows_request_without_router_hint(monkeypatch): |
| 321 | """Keep router hints optional for requests from non-hint-aware routers.""" |
| 322 | kvcr = RecordingKVCR() |
| 323 | tier = _make_tier(monkeypatch, kvcr) |
| 324 | |
| 325 | tier.on_new_request(ReqContext(req_id="req")) |
| 326 | |
| 327 | assert kvcr.submit_hint_calls == [] |
| 328 | |
| 329 | |
| 330 | @pytest.mark.parametrize( |
| 331 | ("status", "expected"), |
| 332 | [ |
| 333 | (QueryStatus.MISS, LookupResult.MISS), |
| 334 | (QueryStatus.HIT, LookupResult.HIT), |
| 335 | (QueryStatus.FETCHABLE, LookupResult.HIT), |
| 336 | (QueryStatus.FETCHING, LookupResult.RETRY), |
| 337 | ], |
| 338 | ) |
| 339 | def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): |
| 340 | """Map KVCR cache states to the scheduler's lookup contract.""" |
| 341 | kvcr = RecordingKVCR() |
| 342 | kvcr.query_status = status |
| 343 | tier = _make_tier(monkeypatch, kvcr) |
| 344 | |
| 345 | assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected |
| 346 | |
| 347 | |
| 348 | def test_kvcr_tier_serves_primary_pin_request(monkeypatch): |
| 349 | """Hold primary-tier hits until KVCR releases the corresponding pin.""" |
| 350 | kvcr = RecordingKVCR() |
| 351 | tier = _make_tier(monkeypatch, kvcr) |
| 352 | keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) |
| 353 | hit_keys = (keys[0], keys[2]) |
| 354 | chunk_ids = {keys[0]: 1, keys[2]: 5} |
| 355 | lifecycle: list[str] = [] |
| 356 | |
| 357 | class Parent: |
| 358 | def on_new_request(self, req_context): |
| 359 | lifecycle.append("new") |
| 360 | return SimpleNamespace() |
| 361 | |
| 362 | def lookup(self, key, req_context): |
| 363 | return LookupResult.HIT if key in hit_keys else LookupResult.MISS |
| 364 | |
| 365 | def create_store_job(self, requested_keys, req_context): |
| 366 | return TransferJob( |
| 367 | job_id=11, |
| 368 | keys=requested_keys, |
| 369 | chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), |
| 370 | is_promotion=True, |
| 371 | req_context=req_context, |
| 372 | ) |
| 373 | |
| 374 | def on_request_finished(self, req_context): |
| 375 | lifecycle.append("finished") |
| 376 | |
| 377 | bindings = kvcr.constructor_bindings |
| 378 | assert bindings is not None |
| 379 | request = bindings.request_pin(keys) |
| 380 | |
| 381 | tier.serve_external_requests(Parent()) |
| 382 | |
| 383 | [(queued_request, result)] = bindings.poll_pin_results() |
| 384 | assert queued_request == request |
| 385 | assert result is not None |
| 386 | pin_handle, descriptors = result |
| 387 | assert descriptors[keys[1]] is None |
| 388 | block_descriptors = descriptors[keys[2]] |
| 389 | assert block_descriptors is not None |
| 390 | (descriptor,) = block_descriptors |
| 391 | assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) |
| 392 | assert lifecycle == ["new", "finished"] |
| 393 | |
| 394 | polls = 0 |
| 395 | |
| 396 | def release_pin(): |
| 397 | nonlocal polls |
| 398 | polls += 1 |
| 399 | bindings.release_pin(pin_handle) |
| 400 | return [] |
| 401 | |
| 402 | monkeypatch.setattr(kvcr, "poll_completed", release_pin) |
| 403 | tier.drain_jobs() |
| 404 | |
| 405 | assert polls == 1 |
| 406 | assert list(tier.get_finished_jobs()) == [JobResult(11, True)] |
| 407 | |
| 408 | |
| 409 | def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch): |
| 410 | """Keep telemetry opt-in and namespace metrics only at the vLLM boundary.""" |
| 411 | assert KVCRSecondaryTierManager.build_metric_definitions({}) == {} |
| 412 | definitions = KVCRSecondaryTierManager.build_metric_definitions( |
| 413 | {"enable_telemetry": True} |
| 414 | ) |
| 415 | assert "vllm:kvcr_duration_seconds" in definitions |
| 416 | assert "vllm:kvcr_transfer_blocks" in definitions |
| 417 | |
| 418 | kvcr = RecordingKVCR() |
| 419 | tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True) |
| 420 | bindings = kvcr.constructor_bindings |
| 421 | assert bindings is not None |
| 422 | assert bindings.stats_factory is not None |
| 423 | stats = bindings.stats_factory() |
| 424 | stats.increase_counter( |
| 425 | "kvcr_transfer_blocks", |
| 426 | 2, |
| 427 | ("remote_deliver",), |
| 428 | ) |
| 429 | kvcr.stats = stats |
| 430 | |
| 431 | returned = tier.get_stats() |
| 432 | |
| 433 | assert returned is stats |
| 434 | assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2} |
| 435 | |
| 436 | |
| 437 | @pytest.mark.parametrize( |
| 438 | ("policy", "expected_type"), |
| 439 | [ |
| 440 | ("fifo", FIFOPolicy), |
| 441 | ("lru", LRUPolicy), |
| 442 | ("g3_fifo", G3FIFOPolicy), |
| 443 | ("g3_lru", G3LRUPolicy), |
| 444 | (f"{__name__}._ExternalPolicy", _ExternalPolicy), |
| 445 | ], |
| 446 | ) |
| 447 | def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type): |
| 448 | """Resolve every built-in and fully qualified external policy.""" |
| 449 | kvcr = RecordingKVCR() |
| 450 | _make_tier(monkeypatch, kvcr, policy=policy) |
| 451 | |
| 452 | bindings = kvcr.constructor_bindings |
| 453 | assert bindings is not None |
| 454 | assert type(bindings.policy) is expected_type |
| 455 | |
| 456 | |
| 457 | @pytest.mark.parametrize( |
| 458 | ("socket_path", "digest"), |
| 459 | [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")], |
| 460 | ) |
| 461 | def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest): |
| 462 | """Reject partial guard configuration before connecting to the service.""" |
| 463 | with pytest.raises(ValueError, match="configured together"): |
| 464 | _make_tier( |
| 465 | monkeypatch, |
| 466 | RecordingKVCR(), |
| 467 | kvcr_service_socket_path=socket_path, |
| 468 | compatibility_digest=digest, |
| 469 | ) |
| 470 | |
| 471 | |
| 472 | def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): |
| 473 | """Cover store descriptors and inventory translation at the KVCR boundary.""" |
| 474 | kvcr = RecordingKVCR() |
| 475 | tier = _make_tier( |
| 476 | monkeypatch, |
| 477 | kvcr, |
| 478 | secondary_g2_slots=2, |
| 479 | enable_kv_cache_events=True, |
| 480 | self_describing_kv_events=True, |
| 481 | ) |
| 482 | local_dram = kvcr.backend_configs.local_dram |
| 483 | assert local_dram is not None |
| 484 | assert [(name, size) for name, _, size in local_dram.pools] == [ |
| 485 | ("", 2 * tier._primary_row_stride) |
| 486 | ] |
| 487 | |
| 488 | key = OffloadKey(b"k0") |
| 489 | tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) |
| 490 | |
| 491 | _, blocks = kvcr.deposit_calls[0] |
| 492 | assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 |
| 493 | assert list(tier.get_finished_jobs()) == [JobResult(11, True)] |
| 494 | |
| 495 | assert kvcr.inventory_sink is not None |
| 496 | kvcr.inventory_sink( |
| 497 | InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False) |
| 498 | ) |
| 499 | kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False)) |
| 500 | events = list(tier.take_events()) |
| 501 | assert [ |
| 502 | (event.keys, event.medium, event.ownership, event.removed) for event in events |
| 503 | ] == [ |
| 504 | ([key], Medium.CPU, "kvcr", False), |
| 505 | ([key], Medium.STORAGE, "kvcr", False), |
| 506 | ] |
| 507 | assert all(event.removal_expected for event in events) |
| 508 | tier.shutdown() |
| 509 | |
| 510 | |
| 511 | def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): |
| 512 | """Require tier-aware events whenever KVCR owns local cache inventory.""" |
| 513 | with pytest.raises(ValueError, match="self_describing_kv_events"): |
| 514 | _make_tier( |
| 515 | monkeypatch, |
| 516 | RecordingKVCR(), |
| 517 | secondary_g2_slots=1, |
| 518 | enable_kv_cache_events=True, |
| 519 | ) |
| 520 | |
| 521 | |
| 522 | def test_kvcr_tier_accumulates_block_results(monkeypatch): |
| 523 | """Wait for every block result and preserve partial success while draining.""" |
| 524 | kvcr = RecordingKVCR() |
| 525 | tier = _make_tier(monkeypatch, kvcr) |
| 526 | keys = [OffloadKey(b"k0"), OffloadKey(b"k1")] |
| 527 | tier.submit_load( |
| 528 | TransferJob( |
| 529 | job_id=13, |
| 530 | keys=keys, |
| 531 | chunk_ids=np.array([0, 1], dtype=np.int64), |
| 532 | is_promotion=True, |
| 533 | req_context=ReqContext(req_id="req"), |
| 534 | ) |
| 535 | ) |
| 536 | |
| 537 | op_handle, _, _ = kvcr.deliver_calls[0] |
| 538 | kvcr.completed.clear() |
| 539 | kvcr.complete(op_handle, {keys[0]: True}) |
| 540 | assert list(tier.get_finished_jobs()) == [] |
| 541 | |
| 542 | kvcr.complete(op_handle, {keys[1]: False}) |
| 543 | tier.drain_jobs() |
| 544 | |
| 545 | [result] = tier.get_finished_jobs() |
| 546 | assert result.job_id == 13 |
| 547 | assert not result.success |
| 548 | assert result.successful_keys == {keys[0]} |