diff --git a/src/kvcr/backends/dangling_ops.py b/src/kvcr/backends/dangling_ops.py index e06c0007..5a424fc6 100644 --- a/src/kvcr/backends/dangling_ops.py +++ b/src/kvcr/backends/dangling_ops.py @@ -98,6 +98,10 @@ def poll_source( def finish_source(self, op: "_SourceWriteOp") -> None: status = self.source_writes.pop((op.route[0], op.op_handle), None) + if status is not None and op.local_sources: + self._backend._kvcr._local_dram._pools.release_readers( + op.source_keys, op.local_sources, success=op.success + ) if status is not None and status.abandoned: self.report_source(op, "quiesced") diff --git a/src/kvcr/backends/local_dram.py b/src/kvcr/backends/local_dram.py index 76f2fcb8..2b330209 100644 --- a/src/kvcr/backends/local_dram.py +++ b/src/kvcr/backends/local_dram.py @@ -517,6 +517,7 @@ def deliver( self._start_deliveries(op, blocks) def discard_fill(self, keys: Collection[BlockKey]) -> None: + state = self._pools._source_state residency_ops: dict[_OpId, _PendingResidencyOp] = {} deliver_ops: dict[_OpId, _PendingDeliverOp] = {} for key in dict.fromkeys(keys): @@ -536,7 +537,9 @@ def discard_fill(self, keys: Collection[BlockKey]) -> None: residency.state is _LocalDramState.FILLING and not residency.pending_deposits ): - residency.state = _LocalDramState.DISCARDING + with state.lock: + state.ready.pop(key, None) + residency.state = _LocalDramState.DISCARDING if residency.remote_indices is not None: residency.remote_cancelled = True for op_id in record.active_op_ids: @@ -571,6 +574,7 @@ def discard_fill(self, keys: Collection[BlockKey]) -> None: self._finish_deliver_if_ready(deliver_op) def poll_main(self, items: Collection[object]) -> list[object]: + self._pools._drain_source_completions() # Progress-side source claims defer framework callbacks to the caller. self._pools._update_capacity_pressure() unhandled: list[object] = [] @@ -614,6 +618,9 @@ def _finish_copy(self, copy: _LocalCopyOp) -> None: copy.success, CacheTier.FW_G2, writer=copy.deposit_op, + dst_descriptors=( + copy.dst_descriptors if self._kvcr.framework_control is not None else () + ), ) def _apply_fill_result( @@ -624,6 +631,7 @@ def _apply_fill_result( source: CacheTier, *, writer: _PendingResidencyOp | None = None, + dst_descriptors: tuple[_TransferRef, ...] = (), ) -> None: committed: list[BlockKey] = [] failed: list[BlockKey] = [] @@ -632,7 +640,10 @@ def _apply_fill_result( affected_deliver_ops: dict[_OpId, _PendingDeliverOp] = {} deliver_keys: dict[_OpId, list[BlockKey]] = {} now = self._kvcr._clock() + offset = 0 for key, indices in zip(ordered_keys, local_indices): + start = offset + offset += len(indices) record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None if record is None or residency is None: @@ -667,6 +678,14 @@ def _apply_fill_result( if len(pending) < pending_count: residency.state = _LocalDramState.READY residency.prepare_label_index() + refs = None + if ( + dst_descriptors + and len(indices) == len(residency.slots) + and indices == tuple(range(len(indices))) + ): + refs = dst_descriptors[start:offset] + self._pools._publish_source(key, residency, refs) record.last_access = now if not was_ready: self._kvcr._on_ingest( @@ -949,6 +968,7 @@ def _resume_capacity_waiters(self) -> None: return self._resuming_capacity_waiters = True try: + self._pools._drain_source_completions() # One pass lets other waiters run while discarded fills retain slots. for _ in range(len(self._capacity_waiters)): waiter = self._capacity_waiters[0] diff --git a/src/kvcr/backends/local_dram_pools.py b/src/kvcr/backends/local_dram_pools.py index 2be06e55..e102e56c 100644 --- a/src/kvcr/backends/local_dram_pools.py +++ b/src/kvcr/backends/local_dram_pools.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Local DRAM pool allocation, residency, claims, and eviction.""" +import queue +import threading from collections import Counter, deque from collections.abc import Callable, Collection, Mapping from contextlib import closing @@ -64,7 +66,8 @@ class _LocalDramResidency(msgspec.Struct, gc=False, eq=False): Reserved pieces become readable only after successful copies. Once all pieces are readable, pending_labels is None; writers is also None when - no copies are in flight. refs caches committed physical-slot references. + no copies are in flight. refs is sparse until a completed residency is + published for remote reads, then holds its immutable descriptor tuple. """ slots: tuple[tuple[str, int], ...] # Full descriptor label and parent-pool slot. @@ -77,11 +80,14 @@ class _LocalDramResidency(msgspec.Struct, gc=False, eq=False): writers: list[int] | None = None remote_indices: tuple[int, ...] | None = None remote_cancelled: bool = False - refs: dict[int, _TransferRef] = msgspec.field(default_factory=dict) + refs: dict[int, _TransferRef] | tuple[_TransferRef, ...] = msgspec.field( + default_factory=dict + ) indices_by_label: dict[str, int | None] | None = None _layout: tuple[str, ...] | None = None _pool_counts: tuple[tuple[str, int], ...] | None = None _cached_size: int | None = None + source_readers: int = 0 # Includes completed reads awaiting access accounting. def __post_init__(self) -> None: self.slots = tuple(self.slots) @@ -163,6 +169,20 @@ def can_wait_for_pieces(self, requested: list[str]) -> bool: ) +class _SourceState: + """Reader state shared with progress; completions return work to the caller.""" + + def __init__(self) -> None: + # Owner paths take _state_lock before this lock; progress takes only this. + self.lock = threading.Lock() + self.ready: dict[BlockKey, _LocalDramResidency] = {} + self.active: dict[BlockKey, _LocalDramResidency] = {} + self.operations = 0 + self.completions: queue.SimpleQueue[ + tuple[tuple[BlockKey, ...], tuple[_LocalDramResidency, ...], bool] + ] = queue.SimpleQueue() + + class _LocalDramPools: """Own pool slots and claims independently of pending operations.""" @@ -181,6 +201,8 @@ def __init__( raise ValueError("local DRAM NIXL backend must be non-empty") self._kvcr = kvcr + self._source_state = _SourceState() + self._draining_source_completions = False self._regions: dict[str, tuple[int, int, int]] = {} self._free_slots: dict[str, deque[int]] = {} for (pool_name, address, length), (_, slot_size) in zip( @@ -226,6 +248,139 @@ def memory_regions(self) -> tuple[RegionDescriptor, ...]: for name, (address, length, size) in self._regions.items() ) + @property + def source_operations(self) -> int: + state = self._source_state + with state.lock: + return state.operations + + def _publish_source( + self, + key: BlockKey, + residency: _LocalDramResidency, + refs: tuple[_TransferRef, ...] | None = None, + ) -> None: + if self._kvcr.framework_control is None or residency.pending_labels: + return + if refs is None: + refs = ( + self.source_refs(residency, range(len(residency.slots))) + if residency.refs + else tuple(self._descriptors(residency.slots)) + ) + state = self._source_state + with state.lock: + if ( + residency.state is _LocalDramState.READY + and not residency.retire_on_release + ): + residency.refs = refs + state.ready[key] = residency + + def acquire_readers( + self, + keys: tuple[BlockKey, ...], + destinations: tuple[tuple[_TransferRef, ...], ...], + key_layout: tuple[str, ...], + ) -> tuple[tuple[_LocalDramResidency, tuple[_TransferRef, ...]], ...] | None: + """Claim every requested piece; None retries contention, () uses the owner.""" + state = self._source_state + if not state.lock.acquire(blocking=False): + return None + try: + claims = [] + for key, destination in zip(keys, destinations): + residency = state.ready.get(key) + if residency is None: + return () + available = cast(tuple[_TransferRef, ...], residency.refs) + if key_layout and key_layout != residency.key_layout: + return () + requested = [ref.label for ref in destination] + if residency.layout == tuple(requested): + refs = available + else: + indices = _layout_indices( + residency.layout, + requested, + indices_by_label=residency.indices_by_label, + ) + if indices is None: + return () + refs = tuple(available[index] for index in indices) + claims.append((residency, refs)) + result = tuple(claims) + for key, (residency, _) in zip(keys, result): + residency.source_readers += 1 + state.active[key] = residency + state.operations += 1 + return result + finally: + state.lock.release() + + def release_readers( + self, + keys: tuple[BlockKey, ...], + sources: tuple[_LocalDramResidency, ...], + *, + success: bool, + ) -> None: + """Queue reader release; retain eviction protection until access accounting.""" + state = self._source_state + with state.lock: + state.operations -= 1 + state.completions.put((keys, sources, success)) + + def _drain_source_completions(self) -> None: + if self._draining_source_completions: + return + self._draining_source_completions = True + state = self._source_state + completions = state.completions + freed = False + try: + while True: + try: + keys, sources, success = completions.get_nowait() + except queue.Empty: + break + if success: + self._kvcr._record_access(keys) + with state.lock: + for key, residency in zip(keys, sources): + if residency.source_readers <= 0: + raise RuntimeError( + "local DRAM source reader ownership lost" + ) + residency.source_readers -= 1 + if not residency.source_readers: + 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.local_dram is residency + and not residency.claim_count + ): + freed |= self._after_last_claim(key, residency) + finally: + self._draining_source_completions = False + if freed and not self._kvcr._closed: + self._resume_capacity_waiters() + + def _reclaimable_slots(self) -> dict[str, int]: + slots = { + name: len(free) + self._evictable_slots[name] + for name, free in self._free_slots.items() + } + state = self._source_state + with state.lock: + for key, residency in state.active.items(): + if self._evictable.score(key) is not None: + for pool, count in residency.pool_counts: + slots[pool] -= count + return slots + def observe_residency( self, observer: Callable[[BlockKey, "_BlockRecord"], None] ) -> None: @@ -239,7 +394,7 @@ def adopt_recovery_slots(self, records: Mapping[BlockKey, "_BlockRecord"]) -> No seen every block first. """ occupied = {pool_name: set() for pool_name in self._regions} - for record in records.values(): + for key, record in records.items(): residency = record.local_dram if residency is None: continue @@ -270,7 +425,9 @@ def adopt_recovery_slots(self, records: Mapping[BlockKey, "_BlockRecord"]) -> No raise ValueError("invalid local DRAM recovery layout") # Byte counts depend on the receiving core's pool geometry. residency._cached_size = None + residency.refs = {} residency.prepare_label_index() + self._publish_source(key, residency) self._free_slots = { pool_name: deque( slot @@ -301,7 +458,8 @@ def telemetry_state(self) -> dict[str, int]: "local_g2_total_slots": total_slots, "local_g2_free_slots": free_slots, "local_g2_allocated_slots": total_slots - free_slots, - "local_g2_evictable_slots": sum(self._evictable_slots.values()), + "local_g2_evictable_slots": sum(self._reclaimable_slots().values()) + - free_slots, } def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: @@ -321,7 +479,6 @@ def acquire_sources( self, keys: Collection[BlockKey], *, - notify_capacity: bool = True, allow_filling: bool = False, ) -> dict[BlockKey, _LocalDramResidency]: sources: dict[BlockKey, _LocalDramResidency] = {} @@ -337,18 +494,20 @@ def acquire_sources( continue self._acquire_claim(key, residency) sources[key] = residency - if notify_capacity: - self._update_capacity_pressure() + self._update_capacity_pressure() return sources def source_refs( - self, residency: _LocalDramResidency, indices: list[int] + self, residency: _LocalDramResidency, indices: Collection[int] ) -> tuple[_TransferRef, ...]: + cached = residency.refs + if isinstance(cached, tuple): + return tuple(cached[index] for index in indices) refs = [] for index in indices: - ref = residency.refs.get(index) + ref = cached.get(index) if ref is None: - ref = residency.refs[index] = self._descriptor(*residency.slots[index]) + ref = cached[index] = self._descriptor(*residency.slots[index]) refs.append(ref) return tuple(refs) @@ -363,6 +522,7 @@ def release_sources(self, keys: Collection[BlockKey]) -> None: def retire_sources(self, keys: Collection[BlockKey]) -> None: """Retire claimed sources when their final internal claim is released.""" + state = self._source_state for key in dict.fromkeys(keys): record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None @@ -372,7 +532,9 @@ def retire_sources(self, keys: Collection[BlockKey]) -> None: or residency.claim_count <= 0 ): raise RuntimeError(f"local DRAM source cannot retire {key!r}") - residency.retire_on_release = True + with state.lock: + residency.retire_on_release = True + state.ready.pop(key, None) def abandon_capacity_eviction(self, key: BlockKey) -> None: """Stop blocking local admission on an eviction that will not land. @@ -385,7 +547,8 @@ def abandon_capacity_eviction(self, key: BlockKey) -> None: self._capacity_eviction_key = None def close(self) -> None: - # Policy state ends with KVCR; teardown emits no per-block removals. + # Drain pending source cleanup before dropping public claims. + self._drain_source_completions() self._public_claims.clear() def _new_public_claim( @@ -430,32 +593,44 @@ def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: raise RuntimeError(f"invalid local DRAM claim for {key!r}") residency.claim_count -= 1 if residency.claim_count == 0: - if residency.state is not _LocalDramState.READY: - record.local_dram = None - self._free(residency.slots) - self._kvcr._prune_block_record(key) - self._resume_capacity_waiters() - return - if residency.retire_on_release: - if self._kvcr._key_history_enabled: - self._kvcr._log_key_history( - "remove", (key,), CacheTier.LOCAL_G2, "retire_on_release" - ) + self._after_last_claim(key, residency) + + def _after_last_claim(self, key: BlockKey, residency: _LocalDramResidency) -> bool: + state = self._source_state + record = self._kvcr._block_record_map[key] + retiring = ( + residency.state is not _LocalDramState.READY or residency.retire_on_release + ) + with state.lock: + if retiring: + state.ready.pop(key, None) + if residency.source_readers: + return False + if retiring: + # Withdraw the source before its slots can be reused. record.local_dram = None - self._residency_observer(key, record) - self._free(residency.slots) - self.abandon_capacity_eviction(key) - self._kvcr._on_remove( - self._kvcr._block_meta(key, record, self._size_bytes(residency)) + if not retiring: + self._make_evictable(key) + return False + if residency.state is _LocalDramState.READY: + if self._kvcr._key_history_enabled: + self._kvcr._log_key_history( + "remove", (key,), CacheTier.LOCAL_G2, "retire_on_release" ) - if not residency.pending_labels: - self._kvcr._publish_inventory( - (key,), CacheTier.LOCAL_G2, removed=True - ) - self._kvcr._prune_block_record(key) - self._resume_capacity_waiters() - else: - self._make_evictable(key) + self._residency_observer(key, record) + self._free(residency.slots) + self.abandon_capacity_eviction(key) + self._kvcr._on_remove( + self._kvcr._block_meta(key, record, self._size_bytes(residency)) + ) + if not residency.pending_labels: + self._kvcr._publish_inventory((key,), CacheTier.LOCAL_G2, removed=True) + else: + self._free(residency.slots) + self._kvcr._prune_block_record(key) + if not self._kvcr._closed and not self._draining_source_completions: + self._resume_capacity_waiters() + return True def _allocate_slots( self, layout: list[str], protected: frozenset[BlockKey], deadline: float @@ -486,6 +661,7 @@ def short() -> set[str]: } deficient = short() + state = self._source_state with closing(self._evictable.candidates(protected)) as candidates: while deficient: key = next(candidates, None) @@ -502,6 +678,9 @@ def short() -> set[str]: raise RuntimeError(f"invalid evictable local DRAM entry {key!r}") if deficient.isdisjoint(pool for pool, _ in residency.pool_counts): continue + with state.lock: + if residency.source_readers: + continue size_bytes = self._size_bytes(residency) decision, eviction_pending = self._kvcr._decide_eviction( self._kvcr._block_meta(key, record, size_bytes), @@ -521,6 +700,13 @@ def short() -> set[str]: freed[pool] += count deficient = short() + with state.lock: + # Policy runs unlocked; a source can claim a selected victim meanwhile. + if any(residency.source_readers for _, _, residency, _ in victims): + # Accept transient admission failure to avoid another policy scan. + return None, [], False + for key, _, _, _ in victims: + state.ready.pop(key, None) for key, record, residency, size_bytes in victims: self._remove_evictable(key, residency) if self._kvcr._key_history_enabled: @@ -550,6 +736,7 @@ def _make_evictable(self, key: BlockKey) -> None: residency is None or residency.state is not _LocalDramState.READY or residency.claim_count + or residency.retire_on_release or residency.pending_labels ): return @@ -605,9 +792,4 @@ def _size_bytes(self, residency: _LocalDramResidency) -> int: def _update_capacity_pressure(self) -> None: if self._kvcr._capacity_needed_callback is None: return - self._kvcr._update_capacity_pressure( - { - name: len(slots) + self._evictable_slots[name] - for name, slots in self._free_slots.items() - } - ) + self._kvcr._update_capacity_pressure(self._reclaimable_slots()) diff --git a/src/kvcr/backends/remote_fw_dram.py b/src/kvcr/backends/remote_fw_dram.py index 27fb2f4f..550c193b 100644 --- a/src/kvcr/backends/remote_fw_dram.py +++ b/src/kvcr/backends/remote_fw_dram.py @@ -43,7 +43,6 @@ OpHandle, ) from .dangling_ops import _DanglingOps, _SourceWriteStatus -from .local_dram_pools import _LocalDramState from .remote_ops import ( _CONTROL_DECODER, _ControlMessage, @@ -866,62 +865,45 @@ def _try_local_source_write( self, progress: _KVCRProgress, source_pin: _SourcePinOp ) -> bool | None: """Return None to retry contention, False when caller work is needed.""" - kvcr = self._kvcr - if kvcr._local_dram is None: + local = self._kvcr._local_dram + if local is None: return False - if not kvcr._state_lock.acquire(blocking=False): + selected = local._pools.acquire_readers( + source_pin.ordered_keys, source_pin.dst_descriptors, source_pin.key_layout + ) + if selected is None: return None - selected = [] + if not selected: + return False + sources = tuple(residency for residency, _ in selected) try: - for key, destination in zip( - source_pin.ordered_keys, source_pin.dst_descriptors - ): - record = 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: - return False - if ( - source_pin.key_layout - and source_pin.key_layout != residency.key_layout - ): - return False - indices = residency.ready_indices([ref.label for ref in destination]) - if indices is None: - return False - selected.append(indices) - sources = kvcr._claim_local_dram_sources( - source_pin.op_id, source_pin.ordered_keys, notify_capacity=False + source_write = _SourceWriteOp( + state=_SourceWriteState.READY_TO_WRITE, + op_id=source_pin.op_id, + keys=source_pin.keys, + started_at=source_pin.started_at, + deadline=source_pin.deadline, + remote_agent=source_pin.remote_agent, + op_handle=source_pin.op_handle, + source_keys=source_pin.ordered_keys, + src_descriptors=tuple(refs for _, refs in selected), + dst_descriptors=source_pin.dst_descriptors, + completed_indices=tuple(range(len(source_pin.ordered_keys))), + route=source_pin.route, + nvtx_state=source_pin.nvtx_state, + local_sources=sources, + _backend=self, ) - source_refs = tuple( - kvcr._local_dram._pools.source_refs(sources[key], indices) - for key, indices in zip(source_pin.ordered_keys, selected) + # Retain cleanup ownership before native submission can start. + progress._in_flight_ops[source_write.op_id] = source_write + except BaseException: + local._pools.release_readers( + source_pin.ordered_keys, sources, success=False ) - kvcr._add_block_dependencies(source_pin, new_operation=True) - finally: - kvcr._state_lock.release() - # Claims keep the sources stable while the write is assembled unlocked. - source_write = _SourceWriteOp( - state=_SourceWriteState.READY_TO_WRITE, - op_id=source_pin.op_id, - keys=source_pin.keys, - started_at=source_pin.started_at, - deadline=source_pin.deadline, - remote_agent=source_pin.remote_agent, - op_handle=source_pin.op_handle, - source_keys=source_pin.ordered_keys, - src_descriptors=source_refs, - dst_descriptors=source_pin.dst_descriptors, - completed_indices=tuple(range(len(source_pin.ordered_keys))), - route=source_pin.route, - nvtx_state=source_pin.nvtx_state, - _backend=self, - ) - # Already on progress: retain cleanup ownership before starting the write. - progress._in_flight_ops[source_write.op_id] = source_write + raise done, _ = source_write.progress(progress, None) if done: progress._in_flight_ops.pop(source_write.op_id) - self._progress_outbound.append(source_write) return True def _handle_malformed_start_write( diff --git a/src/kvcr/backends/remote_ops.py b/src/kvcr/backends/remote_ops.py index de1ba0c9..96d37007 100644 --- a/src/kvcr/backends/remote_ops.py +++ b/src/kvcr/backends/remote_ops.py @@ -17,6 +17,7 @@ from ..types import BlockKey, OpHandle, PinHandle if TYPE_CHECKING: + from .local_dram_pools import _LocalDramResidency from .remote_fw_dram import _RemoteFWDram @@ -373,6 +374,7 @@ class _SourceWriteOp(_RemoteOp): completed_indices: tuple[int, ...] = () route: tuple[str, int] = ("", 0) nvtx_state: _NvtxState | None = None + local_sources: tuple["_LocalDramResidency", ...] = () def progress( self, progress: _KVCRProgress, _event: object | None diff --git a/src/kvcr/backends/remote_sources.py b/src/kvcr/backends/remote_sources.py index 99e5d5e7..99a92e8f 100644 --- a/src/kvcr/backends/remote_sources.py +++ b/src/kvcr/backends/remote_sources.py @@ -27,7 +27,7 @@ class _FwMemResidency: @dataclass class _SourcePinOp(_Op): - """Main-thread source acquisition for a pending write.""" + """Source acquisition for a pending write.""" started_at: float | None deadline: float @@ -84,6 +84,11 @@ def finish_source_write(self, item: _SourceWriteOp) -> None: _SourceWriteState.FINISHED, _SourceWriteState.CANCEL_PENDING, ): + if item.local_sources: + # local.poll_main drains reader releases with their access + # accounting and reclamation. + # TODO: Skip queued local handoffs if useful; keep TransferError events. + return # Framework ownership can be handed back after uncertainty; # KVCR-owned contents stay claimed until native quiescence. pins = self._fw_pins_by_op.pop(item.op_id, None) diff --git a/src/kvcr/runtime/core.py b/src/kvcr/runtime/core.py index 0cfcfd9b..715c1b4f 100644 --- a/src/kvcr/runtime/core.py +++ b/src/kvcr/runtime/core.py @@ -670,6 +670,7 @@ def deposit( dict.fromkeys(blocks, _FAILED_RESULT), ) else: + self._local_dram._pools._drain_source_completions() if self._g3 is not None: for key, refs in blocks.items(): if self._g3.is_ready(key): @@ -716,6 +717,7 @@ def fetch( ) return op_handle + local_dram._pools._drain_source_completions() sources = {} for key in ordered_keys: record = self._block_record_map.get(key) @@ -821,6 +823,7 @@ def get_stats(self) -> TelemetryStats | None: "completed": len(self._completion_queue), } if self._local_dram is not None: + resources["in_flight_ops"] += self._local_dram._pools.source_operations resources.update(self._local_dram._pools.telemetry_state()) if self._g3 is not None: resources.update(self._g3.telemetry_state()) @@ -1072,14 +1075,12 @@ def _claim_local_dram_sources( op_id: _OpId, keys: Collection[BlockKey], *, - notify_capacity: bool = True, allow_filling: bool = False, ) -> Mapping[BlockKey, _LocalDramResidency]: sources = self._local_dram_sources_by_op.get(op_id, {}) if self._local_dram is not None: claimed = self._local_dram._pools.acquire_sources( tuple(key for key in keys if key not in sources), - notify_capacity=notify_capacity, allow_filling=allow_filling, ) if claimed: diff --git a/src/kvcr/service/recovery_journal.py b/src/kvcr/service/recovery_journal.py index d62ab201..0274061a 100644 --- a/src/kvcr/service/recovery_journal.py +++ b/src/kvcr/service/recovery_journal.py @@ -392,9 +392,10 @@ def adopt(self, records: dict[BlockKey, _BlockRecord]) -> None: local_dram.remote_indices = None local_dram.remote_cancelled = False local_dram.claim_count = 0 + local_dram.source_readers = 0 local_dram.retire_on_release = False # References name the previous agent, unlike the slot layout. - local_dram.refs.clear() + local_dram.refs = {} if record.g3 is not None: record.g3.claim_count = 0 record.fw_mem = None diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 29530333..0d00e6ec 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -172,7 +172,15 @@ def _poll_until( def _has_outstanding_operations(kvcr: KVCR) -> bool: - return kvcr._core._outstanding_operations > 0 + core = kvcr._core + local = core._local_dram + return core._outstanding_operations > 0 or ( + local is not None + and ( + local._pools.source_operations > 0 + or not local._pools._source_state.completions.empty() + ) + ) def _wait_until(predicate, *, timeout: float = 1) -> None: diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 9f072aae..06dbae91 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -76,7 +76,7 @@ def _two_pool_kvcr( return _new_kvcr( agent, FakePrimaryPinning(), - FakeBytesControl(), + None, # Exercise lazy descriptor caching without remote source publication. config, local_dram=LocalDramOptions( [ @@ -186,6 +186,44 @@ def forbidden(*args, **kwargs): assert {r["event"] for r in records} >= {"ready", "inventory_accepted"} +def test_completed_deposit_reuses_destination_descriptors(monkeypatch): + payload = b"a" * 8 + b"b" * 8 + b"c" * 8 + b"d" * 8 + b"e" * 8 + b"f" * 8 + primary = ctypes.create_string_buffer(payload, len(payload)) + local = ctypes.create_string_buffer(len(payload)) + agent = FakeNixlAgent() + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig(nixl_agent_name="target", pool_layouts=[("", 8)]), + local_dram=LocalDramOptions([("", ctypes.addressof(local), len(local))]), + framework_regions=[_buffer_region(primary, 8, ":*")], + ) + layout = [":a", ":b", ":c"] + keys = (BlockKey(b"reordered"), BlockKey(b"ordered")) + expected = _op_entries(dict.fromkeys(keys, True)) + blocks = { + key: [_mem_descriptor(3 * k + i, label) for i, label in enumerate(layout)] + for k, key in enumerate(keys) + } + blocks[keys[0]].reverse() + submit = Mock(wraps=kvcr._core._progress.submit) + monkeypatch.setattr(kvcr._core._progress, "submit", submit) + deposit = kvcr.deposit(blocks, expected_layout=layout) + copy = submit.call_args.args[0] + agent.state = "DONE" + assert dict(_poll_until(kvcr, bool))[deposit] == expected + residency = kvcr._core._block_record_map[keys[1]].local_dram + assert all( + ref is dst + for ref, dst in zip(residency.refs, copy.dst_descriptors[3:], strict=True) + ) + primary.raw = b"\0" * len(payload) + delivery = kvcr.deliver(blocks) + assert dict(_poll_until(kvcr, bool))[delivery] == expected + assert primary.raw == payload + + def test_parallel_deposits_preserve_committed_pieces_and_claims(monkeypatch) -> None: pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(24)] source = ctypes.create_string_buffer( @@ -763,6 +801,119 @@ def test_builtin_policy_eviction_order( kvcr.release((second_claim,)) +def test_eviction_preserves_reader_claimed_during_policy_decision(monkeypatch): + primary = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16, 32) + memory = ctypes.create_string_buffer(16) + agent = FakeNixlAgent() + agent.state = "DONE" + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + framework_regions=[_buffer_region(primary, 16)], + local_dram=LocalDramOptions([("", ctypes.addressof(memory), len(memory))]), + ) + local = kvcr._core._local_dram._pools + resident, incoming = BlockKey(b"resident"), BlockKey(b"incoming") + deposit = kvcr.deposit({resident: [_mem_descriptor()]}) + assert _poll_until(kvcr, bool) == [(deposit, _op_entries({resident: True}))] + claimed = [] + + def claim_selected_victim(meta, source): + destination = ((progress_module._TransferRef("target", 0),),) + claimed.extend(local.acquire_readers((meta.block_key,), destination, ())) + return PlacementAction.DROP, None + + with monkeypatch.context() as patch: + patch.setattr(kvcr._core._policy, "decide_eviction", claim_selected_victim) + rejected = kvcr.deposit({incoming: [_mem_descriptor(1)]}) + try: + assert _poll_until(kvcr, bool) == [(rejected, _op_entries({incoming: False}))] + assert len(agent.transfers) == 1 + assert claimed[0][0].source_readers == 1 + assert memory.raw == b"a" * 16 + assert kvcr.query((resident,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + finally: + local.release_readers((resident,), (claimed[0][0],), success=True) + + replacement = kvcr.deposit({incoming: [_mem_descriptor(1)]}) + assert _poll_until(kvcr, bool) == [(replacement, _op_entries({incoming: True}))] + assert memory.raw == b"b" * 16 + + +@pytest.mark.parametrize( + "completion_order", ["reader_first", "spill_first", "during_admission"] +) +def test_source_completions_refresh_lru_before_resuming_capacity_waiters( + completion_order, + monkeypatch, +): + primary = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16 + b"c" * 16, 48) + memory = ctypes.create_string_buffer(48) + agent = FakeNixlAgent() + agent.state = "DONE" + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + framework_regions=[_buffer_region(primary, 16)], + local_dram=LocalDramOptions([("", ctypes.addressof(memory), len(memory))]), + policy=LRUPolicy(), + ) + local = kvcr._core._local_dram._pools + retiring, recent, cold, incoming = map( + BlockKey, (b"retiring", b"recent", b"cold", b"incoming") + ) + now = 0.0 + kvcr._core._clock = lambda: now + for index, key in enumerate((retiring, recent, cold)): + now = index / 10 + deposit = kvcr.deposit({key: [_mem_descriptor(index)]}) + assert _poll_until(kvcr, bool) == [(deposit, _op_entries({key: True}))] + + destination = ((progress_module._TransferRef("target", 0),),) + retiring_source = local.acquire_readers((retiring,), destination, ()) + recent_source = local.acquire_readers((recent,), destination, ()) + assert retiring_source and recent_source + # A pending spill holds the retirement and queues a two-slot deposit. + local.acquire_sources((retiring,)) + local._capacity_eviction_key = retiring + blocks = {incoming: [_mem_descriptor(0), _mem_descriptor(1)]} + deposit = kvcr.deposit(blocks) + assert list(kvcr.poll_completed()) == [] + local.retire_sources((retiring,)) + if completion_order == "spill_first": + local.release_sources((retiring,)) + kvcr.align_sequence([retiring]) + + now = 0.3 + local.release_readers((retiring,), (retiring_source[0][0],), success=False) + if completion_order == "during_admission": + decide_ingest = kvcr._core._policy.decide_ingest + + def complete_reader(*args, **kwargs): + # Finish after the admission's initial completion drain. + monkeypatch.setattr(kvcr._core._policy, "decide_ingest", decide_ingest) + local.release_readers((recent,), (recent_source[0][0],), success=True) + return decide_ingest(*args, **kwargs) + + monkeypatch.setattr(kvcr._core._policy, "decide_ingest", complete_reader) + else: + local.release_readers((recent,), (recent_source[0][0],), success=True) + if completion_order != "spill_first": + local.release_sources((retiring,)) + duplicate = kvcr.deposit(blocks) + completed = dict(_poll_until(kvcr, lambda results: len(results) == 2)) + assert completed == dict.fromkeys( + (deposit, duplicate), _op_entries({incoming: True}) + ) + assert kvcr.query((recent, cold, incoming)) == [ + (QueryStatus.HIT, CacheTier.LOCAL_G2), + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + + @pytest.mark.parametrize( ("access_first", "use_current_time", "expected_time"), [(False, False, None), (True, False, 3.0), (False, True, 5.0), (True, True, 5.0)], diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index b8076456..06a7028e 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -6,7 +6,6 @@ import logging import threading import time -from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -26,6 +25,7 @@ _new_kvcr, _op_entries, _poll_until, + _recovered_record, _start_write_message, _wait_until, ) @@ -33,7 +33,7 @@ from kvcr import DURATION_METRIC, TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC from kvcr.backends.local_dram_pools import _LocalDramResidency, _LocalDramState from kvcr.backends.remote_fw_dram import _RemoteFWDram -from kvcr.backends.remote_ops import _CONTROL_DECODER, _SourceWriteOp, _TargetMetadata +from kvcr.backends.remote_ops import _CONTROL_DECODER, _TargetMetadata from kvcr.backends.remote_sources import _FwMemResidency, _RemoteSources, _SourcePinOp from kvcr.config import KVCRConfig, LocalDramOptions from kvcr.runtime.core import _BlockRecord, _KVCRCore @@ -87,9 +87,7 @@ def _write_probe_message(op_handle: int, incarnation=None) -> bytes: @pytest.mark.parametrize("transfer_state,recovered", [("PROC", False), ("DONE", True)]) -def test_local_source_starts_inline_and_holds_its_slot( - transfer_state, recovered, monkeypatch -): +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) descriptors = [ _mem_descriptor(i, label=label, end_point_name="source") @@ -127,22 +125,25 @@ def test_local_source_starts_inline_and_holds_its_slot( { key: _BlockRecord( local_dram=_LocalDramResidency( - [("", 0), (":b", 1), (":c", 2)], _LocalDramState.READY + [("", 0), (":b", 1), (":c", 2)], + _LocalDramState.READY, + refs=tuple( + _TransferRef("old-agent", i, label, framework=False) + for i, label in enumerate(("", ":b", ":c")) + ), ) ) }, ) else: - deposit = source.deposit({key: descriptors}) + deposit = source.deposit( + {key: [descriptors[0], descriptors[2]]}, + expected_layout=["", ":b", ":c"], + ) assert _poll_until(source, bool) == [(deposit, _op_entries({key: True}))] prior_transfers = len(agent.xfers) - local = source._core._local_dram residency = source._core._block_record_map[key].local_dram assert residency.indices_by_label == {"": 0, ":b": 1, ":c": 2} - make_ref = Mock(wraps=local._pools._descriptor) - monkeypatch.setattr(local._pools, "_descriptor", make_ref) - callbacks.clear() - agent.state = transfer_state request = msgspec.msgpack.decode(_start_write_message(12, key)) request.update( target_regions=( @@ -163,37 +164,43 @@ def test_local_source_starts_inline_and_holds_its_slot( [_TransferRef("target", i, label) for i, label in enumerate((":c", ""))] ], ) - preparing, resume = threading.Event(), threading.Event() - initialize_write = _SourceWriteOp.__init__ - - def pause_write_preparation(op, *args, **kwargs): - preparing.set() - assert resume.wait(timeout=2) - initialize_write(op, *args, **kwargs) - + if not recovered: + # Ready pieces still read through the caller until the block is complete. + control.incoming.extend( + [ + msgspec.msgpack.encode(dict(request, op_handle=11)), + _write_probe_message(99), + ] + ) + _wait_until(lambda: bool(control.sent)) + assert len(agent.xfers) == prior_transfers + _poll_until( + source, + lambda _: ( + len(agent.released_xfers) == prior_transfers + 1 + and not _has_outstanding_operations(source) + ), + ) + assert [ + ctypes.string_at(addr, size) for addr, size, _ in agent.xfers[-1][1] + ] == [b"c" * 16, b"a" * 16] + deposit = source.deposit({key: descriptors[1:2]}) + assert _poll_until(source, bool) == [(deposit, _op_entries({key: True}))] + prior_transfers = len(agent.xfers) + callbacks.clear() + agent.state = transfer_state with ( - patch.object(_SourceWriteOp, "__init__", pause_write_preparation), # A resident-source request must not re-enter the submission queue. patch.object(source._core._progress, "submit", side_effect=AssertionError), + source._core._state_lock, ): control.incoming.append(msgspec.msgpack.encode(request)) - try: - assert preparing.wait(timeout=1) - # Claims and dependencies protect the source while P builds the - # write, but M must be able to acquire the metadata lock already. - assert source._core._state_lock.acquire(blocking=False) - try: - record = source._core._block_record_map[key] - assert record.local_dram.claim_count == 1 - assert record.in_flight_ops - finally: - source._core._state_lock.release() - finally: - resume.set() _wait_until(lambda: len(agent.xfers) == prior_transfers + 1) + if transfer_state == "PROC": + assert residency.source_readers == 1 # The peer request must be served without another caller-side poll. - assert residency.claim_count == 1 + assert residency.claim_count == 0 assert [ ctypes.string_at(addr, size) for addr, size, _ in agent.xfers[-1][1] ] == [ @@ -202,7 +209,6 @@ def pause_write_preparation(op, *args, **kwargs): ] assert agent.xfers[-1][3] == [(4096, 16, 7), (4112, 16, 7)] assert pinning.searches == [] - assert make_ref.call_count == 2 assert not source._core._remote_fw_dram._sources._source_pin_ops assert callbacks == [] if transfer_state == "PROC": @@ -211,8 +217,7 @@ def pause_write_preparation(op, *args, **kwargs): agent.state = "DONE" _wait_until(lambda: len(agent.released_xfers) == prior_transfers + 1) - assert residency.claim_count == 1 # Native completion does not release it. - assert _poll_until(source, lambda _: residency.claim_count == 0) == [] + assert _poll_until(source, lambda _: residency.source_readers == 0) == [] assert not source._core._local_dram_sources_by_op # The caller-queued path reuses the same parts without pinning the framework. @@ -228,7 +233,6 @@ def pause_write_preparation(op, *args, **kwargs): and residency.claim_count == 0 ), ) - assert make_ref.call_count == 2 assert agent.xfers[-1][1] == agent.xfers[-2][1] assert pinning.searches == [] finally: @@ -342,7 +346,7 @@ def test_local_source_waits_for_declared_pieces( @pytest.mark.parametrize( "outcome", ["write", "framework", "deadline", "cancel", "route", "shutdown"] ) -def test_local_source_retries_after_fill_publication(outcome): +def test_local_source_retries_after_reader_contention(outcome): memory = ctypes.create_string_buffer(16) descriptor = _mem_descriptor(end_point_name="source") agent, pinning, control = FakeNixlAgent(), FakePrimaryPinning(), FakeBytesControl() @@ -356,73 +360,53 @@ def test_local_source_retries_after_fill_publication(outcome): ) now = 0.0 source._core._clock = lambda: now - local = source._core._local_dram + local = source._core._local_dram._pools backend = source._core._remote_fw_dram key = BlockKey(b"local") requested_key = BlockKey(b"framework-only") if outcome == "framework" else key - publishing, resume = threading.Event(), threading.Event() - make_evictable = local._pools._make_evictable - - def pause_publication(key): - # READY is already set, but the caller's fill transaction is unfinished. - publishing.set() - assert resume.wait(timeout=2) - make_evictable(key) - try: agent.state = "DONE" deposit = source.deposit({key: [descriptor]}) - _wait_until(lambda: not source._core._progress._completed.empty()) - with ( - patch.object( - local._pools, "_make_evictable", side_effect=pause_publication - ), - ThreadPoolExecutor(max_workers=1) as caller, - ): - completion = caller.submit(source.poll_completed) - try: - assert publishing.wait(timeout=1) - agent.state = "PROC" - control.incoming.extend( - [ - _start_write_message(12, requested_key, target_agent="target"), - _write_probe_message(99), - ] - ) - _wait_until( - lambda: any( - _decode_control_message(raw).get("type") == "write_probe_ack" - for _, raw in control.sent - ) + assert _poll_until(source, bool) == [(deposit, _op_entries({key: True}))] + agent.state = "PROC" + with local._source_state.lock: + control.incoming.extend( + [ + _start_write_message(12, requested_key, target_agent="target"), + _write_probe_message(99), + ] + ) + _wait_until( + lambda: any( + _decode_control_message(raw).get("type") == "write_probe_ack" + for _, raw in control.sent ) - _wait_until(lambda: bool(backend._pending_local_sources)) + ) + _wait_until(lambda: bool(backend._pending_local_sources)) + assert len(agent.xfers) == 1 + assert source._core._block_record_map[key].local_dram.claim_count == 0 + if outcome == "deadline": + now = 2.0 + elif outcome == "cancel": + control.incoming.append(_write_probe_message(12)) + elif outcome == "route": + backend._route_generation["target"] = 1 + elif outcome == "shutdown": + source._core._progress._submissions.put(_STOP) + if outcome not in ("write", "framework"): + _wait_until(lambda: bool(agent.sent_notifs)) + assert _decode_notif(agent.sent_notifs[0][1]) == { + "type": "write_done", + "op_handle": 12, + "success": False, + } assert len(agent.xfers) == 1 - assert source._core._block_record_map[key].local_dram.claim_count == 0 - if outcome == "deadline": - now = 2.0 - elif outcome == "cancel": - control.incoming.append(_write_probe_message(12)) - elif outcome == "route": - backend._route_generation["target"] = 1 - elif outcome == "shutdown": - source._core._progress._submissions.put(_STOP) - if outcome not in ("write", "framework"): - _wait_until(lambda: bool(agent.sent_notifs)) - assert _decode_notif(agent.sent_notifs[0][1]) == { - "type": "write_done", - "op_handle": 12, - "success": False, - } - assert len(agent.xfers) == 1 - finally: - resume.set() - assert completion.result(timeout=1) == [(deposit, _op_entries({key: True}))] # Ready local sources retry without another caller-side poll. if outcome == "write": _wait_until(lambda: len(agent.xfers) == 2) - assert source._core._block_record_map[key].local_dram.claim_count == 1 - assert local._pools.telemetry_state()["local_g2_evictable_slots"] == 0 + assert source._core._block_record_map[key].local_dram.source_readers == 1 + assert local.telemetry_state()["local_g2_evictable_slots"] == 0 elif outcome == "framework": assert _poll_until(source, lambda _: len(agent.xfers) == 2) == [] assert pinning.searches == ( @@ -435,7 +419,6 @@ def pause_publication(key): assert source._core._outstanding_operations == 0 assert source._core._block_record_map[key].local_dram.claim_count == 0 finally: - resume.set() agent.state = "DONE" source.close() assert not backend._pending_local_sources @@ -777,8 +760,10 @@ def send_notif(self, agent_name, notif_msg): @pytest.mark.parametrize("failure", ["initialize", "error", "exception", "async"]) +@pytest.mark.parametrize("local_source", [False, True], ids=["framework", "local"]) def test_kvcr_source_transfer_error_notifies_failure_and_cleans_up( failure: str, + local_source: bool, ) -> None: """Every source-transfer failure reports failure and releases its state.""" @@ -798,9 +783,23 @@ def transfer(self, handle): pinning = FakePrimaryPinning() control = FakeBytesControl() key = BlockKey(b"k0") - control.incoming.append(_start_write_message(5, key)) - kvcr = _new_kvcr(source_agent, pinning, control, name="source") + memory = ctypes.create_string_buffer(16) + kvcr = _new_kvcr( + source_agent, + pinning, + control, + name="source", + local_dram=( + LocalDramOptions([("", ctypes.addressof(memory), len(memory))]) + if local_source + else None + ), + ) + if local_source: + install_recovery_records(kvcr._core, {key: _recovered_record(g2=[("", 0)])}) + residency = kvcr._core._block_record_map[key].local_dram kvcr._core._clock = lambda: 0.0 # Failure must come from ERR, not a timeout. + control.incoming.append(_start_write_message(5, key)) if failure == "async": assert _poll_until(kvcr, lambda _: bool(source_agent.xfers)) == [] @@ -811,9 +810,20 @@ def transfer(self, handle): assert _decode_notif(source_agent.sent_notifs[0][1])["terminal"] is False assert pinning.unpins == [] assert source_agent.released_xfers == [] + if local_source: + assert residency.source_readers == 1 + assert kvcr._core._local_dram._pools.source_operations == 1 source_agent.state = "DONE" - assert _poll_until(kvcr, lambda _: pinning.unpins == ["pin"]) == [] + completed = _poll_until( + kvcr, + lambda _: source_agent.sent_notifs and not _has_outstanding_operations(kvcr), + ) + assert completed == [] + assert pinning.unpins == ([] if local_source else ["pin"]) + if local_source: + assert residency.source_readers == 0 + assert kvcr._core._local_dram._pools.source_operations == 0 assert source_agent.transfers == ([] if failure == "initialize" else [1]) agent_name, notif = source_agent.sent_notifs[-1] assert agent_name == "remote-1" @@ -937,23 +947,37 @@ def on_resilience_event(error): kvcr.close() -def test_source_lifecycles_distinguish_targets_reusing_the_same_handle(): +@pytest.mark.parametrize("local_source", [False, True], ids=["framework", "local"]) +def test_source_lifecycles_distinguish_targets_reusing_the_same_handle(local_source): now = 0.0 agent, control, errors, done = FakeNixlAgent(), FakeBytesControl(), [], set() agent.check_xfer_state = lambda handle: "DONE" if handle in done else "PROC" + memory = ctypes.create_string_buffer(16) source = _new_kvcr( agent, FakePrimaryPinning(), control, name="source", on_resilience_event=errors.append, + framework_regions=[_buffer_region(memory, 16)], + local_dram=( + LocalDramOptions([("", ctypes.addressof(memory), len(memory))]) + if local_source + else None + ), ) source._core._clock = lambda: now key = BlockKey(b"shared") + if local_source: + install_recovery_records(source._core, {key: _recovered_record(g2=[("", 0)])}) + residency = source._core._block_record_map[key].local_dram try: for target in ("target-a", "target-b"): control.incoming.append(_start_write_message(12, key, target_agent=target)) _poll_until(source, lambda _: len(agent.xfers) == 2) + if local_source: + assert residency.source_readers == 2 + assert source._core._local_dram._pools.source_operations == 2 now = 5.0 _poll_until(source, lambda _: len(errors) == 2) pending = {error.op_handle for error in errors} @@ -964,7 +988,8 @@ def test_source_lifecycles_distinguish_targets_reusing_the_same_handle(): == errors[1].source_blocks == {key: [_mem_descriptor(end_point_name="source")]} ) - assert errors[0].source_owners == errors[1].source_owners == {key: "framework"} + owner = "kvcr" if local_source else "framework" + assert errors[0].source_owners == errors[1].source_owners == {key: owner} for native_handle in (1, 2): done.add(native_handle) _poll_until(source, lambda _: len(errors) == 2 + native_handle) @@ -974,13 +999,29 @@ def test_source_lifecycles_distinguish_targets_reusing_the_same_handle(): assert event.source_owners == errors[0].source_owners pending.remove(event.op_handle) assert len(pending) == 2 - native_handle + if local_source: + assert residency.source_readers == 2 - native_handle + assert ( + source._core._local_dram._pools.source_operations + == 2 - native_handle + ) + if native_handle == 1: + replacement = BlockKey(b"replacement") + blocked = source.deposit( + {replacement: [_mem_descriptor(end_point_name="source")]} + ) + assert source.poll_completed() == [ + (blocked, _op_entries({replacement: False})) + ] + assert len(agent.xfers) == 2 assert not _has_outstanding_operations(source) finally: done.update((1, 2)) source.close() -def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(): +@pytest.mark.parametrize("mixed", [False, True]) +def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(mixed): now = 0.0 memory = ctypes.create_string_buffer(16) descriptor = _mem_descriptor(end_point_name="source") @@ -997,10 +1038,11 @@ def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(): source._core._clock = lambda: now key, replacement = BlockKey(b"k0"), BlockKey(b"k1") missing, framework_hit = BlockKey(b"missing"), BlockKey(b"framework-hit") - expected_sources = { - key: [descriptor], - framework_hit: [descriptor], - } + expected_sources = {key: [descriptor]} + expected_owners = {key: "kvcr"} + if mixed: + expected_sources[framework_hit] = [descriptor] + expected_owners[framework_hit] = "framework" try: agent.state = "DONE" deposit = source.deposit({key: [descriptor]}) @@ -1009,9 +1051,9 @@ def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(): payload = msgspec.msgpack.decode( _start_write_message(12, key, target_agent="target") ) - payload["keys"] = [key, missing, framework_hit] + payload["keys"] = [key, missing, framework_hit] if mixed else [key] payload["dst_descriptors"] = [ - [_TransferRef("target", index)] for index in range(3) + [_TransferRef("target", index)] for index in range(len(payload["keys"])) ] control.incoming.append(msgspec.msgpack.encode(payload)) _poll_until(source, lambda _: len(agent.xfers) == 2) @@ -1023,14 +1065,15 @@ def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(): assert list(source.poll_completed()) == [ (blocked, _op_entries({replacement: False})) ] - assert source._core._block_record_map[key].local_dram.claim_count == 1 + residency = source._core._block_record_map[key].local_dram + assert residency.claim_count + residency.source_readers == 1 agent.state = "DONE" assert _poll_until(source, lambda _: len(errors) == 2) == [] assert [error.state for error in errors] == ["uncertain", "quiesced"] - assert source._core._block_record_map[key].local_dram.claim_count == 0 + assert residency.claim_count + residency.source_readers == 0 assert errors[1].source_blocks == expected_sources for error in errors: - assert error.source_owners == {key: "kvcr", framework_hit: "framework"} + assert error.source_owners == expected_owners assert error.destination_regions is error.destination_owner is None deposit = source.deposit({replacement: [descriptor]}) assert _poll_until(source, bool) == [ diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index 950bd638..25a88e37 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -232,6 +232,7 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: [("full", 0), ("swa", 10)], _LocalDramState.READY, claim_count=2, + source_readers=2, retire_on_release=True, pending_deposits=1, writers=[1, 1], @@ -281,7 +282,10 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: } # A kept mirror must match exactly what the handback frames carry. residency = served[ready].local_dram - residency.refs[0] = _TransferRef("old-agent", 0, "full", framework=False) + residency.refs = ( + _TransferRef("old-agent", 0, "full", framework=False), + _TransferRef("old-agent", 10, "swa", framework=False), + ) framed = { BlockKey(key): _decode_recovery_record(payload, _TWO_POOLS) for _, key, payload in _recovery_frames(served) @@ -291,6 +295,7 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: assert mirror._records is served assert mirror._records[ready].local_dram is residency + assert residency.source_readers == 0 assert not residency.refs # A new agent must regenerate its own references. assert partially_ready not in mirror._records assert mirror._records == framed