From 419ac8be3475f1548083735502b998e00ec7863e Mon Sep 17 00:00:00 2001 From: Weiyang Wang Date: Thu, 29 Oct 2026 02:55:21 +0200 Subject: [PATCH] feat: add source-owned G2 removal for move hints Let fetch honor copy/move mode in request-scoped kv.fetch hints. Capture usable source advice and matching full keys at admission; move hints always require hash matching. Hint replacement does not change admitted work, failed advice preserves local hits, and deliver remains a copy. Carry move intent in the initial transfer request. The source admits moves on its owner thread, withdraws eligible G2 keys, and keeps its transfer claims through native completion. Successful writes permit removal after earlier operations drain; failed writes reopen only keys the move owns. Complete destination cache hits send an initial one-way invalidation. Reject source keys with incomplete G2 copies or unfinished deposits before closing admission, allowing those deposits to finish normally. Use the existing abandonment budget for destination move completion while retaining the source's normal transfer timeout. Exceptionally delayed destinations can still fail after source removal. No post-copy confirmation is required, and destination completion does not wait for source cleanup. Preserve no_evict protection and incompatible source layouts. Framework and G3 copies are untouched. Existing inventory hooks report actual source removals and newly complete destination copies. Rejecting a key relinquishes the move's removal ownership before later retries or callbacks. Validation: 920 unit tests passed, 2 environment-dependent skips; lint and formatting passed. Regression coverage includes delayed completion, partial/cached blocks, retention, failures, and overlapping invalidations. Signed-off-by: Weiyang Wang Co-Authored-By: Claude Fable 5.1 --- docs/design_overview.md | 22 +- src/kvcr/api.py | 10 +- src/kvcr/backends/local_dram.py | 39 +- src/kvcr/backends/local_dram_pools.py | 74 +- src/kvcr/backends/moves.py | 86 ++ src/kvcr/backends/remote_fw_dram.py | 88 +- src/kvcr/backends/remote_ops.py | 13 + src/kvcr/backends/remote_sources.py | 52 +- src/kvcr/hints/fetch.py | 12 +- src/kvcr/runtime/core.py | 81 +- src/kvcr/runtime/invalidation.py | 161 ++++ src/kvcr/service/recovery_journal.py | 2 + tests/unit/_kvcr_test_utils.py | 7 + tests/unit/nvtx/test_nvtx_hints.py | 49 + tests/unit/test_g3.py | 41 + tests/unit/test_hints_fetch.py | 37 +- tests/unit/test_hints_parser.py | 29 +- tests/unit/test_kvcr_invalidate.py | 314 +++++++ tests/unit/test_kvcr_move.py | 1205 +++++++++++++++++++++++++ tests/unit/test_kvcr_remote_source.py | 62 ++ 20 files changed, 2321 insertions(+), 63 deletions(-) create mode 100644 src/kvcr/backends/moves.py create mode 100644 src/kvcr/runtime/invalidation.py create mode 100644 tests/unit/test_kvcr_invalidate.py create mode 100644 tests/unit/test_kvcr_move.py diff --git a/docs/design_overview.md b/docs/design_overview.md index f005a2d9..49150f16 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -249,22 +249,36 @@ kvcr.discard_hint(request_id) # request is over or its hints should be discarded ### Inventory and block identity -The router maintains an eventually-consistent global prefix inventory for cache-aware request routing. The router is the routing-key authority and supplies the request's content-derived `BlockKey` values. A `BlockKey` identifies one cacheable KV unit, while a `BlockKeyList` is an ordered exact request path rather than an implicit prefix expansion. The framework normally passes these keys unchanged to the KVCR; an integration with distinct local hashes translates them at its boundary. The key space may be namespaced by model version, quantization, KV layout, tenant, or hash algorithm, and the KVCR treats each full key as opaque. +The router maintains an eventually-consistent global prefix inventory for cache-aware request routing. Its advertised hashes may differ from the full keys used by the engine's cache. The framework supplies exact `BlockKey` values to KVCR, including any representation-specific identity such as a cache-group ID. A `BlockKey` identifies one cacheable KV unit, while a `BlockKeyList` is an ordered exact request path rather than an implicit prefix expansion. The key space may be namespaced by model version, quantization, KV layout, tenant, or hash algorithm, and KVCR treats each full key as opaque. Framework-owned memory is reported through the framework's existing router event path. If needed, `tier_enter`/`tier_exit` reporting or `inventory_snapshot()` can be added to the KVCR API later. -KVCR-owned inventory events use the same `BlockKey` values as fetch and `submit_hint` in the normal case. KVCR deliberately reports only the affected keys, tier, and whether the residency was removed. KVCR delivers each event to `inventory_sink`; the framework or local indexer owns any buffering or batching for publication. Each event may contain multiple keys and is tier-specific because one block may be present in several tiers at once. +KVCR-owned inventory events use the same full `BlockKey` values as fetch. Request-scoped hints instead carry hashes interpreted by the configured key adapter. KVCR deliberately reports only the affected keys, tier, and whether the residency was removed. KVCR delivers each event to `inventory_sink`; the framework or local indexer owns any buffering or batching for publication. Each event may contain multiple keys and is tier-specific because one block may be present in several tiers at once. ### Request-scoped hints For data shared among engines on the same node, the router can use one KVCR as the source in hints sent to the other local KVCRs. A destination may retain a local copy when replication is worthwhile. Otherwise, the router includes `no_retain` in `submit_hint`. The destination still satisfies any committed framework request; `no_retain` only advises it not to keep an additional KVCR-owned copy after the framework no longer needs it. -KVCR-to-KVCR movement is destination-initiated. A router hint gives the destination the source node and KVCR control endpoint. `mode=copy` retains the source copy, while `mode=move` permits source eviction only after successful completion. Timeout or cancellation keeps the source copy. When the hint has no source control endpoint, the destination uses its local storage information or checks object storage when enabled. The router does not need object-store inventory for locality decisions, but may track object presence to issue source hints. +KVCR-to-KVCR movement is destination-initiated. A router hint gives the destination the source node and KVCR control endpoint. `mode=copy` retains the source copy, while `mode=move` lets the source remove its G2 copy after its native transfer succeeds. When the hint has no source control endpoint, the destination uses its local storage information or checks object storage when enabled. The router does not need object-store inventory for locality decisions, but may track object presence to issue source hints. `submit_hint` may establish the peer connection, but does not start data movement or remote pinning yet. Each hint in an envelope is parsed and handled independently; proactive fetching, staging, or pinning may be added later as new hint types. `discard_hint` normally reports that the relevant request is over, and also permits an integration to discard the request-scoped hint early. A route-time `submit_hint` may carry a whole `BlockKeyList` from one source. In the future, if necessary, multi-source assembly can be added by splitting the list into multiple hinted operations. The router may choose not to send a hint based on its internal cost model. Router hints may later be extended to request proactive copy or move for cache rebalancing, which may require periodic utilization reports from KVCR. The router also considers sending hints for matching TP layouts, meaning prefill-to-prefill, decode-to-decode, and aggregated deployments, which cover many use cases. A TP-independent host representation further removes that constraint. +In the current implementation, `kv.fetch` requires a source control endpoint and carries advertised block hashes. The framework supplies complete `BlockKey` values with the same request ID; the key adapter matches those keys against the hint hashes. There is no hash-to-key index for autonomous operations on unrelated hashes, and `no_retain` and source-less hints are not implemented. + +`mode="copy"` is the default. With usable `mode="move"` advice, `fetch` requests source G2 removal for keys whose hashes match the hint, including destination cache hits. Move hints always require hash matching, even when opportunistic querying is enabled. Other keys retain ordinary local fetch behavior, and failed or conflicting advice does not prevent local cache hits. `deliver` always copies; a request using both calls can therefore move fetched keys while copying delivered keys. `discard_hint` clears advice without cancelling already-started operations. + +### Move and source invalidation + +The move handler captures the source advice and matching keys when `fetch` is admitted. The transfer request carries a move flag. Source keys with incomplete G2 copies or unfinished deposits fail individually before withdrawal; their deposits can still finish. The source admits moves on its owner thread, where the affected keys stop accepting new reads and deposits while existing operations retain their claims; the write itself still runs on the progress thread. Successful native transfer leaves removal pending before the source transfer claim is released; failed transfer reopens admission and keeps the copy. There is no destination completion confirmation or wait for source cleanup. + +A destination reuses complete copies it already holds and immediately sends the source one `invalidate` message naming them. For partly cached keys, unfinished destination writes settle before the source transfers the remaining pieces. A move joining an existing copy waits for that copy, then handles the complete key as a cache hit. Results contain ordinary fetch release handles protecting the destination copies. + +The destination uses `abandon_timeout_ms` as its move completion allowance; the source retains its normal transfer timeout. This tolerates delayed completion processing, but a destination that exceeds the longer allowance can still fail after the source has removed its copy. Existing cancellation and buffer-quarantine rules continue to protect memory until native access stops. + +Removal affects only KVCR-owned G2 copies in the requested complete layout, after earlier operations and readers drain. `deposit(no_evict=True)` protection is honored; framework-owned memory and G3 copies are retained. Actual source removal and newly complete destination copies publish the corresponding `inventory_sink` events. Removal remains best effort, with no separate public outcome. Control messages assume trusted peers, and no removal is requested when the hint names this node as its own source. + ### Peer control and liveness The router is never on the data path: KVCR instances execute transfers peer-to-peer via NIXL. Source endpoints remain request-scoped rather than becoming part of global block state, so this coordination adds no KVCR-side peer inventory. @@ -313,7 +327,7 @@ At initialization, an integration selects one built-in policy or supplies an ext Eviction scores are finite floats. LRU favors evicting later sequence keys at comparable recency. Custom policies may use ordering position or ignore it. -Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims are hard mechanism constraints: while they exist, policy cannot evict the entry. `no_retain` is an advisory router retention hint. If it conflicts with a framework `no_evict` claim, `no_evict` wins until the framework calls `release` with the returned handle; policy may then honor or override `no_retain`. +Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims prevent eviction while held. Move uses the source invalidation rules above rather than a policy eviction decision. `no_retain` is not currently supported. The policy methods are invoked as follows: diff --git a/src/kvcr/api.py b/src/kvcr/api.py index 82313cb4..e30cdbbe 100644 --- a/src/kvcr/api.py +++ b/src/kvcr/api.py @@ -187,7 +187,15 @@ def fetch( expected_layout: list[str] | None = None, hints: object | None = None, ) -> OpHandle: - """Asynchronously fetch blocks into KVCR-managed storage.""" + """Asynchronously fetch blocks into KVCR storage and return release claims. + + A usable kv.fetch hint with mode="move" requests best-effort source G2 + removal for matching successful keys, including destination cache hits. + Other keys retain ordinary fetch behavior. Failed advice does not + prevent local cache hits. Completion reports destination readiness + without waiting for source cleanup, which respects no_evict protection + and leaves G3 and framework memory intact. + """ return self._core.fetch(keys, request_id, expected_layout, hints) def release( diff --git a/src/kvcr/backends/local_dram.py b/src/kvcr/backends/local_dram.py index fa9c26e5..fa5dd5b6 100644 --- a/src/kvcr/backends/local_dram.py +++ b/src/kvcr/backends/local_dram.py @@ -38,6 +38,7 @@ if TYPE_CHECKING: from ..runtime.core import _KVCRCore + from .moves import _PendingMove logger = logging.getLogger(__name__) @@ -59,6 +60,8 @@ class _PendingResidencyOp(_Op): key_layout: tuple[str, ...] | None = None fill_sources: Mapping[BlockKey, CacheTier] = field(default_factory=dict) claimed_keys: set[BlockKey] = field(default_factory=set) + # Move advice captured at admission; see _PendingMove. + move: "_PendingMove | None" = None @dataclass @@ -247,9 +250,23 @@ def deposit( key_layout=key_layout, ) self._pending_residency_ops[op.op_id] = op - self._kvcr._add_block_dependencies(op, new_operation=True) + self._admit(op) self._start_deposits(op, blocks) + def _admit(self, op: _PendingResidencyOp) -> None: + """Decide admission once: keys pending removal fail now and never hold a + dependency, so work on other keys cannot delay their removal. Accepted + keys may finish even if removal becomes pending later. + """ + rejected = self._kvcr._invalidation_handler.rejected_entries(op.keys) + if rejected: + op.results.update(rejected) + self._kvcr._add_block_dependencies( + op, + new_operation=True, + keys=op.keys - op.results.keys() if rejected else None, + ) + def _start_deposits( self, op: _PendingResidencyOp, keys: Collection[BlockKey] ) -> None: @@ -394,6 +411,7 @@ def fetch( hints: object | None, layout: list[str], key_layout: tuple[str, ...] | None = None, + move: "_PendingMove | None" = None, ) -> dict[BlockKey, list[_TransferRef]]: ordered_keys = tuple(dict.fromkeys(keys)) if not ordered_keys: @@ -409,9 +427,10 @@ def fetch( key_layout=key_layout, fill_sources=sources, hints=hints, + move=move, ) self._pending_residency_ops[op.op_id] = op - self._kvcr._add_block_dependencies(op, new_operation=True) + self._admit(op) destinations = {} for key in ordered_keys: refs = self._prepare_fetch(op, key) @@ -423,6 +442,8 @@ def fetch( def _prepare_fetch( self, op: _PendingResidencyOp, key: BlockKey ) -> list[_TransferRef] | None: + if key in op.results: + return None layout = op.requested_layout record = self._kvcr._block_record(key) residency = record.local_dram @@ -463,12 +484,22 @@ def _prepare_fetch( ready = [residency.slots[index] for index in indices] self._kvcr._record_access((key,)) op.results[key] = self._pools._new_public_claim(key, residency, slots=ready) + if op.move is not None and key in op.move.keys: + op.move.keys.remove(key) + op.move.cached_keys.append(key) else: if key not in op.claimed_keys: self._pools._acquire_claim(key, residency) op.claimed_keys.add(key) if source is None or residency.remote_indices is not None: return None + if ( + op.move is not None + and key in op.move.keys + and residency.pending_deposits + ): + # Source removal requires every piece outside this write to be ready. + return None # One remote fill owns these destinations until its native terminal. missing = [ index @@ -783,6 +814,7 @@ def _resolve_fetch_waiter( {key: refs}, residency_op.request_id, residency_op.deadline, + move=residency_op.move, ) return key in residency_op.results @@ -898,6 +930,9 @@ def _finish_delivery_copy(self, copy: _LocalCopyOp) -> None: self._finish_deliver_if_ready(op) def _finish_residency_if_ready(self, op: _PendingResidencyOp) -> None: + if op.move is not None: + # Complete local hits ask the source to drop its copies right away. + self._kvcr._move_handler.invalidate_cached(op.move) if ( len(op.results) != len(op.keys) or self._pending_residency_ops.get(op.op_id) is not op diff --git a/src/kvcr/backends/local_dram_pools.py b/src/kvcr/backends/local_dram_pools.py index ed10bc8f..6107b579 100644 --- a/src/kvcr/backends/local_dram_pools.py +++ b/src/kvcr/backends/local_dram_pools.py @@ -15,7 +15,7 @@ from ..config import LocalDramOptions from ..runtime.policy_runtime import _EvictionQueue -from ..runtime.progress import _TransferRef +from ..runtime.progress import _OpId, _TransferRef from ..types import ( BlockKey, CacheTier, @@ -86,6 +86,9 @@ class _LocalDramResidency(msgspec.Struct, gc=False, eq=False): refs: dict[int, _TransferRef] | tuple[_TransferRef, ...] = msgspec.field( default_factory=dict ) + # no_evict deposit claims, included in claim_count. Invalidation honours + # these guarantees but waits for the other, temporary claims to drain. + retention_claims: int = 0 indices_by_label: dict[str, int | None] | None = None _layout: tuple[str, ...] | None = None _pool_counts: tuple[tuple[str, int], ...] | None = None @@ -233,8 +236,9 @@ def __init__( self._evictable_slots: Counter[str] = Counter() self._unscored: set[BlockKey] = set() self._capacity_eviction_key: BlockKey | None = None + # Handle -> (key, residency, retention); retention marks a no_evict claim. self._public_claims: dict[ - ReleaseHandle, tuple[BlockKey, _LocalDramResidency] + ReleaseHandle, tuple[BlockKey, _LocalDramResidency, bool] ] = {} self._next_release_handle = 1 # A no-op until something attaches: the tiers publish residency @@ -269,7 +273,12 @@ def _publish_source( residency: _LocalDramResidency, refs: tuple[_TransferRef, ...] | None = None, ) -> None: - if self._kvcr.framework_control is None or residency.pending_labels: + # Finishing an existing fill must not reopen a source being drained. + if ( + self._kvcr.framework_control is None + or residency.pending_labels + or self._kvcr._block_record_map[key].pending_invalidation + ): return if refs is None: refs = ( @@ -381,6 +390,8 @@ def _drain_source_completions(self) -> None: state.active.pop(key) for key, residency in zip(keys, sources): record = self._kvcr._block_record_map.get(key) + if record is not None and record.pending_invalidation: + self._kvcr._invalidation_handler.notify(key) if ( record is not None and record.local_dram is residency @@ -493,7 +504,9 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: if claim is None: results.append((handle, False)) continue - key, residency = claim + key, residency, retention = claim + if retention: + residency.retention_claims -= 1 self._release_claim(key, residency) results.append((handle, True)) self._update_capacity_pressure() @@ -504,6 +517,7 @@ def acquire_sources( keys: Collection[BlockKey], *, allow_filling: bool = False, + op_id: _OpId | None = None, ) -> dict[BlockKey, _LocalDramResidency]: sources: dict[BlockKey, _LocalDramResidency] = {} for key in keys: @@ -511,8 +525,11 @@ def acquire_sources( continue record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None - if residency is None or ( - residency.state is not _LocalDramState.READY + if ( + residency is None + or record.pending_invalidation + and op_id not in (record.in_flight_ops or ()) + or residency.state is not _LocalDramState.READY and not (allow_filling and residency.state is _LocalDramState.FILLING) ): continue @@ -586,7 +603,11 @@ def _new_public_claim( self._acquire_claim(key, residency) handle = ReleaseHandle(self._next_release_handle) self._next_release_handle += 1 - self._public_claims[handle] = (key, residency) + # Fetch passes readable slots; a protected deposit claims only retention. + retention = slots is None + self._public_claims[handle] = (key, residency, retention) + if retention: + residency.retention_claims += 1 return OpEntryResult( OpEntryStatus.SUCCESS, [ @@ -617,6 +638,8 @@ def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: ): raise RuntimeError(f"invalid local DRAM claim for {key!r}") residency.claim_count -= 1 + if record.pending_invalidation: + self._kvcr._invalidation_handler.notify(key) if residency.claim_count == 0: self._after_last_claim(key, residency) @@ -634,7 +657,9 @@ def _after_last_claim(self, key: BlockKey, residency: _LocalDramResidency) -> bo state.index.pop(key, None) record.local_dram = None if not retiring: - self._make_evictable(key) + if not record.pending_invalidation: + self._make_evictable(key) + # Pending removal owns cleanup; avoid scoring a copy about to drop. return False if residency.state is _LocalDramState.READY: if self._kvcr._key_history_enabled: @@ -656,6 +681,39 @@ def _after_last_claim(self, key: BlockKey, residency: _LocalDramResidency) -> bo self._resume_capacity_waiters() return True + def _drop(self, key: BlockKey) -> None: + """Reclaim G2 slots after source withdrawal and all readers/operations drain. + Unlike reader release, this removes the copy rather than retaining it. + """ + core = self._kvcr + record = core._block_record_map[key] + residency = record.local_dram + if residency is None or residency.claim_count or record.in_flight_ops: + raise RuntimeError(f"local DRAM entry is not ready to drop {key!r}") + self._remove_evictable(key, residency) + if core._key_history_enabled: + core._log_key_history("remove", (key,), CacheTier.LOCAL_G2, "invalidate") + state = self._source_state + # The owner lock excludes index writers; an admitted fill may reinsert None. + if key in state.index: + with state.lock: + state.index.pop(key, None) + record.local_dram = None + self._free(residency.slots) + self.abandon_capacity_eviction(key) + self._residency_observer(key, record) + core._prune_block_record(key) + core._on_remove(core._block_meta(key, record, self._size_bytes(residency))) + # A callback's redeposit must not be followed by a stale removal event. + current = core._block_record_map.get(key) + replacement = current.local_dram if current is not None else None + if not residency.pending_labels and not ( + replacement is not None + and replacement.state is _LocalDramState.READY + and not replacement.pending_labels + ): + core._publish_inventory((key,), CacheTier.LOCAL_G2, removed=True) + def _allocate_slots( self, layout: list[str], protected: frozenset[BlockKey], deadline: float ) -> tuple[list[tuple[str, int]] | None, list[BlockKey], bool]: diff --git a/src/kvcr/backends/moves.py b/src/kvcr/backends/moves.py new file mode 100644 index 00000000..cf9e90be --- /dev/null +++ b/src/kvcr/backends/moves.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Capture move advice and invalidate copies already held at the destination.""" + +import logging +from collections.abc import Collection +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from ..types import BlockKey + +if TYPE_CHECKING: + from ..runtime.core import _KVCRCore + from .remote_fw_dram import _RequestHint + +logger = logging.getLogger(__name__) + + +@dataclass +class _PendingMove: + """Admission snapshot: keys still to transfer, and complete local hits + awaiting one invalidation message to the source. + """ + + hint: "_RequestHint" + layout: tuple[str, ...] + keys: set[BlockKey] + cached_keys: list[BlockKey] = field(default_factory=list) + + +class _MoveHandler: + """Keep move advice stable and send invalidation for complete local hits.""" + + def __init__(self, core: "_KVCRCore") -> None: + self._core = core + + def prepare( + self, + keys: Collection[BlockKey], + request_id: str | None, + layout: tuple[str, ...], + ) -> _PendingMove | None: + """Capture usable move advice; other keys keep ordinary fetch behavior.""" + remote = self._core._remote_fw_dram + hint = remote._request_hints.get(request_id) if request_id is not None else None + adapter = remote._key_adapter + if ( + hint is None + or not hint.move + or hint.source is None + or hint.failed + or adapter is None + or hint.source == getattr(self._core.framework_control, "endpoint", None) + ): + return None + matched = { + key + for key in keys + if key not in hint.missing_keys + and adapter.encode_to_router(key) in hint.block_hashes + } + return _PendingMove(hint, layout, matched) if matched else None + + def invalidate_cached(self, move: _PendingMove) -> None: + """Send one invalidation for the complete local hits found so far.""" + from .remote_ops import _Invalidate + + if not move.cached_keys: + return + keys = tuple(move.cached_keys) + move.cached_keys.clear() + try: + self._core._progress.submit( + _Invalidate( + keys=keys, + key_layout=move.layout, + source_control_endpoint=move.hint.source, + ) + ) + except Exception: + # Removal is best effort; the destination result stands. + logger.warning( + "KVCR cached-key invalidation could not reach %s", + move.hint.source, + exc_info=True, + ) diff --git a/src/kvcr/backends/remote_fw_dram.py b/src/kvcr/backends/remote_fw_dram.py index 935ee0d8..dc9fa519 100644 --- a/src/kvcr/backends/remote_fw_dram.py +++ b/src/kvcr/backends/remote_fw_dram.py @@ -47,6 +47,7 @@ _CONTROL_DECODER, _ControlMessage, _decode_notif, + _Invalidate, _NvtxState, _PeerMessage, _SourceWriteOp, @@ -65,6 +66,7 @@ if TYPE_CHECKING: from ..runtime.core import _KVCRCore + from .moves import _PendingMove @dataclass(frozen=True) @@ -75,6 +77,7 @@ class _RequestHint: failed: bool = False missing_keys: frozenset[BlockKey] = frozenset() trace: _nvtx._LifecycleTrace | None = field(default=None, repr=False, compare=False) + move: bool = False @dataclass(frozen=True) @@ -162,6 +165,8 @@ def submit_hint( src: str | None, block_hashes: frozenset[int], request_id: str | None, + *, + move: bool = False, ) -> None: kvcr = self._kvcr if request_id is not None: @@ -193,6 +198,7 @@ def submit_hint( else kvcr._timer() ), trace=trace, + move=move, ) if src is not None and self._options.eager_ctrl_connect: kvcr._progress.submit(_TargetMetadataRequest(src)) @@ -210,9 +216,8 @@ def query(self, key: BlockKey, request_id: str) -> bool: ): return False if ( - not self._options.opportunistic_query - and adapter.encode_to_router(key) not in request_hint.block_hashes - ): + request_hint.move or not self._options.opportunistic_query + ) and adapter.encode_to_router(key) not in request_hint.block_hashes: return False return True @@ -224,26 +229,32 @@ def _start_target_pull( op_handle: OpHandle, *, local_fill: bool, + move: "_PendingMove | None" = None, ) -> bool: kvcr = self._kvcr started_at = kvcr._timer() scope = "remote_fetch" if local_fill else "remote_deliver" - current_hint = ( - self._request_hints.get(request_id) if request_id is not None else None - ) - if current_hint is not None and ( - current_hint.failed or current_hint.missing_keys.issuperset(blocks) + # Moves use their admission snapshot, independent of later advice. + hint = move.hint if move is not None else self._request_hints.get(request_id) + if ( + hint is None + or hint.source is None + or hint.failed + or (move is None and hint.missing_keys.issuperset(blocks)) ): kvcr._record_duration(scope, started_at, "failed") return False - - if current_hint is None or current_hint.source is None: - kvcr._record_duration(scope, started_at, "failed") - return False - keys = tuple(key for key in blocks if key not in current_hint.missing_keys) - kvcr._record_duration("hint_wait", current_hint.submitted_at, "complete") - if request_id is not None: - self._request_hints[request_id] = replace(current_hint, submitted_at=None) + keys = ( + tuple(key for key in blocks if key not in hint.missing_keys) + if move is None and hint.missing_keys + else tuple(blocks) + ) + # Consume timing only while this exact hint is still current. + if hint.submitted_at is not None and ( + move is None or self._request_hints.get(request_id) is hint + ): + self._request_hints[request_id] = replace(hint, submitted_at=None) + kvcr._record_duration("hint_wait", hint.submitted_at, "complete") op = _TargetPullOp( state=_TargetPullState.START_WRITE, @@ -252,17 +263,22 @@ def _start_target_pull( started_at=started_at, deadline=deadline, op_id=("target", op_handle), - remote_ctrl_ep=current_hint.source, + remote_ctrl_ep=hint.source, _backend=self, ordered_keys=keys, dst_descriptors=tuple(tuple(blocks[key]) for key in keys), - request_id=request_id, + # Explicit-source fills must not update discarded or replacement advice. + request_id=request_id if move is None else None, + move=move is not None, key_layout=( kvcr._block_record_map[keys[0]].local_dram.key_layout if local_fill else () ), ) + if move is not None: + # Keys sent natively need no later invalidation message. + move.keys.difference_update(keys) if self._nvtx is not None: op.trace = self._nvtx.operation( kvcr, @@ -270,13 +286,11 @@ def _start_target_pull( blocks, request_id, local_fill=local_fill, - hint_trace_id=current_hint.trace.trace_id - if current_hint.trace is not None - else 0, + hint_trace_id=hint.trace.trace_id if hint.trace is not None else 0, range_name="target.remote", ) - if current_hint.trace is not None: - current_hint.trace.mark("hint.used", op_handle=op_handle) + if hint.trace is not None: + hint.trace.mark("hint.used", op_handle=op_handle) if op.trace is not None: op.trace.mark("target.queued") if logger.isEnabledFor(logging.DEBUG): @@ -286,7 +300,7 @@ def _start_target_pull( scope, request_id, op_handle, - current_hint.source, + hint.source, len(keys), kvcr._descriptor_bytes(chain.from_iterable(op.dst_descriptors)), ) @@ -337,6 +351,7 @@ def fetch( deadline: float, *, op_handle: OpHandle, + move: "_PendingMove | None" = None, ) -> bool: return self._start_target_pull( blocks, @@ -344,6 +359,7 @@ def fetch( deadline, op_handle, local_fill=True, + move=move, ) def poll_main(self, items: Collection[object]) -> None: @@ -375,6 +391,11 @@ def poll_main(self, items: Collection[object]) -> None: ) elif isinstance(item, _ProgressUpdate): self._apply_progress_update(item) + elif isinstance(item, _Invalidate): + # The destination holds these keys; drop eligible local copies. + self._kvcr._invalidation_handler.submit( + item.keys, key_layout=item.key_layout + ) else: raise TypeError(f"unsupported KVCR main item: {type(item)!r}") if not self._closed: @@ -490,6 +511,13 @@ def poll_progress( ): continue self._send_control(progress, item.endpoint, _TargetMetadata()) + elif isinstance(item, _Invalidate): + # One-way: no reply is awaited and no public result depends on it. + if not self._send_control(progress, item.source_control_endpoint, item): + logger.warning( + "KVCR invalidation could not reach %s", + item.source_control_endpoint, + ) else: raise TypeError(f"unsupported KVCR progress item: {type(item)!r}") @@ -557,6 +585,8 @@ def _process_control_messages(self, progress: _KVCRProgress) -> bool: handled_at = self._kvcr._timer() if isinstance(payload, _StartWrite): self._handle_start_write(progress, payload) + elif isinstance(payload, _Invalidate): + self._progress_outbound.append(payload) elif isinstance(payload, _TargetMetadata): self._handle_target_metadata(progress, payload) elif isinstance(payload, _TargetMetadataAck): @@ -593,7 +623,7 @@ def _send_control( if not isinstance(payload, _TargetMetadataAck): payload.source_control_endpoint = endpoint includes_metadata = ( - not isinstance(payload, (_TargetMetadataAck, _WriteRefused)) + not isinstance(payload, (_TargetMetadataAck, _WriteRefused, _Invalidate)) and endpoint not in self._metadata_acked_sources and ( isinstance(payload, (_TargetMetadata, _StartWrite)) @@ -832,6 +862,11 @@ def _handle_start_write( if self._nvtx is not None else None, ) + if payload.move: + # A move closes owner-side admission for its keys, so the owner + # thread admits it; the write itself still runs on progress. + self._queue_source_pin(source_write, payload.key_layout, move=True) + return prepared = self._try_local_source_write( progress, source_write, payload.key_layout ) @@ -841,7 +876,7 @@ def _handle_start_write( self._queue_source_pin(source_write, payload.key_layout) def _queue_source_pin( - self, op: _SourceWriteOp, key_layout: tuple[str, ...] + self, op: _SourceWriteOp, key_layout: tuple[str, ...], *, move: bool = False ) -> None: self._progress_outbound.append( _SourcePinOp( @@ -856,6 +891,7 @@ def _queue_source_pin( key_layout=key_layout, route=op.route, nvtx_state=op.nvtx_state, + move=move, ) ) diff --git a/src/kvcr/backends/remote_ops.py b/src/kvcr/backends/remote_ops.py index d31ab01e..84e76c72 100644 --- a/src/kvcr/backends/remote_ops.py +++ b/src/kvcr/backends/remote_ops.py @@ -37,6 +37,14 @@ class _StartWrite(_PeerMessage, tag="start_write"): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] key_layout: tuple[str, ...] = () source_incarnation: str | None = None + move: bool = False + + +class _Invalidate(_PeerMessage, tag="invalidate"): + """Remove eligible source G2 copies already held completely at the destination.""" + + keys: tuple[BlockKey, ...] + key_layout: tuple[str, ...] class _TargetMetadata(_PeerMessage, tag="target_metadata"): @@ -65,6 +73,7 @@ class _WriteProbeAck(_PeerMessage, tag="write_probe_ack"): _ControlMessage = ( _StartWrite + | _Invalidate | _TargetMetadata | _TargetMetadataAck | _WriteRefused @@ -111,6 +120,7 @@ class _TargetPullOp(_RemoteOp): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] = () request_id: str | None = None key_layout: tuple[str, ...] = () + move: bool = False success: bool = False completed_keys: set[BlockKey] = field(default_factory=set) probe_sent: bool = False @@ -180,6 +190,7 @@ def _progress( keys=self.ordered_keys, dst_descriptors=self.dst_descriptors, key_layout=self.key_layout if self.local_fill else (), + move=self.move, ), ) if not sent: @@ -373,6 +384,8 @@ class _SourceWriteOp(_RemoteOp): completed_indices: tuple[int, ...] = () route: tuple[str, int] = ("", 0) nvtx_state: _NvtxState | None = None + # Keys this move marked pending at the source; reopened if the write fails. + moving: tuple[BlockKey, ...] = () # None uses caller cleanup; a tuple (even empty) is owned by progress. local_sources: tuple["_LocalDramResidency", ...] | None = None diff --git a/src/kvcr/backends/remote_sources.py b/src/kvcr/backends/remote_sources.py index 4d604cc0..1c511cce 100644 --- a/src/kvcr/backends/remote_sources.py +++ b/src/kvcr/backends/remote_sources.py @@ -38,6 +38,9 @@ class _SourcePinOp(_Op): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] route: tuple[str, int] = ("", 0) key_layout: tuple[str, ...] = () + move: bool = False + # Keys this move marked pending; reopened if they are not written. + moving: tuple[BlockKey, ...] = () framework_pins: set[PinHandle] = field(default_factory=set) pending_pin_ids: set[PinRequestId] = field(default_factory=set) framework_acquire_attempted: bool = False @@ -93,6 +96,8 @@ def finish_source_write(self, item: _SourceWriteOp) -> None: # KVCR-owned contents stay claimed until native quiescence. pins = self._fw_pins_by_op.pop(item.op_id, None) if item.state is _SourceWriteState.FINISHED: + if item.moving and not item.success: + self._kvcr._invalidation_handler.reopen(item.moving) self._kvcr._remove_block_dependencies(item) if item.success: self._kvcr._record_access( @@ -226,6 +231,14 @@ def _submit_prepared_source_write( ) retained_keys = source_pin.keys if waiting else set(completed_keys) + if source_pin.moving: + omitted = [key for key in source_pin.moving if key not in retained_keys] + # Forget rejected keys before reopening them; later retries must not + # cancel a newer invalidation for the same key. + source_pin.moving = tuple( + key for key in source_pin.moving if key in retained_keys + ) + kvcr._invalidation_handler.reopen(omitted) kvcr._release_local_dram_sources(op_id, local_sources.keys() - retained_keys) relevant_pins = { record.fw_mem.pin_handle @@ -270,6 +283,7 @@ def _submit_prepared_source_write( source_keys=completed_keys, src_descriptors=tuple(selected_sources), completed_indices=tuple(completed_indices), + moving=source_pin.moving, ) kvcr._add_block_dependencies(source_write, new_operation=True) self._fw_pins_by_op[source_write.op_id] = framework_pins @@ -383,9 +397,44 @@ def _expire_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: def start_source_pin(self, op: _SourcePinOp) -> None: kvcr = self._kvcr + # Decide admission once; retries retain the original request positions. + pending = kvcr._invalidation_handler._pending + if pending: + op.keys = op.keys.difference(pending) + if op.move and kvcr._local_dram is not None: + # Closing admission must not strand a producer of missing pieces. + # Reject busy keys before claiming them or marking removal pending. + busy = set() + producers = kvcr._local_dram._pending_residency_ops + for key in op.keys: + record = kvcr._block_record_map.get(key) + if record is None: + continue + residency = record.local_dram + if residency is not None: + if residency.pending_labels or residency.pending_deposits: + busy.add(key) + elif any( + producer_id[0] == "deposit" + and (producer := producers.get(producer_id)) is not None + and key not in producer.results + for producer_id in record.active_op_ids + ): + busy.add(key) + if busy: + op.keys = op.keys.difference(busy) kvcr._add_block_dependencies(op, new_operation=True) + if op.move: + # Reject new reads and deposits of these keys while they move. + op.moving = kvcr._invalidation_handler.submit( + op.keys, key_layout=op.key_layout + ) self._source_pin_ops[op.op_id] = op - self._resume_source_pin(op.op_id, op) + if op.ordered_keys and not op.keys: + # Every requested key was rejected; send failure without a transfer. + self._submit_prepared_source_write(op.op_id, op, force_failure=True) + else: + self._resume_source_pin(op.op_id, op) def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: if self._source_pin_ops.get(op_id) is not op: @@ -427,6 +476,7 @@ def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: def _remove_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: self._cancel_pending_pin_for_op(op_id, op) self._source_pin_ops.pop(op_id, None) + self._kvcr._invalidation_handler.reopen(op.moving) self._kvcr._remove_block_dependencies(op) self._kvcr._release_local_dram_sources(op.op_id) framework_pins = set(op.framework_pins) diff --git a/src/kvcr/hints/fetch.py b/src/kvcr/hints/fetch.py index 18330ea1..a8961a1a 100644 --- a/src/kvcr/hints/fetch.py +++ b/src/kvcr/hints/fetch.py @@ -33,18 +33,19 @@ class _KVFetchHint(msgspec.Struct, frozen=True): action_id: str | None source_endpoint: str block_hashes: frozenset[int] + move: bool = False def _parse_kv_fetch( action_id: str | None, payload: Mapping[str, object] ) -> _KVFetchHint: - # TODO: Add mode and no_retain when KVCR starts consuming them. if "no_retain" in payload: raise ValueError("no_retain is not currently supported") + # Fetch uses the mode for matching keys; deliver always copies. mode = payload.get("mode", "copy") - if mode != "copy": - raise ValueError("only copy mode is currently supported") + if mode not in ("copy", "move"): + raise ValueError("kv.fetch mode must be copy or move") source = payload.get("source_control_endpoint") if source is None: @@ -66,8 +67,7 @@ def _parse_kv_fetch( if not hashes: raise ValueError("invalid router hint") - # Add other fields to the parsed result when KVCR starts consuming them. - return _KVFetchHint(action_id, source, frozenset(hashes)) + return _KVFetchHint(action_id, source, frozenset(hashes), move=mode == "move") class _FetchHintHandler: @@ -89,7 +89,7 @@ def submit( # One store call per hint; the store's own rules apply (a later hint # replaces the request's hashes, a different source fails the request). self._core._remote_fw_dram.submit_hint( - hint.source_endpoint, hint.block_hashes, request_id + hint.source_endpoint, hint.block_hashes, request_id, move=hint.move ) self._dispatcher.report( HintResult( diff --git a/src/kvcr/runtime/core.py b/src/kvcr/runtime/core.py index 715c1b4f..2d8b38b0 100644 --- a/src/kvcr/runtime/core.py +++ b/src/kvcr/runtime/core.py @@ -24,6 +24,7 @@ _LocalDramResidency, _LocalDramState, ) +from ..backends.moves import _MoveHandler, _PendingMove from ..config import ( KVCRBackendConfigs, KVCRConfig, @@ -49,6 +50,7 @@ ReleaseHandle, ReleaseResult, ) +from .invalidation import _InvalidationHandler from .policy_runtime import _PolicyInvoker from .progress import ( _FAILED_RESULT, @@ -139,6 +141,8 @@ class _BlockRecord(msgspec.Struct, gc=False): access_count: int = 0 last_access: float | None = None position: int = -1 + # Close G2 to new reads; retain the record until existing users drain. + pending_invalidation: bool = False def add_in_flight_op(self, op_id: _OpId) -> None: if self.in_flight_ops is None: @@ -289,6 +293,7 @@ def __init__( from ..backends.dangling_ops import _log_resilience_event from ..backends.remote_fw_dram import _RemoteFWDram + self._invalidation_handler = _InvalidationHandler(self) self._on_resilience_event_callback = ( _log_resilience_event if bindings.on_resilience_event is None @@ -318,6 +323,7 @@ def __init__( ) self._nvtx = self._remote_fw_dram._nvtx self._nvtx_operations = {} + self._move_handler = _MoveHandler(self) self._g3 = ( _G3( self, @@ -435,6 +441,7 @@ def _has_complete_local(self, key: BlockKey) -> bool: residency = record.local_dram if record is not None else None return ( residency is not None + and not record.pending_invalidation and residency.state is _LocalDramState.READY and not residency.pending_labels and len(residency.slots) == len(residency.key_layout) @@ -456,6 +463,9 @@ def query( remote_keys = self._remote_delivery_keys.get(request_id, ()) for key in keys: record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + statuses.append((QueryStatus.MISS, None)) + continue residency = record.local_dram if record is not None else None if key in remote_keys: fetching_locally = ( @@ -582,9 +592,14 @@ def deliver( local_blocks: dict[BlockKey, list[_TransferRef]] = {} g3_blocks: dict[BlockKey, list[_TransferRef]] = {} remote_blocks: dict[BlockKey, list[_TransferRef]] = {} + # Record pending-removal failures while other keys proceed. + rejected: dict[BlockKey, OpEntryResult] = {} remote_keys = self._remote_delivery_keys.get(request_id, ()) for key, destination in normalized.items(): - if key in remote_keys and ( + record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + rejected[key] = _FAILED_RESULT + elif key in remote_keys and ( not self._has_complete_local(key) or self._remote_fw_dram.query(key, request_id) ): @@ -603,8 +618,13 @@ def deliver( remote_blocks=len(remote_blocks), g3_blocks=len(g3_blocks), ) - if sum(map(bool, (local_blocks, g3_blocks, remote_blocks))) > 1: + if ( + sum(map(bool, (local_blocks, g3_blocks, remote_blocks, rejected))) + > 1 + ): self._joined_completions[op_handle] = (set(blocks), {}) + if rejected: + self._complete(op_handle, rejected) if local_dram is not None and local_blocks: local_dram.deliver(op_handle, local_blocks, deadline=deadline) if g3_blocks and ( @@ -671,6 +691,7 @@ def deposit( ) else: self._local_dram._pools._drain_source_completions() + self._drain_invalidation() if self._g3 is not None: for key, refs in blocks.items(): if self._g3.is_ready(key): @@ -718,6 +739,12 @@ def fetch( return op_handle local_dram._pools._drain_source_completions() + self._drain_invalidation() + intent = ( + self._move_handler.prepare(ordered_keys, request_id, declaration) + if request_id is not None + else None + ) sources = {} for key in ordered_keys: record = self._block_record_map.get(key) @@ -727,7 +754,9 @@ def fetch( and residency.ready_indices(requested) is not None ): continue - if self._g3 is not None and self._g3.is_ready(key): + if intent is not None and key in intent.keys: + sources[key] = CacheTier.REMOTE_G2 + elif self._g3 is not None and self._g3.is_ready(key): if not self._g3.is_ready(key, declaration): raise ValueError( "fetch layout must match existing G3 key layout" @@ -738,6 +767,12 @@ def fetch( ): sources[key] = CacheTier.REMOTE_G2 deadline = self._operation_deadline() + if intent is not None: + # Accept delayed completion using the existing abandonment budget. + # The source independently caps its native transfer timeout. + deadline += ( + self.config.abandon_timeout_ms - self.config.operation_timeout_ms + ) / 1000 destinations = local_dram.fetch( op_handle, ordered_keys, @@ -747,6 +782,7 @@ def fetch( hints=hints, layout=requested, key_layout=declaration, + move=intent, ) for source in (CacheTier.G3, CacheTier.REMOTE_G2): self._start_local_fill( @@ -758,6 +794,7 @@ def fetch( }, request_id, deadline, + move=intent, ) return op_handle @@ -765,7 +802,9 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: if self._local_dram is None: return [(handle, False) for handle in handles] with self._state_lock: - return self._local_dram._pools.release(handles) + results = self._local_dram._pools.release(handles) + self._drain_invalidation() + return results # TODO: Expose individual entry completions as they become available. def poll_completed(self) -> Iterable[OpResult]: @@ -780,6 +819,7 @@ def poll_completed(self) -> Iterable[OpResult]: if self._local_dram is not None: progress_items = self._local_dram.poll_main(progress_items) self._remote_fw_dram.poll_main(progress_items) + self._drain_invalidation() completed = self._completion_queue self._completion_queue = [] if self._nvtx_operations: @@ -789,6 +829,12 @@ def poll_completed(self) -> Iterable[OpResult]: trace.returned() return completed + def _drain_invalidation(self) -> None: + # Cheap when nothing is pending; otherwise a released or finished user + # lets its source copy drop (and free capacity) without waiting for a poll. + if self._invalidation_handler._dirty: + self._invalidation_handler.drain() + def _notify_transfer_errors(self, progress_items: list[object]) -> None: # Notify before cleanup can return a buffer to its allocator. Retain the # batch if user code raises, consuming each notification exactly once. @@ -898,6 +944,8 @@ def run(stage: Callable[[], None]) -> None: ) run(self._close_hint_handlers) + # Native work has stopped; abandon pending removals without waiting. + run(self._invalidation_handler.close) if self._g3 is not None: run(self._g3.close_main) if self._local_dram is not None: @@ -1028,10 +1076,17 @@ def _update_capacity_pressure(self, reclaimable_slots: Mapping[str, int]) -> Non # Generic operation/block dependency bookkeeping. - def _add_block_dependencies(self, op: _Op, *, new_operation: bool) -> None: + def _add_block_dependencies( + self, + op: _Op, + *, + new_operation: bool, + keys: Collection[BlockKey] | None = None, + ) -> None: if new_operation: self._outstanding_operations += 1 - for key in op.keys: + # Exclude rejected keys so work on other keys cannot delay their removal. + for key in op.keys if keys is None else keys: self._block_record(key).add_in_flight_op(op.op_id) def _remove_block_dependencies(self, op: _Op) -> None: @@ -1050,6 +1105,8 @@ def _start_local_fill( blocks: Mapping[BlockKey, list[_TransferRef]], request_id: str | None, deadline: float, + *, + move: _PendingMove | None = None, ) -> None: if not blocks: return @@ -1062,8 +1119,13 @@ def _start_local_fill( deadline, ) elif source is CacheTier.REMOTE_G2: + # A move keeps the advice it was admitted with (see _PendingMove). started = self._remote_fw_dram.fetch( - blocks, request_id, deadline, op_handle=fill_handle + blocks, + request_id, + deadline, + op_handle=fill_handle, + move=move, ) else: raise ValueError(f"unsupported local fill source {source.value}") @@ -1082,6 +1144,7 @@ def _claim_local_dram_sources( claimed = self._local_dram._pools.acquire_sources( tuple(key for key in keys if key not in sources), allow_filling=allow_filling, + op_id=op_id, ) if claimed: sources.update(claimed) @@ -1217,6 +1280,7 @@ def _is_local_resident(self, key: BlockKey) -> bool: record = self._block_record_map.get(key) return ( record is not None + and not record.pending_invalidation and record.local_dram is not None and record.local_dram.state is not _LocalDramState.DISCARDING ) @@ -1270,12 +1334,15 @@ def _record_access(self, keys: Collection[BlockKey]) -> None: def _prune_block_record(self, key: BlockKey) -> None: record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + self._invalidation_handler.notify(key) if ( record is not None and record.fw_mem is None and record.local_dram is None and record.g3 is None and not record.in_flight_ops + and not record.pending_invalidation ): self._block_record_map.pop(key) diff --git a/src/kvcr/runtime/invalidation.py b/src/kvcr/runtime/invalidation.py new file mode 100644 index 00000000..bc901318 --- /dev/null +++ b/src/kvcr/runtime/invalidation.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Drop moved source G2 copies after readers and transfers drain.""" + +from collections.abc import Collection +from typing import TYPE_CHECKING + +from ..types import BlockKey, OpEntryResult +from .progress import _FAILED_RESULT + +if TYPE_CHECKING: + from .core import _BlockRecord, _KVCRCore + + +class _InvalidationHandler: + """Remove sources under the owner lock after existing readers drain. + Submission rejects new reads and deposits; a failed move reopens its keys. + """ + + def __init__(self, core: "_KVCRCore") -> None: + self._core = core + # The layout the destination holds; None means unrestricted removal. + self._pending: dict[ + BlockKey, tuple["_BlockRecord", tuple[str, ...] | None] + ] = {} + # Recheck only on submission or release, not every idle poll. + self._dirty: dict[BlockKey, None] = {} + + def has_pending(self) -> bool: + return bool(self._pending) + + def rejected_entries( + self, keys: Collection[BlockKey] + ) -> dict[BlockKey, OpEntryResult]: + """Check new requests once; admitted operations may resume during removal.""" + if not self._pending: + return {} + return {key: _FAILED_RESULT for key in keys if key in self._pending} + + def close(self) -> None: + """Abandon pending removals after progress stops.""" + for record, _ in self._pending.values(): + record.pending_invalidation = False + self._core._outstanding_operations -= len(self._pending) + self._pending.clear() + self._dirty.clear() + + def submit( + self, + keys: Collection[BlockKey], + *, + key_layout: tuple[str, ...] | None = None, + ) -> tuple[BlockKey, ...]: + """Queue removal of G2 copies held in key_layout; return the keys marked. + + Keys protected by no_evict or held in another complete layout are left + alone. Keys without a G2 copy are skipped too unless an admitted deposit + or fetch may still create one: framework and G3 copies are never touched. + """ + core = self._core + local_dram = core._local_dram + if local_dram is None: + return () + state = local_dram._pools._source_state + marked = [] + for key in keys: + record = core._block_record_map.get(key) + if record is None or key in self._pending: + continue + local = record.local_dram + if local is None: + # A batch dependency can outlive this key's terminal result. + if not any( + (op := local_dram._pending_residency_ops.get(op_id)) is not None + and key not in op.results + for op_id in record.active_op_ids + ): + continue + elif local.retention_claims or ( + key_layout is not None and local.key_layout != key_layout + ): + continue + self._pending[key] = record, key_layout + # Close admission atomically with progress acquiring readers. + # No slot reclamation or callbacks run under this lock. + with state.lock: + record.pending_invalidation = True + state.index.pop(key, None) + core._outstanding_operations += 1 + self._dirty[key] = None + marked.append(key) + return tuple(marked) + + def reopen(self, keys: Collection[BlockKey]) -> None: + """Cancel a failed move's own pending keys; their copies stay readable.""" + for key in keys: + pending = self._pending.get(key) + if pending is not None: + self._settle(key, pending[0], drop=False) + + def _settle(self, key: BlockKey, record: "_BlockRecord", *, drop: bool) -> bool: + """Clear the pending mark, then drop or republish the copy; True if + capacity changed.""" + core = self._core + pools = core._local_dram._pools + # Clear intent before callbacks can redeposit this key or poll again. + del self._pending[key] + self._dirty.pop(key, None) + record.pending_invalidation = False + core._outstanding_operations -= 1 + local = record.local_dram + if local is None: + core._prune_block_record(key) + return False + if drop: + pools._drop(key) + return True + pools._publish_source(key, local) + if local.retention_claims or local.source_readers: + return False + pools._make_evictable(key) + return True + + def notify(self, key: BlockKey) -> None: + if key in self._pending: + self._dirty[key] = None + + @staticmethod + def _is_ready(record: "_BlockRecord") -> bool: + # Submission withdrew the source under the reader lock; its reader count + # can now only decrease during owner-side completion processing. + local = record.local_dram + return not ( + record.in_flight_ops + or local is not None + and (local.claim_count > local.retention_claims or local.source_readers) + ) + + def drain(self) -> None: + pools = self._core._local_dram._pools + dirty, self._dirty = self._dirty, {} + capacity_changed = False + for key in dirty: + # A preceding callback may have removed this key. + pending = self._pending.get(key) + if pending is None: + continue + record, layout = pending + if not self._is_ready(record): + continue + local = record.local_dram + # An earlier producer can grant retention or create a layout the + # destination never held; keep that copy and reopen admission. + keep = local is not None and ( + local.retention_claims + or (layout is not None and local.key_layout != layout) + ) + capacity_changed |= self._settle(key, record, drop=not keep) + if capacity_changed: + pools._resume_capacity_waiters() + pools._update_capacity_pressure() diff --git a/src/kvcr/service/recovery_journal.py b/src/kvcr/service/recovery_journal.py index 0274061a..d7452393 100644 --- a/src/kvcr/service/recovery_journal.py +++ b/src/kvcr/service/recovery_journal.py @@ -394,10 +394,12 @@ def adopt(self, records: dict[BlockKey, _BlockRecord]) -> None: local_dram.claim_count = 0 local_dram.source_readers = 0 local_dram.retire_on_release = False + local_dram.retention_claims = 0 # References name the previous agent, unlike the slot layout. local_dram.refs = {} if record.g3 is not None: record.g3.claim_count = 0 + record.pending_invalidation = False record.fw_mem = None record.in_flight_ops = None record.access_count = 0 diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 04fef6e2..a602a261 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -155,6 +155,13 @@ def create_agent(name, *_): yield +def _invalidate(kvcr: KVCR, keys: Collection[BlockKey]) -> None: + """Exercise the internal invalidation mechanism under its owner lock.""" + core = kvcr._core + with core._state_lock: + core._invalidation_handler.submit(keys) + + def _poll_until( kvcr: KVCR, predicate, diff --git a/tests/unit/nvtx/test_nvtx_hints.py b/tests/unit/nvtx/test_nvtx_hints.py index a5bd28ed..2e627666 100644 --- a/tests/unit/nvtx/test_nvtx_hints.py +++ b/tests/unit/nvtx/test_nvtx_hints.py @@ -5,6 +5,7 @@ import msgspec import pytest from _kvcr_test_utils import ( + _has_outstanding_operations, _mem_descriptor, _poll_until, _record_hint_results, @@ -13,6 +14,7 @@ _write_done_notification, ) from _nvtx_test_utils import make_target, payloads +from test_kvcr_move import _Pair from kvcr import _nvtx from kvcr.hints.results import HintStatus @@ -85,3 +87,50 @@ def test_batch_hint_traces(recording, monkeypatch, level, scenario): == int(submitted[-1]["trace_id"]) == stored.trace.trace_id ) + + +@pytest.mark.parametrize("advice", ["unchanged", "discarded", "replaced"]) +def test_move_preserves_captured_hint_trace(recording, monkeypatch, advice): + pair = _Pair() + key = BlockKey(b"move-trace") + pair.deposit(key, no_evict=False) + core = pair.target._core + remote = core._remote_fw_dram + durations = [] + monkeypatch.setattr(core, "_timer", lambda: 1.0) + monkeypatch.setattr( + core, "_record_duration", lambda *event: durations.append(event) + ) + pair.advise("move") + original_hint = remote._request_hints["req"] + start_pull = remote._start_target_pull + + def start_after_advice_change(*args, **kwargs): + # Admission already captured the source; the first fill can start later. + if advice != "unchanged": + pair.target.discard_hint("req") + if advice == "replaced": + pair.advise(hashes=(999,), source="tcp://replacement:1") + return start_pull(*args, **kwargs) + + monkeypatch.setattr(remote, "_start_target_pull", start_after_advice_change) + moved = pair.target.fetch((key,), "req") + future_hint = remote._request_hints.get("req") + [used] = payloads(recording, "hint.used") + [queued] = payloads(recording, "target.queued") + assert int(used["trace_id"]) == original_hint.trace.trace_id + assert int(queued["hint_trace_id"]) == original_hint.trace.trace_id + assert int(used["op_handle"]) == int(queued["op_handle"]) + assert (("hint_wait", 1.0, "complete") in durations) is (advice == "unchanged") + if advice == "unchanged": + assert future_hint.submitted_at is None + elif advice == "replaced": + assert future_hint.submitted_at == 1.0 + + pair.until(lambda: len(pair.source_agent.xfers) == 2) + pair.source_agent.state = "ERR" + assert not pair.result(moved)[key].success + pair.source_agent.state = "DONE" + pair.until(lambda: not _has_outstanding_operations(pair.source)) + assert remote._request_hints.get("req") is future_hint + assert future_hint is None or not future_hint.failed diff --git a/tests/unit/test_g3.py b/tests/unit/test_g3.py index c3f78b29..e53ed3c7 100644 --- a/tests/unit/test_g3.py +++ b/tests/unit/test_g3.py @@ -17,6 +17,7 @@ _ConstantHashAdapter, _decode_control_message, _has_outstanding_operations, + _invalidate, _mem_descriptor, _new_kvcr, _poll_until, @@ -360,6 +361,46 @@ def test_g3_multi_pool_partial_read_and_layout_mismatch(tmp_path): kvcr.release((result.release_handle,)) +@pytest.mark.parametrize("operation", ["fetch", "deposit"]) +def test_invalidation_skips_g3_key_finished_in_pending_batch(tmp_path, operation): + disk, slow = BlockKey(b"disk"), BlockKey(b"slow") + page = os.sysconf("SC_PAGE_SIZE") + local, framework = (ctypes.create_string_buffer(page) for _ in range(2)) + agent = _FakeG3Agent() + kvcr = _new_g3_kvcr( + tmp_path, + local, + agent=agent, + framework_regions=[_buffer_region(framework, page)], + ) + install_recovery_records(kvcr._core, {disk: _recovered_record(g3=0)}) + agent.state = "PROC" + pending = kvcr.deposit({slow: [_mem_descriptor(0)]}) + mixed = ( + kvcr.fetch((disk, slow)) + if operation == "fetch" + else kvcr.deposit({key: [_mem_descriptor(0)] for key in (disk, slow)}) + ) + op = kvcr._core._local_dram._pending_residency_ops[(operation, mixed)] + try: + # Only the sibling is unfinished; this key cannot create a G2 copy. + assert not op.results[disk].success and slow not in op.results + assert kvcr._core._block_record_map[disk].local_dram is None + _invalidate(kvcr, (disk,)) + assert not kvcr._core._invalidation_handler.has_pending() + assert kvcr.query((disk,)) == [(QueryStatus.FETCHABLE, CacheTier.G3)] + assert not kvcr.poll_completed() + assert kvcr.query((disk,)) == [(QueryStatus.FETCHABLE, CacheTier.G3)] + finally: + agent.state = "DONE" + results = dict(_poll_until(kvcr, lambda items: len(items) == 2)) + assert results[pending][slow].success and results[mixed][slow].success + assert not results[mixed][disk].success + handle = results[mixed][slow].release_handle + if handle is not None: + assert kvcr.release((handle,)) == [(handle, True)] + + def test_g3_multi_pool_impossible_spill_preserves_disk_entries(tmp_path): kvcr, primary, local, keys = _multi_pool_g3(tmp_path, disk_units=6, deposits=0) small, pinned = BlockKey(b"small"), BlockKey(b"pinned") diff --git a/tests/unit/test_hints_fetch.py b/tests/unit/test_hints_fetch.py index bfedbbf4..0f8d5385 100644 --- a/tests/unit/test_hints_fetch.py +++ b/tests/unit/test_hints_fetch.py @@ -94,27 +94,38 @@ def release(self): return self._inner.release() -def test_gpu_fetch_hint_takes_inline_path_without_lock(monkeypatch): +@pytest.mark.parametrize("mode", [None, "copy", "move"]) +def test_gpu_fetch_hint_takes_inline_path_without_lock(monkeypatch, mode): kvcr, results = _fetch_kvcr(monkeypatch) lock = _CountingLock(kvcr._core._state_lock) kvcr._core._state_lock = lock - kvcr.submit_hint(_router_hint("tcp://a:1", (1,)), "req") + hint = _router_hint("tcp://a:1", (1,)) + if mode is not None: + hint.actions[0].payload["mode"] = mode + kvcr.submit_hint(hint, "req") assert lock.entries == 0 + assert kvcr._core._remote_fw_dram._request_hints["req"].move is (mode == "move") assert kvcr.query((_key(1),), "req") == [_FETCHABLE] assert [(r.status, r.request_id) for r in results] == [(HintStatus.ACCEPTED, "req")] @pytest.mark.parametrize("one_envelope", [True, False], ids=["envelope", "calls"]) +@pytest.mark.parametrize("modes", [(None, None), ("copy", "move"), ("move", "copy")]) @pytest.mark.parametrize( ("second_source", "expected"), [("tcp://a:1", [_MISS, _FETCHABLE]), ("tcp://b:1", [_MISS, _MISS])], ids=["same_source_replaces", "other_source_fails"], ) -def test_second_fetch_for_a_request(monkeypatch, one_envelope, second_source, expected): +def test_second_fetch_for_a_request( + monkeypatch, one_envelope, second_source, expected, modes +): kvcr, results = _fetch_kvcr(monkeypatch) first, second = _fetch("a1", "tcp://a:1", [1]), _fetch("a2", second_source, [2]) + for hint, mode in zip((first, second), modes): + if mode is not None: + hint.payload["mode"] = mode for hints in [(first, second)] if one_envelope else [(first,), (second,)]: kvcr.submit_hint(_fetch_envelope(*hints), "req") # The store's rules apply within an envelope as across calls, and it reports @@ -124,6 +135,9 @@ def test_second_fetch_for_a_request(monkeypatch, one_envelope, second_source, ex ("a2", HintStatus.ACCEPTED, None), ] assert kvcr.query((_key(1), _key(2)), "req") == expected + # A same-source hint replaces the mode; a rejected conflicting hint does not. + kept = modes[1] if second_source == "tcp://a:1" else modes[0] + assert kvcr._core._remote_fw_dram._request_hints["req"].move is (kept == "move") def test_mixed_valid_and_unknown_hints(monkeypatch): @@ -154,3 +168,20 @@ def test_fetch_hint_without_request_id_is_accepted(monkeypatch): kvcr.submit_hint(_router_hint("tcp://a:1", (1,))) assert [(r.status, r.request_id) for r in results] == [(HintStatus.ACCEPTED, None)] assert not kvcr._core._remote_fw_dram._request_hints + + +def test_later_hint_replaces_move_mode_and_discard_clears_it(monkeypatch): + kvcr, results = _fetch_kvcr(monkeypatch) + moved = _fetch("move-1", "tcp://a:1", [1]) + moved.payload["mode"] = "move" + copied = _fetch("copy-1", "tcp://a:1", [2]) + hints = kvcr._core._remote_fw_dram._request_hints + + kvcr.submit_hint(_fetch_envelope(moved), "req") + assert hints["req"].move + kvcr.submit_hint(_fetch_envelope(copied), "req") + assert not hints["req"].move + assert kvcr.query((_key(1), _key(2)), "req") == [_MISS, _FETCHABLE] + kvcr.discard_hint("req") + assert "req" not in hints + assert all(r.status is HintStatus.ACCEPTED for r in results) diff --git a/tests/unit/test_hints_parser.py b/tests/unit/test_hints_parser.py index df2fa2e9..366310d9 100644 --- a/tests/unit/test_hints_parser.py +++ b/tests/unit/test_hints_parser.py @@ -61,6 +61,30 @@ def test_parse_hints_accepts_struct(): assert hints[0].block_hashes == frozenset({123, 456}) +@pytest.mark.parametrize("mode", [None, "copy", "move"]) +def test_fetch_hint_parses_optional_copy_or_move_mode(mode): + action = _fetch_hint("a", "tcp://source:1", (1, 2)) + if mode is not None: + action.payload["mode"] = mode + hints, rejected = parse_hints(_HintEnvelope("0.1", "m", [action]), _types()) + + assert not rejected + assert hints == [ + _KVFetchHint("a", "tcp://source:1", frozenset({1, 2}), move=mode == "move") + ] + + +@pytest.mark.parametrize("mode", [None, True, 1, "", "MOVE", "delete", [], {}]) +def test_fetch_hint_rejects_invalid_mode(mode): + action = _fetch_hint("a", "tcp://source:1", (1,), mode=mode) + hints, rejected = parse_hints(_HintEnvelope("0.1", "m", [action]), _types()) + + assert not hints + assert [(r.status, r.reason) for r in rejected] == [ + (HintStatus.REJECTED, "kv.fetch mode must be copy or move") + ] + + def test_parse_hints_rejects_invalid_actions_and_payloads(): hint = _HintEnvelope( "0.1", @@ -178,11 +202,6 @@ def parse(action_id, payload): {"source_control_endpoint": "tcp://source:1", "block_hashes": [True]}, {"source_control_endpoint": "tcp://source:1", "block_hashes": [-1]}, {"source_control_endpoint": "tcp://source:1", "block_hashes": [1 << 64]}, - { - "source_control_endpoint": "tcp://source:1", - "block_hashes": [1], - "mode": "move", - }, { "source_control_endpoint": "tcp://source:1", "block_hashes": [1], diff --git a/tests/unit/test_kvcr_invalidate.py b/tests/unit/test_kvcr_invalidate.py new file mode 100644 index 00000000..2eb31552 --- /dev/null +++ b/tests/unit/test_kvcr_invalidate.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Internal removal preserves existing readers and application retention.""" + +import ctypes + +import pytest +from _kvcr_test_utils import ( + _buffer_region, + _invalidate, + _mem_descriptor, + _new_local_kvcr, + _poll_until, + _wait_until, +) +from test_kvcr_move import _Pair +from test_kvcr_remote_target import _CopyingWriteAgent + +from kvcr.policy import FIFOPolicy +from kvcr.types import BlockKey, CacheTier, InventoryEvent, QueryStatus + + +class _Local: + def __init__(self, *, policy=None, slots=1): + self.agent = _CopyingWriteAgent() + self.local = ctypes.create_string_buffer(16 * slots) + self.framework = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16, 32) + self.events = [] + self.kvcr = _new_local_kvcr( + self.agent, + self.local, + slots, + inventory_sink=self.events.append, + framework_regions=[_buffer_region(self.framework, 16)], + policy=policy, + ) + + def deposit(self, key, *, index=0, no_evict=False): + self.agent.state = "DONE" + handle = self.kvcr.deposit({key: [_mem_descriptor(index)]}, no_evict=no_evict) + result = dict(_poll_until(self.kvcr, bool))[handle][key] + self.agent.state = "PROC" + assert result.success + return result + + +def test_invalidate_drains_existing_readers_and_rejects_new_reads(): + fixture = _Local() + kvcr = fixture.kvcr + key = BlockKey(b"active-reader") + fixture.deposit(key) + fetch = kvcr.fetch((key,)) + reader = dict(kvcr.poll_completed())[fetch][key].release_handle + assert reader is not None + + _invalidate(kvcr, (key, key)) + + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + another_fetch = kvcr.fetch((key,)) + another_delivery = kvcr.deliver({key: [_mem_descriptor(1)]}) + completed = dict(kvcr.poll_completed()) + assert kvcr._core._invalidation_handler.has_pending() + for handle in (another_fetch, another_delivery): + assert not completed[handle][key].success + assert completed[handle][key].release_handle is None + assert kvcr._core._block_record_map[key].local_dram.claim_count == 1 + assert fixture.local.raw == b"a" * 16 + assert fixture.framework.raw[16:] == b"b" * 16 + assert not any(event.removed for event in fixture.events) + pools = kvcr._core._local_dram._pools + assert pools.telemetry_state()["local_g2_free_slots"] == 0 + + assert kvcr.release((reader,)) == [(reader, True)] + assert not kvcr.poll_completed() + assert not kvcr._core._invalidation_handler.has_pending() + assert pools.telemetry_state()["local_g2_free_slots"] == 1 + assert pools.telemetry_state()["local_g2_evictable_slots"] == 0 + assert fixture.events.count(InventoryEvent((key,), CacheTier.LOCAL_G2, True)) == 1 + # Freed slots can be reused without resurrecting the removed key. + fixture.deposit(BlockKey(b"replacement"), index=1) + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + assert fixture.local.raw == b"b" * 16 + + +def test_existing_retention_skips_invalidation(): + fixture = _Local() + kvcr = fixture.kvcr + key = BlockKey(b"retained") + retention = fixture.deposit(key, no_evict=True).release_handle + fetch = kvcr.fetch((key,)) + reader = dict(kvcr.poll_completed())[fetch][key].release_handle + _invalidate(kvcr, (key,)) + assert not kvcr.poll_completed() + handler = kvcr._core._invalidation_handler + assert retention is not None + assert not handler.has_pending() + assert not kvcr._core._block_record_map[key].pending_invalidation + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + reopened = kvcr.fetch((key,)) + result = dict(kvcr.poll_completed())[reopened][key] + assert result.success and result.release_handle is not None + assert kvcr.release((result.release_handle,)) == [(result.release_handle, True)] + assert kvcr.release((reader,)) == [(reader, True)] + _invalidate(kvcr, (key,)) + assert not kvcr.poll_completed() + assert not handler.has_pending() + assert not any(event.removed for event in fixture.events) + assert fixture.local.raw == b"a" * 16 + + # Releasing retention does not revive skipped removal requests. + assert kvcr.release((retention,)) == [(retention, True)] + assert not kvcr.poll_completed() + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + _invalidate(kvcr, (key,)) + assert not kvcr.poll_completed() + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + assert fixture.events.count(InventoryEvent((key,), CacheTier.LOCAL_G2, True)) == 1 + + +@pytest.mark.parametrize("no_evict", [False, True]) +def test_pending_invalidation_rejects_only_affected_deposit_keys(no_evict): + fixture = _Local(slots=2) + kvcr = fixture.kvcr + key, other = BlockKey(b"removing"), BlockKey(b"unrelated") + fixture.deposit(key) + fetch = kvcr.fetch((key,)) + reader = dict(kvcr.poll_completed())[fetch][key].release_handle + _invalidate(kvcr, (key,)) + + deposit = kvcr.deposit( + {key: [_mem_descriptor()], other: [_mem_descriptor(1)]}, no_evict=no_evict + ) + _wait_until(lambda: len(fixture.agent.transfers) == 2) + record = kvcr._core._block_record_map[key] + assert record.pending_invalidation + assert record.local_dram.retention_claims == 0 + assert record.local_dram.claim_count == 1 + assert not any(event.removed for event in fixture.events) + + assert kvcr.release((reader,)) == [(reader, True)] + # A rejected key must not wait for the unrelated admitted copy to finish. + assert not kvcr.poll_completed() + assert not kvcr._core._invalidation_handler.has_pending() + assert key not in kvcr._core._block_record_map + assert fixture.events.count(InventoryEvent((key,), CacheTier.LOCAL_G2, True)) == 1 + fixture.agent.state = "DONE" + results = dict(_poll_until(kvcr, bool))[deposit] + assert not results[key].success and results[key].release_handle is None + assert results[other].success + assert (results[other].release_handle is not None) is no_evict + assert kvcr.query((key, other)) == [ + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + assert fixture.local.raw == b"a" * 16 + b"b" * 16 + retained = results[other].release_handle + if retained is not None: + assert kvcr.release((retained,)) == [(retained, True)] + fixture.deposit(key) + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + + +@pytest.mark.parametrize( + "no_evict,duplicate,succeeds", + [ + (False, False, True), + (False, True, True), + (True, False, True), + (True, True, True), + (True, False, False), + ], +) +def test_inflight_deposits_obey_invalidation_order(no_evict, duplicate, succeeds): + fixture = _Local() + kvcr = fixture.kvcr + key = BlockKey(b"copying") + blocks = {key: [_mem_descriptor()]} + writer = kvcr.deposit(blocks, no_evict=no_evict and not duplicate) + _wait_until(lambda: bool(fixture.agent.transfers)) + deposit = kvcr.deposit(blocks, no_evict=no_evict) if duplicate else writer + _invalidate(kvcr, (key,)) + assert not kvcr.poll_completed() + assert kvcr._core._invalidation_handler.has_pending() + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + residency = kvcr._core._block_record_map[key].local_dram + assert residency.claim_count == 1 + + if not succeeds: + states = iter(("ERR", "DONE")) + fixture.agent.check_xfer_state = lambda _: next(states, "DONE") + else: + fixture.agent.state = "DONE" + completed = dict(_poll_until(kvcr, lambda results: len(results) == 1 + duplicate)) + assert completed[writer][key].success is succeeds + result = completed[deposit][key] + assert result.success is succeeds + assert len(fixture.agent.transfers) == 1 + if no_evict and succeeds: + retention = result.release_handle + assert retention is not None + assert not kvcr._core._invalidation_handler.has_pending() + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert kvcr.release((retention,)) == [(retention, True)] + assert not kvcr.poll_completed() + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + _invalidate(kvcr, (key,)) + elif not succeeds: + assert result.release_handle is None + assert not kvcr._core._local_dram._pools._public_claims + # Failed protected operations must not block later source removal. + fixture.deposit(key) + _invalidate(kvcr, (key,)) + else: + assert result.release_handle is None + assert not kvcr.poll_completed() + assert residency.retention_claims == 0 + assert not kvcr._core._local_dram._pools._public_claims + assert not kvcr._core._invalidation_handler.has_pending() + assert key not in kvcr._core._block_record_map + assert fixture.events.count(InventoryEvent((key,), CacheTier.LOCAL_G2, True)) == 1 + + +def test_releasing_the_last_claim_frees_a_pending_slot_immediately(): + fixture = _Local() + kvcr = fixture.kvcr + key, other = BlockKey(b"released"), BlockKey(b"next") + fixture.deposit(key) + fetch = kvcr.fetch((key,)) + reader = dict(kvcr.poll_completed())[fetch][key].release_handle + _invalidate(kvcr, (key,)) + pools = kvcr._core._local_dram._pools + assert pools.telemetry_state()["local_g2_free_slots"] == 0 + + # No poll in between: release itself drops the copy and frees its slot. + assert kvcr.release((reader,)) == [(reader, True)] + assert pools.telemetry_state()["local_g2_free_slots"] == 1 + assert key not in kvcr._core._block_record_map + assert fixture.deposit(other, index=1).success + assert fixture.local.raw == b"b" * 16 + + +def test_fill_finishing_during_pending_invalidation_stays_withdrawn_from_readers(): + pair = _Pair() + key = BlockKey(b"republish-guard") + source = pair.source + core = source._core + pools = core._local_dram._pools + deposit = source.deposit( + {key: pair.descriptors()}, no_evict=False, expected_layout=list(pair.labels) + ) + _wait_until(lambda: len(pair.source_agent.xfers) == 1) + read = source.fetch((key,)) + _invalidate(source, (key,)) + assert core._block_record_map[key].pending_invalidation + + pair.source_agent.state = "DONE" + results = dict(_poll_until(source, lambda done: len(done) == 2)) + pair.source_agent.state = "PROC" + assert results[deposit][key].success and results[read][key].success + reader = results[read][key].release_handle + assert reader is not None + # The admitted fill landed, but the key must not reopen for remote readers. + assert core._invalidation_handler.has_pending() + assert key not in pools._source_state.index + assert source.query((key,)) == [(QueryStatus.MISS, None)] + + before = len(pair.source_agent.xfers) + copied = pair.result(pair.start((key,), move=False))[key] + assert not copied.success and copied.release_handle is None + assert len(pair.source_agent.xfers) == before + assert key not in pools._source_state.index + + assert source.release((reader,)) == [(reader, True)] + assert key not in core._block_record_map + assert not core._invalidation_handler.has_pending() + assert pools.telemetry_state()["local_g2_free_slots"] == 1 + + +def test_removal_callback_can_redeposit_without_a_stale_removal_event(): + class RedepositPolicy(FIFOPolicy): + callback = None + + def on_remove(self, meta): + super().on_remove(meta) + callback, self.callback = self.callback, None + if callback is not None: + callback(meta) + + policy = RedepositPolicy() + fixture = _Local(policy=policy) + kvcr = fixture.kvcr + key = BlockKey(b"replaced-during-removal") + fixture.deposit(key) + fixture.events.clear() + replacement = [] + + def redeposit(meta): + assert meta.block_key == key + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + # deposit() polls recursively; the new residency must survive that poll. + replacement.append(fixture.deposit(key, index=1, no_evict=True)) + + policy.callback = redeposit + _invalidate(kvcr, (key,)) + assert not kvcr.poll_completed() + + assert not kvcr._core._invalidation_handler.has_pending() + assert len(replacement) == 1 and replacement[0].success + retention = replacement[0].release_handle + assert retention is not None + assert kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert fixture.local.raw == b"b" * 16 + assert fixture.events == [InventoryEvent((key,), CacheTier.LOCAL_G2, False)] + assert kvcr.release((retention,)) == [(retention, True)] diff --git a/tests/unit/test_kvcr_move.py b/tests/unit/test_kvcr_move.py new file mode 100644 index 00000000..cea1eee4 --- /dev/null +++ b/tests/unit/test_kvcr_move.py @@ -0,0 +1,1205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Move transfers, early source withdrawal, and cached-key invalidation.""" + +import ctypes +import os +import time +from collections import Counter +from dataclasses import replace + +import msgspec +import pytest +from _kvcr_test_utils import ( + FakeBytesControl, + FakePrimaryPinning, + _buffer_region, + _ConstantHashAdapter, + _decode_control_message, + _has_outstanding_operations, + _HintAction, + _HintEnvelope, + _invalidate, + _mem_descriptor, + _new_kvcr, + _poll_until, + _wait_until, +) +from test_g3 import _FakeG3Agent, _MoveLocalToG3Policy +from test_kvcr_remote_target import _CopyingWriteAgent + +from kvcr.backends.remote_ops import _Invalidate +from kvcr.config import G3Options, KVCRConfig, LocalDramOptions, RemoteFWDramOptions +from kvcr.policy import FIFOPolicy +from kvcr.types import BlockKey, CacheTier, PlacementAction, QueryStatus + + +class _KeepPolicy(FIFOPolicy): + def decide_eviction(self, meta, source): + return PlacementAction.KEEP, None + + +class _IndividualWriteAgent(_CopyingWriteAgent): + def __init__(self): + super().__init__(metadata=b"source-md") + self.states = {} + + def check_xfer_state(self, handle): + return self.states.get(handle, self.state) + + +class _IndividualG3WriteAgent(_FakeG3Agent): + def __init__(self): + super().__init__() + self.states = {} + + def transfer(self, handle): + if self._xfer_backends[handle] == ("MOCK",): + return super().transfer(handle) + return _CopyingWriteAgent.transfer(self, handle) + + def check_xfer_state(self, handle): + return self.states.get(handle, self.state) + + +class _Pair: + """Real buffers with manually forwarded control messages and completions.""" + + def __init__( + self, + *, + slots=1, + labels=("",), + source_agent=None, + target_labels=None, + opportunistic_query=False, + block_size=16, + source_g3=None, + source_policy=None, + ): + self.labels = labels + target_labels = labels if target_labels is None else target_labels + self.results = [{}, {}] + self.control_offsets = [0, 0] + self.notif_offsets = [0, 0] + self.delivered_native = set() + self.source_now = 0.0 + self.source_events = [] + self.target_events = [] + self.source_agent = source_agent or _CopyingWriteAgent(metadata=b"source-md") + self.target_agent = _CopyingWriteAgent(metadata=b"target-md") + self.source_control = FakeBytesControl("tcp://source:1") + self.target_control = FakeBytesControl("tcp://target:1") + self.source_pinning = FakePrimaryPinning() + self.source_primary = {} + self.target_primary = {} + self.source_local = {} + self.target_local = {} + for index, label in enumerate(labels): + self.source_primary[label] = ctypes.create_string_buffer( + (slots + 1) * block_size + ) + self.source_primary[label].raw = bytes([97 + index]) * ( + (slots + 1) * block_size + ) + self.source_local[label] = ctypes.create_string_buffer(slots * block_size) + if label in target_labels: + self.target_primary[label] = ctypes.create_string_buffer( + (slots + 1) * block_size + ) + self.target_local[label] = ctypes.create_string_buffer( + slots * block_size + ) + config = KVCRConfig( + nixl_agent_name="source", + pool_layouts=[(label, block_size) for label in labels], + ) + + def dram(memories): + return LocalDramOptions( + [ + (label, ctypes.addressof(memory), len(memory)) + for label, memory in memories.items() + ] + ) + + self.source = _new_kvcr( + self.source_agent, + self.source_pinning, + self.source_control, + config, + name="source", + local_dram=dram(self.source_local), + g3=source_g3, + framework_regions=[ + _buffer_region(memory, block_size, label) + for label, memory in self.source_primary.items() + ], + policy=source_policy or _KeepPolicy(), + key_adapter=_ConstantHashAdapter(), + inventory_sink=self.source_events.append, + ) + self.source._core._clock = lambda: self.source_now + self.target = _new_kvcr( + self.target_agent, + FakePrimaryPinning(), + self.target_control, + KVCRConfig( + nixl_agent_name="target", + pool_layouts=[(label, block_size) for label in target_labels], + ), + local_dram=dram(self.target_local), + framework_regions=[ + _buffer_region(memory, block_size, label) + for label, memory in self.target_primary.items() + ], + key_adapter=_ConstantHashAdapter(), + remote_options=RemoteFWDramOptions( + eager_ctrl_connect=False, opportunistic_query=opportunistic_query + ), + inventory_sink=self.target_events.append, + ) + self.target._core._clock = lambda: 0.0 + + def descriptors(self, index=0, *, labels=None, endpoint="source"): + return [ + _mem_descriptor(index, label=label, end_point_name=endpoint) + for label in self.labels + if labels is None or label in labels + ] + + def deposit(self, key, *, index=0, no_evict=True): + self.source_agent.state = "DONE" + handle = self.source.deposit( + {key: self.descriptors(index)}, + no_evict=no_evict, + expected_layout=list(self.labels), + ) + result = dict(_poll_until(self.source, bool))[handle][key] + self.source_agent.state = "PROC" + return result + + def start(self, keys, *, move=True, destinations=None): + self.advise("move" if move else "copy") + if destinations is not None: + return self.target.deliver(destinations, "req") + return self.target.fetch(keys, "req", expected_layout=list(self.labels)) + + def advise(self, mode="copy", *, hashes=(123,), source="tcp://source:1"): + self.target.submit_hint( + _HintEnvelope( + "0.1", + "move-message", + [ + _HintAction( + "move-action", + "kv.fetch", + "1.0", + { + "mode": mode, + "source_control_endpoint": source, + "block_hashes": list(hashes), + }, + ) + ], + ), + "req", + ) + + def pump(self, *, target_requests=True, source_poll=True): + for index, (worker, control, agent, other_control, other_agent) in enumerate( + ( + ( + self.source, + self.source_control, + self.source_agent, + self.target_control, + self.target_agent, + ), + ( + self.target, + self.target_control, + self.target_agent, + self.source_control, + self.source_agent, + ), + ) + ): + if index or source_poll: + self.results[index].update(worker.poll_completed()) + sent = control.sent[self.control_offsets[index] :] + self.control_offsets[index] += len(sent) + for _, raw in sent: + if index == 1 and not target_requests: + continue + other_control.incoming.append(raw) + notifications = agent.sent_notifs[self.notif_offsets[index] :] + self.notif_offsets[index] += len(notifications) + other_agent.notifs.setdefault("peer", []).extend( + raw for _, raw in notifications + ) + for transfer, data in enumerate(agent.xfers, 1): + native_id = index, transfer + if ( + native_id in self.delivered_native + or transfer not in agent.released_xfers + ): + continue + self.delivered_native.add(native_id) + if data[5]: + other_agent.notifs.setdefault("peer", []).append(data[5]) + time.sleep(0.001) + + def until(self, predicate, **kwargs): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + self.pump(**kwargs) + if predicate(): + return + raise AssertionError("Move did not reach the expected state") + + def result(self, handle): + self.until(lambda: handle in self.results[1]) + return self.results[1].pop(handle) + + def copy_done(self): + self.source_agent.state = "DONE" + + +def _requests(pair, worker="target", message_type=None): + messages = [ + _decode_control_message(raw) + for _, raw in getattr(pair, f"{worker}_control").sent + ] + return [ + item + for item in messages + if message_type is None or item["type"] == message_type + ] + + +def _invalidations(pair): + """Cached keys the destination asked the source to drop, per message.""" + return [tuple(item["keys"]) for item in _requests(pair, message_type="invalidate")] + + +def _inventory_keys(events, *, removed): + assert all(event.tier is CacheTier.LOCAL_G2 for event in events) + return Counter( + key for event in events if event.removed is removed for key in event.keys + ) + + +def _cache_at_target(pair, key, *, index=0): + pair.target_agent.state = "DONE" + deposited = pair.target.deposit( + {key: pair.descriptors(index, endpoint="target")}, + expected_layout=list(pair.labels), + ) + result = dict(_poll_until(pair.target, bool))[deposited][key] + assert result.success and result.release_handle is None + + +@pytest.mark.parametrize("moving", [False, True], ids=["copy", "move"]) +@pytest.mark.parametrize("cached_count", [0, 1, 2]) +@pytest.mark.parametrize("protected", [False, True]) +def test_copy_and_move_with_destination_cache_hits(moving, cached_count, protected): + pair = _Pair(slots=2) + keys = (BlockKey(b"first"), BlockKey(b"second")) + pair.source_primary[""].raw = b"a" * 16 + b"b" * 16 + b"c" * 16 + pair.target_primary[""].raw = pair.source_primary[""].raw + retentions = [] + for index, key in enumerate(keys): + result = pair.deposit(key, index=index, no_evict=protected) + if result.release_handle is not None: + retentions.append(result.release_handle) + if index < cached_count: + _cache_at_target(pair, key, index=index) + pair.source_events.clear() + pair.target_events.clear() + before = len(pair.source_agent.xfers) + retry = moving and cached_count == 1 and not protected + if retry: + # The target's metadata acknowledgment outlived the source's route cache. + pair.target._core._remote_fw_dram._metadata_acked_sources.add("tcp://source:1") + handle = pair.start(keys, move=moving) + if cached_count == 2: + results = dict(pair.target.poll_completed())[handle] + else: + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + if moving and not protected: + # Moving keys stop serving at admission; cached keys are dropped + # while the physical transfer is still running. + pair.until( + lambda: pair.source.query(keys) == [(QueryStatus.MISS, None)] * 2 + ) + pair.until( + lambda: all( + key not in pair.source._core._block_record_map + for key in keys[:cached_count] + ) + ) + assert _inventory_keys(pair.source_events, removed=True) == Counter( + keys[:cached_count] + ) + assert pair.source._core._invalidation_handler.has_pending() + elif moving: + # A no_evict guarantee keeps protected copies readable throughout. + for _ in range(3): + pair.pump() + assert ( + pair.source.query(keys) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] * 2 + ) + assert not pair.source._core._invalidation_handler.has_pending() + pair.copy_done() + results = pair.result(handle) + assert all(result.success for result in results.values()) + assert pair.target_local[""].raw == b"a" * 16 + b"b" * 16 + expected = ( + (QueryStatus.MISS, None) + if moving and not protected + else (QueryStatus.HIT, CacheTier.LOCAL_G2) + ) + pair.until(lambda: pair.source.query(keys) == [expected, expected]) + pair.until(lambda: not pair.source._core._invalidation_handler.has_pending()) + requests = _requests(pair, message_type="start_write") + assert len(requests) == int(cached_count < 2) + int(retry) + if retry: + assert "target_agent_metadata" not in requests[0] + assert requests[1]["target_agent_metadata"] == b"target-md" + for field in ("keys", "dst_descriptors", "move"): + assert requests[0][field] == requests[1][field] + for request in requests: + assert len(request["keys"]) == len(keys) - cached_count + assert request.get("move", False) is moving + assert all(request["dst_descriptors"]) + standalone = _requests(pair, message_type="invalidate") + assert len(standalone) == int(moving and cached_count > 0) + if standalone: + assert standalone[0]["keys"] == list(keys[:cached_count]) + assert standalone[0]["key_layout"] == list(pair.labels) + assert len(pair.source_agent.xfers) == before + int(cached_count < 2) + assert _inventory_keys(pair.source_events, removed=True) == Counter( + keys if moving and not protected else () + ) + assert _inventory_keys(pair.target_events, removed=False) == Counter( + keys[cached_count:] + ) + assert not _inventory_keys(pair.target_events, removed=True) + handles = tuple(result.release_handle for result in results.values()) + assert all(handle is not None for handle in handles) + assert pair.target.release(handles) == [(handle, True) for handle in handles] + assert pair.source.release(retentions) == [(handle, True) for handle in retentions] + + +@pytest.mark.parametrize("replace_advice", [False, True]) +def test_failed_move_preserves_source_and_successful_cached_result(replace_advice): + pair = _Pair(slots=2) + cached, physical = BlockKey(b"cached"), BlockKey(b"physical") + pair.target_primary[""].raw = pair.source_primary[""].raw + pair.deposit(cached, no_evict=False) + pair.deposit(physical, index=1, no_evict=False) + _cache_at_target(pair, cached) + before = len(pair.source_agent.xfers) + moved = pair.start((cached, physical)) + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + if replace_advice: + pair.target.discard_hint("req") + pair.advise(source="tcp://replacement:1") + replacement = pair.target._core._remote_fw_dram._request_hints["req"] + pair.source_agent.state = "ERR" + results = pair.result(moved) + assert results[cached].success and not results[physical].success + assert results[physical].release_handle is None + pair.source_agent.state = "DONE" + pair.until(lambda: not _has_outstanding_operations(pair.source)) + pair.until(lambda: not pair.source._core._invalidation_handler.has_pending()) + assert _invalidations(pair) == [(cached,)] + assert not _requests(pair, message_type="move_done") + assert pair.source.query((cached, physical)) == [ + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + assert pair.source_local[""].raw == b"a" * 32 + assert pair.target.release((results[cached].release_handle,)) == [ + (results[cached].release_handle, True) + ] + if replace_advice: + assert pair.target._core._remote_fw_dram._request_hints["req"] == replacement + + +@pytest.mark.parametrize("moving", [False, True], ids=["copy", "move"]) +@pytest.mark.parametrize( + "local_piece,advice_change", + [ + ("ready", None), + ("inflight", None), + ("inflight", "discard"), + ("inflight", "replace"), + ("inflight", "conflict"), + ("fails", None), + ], +) +def test_partial_move_settles_local_writers_before_source_removal( + moving, local_piece, advice_change, monkeypatch +): + pair = _Pair(labels=("a", "b")) + key = BlockKey(b"partial") + pair.deposit(key, no_evict=False) + pair.target_primary["a"].raw = b"a" * 32 + pair.target_agent.state = "DONE" if local_piece == "ready" else "PROC" + deposited = pair.target.deposit( + {key: pair.descriptors(labels=("a",), endpoint="target")}, + expected_layout=["a", "b"], + ) + _wait_until(lambda: bool(pair.target_agent.xfers)) + if local_piece == "ready": + assert dict(_poll_until(pair.target, bool))[deposited][key].success + reported = [] + poll_completed = pair.target.poll_completed + + def observe_completions(): + for item in poll_completed(): + reported.append(item[0]) + yield item + + monkeypatch.setattr(pair.target, "poll_completed", observe_completions) + handle = pair.start((key,), move=moving) + if advice_change in ("discard", "replace"): + pair.target.discard_hint("req") + if advice_change in ("replace", "conflict"): + pair.advise(hashes=(999,), source="tcp://replacement:1") + future_hint = pair.target._core._remote_fw_dram._request_hints.get("req") + if moving and local_piece != "ready": + # Source success will remove the whole key, so a local piece must first + # succeed or become part of the same remote transfer. + for _ in range(3): + pair.pump() + assert not _requests(pair, message_type="start_write") + assert handle not in pair.results[1] + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + if local_piece == "fails": + states = iter(("ERR", "DONE")) + pair.target_agent.check_xfer_state = lambda _: next(states, "DONE") + else: + pair.target_agent.state = "DONE" + pair.until(lambda: deposited in pair.results[1]) + assert pair.results[1].pop(deposited)[key].success is (local_piece != "fails") + pair.until(lambda: bool(_requests(pair, message_type="start_write"))) + request = _requests(pair, message_type="start_write")[0] + assert request["source_control_endpoint"] == "tcp://source:1" + assert request["key_layout"] == ["a", "b"] + assert request.get("move", False) is moving + assert [ref["label"] for ref in request["dst_descriptors"][0]] == ( + ["a", "b"] if moving and local_piece == "fails" else ["b"] + ) + if local_piece == "fails" and not moving: + # Ordinary copying can still fill its missing pieces independently. + states = iter(("ERR", "DONE")) + pair.target_agent.check_xfer_state = lambda _: next(states, "DONE") + pair.until(lambda: deposited in pair.results[1]) + assert not pair.results[1].pop(deposited)[key].success + pair.copy_done() + pair.until(lambda: len(_requests(pair, message_type="start_write")) == 2) + second = _requests(pair, message_type="start_write")[1] + assert [ref["label"] for ref in second["dst_descriptors"][0]] == ["a"] + else: + pair.copy_done() + if local_piece == "inflight" and not moving: + for _ in range(3): + pair.pump() + assert handle not in pair.results[1] + assert not _invalidations(pair) + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + pair.target_agent.state = "DONE" + result = pair.result(handle)[key] + assert result.success + expected = ( + (QueryStatus.MISS, None) if moving else (QueryStatus.HIT, CacheTier.LOCAL_G2) + ) + pair.until(lambda: pair.source.query((key,)) == [expected]) + pair.until(lambda: not pair.source._core._invalidation_handler.has_pending()) + assert not _invalidations(pair) and not _requests(pair, message_type="move_done") + assert _inventory_keys(pair.target_events, removed=False) == Counter((key,)) + assert pair.source._core._local_dram._pools.telemetry_state()[ + "local_g2_free_slots" + ] == (2 if moving else 0) + assert pair.target_local["a"].raw == b"a" * 16 + assert pair.target_local["b"].raw == b"b" * 16 + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + pair.until( + lambda: ( + not _has_outstanding_operations(pair.source) + and not _has_outstanding_operations(pair.target) + ) + ) + assert not pair.target._core._local_dram._pending_residency_ops + assert reported.count(handle) == 1 + if advice_change is not None: + assert ( + pair.target._core._remote_fw_dram._request_hints.get("req") == future_hint + ) + + +@pytest.mark.parametrize( + "cached,slow_source", [(False, False), (False, True), (True, False)] +) +def test_move_withdraws_before_transfer_then_drains_existing_readers( + cached, slow_source, monkeypatch +): + pair = _Pair() + key = BlockKey(b"draining") + pair.deposit(key, no_evict=False) + if cached: + pair.target_primary[""].raw = pair.source_primary[""].raw + _cache_at_target(pair, key) + read = pair.source.fetch((key,)) + guard = dict(pair.source.poll_completed())[read][key].release_handle + core = pair.source._core + pools = core._local_dram._pools + residency = core._block_record_map[key].local_dram + if slow_source: + monkeypatch.setattr( + core._remote_fw_dram, + "_try_local_source_write", + lambda *args, **kwargs: False, + ) + before = len(pair.source_agent.xfers) + moved = pair.start((key,)) + pair.until(lambda: core._invalidation_handler.has_pending()) + if not cached: + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + assert pair.source.query((key,)) == [(QueryStatus.MISS, None)] + assert core._invalidation_handler.has_pending() + assert ( + pools.acquire_readers((key,), (residency.refs,), residency.key_layout) is False + ) + rejected = [ + pair.source.fetch((key,)), + pair.source.deliver({key: pair.descriptors()}), + pair.source.deposit({key: pair.descriptors()}, no_evict=False), + pair.source.deposit({key: pair.descriptors()}, no_evict=True), + ] + completed = dict(pair.source.poll_completed()) + assert all(not completed[handle][key].success for handle in rejected) + assert len(pair.source_agent.xfers) == before + int(not cached) + assert _invalidations(pair) == ([(key,)] if cached else []) + with core._state_lock: + pair.copy_done() + pair.until(lambda: moved in pair.results[1], source_poll=False) + result = pair.results[1].pop(moved)[key] + assert result.success and pair.target_local[""].raw == b"a" * 16 + assert core._block_record_map[key].local_dram is residency + assert pair.source.query((key,)) == [(QueryStatus.MISS, None)] + assert core._block_record_map[key].local_dram is residency + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + _invalidate(pair.target, (key,)) + pair.until(lambda: not pair.target._core._invalidation_handler.has_pending()) + pair.source_pinning.prefix_length = 0 + before = len(pair.source_agent.xfers) + copied = pair.result(pair.start((key,), move=False))[key] + assert not copied.success and len(pair.source_agent.xfers) == before + assert pair.source_local[""].raw == b"a" * 16 + dropped = [] + original_drop = pools._drop + + def drop_without_reader_lock(key): + assert pools._source_state.lock.acquire(blocking=False) + pools._source_state.lock.release() + original_drop(key) + dropped.append(key) + + monkeypatch.setattr(pools, "_drop", drop_without_reader_lock) + assert pair.source.release((guard,)) == [(guard, True)] + pair.until(lambda: bool(dropped)) + assert dropped == [key] and key not in core._block_record_map + assert pools.telemetry_state()["local_g2_free_slots"] == 1 + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize( + "advice", + [ + "missing", + "discarded", + "conflicting", + "unmatched", + "copy", + "no_adapter", + "known_missing", + ], +) +def test_fetch_preserves_copy_semantics_without_usable_move_advice(cached, advice): + pair = _Pair() + key = BlockKey(b"copy-fallback") + pair.target_primary[""].raw = pair.source_primary[""].raw + pair.deposit(key, no_evict=False) + if cached: + _cache_at_target(pair, key) + if advice == "discarded": + pair.advise("move") + pair.target.discard_hint("req") + elif advice == "conflicting": + pair.advise("move") + pair.advise("move", source="tcp://different-source:1") + elif advice == "unmatched": + pair.advise("move", hashes=(999,)) + elif advice == "copy": + pair.advise("copy") + elif advice == "no_adapter": + pair.advise("move") + pair.target._core._remote_fw_dram._key_adapter = None + elif advice == "known_missing": + pair.advise("move") + remote = pair.target._core._remote_fw_dram + remote._request_hints["req"] = replace( + remote._request_hints["req"], missing_keys=frozenset((key,)) + ) + handle = pair.target.fetch((key,), "req") + pair.copy_done() + result = pair.result(handle)[key] + assert result.success is (cached or advice == "copy") + for _ in range(3): + pair.pump() + assert not _invalidations(pair) + assert len(_requests(pair, message_type="start_write")) == int( + not cached and advice == "copy" + ) + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + if result.success: + assert result.release_handle is not None + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + else: + assert result.release_handle is None + + +@pytest.mark.parametrize( + "cached,opportunistic", [(True, False), (False, False), (False, True)] +) +def test_fetch_moves_only_hash_matched_keys(cached, opportunistic, monkeypatch): + pair = _Pair(slots=2, opportunistic_query=opportunistic) + moved, copied = BlockKey(b"moved"), BlockKey(b"copied") + pair.source_primary[""].raw = b"a" * 16 + b"b" * 16 + b"c" * 16 + pair.target_primary[""].raw = pair.source_primary[""].raw + for index, key in enumerate((moved, copied)): + pair.deposit(key, index=index, no_evict=False) + if cached: + _cache_at_target(pair, key, index=index) + monkeypatch.setattr( + pair.target._core._remote_fw_dram._key_adapter, + "encode_to_router", + lambda key: 123 if key == moved else 456, + ) + pair.advise("move", hashes=(123,)) + handle = pair.target.fetch((moved, copied), "req") + pair.copy_done() + results = pair.result(handle) + assert results[moved].success + assert results[copied].success is cached + pair.until(lambda: pair.source.query((moved,)) == [(QueryStatus.MISS, None)]) + pair.until(lambda: not pair.source._core._invalidation_handler.has_pending()) + assert _invalidations(pair) == ([(moved,)] if cached else []) + for request in _requests(pair, message_type="start_write"): + assert request["keys"] == [moved] and request["move"] + assert pair.source.query((moved, copied)) == [ + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + for key, result in results.items(): + if not result.success: + continue + address = ( + ctypes.addressof(pair.target_local[""]) + + 16 * result.descriptors[0].element_index + ) + assert ctypes.string_at(address, 16) == (b"a" if key == moved else b"b") * 16 + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + + +@pytest.mark.parametrize("succeeds", [False, True]) +def test_move_joining_a_copy_invalidates_only_after_complete_residency(succeeds): + pair = _Pair() + key = BlockKey(b"joined-copy") + pair.deposit(key, no_evict=False) + before = len(pair.source_agent.xfers) + copied = pair.start((key,), move=False) + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + moved = pair.start((key,)) + for _ in range(3): + pair.pump() + assert not _invalidations(pair) + requests = _requests(pair, message_type="start_write") + assert len(requests) == 1 and not requests[0].get("move", False) + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + + pair.source_agent.state = "DONE" if succeeds else "ERR" + pair.until(lambda: copied in pair.results[1] and moved in pair.results[1]) + results = [pair.results[1].pop(handle)[key] for handle in (copied, moved)] + assert all(result.success is succeeds for result in results) + pair.source_agent.state = "DONE" + pair.until(lambda: not _has_outstanding_operations(pair.source)) + if succeeds: + pair.until(lambda: key not in pair.source._core._block_record_map) + assert _invalidations(pair) == [(key,)] + else: + assert not _invalidations(pair) + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert not _requests(pair, message_type="move_done") + assert len(pair.source_agent.xfers) == before + 1 + assert _inventory_keys(pair.target_events, removed=False) == Counter( + (key,) if succeeds else () + ) + handles = tuple(result.release_handle for result in results if result.success) + assert pair.target.release(handles) == [(handle, True) for handle in handles] + + +def test_fetch_ignores_self_move_advice_for_local_hit(): + pair = _Pair() + key = BlockKey(b"self-source") + _cache_at_target(pair, key) + pair.advise("move", source=pair.target_control.endpoint) + handle = pair.target.fetch((key,), "req", expected_layout=[""]) + result = dict(pair.target.poll_completed())[handle][key] + assert result.success and result.release_handle is not None + for _ in range(3): + pair.pump() + assert not _requests(pair) + assert pair.target.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + + +def test_deliver_copies_under_move_advice(): + pair = _Pair() + key = BlockKey(b"delivery") + pair.deposit(key, no_evict=False) + handle = pair.start((key,), destinations={key: pair.descriptors(endpoint="target")}) + pair.copy_done() + result = pair.result(handle)[key] + assert result.success and result.release_handle is None + assert pair.target_primary[""].raw[:16] == b"a" * 16 + assert pair.source.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert not _invalidations(pair) + + +@pytest.mark.parametrize("after_allowance", [False, True]) +def test_move_notification_uses_extended_destination_allowance(after_allowance): + pair = _Pair() + key = BlockKey(b"late") + pair.deposit(key, no_evict=False) + held = [] + real_get_new_notifs = pair.target_agent.get_new_notifs + + def hold_notifs(): + for raw in real_get_new_notifs().values(): + held.extend(raw) + return {} + + pair.target_agent.get_new_notifs = hold_notifs + before = len(pair.source_agent.xfers) + handle = pair.start((key,)) + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + source_writes = [ + op + for op in pair.source._core._progress._in_flight_ops.values() + if op.op_id[0] == "source" + ] + assert source_writes and all( + op.deadline <= pair.source._core.config.operation_timeout_ms / 1000 + for op in source_writes + ) + pair.copy_done() + # Source cleanup finishes even while the destination cannot see write_done. + pair.until(lambda: key not in pair.source._core._block_record_map) + _wait_until(lambda: bool(held)) + assert handle not in pair.results[1] + config = pair.target._core.config + elapsed = ( + config.abandon_timeout_ms / 1000 + 0.1 + if after_allowance + else config.operation_timeout_ms / 1000 + 0.1 + ) + pair.target._core._clock = lambda: elapsed + if not after_allowance: + for _ in range(3): + pair.pump() + assert handle not in pair.results[1] + pair.target_agent.get_new_notifs = real_get_new_notifs + pair.target_agent.notifs.setdefault("peer", []).extend(held) + result = pair.result(handle)[key] + assert result.success is not after_allowance + assert (result.release_handle is None) is after_allowance + pair.until( + lambda: ( + not _has_outstanding_operations(pair.source) + and not _has_outstanding_operations(pair.target) + ) + ) + assert not _invalidations(pair) and not _requests(pair, message_type="move_done") + assert pair.source.query((key,)) == [(QueryStatus.MISS, None)] + assert pair.source_local[""].raw == b"a" * 16 + assert pair.target.query((key,)) == [ + (QueryStatus.MISS, None) + if after_allowance + else (QueryStatus.HIT, CacheTier.LOCAL_G2) + ] + assert _inventory_keys(pair.source_events, removed=True) == Counter((key,)) + assert _inventory_keys(pair.target_events, removed=False) == Counter( + () if after_allowance else (key,) + ) + if result.release_handle is not None: + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + + +def test_mixed_move_reports_missing_source_against_original_key_indices(): + pair = _Pair(slots=3) + cached, missing, physical = map(BlockKey, (b"cached", b"missing", b"physical")) + pair.source_primary[""].raw = b"a" * 16 + b"b" * 16 + b"c" * 16 + b"d" * 16 + pair.target_primary[""].raw = pair.source_primary[""].raw + pair.deposit(cached, no_evict=False) + pair.deposit(physical, index=2, no_evict=False) + pair.source_pinning.prefix_length = 0 + _cache_at_target(pair, cached) + handle = pair.start((cached, missing, physical)) + pair.copy_done() + results = pair.result(handle) + assert results[cached].success and results[physical].success + assert not results[missing].success and results[missing].release_handle is None + address = ( + ctypes.addressof(pair.target_local[""]) + + results[physical].descriptors[0].element_index * 16 + ) + assert ctypes.string_at(address, 16) == b"c" * 16 + requests = _requests(pair, message_type="start_write") + assert len(requests) == 1 and requests[0]["keys"] == [missing, physical] + assert requests[0]["move"] + assert _invalidations(pair) == [(cached,)] + pair.until( + lambda: pair.source.query((cached, physical)) == [(QueryStatus.MISS, None)] * 2 + ) + pair.until(lambda: not pair.source._core._invalidation_handler.has_pending()) + assert _inventory_keys(pair.source_events, removed=True) == Counter( + (cached, physical) + ) + handles = (results[cached].release_handle, results[physical].release_handle) + assert pair.target.release(handles) == [(handle, True) for handle in handles] + + +@pytest.mark.parametrize("moving", [False, True], ids=["copy", "move"]) +@pytest.mark.parametrize("writer_running", [False, True], ids=["missing", "writing"]) +def test_move_rejects_busy_source_without_blocking_its_producer(moving, writer_running): + pair = _Pair(slots=2, labels=("a", "b"), source_agent=_IndividualWriteAgent()) + ready, busy = BlockKey(b"ready"), BlockKey(b"busy") + pair.deposit(ready, index=1, no_evict=False) + pair.source_agent.state = "DONE" + partial = pair.source.deposit( + {busy: pair.descriptors(labels=("a",))}, expected_layout=["a", "b"] + ) + assert dict(_poll_until(pair.source, bool))[partial][busy].success + pair.source_agent.state = "PROC" + writer = None + if writer_running: + before = len(pair.source_agent.xfers) + writer = pair.source.deposit( + {busy: pair.descriptors(labels=("b",))}, expected_layout=["a", "b"] + ) + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + before = len(pair.source_agent.xfers) + handle = pair.start((ready, busy), move=moving) + core = pair.source._core + if moving: + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + pair.source_agent.states[before + 1] = "DONE" + results = pair.result(handle) + assert results[ready].success and not results[busy].success + assert results[busy].release_handle is None + pair.until(lambda: not core._invalidation_handler.has_pending()) + else: + pair.until(lambda: bool(core._remote_fw_dram._sources._source_pin_ops)) + for _ in range(3): + pair.pump() + assert handle not in pair.results[1] + assert len(pair.source_agent.xfers) == before + assert not core._block_record_map[busy].pending_invalidation + assert busy not in _inventory_keys(pair.source_events, removed=True) + # A rejected move must leave the producer free to finish missing pieces. + followup = pair.source.deposit( + {busy: pair.descriptors(labels=("b",))}, expected_layout=["a", "b"] + ) + pair.source_agent.states.clear() + pair.copy_done() + pair.until(lambda: followup in pair.results[0]) + assert pair.results[0][followup][busy].success + if writer is not None: + assert pair.results[0][writer][busy].success + if not moving: + results = pair.result(handle) + assert all(entry.success for entry in results.values()) + assert pair.source_now == pair.target._core._clock() == 0 + read = pair.source.fetch((busy,), expected_layout=["a", "b"]) + result = dict(pair.source.poll_completed())[read][busy] + assert result.success + for descriptor in result.descriptors: + address = ctypes.addressof(pair.source_local[descriptor.label]) + assert ( + ctypes.string_at(address + descriptor.element_index * 16, 16) + == (pair.source_primary[descriptor.label].raw[:16]) + ) + assert pair.source.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + claims = tuple(entry.release_handle for entry in results.values() if entry.success) + assert pair.target.release(claims) == [(claim, True) for claim in claims] + assert _inventory_keys(pair.source_events, removed=True) == Counter( + (ready,) if moving else () + ) + + +@pytest.mark.parametrize("matching_layout", [False, True]) +@pytest.mark.parametrize("declared_layout", [False, True]) +@pytest.mark.parametrize("no_evict", [False, True]) +@pytest.mark.parametrize("destination_cached", [False, True]) +def test_move_checks_layout_of_capacity_waiting_deposit( + tmp_path, matching_layout, declared_layout, no_evict, destination_cached +): + page = os.sysconf("SC_PAGE_SIZE") + pair = _Pair( + labels=("a", "b"), + target_labels=("a",), + block_size=page, + source_agent=_IndividualG3WriteAgent(), + source_policy=_MoveLocalToG3Policy(), + source_g3=G3Options( + paths=(tmp_path / "source.g3",), + capacity_bytes_per_file=page * 8, + backend="MOCK", + ), + ) + key, occupant = BlockKey(b"waiting-deposit"), BlockKey(b"occupant") + pair.source_primary["a"].raw = b"a" * page + b"A" * page + pair.source_primary["b"].raw = b"b" * page + b"B" * page + pair.target_primary["a"].raw = pair.source_primary["a"].raw + pair.deposit(occupant, no_evict=False) + source_layout = ["a"] if matching_layout else ["a", "b"] + deposited = pair.source.deposit( + {key: pair.descriptors(1, labels=source_layout)}, + expected_layout=source_layout if declared_layout else None, + no_evict=no_evict, + ) + source = pair.source._core + assert source._block_record_map[key].local_dram is None + assert len(source._local_dram._capacity_waiters) == 1 + pair.until(lambda: len(pair.source_agent.xfers) == 2) + + def poll_framework_a(): + for request, (handle, blocks) in pair.source_pinning.poll_pin_results(): + yield ( + request, + ( + handle, + {key: pair.descriptors(1, labels=("a",)) for key in blocks}, + ), + ) + + source._poll_pin_results_callback = poll_framework_a + if destination_cached: + pair.target_agent.state = "DONE" + cached = pair.target.deposit( + {key: pair.descriptors(1, labels=("a",), endpoint="target")}, + expected_layout=["a"], + ) + assert dict(_poll_until(pair.target, bool))[cached][key].success + pair.advise("move") + moved = pair.target.fetch((key,), "req", expected_layout=["a"]) + if destination_cached: + target_result = dict(pair.target.poll_completed())[moved][key] + pair.until(lambda: source._invalidation_handler.has_pending()) + assert not pair.source_pinning.searches + else: + # A capacity-waiting deposit is already admitted; a move must not fence it. + target_result = pair.result(moved)[key] + assert not source._invalidation_handler.has_pending() + assert not pair.source_pinning.searches + assert len(pair.source_agent.xfers) == 2 + assert target_result.success is destination_cached + if destination_cached: + assert target_result.release_handle is not None + assert pair.target_local["a"].raw == b"A" * page + else: + assert target_result.release_handle is None + # The source has accepted the deposit but has no residency to check yet. + assert source._block_record_map[key].local_dram is None + assert deposited not in pair.results[0] + + pair.copy_done() + pair.until(lambda: deposited in pair.results[0]) + source_result = pair.results[0][deposited][key] + assert source_result.success + assert (source_result.release_handle is not None) is no_evict + for label in source_layout: + assert pair.source_local[label].raw == pair.source_primary[label].raw[page:] + pair.until(lambda: not source._invalidation_handler.has_pending()) + preserved = not destination_cached or not matching_layout or no_evict + assert pair.source.query((key,)) == [ + (QueryStatus.HIT, CacheTier.LOCAL_G2) if preserved else (QueryStatus.MISS, None) + ] + if preserved: + assert source._block_record_map[key].local_dram.key_layout == tuple( + source_layout + ) + else: + # An admitted deposit allocated after withdrawal; dropping clears its entry. + assert key not in source._local_dram._pools._source_state.index + if not any((matching_layout, no_evict, declared_layout, destination_cached)): + # Preserving an unprotected copy must leave its slots available to policy. + replacement = BlockKey(b"capacity-reuse") + assert pair.deposit(replacement, no_evict=False).success + for label in source_layout: + assert pair.source_local[label].raw == pair.source_primary[label].raw[:page] + if target_result.release_handle is not None: + assert pair.target.release((target_result.release_handle,)) == [ + (target_result.release_handle, True) + ] + if no_evict: + assert pair.source.release((source_result.release_handle,)) == [ + (source_result.release_handle, True) + ] + pair.until( + lambda: ( + not _has_outstanding_operations(pair.source) + and not _has_outstanding_operations(pair.target) + ) + ) + assert moved not in pair.results[1] + for core in (source, pair.target._core): + pools = core._local_dram._pools + assert not pools._public_claims + assert not pools._source_state.active + assert not core._local_dram_sources_by_op + assert all( + not record.in_flight_ops for record in core._block_record_map.values() + ) + + +def test_waiting_move_preserves_later_invalidation_of_rejected_key( + tmp_path, monkeypatch +): + page = os.sysconf("SC_PAGE_SIZE") + pair = _Pair( + slots=2, + labels=("a", "b"), + target_labels=("a",), + block_size=page, + source_agent=_IndividualG3WriteAgent(), + source_policy=_MoveLocalToG3Policy(), + source_g3=G3Options( + paths=(tmp_path / "source.g3",), + capacity_bytes_per_file=page * 12, + backend="MOCK", + ), + ) + occupant, waiting, mismatched = map( + BlockKey, (b"occupant", b"waiting", b"mismatched") + ) + pair.deposit(occupant, no_evict=False) + pair.source.deposit( + {waiting: pair.descriptors(1, labels=("a",))}, expected_layout=["a"] + ) + pair.until(lambda: len(pair.source_agent.xfers) == 2) + deposited = pair.source.deposit( + {mismatched: pair.descriptors(2)}, expected_layout=["a", "b"] + ) + core = pair.source._core + pair.until(lambda: len(pair.source_agent.xfers) == 3) + assert core._block_record_map[mismatched].local_dram is None + + def poll_framework_full_layout(): + for request, (pin, blocks) in pair.source_pinning.poll_pin_results(): + yield request, (pin, {key: pair.descriptors(2) for key in blocks}) + + core._poll_pin_results_callback = poll_framework_full_layout + sources = core._remote_fw_dram._sources + + def admit_source(op): + # Exercise retry ownership after admission; public move now rejects + # this fixture's busy writer before reaching source selection. + core._add_block_dependencies(op, new_operation=True) + op.moving = core._invalidation_handler.submit(op.keys, key_layout=op.key_layout) + sources._source_pin_ops[op.op_id] = op + sources._resume_source_pin(op.op_id, op) + + monkeypatch.setattr(sources, "start_source_pin", admit_source) + pair.advise("move") + moved = pair.target.fetch((mismatched, waiting), "req", expected_layout=["a"]) + pair.until( + lambda: ( + bool(sources._source_pin_ops) + and all(op.keys == {waiting} for op in sources._source_pin_ops.values()) + ) + ) + # The move rejected this layout while waiting for the other key's writer. + # Finish its admitted deposit, then independently invalidate the new copy. + pair.source_agent.states[3] = "DONE" + pair.until(lambda: len(pair.source_agent.xfers) == 4) + pair.source_agent.states[4] = "DONE" + pair.until(lambda: deposited in pair.results[0]) + assert pair.results[0][deposited][mismatched].success + pair.until(lambda: not core._block_record_map[mismatched].pending_invalidation) + read = pair.source.fetch((mismatched,), expected_layout=["a", "b"]) + guard = dict(pair.source.poll_completed())[read][mismatched].release_handle + assert guard is not None + pair.source_control.incoming.append( + msgspec.msgpack.encode(_Invalidate(keys=(mismatched,), key_layout=("a", "b"))) + ) + pair.until(lambda: core._block_record_map[mismatched].pending_invalidation) + + pair.source_agent.states[2] = "DONE" + pair.until(lambda: not sources._source_pin_ops) + preserved = core._block_record_map[mismatched].pending_invalidation + assert pair.source.release((guard,)) == [(guard, True)] + pair.source_agent.states.clear() + pair.copy_done() + result = pair.result(moved) + assert not result[mismatched].success and result[waiting].success + claim = result[waiting].release_handle + assert pair.target.release((claim,)) == [(claim, True)] + assert preserved, "Old move cancelled a newer invalidation" + pair.until(lambda: mismatched not in core._block_record_map) + + +def test_move_source_close_after_immediate_native_completion(): + pair = _Pair() + key = BlockKey(b"immediate-move") + assert pair.deposit(key, no_evict=False).success + core = pair.source._core + pools = core._local_dram._pools + residency = core._block_record_map[key].local_dram + pair.source_agent.state = "DONE" + before = len(pair.source_agent.xfers) + handle = pair.start((key,)) + # The owner admits the move and the write completes natively at once; + # close() must settle it without another application poll. + pair.until(lambda: len(pair.source_agent.xfers) == before + 1) + assert residency.claim_count == 1 + assert core._invalidation_handler.has_pending() + pair.source.close() + assert not pools._source_state.active + assert not core._local_dram_sources_by_op + assert not core._invalidation_handler.has_pending() + pair.until(lambda: handle in pair.results[1], source_poll=False) + result = pair.results[1].pop(handle)[key] + assert result.success and pair.target_local[""].raw == b"a" * 16 + assert pair.target.release((result.release_handle,)) == [ + (result.release_handle, True) + ] diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index 93b73bd7..a6c62ba5 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -21,6 +21,7 @@ _decode_control_message, _decode_notif, _has_outstanding_operations, + _invalidate, _mem_descriptor, _new_kvcr, _op_entries, @@ -161,6 +162,67 @@ def test_framework_pinning_disabled(keys, completed_indices, source_case, kvcr_c ) +def test_invalidation_without_framework_pinning_preserves_admitted_readers(): + memory = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16, 32) + agent, control = FakeNixlAgent(), FakeBytesControl() + source = _new_kvcr( + agent, + SimpleNamespace(request_pin=None, poll_pin_results=None, release_pin=None), + control, + name="source", + framework_regions=[_buffer_region(memory, 16)], + local_dram=LocalDramOptions([("", ctypes.addressof(memory), 32)]), + ) + key, other = BlockKey(b"removing"), BlockKey(b"unrelated") + pools = source._core._local_dram._pools + try: + agent.state = "DONE" + deposit = source.deposit( + { + key: [_mem_descriptor(end_point_name="source")], + other: [_mem_descriptor(1, end_point_name="source")], + } + ) + assert dict(_poll_until(source, bool))[deposit] == _op_entries( + {key: True, other: True} + ) + prior = len(agent.xfers) + agent.state = "PROC" + control.incoming.append(_start_write_message(12, key)) + _wait_until(lambda: len(agent.transfers) == prior + 1) + residency = source._core._block_record_map[key].local_dram + assert residency.source_readers == 1 + _invalidate(source, (key,)) + + request = msgspec.msgpack.decode(_start_write_message(13, key)) + request.update( + keys=[key, other], + dst_descriptors=[[_TransferRef("target", i)] for i in range(2)], + ) + # A pending removal is a miss; unrelated sources still run on progress. + with source._core._state_lock: + control.incoming.append(msgspec.msgpack.encode(request)) + _wait_until(lambda: len(agent.transfers) == prior + 2) + assert residency.source_readers == 1 + assert pools.telemetry_state()["local_g2_free_slots"] == 0 + assert _decode_notif(agent.xfers[-1][5])["completed_indices"] == [1] + assert agent.xfers[-1][3] == [(144, 16, 0)] + + agent.state = "DONE" + _poll_until(source, lambda _: not _has_outstanding_operations(source)) + admitted = _decode_notif(agent.xfers[prior][5]) + assert admitted["success"] and admitted["completed_indices"] == [0] + assert not agent.sent_notifs # No cancellation/failure replaced native success. + assert residency.source_readers == 0 + assert key not in pools._source_state.index + assert key not in source._core._block_record_map + assert other in pools._source_state.index + assert pools.telemetry_state()["local_g2_free_slots"] == 1 + finally: + agent.state = "DONE" + source.close() + + @pytest.mark.parametrize("transfer_state,recovered", [("PROC", False), ("DONE", True)]) def test_local_source_starts_inline_and_holds_its_slot(transfer_state, recovered): memory = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16 + b"c" * 16, 48)