From e5b25a1b1e5d2bc0fa06a4e8901c51ffc1eddf2e Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Wed, 9 Sep 2026 06:41:24 +0300 Subject: [PATCH 01/16] feat(core)!: support multiple pools BREAKING CHANGE: local DRAM residencies and transfer results carry ordered pool-and-slot locations. Signed-off-by: Kapil Arya --- src/kvcr/core.py | 65 +++-- src/kvcr/guard.py | 16 +- src/kvcr/local_disk.py | 2 +- src/kvcr/local_dram.py | 376 ++++++++++++++++---------- src/kvcr/policy_runtime.py | 16 +- src/kvcr/recovery_journal.py | 72 +++-- src/kvcr/remote_fw_dram.py | 67 +++-- tests/unit/_kvcr_test_utils.py | 6 +- tests/unit/test_g3.py | 34 +++ tests/unit/test_guard.py | 46 ++-- tests/unit/test_kvcr.py | 33 +-- tests/unit/test_kvcr_local_dram.py | 164 ++++++++++- tests/unit/test_kvcr_remote_source.py | 11 +- tests/unit/test_recovery_journal.py | 47 ++-- tests/unit/test_recovery_mirror.py | 110 +++++--- 15 files changed, 730 insertions(+), 335 deletions(-) diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 7c6f130..60c4506 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -118,9 +118,7 @@ def __init__( self.config = config self.pool_layouts = list(config.pool_layouts) _validate_pool_layouts(self.pool_layouts) - # TODO: Support multiple pools after remote fetch and G3 discover layouts. - if len(self.pool_layouts) != 1: - raise ValueError("only a single pool is currently supported") + self._block_sizes = dict(self.pool_layouts) self.block_size_bytes = self.pool_layouts[0][1] if self.config.operation_timeout_ms <= 0: raise ValueError("operation_timeout_ms must be positive") @@ -140,6 +138,8 @@ def __init__( self._stats_factory = bindings.stats_factory local_dram_config = backend_configs.local_dram g3_config = backend_configs.g3 + if g3_config is not None and len(self.pool_layouts) != 1: + raise ValueError("G3 does not support multiple pools") policy = bindings.policy if policy is None: policy = G3LRUPolicy() if g3_config is not None else LRUPolicy() @@ -167,7 +167,9 @@ def __init__( self._closed = False self._outstanding_operations = 0 self._framework_pin_keys: dict[PinHandle, set[BlockKey]] = {} - self._local_dram_sources_by_op: dict[_OpId, dict[BlockKey, MemDescriptor]] = {} + self._local_dram_sources_by_op: dict[ + _OpId, dict[BlockKey, list[MemDescriptor]] + ] = {} self._completion_queue: list[OpResult] = [] self._joined_completions: dict[ @@ -231,7 +233,7 @@ def __init__( if framework_dram is not None: memory_regions.append((framework_dram.address, framework_dram.length)) if self._local_dram is not None: - memory_regions.append(self._local_dram.memory_region) + memory_regions.extend(self._local_dram.memory_regions) dram_backends: set[str] = set() if self._local_dram is not None: dram_backends.add(local_dram_config.backend) @@ -297,7 +299,8 @@ def adopt_recovery_records(self, records: dict[BlockKey, _BlockRecord]) -> None: # so routing through it would skip everything recovered into both. for key, record in records.items(): if record.local_dram is not None or g3 is None: - source, slot_size = CacheTier.LOCAL_G2, local_dram._slot_size + source = CacheTier.LOCAL_G2 + slot_size = local_dram._size_bytes(record.local_dram.slots) else: source, slot_size = CacheTier.G3, g3._slot_size self._policy.on_ingest(self._block_meta(key, record, slot_size), source) @@ -366,14 +369,14 @@ def deliver( key: self._normalize_descriptors(descriptors) for key, descriptors in blocks.items() } - local_blocks: dict[BlockKey, MemDescriptor] = {} + local_blocks: dict[BlockKey, list[MemDescriptor]] = {} g3_blocks: dict[BlockKey, MemDescriptor] = {} - remote_blocks: dict[BlockKey, MemDescriptor] = {} + remote_blocks: dict[BlockKey, list[MemDescriptor]] = {} for key, destination in normalized.items(): if self._is_local_resident(key): local_blocks[key] = destination elif self._g3 is not None and self._g3.is_ready(key): - g3_blocks[key] = destination + g3_blocks[key] = destination[0] else: remote_blocks[key] = destination @@ -429,9 +432,13 @@ def fetch( expected_layout: list[str] | None = None, hints: object | None = None, ) -> OpHandle: - expected_layout = [""] if expected_layout is None else expected_layout - if expected_layout != [self.pool_layouts[0][0]]: - raise ValueError("expected layout must match the configured pool_layouts") + expected_layout = [""] if expected_layout is None else list(expected_layout) + if not expected_layout or any( + name not in self._block_sizes for name in expected_layout + ): + raise ValueError("expected layout must use configured pools") + if self._g3 is not None and len(expected_layout) != 1: + raise ValueError("G3 does not support multi-block layouts") op_handle = self._next_op_handle self._next_op_handle += 1 local_dram = self._local_dram @@ -459,6 +466,7 @@ def fetch( request_id, deadline, hints=hints, + layout=expected_layout, ) for source in (CacheTier.G3, CacheTier.REMOTE_G2): self._start_local_fill( @@ -661,7 +669,7 @@ def _remove_block_dependencies(self, op: _Op) -> None: def _start_local_fill( self, source: CacheTier, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], request_id: str | None, deadline: float, ) -> None: @@ -671,7 +679,9 @@ def _start_local_fill( self._next_fill_handle -= 1 if source is CacheTier.G3: started = self._g3 is not None and self._g3.start_fill( - fill_handle, dict(blocks), deadline + fill_handle, + {key: descriptors[0] for key, descriptors in blocks.items()}, + deadline, ) elif source is CacheTier.REMOTE_G2: started = self._remote_fw_dram.fetch( @@ -684,7 +694,7 @@ def _start_local_fill( def _claim_local_dram_sources( self, op_id: _OpId, keys: Collection[BlockKey] - ) -> Mapping[BlockKey, MemDescriptor]: + ) -> Mapping[BlockKey, list[MemDescriptor]]: sources = self._local_dram_sources_by_op.get(op_id, {}) if self._local_dram is not None: claimed = self._local_dram.acquire_sources( @@ -695,17 +705,24 @@ def _claim_local_dram_sources( self._local_dram_sources_by_op[op_id] = sources return sources - def _normalize_descriptors(self, descriptors: list[MemDescriptor]) -> MemDescriptor: + def _normalize_descriptors( + self, descriptors: list[MemDescriptor] + ) -> list[MemDescriptor]: if not isinstance(descriptors, list): raise TypeError("block descriptors must be a list") - if len(descriptors) != 1 or not isinstance(descriptors[0], MemDescriptor): - raise ValueError("each block requires exactly one descriptor") - descriptor = descriptors[0] - if descriptor.info != self.pool_layouts[0][0]: - raise ValueError(f"unknown descriptor pool {descriptor.info!r}") - if descriptor.size != self.block_size_bytes: - raise ValueError("block descriptor has the wrong byte count") - return descriptor + if not descriptors or not all( + isinstance(descriptor, MemDescriptor) for descriptor in descriptors + ): + raise ValueError("each block requires at least one descriptor") + for descriptor in descriptors: + block_size = self._block_sizes.get(descriptor.info) + if block_size is None: + raise ValueError(f"unknown descriptor pool {descriptor.info!r}") + if descriptor.size != block_size: + raise ValueError("block descriptor has the wrong byte count") + if self._g3 is not None and len(descriptors) != 1: + raise ValueError("G3 does not support multi-block layouts") + return list(descriptors) def _release_local_dram_sources( self, diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index 9b37f49..61f3309 100644 --- a/src/kvcr/guard.py +++ b/src/kvcr/guard.py @@ -193,10 +193,10 @@ def recover(self, pool_layouts: PoolBlockLayouts) -> _RecoveryMirror: return self.mirror return read_handback(self.attachment, self._compatibility_digest, pool_layouts) - def start_primary(self) -> None: + def start_primary(self, pool_layouts: PoolBlockLayouts) -> None: """Arm recovery for the accepted primary and reset its journal.""" if self.mirror is None: - self.mirror = _RecoveryMirror() + self.mirror = _RecoveryMirror(tuple(name for name, _ in pool_layouts)) self._journal.reset() def poll(self) -> bool: @@ -220,7 +220,9 @@ def invalidate_journal(self) -> None: with suppress(Exception): self._journal.invalidate() - def take_for_promotion(self) -> dict[BlockKey, _BlockRecord]: + def take_for_promotion( + self, pool_layouts: PoolBlockLayouts + ) -> dict[BlockKey, _BlockRecord]: """Drain and transfer recovered records, leaving a fresh mirror.""" records: dict[BlockKey, _BlockRecord] = {} mirror = self.mirror @@ -235,7 +237,7 @@ def take_for_promotion(self) -> dict[BlockKey, _BlockRecord]: except RecoveryJournalError as error: self._drop_recovery(error) # A handover still needs somewhere to put the core's eventual records. - self.mirror = _RecoveryMirror() + self.mirror = _RecoveryMirror(tuple(name for name, _ in pool_layouts)) return records def prepare_to_serve( @@ -316,7 +318,7 @@ def _write_handback( write_recovery_snapshot( self.attachment, canonical_pool_terms(self._compatibility_digest, pool_layouts, self._spec), - _recovery_frames(records), + _recovery_frames(records, tuple(name for name, _ in pool_layouts)), ) def close(self) -> None: @@ -754,7 +756,7 @@ def _adopt(self, control: ZmqPeerControlChannel, tier_config: _TierConfig) -> No self._hand_back(served_under) self._resumable = True # A refused handback is cold for the new lease, not unmirrored. - self._recovery.start_primary() + self._recovery.start_primary(tier_config.pool_layouts) # The old channel is the last reference to the prior primary's listener. if self._control is not None: self._control.close() @@ -842,7 +844,7 @@ def _promote(self) -> None: self._resumable = False if self._failure is not None: raise self._failure - self._serve(self._recovery.take_for_promotion()) + self._serve(self._recovery.take_for_promotion(self._configured.pool_layouts)) def _serve(self, records: dict[BlockKey, _BlockRecord]) -> None: """Answer on this pool's endpoint, with whatever came back from it. diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index c480976..3a376c3 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -293,7 +293,7 @@ def resolve_eviction( if key not in sources: return (PlacementAction.KEEP, None), False try: - if not self._start_store(op_id, sources, deadline): + if not self._start_store(op_id, {key: sources[key][0]}, deadline): self._recover_store_failure(key, "G3 destination unavailable") return (PlacementAction.KEEP, None), False except Exception: diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index f9e5a05..44a38a5 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -3,7 +3,7 @@ """KVCR-owned local DRAM slots, claims, and transfers.""" import logging -from collections import deque +from collections import Counter, deque from collections.abc import Callable, Collection, Mapping from dataclasses import dataclass, field from enum import Enum, auto @@ -40,7 +40,7 @@ class _LocalDramState(Enum): @dataclass(slots=True) class _LocalDramResidency: - slot: int + slots: list[tuple[str, int]] state: _LocalDramState claim_count: int = 0 retire_on_release: bool = False @@ -59,7 +59,7 @@ class _PendingResidencyOp(_Op): @dataclass class _PendingDeliverOp(_Op): deadline: float - destinations: Mapping[BlockKey, MemDescriptor] + destinations: Mapping[BlockKey, list[MemDescriptor]] results: dict[BlockKey, OpEntryResult] = field(default_factory=dict) active_keys: set[BlockKey] = field(default_factory=set) @@ -68,14 +68,15 @@ class _PendingDeliverOp(_Op): class _CapacityWaiter: op: _PendingResidencyOp key: BlockKey - source: MemDescriptor | CacheTier + source: list[MemDescriptor] | CacheTier + layout: list[str] @dataclass class _LocalCopyOp(_ProgressOp): deliver_op_id: _OpId | None ordered_keys: tuple[BlockKey, ...] - local_slots: tuple[int, ...] + local_slots: tuple[tuple[tuple[str, int], ...], ...] src_descriptors: tuple[MemDescriptor, ...] dst_descriptors: tuple[MemDescriptor, ...] deadline: float @@ -135,34 +136,36 @@ def close(self, progress: _KVCRProgress) -> bool: class _LocalDram: - """Main-thread metadata for one externally allocated DRAM region.""" + """Main-thread metadata for KVCR-owned DRAM pools.""" def __init__( self, kvcr: "_KVCRCore", region: LocalDramOptions, ) -> None: - if len(region.pools) != 1: - raise ValueError("local DRAM supports only a single pool") - pool_name, address, length = region.pools[0] - if address <= 0: - raise ValueError("local DRAM address must be positive") - if type(length) is not int or length <= 0: - raise ValueError("local DRAM pool size must be a positive integer") - if pool_name != kvcr.pool_layouts[0][0]: - raise ValueError("local DRAM pool name must match pool_layouts") + if [pool[0] for pool in region.pools] != [ + pool[0] for pool in kvcr.pool_layouts + ]: + raise ValueError("local DRAM pools must match pool_layouts") if not region.backend: raise ValueError("local DRAM NIXL backend must be non-empty") self._kvcr = kvcr self._backend = region.backend - self._address = address - self._length = length - self._slot_size = kvcr.block_size_bytes - slot_count = length // self._slot_size - if not slot_count: - raise ValueError("local DRAM pool must hold at least one block") - self._free_slots = deque(range(slot_count)) + self._pools: dict[str, tuple[int, int, int]] = {} + self._free_slots: dict[str, deque[int]] = {} + for (pool_name, address, length), (_, slot_size) in zip( + region.pools, kvcr.pool_layouts + ): + if address <= 0: + raise ValueError("local DRAM address must be positive") + if type(length) is not int or length <= 0: + raise ValueError("local DRAM pool size must be a positive integer") + slot_count = length // slot_size + if not slot_count: + raise ValueError("local DRAM pool must hold at least one block") + self._pools[pool_name] = (address, length, slot_size) + self._free_slots[pool_name] = deque(range(slot_count)) self._evictable = _EvictionQueue() self._unscored: set[BlockKey] = set() self._pending_residency_ops: dict[_OpId, _PendingResidencyOp] = {} @@ -182,12 +185,12 @@ def __init__( ) @property - def memory_region(self) -> tuple[int, int]: - return self._address, self._length + def memory_regions(self) -> tuple[tuple[int, int], ...]: + return tuple((address, length) for address, length, _ in self._pools.values()) @property def _total_slots(self) -> int: - return self._length // self._slot_size + return sum(length // slot_size for _, length, slot_size in self._pools.values()) def observe_residency( self, observer: Callable[[BlockKey, "_BlockRecord"], None] @@ -201,24 +204,31 @@ def adopt_recovery_slots(self, records: Mapping[BlockKey, "_BlockRecord"]) -> No with them. Ranking them is rank_recovered, which needs the policy to have seen every block first. """ - slot_count = self._total_slots - occupied: set[int] = set() + occupied = {pool_name: set() for pool_name in self._pools} for record in records.values(): residency = record.local_dram if residency is None: continue - slot = residency.slot - if ( - residency.state is not _LocalDramState.READY - or type(slot) is not int - or not 0 <= slot < slot_count - or slot in occupied - ): + if residency.state is not _LocalDramState.READY or not residency.slots: raise ValueError("invalid local DRAM recovery slots") - occupied.add(slot) - self._free_slots = deque( - slot for slot in range(slot_count) if slot not in occupied - ) + for pool_name, slot in residency.slots: + pool = self._pools.get(pool_name) + if ( + pool is None + or type(slot) is not int + or not 0 <= slot < pool[1] // pool[2] + or slot in occupied[pool_name] + ): + raise ValueError("invalid local DRAM recovery slots") + occupied[pool_name].add(slot) + self._free_slots = { + pool_name: deque( + slot + for slot in range(length // slot_size) + if slot not in occupied[pool_name] + ) + for pool_name, (_, length, slot_size) in self._pools.items() + } def rank_recovered(self, records: Mapping[BlockKey, "_BlockRecord"]) -> None: """Make recovered rows evictable, once the policy can score them. @@ -234,17 +244,18 @@ def rank_recovered(self, records: Mapping[BlockKey, "_BlockRecord"]) -> None: def telemetry_state(self) -> dict[str, int]: total_slots = self._total_slots + free_slots = sum(map(len, self._free_slots.values())) return { "local_g2_total_slots": total_slots, - "local_g2_free_slots": len(self._free_slots), - "local_g2_allocated_slots": total_slots - len(self._free_slots), - "local_g2_evictable_slots": len(self._evictable), + "local_g2_free_slots": free_slots, + "local_g2_allocated_slots": total_slots - free_slots, + "local_g2_evictable_slots": self._evictable.total_weight, } def deposit( self, op_handle: OpHandle, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], *, no_evict: bool, hints: object | None, @@ -265,15 +276,17 @@ def deposit( self._kvcr._add_block_dependencies(op, new_operation=True) copy_keys: list[BlockKey] = [] - slots: list[int] = [] + slots: list[tuple[tuple[str, int], ...]] = [] src_descriptors: list[MemDescriptor] = [] dst_descriptors: list[MemDescriptor] = [] evicted: list[BlockKey] = [] - for key, src in blocks.items(): + for key, sources in blocks.items(): record = self._kvcr._block_record(key) residency = record.local_dram if residency is not None: - if residency.state is _LocalDramState.READY: + if not self._same_layout(residency.slots, sources): + op.results[key] = OpEntryResult(OpEntryStatus.FAILED) + elif residency.state is _LocalDramState.READY: op.results[key] = ( self._new_public_claim( key, residency, include_descriptors=False @@ -284,11 +297,9 @@ def deposit( elif residency.state is _LocalDramState.DISCARDING: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) continue - if src.size != self._slot_size: - op.results[key] = OpEntryResult(OpEntryStatus.FAILED) - continue + size_bytes = sum(source.size for source in sources) decision = self._kvcr._policy.decide_ingest( - self._kvcr._block_meta(key, record, self._slot_size), + self._kvcr._block_meta(key, record, size_bytes), CacheTier.FW_G2, required_local=no_evict, framework_hints=hints, @@ -296,22 +307,25 @@ def deposit( if decision[0] is PlacementAction.DROP: op.results[key] = OpEntryResult(OpEntryStatus.DROPPED) continue - slot, evicted_key, eviction_pending = self._allocate_slot(keys, deadline) - if slot is None: + locations, evicted_keys, eviction_pending = self._allocate_slots( + [source.info for source in sources], keys, deadline + ) + evicted.extend(evicted_keys) + if locations is None: if eviction_pending: - self._enqueue_capacity_waiter(op, key, src) + self._enqueue_capacity_waiter( + op, key, sources, [source.info for source in sources] + ) else: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) continue - if evicted_key is not None: - evicted.append(evicted_key) self._kvcr._block_record(key).local_dram = _LocalDramResidency( - slot, _LocalDramState.FILLING + locations, _LocalDramState.FILLING ) copy_keys.append(key) - slots.append(slot) - src_descriptors.append(src) - dst_descriptors.append(self._descriptor(slot)) + slots.append(tuple(locations)) + src_descriptors.extend(sources) + dst_descriptors.extend(self._descriptors(locations)) self._update_capacity_pressure() self._kvcr._publish_inventory(evicted, CacheTier.LOCAL_G2, removed=True) @@ -343,7 +357,8 @@ def fetch( deadline: float, *, hints: object | None, - ) -> dict[BlockKey, MemDescriptor]: + layout: list[str], + ) -> dict[BlockKey, list[MemDescriptor]]: ordered_keys = tuple(dict.fromkeys(keys)) key_set = set(ordered_keys) if not key_set: @@ -369,16 +384,21 @@ def fetch( else: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) elif residency.state is _LocalDramState.READY: - self._kvcr._record_access((key,)) - op.results[key] = self._new_public_claim( - key, residency, include_descriptors=True - ) + if [name for name, _ in residency.slots] == layout: + self._kvcr._record_access((key,)) + op.results[key] = self._new_public_claim( + key, residency, include_descriptors=True + ) + else: + op.results[key] = OpEntryResult(OpEntryStatus.FAILED) + elif [name for name, _ in residency.slots] != layout: + op.results[key] = OpEntryResult(OpEntryStatus.FAILED) elif residency.state is _LocalDramState.DISCARDING: # A discarded fill still owns its slot, so this block cannot be # reserved yet. Wait for the slot instead of failing a key a # lower tier can still serve. if key in sources: - self._enqueue_capacity_waiter(op, key, sources[key]) + self._enqueue_capacity_waiter(op, key, sources[key], layout) else: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) destinations, eviction_pending = self.reserve_fill( @@ -387,10 +407,11 @@ def fetch( required_local=True, deadline=deadline, framework_hints=hints, + layouts={key: layout for key in to_reserve}, ) op.remote_fill_keys.update(destinations) for key in eviction_pending: - self._enqueue_capacity_waiter(op, key, sources[key]) + self._enqueue_capacity_waiter(op, key, sources[key], layout) for key in to_reserve: if key not in destinations and key not in eviction_pending: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) @@ -399,7 +420,7 @@ def fetch( def complete_fill(self, keys: Collection[BlockKey], *, success: bool) -> None: ordered_keys = tuple(keys) - slots: list[int] = [] + slots: list[tuple[tuple[str, int], ...]] = [] for key in ordered_keys: record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None @@ -413,7 +434,7 @@ def complete_fill(self, keys: Collection[BlockKey], *, success: bool) -> None: or (success and residency.state is not _LocalDramState.FILLING) ): raise RuntimeError(f"local DRAM fill state lost for {key!r}") - slots.append(residency.slot) + slots.append(tuple(residency.slots)) self._apply_fill_result( ordered_keys, tuple(slots), success, CacheTier.REMOTE_G2 ) @@ -421,7 +442,7 @@ def complete_fill(self, keys: Collection[BlockKey], *, success: bool) -> None: def deliver( self, op_handle: OpHandle, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], *, deadline: float, ) -> None: @@ -450,8 +471,8 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: def acquire_sources( self, keys: Collection[BlockKey] - ) -> dict[BlockKey, MemDescriptor]: - sources: dict[BlockKey, MemDescriptor] = {} + ) -> dict[BlockKey, list[MemDescriptor]]: + sources: dict[BlockKey, list[MemDescriptor]] = {} for key in keys: if key in sources: continue @@ -460,7 +481,7 @@ def acquire_sources( if residency is None or residency.state is not _LocalDramState.READY: continue self._acquire_claim(key, residency) - sources[key] = self._descriptor(residency.slot) + sources[key] = self._descriptors(residency.slots) self._update_capacity_pressure() return sources @@ -549,12 +570,13 @@ def poll_main(self, items: Collection[object]) -> list[object]: return unhandled def _finish_copy(self, copy: _LocalCopyOp) -> None: + byte_count = sum(descriptor.size for descriptor in copy.src_descriptors) self._kvcr._record_transfer( "local_deliver" if copy.deliver_op_id is not None else "local_fill", copy.started_at, copy.success, len(copy.ordered_keys), - len(copy.ordered_keys) * self._slot_size, + byte_count, ) if copy.deliver_op_id is not None: self._finish_delivery_copy(copy) @@ -570,7 +592,7 @@ def _finish_copy(self, copy: _LocalCopyOp) -> None: def _apply_fill_result( self, ordered_keys: tuple[BlockKey, ...], - local_slots: tuple[int, ...], + local_slots: tuple[tuple[tuple[str, int], ...], ...], success: bool, source: CacheTier, ) -> None: @@ -579,13 +601,13 @@ def _apply_fill_result( affected_deliver_ops: dict[_OpId, _PendingDeliverOp] = {} deliver_keys: dict[_OpId, list[BlockKey]] = {} now = self._kvcr._clock() - for key, slot in zip(ordered_keys, local_slots): + for key, slots in zip(ordered_keys, local_slots): 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 - or residency.slot != slot + or residency.slots != list(slots) or residency.state not in ( _LocalDramState.FILLING, @@ -598,13 +620,15 @@ def _apply_fill_result( record.last_access = now residency.state = _LocalDramState.READY self._residency_observer(key, record) - meta = self._kvcr._block_meta(key, record, self._slot_size) + meta = self._kvcr._block_meta( + key, record, self._size_bytes(residency.slots) + ) self._kvcr._on_ingest(meta, source) self._make_evictable(key) committed.append(key) else: record.local_dram = None - self._free_slots.append(slot) + self._free(residency.slots) for op_id in record.active_op_ids: residency_op = self._pending_residency_ops.get(op_id) @@ -658,10 +682,11 @@ def reserve_fill( required_local: bool, deadline: float, framework_hints: object | None = None, - ) -> tuple[dict[BlockKey, MemDescriptor], set[BlockKey]]: + layouts: Mapping[BlockKey, list[str]], + ) -> tuple[dict[BlockKey, list[MemDescriptor]], set[BlockKey]]: keys = tuple(dict.fromkeys(keys)) protected = set(keys) - destinations: dict[BlockKey, MemDescriptor] = {} + destinations: dict[BlockKey, list[MemDescriptor]] = {} eviction_pending: set[BlockKey] = set() evicted: list[BlockKey] = [] for key in keys: @@ -671,24 +696,29 @@ def reserve_fill( if record.local_dram is not None: continue decision = self._kvcr._policy.decide_ingest( - self._kvcr._block_meta(key, record, self._slot_size), + self._kvcr._block_meta( + key, + record, + sum(self._pools[name][2] for name in layouts[key]), + ), sources[key], required_local, framework_hints=framework_hints, ) if decision[0] is PlacementAction.DROP: continue - slot, evicted_key, waiting = self._allocate_slot(protected, deadline) - if slot is None: + locations, evicted_keys, waiting = self._allocate_slots( + layouts[key], protected, deadline + ) + evicted.extend(evicted_keys) + if locations is None: if waiting: eviction_pending.add(key) continue - if evicted_key is not None: - evicted.append(evicted_key) self._kvcr._block_record(key).local_dram = _LocalDramResidency( - slot, _LocalDramState.FILLING + locations, _LocalDramState.FILLING ) - destinations[key] = self._descriptor(slot) + destinations[key] = self._descriptors(locations) self._update_capacity_pressure() self._kvcr._publish_inventory(evicted, CacheTier.LOCAL_G2, removed=True) return destinations, eviction_pending @@ -697,7 +727,7 @@ def _start_deliveries( self, op: _PendingDeliverOp, keys: Collection[BlockKey] ) -> None: copy_keys: list[BlockKey] = [] - local_slots: list[int] = [] + local_slots: list[tuple[tuple[str, int], ...]] = [] src_descriptors: list[MemDescriptor] = [] dst_descriptors: list[MemDescriptor] = [] now = self._kvcr._clock() @@ -712,7 +742,7 @@ def _start_deliveries( continue elif ( residency.state is _LocalDramState.DISCARDING - or op.destinations[key].size != self._slot_size + or not self._same_layout(residency.slots, op.destinations[key]) or now >= op.deadline ): op.results[key] = OpEntryResult(OpEntryStatus.FAILED) @@ -720,9 +750,9 @@ def _start_deliveries( self._acquire_claim(key, residency) op.active_keys.add(key) copy_keys.append(key) - local_slots.append(residency.slot) - src_descriptors.append(self._descriptor(residency.slot)) - dst_descriptors.append(op.destinations[key]) + local_slots.append(tuple(residency.slots)) + src_descriptors.extend(self._descriptors(residency.slots)) + dst_descriptors.extend(op.destinations[key]) self._update_capacity_pressure() if copy_keys: @@ -748,12 +778,12 @@ def _finish_delivery_copy(self, copy: _LocalCopyOp) -> None: if copy.deliver_op_id is None: raise RuntimeError("local delivery has no owning operation") op = self._pending_deliver_ops[copy.deliver_op_id] - for key, slot in zip(copy.ordered_keys, copy.local_slots): + for key, slots in zip(copy.ordered_keys, copy.local_slots): 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.slot != slot + or residency.slots != list(slots) or residency.state is not _LocalDramState.READY ): raise RuntimeError(f"local DRAM delivery state lost for {key!r}") @@ -818,11 +848,12 @@ def _enqueue_capacity_waiter( self, op: _PendingResidencyOp, key: BlockKey, - source: MemDescriptor | CacheTier, + source: list[MemDescriptor] | CacheTier, + layout: list[str], ) -> None: if key in op.capacity_waiters: raise RuntimeError(f"duplicate local capacity waiter for {key!r}") - self._capacity_waiters.append(_CapacityWaiter(op, key, source)) + self._capacity_waiters.append(_CapacityWaiter(op, key, source, layout)) op.capacity_waiters.add(key) def _resume_capacity_waiters(self) -> None: @@ -855,7 +886,9 @@ def _resume_capacity_waiters(self) -> None: if residency is not None: self._capacity_waiters.popleft() op.capacity_waiters.remove(waiter.key) - if residency.state is _LocalDramState.READY: + if [name for name, _ in residency.slots] != waiter.layout: + op.results[waiter.key] = OpEntryResult(OpEntryStatus.FAILED) + elif residency.state is _LocalDramState.READY: op.results[waiter.key] = ( self._new_public_claim( waiter.key, @@ -870,36 +903,32 @@ def _resume_capacity_waiters(self) -> None: self._finish_residency_if_ready(op) continue - evicted_key: BlockKey | None = None - if self._free_slots: - slot = self._free_slots.popleft() - elif self._capacity_eviction_key is not None: - break - else: - slot, evicted_key, eviction_pending = self._allocate_slot( - op.keys, op.deadline + locations, evicted_keys, eviction_pending = self._allocate_slots( + waiter.layout, op.keys, op.deadline + ) + if evicted_keys: + self._kvcr._publish_inventory( + evicted_keys, CacheTier.LOCAL_G2, removed=True ) - if slot is None: - if eviction_pending: - break - self._capacity_waiters.popleft() - op.capacity_waiters.remove(waiter.key) - op.results[waiter.key] = OpEntryResult(OpEntryStatus.FAILED) - self._finish_residency_if_ready(op) - continue + if locations is None: + if eviction_pending: + break + self._capacity_waiters.popleft() + op.capacity_waiters.remove(waiter.key) + op.results[waiter.key] = OpEntryResult(OpEntryStatus.FAILED) + self._finish_residency_if_ready(op) + continue self._capacity_waiters.popleft() op.capacity_waiters.remove(waiter.key) - if evicted_key is not None: - self._kvcr._publish_inventory( - (evicted_key,), CacheTier.LOCAL_G2, removed=True - ) - record.local_dram = _LocalDramResidency(slot, _LocalDramState.FILLING) + record.local_dram = _LocalDramResidency( + locations, _LocalDramState.FILLING + ) if isinstance(waiter.source, CacheTier): op.remote_fill_keys.add(waiter.key) self._kvcr._start_local_fill( waiter.source, - {waiter.key: self._descriptor(slot)}, + {waiter.key: self._descriptors(locations)}, op.request_id, op.deadline, ) @@ -910,9 +939,9 @@ def _resume_capacity_waiters(self) -> None: keys={waiter.key}, deliver_op_id=None, ordered_keys=(waiter.key,), - local_slots=(slot,), - src_descriptors=(waiter.source,), - dst_descriptors=(self._descriptor(slot),), + local_slots=(tuple(locations),), + src_descriptors=tuple(waiter.source), + dst_descriptors=tuple(self._descriptors(locations)), deadline=op.deadline, backend=self._backend, clock=self._kvcr._clock, @@ -937,7 +966,7 @@ def _new_public_claim( self._public_claims[handle] = (key, residency) return OpEntryResult( OpEntryStatus.SUCCESS, - [self._descriptor(residency.slot)] if include_descriptors else None, + self._descriptors(residency.slots) if include_descriptors else None, handle, ) @@ -960,10 +989,12 @@ def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: if residency.retire_on_release: record.local_dram = None self._residency_observer(key, record) - self._free_slots.append(residency.slot) + self._free(residency.slots) self.abandon_capacity_eviction(key) self._kvcr._on_remove( - self._kvcr._block_meta(key, record, self._slot_size) + self._kvcr._block_meta( + key, record, self._size_bytes(residency.slots) + ) ) self._kvcr._publish_inventory((key,), CacheTier.LOCAL_G2, removed=True) self._kvcr._prune_block_record(key) @@ -971,15 +1002,22 @@ def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: else: self._make_evictable(key) - def _allocate_slot( - self, protected: set[BlockKey], deadline: float - ) -> tuple[int | None, BlockKey | None, bool]: - if self._free_slots: - return self._free_slots.popleft(), None, False + def _allocate_slots( + self, pool_names: list[str], protected: set[BlockKey], deadline: float + ) -> tuple[list[tuple[str, int]] | None, list[BlockKey], bool]: + required = Counter(pool_names) + available = {name: len(self._free_slots[name]) for name in required} + if all(available[name] >= count for name, count in required.items()): + return ( + [(name, self._free_slots[name].popleft()) for name in pool_names], + [], + False, + ) if self._capacity_eviction_key is not None: - return None, None, True + return None, [], True self._retry_unscored() skipped = set(protected) + victims: list[tuple[BlockKey, "_BlockRecord", _LocalDramResidency, int]] = [] while (key := self._evictable.select(skipped)) is not None: record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None @@ -990,40 +1028,69 @@ def _allocate_slot( or residency.claim_count ): raise RuntimeError(f"invalid evictable local DRAM entry {key!r}") + if not any( + name in required and available[name] < required[name] + for name, _ in residency.slots + ): + skipped.add(key) + continue + size_bytes = self._size_bytes(residency.slots) + free_before = { + name: len(self._free_slots[name]) for name in required + } decision, eviction_pending = self._kvcr._decide_eviction( - self._kvcr._block_meta(key, record, self._slot_size), + self._kvcr._block_meta(key, record, size_bytes), CacheTier.LOCAL_G2, deadline, ) - if self._free_slots: - return self._free_slots.popleft(), None, False + for name in required: + available[name] += len(self._free_slots[name]) - free_before[name] + if all(available[name] >= count for name, count in required.items()): + break if eviction_pending: self._capacity_eviction_key = key - return None, None, True + return None, [], True if decision[0] is PlacementAction.KEEP: skipped.add(key) continue + victims.append((key, record, residency, size_bytes)) + skipped.add(key) + for name, _ in residency.slots: + if name in available: + available[name] += 1 + if all(available[name] >= count for name, count in required.items()): + break + else: + return None, [], False + + for key, record, residency, size_bytes in victims: self._evictable.remove(key) record.local_dram = None self._residency_observer(key, record) - self._kvcr._on_remove(self._kvcr._block_meta(key, record, self._slot_size)) + self._kvcr._on_remove(self._kvcr._block_meta(key, record, size_bytes)) self._kvcr._prune_block_record(key) - return residency.slot, key, False - return None, None, False + self._free(residency.slots) + return ( + [(name, self._free_slots[name].popleft()) for name in pool_names], + [victim[0] for victim in victims], + False, + ) def _make_evictable(self, key: BlockKey) -> None: record = self._kvcr._block_record_map.get(key) if record is None: raise RuntimeError(f"missing block record for {key!r}") score = self._kvcr._policy.eviction_score( - self._kvcr._block_meta(key, record, self._slot_size), + self._kvcr._block_meta( + key, record, self._size_bytes(record.local_dram.slots) + ), CacheTier.LOCAL_G2, ) if score is None: self._unscored.add(key) return self._unscored.discard(key) - self._evictable.insert(key, score) + self._evictable.insert(key, score, len(record.local_dram.slots)) def _remove_evictable(self, key: BlockKey) -> None: self._unscored.discard(key) @@ -1033,17 +1100,38 @@ def _retry_unscored(self) -> None: for key in tuple(self._unscored): self._make_evictable(key) - def _descriptor(self, slot: int) -> MemDescriptor: + def _descriptors( + self, locations: Collection[tuple[str, int]] + ) -> list[MemDescriptor]: + return [self._descriptor(pool_name, slot) for pool_name, slot in locations] + + def _descriptor(self, pool_name: str, slot: int) -> MemDescriptor: + address, _, slot_size = self._pools[pool_name] return MemDescriptor( end_point_name=self._kvcr.nixl_agent_name, mem_type="DRAM", - addr=self._address + slot * self._slot_size, - size=self._slot_size, + addr=address + slot * slot_size, + size=slot_size, device_Id=0, - info=self._kvcr.pool_layouts[0][0], + info=pool_name, ) + def _free(self, locations: Collection[tuple[str, int]]) -> None: + for pool_name, slot in locations: + self._free_slots[pool_name].append(slot) + + def _size_bytes(self, locations: Collection[tuple[str, int]]) -> int: + return sum(self._pools[pool_name][2] for pool_name, _ in locations) + + @staticmethod + def _same_layout( + locations: Collection[tuple[str, int]], descriptors: Collection[MemDescriptor] + ) -> bool: + return [name for name, _ in locations] == [ + descriptor.info for descriptor in descriptors + ] + def _update_capacity_pressure(self) -> None: self._kvcr._update_capacity_pressure( - len(self._free_slots) + len(self._evictable) + sum(map(len, self._free_slots.values())) + self._evictable.total_weight ) diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index cea2ce1..4a747a6 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -148,6 +148,7 @@ def on_remove(self, meta: BlockMeta) -> None: class _Entry: score: float sequence: int + weight: int class _EvictionQueue: @@ -155,18 +156,25 @@ def __init__(self) -> None: self._heap: list[tuple[float, int, BlockKey]] = [] self._live: dict[BlockKey, _Entry] = {} self._next_sequence = 0 + self.total_weight = 0 def __len__(self) -> int: return len(self._live) - def insert(self, key: BlockKey, score: float) -> None: - entry = _Entry(score, self._next_sequence) + def insert(self, key: BlockKey, score: float, weight: int = 1) -> None: + previous = self._live.get(key) + if previous is not None: + self.total_weight -= previous.weight + entry = _Entry(score, self._next_sequence, weight) self._next_sequence += 1 self._live[key] = entry + self.total_weight += weight heapq.heappush(self._heap, (entry.score, entry.sequence, key)) def remove(self, key: BlockKey) -> None: - self._live.pop(key, None) + entry = self._live.pop(key, None) + if entry is not None: + self.total_weight -= entry.weight def select(self, excluded: set[BlockKey]) -> BlockKey | None: skipped: list[tuple[float, int, BlockKey]] = [] @@ -174,7 +182,7 @@ def select(self, excluded: set[BlockKey]) -> BlockKey | None: while self._heap: score, sequence, key = self._heap[0] entry = self._live.get(key) - if entry != _Entry(score, sequence): + if entry is None or (entry.score, entry.sequence) != (score, sequence): heapq.heappop(self._heap) continue if key not in excluded: diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index a4de60c..d998f30 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -78,8 +78,8 @@ def store_release(self, value: int) -> None: # # Field order is the format. Append only -- never reorder or remove. class _RecoveryBlock(msgspec.Struct, frozen=True, array_like=True): - # A slot per tier, or nothing. Bare ints: wrapping one costs a byte each. - g2: Annotated[int, msgspec.Meta(ge=0)] | None = None + # Ordered pool locations, or nothing. Pool names may repeat. + g2: list[tuple[str, int]] | None = None g3: Annotated[int, msgspec.Meta(ge=0)] | None = None @@ -104,23 +104,40 @@ def _is_recoverable(record: _BlockRecord) -> bool: ) -def _project_recovery_record(record: _BlockRecord) -> _RecoveryBlock: +def _g2_locations( + slots: list[tuple[str, int]], pool_names: tuple[str, ...] +) -> list[tuple[str, int]]: + if type(slots) is not list or not slots: + raise ValueError("G2 recovery locations must be a non-empty list") + allowed = set(pool_names) + for location in slots: + if type(location) is not tuple or len(location) != 2: + raise ValueError("G2 recovery location must be a pool and slot pair") + pool_name, pool_slot = location + if pool_name not in allowed or type(pool_slot) is not int or pool_slot < 0: + raise ValueError("G2 recovery location does not match the pool group") + return slots + + +def _project_recovery_record( + record: _BlockRecord, pool_names: tuple[str, ...] +) -> _RecoveryBlock: + g3 = record.g3.slot if record.g3 is not None else None local_dram = record.local_dram - return _RecoveryBlock( - g2=( - local_dram.slot - if local_dram is not None and local_dram.state is _LocalDramState.READY - else None - ), - g3=record.g3.slot if record.g3 is not None else None, - ) + if local_dram is None or local_dram.state is not _LocalDramState.READY: + return _RecoveryBlock(g3=g3) + return _RecoveryBlock(g2=_g2_locations(local_dram.slots, pool_names), g3=g3) -def _decode_recovery_record(payload: bytes) -> _BlockRecord: +def _decode_recovery_record( + payload: bytes, pool_names: tuple[str, ...] +) -> _BlockRecord: recovered = _RECOVERY_DECODER.decode(payload) return _BlockRecord( local_dram=( - _LocalDramResidency(recovered.g2, _LocalDramState.READY) + _LocalDramResidency( + _g2_locations(recovered.g2, pool_names), _LocalDramState.READY + ) if recovered.g2 is not None else None ), @@ -309,7 +326,8 @@ def _read_ring(self, mapping: object, position: int, length: int) -> bytes: class _RecoveryMirror: - def __init__(self) -> None: + def __init__(self, pool_names: tuple[str, ...]) -> None: + self._pool_names = pool_names self._records: dict[BlockKey, _BlockRecord] = {} def apply(self, record_type: int, key: bytes, payload: bytes) -> None: @@ -317,7 +335,7 @@ def apply(self, record_type: int, key: bytes, payload: bytes) -> None: # where frames are published and read, not again here. del record_type try: - record = _decode_recovery_record(payload) + record = _decode_recovery_record(payload, self._pool_names) except (TypeError, ValueError, msgspec.DecodeError) as error: raise RecoveryMirrorError("recovery record is malformed") from error block_key = BlockKey(key) @@ -340,6 +358,7 @@ def adopt(self, records: dict[BlockKey, _BlockRecord]) -> None: if local_dram.state is not _LocalDramState.READY: record.local_dram = None else: + _g2_locations(local_dram.slots, self._pool_names) local_dram.claim_count = 0 local_dram.retire_on_release = False if record.g3 is not None: @@ -366,7 +385,10 @@ def take_records(self) -> dict[BlockKey, _BlockRecord]: def _attach_journal( - local_dram: _LocalDram, journal: RecoveryJournal, g3: _G3 | None = None + local_dram: _LocalDram, + journal: RecoveryJournal, + pool_names: tuple[str, ...], + g3: _G3 | None = None, ) -> None: """Attach stable G2/G3 residency publication to one journal.""" enabled = True @@ -389,7 +411,7 @@ def publish_frame(record_type: int, key: bytes, payload: bytes) -> None: def publish(key: BlockKey, record: _BlockRecord) -> None: # TODO: Publish per-tier deltas if full-record journal traffic is material. - recovered = _project_recovery_record(record) + recovered = _project_recovery_record(record, pool_names) publish_frame(_RECORD_BLOCK, bytes(key), _RECOVERY_ENCODER.encode(recovered)) if g3 is not None: @@ -495,7 +517,12 @@ def adopt_claimed_pool(core: _KVCRCore, claimed: ClaimedPool) -> None: hold = claimed.hold if core._local_dram is None: raise ValueError("a claimed pool must give the core its local DRAM tier") - _attach_journal(core._local_dram, RecoveryJournal(hold._attachment), core._g3) + _attach_journal( + core._local_dram, + RecoveryJournal(hold._attachment), + tuple(pool[0] for pool in hold.local_dram.pools), + core._g3, + ) install_recovery_records(core, claimed.recovered.take_records()) hold.hand_listener_to(claimed.adopt_listener) @@ -528,13 +555,13 @@ def _pack_frame(record_type: int, key: bytes, payload: bytes, size: int) -> byte def _recovery_frames( - records: Mapping[BlockKey, _BlockRecord], + records: Mapping[BlockKey, _BlockRecord], pool_names: tuple[str, ...] ) -> Iterator[tuple[int, bytes, bytes]]: """Every frame a returning primary needs to rebuild this state.""" for key, record in records.items(): if not _is_recoverable(record): continue - payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record)) + payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record, pool_names)) yield _RECORD_BLOCK, bytes(key), payload @@ -672,7 +699,8 @@ def read_handback( write never finished is this service's own, and is thrown away -- nothing else ever would, and it would refuse every later claim on this pool too. """ - mirror = _RecoveryMirror() + pool_names = tuple(name for name, _ in pool_layouts) + mirror = _RecoveryMirror(pool_names) terms = canonical_pool_terms(compatibility_digest, pool_layouts, pool._spec) try: for frame in read_recovery_snapshot(pool, terms): @@ -682,7 +710,7 @@ def read_handback( "KVCR discarding a handback region that was never finished", exc_info=True ) pool.release_snapshot_region() - return _RecoveryMirror() + return _RecoveryMirror(pool_names) return mirror diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index 2cd2255..9424030 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -43,12 +43,12 @@ from .core import _KVCRCore -_MEM_DESCRIPTORS_TYPE = tuple[MemDescriptor, ...] +_MEM_DESCRIPTOR_LISTS_TYPE = tuple[tuple[MemDescriptor, ...], ...] @dataclass(slots=True) class _FwMemResidency: - descriptor: MemDescriptor + descriptors: list[MemDescriptor] pin_handle: PinHandle @@ -92,7 +92,7 @@ class _TargetPullOp(_RemoteOp): remote_ctrl_ep: str _backend: "_RemoteFWDram" = field(repr=False, compare=False) ordered_keys: tuple[BlockKey, ...] = () - dst_descriptors: tuple[MemDescriptor, ...] = () + dst_descriptors: tuple[tuple[MemDescriptor, ...], ...] = () request_id: str | None = None success: bool = False completed_keys: set[BlockKey] = field(default_factory=set) @@ -202,7 +202,7 @@ class _SourcePinOp(_Op): remote_agent: bytes op_handle: int ordered_keys: tuple[BlockKey, ...] - dst_descriptors: tuple[MemDescriptor, ...] + dst_descriptors: tuple[tuple[MemDescriptor, ...], ...] route: tuple[str, int] = ("", 0) framework_pins: set[PinHandle] = field(default_factory=set) pending_pin_ids: set[PinRequestId] = field(default_factory=set) @@ -225,10 +225,10 @@ class _SourceWriteOp(_RemoteOp): remote_agent: bytes op_handle: int ordered_keys: tuple[BlockKey, ...] - dst_descriptors: tuple[MemDescriptor, ...] + dst_descriptors: tuple[tuple[MemDescriptor, ...], ...] _backend: "_RemoteFWDram" = field(repr=False, compare=False) framework_pins: set[PinHandle] = field(default_factory=set) - src_descriptors: tuple[MemDescriptor, ...] = () + src_descriptors: tuple[tuple[MemDescriptor, ...], ...] = () transfer_id: int | None = None success: bool = False completed_count: int = 0 @@ -269,8 +269,16 @@ def progress( try: transfer_id, submitted = progress.submit_transfer( "WRITE", - self.src_descriptors, - self.dst_descriptors[: self.completed_count], + tuple( + descriptor + for descriptors in self.src_descriptors + for descriptor in descriptors + ), + tuple( + descriptor + for descriptors in self.dst_descriptors[: self.completed_count] + for descriptor in descriptors + ), remote_side_agent=self.remote_agent, backend=backend._options.backend, notif_msg=_write_done_notif( @@ -471,7 +479,7 @@ def query(self, key: BlockKey, request_id: str) -> bool: def _start_target_pull( self, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], request_id: str | None, deadline: float, op_handle: OpHandle, @@ -506,7 +514,7 @@ def _start_target_pull( remote_ctrl_ep=current_hint.source, _backend=self, ordered_keys=keys, - dst_descriptors=tuple(blocks[key] for key in keys), + dst_descriptors=tuple(tuple(blocks[key]) for key in keys), request_id=request_id, ) kvcr._add_block_dependencies(op, new_operation=True) @@ -516,7 +524,7 @@ def _start_target_pull( def deliver( self, op_handle: OpHandle, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], request_id: str | None, *, deadline: float, @@ -540,7 +548,7 @@ def deliver( def fetch( self, - blocks: Mapping[BlockKey, MemDescriptor], + blocks: Mapping[BlockKey, list[MemDescriptor]], request_id: str | None, deadline: float, *, @@ -893,13 +901,13 @@ def _handle_start_write( raise TypeError("invalid remaining_timeout_ms") keys = _message_keys(payload) dst_descriptors = tuple( - self._kvcr._normalize_descriptors([descriptor]) - for descriptor in msgspec.convert( - payload["dst_descriptors"], type=_MEM_DESCRIPTORS_TYPE + tuple(self._kvcr._normalize_descriptors(list(descriptors))) + for descriptors in msgspec.convert( + payload["dst_descriptors"], type=_MEM_DESCRIPTOR_LISTS_TYPE ) ) - except (KeyError, TypeError, ValueError, msgspec.ValidationError): - logger.warning("KVCR malformed start_write op=%d", op_handle) + except (KeyError, TypeError, ValueError, msgspec.ValidationError) as error: + logger.warning("KVCR malformed start_write op=%d: %s", op_handle, error) self._notify_start_write_failure(progress, payload, op_handle) return if not keys or len(keys) != len(dst_descriptors): @@ -913,7 +921,7 @@ def _handle_start_write( ) deadline = received_at + remaining_timeout_ms / 1000 try: - fallback_target = dst_descriptors[0].end_point_name + fallback_target = dst_descriptors[0][0].end_point_name target_agent, remote_agent = self._remote_agent( progress, payload, fallback_target=fallback_target ) @@ -961,7 +969,7 @@ def _submit_prepared_source_write( source_pin.op_id, source_pin.ordered_keys ) framework_sources = { - key: record.fw_mem.descriptor + key: record.fw_mem.descriptors for key in source_pin.ordered_keys if key not in local_sources and (record := kvcr._block_record_map.get(key)) is not None @@ -970,7 +978,18 @@ def _submit_prepared_source_write( sources = {} if force_failure else {**framework_sources, **local_sources} completed_keys: tuple[BlockKey, ...] = () for index, key in enumerate(source_pin.ordered_keys): - if key not in sources: + source = sources.get(key) + destination = source_pin.dst_descriptors[index] + if source is None: + break + if [ + (descriptor.info, descriptor.size) for descriptor in source + ] != [(descriptor.info, descriptor.size) for descriptor in destination]: + logger.warning( + "KVCR start_write layout mismatch op=%d key=%r", + source_pin.op_handle, + key, + ) break completed_keys = source_pin.ordered_keys[: index + 1] @@ -1010,7 +1029,7 @@ def _submit_prepared_source_write( route=source_pin.route, _backend=self, framework_pins=framework_pins, - src_descriptors=tuple(sources[key] for key in completed_keys), + src_descriptors=tuple(tuple(sources[key]) for key in completed_keys), completed_count=len(completed_keys), ) kvcr._add_block_dependencies(source_write, new_operation=True) @@ -1332,7 +1351,7 @@ def _acquire_framework_sources( self, keys: tuple[BlockKey, ...], ) -> ( - tuple[dict[BlockKey, MemDescriptor], set[PinHandle]] + tuple[dict[BlockKey, list[MemDescriptor]], set[PinHandle]] | _PendingFrameworkSources | None ): @@ -1365,14 +1384,14 @@ def _acquire_framework_sources( framework_pins=held_framework_pins, ) - descriptors: dict[BlockKey, MemDescriptor] = {} + descriptors: dict[BlockKey, list[MemDescriptor]] = {} framework_pins: set[PinHandle] = set() for key in keys: record = kvcr._block_record_map.get(key) residency = record.fw_mem if record is not None else None if residency is None: continue - descriptors[key] = residency.descriptor + descriptors[key] = residency.descriptors framework_pins.add(residency.pin_handle) if not descriptors: return None diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index e29c444..d173e95 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -438,7 +438,7 @@ def _start_write_message( "remaining_timeout_ms": remaining_timeout_ms, "target_agent_metadata": b"target-md", "keys": [key], - "dst_descriptors": [_mem_descriptor().__dict__], + "dst_descriptors": [[_mem_descriptor().__dict__]], } if target_agent is not None: payload["target_agent"] = target_agent @@ -560,7 +560,9 @@ def decode(self, key): return 123 -def _recovered_record(*, g2: int | None = None, g3: int | None = None) -> _BlockRecord: +def _recovered_record( + *, g2: int | list[tuple[str, int]] | None = None, g3: int | None = None +) -> _BlockRecord: """A block record as recovery rebuilds one: settled residencies, nothing live.""" return _BlockRecord( local_dram=( diff --git a/tests/unit/test_g3.py b/tests/unit/test_g3.py index 2cae9ff..0ccd540 100644 --- a/tests/unit/test_g3.py +++ b/tests/unit/test_g3.py @@ -119,9 +119,12 @@ def transfer(self, handle): class _MoveLocalToG3Policy(FIFOPolicy): def __init__(self): self.move = True + self.keep_g3 = False self.failures = [] def decide_eviction(self, meta, source): + if self.keep_g3 and source is CacheTier.G3: + return (PlacementAction.KEEP, None) if self.move and source is CacheTier.LOCAL_G2: return (PlacementAction.MOVE_TO, CacheTier.G3) return super().decide_eviction(meta, source) @@ -774,6 +777,37 @@ def test_failed_g3_spill_recovers_by_dropping_source(tmp_path, caplog) -> None: assert metrics[("counter", TRANSFER_BLOCKS_METRIC, "g3_store")] == 1 +def test_full_g3_does_not_hide_a_synchronously_freed_local_slot(tmp_path) -> None: + page_size = os.sysconf("SC_PAGE_SIZE") + primary = ctypes.create_string_buffer(3 * page_size) + local = ctypes.create_string_buffer(page_size) + policy = _MoveLocalToG3Policy() + kvcr = _new_g3_kvcr(tmp_path, local, policy=policy, g3_slot_count=1) + first, second, third = (BlockKey(bytes((index,))) for index in range(3)) + + for index, key in enumerate((first, second)): + assert _deposit( + kvcr, + key, + ctypes.addressof(primary) + index * page_size, + page_size, + ).success + + policy.keep_g3 = True + assert _deposit( + kvcr, + third, + ctypes.addressof(primary) + 2 * page_size, + page_size, + ).success + + assert kvcr.query((first, second, third)) == [ + (QueryStatus.FETCHABLE, CacheTier.G3), + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + + def test_g3_spill_waits_until_local_source_claim_is_released(tmp_path) -> None: page_size = os.sysconf("SC_PAGE_SIZE") primary = ctypes.create_string_buffer(page_size * 2) diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 1bd4d0f..155aa99 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -86,13 +86,15 @@ def drain(self): def _frame(key: BlockKey, record: _BlockRecord) -> tuple[int, bytes, bytes]: """One journal frame, exactly as a primary would publish it.""" - payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record)) + payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record, ("",))) return (_RECORD_BLOCK, bytes(key), payload) def _give_serving_core(guard: _Guard) -> Mock: """A serving core still holding one READY G2 block.""" - record = _BlockRecord(local_dram=_LocalDramResidency(0, _LocalDramState.READY)) + record = _BlockRecord( + local_dram=_LocalDramResidency([("", 0)], _LocalDramState.READY) + ) core = Mock(_block_record_map={BlockKey(b"warm"): record}) guard._core = core guard._serving = True @@ -235,8 +237,8 @@ def test_guard_lives_out_adopt_promote_and_readopt_in_ownership_order( ) journal = _Journal( [ - _frame(first, _recovered_record(g2=0, g3=7)), - _frame(second, _recovered_record(g2=1)), + _frame(first, _recovered_record(g2=[("", 0)], g3=7)), + _frame(second, _recovered_record(g2=[("", 1)])), _frame(g3_only, _recovered_record(g3=9)), ] ) @@ -332,10 +334,10 @@ def new_channel() -> Mock: retained_g3 = guard._recovery._g3_records[first] records = { first: _BlockRecord( - local_dram=_LocalDramResidency(0, _LocalDramState.FILLING), + local_dram=_LocalDramResidency([("", 0)], _LocalDramState.FILLING), g3=_G3Residency(7), ), - second: _recovered_record(g2=1), + second: _recovered_record(g2=[("", 1)]), } cores[0]._block_record_map = records write_handback = Mock() @@ -346,7 +348,7 @@ def new_channel() -> Mock: assert guard._recovery.mirror._records[first].g3 is retained_g3 assert guard._recovery.mirror._records == { first: _recovered_record(g3=7), - second: _recovered_record(g2=1), + second: _recovered_record(g2=[("", 1)]), g3_only: _recovered_record(g3=9), } assert guard._recovery._g3_records == {} @@ -357,7 +359,7 @@ def new_channel() -> Mock: [ _frame(first, _BlockRecord()), _frame(g3_only, _BlockRecord()), - _frame(fresh, _recovered_record(g2=0, g3=7)), + _frame(fresh, _recovered_record(g2=[("", 0)], g3=7)), ] ) guard._promote() @@ -383,7 +385,7 @@ def test_a_pool_that_lost_its_recovery_stays_claimable_on_every_path( guard = _configurable_guard() # Every one of these readers runs on a pool a primary has already claimed. guard._configured = _TierConfig([("", 16)], None) - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) guard._recovery.attachment = Mock() guard._control = None journal = Mock() @@ -501,7 +503,7 @@ def test_recovery_close_error_stays_first_while_lease_cleanup_continues() -> Non owner = Mock() guard = _Guard(_TEST_SPEC, compatibility_digest=_TEST_DIGEST, owner=owner) guard._recovery.attachment = attachment - mirror = guard._recovery.mirror = _RecoveryMirror() + mirror = guard._recovery.mirror = _RecoveryMirror(("",)) g3_records = guard._recovery._g3_records = {BlockKey(b"g3"): _G3Residency(0)} guard._pool_lease.current = holder @@ -573,7 +575,7 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( # A hand-back that fails cannot be reported as a refused claim. guard._configured = _TierConfig([("", 16)], None) guard._serving = True - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) guard._core = Mock(_block_record_map={}) failure = OSError("no space left on device") guard._hand_back = Mock(side_effect=failure) @@ -590,7 +592,7 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( if refused_by == "geometry": expected: type[Exception] = ValueError tier_config = _TierConfig([("", _TEST_SPEC.mapping_bytes)], None) - handback = Mock(return_value=_RecoveryMirror()) + handback = Mock(return_value=_RecoveryMirror(("",))) else: expected = RecoveryJournalError tier_config = _TierConfig([("", 16)], None) @@ -614,7 +616,7 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( def test_a_handback_with_an_unexpected_storage_error_fails() -> None: """Only capacity errors (ENOSPC/EDQUOT) are survivable at the handback writer.""" guard = _configurable_guard() - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) _give_serving_core(guard) error = OSError(errno.EIO, "I/O error") guard._recovery._write_handback = Mock(side_effect=error) @@ -639,7 +641,7 @@ def test_a_handback_without_a_mirror_does_not_close_the_core() -> None: def test_a_handback_the_filesystem_refuses_leaves_a_cold_pool() -> None: """ENOSPC at the pool tail drops the mirror with the handback it refused.""" guard = _configurable_guard() - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) _give_serving_core(guard) guard._recovery._write_handback = Mock( side_effect=OSError(errno.ENOSPC, "No space left on device") @@ -659,7 +661,7 @@ def test_a_dropped_handback_still_leaves_the_new_lease_mirrored(code: int) -> No guard._control = None guard._failure_callback = lambda *_args: None guard._configured = _TierConfig([("", 16)], None) - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) _give_serving_core(guard) guard._recovery._journal = _Journal() guard._recovery._write_handback = Mock(side_effect=OSError(code, "No space left")) @@ -675,7 +677,11 @@ def test_a_dropped_handback_still_leaves_the_new_lease_mirrored(code: int) -> No # And the grant is retractable: the Guard it stood down can resume. assert guard._resumable is True guard._recovery._journal.pending = [ - (_RECORD_BLOCK, b"fresh", _RECOVERY_ENCODER.encode(_RecoveryBlock(g2=1))) + ( + _RECORD_BLOCK, + b"fresh", + _RECOVERY_ENCODER.encode(_RecoveryBlock(g2=[("", 1)])), + ) ] guard._poll() assert BlockKey(b"fresh") in guard._recovery.mirror._records @@ -685,7 +691,7 @@ def test_a_grant_that_never_arrived_resumes_the_guard_it_stood_down() -> None: """An aborted grant re-promotes after a hand-back; otherwise it releases.""" guard = _configurable_guard() guard._resumable = True - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) outcomes: list[str] = [] guard._promote = lambda: outcomes.append("promote") guard._release = lambda: outcomes.append("release") @@ -715,14 +721,14 @@ def test_a_release_drops_its_mirror_after_handing_back_what_it_can( control = Mock() guard._control = control guard._configured = _TierConfig([("", 16)], None) - guard._recovery.mirror = _RecoveryMirror() + guard._recovery.mirror = _RecoveryMirror(("",)) if mode == "serving": _give_serving_core(guard) else: guard._recovery.mirror.apply( - *_frame(BlockKey(b"published"), _recovered_record(g2=0)) + *_frame(BlockKey(b"published"), _recovered_record(g2=[("", 0)])) ) - tail = _frame(BlockKey(b"tail"), _recovered_record(g2=1)) + tail = _frame(BlockKey(b"tail"), _recovered_record(g2=[("", 1)])) guard._recovery._journal = _Journal(pending=[tail]) guard._recovery._write_handback = Mock( side_effect=OSError(errno.ENOSPC, "No space left") diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index 8a5338f..7968830 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -65,7 +65,7 @@ def test_local_dram_observer_reports_only_stable_slot_changes() -> None: backend = kvcr._core._local_dram assert backend is not None keys = tuple(BlockKey(f"k{index}".encode()) for index in range(3)) - observed: list[tuple[BlockKey, int | None]] = [] + observed: list[tuple[BlockKey, list[tuple[str, int]] | None]] = [] def observe(key: BlockKey, record: _BlockRecord) -> None: residency = record.local_dram @@ -74,14 +74,14 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: else: assert residency.state is _LocalDramState.READY assert local.raw == bytes((ord("a") + keys.index(key),)) * block_size - observed.append((key, None if residency is None else residency.slot)) + observed.append((key, None if residency is None else residency.slots)) backend.observe_residency(observe) address = ctypes.addressof(primary) first = kvcr.deposit({keys[0]: [_mem_descriptor(address, block_size)]}) _poll_until(kvcr, lambda done: first in dict(done)) - assert observed == [(keys[0], 0)] + assert observed == [(keys[0], [("", 0)])] agent.state = "ERR" failed = kvcr.deposit( @@ -90,7 +90,7 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: failed_result = dict(_poll_until(kvcr, lambda done: failed in dict(done)))[failed] assert not failed_result[keys[1]].success assert observed == [ - (keys[0], 0), + (keys[0], [("", 0)]), (keys[0], None), ] @@ -104,7 +104,7 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: assert len(observed) == 3 backend.release_sources((keys[2],)) assert observed[2:] == [ - (keys[2], 0), + (keys[2], [("", 0)]), (keys[2], None), ] @@ -338,8 +338,13 @@ def make_journal(pool) -> object: events.append("journal") return journal - def attach_journal(local, configured_journal, disk) -> None: - assert (local, configured_journal, disk) == (local_dram, journal, g3) + def attach_journal(local, configured_journal, pool_names, disk) -> None: + assert (local, configured_journal, pool_names, disk) == ( + local_dram, + journal, + ("",), + g3, + ) events.append("attach") monkeypatch.setattr( @@ -431,18 +436,6 @@ def test_kvcr_rejects_no_dram_backends() -> None: ) -def test_kvcr_rejects_multi_pool_layouts() -> None: - with pytest.raises(ValueError, match="only a single pool"): - KVCR( - KVCRConfig( - nixl_agent_name="target", - pool_layouts=[("full", 8), ("swa", 8)], - ), - KVCRBindings(Mock(), Mock(), Mock()), - KVCRBackendConfigs(), - ) - - def test_kvcr_rejects_ambiguous_pool_names() -> None: bindings = KVCRBindings(Mock(), Mock(), Mock()) for pool_layouts, message in ( @@ -664,7 +657,7 @@ def test_resident_records_carry_no_instance_dictionary() -> None: """Every record a resident block can hold, so none of them grows one back.""" for residency in ( _BlockRecord(), - _LocalDramResidency(0, _LocalDramState.READY), + _LocalDramResidency([("", 0)], _LocalDramState.READY), _G3Residency(0), _FwMemResidency(_mem_descriptor(), object()), ): diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 530c3be..37ce8c4 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -8,8 +8,11 @@ import pytest from _kvcr_test_utils import ( + FakeBytesControl, FakeNixlAgent, + FakePrimaryPinning, _mem_descriptor, + _new_kvcr, _new_local_kvcr, _op_entries, _poll_until, @@ -17,6 +20,7 @@ _wait_until, ) +from kvcr.config import KVCRConfig, LocalDramOptions from kvcr.core import _BlockRecord from kvcr.local_dram import _LocalDramResidency, _LocalDramState from kvcr.policy import FIFOPolicy, LRUPolicy @@ -107,11 +111,155 @@ def test_local_deposit_deduplicates_and_evicts_fifo() -> None: ] -def test_local_transfer_rejects_multiple_descriptors() -> None: - kvcr = _new_local_kvcr(FakeNixlAgent(), ctypes.create_string_buffer(16), 1) +def test_local_transfer_accepts_multiple_blocks_in_one_pool() -> None: + agent = FakeNixlAgent() + agent.state = "DONE" + source = ctypes.create_string_buffer(32) + kvcr = _new_local_kvcr(agent, ctypes.create_string_buffer(32), 2) + + operation = kvcr.deposit( + { + BlockKey(b"key"): [ + _mem_descriptor(ctypes.addressof(source)), + _mem_descriptor(ctypes.addressof(source) + 16), + ] + } + ) + + assert dict(_poll_until(kvcr, lambda results: bool(results)))[operation][ + BlockKey(b"key") + ].success + + +def test_multi_pool_residency_moves_and_evicts_as_one_key() -> None: + full = ctypes.create_string_buffer(16) + swa = ctypes.create_string_buffer(16) + source = ctypes.create_string_buffer(32) + agent = FakeNixlAgent() + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig( + nixl_agent_name="target", + pool_layouts=[("full", 16), ("swa", 8)], + ), + local_dram=LocalDramOptions( + [ + ("full", ctypes.addressof(full), 16), + ("swa", ctypes.addressof(swa), 16), + ] + ), + ) + descriptors = [ + _mem_descriptor(ctypes.addressof(source), 16, info="full"), + _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"), + _mem_descriptor(ctypes.addressof(source) + 24, 8, info="swa"), + ] + first, second = BlockKey(b"first"), BlockKey(b"second") + + operation = kvcr.deposit({first: descriptors}) + _wait_until(lambda: bool(agent.transfers)) + wrong_layout = kvcr.fetch((first,), expected_layout=["swa"]) + assert ( + dict(kvcr.poll_completed())[wrong_layout][first].status is OpEntryStatus.FAILED + ) + agent.state = "DONE" + _poll_until(kvcr, lambda done: operation in dict(done)) + claim = kvcr.fetch((first,), expected_layout=["full", "swa", "swa"]) + result = dict(_poll_until(kvcr, lambda done: claim in dict(done)))[claim][first] + assert [(item.info, item.addr) for item in result.descriptors or ()] == [ + ("full", ctypes.addressof(full)), + ("swa", ctypes.addressof(swa)), + ("swa", ctypes.addressof(swa) + 8), + ] + kvcr.release([result.release_handle]) + assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 3 + + operation = kvcr.deposit({second: descriptors}) + _poll_until(kvcr, lambda done: operation in dict(done)) + assert kvcr.query((first, second)) == [ + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + - with pytest.raises(ValueError, match="exactly one descriptor"): - kvcr.deposit({BlockKey(b"key"): [_mem_descriptor(), _mem_descriptor()]}) +def test_failed_group_reservation_does_not_evict_a_partial_group() -> None: + pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(8)] + source = ctypes.create_string_buffer(16) + agent = FakeNixlAgent() + agent.state = "DONE" + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig(nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)]), + local_dram=LocalDramOptions( + [ + ("full", ctypes.addressof(pools[0]), 8), + ("swa", ctypes.addressof(pools[1]), 8), + ] + ), + ) + full, swa, grouped = (BlockKey(name) for name in (b"full", b"swa", b"grouped")) + descriptors = [ + _mem_descriptor(ctypes.addressof(source), 8, info="full"), + _mem_descriptor(ctypes.addressof(source) + 8, 8, info="swa"), + ] + + for key, descriptor in ((full, descriptors[0]), (swa, descriptors[1])): + operation = kvcr.deposit({key: [descriptor]}, no_evict=key == swa) + _poll_until(kvcr, lambda done: operation in dict(done)) + + operation = kvcr.deposit({grouped: descriptors}) + result = dict(_poll_until(kvcr, lambda done: operation in dict(done)))[operation] + assert result[grouped].status is OpEntryStatus.FAILED + assert kvcr.query((full, swa)) == [ + (QueryStatus.HIT, CacheTier.LOCAL_G2), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + + +def test_group_allocation_evicts_enough_whole_keys() -> None: + pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(16)] + source = ctypes.create_string_buffer(24) + agent = FakeNixlAgent() + agent.state = "DONE" + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig(nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)]), + local_dram=LocalDramOptions( + [ + ("full", ctypes.addressof(pools[0]), 8), + ("swa", ctypes.addressof(pools[1]), 16), + ] + ), + ) + full, swa0, swa1, grouped = ( + BlockKey(name) for name in (b"full", b"swa0", b"swa1", b"grouped") + ) + descriptors = [ + _mem_descriptor(ctypes.addressof(source), 8, info="full"), + _mem_descriptor(ctypes.addressof(source) + 8, 8, info="swa"), + _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"), + ] + + for key, descriptor in zip((full, swa0, swa1), descriptors, strict=True): + operation = kvcr.deposit({key: [descriptor]}) + _poll_until(kvcr, lambda done: operation in dict(done)) + operation = kvcr.deposit({grouped: descriptors}) + result = dict(_poll_until(kvcr, lambda done: operation in dict(done)))[operation] + + assert result[grouped].success + assert kvcr.query((full, swa0, swa1, grouped)) == [ + (QueryStatus.MISS, None), + (QueryStatus.MISS, None), + (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), + ] + assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 3 @pytest.mark.parametrize( @@ -432,7 +580,7 @@ def initialize_xfer(self, *args, **kwargs): def _g2_recovered(**slots: int) -> dict[BlockKey, _BlockRecord]: return { BlockKey(name.encode()): _BlockRecord( - local_dram=_LocalDramResidency(slot, _LocalDramState.READY) + local_dram=_LocalDramResidency([("", slot)], _LocalDramState.READY) ) for name, slot in slots.items() } @@ -473,7 +621,7 @@ def test_a_recovered_pool_deposits_into_free_rows_then_evicts_to_admit_more() -> # Every row is now occupied, so the deposit below can only land by evicting # a recovered row -- a pool recovered full has to stay writable. - assert not local_dram._free_slots + assert not local_dram._free_slots[""] extra = BlockKey(b"extra") operation = kvcr.deposit( {extra: [_mem_descriptor(ctypes.addressof(primary) + 2 * block_size)]} @@ -498,7 +646,7 @@ def test_installing_records_into_a_core_that_holds_some_is_refused() -> None: local = ctypes.create_string_buffer(64) kvcr = _new_local_kvcr(FakeNixlAgent(), local, 4) kvcr._core._block_record_map[BlockKey(b"held")] = _BlockRecord( - local_dram=_LocalDramResidency(1, _LocalDramState.READY) + local_dram=_LocalDramResidency([("", 1)], _LocalDramState.READY) ) with pytest.raises(RecoveryMirrorError, match="holds none"): @@ -512,7 +660,7 @@ def test_installing_records_into_a_core_that_holds_some_is_refused() -> None: _g2_recovered(first=0, second=4), { BlockKey(b"first"): _BlockRecord( - local_dram=_LocalDramResidency(0, _LocalDramState.FILLING) + local_dram=_LocalDramResidency([("", 0)], _LocalDramState.FILLING) ) }, ], diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index f69f3de..18eb767 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -134,6 +134,7 @@ def test_kvcr_malformed_start_write_notifies_failure(kvcr_caplog): } assert any( "malformed start_write" in record.getMessage() + and "remaining_timeout_ms" in record.getMessage() for record in kvcr_caplog.records if record.levelno == logging.WARNING ) @@ -442,7 +443,7 @@ def test_pending_pin_waiters_share_partial_results_and_request_uncovered_keys( "target_agent_metadata": b"target-md", "keys": list(op_keys), "dst_descriptors": [ - _mem_descriptor(addr=128 + index * 16).__dict__ + [_mem_descriptor(addr=128 + index * 16).__dict__] for index in range(len(op_keys)) ], } @@ -582,8 +583,10 @@ def test_a_resumed_write_holds_a_pin_another_operation_acquired() -> None: key = BlockKey(b"shared") borrowed = PinHandle("pinned-by-the-other-operation") + sources = [_mem_descriptor(info="full"), _mem_descriptor(info="swa")] + destinations = tuple(_mem_descriptor(info=item.info) for item in sources) kvcr._block_record_map[key] = _BlockRecord( - fw_mem=_FwMemResidency(_mem_descriptor(), borrowed) + fw_mem=_FwMemResidency(sources, borrowed) ) # This operation acquired a pin of its own for a key it no longer needs. @@ -594,7 +597,7 @@ def test_a_resumed_write_holds_a_pin_another_operation_acquired() -> None: remote_agent=b"peer", op_handle=1, ordered_keys=(key,), - dst_descriptors=(_mem_descriptor(),), + dst_descriptors=(destinations,), op_id=("source", 1), keys={key}, framework_pins={stale}, @@ -604,6 +607,8 @@ def test_a_resumed_write_holds_a_pin_another_operation_acquired() -> None: backend._submit_prepared_source_write(("source", 1), waiting) submitted = kvcr._progress.submit.call_args.args[0] + assert submitted.src_descriptors == (tuple(sources),) + assert submitted.dst_descriptors == (destinations,) assert borrowed in submitted.framework_pins, ( "the resumed write reads through this pin but does not hold it" ) diff --git a/tests/unit/test_recovery_journal.py b/tests/unit/test_recovery_journal.py index 1087b98..76acd9e 100644 --- a/tests/unit/test_recovery_journal.py +++ b/tests/unit/test_recovery_journal.py @@ -138,18 +138,20 @@ def test_publisher_streams_mutations_until_the_journal_refuses_or_fails( journal, _ = journal_and_mapping local_dram, g3 = _Source(), _Source() key = BlockKey(b"block") - _attach_journal(local_dram, journal, g3) + _attach_journal(local_dram, journal, ("pool0",), g3) caplog.set_level("WARNING", logger="kvcr.recovery_journal") - local_dram.emit(key, _recovered_record(g2=2)) - g3.emit(key, _recovered_record(g2=2, g3=7)) + local_dram.emit(key, _recovered_record(g2=[("pool0", 2)])) + g3.emit(key, _recovered_record(g2=[("pool0", 2)], g3=7)) local_dram.emit(key, _recovered_record(g3=7)) g3.emit(key, _BlockRecord()) frames = [journal.read_next() for _ in range(4)] assert journal.read_next() is None - assert [_decode_recovery_record(payload) for _, _, payload in frames] == [ - _recovered_record(g2=2), - _recovered_record(g2=2, g3=7), + assert [ + _decode_recovery_record(payload, ("pool0",)) for _, _, payload in frames + ] == [ + _recovered_record(g2=[("pool0", 2)]), + _recovered_record(g2=[("pool0", 2)], g3=7), _recovered_record(g3=7), _BlockRecord(), ] @@ -160,8 +162,8 @@ def test_publisher_streams_mutations_until_the_journal_refuses_or_fails( journal.invalidate() caplog.clear() with patch.object(journal, "publish", wraps=journal.publish) as publish: - local_dram.emit(key, _recovered_record(g2=0)) - local_dram.emit(key, _recovered_record(g2=1)) + local_dram.emit(key, _recovered_record(g2=[("pool0", 0)])) + local_dram.emit(key, _recovered_record(g2=[("pool0", 1)])) assert publish.call_count == 1 assert len(caplog.messages) == 1 @@ -170,10 +172,13 @@ def test_publisher_streams_mutations_until_the_journal_refuses_or_fails( fresh = RecoveryJournal(_attachment(fresh_mapping, _TEST_JOURNAL_BYTES)) fresh.reset() source = _Source() - _attach_journal(source, fresh) + _attach_journal(source, fresh, ("pool0",)) assert not fresh.is_invalid() with patch.object(fresh, "publish", side_effect=RuntimeError("publish failed")): - source.emit(BlockKey(b"still-serving"), _recovered_record(g2=0)) + source.emit( + BlockKey(b"still-serving"), + _recovered_record(g2=[("pool0", 0)]), + ) assert fresh.is_invalid() @@ -271,7 +276,10 @@ def _attached(tmp_path: Path) -> Iterator[KVCRPoolAttachment]: def _write_slot(pool: KVCRPoolAttachment, terms: bytes, key: bytes, slot: int) -> None: """One-slot handback region: the smallest finished snapshot.""" - frames = _recovery_frames({BlockKey(key * 32): _recovered_record(g2=slot)}) + frames = _recovery_frames( + {BlockKey(key * 32): _recovered_record(g2=[("pool0", slot)])}, + ("pool0",), + ) write_recovery_snapshot(pool, terms, frames) @@ -279,30 +287,31 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - """Replayed whole under its own terms, discardable when torn, gone once released.""" with _attached(tmp_path) as pool: path = Path(pool._spec.path) - terms = canonical_pool_terms(_TEST_DIGEST, [("", 4096)], pool._spec) + pool_layouts = [("pool0", 4096)] + terms = canonical_pool_terms(_TEST_DIGEST, pool_layouts, pool._spec) assert list(read_recovery_snapshot(pool, terms)) == [] records = { - BlockKey(b"a" * 32): _recovered_record(g2=3), + BlockKey(b"a" * 32): _recovered_record(g2=[("pool0", 3)]), # One with both halves, one only on disk. - BlockKey(b"b" * 32): _recovered_record(g2=4, g3=9), + BlockKey(b"b" * 32): _recovered_record(g2=[("pool0", 4)], g3=9), BlockKey(b"c" * 32): _recovered_record(g3=2), } - write_recovery_snapshot(pool, terms, _recovery_frames(records)) + write_recovery_snapshot(pool, terms, _recovery_frames(records, ("pool0",))) # Inside the pool file, so it has no name of its own to be found under. assert set(tmp_path.iterdir()) == {path} assert path.stat().st_size > pool._spec.mapping_bytes # The mirror the ring feeds is also what replays the region. - mirror = _RecoveryMirror() + mirror = _RecoveryMirror(("pool0",)) for frame in read_recovery_snapshot(pool, terms): mirror.apply(*frame) assert mirror.take_records() == records # A slot number only means the same bytes under the same geometry. for other in ( - canonical_pool_terms("another-digest", [("", 4096)], pool._spec), - canonical_pool_terms(_TEST_DIGEST, [("", 8192)], pool._spec), + canonical_pool_terms("another-digest", pool_layouts, pool._spec), + canonical_pool_terms(_TEST_DIGEST, [("pool0", 8192)], pool._spec), ): with pytest.raises(RecoveryJournalError, match="other terms"): list(read_recovery_snapshot(pool, other)) @@ -328,7 +337,7 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - region[: _SNAPSHOT_HEADER.size] = bytes(_SNAPSHOT_HEADER.size) with pytest.raises(RecoveryJournalTornError, match="unfinished"): list(read_recovery_snapshot(pool, terms)) - assert read_handback(pool, _TEST_DIGEST, [("", 4096)])._records == {} + assert read_handback(pool, _TEST_DIGEST, pool_layouts)._records == {} assert list(read_recovery_snapshot(pool, terms)) == [] # A released region is truncated away, so it replays nothing. diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index cf98375..9c90011 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -19,16 +19,19 @@ ) from kvcr.types import BlockKey +_ONE_POOL = ("",) +_TWO_POOLS = ("full", "swa") -def _payload(record: _BlockRecord) -> bytes: - return _RECOVERY_ENCODER.encode(_project_recovery_record(record)) + +def _payload(record: _BlockRecord, pool_names: tuple[str, ...] = _ONE_POOL) -> bytes: + return _RECOVERY_ENCODER.encode(_project_recovery_record(record, pool_names)) # Every live-only field set, to prove projection strips all of it. _FULLY_LOADED_RECORD = _BlockRecord( fw_mem=object(), local_dram=_LocalDramResidency( - 3, + [("", 3)], _LocalDramState.READY, claim_count=2, retire_on_release=True, @@ -41,21 +44,42 @@ def _payload(record: _BlockRecord) -> bytes: @pytest.mark.parametrize( - ("record", "wire", "recovered"), + ("record", "pool_names", "wire", "recovered"), [ - (_FULLY_LOADED_RECORD, [3, 5], _recovered_record(g2=3, g3=5)), - # An absent tier still occupies its slot, because position is the name. - (_BlockRecord(), [None, None], _BlockRecord()), - (_recovered_record(g2=3), [3, None], _recovered_record(g2=3)), - (_recovered_record(g3=5), [None, 5], _recovered_record(g3=5)), + ( + _FULLY_LOADED_RECORD, + _ONE_POOL, + [[["", 3]], 5], + _recovered_record(g2=[("", 3)], g3=5), + ), + ( + _recovered_record(g2=[("full", 7), ("full", 2), ("swa", 9)], g3=5), + _TWO_POOLS, + [[["full", 7], ["full", 2], ["swa", 9]], 5], + _recovered_record(g2=[("full", 7), ("full", 2), ("swa", 9)], g3=5), + ), + (_BlockRecord(), _ONE_POOL, [None, None], _BlockRecord()), + ( + _recovered_record(g2=[("", 3)]), + _ONE_POOL, + [[["", 3]], None], + _recovered_record(g2=[("", 3)]), + ), + (_recovered_record(g3=5), _ONE_POOL, [None, 5], _recovered_record(g3=5)), # A G2 slot still FILLING or DISCARDING never settled, so it must not wire. ( - _BlockRecord(local_dram=_LocalDramResidency(0, _LocalDramState.FILLING)), + _BlockRecord( + local_dram=_LocalDramResidency([("", 0)], _LocalDramState.FILLING) + ), + _ONE_POOL, [None, None], _BlockRecord(), ), ( - _BlockRecord(local_dram=_LocalDramResidency(0, _LocalDramState.DISCARDING)), + _BlockRecord( + local_dram=_LocalDramResidency([("", 0)], _LocalDramState.DISCARDING) + ), + _ONE_POOL, [None, None], _BlockRecord(), ), @@ -63,32 +87,30 @@ def _payload(record: _BlockRecord) -> bytes: ) def test_recovery_wire_round_trip_keeps_only_settled_slots( record: _BlockRecord, + pool_names: tuple[str, ...], wire: list[object], recovered: _BlockRecord, ) -> None: """Only settled G2/G3 slots reach the wire; decode rebuilds fresh live state.""" - encoded = _payload(record) + encoded = _payload(record, pool_names) - # Positional, so no field names ride along in every record. + # The outer record stays positional; G2 locations carry their pool names. assert msgspec.msgpack.decode(encoded) == wire - assert len(encoded) == 3 - assert _decode_recovery_record(encoded) == recovered + assert _decode_recovery_record(encoded, pool_names) == recovered def test_recovery_encoding_accepts_a_field_appended_later() -> None: """Appending is the one change this format allows, and it has to work.""" class _RecoveryBlockV2(msgspec.Struct, frozen=True, array_like=True): - g2: int | None = None + g2: list[tuple[str, int]] | None = None g3: int | None = None appended: int = 0 - today = _RECOVERY_ENCODER.encode( - _project_recovery_record(_recovered_record(g2=3, g3=5)) - ) + today = _payload(_recovered_record(g2=[("", 3)], g3=5)) upgraded = msgspec.msgpack.Decoder(_RecoveryBlockV2).decode(today) - assert upgraded.g2 == 3 + assert upgraded.g2 == [("", 3)] assert upgraded.g3 == 5 assert upgraded.appended == 0 @@ -101,25 +123,32 @@ class _RecoveryBlockV2(msgspec.Struct, frozen=True, array_like=True): msgspec.msgpack.encode({"g2": {"slot": 0, "state": "ready"}}), msgspec.msgpack.encode({"g3": {"slot": -1}}), msgspec.msgpack.encode({"g2": {"slot": "0"}}), + msgspec.msgpack.encode([[["other", 0]], None]), + msgspec.msgpack.encode([[["", -1]], None]), + msgspec.msgpack.encode([0, None]), + msgspec.msgpack.encode([[[""]], None]), + msgspec.msgpack.encode([[], None]), ], ) def test_mirror_rejects_malformed_or_unknown_wire_state(payload: bytes) -> None: """A frame that does not decode to valid wire state is refused, not applied.""" - mirror = _RecoveryMirror() + mirror = _RecoveryMirror(_ONE_POOL) with pytest.raises(RecoveryMirrorError, match="malformed"): mirror.apply(_RECORD_BLOCK, b"block", payload) + assert mirror._records == {} + def test_mirror_replaces_blocks_whole_and_hands_them_over_uncopied() -> None: """Frames replace blocks whole in _records (mirrored table); take transfers it.""" - mirror = _RecoveryMirror() + mirror = _RecoveryMirror(_ONE_POOL) key = BlockKey(b"spilled") - mirror.apply(_RECORD_BLOCK, key, _payload(_recovered_record(g2=1))) - mirror.apply(_RECORD_BLOCK, key, _payload(_recovered_record(g2=1, g3=7))) + mirror.apply(_RECORD_BLOCK, key, _payload(_recovered_record(g2=[("", 1)]))) + mirror.apply(_RECORD_BLOCK, key, _payload(_recovered_record(g2=[("", 1)], g3=7))) - assert mirror._records == {key: _recovered_record(g2=1, g3=7)} + assert mirror._records == {key: _recovered_record(g2=[("", 1)], g3=7)} mirror.apply(_RECORD_BLOCK, key, _payload(_recovered_record(g3=7))) @@ -130,14 +159,14 @@ def test_mirror_replaces_blocks_whole_and_hands_them_over_uncopied() -> None: assert mirror._records == {} - mirror.apply(_RECORD_BLOCK, b"resident", _payload(_recovered_record(g2=1))) + mirror.apply(_RECORD_BLOCK, b"resident", _payload(_recovered_record(g2=[("", 1)]))) held = mirror._records taken = mirror.take_records() # Sole ownership: copying would leave two live populations of the set. assert taken is held - assert taken == {BlockKey(b"resident"): _recovered_record(g2=1)} + assert taken == {BlockKey(b"resident"): _recovered_record(g2=[("", 1)])} assert mirror._records == {} @@ -150,7 +179,10 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: served = { ready: _BlockRecord( local_dram=_LocalDramResidency( - 0, _LocalDramState.READY, claim_count=1, retire_on_release=True + [("full", 0), ("swa", 10)], + _LocalDramState.READY, + claim_count=1, + retire_on_release=True, ), in_flight_ops={("target", 7)}, access_count=12, @@ -158,33 +190,37 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: ), spilled: _BlockRecord(g3=_G3Residency(3, claim_count=2)), filling: _BlockRecord( - local_dram=_LocalDramResidency(1, _LocalDramState.FILLING) + local_dram=_LocalDramResidency( + [("full", 1), ("swa", 11)], _LocalDramState.FILLING + ) ), forgotten: _BlockRecord(), # A good G3 residency must not carry a half-written G2 slot with it. filling_spill: _BlockRecord( - local_dram=_LocalDramResidency(7, _LocalDramState.FILLING), + local_dram=_LocalDramResidency( + [("full", 7), ("swa", 17)], _LocalDramState.FILLING + ), g3=_G3Residency(4), ), discarding_spill: _BlockRecord( - local_dram=_LocalDramResidency(8, _LocalDramState.DISCARDING), + local_dram=_LocalDramResidency( + [("full", 8), ("swa", 18)], _LocalDramState.DISCARDING + ), g3=_G3Residency(5), ), } # A kept mirror must match exactly what the handback frames carry. framed = { - BlockKey(key): _decode_recovery_record(payload) - for _, key, payload in _recovery_frames(served) + BlockKey(key): _decode_recovery_record(payload, _TWO_POOLS) + for _, key, payload in _recovery_frames(served, _TWO_POOLS) } - assert set(framed) == {ready, spilled, filling_spill, discarding_spill} - - mirror = _RecoveryMirror() + mirror = _RecoveryMirror(_TWO_POOLS) mirror.adopt(served) assert mirror._records is served assert mirror._records == framed assert mirror._records == { - ready: _recovered_record(g2=0), + ready: _recovered_record(g2=[("full", 0), ("swa", 10)]), spilled: _recovered_record(g3=3), filling_spill: _recovered_record(g3=4), discarding_spill: _recovered_record(g3=5), From 2a0fa689eb7304a13ad5a9fc50939312853852be Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Wed, 9 Sep 2026 06:41:58 +0300 Subject: [PATCH 02/16] feat(guard)!: manage multiple pools BREAKING CHANGE: Guard claims and grants describe one jointly owned group of pool regions. Signed-off-by: Kapil Arya --- docs/dev-guide.md | 123 ++++++++++-------- src/kvcr/guard.py | 159 +++++++++++------------ src/kvcr/guard_protocol.py | 95 +++++++++----- src/kvcr/kvcr_service.py | 61 +++++---- src/kvcr/local_disk.py | 18 +-- src/kvcr/recovery_journal.py | 56 +++++--- tests/unit/test_guard.py | 144 +++++++++++--------- tests/unit/test_guard_integration.py | 108 ++++++++++++++- tests/unit/test_guard_protocol.py | 154 +++++++++++++++------- tests/unit/test_kvcr.py | 51 ++++++-- tests/unit/test_kvcr_service.py | 90 ++++++++----- tests/unit/test_kvcr_service_workflow.py | 47 +++---- tests/unit/test_recovery_journal.py | 53 ++++++-- 13 files changed, 743 insertions(+), 416 deletions(-) diff --git a/docs/dev-guide.md b/docs/dev-guide.md index d82b7af..fb6fcd9 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -316,10 +316,11 @@ IDs, or raw endpoints. ### KVCR service daemon -The KVCR service daemon owns pool lifecycle. It pre-allocates `--guard-count` -contiguous pool allocations before exposing its socket, one per Guard. Each has -the same ordered pool sizes. A worker claims a Guard by index; its allocation -outlives that worker but not the service: +The KVCR service daemon owns pool lifecycle. It pre-allocates +`--guard-count` Guard-owned pool groups before exposing its socket. Every +group has the same ordered set of usable pool sizes from `--pool-sizes-gb`. +A worker claims a whole group by Guard index; its pools outlive that worker +but not the service: ```bash python -m kvcr.kvcr_service \ @@ -334,61 +335,72 @@ python -m kvcr.kvcr_service \ | --- | --- | --- | | `--socket-path` | *(required)* | Unix socket the workers connect to | | `--pool-dir` | *(required)* | Writable directory holding the pool files | -| `--guard-count` | *(required)* | Number of Guards available by index | -| `--pool-sizes-gb` | *(required)* | Comma-separated usable pool sizes in each Guard allocation | +| `--guard-count` | *(required)* | Number of Guard-owned pool groups available by index | +| `--pool-sizes-gb` | *(required)* | Comma-separated usable sizes of the ordered pools in every group | | `--compatibility-digest` | *(required)* | Exact digest every claimant must provide | -The service rounds each pool size down to a memory-page boundary and adds one -fixed 100 MiB journal to each Guard allocation. For example, `48,16` maps 64 -GiB plus the journal for every Guard. The current claim path exposes those -regions as one combined data area. +Each Guard gets one fixed 100 MiB recovery-journal region, added on top of the +listed usable sizes. The example therefore creates one mapping of 64 GiB plus +100 MiB. Its layout is `[journal header + journal payload][pool 0][pool 1]`; +additional pools follow in list order. Each listed size is rounded down to the +native memory-page boundary; a value smaller than one page is rejected. The pre-release wire protocol remains version 1. A worker calls `KVCRClient.claim(guard_index, pool_layouts, compatibility_digest, control_bind)`, -naming the address its Guard will answer on. The digest must match the service -exactly, and each pool-layout entry is `(pool_name, block_size_bytes)`. Callers must -change the digest whenever the pool layout or any other KV-cache term changes. -The returned `KVCRPoolHold` describes the mapped local DRAM and owns an exclusive -lease on the pool. Only one pool-layout entry is currently supported; an empty -string is a valid pool name. - -**A pool's configuration is fixed by its first claim.** Every later claim on -that pool must name the same pool layout and, when G3 is configured, the same -G3 paths in the same order, the same per-file capacity, and the same backend -and backend options. It must also name the same remote framework DRAM backend. -A mismatch is refused for the life of the service. Change the configuration by -restarting the service, which recreates the pools. - -The service grants a pool to one live claimant at a time, and pool mappings are -not inherited by forked children. The `KVCRPoolHold` remains owned by the +naming the address its Guard will answer on and each ordered pool's name and +block size. The digest must match the service exactly, and callers must change +it whenever a pool layout or any other KV-cache layout term changes. The +returned `KVCRPoolHold` owns the group's exclusive lease and exposes every pool +through `local_dram.pools` as `(name, address, size_bytes)`. The client maps the +allocation once; each pool's geometry is still validated independently, so +pools do not have to agree on a block size. + +`KVCRConfig.pool_layouts` supplies the same ordered layouts to direct and +`KVCRGuardConfig`-driven construction. Remote-transfer peers must use the same +pool names, block sizes, and order; a mismatch fails that operation. G3 remains +limited to a single-pool layout. + +**A pool group's configuration is fixed by its first claim.** Every later +claim on that Guard must name the same ordered pool layout and, when G3 is +configured, the same G3 paths in the same order, the same per-file capacity, +the same backend and backend options, and the same remote framework DRAM +backend; one that does not is refused for the life of the service, because a +different layout renames the blocks and slots the recovered records describe. +Change the layout by restarting the service, which recreates the groups. + +The service grants a whole pool group to one live claimant at a time; its pools +are allocated, claimed, promoted, and freed together. Pool mappings are not +inherited by forked children. The `KVCRPoolHold` remains owned by the claiming process and must not be used by a forked child. Applications must also create the shareable framework-control listener after their final fork. A second claim is rejected while the claimant's pidfd reports it alive. The lease socket -is close-on-exec, and the service continues fencing the pool by that pidfd until +is close-on-exec, and the service continues fencing the group by that pidfd until the process exits. Closing the claim connection, including an EOF, does not release a live claimant's lease. `KVCRPoolHold.release()` first unmaps the pool -locally, then explicitly releases the lease and waits for the service's +group locally, then explicitly releases the lease and waits for the service's acknowledgement. #### Recovery across a claimant's death -A `KVCRGuardConfig` opts into the service pool and its Guard together. A +A `KVCRGuardConfig` opts into a service pool group and its Guard together. A claimant whose framework control cannot share a listener is refused rather than granted an unguarded pool -- recovery asked for and silently not provided is worse than a failed startup. Without a `KVCRGuardConfig`, KVCR neither contacts the service nor builds a Guard. -The service binds the pool's control endpoint and hands the claimant a -duplicate of it. When that claimant dies, the pool's Guard takes over the same -address with the cache still in place; no second port is configured, and the -pool stays busy to any claimant that cannot inherit the endpoint. A clean -release instead returns the Guard to standby and the pool to claimable, and a -replacement primary takes a served pool back keeping the recovered records -rather than rebuilding them. Either handover costs time linear in the number of -recovered blocks, so size it against how much cache a pool actually holds. - -Recovered blocks are ranked for eviction as they are installed, so a pool -recovered full still accepts new deposits. They carry no access history, so a +The service binds the pool group's control endpoint and hands the claimant a +duplicate of it. When that claimant dies, the whole group transfers to its +Guard, which takes over the same address with every pool retained; no second +port is configured, and the group stays busy to any claimant that cannot +inherit the endpoint. The promoted Guard serves recovered G2 data from every +configured pool. A clean release returns the Guard to standby and the group to +claimable, and a replacement primary takes the entire served group back with +its recovered records rather than rebuilding them. Either handover costs time +linear in the number of recovered blocks, so size it against how much cache +the group holds. + +Recovered blocks are ranked for eviction as they are installed, so a fully +recovered group still accepts new deposits. They carry no access history, so a recovered block ranks below anything this process has served and is evicted first. @@ -398,25 +410,24 @@ recovery has to retry it. A promoted Guard always answers a stale request -- serving it, or failing it, even when it was promoted with nothing to serve -- so the peer retries instead of waiting on a completion nobody will send. -Every pool has a Guard for its whole life, and there is no per-pool -containment. Any Guard failure stops the service, on the grounds that a pool +Every pool group has a Guard for its whole life, and there is no per-Guard +containment. Any Guard failure stops the service, on the grounds that a group which can no longer be recovered, and may still hold an endpoint the service cannot reach, is not something to limp on with. One case is deliberately not a Guard failure: a primary publishing faster than its Guard can mirror fills the ring. Both sides treat that as survivable -- the -primary stops publishing, the Guard drops what it holds -- and the pool becomes -claimable but cold if that primary dies. Recovery is lost for that pool only. +primary stops publishing, the Guard drops what it holds -- and the group becomes +claimable but cold if that primary dies. Recovery is lost for that group only. Watch for `KVCR pool recovery disabled` if failovers stop coming back warm. The -journal is a fixed 100 MiB whatever `--pool-sizes-gb` is, so the only levers are -larger blocks, which publish fewer residency changes, or accepting a cold -failover for that pool. +journal is a fixed 100 MiB whatever `--pool-sizes-gb` is, so the only levers +are larger blocks, shorter pool names, fewer pool locations per key, or +accepting a cold failover for that group. -A Guard serves only the recovered G2 half; it opens no G3. A block that lived -only on disk is unavailable until a replacement primary claims the pool. The -records naming it are carried across, so the replacement reopens the tier with -its disk cache rather than a cold one -- the files themselves are not held in -the meantime, which is the limitation described below. +A Guard opens no G3. A block that lived only on disk is unavailable until a +replacement primary claims the group. The records naming it are carried across, +so the replacement reopens the tier with its disk cache rather than a cold one -- +the files themselves are not held in the meantime, which is the limitation below. **Deployment prerequisite.** Make the configured NIXL backends available in each process that uses them. Nothing checks plugin availability across processes @@ -429,7 +440,7 @@ what those records name. Nothing holds those files while the Guard serves either -- a tier's exclusive lock lives with the tier, and a Guard opens no G3. Pointing a second KVCR at the same G3 paths is therefore not a supported configuration: it is not detected, and the replacement will serve whatever is -in the slots. The intended first step -- having the service refuse two pools +in the slots. The intended first step -- having the service refuse two Guards that name the same paths -- is not implemented. The same applies to a file that is simply gone. A tier recreates a missing G3 @@ -880,14 +891,14 @@ Verify that: - the socket parent and pool directory exist and are writable; - the pool directory has capacity for every Guard's full allocation: the sum - of `--pool-sizes-gb` plus its 100 MiB journal. A pool changing hands briefly + of `--pool-sizes-gb` plus one 100 MiB journal. A group changing hands briefly appends its handback snapshot past that size; where there is no room for it, that handover comes back cold and the service carries on; - another process is not listening on the socket; - `--guard-count` is at least one; and -- every comma-separated `--pool-sizes-gb` value is positive, finite, and at - least one memory page. +- every `--pool-sizes-gb` item is positive, finite, and at least one memory + page. The service removes a stale socket only after confirming no live service is listening. It refuses to replace a socket owned by another live service. diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index 61f3309..453e7bf 100644 --- a/src/kvcr/guard.py +++ b/src/kvcr/guard.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Private journal-backed Guard for one service-owned pool.""" +"""Private journal-backed Guard for one service-owned pool group.""" import concurrent.futures import enum @@ -26,7 +26,7 @@ ) from .control_channels import KVCRServiceError, ZmqPeerControlChannel from .core import _BlockRecord, _KVCRCore -from .guard_protocol import PidfdLiveness, _TierConfig +from .guard_protocol import PidfdLiveness, _PoolDescriptor, _TierConfig from .local_disk import _G3Residency from .memory import ( KVCRPoolAttachment, @@ -59,7 +59,7 @@ class _PoolLease: - """One pool's holder identity and persistent control listener.""" + """One Guard's holder identity and persistent control listener.""" def __init__(self, guard_index: int) -> None: self._guard_index = guard_index @@ -171,11 +171,18 @@ def _with_g3( class _RecoveryState: - """One pool's attachment, journal, and mutable recovery ownership.""" + """One pool group's attachment, descriptors, and mutable recovery.""" - def __init__(self, spec: KVCRPoolSpec, compatibility_digest: str) -> None: + def __init__( + self, + spec: KVCRPoolSpec, + compatibility_digest: str, + pool_sizes_bytes: tuple[int, ...], + ) -> None: self._spec = spec self._compatibility_digest = compatibility_digest + self._pool_sizes_bytes = pool_sizes_bytes + self.pools: tuple[_PoolDescriptor, ...] = () self.attachment: KVCRPoolAttachment | None = None self._journal: RecoveryJournal | None = None self.mirror: _RecoveryMirror | None = None @@ -187,16 +194,30 @@ def prepare(self) -> None: self.attachment = KVCRPoolAttachment.attach(self._spec) self._journal = RecoveryJournal(self.attachment) - def recover(self, pool_layouts: PoolBlockLayouts) -> _RecoveryMirror: - """Return held recovery or read the prior handback under this pool layout.""" - if self.mirror is not None: - return self.mirror - return read_handback(self.attachment, self._compatibility_digest, pool_layouts) + def configure(self, pool_layouts: PoolBlockLayouts) -> None: + """Take up an ordered layout and its prior recovery atomically.""" + if len(pool_layouts) != len(self._pool_sizes_bytes): + raise ValueError("pool layout must match the number of pools in the Guard") + offset = self._spec.journal_bytes + descriptors = [] + for size_bytes, (name, block_size_bytes) in zip( + self._pool_sizes_bytes, pool_layouts + ): + descriptors.append( + _PoolDescriptor(name, size_bytes, block_size_bytes, offset) + ) + offset += size_bytes + pools = tuple(descriptors) + mirror = self.mirror + if mirror is None: + mirror = read_handback(self.attachment, self._compatibility_digest, pools) + self.pools = pools + self.mirror = mirror - def start_primary(self, pool_layouts: PoolBlockLayouts) -> None: + def start_primary(self) -> None: """Arm recovery for the accepted primary and reset its journal.""" if self.mirror is None: - self.mirror = _RecoveryMirror(tuple(name for name, _ in pool_layouts)) + self.mirror = _RecoveryMirror(tuple(pool.name for pool in self.pools)) self._journal.reset() def poll(self) -> bool: @@ -220,9 +241,7 @@ def invalidate_journal(self) -> None: with suppress(Exception): self._journal.invalidate() - def take_for_promotion( - self, pool_layouts: PoolBlockLayouts - ) -> dict[BlockKey, _BlockRecord]: + def take_for_promotion(self) -> dict[BlockKey, _BlockRecord]: """Drain and transfer recovered records, leaving a fresh mirror.""" records: dict[BlockKey, _BlockRecord] = {} mirror = self.mirror @@ -237,7 +256,7 @@ def take_for_promotion( except RecoveryJournalError as error: self._drop_recovery(error) # A handover still needs somewhere to put the core's eventual records. - self.mirror = _RecoveryMirror(tuple(name for name, _ in pool_layouts)) + self.mirror = _RecoveryMirror(tuple(pool.name for pool in self.pools)) return records def prepare_to_serve( @@ -249,17 +268,6 @@ def prepare_to_serve( } return _without_g3(records) - def local_dram_info( - self, - effective_bytes: int, - pool_name: str, - backend: str, - ) -> LocalDramOptions: - return LocalDramOptions( - [(pool_name, self.attachment.data_address, effective_bytes)], - backend, - ) - def release_snapshot_region(self) -> None: self.attachment.release_snapshot_region() @@ -269,16 +277,12 @@ def release_snapshot_region(self) -> None: # replacement serves whatever is in the slots. Refusing two pools that name # the same paths is the cheap first step; it does not cover a second # service, or a KVCR using G3 with no pool at all. - def hand_back( - self, - records: dict[BlockKey, _BlockRecord], - pool_layouts: PoolBlockLayouts, - ) -> None: - """Write and mirror a closed core's map under its pool layout.""" + def hand_back(self, records: dict[BlockKey, _BlockRecord]) -> None: + """Write and mirror a closed core's mutable record table in place.""" mirror = self.mirror records = _with_g3(records, self._g3_records) try: - self._write_handback(records, pool_layouts) + self._write_handback(records) except OSError as error: if error.errno not in _RECOVERY_CAPACITY_ERRORS: raise @@ -288,13 +292,13 @@ def hand_back( mirror.adopt(records) self._g3_records = {} - def release(self, pool_layouts: PoolBlockLayouts) -> None: + def release(self) -> None: """Write the current primary's journal tail, then drop its mirror.""" mirror = self.mirror try: while (frame := self._journal.read_next()) is not None: mirror.apply(*frame) - self._write_handback(mirror.take_records(), pool_layouts) + self._write_handback(mirror.take_records()) except RecoveryJournalError as error: self._drop_recovery(error) except OSError as error: @@ -310,15 +314,11 @@ def _drop_recovery(self, error: RecoveryJournalError | OSError) -> None: ) self.mirror = None - def _write_handback( - self, - records: Mapping[BlockKey, _BlockRecord], - pool_layouts: PoolBlockLayouts, - ) -> None: + def _write_handback(self, records: Mapping[BlockKey, _BlockRecord]) -> None: write_recovery_snapshot( self.attachment, - canonical_pool_terms(self._compatibility_digest, pool_layouts, self._spec), - _recovery_frames(records, tuple(name for name, _ in pool_layouts)), + canonical_pool_terms(self._compatibility_digest, self.pools, self._spec), + _recovery_frames(records, tuple(pool.name for pool in self.pools)), ) def close(self) -> None: @@ -329,10 +329,11 @@ def close(self) -> None: self._journal = None self.mirror = None self._g3_records = {} + self.pools = () class _Command: - """One request on the pool's mailbox, and the future its answer arrives on.""" + """One request on the Guard's mailbox, and the future its answer arrives on.""" def __init__(self, operation: str, args: tuple[Any, ...] = ()) -> None: self.operation, self.args = operation, args @@ -356,8 +357,8 @@ class _Phase(enum.Enum): class _Guard: - """One pool's lifecycle actor: owns the pool file, control endpoint, holder - pidfd, and recovery state, alive as long as the service owns the pool. + """One pool group's lifecycle actor: owns its allocation, control endpoint, + holder pidfd, and recovery state, alive as long as the service owns the group. Outlives every primary: a claim is reported to it, not what creates it. """ @@ -368,12 +369,13 @@ def __init__( *, compatibility_digest: str, guard_index: int = 0, + pool_sizes_bytes: tuple[int, ...], owner: _KVCRPoolOwner | None = None, refusing: Callable[[], bool] = lambda: False, ) -> None: self._spec = spec self._guard_index = guard_index - # Owned here, not by the registry: one thread owns one pool, so a + # Owned here, not by the registry: one thread owns one pool group, so a # claim needs no lock -- the mailbox is the reservation. self._owner = owner self._refusing = refusing @@ -381,7 +383,7 @@ def __init__( # Owned by the current primary. self._control: ZmqPeerControlChannel | None = None self._configured: _TierConfig | None = None - self._recovery = _RecoveryState(spec, compatibility_digest) + self._recovery = _RecoveryState(spec, compatibility_digest, pool_sizes_bytes) self._core = None self._commands: queue.Queue[_Command] = queue.Queue() self._ops = { @@ -421,8 +423,8 @@ def start(self) -> None: def claim( self, liveness: PidfdLiveness, tier_config: _TierConfig, bind: tuple[str, int] - ) -> "tuple[KVCRPoolSpec, int, _Lease]": - """Give this pool to a primary, with the endpoint its Guard answers on. + ) -> "tuple[KVCRPoolSpec, tuple[_PoolDescriptor, ...], int, _Lease]": + """Give this pool group to a primary, with its Guard's endpoint. Reserved on the requesting thread: a mid-transition pool answers busy immediately instead of queueing the claimant. """ @@ -630,7 +632,7 @@ def _observe_holder(self) -> None: def _claim( self, liveness: PidfdLiveness, tier_config: _TierConfig, bind: tuple[str, int] - ) -> "tuple[KVCRPoolSpec, int, _Lease]": + ) -> "tuple[KVCRPoolSpec, tuple[_PoolDescriptor, ...], int, _Lease]": """Give the pool to a primary, and take up the endpoint it named. All fallible work runs before the lease exists, and commit and refusal share one lock: a lease is never half-granted. @@ -651,7 +653,7 @@ def _claim( if not self._closing and not self._refusing(): self._pool_lease.current = liveness self._phase = _Phase.PRIMARY - return self._spec, granted_fd, liveness + return self._spec, self._recovery.pools, granted_fd, liveness # Refused at the commit: a closing service must not grant a pool. # Everything adopted goes back as a release would have put it. self._release() @@ -737,26 +739,24 @@ def _adopt(self, control: ZmqPeerControlChannel, tier_config: _TierConfig) -> No if self._failure is not None: raise self._failure self._refuse_incompatible(tier_config) - served_under = self._configured.pool_layouts if self._configured else () # The prior handback is this lease's baseline. Read now, under - # the claim's pool layout: refusing at promotion stops the service, + # the claim's layout: refusing at promotion stops the service, # and a claimant dying in between takes everything with it. - recovered = self._recovery.recover(tier_config.pool_layouts) + self._recovery.configure(tier_config.pool_layouts) # Last, once nothing left can refuse this claim: a pool whose handback # would not replay has not chosen anything, and a corrected claim can # still have it. - self._configure(tier_config) - self._recovery.mirror = recovered + self._configured = tier_config except BaseException: control.close() raise try: if self._serving: - self._hand_back(served_under) + self._hand_back() self._resumable = True # A refused handback is cold for the new lease, not unmirrored. - self._recovery.start_primary(tier_config.pool_layouts) + self._recovery.start_primary() # The old channel is the last reference to the prior primary's listener. if self._control is not None: self._control.close() @@ -776,32 +776,25 @@ def _release(self) -> None: if self._failure is not None: raise self._failure if self._serving: - self._hand_back(self._configured.pool_layouts) + self._hand_back() # Re-adopt lets start_primary() retain or replace the mirror. self._recovery.mirror = None elif self._recovery.mirror is not None: - self._recovery.release(self._configured.pool_layouts) + self._recovery.release() if self._control is not None: self._control.close() self._control = None def _refuse_incompatible(self, tier_config: _TierConfig) -> None: - """Refuse tiers other than the ones this pool was claimed with. + """Refuse tiers other than the ones this pool group was claimed with. The first claim fixes configuration for the service's lifetime: the - bytes stay, and a changed pool layout or G3 path order misnames every slot. + bytes stay, and a changed block size or G3 path order misnames every slot. """ if self._configured is not None and self._configured != tier_config: raise RecoveryMirrorError( "KVCR pool was claimed with another tier configuration" ) - def _configure(self, tier_config: _TierConfig) -> None: - """Take up this primary's tiers: the geometry check runs before the - assignment, so a bad configuration leaves the old one intact. - """ - _compute_pool_geometry(self._spec.data_bytes, tier_config.pool_layouts[0][1]) - self._configured = tier_config - def _poll(self) -> bool: """Mirror what is waiting; True if more remains. The batch bounds a command's wait, not how much drains. @@ -844,10 +837,10 @@ def _promote(self) -> None: self._resumable = False if self._failure is not None: raise self._failure - self._serve(self._recovery.take_for_promotion(self._configured.pool_layouts)) + self._serve(self._recovery.take_for_promotion()) def _serve(self, records: dict[BlockKey, _BlockRecord]) -> None: - """Answer on this pool's endpoint, with whatever came back from it. + """Answer on this pool group's endpoint, with whatever came back from it. Serving nothing is still serving: answering refuses peers that staying bound would leave hanging. G2 only, no G3: that half is kept whole for @@ -858,17 +851,21 @@ def _serve(self, records: dict[BlockKey, _BlockRecord]) -> None: def reject_pin(keys: object) -> int: raise RuntimeError("Guard has no framework-owned memory") - pool_name, block_size = self._configured.pool_layouts[0] - effective_bytes, _ = _compute_pool_geometry(self._spec.data_bytes, block_size) - dram = self._recovery.local_dram_info( - effective_bytes, - pool_name, + dram = LocalDramOptions( + [ + ( + pool.name, + self._recovery.attachment.address + pool.offset_bytes, + _compute_pool_geometry(pool.size_bytes, pool.block_size_bytes)[0], + ) + for pool in self._recovery.pools + ], self._configured.remote_fw_dram_backend, ) core = _KVCRCore( KVCRConfig( nixl_agent_name=f"KVCR-Guard-{uuid.uuid4()}", - pool_layouts=list(self._configured.pool_layouts), + pool_layouts=self._configured.pool_layouts, inventory_report_interval_ms=0, nixl_listen_port=0, ), @@ -894,8 +891,8 @@ def reject_pin(keys: object) -> int: core.start() self._serving = True - def _hand_back(self, pool_layouts: PoolBlockLayouts) -> None: - """Stop serving, leaving this pool's state where the next primary looks. + def _hand_back(self) -> None: + """Stop serving, leaving this pool group's state where the next primary looks. The core closes first: the Guard stops answering, and region and records both come from the map close leaves behind. """ @@ -903,7 +900,7 @@ def _hand_back(self, pool_layouts: PoolBlockLayouts) -> None: if core is None or self._recovery.mirror is None: raise RecoveryMirrorError("a serving Guard has no state to hand back") core.close() - self._recovery.hand_back(core._block_record_map, pool_layouts) + self._recovery.hand_back(core._block_record_map) self._core = None self._serving = False diff --git a/src/kvcr/guard_protocol.py b/src/kvcr/guard_protocol.py index 24985a5..8520fbe 100644 --- a/src/kvcr/guard_protocol.py +++ b/src/kvcr/guard_protocol.py @@ -46,6 +46,23 @@ class _G3Config(msgspec.Struct, frozen=True, forbid_unknown_fields=True): def __post_init__(self) -> None: if not all(os.path.isabs(path) for path in self.paths): raise ValueError("G3 paths must be absolute") + resolved = {os.path.realpath(path) for path in self.paths} + if len(resolved) != len(self.paths): + raise ValueError("G3 file paths must be unique") + + +class _PoolDescriptor(msgspec.Struct, frozen=True, forbid_unknown_fields=True): + """One ordered pool region within a Guard-owned allocation.""" + + name: str + size_bytes: Annotated[int, msgspec.Meta(gt=0)] + block_size_bytes: Annotated[int, msgspec.Meta(gt=0)] + offset_bytes: Annotated[int, msgspec.Meta(ge=0)] = 0 + + def __post_init__(self) -> None: + _compute_pool_geometry(self.size_bytes, self.block_size_bytes) + if type(self.offset_bytes) is not int or self.offset_bytes < 0: + raise ValueError("pool offset must be a non-negative integer") class _TierConfig(msgspec.Struct, frozen=True, forbid_unknown_fields=True): @@ -55,22 +72,15 @@ class _TierConfig(msgspec.Struct, frozen=True, forbid_unknown_fields=True): def __post_init__(self) -> None: _validate_pool_layouts(self.pool_layouts) - # TODO: Support multiple pools after fetch and storage can discover layouts. - if len(self.pool_layouts) != 1: - raise ValueError("only a single pool is currently supported") - block_size_bytes = self.pool_layouts[0][1] - # Mirrors what the claimant's _G3 will enforce. The first claim fixes - # the pool's tiers forever, so a config no claimant could ever open - # must be refused here, before it binds. - if self.g3 is None: - return - if block_size_bytes % mmap.PAGESIZE: - raise ValueError("G3 slot size must be page aligned") - if self.g3.capacity_bytes_per_file % block_size_bytes: - raise ValueError("G3 file capacity must contain complete slots") - resolved = {os.path.realpath(path) for path in self.g3.paths} - if len(resolved) != len(self.g3.paths): - raise ValueError("G3 file paths must be unique") + if self.g3 is not None: + if len(self.pool_layouts) != 1: + raise ValueError("G3 does not support multiple pools") + block_size_bytes = self.pool_layouts[0][1] + if ( + block_size_bytes % mmap.PAGESIZE + or self.g3.capacity_bytes_per_file % block_size_bytes + ): + raise ValueError("G3 requires page-aligned complete slots") class _Claim(msgspec.Struct, frozen=True, tag="claim"): @@ -105,6 +115,7 @@ class _Granted(msgspec.Struct, frozen=True, tag="granted"): guard_index: int spec: KVCRPoolSpec tier_config: _TierConfig + pools: tuple[_PoolDescriptor, ...] version: ProtocolVersion @@ -167,9 +178,10 @@ def close(self) -> None: @dataclass class KVCRPoolHold: - """A mapped pool and the connection holding its lease.""" + """One mapped pool group and the connection holding its lease.""" local_dram: LocalDramOptions + _pools: tuple[_PoolDescriptor, ...] _attachment: KVCRPoolAttachment _connection: FramedConnection _control_listener_fd: int | None = None @@ -187,7 +199,7 @@ def release(self, *, activated: bool = True) -> None: """Stop local access before releasing the connection-scoped lease. ``activated=False`` tells the service this lease never served: the - Guard it stood down may resume instead of leaving the pool idle. + Guard it stood down may resume instead of leaving the pool group idle. """ if self._release_attempted: return @@ -225,7 +237,7 @@ def claim( g3: G3Options | None = None, remote_fw_dram_backend: str = "UCX", ) -> KVCRPoolHold: - """Claim and map one service-owned pool.""" + """Claim and map one Guard-owned pool group.""" g3_config = g3 and { "paths": [str(path.expanduser().resolve()) for path in g3.paths], "capacity_bytes_per_file": g3.capacity_bytes_per_file, @@ -259,27 +271,30 @@ def claim( if isinstance(response, _Error): raise KVCRServiceError(response.message) grant_received = True - spec = _grant_spec(response, guard_index, request.tier_config) + spec, granted_pools = _grant_layout( + response, guard_index, request.tier_config + ) if listener_fd is None: # Every pool has a Guard, and a Guard answers on the endpoint # this claimant named. A grant without it means the two sides - # disagree about who serves this pool. + # disagree about who serves this pool group. raise KVCRGuardProtocolError( "claim was granted without the endpoint it answers on" ) - try: - pool_name, block_size = request.tier_config.pool_layouts[0] - effective_bytes, _ = _compute_pool_geometry(spec.data_bytes, block_size) - except ValueError as geometry_error: - raise KVCRGuardProtocolError( - "invalid pool grant: no room for one KV block" - ) from geometry_error attachment = KVCRPoolAttachment.attach(spec) return KVCRPoolHold( local_dram=LocalDramOptions( - [(pool_name, attachment.data_address, effective_bytes)], + [ + ( + pool.name, + attachment.address + pool.offset_bytes, + pool.size_bytes, + ) + for pool in granted_pools + ], request.tier_config.remote_fw_dram_backend, ), + _pools=granted_pools, _attachment=attachment, _connection=connection, _control_listener_fd=listener_fd, @@ -309,11 +324,11 @@ def claim( raise -def _grant_spec( +def _grant_layout( response: _Granted, requested_guard_index: int, requested_tier_config: _TierConfig, -) -> KVCRPoolSpec: +) -> tuple[KVCRPoolSpec, tuple[_PoolDescriptor, ...]]: """Take the grant apart, refusing one that answers a different request.""" if response.guard_index != requested_guard_index: raise KVCRGuardProtocolError( @@ -322,7 +337,23 @@ def _grant_spec( ) if response.tier_config != requested_tier_config: raise KVCRGuardProtocolError("claim tier configuration mismatch") - return response.spec + pools = response.pools + if len(pools) != len(requested_tier_config.pool_layouts): + raise KVCRGuardProtocolError("claim pool count mismatch") + expected_offset = response.spec.journal_bytes + for index, (pool, layout) in enumerate( + zip(pools, requested_tier_config.pool_layouts) + ): + if ( + (pool.name, pool.block_size_bytes) != layout + or pool.offset_bytes != expected_offset + or pool.size_bytes % mmap.PAGESIZE + ): + raise KVCRGuardProtocolError(f"claim pool {index} layout mismatch") + expected_offset += pool.size_bytes + if expected_offset != response.spec.mapping_bytes: + raise KVCRGuardProtocolError("claim pool sizes do not fill the allocation") + return response.spec, pools def _send_release(connection: FramedConnection, *, activated: bool = True) -> None: diff --git a/src/kvcr/kvcr_service.py b/src/kvcr/kvcr_service.py index bcd62ba..8111918 100644 --- a/src/kvcr/kvcr_service.py +++ b/src/kvcr/kvcr_service.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""KVCR-Service: the process that owns what outlives a worker. - -Today that is shared-memory pools and their recovery Guards; a worker claims one. -""" +"""KVCR-Service owns Guarded shared-memory pool groups beyond a worker's life.""" import argparse import contextlib @@ -39,6 +36,7 @@ _Claim, _Error, _Granted, + _PoolDescriptor, _Released, _TierConfig, ) @@ -64,10 +62,10 @@ class _PoolRegistry: - """A directory of pools, each owned end to end by its own Guard thread. + """A directory of pool groups, each owned by one Guard thread. - No locks: each pool's mailbox orders its claims, releases and deaths; - pools share nothing but the refusal flag. + No locks: each Guard's mailbox orders its claims, releases and deaths; + Guards share nothing but the refusal flag. """ def __init__( @@ -98,20 +96,21 @@ def __init__( journal_bytes=journal_bytes, pool_dir=self._pool_dir, ) - # Built with the pool, not a claim: a Guard that cannot attach its - # pool is better discovered at startup than when a worker dies. + # Built with the group, not a claim: a Guard that cannot attach its + # allocation is better discovered at startup than when a worker dies. try: guard = _Guard( owner.spec, functools.partial(self._guard_failed, rank), compatibility_digest=compatibility_digest, guard_index=rank, + pool_sizes_bytes=pool_sizes_bytes, owner=owner, refusing=self._refusing.is_set, ) except BaseException: - # Nothing has recorded this pool yet, so the sweep below cannot - # reach it and its file would outlive the process. + # Nothing has recorded this pool group yet, so the sweep + # below cannot reach it and its file would outlive the process. owner.close() raise # Recorded before it starts, so a failed preparation is rolled back by @@ -123,14 +122,14 @@ def __init__( raise def _release_pools(self) -> None: - """Give back every pool built so far, for a startup that cannot finish.""" + """Give back every pool group built so far after a failed startup.""" for guard_index, guard in list(self._guards.items()): # Even an interrupt must not stop the sweep: the startup failure is - # already propagating, and every pool left behind is committed RAM. + # already propagating, and every group left behind is committed RAM. try: guard.close() except BaseException: - # The Guard's thread may still hold this pool's mapping, and + # The Guard's thread may still hold this group's mapping, and # unlinking under it would fault the process. Leave the file # for the next start's purge, and keep the pool visible. logger.warning( @@ -186,8 +185,8 @@ def claim( tier_config: _TierConfig, liveness: PidfdLiveness, control_bind: tuple[str, int], - ) -> "tuple[KVCRPoolSpec, int, _Lease]": - """Give a pool to a primary, and hand back the endpoint it answers on. + ) -> "tuple[KVCRPoolSpec, tuple[_PoolDescriptor, ...], int, _Lease]": + """Give a pool group to a primary and return its Guard endpoint. The refusal check is a fast path only; the grant commits on the pool's actor under the same lock refuse_claims reads, so no grant follows it. @@ -204,27 +203,27 @@ def abort_grant(self, guard_index: int, lease: "_Lease") -> None: self._guard(guard_index).abort_grant(lease) def refuse_claims(self) -> None: - """Stop granting pools without waiting for the close path to run. + """Stop granting pool groups without waiting for the close path to run. - Each pool's phase lock is the barrier: after this returns, no grant can + Each Guard's phase lock is the barrier: after this returns, no grant can commit. Pre-barrier grants may still deliver; those leases are fenced. """ self._refusing.set() - # Snapshot: close() deletes pools from the dict on other threads. + # Snapshot: close() deletes groups from the dict on other threads. for guard in list(self._guards.values()): with guard._phase_lock: pass def close(self) -> None: - """Give every pool back, keeping the first reason one would not go. + """Give every pool group back, keeping the first reason one would not go. - A pool that will not close keeps only its own file and endpoint and + A group that will not close keeps only its own file and endpoint and stays listed, so a later close can try it again; failing that, the flock dies with the process and the next start reclaims. """ self._refusing.set() failure: BaseException | None = None - # Tell all pools before waiting on any: a wedged one must not block the rest. + # Tell all Guards before waiting on any: a wedged one must not block the rest. for guard in self._guards.values(): try: guard.begin_close() @@ -241,7 +240,7 @@ def close(self) -> None: except BaseException as error: # noqa: BLE001 - raised below failure = failure or error kept.add(guard_index) - # Wedged pools stay visible; drained ones stay listed until the whole + # Wedged groups stay visible; drained ones stay listed until the whole # drain finished, so a release racing shutdown is absorbed. for guard_index in [index for index in self._guards if index not in kept]: del self._guards[guard_index] @@ -257,8 +256,8 @@ def _guard_failed( ) -> None: """A Guard has stopped being one, which the service cannot survive. - TODO: no per-pool containment. Its pool can no longer be recovered and may - still hold an endpoint the service cannot reach. One pool takes the others' + TODO: no per-Guard containment. Its group can no longer be recovered and may + still hold an endpoint the service cannot reach. One Guard takes the others' workers with it; add isolation back if that stops being acceptable. """ logger.critical("KVCR Guard %d failed", guard_index) @@ -321,7 +320,7 @@ def handle(self) -> None: def _await_release(self, guard_index: int, lease: "_Lease") -> None: """Wait for the one message a held connection may send: its release. - The pool's actor watches the pidfd, not this thread. EOF only ends the + The Guard actor watches the pidfd, not this thread. EOF only ends the connection; the lease outlives it, and a death still promotes. """ while True: @@ -406,14 +405,20 @@ def dispatch( raise KVCRServiceError( "KVCR compatibility digest does not match the service" ) - spec, listener_fd, lease = self.registry.claim( + spec, pools, listener_fd, lease = self.registry.claim( request.guard_index, request.tier_config, liveness, (request.control_host, request.control_port), ) return ( - _Granted(request.guard_index, spec, request.tier_config, _PROTOCOL_VERSION), + _Granted( + request.guard_index, + spec, + request.tier_config, + pools, + _PROTOCOL_VERSION, + ), (request.guard_index, listener_fd, lease), ) diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index 3a376c3..d5555f9 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -33,6 +33,15 @@ logger = logging.getLogger(__name__) +def _validate_g3_slot_geometry(config: G3Options, slot_size: int) -> None: + """Validate the scalar data-plane relationship Guard protocol ignores.""" + if slot_size <= 0 or slot_size % os.sysconf("SC_PAGE_SIZE"): + raise ValueError("G3 slot size must be positive and page aligned") + capacity = config.capacity_bytes_per_file + if capacity <= 0 or capacity % slot_size: + raise ValueError("G3 file capacity must contain complete slots") + + @dataclass(slots=True) class _G3Residency: slot: int @@ -120,19 +129,12 @@ class _G3: """Own bounded files and the metadata needed to use them as G3 cache.""" def __init__(self, kvcr: "_KVCRCore", config: G3Options, slot_size: int) -> None: - page_size = os.sysconf("SC_PAGE_SIZE") paths = tuple(Path(path).expanduser().resolve() for path in config.paths) if not paths: raise ValueError("G3 requires at least one file path") if len(paths) != len(set(paths)): raise ValueError("G3 file paths must be unique") - if slot_size <= 0 or slot_size % page_size: - raise ValueError("G3 slot size must be positive and page aligned") - if ( - config.capacity_bytes_per_file <= 0 - or config.capacity_bytes_per_file % slot_size - ): - raise ValueError("G3 file capacity must contain complete slots") + _validate_g3_slot_geometry(config, slot_size) if not config.backend: raise ValueError("G3 NIXL backend must be non-empty") if not all( diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index d998f30..c9b3946 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -17,13 +17,18 @@ import msgspec -from .config import KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig +from .config import ( + KVCRBackendConfigs, + KVCRConfig, + KVCRGuardConfig, + _validate_pool_layouts, +) from .core import _BlockRecord, _KVCRCore -from .guard_protocol import KVCRClient, KVCRPoolHold -from .local_disk import _G3, _G3Residency +from .guard_protocol import KVCRClient, KVCRPoolHold, _PoolDescriptor +from .local_disk import _G3, _G3Residency, _validate_g3_slot_geometry from .local_dram import _LocalDram, _LocalDramResidency, _LocalDramState from .memory import _JOURNAL_HEADER_BYTES, KVCRPoolAttachment, KVCRPoolSpec -from .types import BlockKey, PoolBlockLayouts, RecoveryMirrorError +from .types import BlockKey, RecoveryMirrorError if TYPE_CHECKING: from .api import KVCRBindings @@ -73,10 +78,9 @@ def store_release(self, value: int) -> None: _RECORD_TYPES = frozenset({_RECORD_BLOCK}) -# Arrays, not maps: repeating field names costs ring space, and the ring -# filling ends recovery. 3 bytes a record instead of 21. -# -# Field order is the format. Append only -- never reorder or remove. +# Arrays avoid repeating field names in the bounded ring. Field order is the +# format: append only, never reorder or remove. G2 cost grows with its location +# count and pool-name lengths. class _RecoveryBlock(msgspec.Struct, frozen=True, array_like=True): # Ordered pool locations, or nothing. Pool names may repeat. g2: list[tuple[str, int]] | None = None @@ -453,6 +457,14 @@ def claim_guarded_pool( """ if backend_configs.local_dram is not None: raise ValueError("guard_config conflicts with backend_configs.local_dram") + if backend_configs.g3 is not None: + _validate_pool_layouts(config.pool_layouts) + if len(config.pool_layouts) != 1: + raise ValueError("G3 does not support multiple pools") + _validate_g3_slot_geometry( + backend_configs.g3, + config.pool_layouts[0][1], + ) # Duck-typed: what matters is whether the framework's control can hand its # endpoint over, not what class it is. framework_control = bindings.framework_control @@ -478,7 +490,7 @@ def claim_guarded_pool( recovered = read_handback( hold._attachment, guard_config.compatibility_digest, - config.pool_layouts, + hold._pools, ) except BaseException: # A failing release must not mask the error that made the claim unusable. @@ -566,16 +578,16 @@ def _recovery_frames( # Bound to the pool and to the geometry: a slot index only means the same -# bytes under the same file and pool layout. The generation stops a replay into a +# bytes under the same file and layout. The generation stops a replay into a # different pool of the same shape; the digest separates finished from filling. _SNAPSHOT_HEADER = struct.Struct("<32sQ") _SNAPSHOT_DOMAIN = b"KVCR-HANDBACK\0" -_SNAPSHOT_TERMS = struct.Struct(" bytes: """Encode what a handback region must not be replayed across.""" @@ -584,13 +596,23 @@ def canonical_pool_terms( + compatibility_digest.encode() + b"\0" + bytes.fromhex(spec.generation) - + msgspec.msgpack.encode(pool_layouts) - + _SNAPSHOT_TERMS.pack( + + _SNAPSHOT_ALLOCATION_TERMS.pack( spec.journal_bytes, spec.mapping_bytes, spec.device, spec.inode, ) + + msgspec.msgpack.encode( + [ + ( + pool.name, + pool.size_bytes, + pool.block_size_bytes, + pool.offset_bytes, + ) + for pool in pools + ] + ) ) @@ -689,7 +711,7 @@ def read_recovery_snapshot( def read_handback( pool: KVCRPoolAttachment, compatibility_digest: str, - pool_layouts: PoolBlockLayouts, + pools: tuple[_PoolDescriptor, ...], ) -> _RecoveryMirror: """Replay whatever the last Guard left for this pool, if anything. @@ -699,9 +721,9 @@ def read_handback( write never finished is this service's own, and is thrown away -- nothing else ever would, and it would refuse every later claim on this pool too. """ - pool_names = tuple(name for name, _ in pool_layouts) + pool_names = tuple(pool.name for pool in pools) mirror = _RecoveryMirror(pool_names) - terms = canonical_pool_terms(compatibility_digest, pool_layouts, pool._spec) + terms = canonical_pool_terms(compatibility_digest, pools, pool._spec) try: for frame in read_recovery_snapshot(pool, terms): mirror.apply(*frame) diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 155aa99..1618b8f 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -28,7 +28,7 @@ ) from kvcr.guard_protocol import _G3Config, _TierConfig from kvcr.local_disk import _G3Residency -from kvcr.local_dram import _LocalDramResidency, _LocalDramState +from kvcr.local_dram import _LocalDramState from kvcr.memory import KVCRPoolSpec from kvcr.recovery_journal import ( _RECORD_BLOCK, @@ -51,17 +51,36 @@ journal_bytes=8192, ) _TEST_DIGEST = "opaque digest: Preserve-Me EXACTLY" -# G3 terms are refused at decode unless a real claimant could open them, so -# tests that carry G3 use page-aligned strides over a page-sized pool. -_PAGE_STRIDE = os.sysconf("SC_PAGE_SIZE") -_PAGE_SPEC = msgspec.structs.replace(_TEST_SPEC, mapping_bytes=8192 + 2 * _PAGE_STRIDE) +# The scalar claim path validates G3 geometry, so tests that carry G3 use +# page-aligned block sizes over a page-sized pool. +_PAGE_BLOCK_SIZE_BYTES = os.sysconf("SC_PAGE_SIZE") +_PAGE_SPEC = msgspec.structs.replace( + _TEST_SPEC, mapping_bytes=8192 + 2 * _PAGE_BLOCK_SIZE_BYTES +) + + +def _tier( + block_size_bytes: int, + g3: _G3Config | None = None, + backend: str = "UCX", +) -> _TierConfig: + return _TierConfig([("", block_size_bytes)], g3, backend) + + +def _guard(spec: KVCRPoolSpec = _TEST_SPEC, failure_callback=None, **kwargs) -> _Guard: + return _Guard( + spec, + failure_callback, + compatibility_digest=_TEST_DIGEST, + pool_sizes_bytes=(spec.data_bytes,), + **kwargs, + ) def _fake_attachment() -> Mock: """A stand-in with the pool-tail surface a Guard reaches for.""" attachment = Mock( address=1234, - data_address=1234 + _TEST_SPEC.journal_bytes, _spec=_TEST_SPEC, ) attachment.mapped_snapshot.return_value = nullcontext(None) @@ -92,10 +111,8 @@ def _frame(key: BlockKey, record: _BlockRecord) -> tuple[int, bytes, bytes]: def _give_serving_core(guard: _Guard) -> Mock: """A serving core still holding one READY G2 block.""" - record = _BlockRecord( - local_dram=_LocalDramResidency([("", 0)], _LocalDramState.READY) - ) - core = Mock(_block_record_map={BlockKey(b"warm"): record}) + records = {BlockKey(b"warm"): _recovered_record(g2=[("", 0)])} + core = Mock(_block_record_map=records) guard._core = core guard._serving = True return core @@ -103,7 +120,7 @@ def _give_serving_core(guard: _Guard) -> Mock: def _configurable_guard() -> _Guard: """A Guard past preparation, with nothing held and no thread running.""" - guard = _Guard(_TEST_SPEC, compatibility_digest=_TEST_DIGEST) + guard = _guard() guard._phase = _Phase.IDLE return guard @@ -176,9 +193,9 @@ def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core(caplog) -> N control = Mock() journal = Mock() failure_callback = Mock() - guard = _Guard(_TEST_SPEC, failure_callback, compatibility_digest=_TEST_DIGEST) + guard = _guard(failure_callback=failure_callback) guard._control = control - guard._configure(_TierConfig([("", 16)], None)) + guard._configured = _tier(16) guard._recovery._journal = journal guard._serving = True guard._core = core @@ -195,8 +212,8 @@ def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core(caplog) -> N assert caplog.records[0].exc_info[1] is error core.close.assert_called_once_with() control.close.assert_not_called() - journal.invalidate.assert_called_once_with() failure_callback.assert_called_once_with(guard, error) + journal.invalidate.assert_called_once_with() def test_standby_guard_failure_releases_adopted_listener() -> None: @@ -210,9 +227,9 @@ def test_standby_guard_failure_releases_adopted_listener() -> None: error = RuntimeError("journal poll failed") journal = Mock() journal.read_next.side_effect = error - guard = _Guard(_TEST_SPEC, failure_callback, compatibility_digest=_TEST_DIGEST) + guard = _guard(failure_callback=failure_callback) guard._control = control - guard._configure(_TierConfig([("", 16)], None)) + guard._configured = _tier(16) guard._recovery._journal = journal guard._recovery.mirror = Mock() @@ -255,9 +272,12 @@ def test_guard_lives_out_adopt_promote_and_readopt_in_ownership_order( def new_core(config, bindings, backends) -> Mock: constructed.append((config, bindings, backends)) core = Mock(_local_dram=Mock(), _g3=None, _block_record_map={}) - core.adopt_recovery_records.side_effect = lambda records: order.append( - ("adopt", tuple(records)) - ) + + def adopt(records) -> None: + core._block_record_map = records + order.append(("adopt", tuple(records))) + + core.adopt_recovery_records.side_effect = adopt core.start.side_effect = lambda: order.append("start") label = f"core{len(cores) + 1}" core.close.side_effect = lambda: closed.append(label) @@ -278,19 +298,19 @@ def new_channel() -> Mock: attachment.release_snapshot_region.side_effect = lambda: order.append("clear") g3_config = _G3Config( paths=(str(tmp_path / "g3.data"),), - capacity_bytes_per_file=10 * _PAGE_STRIDE, + capacity_bytes_per_file=10 * _PAGE_BLOCK_SIZE_BYTES, backend="FILE", backend_options={}, ) - tier = _TierConfig([("", _PAGE_STRIDE)], g3_config, "REMOTE") - guard = _Guard(_PAGE_SPEC, compatibility_digest=_TEST_DIGEST) + tier = _tier(_PAGE_BLOCK_SIZE_BYTES, g3_config, "REMOTE") + guard = _guard(_PAGE_SPEC) # Driven directly, then the thread starts already busy: the actor blocks # on an empty mailbox when idle, so mutating around a sleeping thread # would race its wakeup instead of testing the ordering. guard._started = True guard._recovery.prepare() # Unclaimed, so any tier shape is still available; the first claim fixes it. - guard._refuse_incompatible(_TierConfig([("", 16)], None)) + guard._refuse_incompatible(_tier(16)) guard._adopt(new_channel(), tier) try: attach.assert_called_once_with(_PAGE_SPEC) @@ -298,7 +318,7 @@ def new_channel() -> Mock: # Adoption only grants; a core exists once a promotion needs one. assert constructed == [] with pytest.raises(RecoveryMirrorError, match="another tier configuration"): - guard._refuse_incompatible(_TierConfig([("", 16)], None)) + guard._refuse_incompatible(_tier(16)) promoted_records = guard._recovery.mirror._records guard._promote() @@ -310,7 +330,7 @@ def new_channel() -> Mock: assert config.nixl_listen_port == 0 assert bindings.framework_control is channels[0] assert backends.local_dram == LocalDramOptions( - [("", 1234 + 8192, 2 * _PAGE_STRIDE)], + [("", 1234 + 8192, 2 * _PAGE_BLOCK_SIZE_BYTES)], "REMOTE", ) assert backends.remote_fw_dram.backend == "REMOTE" @@ -332,14 +352,10 @@ def new_channel() -> Mock: # carried whole -- g3_only no longer names any live record, so a # rebuild from the core's map could not produce it. retained_g3 = guard._recovery._g3_records[first] - records = { - first: _BlockRecord( - local_dram=_LocalDramResidency([("", 0)], _LocalDramState.FILLING), - g3=_G3Residency(7), - ), - second: _recovered_record(g2=[("", 1)]), - } - cores[0]._block_record_map = records + records = cores[0]._block_record_map + first_residency = records[first].local_dram + assert first_residency is not None + first_residency.state = _LocalDramState.FILLING write_handback = Mock() guard._recovery._write_handback = write_handback guard._adopt(new_channel(), tier) @@ -384,7 +400,7 @@ def test_a_pool_that_lost_its_recovery_stays_claimable_on_every_path( """Which reader finds the invalid journal is a race; none may take the service.""" guard = _configurable_guard() # Every one of these readers runs on a pool a primary has already claimed. - guard._configured = _TierConfig([("", 16)], None) + guard._configured = _tier(16) guard._recovery.mirror = _RecoveryMirror(("",)) guard._recovery.attachment = Mock() guard._control = None @@ -394,7 +410,7 @@ def test_a_pool_that_lost_its_recovery_stays_claimable_on_every_path( journal.drain.side_effect = error guard._recovery._journal = journal written: list[object] = [] - guard._recovery._write_handback = lambda records, stride: written.append(records) + guard._recovery._write_handback = lambda records: written.append(records) served: list[dict] = [] guard._serve = served.append reported: list[BaseException] = [] @@ -435,15 +451,15 @@ def test_a_pool_that_lost_its_recovery_stays_claimable_on_every_path( def test_the_same_g3_paths_in_another_order_are_another_configuration() -> None: """A slot names its file by position, so reordering renames every slot.""" guard = _configurable_guard() - guard._configured = _TierConfig( - [("", _PAGE_STRIDE)], - _G3Config(("/a", "/b"), _PAGE_STRIDE, "MOCK", {}), + guard._configured = _tier( + _PAGE_BLOCK_SIZE_BYTES, + _G3Config(("/a", "/b"), _PAGE_BLOCK_SIZE_BYTES, "MOCK", {}), ) with pytest.raises(RecoveryMirrorError, match="another tier configuration"): guard._refuse_incompatible( - _TierConfig( - [("", _PAGE_STRIDE)], - _G3Config(("/b", "/a"), _PAGE_STRIDE, "MOCK", {}), + _tier( + _PAGE_BLOCK_SIZE_BYTES, + _G3Config(("/b", "/a"), _PAGE_BLOCK_SIZE_BYTES, "MOCK", {}), ) ) @@ -456,9 +472,9 @@ def test_guard_closes_control_when_its_thread_does_not_start(monkeypatch) -> Non "kvcr.guard.KVCRPoolAttachment.attach", Mock(return_value=attachment) ) monkeypatch.setattr("kvcr.guard.RecoveryJournal", Mock()) - guard = _Guard(_TEST_SPEC, compatibility_digest=_TEST_DIGEST) + guard = _guard() guard._control = control - guard._configure(_TierConfig([("", 16)], None)) + guard._configured = _tier(16) guard._thread.start = Mock(side_effect=RuntimeError("thread start failed")) with pytest.raises(RuntimeError, match="thread start failed"): @@ -501,7 +517,7 @@ def test_recovery_close_error_stays_first_while_lease_cleanup_continues() -> Non attachment = Mock(close=Mock(side_effect=[attachment_error, None])) holder = Mock(close=Mock(side_effect=RuntimeError("holder close failed"))) owner = Mock() - guard = _Guard(_TEST_SPEC, compatibility_digest=_TEST_DIGEST, owner=owner) + guard = _guard(owner=owner) guard._recovery.attachment = attachment mirror = guard._recovery.mirror = _RecoveryMirror(("",)) g3_records = guard._recovery._g3_records = {BlockKey(b"g3"): _G3Residency(0)} @@ -534,9 +550,9 @@ def test_a_close_refused_by_a_moving_core_retains_the_pool_until_quiescent( error = RuntimeError("close failed") core.close.side_effect = error core.is_quiescent.return_value = False - guard = _Guard(_TEST_SPEC, compatibility_digest=_TEST_DIGEST) + guard = _guard() guard._control = control - guard._configure(_TierConfig([("", 16)], None)) + guard._configured = _tier(16) guard._core = core guard._recovery.attachment = attachment caplog.set_level(logging.WARNING, logger="kvcr.guard") @@ -566,14 +582,14 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( """Refusals before the pool moves leave it choosable; failures after are fatal.""" reported: list[BaseException] = [] guard = _configurable_guard() - guard._recovery.attachment = Mock() + guard._recovery.attachment = _fake_attachment() guard._recovery._journal = Mock() guard._failure_callback = lambda _guard, error: reported.append(error) control = Mock() if refused_by == "hand-over": # A hand-back that fails cannot be reported as a refused claim. - guard._configured = _TierConfig([("", 16)], None) + guard._configured = _tier(16) guard._serving = True guard._recovery.mirror = _RecoveryMirror(("",)) guard._core = Mock(_block_record_map={}) @@ -581,7 +597,7 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( guard._hand_back = Mock(side_effect=failure) with pytest.raises(OSError, match="no space left"): - guard._adopt(control, _TierConfig([("", 16)], None)) + guard._adopt(control, _tier(16)) assert reported == [failure] assert guard._failure is failure @@ -591,11 +607,11 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( ) if refused_by == "geometry": expected: type[Exception] = ValueError - tier_config = _TierConfig([("", _TEST_SPEC.mapping_bytes)], None) + tier_config = _tier(_TEST_SPEC.mapping_bytes) handback = Mock(return_value=_RecoveryMirror(("",))) else: expected = RecoveryJournalError - tier_config = _TierConfig([("", 16)], None) + tier_config = _tier(16) handback = Mock(side_effect=RecoveryJournalError("written for other terms")) monkeypatch.setattr("kvcr.guard.read_handback", handback) @@ -609,7 +625,7 @@ def test_only_a_claim_refused_before_the_pool_moves_costs_nothing( guard._hand_back.assert_not_called() # Nothing was chosen, so a corrected claim can still have this pool. assert guard._configured is None - guard._refuse_incompatible(_TierConfig([("", 32)], None)) + guard._refuse_incompatible(_tier(32)) control.close.assert_called_once_with() @@ -622,7 +638,7 @@ def test_a_handback_with_an_unexpected_storage_error_fails() -> None: guard._recovery._write_handback = Mock(side_effect=error) with pytest.raises(OSError) as raised: - guard._hand_back(16) + guard._hand_back() assert raised.value is error @@ -633,7 +649,7 @@ def test_a_handback_without_a_mirror_does_not_close_the_core() -> None: core = _give_serving_core(guard) with pytest.raises(RecoveryMirrorError, match="no state to hand back"): - guard._hand_back(16) + guard._hand_back() core.close.assert_not_called() @@ -647,7 +663,7 @@ def test_a_handback_the_filesystem_refuses_leaves_a_cold_pool() -> None: side_effect=OSError(errno.ENOSPC, "No space left on device") ) - guard._hand_back(16) + guard._hand_back() assert guard._serving is False assert guard._core is None @@ -660,13 +676,14 @@ def test_a_dropped_handback_still_leaves_the_new_lease_mirrored(code: int) -> No guard = _configurable_guard() guard._control = None guard._failure_callback = lambda *_args: None - guard._configured = _TierConfig([("", 16)], None) + guard._configured = _tier(16) guard._recovery.mirror = _RecoveryMirror(("",)) + guard._recovery.attachment = _fake_attachment() _give_serving_core(guard) guard._recovery._journal = _Journal() guard._recovery._write_handback = Mock(side_effect=OSError(code, "No space left")) - guard._adopt(Mock(), _TierConfig([("", 16)], None)) + guard._adopt(Mock(), _tier(16)) # The pool went cold, not fatal: the Guard stood down and dropped the core. assert guard._serving is False @@ -696,19 +713,19 @@ def test_a_grant_that_never_arrived_resumes_the_guard_it_stood_down() -> None: guard._promote = lambda: outcomes.append("promote") guard._release = lambda: outcomes.append("release") - lease = Mock() + lease = Mock(close=Mock(side_effect=lambda: outcomes.append("close"))) guard._pool_lease.current = lease guard._abort(lease) - assert outcomes == ["promote"] + assert outcomes == ["promote", "close"] lease.close.assert_called_once_with() assert guard._pool_lease.current is None assert guard._phase is _Phase.STANDBY guard._resumable = False - stale = Mock() + stale = Mock(close=Mock(side_effect=lambda: outcomes.append("close"))) guard._pool_lease.current = stale guard._abort(stale) - assert outcomes == ["promote", "release"] + assert outcomes == ["promote", "close", "release", "close"] assert guard._phase is _Phase.IDLE @@ -720,7 +737,7 @@ def test_a_release_drops_its_mirror_after_handing_back_what_it_can( guard = _configurable_guard() control = Mock() guard._control = control - guard._configured = _TierConfig([("", 16)], None) + guard._configured = _tier(16) guard._recovery.mirror = _RecoveryMirror(("",)) if mode == "serving": _give_serving_core(guard) @@ -741,9 +758,8 @@ def test_a_release_drops_its_mirror_after_handing_back_what_it_can( assert guard._recovery.mirror is None control.close.assert_called_once_with() if mode == "accepts": - records, pool_layouts = guard._recovery._write_handback.call_args.args + (records,) = guard._recovery._write_handback.call_args.args assert set(records) == {BlockKey(b"published"), BlockKey(b"tail")} - assert pool_layouts == [("", 16)] elif mode == "serving": assert guard._serving is False assert guard._core is None diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 3469f8b..68949ab 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -28,7 +28,7 @@ free_port, ) -from kvcr import KVCR, KVCRBindings +from kvcr import KVCR, KVCRBindings, KVCRClient from kvcr import progress as kvcr_progress from kvcr.config import ( FrameworkDramInput, @@ -39,8 +39,11 @@ RemoteFWDramOptions, ) from kvcr.control_channels import ZmqPeerControlChannel +from kvcr.core import _BlockRecord from kvcr.guard import _Guard from kvcr.kvcr_service import _KVCRService +from kvcr.local_dram import _LocalDramResidency, _LocalDramState +from kvcr.recovery_journal import RecoveryJournal, _recovery_frames, read_handback from kvcr.types import BlockKey, CacheTier, QueryStatus _TIMEOUT_SECONDS = 5 @@ -192,6 +195,32 @@ def _primary_child( time.sleep(60) +def _group_primary_child(socket_path: str, control_port: str) -> None: + """Claim one pool group, fill every pool, and publish one grouped slot.""" + page_size = os.sysconf("SC_PAGE_SIZE") + hold = KVCRClient(socket_path).claim( + 0, + [("pool0", page_size + page_size // 2), ("pool1", page_size)], + _DIGEST, + ("127.0.0.1", int(control_port)), + ) + for index, (_name, address, size_bytes) in enumerate(hold.local_dram.pools): + ctypes.memset(address, ord("A") + index, size_bytes) + record = _BlockRecord( + local_dram=_LocalDramResidency( + [("pool0", 0), ("pool1", 0)], _LocalDramState.READY + ) + ) + journal = RecoveryJournal(hold._attachment) + journal.publish( + *next( + iter(_recovery_frames({BlockKey(b"grouped"): record}, ("pool0", "pool1"))) + ) + ) + print("ready", flush=True) + time.sleep(60) + + def _stale_peer_child(control_port: str, probe_port: str) -> None: """A dead primary's peer: it sends into the pool's endpoint and must get a terminal refusal back, not silence until its operation deadline.""" @@ -226,16 +255,22 @@ def _stale_peer_child(control_port: str, probe_port: str) -> None: @pytest.fixture def live_service( tmp_path: Path, + request: pytest.FixtureRequest, ) -> Iterator[tuple[_KVCRService, Callable[..., subprocess.Popen[str]]]]: - """A one-pool service on its own thread; children it spawns die with it.""" + """A service on its own thread; children it spawns die with it.""" + pool_count = getattr(request, "param", 1) pool_dir = tmp_path / "pools" pool_dir.mkdir() + page_size = os.sysconf("SC_PAGE_SIZE") + pool_sizes = ( + (2 * page_size, page_size) if pool_count == 2 else (page_size,) * pool_count + ) service = _KVCRService( tmp_path / "service.sock", pool_dir, guard_count=1, - pool_sizes_bytes=(os.sysconf("SC_PAGE_SIZE"),), - journal_bytes=8192, + pool_sizes_bytes=pool_sizes, + journal_bytes=2 * page_size, compatibility_digest=_DIGEST, ) server_thread = threading.Thread(target=service.serve_forever) @@ -406,6 +441,71 @@ def test_a_promoted_guard_serves_real_nixl_transfers( replacement.close() +@pytest.mark.parametrize("live_service", [2], indirect=True) +def test_two_pool_group_survives_guard_failover_and_reclaim( + monkeypatch: pytest.MonkeyPatch, + live_service: tuple[_KVCRService, Callable[..., subprocess.Popen[str]]], +) -> None: + """One crash moves both pools to the Guard and one claim takes both back.""" + page_size = os.sysconf("SC_PAGE_SIZE") + control_port = free_port() + guard_agent = _FileBackedNixlAgent() + guard_agent.state = "DONE" + monkeypatch.setattr(kvcr_progress, "nixl_agent_config", lambda **kwargs: kwargs) + monkeypatch.setattr(kvcr_progress, "nixl_agent", lambda _name, _config: guard_agent) + service, spawn = live_service + + primary = spawn("_group_primary_child", service.socket_path, control_port) + _await_marker(primary, "ready") + guard = service._registry._guards[0] + pools = guard._recovery.pools + + primary.kill() + primary.wait(timeout=_TIMEOUT_SECONDS) + _wait_until(lambda: guard._serving, timeout=_TIMEOUT_SECONDS) + + key = BlockKey(b"grouped") + record = guard._core._block_record_map[key] + assert record.local_dram == _LocalDramResidency( + [("pool0", 0), ("pool1", 0)], _LocalDramState.READY + ) + assert guard._core._local_dram.memory_regions == ( + ( + guard._recovery.attachment.address + pools[0].offset_bytes, + page_size + page_size // 2, + ), + ( + guard._recovery.attachment.address + pools[1].offset_bytes, + page_size, + ), + ) + + replacement = KVCRClient(service.socket_path).claim( + 0, + [("pool0", page_size + page_size // 2), ("pool1", page_size)], + _DIGEST, + ("127.0.0.1", control_port), + ) + try: + recovered = read_handback( + replacement._attachment, + _DIGEST, + replacement._pools, + ).take_records() + recovered_record = recovered[key] + assert recovered_record.local_dram is not None + assert recovered_record.local_dram.slots == [("pool0", 0), ("pool1", 0)] + for index, (_name, address, _size_bytes) in enumerate( + replacement.local_dram.pools + ): + assert ( + ctypes.string_at(address, page_size) + == bytes((ord("A") + index,)) * page_size + ) + finally: + replacement.release() + + @pytest.mark.parametrize("recovery", ["kept", "given-up"]) def test_request_timeout_during_promotion_then_retry_uses_guard( tmp_path: Path, diff --git a/tests/unit/test_guard_protocol.py b/tests/unit/test_guard_protocol.py index 604c3f0..92de3f1 100644 --- a/tests/unit/test_guard_protocol.py +++ b/tests/unit/test_guard_protocol.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import errno +import mmap import os import select import socket @@ -25,6 +26,7 @@ _Error, _G3Config, _Granted, + _PoolDescriptor, _Release, _Released, _TierConfig, @@ -32,15 +34,22 @@ from kvcr.memory import _JOURNAL_HEADER_BYTES, KVCRPoolSpec _GUARD_INDEX = 3 -_BLOCK_SIZE_BYTES = 1024 +_POOL_LAYOUTS = [("pool0", 1024), ("pool1", 3072)] +_POOL_SIZES = (mmap.PAGESIZE, 2 * mmap.PAGESIZE) _GENERATION = "a" * 32 _DEVICE = 2049 _INODE = 42 _DIGEST = "opaque digest: leave unchanged" _JOURNAL_BYTES = 2 * _JOURNAL_HEADER_BYTES -_MAPPING_BYTES = _JOURNAL_BYTES + 8195 -_POOL_LAYOUTS = [("", _BLOCK_SIZE_BYTES)] +_POOL_OFFSETS = (_JOURNAL_BYTES, _JOURNAL_BYTES + _POOL_SIZES[0]) +_MAPPING_BYTES = _JOURNAL_BYTES + sum(_POOL_SIZES) _TIER_CONFIG = _TierConfig(_POOL_LAYOUTS, None) +_WIRE_POOLS = tuple( + _PoolDescriptor(name, size_bytes, block_size_bytes, offset_bytes) + for (name, block_size_bytes), size_bytes, offset_bytes in zip( + _POOL_LAYOUTS, _POOL_SIZES, _POOL_OFFSETS, strict=True + ) +) def test_close_swaps_the_pidfd_under_its_lock() -> None: @@ -150,10 +159,6 @@ def __init__( def address(self) -> int: return 1234 - @property - def data_address(self) -> int: - return self.address + _JOURNAL_BYTES - def close(self) -> None: if self._events is not None: self._events.append("attachment.close") @@ -166,6 +171,7 @@ def _grant( guard_index: int = _GUARD_INDEX, mapping_bytes: int = _MAPPING_BYTES, tier_config: _TierConfig = _TIER_CONFIG, + pools: tuple[_PoolDescriptor, ...] = _WIRE_POOLS, ) -> _Granted: return _Granted( guard_index, @@ -179,10 +185,26 @@ def _grant( journal_bytes=_JOURNAL_BYTES, ), tier_config, + pools, 1, ) +def _grant_with_pool(index: int = 0, **changes) -> _Granted: + pools = list(_WIRE_POOLS) + pools[index] = msgspec.structs.replace(pools[index], **changes) + return _grant(pools=tuple(pools)) + + +def _local_dram(address: int = 1234) -> LocalDramOptions: + return LocalDramOptions( + [ + (pool.name, address + pool.offset_bytes, pool.size_bytes) + for pool in _WIRE_POOLS + ] + ) + + def _connect_with( monkeypatch: pytest.MonkeyPatch, connection: _RecordingConnection, @@ -194,30 +216,62 @@ def _connect_with( ) -def test_unsupported_tier_configuration_is_refused_at_decode() -> None: - """Refuse unsupported pool and G3 layouts before they bind the pool.""" - page = os.sysconf("SC_PAGE_SIZE") +def test_pool_descriptor_constraints_are_part_of_the_wire_contract() -> None: + with pytest.raises(ValueError, match="positive"): + _PoolDescriptor("pool", 0, 1) + with pytest.raises(ValueError, match="non-negative"): + _PoolDescriptor("pool", 1, 1, -1) + with pytest.raises(ValueError, match="complete KV block"): + _PoolDescriptor("pool", 1023, 1024) + + +def test_collection_wire_shapes_are_nonempty_and_have_no_scalar_aliases() -> None: + claim_wire = msgspec.to_builtins( + _Claim(_GUARD_INDEX, _DIGEST, _TIER_CONFIG, "127.0.0.1", 5555, 1) + ) + empty_claim = {**claim_wire, "tier_config": {"pool_layouts": [], "g3": None}} + negative_claim = {**claim_wire, "guard_index": -1} + scalar_tier = {"row_stride": _POOL_LAYOUTS[0][1], "g3": None} + scalar_claim = { + **claim_wire, + "tier_config": scalar_tier, + } + scalar_claim["pool_index"] = scalar_claim.pop("guard_index") + scalar_grant = {**msgspec.to_builtins(_grant()), "tier_config": scalar_tier} + scalar_grant["pool_index"] = scalar_grant.pop("guard_index") + scalar_grant.pop("pools") + for decoder, wire in ( + (protocol_module._CLAIM_DECODER, empty_claim), + (protocol_module._CLAIM_DECODER, negative_claim), + (protocol_module._CLAIM_DECODER, scalar_claim), + (protocol_module._CLAIM_RESPONSE_DECODER, scalar_grant), + ): + with pytest.raises(msgspec.ValidationError): + decoder.decode(msgspec.msgpack.encode(wire)) + + +def test_g3_config_keeps_its_intrinsic_path_checks() -> None: good = { "paths": ("/g3/a",), - "capacity_bytes_per_file": page, + "capacity_bytes_per_file": 1, "backend": "FILE", "backend_options": {}, } - with pytest.raises(ValueError, match="only a single pool"): - _TierConfig([("full", page), ("swa", page)], None) - with pytest.raises(ValueError, match="page aligned"): - _TierConfig([("", page // 2)], _G3Config(**good)) - with pytest.raises(ValueError, match="complete slots"): + with pytest.raises(ValueError, match="absolute"): + _G3Config(**{**good, "paths": ("g3/a",)}) + with pytest.raises(ValueError, match="unique"): + _G3Config(**{**good, "paths": ("/g3/a", "/g3//a")}) + + with pytest.raises(ValueError, match="does not support multiple pools"): _TierConfig( - [("", page)], - _G3Config(**{**good, "capacity_bytes_per_file": page + 1}), + _POOL_LAYOUTS, + _G3Config(**{**good, "capacity_bytes_per_file": mmap.PAGESIZE}), ) - with pytest.raises(ValueError, match="unique"): + with pytest.raises(ValueError, match="complete slots"): _TierConfig( - [("", page)], - _G3Config(**{**good, "paths": ("/g3/a", "/g3//a")}), + [("pool", mmap.PAGESIZE)], + _G3Config(**{**good, "capacity_bytes_per_file": mmap.PAGESIZE + 1}), ) - _TierConfig([("", page)], _G3Config(**good)) def test_claim_and_release_round_trip_typed_messages_and_geometry( @@ -240,7 +294,7 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( "guard_index": _GUARD_INDEX, "compatibility_digest": _DIGEST, "tier_config": { - "pool_layouts": [("", _BLOCK_SIZE_BYTES)], + "pool_layouts": _POOL_LAYOUTS, "g3": None, "remote_fw_dram_backend": "UCX", }, @@ -248,27 +302,18 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( "control_port": 5555, "version": 1, } - assert msgspec.to_builtins(_grant()) == { - "type": "granted", - "guard_index": _GUARD_INDEX, - "spec": { - "pool_id": f"pool_{_GUARD_INDEX}", - "path": f"/tmp/kvcr-pool_{_GUARD_INDEX}-{_GENERATION}", - "generation": _GENERATION, - "device": _DEVICE, - "inode": _INODE, - "mapping_bytes": _MAPPING_BYTES, - "journal_bytes": _JOURNAL_BYTES, - }, - "tier_config": { - "pool_layouts": [("", _BLOCK_SIZE_BYTES)], - "g3": None, - "remote_fw_dram_backend": "UCX", - }, - "version": 1, + grant_wire = msgspec.to_builtins(_grant()) + assert grant_wire["type"] == "granted" + assert grant_wire["version"] == 1 + assert grant_wire["guard_index"] == _GUARD_INDEX + assert grant_wire["tier_config"] == { + "pool_layouts": _POOL_LAYOUTS, + "g3": None, + "remote_fw_dram_backend": "UCX", } + assert grant_wire["pools"] == msgspec.to_builtins(_WIRE_POOLS) attach.assert_called_once_with(_grant().spec) - assert hold.local_dram == LocalDramOptions([("", 1234 + _JOURNAL_BYTES, 8192)]) + assert hold.local_dram == _local_dram() # The endpoint a Guard will answer on, handed over with the grant. assert hold._control_listener_fd == connection.handed_fd @@ -302,14 +347,28 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( [ pytest.param(_grant(guard_index=_GUARD_INDEX + 1), None, id="wrong-guard"), pytest.param( - _grant(tier_config=_TierConfig([("", _BLOCK_SIZE_BYTES * 2)], None)), + _grant( + tier_config=_TierConfig( + [("pool0", _POOL_LAYOUTS[0][1] * 2), _POOL_LAYOUTS[1]], None + ) + ), None, - id="wrong-pool-layout", + id="wrong-tier-layout", ), + pytest.param(_grant(pools=_WIRE_POOLS[:1]), None, id="wrong-pool-count"), pytest.param( - _grant(mapping_bytes=_JOURNAL_BYTES + _BLOCK_SIZE_BYTES - 1), + _grant_with_pool(block_size_bytes=_POOL_LAYOUTS[0][1] * 2), None, - id="short-mapping", + id="wrong-pool-block-size", + ), + pytest.param( + _grant_with_pool(1, offset_bytes=_POOL_OFFSETS[1] + mmap.PAGESIZE), + None, + id="noncontiguous-offset", + ), + pytest.param(_grant_with_pool(size_bytes=_POOL_SIZES[0] + 1), None, id="size"), + pytest.param( + _grant(mapping_bytes=_MAPPING_BYTES + mmap.PAGESIZE), None, id="mapping" ), pytest.param( KVCRGuardProtocolError("invalid granted message"), @@ -363,7 +422,8 @@ def test_release_failures_leave_a_retry_and_report_a_lost_acknowledgement() -> N [ConnectionResetError("release acknowledgement was lost")], events ) hold = KVCRPoolHold( - local_dram=LocalDramOptions([("", 1234, 8192)]), + local_dram=_local_dram(attachment.address), + _pools=_WIRE_POOLS, _attachment=attachment, _connection=connection, ) diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index 7968830..865bbde 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -4,6 +4,7 @@ import ctypes import logging +import mmap import threading from contextlib import nullcontext, suppress from functools import partial @@ -48,6 +49,7 @@ def _fake_hold(**fields: Any) -> SimpleNamespace: """A hold double that hands its listener over exactly like the real one.""" + fields.setdefault("_pools", ()) hold = SimpleNamespace(**fields) hold.hand_listener_to = partial(KVCRPoolHold.hand_listener_to, hold) return hold @@ -65,7 +67,7 @@ def test_local_dram_observer_reports_only_stable_slot_changes() -> None: backend = kvcr._core._local_dram assert backend is not None keys = tuple(BlockKey(f"k{index}".encode()) for index in range(3)) - observed: list[tuple[BlockKey, list[tuple[str, int]] | None]] = [] + observed: list[tuple[BlockKey, int | None]] = [] def observe(key: BlockKey, record: _BlockRecord) -> None: residency = record.local_dram @@ -135,6 +137,8 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: [ ("control-absent", ValueError, "share its control endpoint", []), ("control-cannot-share", ValueError, "share its control endpoint", []), + ("g3-invalid", ValueError, "page aligned", []), + ("g3-multi-pool", ValueError, "does not support multiple pools", []), ( "handback-unreadable", RuntimeError, @@ -151,12 +155,14 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: ids=[ "control-absent", "control-cannot-share", + "g3-invalid", + "g3-multi-pool", "handback-unreadable", "install-fails", ], ) def test_a_guarded_startup_that_fails_gives_back_everything_it_took( - monkeypatch, stage, error, match, expected_events + tmp_path, monkeypatch, stage, error, match, expected_events ) -> None: """Refused before the claim, or unwound after it: core closed, pool returned.""" events: list[str] = [] @@ -179,10 +185,18 @@ def claim(*_args, **_kwargs) -> SimpleNamespace: # taken. control: Any = Mock() control.control_bind_address.return_value = ("127.0.0.1", 5555) + backend_configs = KVCRBackendConfigs() if stage == "control-absent": control = None elif stage == "control-cannot-share": control = SimpleNamespace(control_bind_address=None, adopt_listener=None) + elif stage in ("g3-invalid", "g3-multi-pool"): + backend_configs = KVCRBackendConfigs( + g3=G3Options( + paths=(tmp_path / "g3",), + capacity_bytes_per_file=mmap.PAGESIZE, + ) + ) elif stage == "handback-unreadable": # The lease is live well before the caller is handed anything. monkeypatch.setattr( @@ -219,11 +233,15 @@ def is_quiescent(self) -> bool: KVCR( KVCRConfig( nixl_agent_name="target", - pool_layouts=[("", 1024)], + pool_layouts=( + [("full", 1024), ("swa", 1024)] + if stage == "g3-multi-pool" + else [("", mmap.PAGESIZE // 2 if stage == "g3-invalid" else 1024)] + ), nixl_listen_port=1, ), KVCRBindings(Mock(), Mock(), Mock(), framework_control=control), - KVCRBackendConfigs(), + backend_configs, _GUARD_CONFIG, ) @@ -338,8 +356,8 @@ def make_journal(pool) -> object: events.append("journal") return journal - def attach_journal(local, configured_journal, pool_names, disk) -> None: - assert (local, configured_journal, pool_names, disk) == ( + def attach_journal(local, configured_journal, pool_name, disk) -> None: + assert (local, configured_journal, pool_name, disk) == ( local_dram, journal, ("",), @@ -368,14 +386,14 @@ def attach_journal(local, configured_journal, pool_names, disk) -> None: primary_control.adopt_listener.side_effect = lambda fd: events.append(f"adopt:{fd}") g3_config = G3Options( paths=(tmp_path / "g3",), - capacity_bytes_per_file=8192, + capacity_bytes_per_file=2 * mmap.PAGESIZE, ) backend_configs = KVCRBackendConfigs( g3=g3_config, remote_fw_dram=RemoteFWDramOptions(backend="REMOTE"), ) controller = KVCR( - KVCRConfig(nixl_agent_name="target", pool_layouts=[("", 1024)]), + KVCRConfig(nixl_agent_name="target", pool_layouts=[("", mmap.PAGESIZE)]), KVCRBindings(Mock(), Mock(), Mock(), framework_control=primary_control), backend_configs, KVCRGuardConfig( @@ -387,7 +405,7 @@ def attach_journal(local, configured_journal, pool_names, disk) -> None: claim.assert_called_once_with( 3, - [("", 1024)], + [("", mmap.PAGESIZE)], "Opaque-Digest", ("127.0.0.1", 5555), g3_config, @@ -436,6 +454,19 @@ def test_kvcr_rejects_no_dram_backends() -> None: ) +def test_kvcr_accepts_multi_pool_layouts() -> None: + kvcr = _new_kvcr( + FakeNixlAgent(), + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig( + nixl_agent_name="target", + pool_layouts=[("full", 8), ("swa", 4)], + ), + ) + assert kvcr._core.pool_layouts == [("full", 8), ("swa", 4)] + + def test_kvcr_rejects_ambiguous_pool_names() -> None: bindings = KVCRBindings(Mock(), Mock(), Mock()) for pool_layouts, message in ( @@ -659,7 +690,7 @@ def test_resident_records_carry_no_instance_dictionary() -> None: _BlockRecord(), _LocalDramResidency([("", 0)], _LocalDramState.READY), _G3Residency(0), - _FwMemResidency(_mem_descriptor(), object()), + _FwMemResidency([_mem_descriptor()], object()), ): assert not hasattr(residency, "__dict__"), type(residency).__name__ diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index 79b6fce..bb760e9 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -48,7 +48,7 @@ def _holders_of(registry) -> dict[int, object]: - """The pools a worker holds, in the shape the old binding map had.""" + """The Guards a worker holds, in the shape the old binding map had.""" return { i: p._pool_lease.current for i, p in registry._guards.items() @@ -58,18 +58,15 @@ def _holders_of(registry) -> dict[int, object]: _SERVER_STOP_TIMEOUT_SECONDS = 5 _CONNECTION_POLL_INTERVAL_SECONDS = 0.001 +_PAGE_BLOCK_SIZE_BYTES = os.sysconf("SC_PAGE_SIZE") _TEST_GUARD_COUNT = 2 -_TEST_JOURNAL_BYTES = 8192 -_TEST_POOL_SIZE_BYTES = 8192 -_TEST_POOL_SIZES_BYTES = (_TEST_POOL_SIZE_BYTES,) +_TEST_JOURNAL_BYTES = 2 * _PAGE_BLOCK_SIZE_BYTES +_TEST_POOL_SIZES_BYTES = (2 * _PAGE_BLOCK_SIZE_BYTES,) _TEST_BLOCK_SIZE_BYTES = 1024 -_TEST_POOL_LAYOUTS = [("", _TEST_BLOCK_SIZE_BYTES)] +_TEST_POOL_LAYOUTS = [("pool0", _TEST_BLOCK_SIZE_BYTES)] _TEST_DIGEST = "opaque digest: Preserve-Me EXACTLY" _TEST_TIER_CONFIG = _TierConfig(_TEST_POOL_LAYOUTS, None) -# G3 terms are refused at decode unless a real claimant could open them, so -# the one claim that carries G3 uses a page-aligned stride. -_PAGE_STRIDE = os.sysconf("SC_PAGE_SIZE") class _FakeLiveness: @@ -96,7 +93,7 @@ def close(self) -> None: def _claim(registry, guard_index, liveness, control_bind=None): """Claim through the registry, closing the granted fd the tests never send.""" - spec, listener_fd, lease = registry.claim( + spec, _pools, listener_fd, lease = registry.claim( guard_index, _TEST_TIER_CONFIG, liveness, @@ -107,7 +104,7 @@ def _claim(registry, guard_index, liveness, control_bind=None): def _kill_and_wait(registry, guard_index, liveness) -> None: - """Die the way a real claimant does: the pool's own actor notices.""" + """Die the way a real claimant does: the Guard's own actor notices.""" liveness.kill() guard = registry._guards[guard_index] _wait_until( @@ -217,12 +214,16 @@ def _take(duplicate: socket.socket): def _stand_in_pool(spec) -> Mock: """The pool-tail surface a Guard reaches for, without a real mapping.""" - attachment = Mock(address=1234, data_address=1234 + spec.journal_bytes, _spec=spec) + attachment = Mock(address=1234, _spec=spec) attachment.mapped_snapshot.return_value = nullcontext(None) return attachment -def _new_registry(tmp_path: Path, guard_count: int = 1) -> _PoolRegistry: +def _new_registry( + tmp_path: Path, + guard_count: int = 1, + pool_sizes_bytes: tuple[int, ...] = _TEST_POOL_SIZES_BYTES, +) -> _PoolRegistry: """A registry of real Guards over stand-in pool mappings.""" journal = Mock() journal.read_next.return_value = None @@ -233,7 +234,7 @@ def _new_registry(tmp_path: Path, guard_count: int = 1) -> _PoolRegistry: return _PoolRegistry( tmp_path, guard_count, - _TEST_POOL_SIZES_BYTES, + pool_sizes_bytes, _TEST_JOURNAL_BYTES, _TEST_DIGEST, ) @@ -257,13 +258,38 @@ def test_socket_is_private(tmp_path: Path) -> None: assert stat.S_IMODE(harness.server.socket_path.stat().st_mode) == 0o600 -def test_each_guard_allocation_contains_all_pool_sizes(tmp_path: Path) -> None: - pool_sizes = (2 * _PAGE_STRIDE, 3 * _PAGE_STRIDE) - with _running_server(tmp_path, pool_sizes_bytes=pool_sizes) as harness: - expected = _TEST_JOURNAL_BYTES + sum(pool_sizes) - assert all( - guard._owner.spec.mapping_bytes == expected - for guard in harness.server._registry._guards.values() +def test_client_claims_one_grouped_allocation_with_independent_strides( + tmp_path: Path, +) -> None: + pool_sizes = (2 * _PAGE_BLOCK_SIZE_BYTES, 3 * _PAGE_BLOCK_SIZE_BYTES) + pool_layouts = [ + ("pool0", _TEST_BLOCK_SIZE_BYTES), + ("pool1", 3 * _TEST_BLOCK_SIZE_BYTES), + ] + with _running_server( + tmp_path, + guard_count=2, + pool_sizes_bytes=pool_sizes, + ) as harness: + with pytest.raises(KVCRServiceError, match="out of range"): + harness.client.claim(2, _TEST_POOL_LAYOUTS, _TEST_DIGEST, _control_bind(2)) + with pytest.raises(KVCRServiceError, match="pool layout"): + harness.client.claim(1, _TEST_POOL_LAYOUTS, _TEST_DIGEST, _control_bind(1)) + + hold = harness.client.claim(1, pool_layouts, _TEST_DIGEST, _control_bind(1)) + hold.release() + guard = harness.server._registry._guards[1] + spec = guard._owner.spec + assert spec.mapping_bytes == _TEST_JOURNAL_BYTES + sum(pool_sizes) + offsets = (_TEST_JOURNAL_BYTES, _TEST_JOURNAL_BYTES + pool_sizes[0]) + assert tuple( + (pool.name, pool.size_bytes, pool.block_size_bytes, pool.offset_bytes) + for pool in guard._recovery.pools + ) == tuple( + (name, size, block_size, offset) + for (name, block_size), size, offset in zip( + pool_layouts, pool_sizes, offsets, strict=True + ) ) @@ -275,9 +301,9 @@ def test_registry_lifecycle_from_independent_leases_to_a_wedged_close( guard = registry._guards[0] first, second, third = _FakeLiveness(), _FakeLiveness(), _FakeLiveness() first_spec, stale = _claim(registry, 0, first) - _spec, _fd, _lease = registry.claim( + _spec, _pools, _fd, _lease = registry.claim( 1, - _TierConfig([("", _TEST_BLOCK_SIZE_BYTES * 2)], None), + _TierConfig([("pool0", _TEST_BLOCK_SIZE_BYTES * 2)], None), second, _control_bind(1), ) @@ -387,7 +413,7 @@ def _take_and_record(duplicate: socket.socket): g3 = G3Options( paths=((tmp_path / "g3").resolve(),), - capacity_bytes_per_file=2 * _PAGE_STRIDE, + capacity_bytes_per_file=2 * _PAGE_BLOCK_SIZE_BYTES, backend="FILE", backend_options={"mode": "direct"}, ) @@ -403,11 +429,11 @@ def _take_and_record(duplicate: socket.socket): assert guard._phase is _Phase.UNCONFIGURED hold = harness.client.claim( - 1, [("", _PAGE_STRIDE)], _TEST_DIGEST, _control_bind(), g3 + 1, [("pool0", _PAGE_BLOCK_SIZE_BYTES)], _TEST_DIGEST, _control_bind(), g3 ) assert guard._phase is _Phase.PRIMARY and guard._control is control assert guard._configured == _TierConfig( - [("", _PAGE_STRIDE)], + [("pool0", _PAGE_BLOCK_SIZE_BYTES)], _G3Config( paths=(str(g3.paths[0]),), capacity_bytes_per_file=g3.capacity_bytes_per_file, @@ -517,7 +543,7 @@ def test_a_standby_survives_failed_claims_and_hands_over_to_a_replacement( with pytest.raises(KVCRServiceError, match="another tier configuration"): registry.claim( 0, - _TierConfig([("", _TEST_BLOCK_SIZE_BYTES * 2)], None), + _TierConfig([("pool0", _TEST_BLOCK_SIZE_BYTES * 2)], None), _FakeLiveness(), control_bind, ) @@ -608,7 +634,7 @@ def test_claim_refusals_and_internal_failures_do_not_bind( with pytest.raises(KVCRServiceError, match="one complete KV block"): harness.client.claim( 0, - [("", _TEST_POOL_SIZE_BYTES + 1)], + [("pool0", _TEST_POOL_SIZES_BYTES[0] + 1)], _TEST_DIGEST, _control_bind(), ) @@ -694,7 +720,8 @@ def test_fork_and_exec_do_not_preserve_claimant_access( "import time", "from kvcr.guard_protocol import KVCRClient", f"hold = KVCRClient({str(harness.server.socket_path)!r}).claim(" - f"0, {_TEST_POOL_LAYOUTS!r}, {_TEST_DIGEST!r}, {_control_bind()!r})", + f"0, {_TEST_POOL_LAYOUTS!r}, {_TEST_DIGEST!r}, " + f"{_control_bind()!r})", "forked_pid = os.fork()", "if forked_pid == 0:", " hold._connection.close()", @@ -1002,7 +1029,6 @@ def test_promotion_failure_fails_the_pool_and_stops_the_whole_service( _claim(registry, 0, liveness) _kill_and_wait(registry, 0, liveness) - assert server._fatal_error is failure server.shutdown.assert_called_once_with() assert registry._refusing.is_set() is True @@ -1253,12 +1279,12 @@ def _service_args(pool_sizes_gb: str) -> list[str]: def test_pool_size_list_preserves_order_and_floors_each_item_to_pages() -> None: - raw_sizes = (2 * _PAGE_STRIDE + 123, 3 * _PAGE_STRIDE + 456) + raw_sizes = (2 * _PAGE_BLOCK_SIZE_BYTES + 123, 3 * _PAGE_BLOCK_SIZE_BYTES + 456) parsed = _parse_args( _service_args(",".join(str(size / (1 << 30)) for size in raw_sizes)) ) - expected = (2 * _PAGE_STRIDE, 3 * _PAGE_STRIDE) + expected = (2 * _PAGE_BLOCK_SIZE_BYTES, 3 * _PAGE_BLOCK_SIZE_BYTES) assert parsed.pool_sizes_bytes == expected @@ -1269,7 +1295,7 @@ def test_pool_size_list_preserves_order_and_floors_each_item_to_pages() -> None: "1,,2", "nan", "0", - str((_PAGE_STRIDE - 1) / (1 << 30)), + str((_PAGE_BLOCK_SIZE_BYTES - 1) / (1 << 30)), "1e1000000", ",".join([str(sys.maxsize // (1 << 30))] * 2), ], diff --git a/tests/unit/test_kvcr_service_workflow.py b/tests/unit/test_kvcr_service_workflow.py index 0e579ab..333046c 100644 --- a/tests/unit/test_kvcr_service_workflow.py +++ b/tests/unit/test_kvcr_service_workflow.py @@ -3,6 +3,7 @@ """Whole-workflow tests for the standalone KVCR service daemon.""" import ctypes +import mmap import signal import subprocess import sys @@ -34,10 +35,11 @@ from kvcr.kvcr_service import _DEFAULT_JOURNAL_BYTES, _KVCRService _BLOCK_SIZE_BYTES = 1024 +_POOL_LAYOUTS = [("pool0", _BLOCK_SIZE_BYTES)] _DIGEST = "opaque workflow digest: Preserve-Me EXACTLY" -_JOURNAL_BYTES = 8192 -_POOL_SIZE_BYTES = 8192 -_CLI_POOL_SIZE_BYTES = 8192 +_JOURNAL_BYTES = 2 * mmap.PAGESIZE +_POOL_SIZES_BYTES = (2 * mmap.PAGESIZE,) +_CLI_POOL_SIZE_BYTES = 2 * mmap.PAGESIZE _CLI_POOL_SIZE_GB = str(_CLI_POOL_SIZE_BYTES / (1 << 30)) _STOP_TIMEOUT_SECONDS = 5.0 _START_TIMEOUT_SECONDS = 60.0 @@ -75,10 +77,7 @@ def _claim_when_ready( while time.monotonic() < deadline: try: return client.claim( - guard_index, - [("", _BLOCK_SIZE_BYTES)], - _DIGEST, - _control_bind(guard_index), + guard_index, _POOL_LAYOUTS, _DIGEST, _control_bind(guard_index) ) except KVCRSocketError: if process.poll() is not None: @@ -133,7 +132,7 @@ def _running_service( socket_path, pool_dir, guard_count=guard_count, - pool_sizes_bytes=(_POOL_SIZE_BYTES,), + pool_sizes_bytes=_POOL_SIZES_BYTES, compatibility_digest=_DIGEST, journal_bytes=_JOURNAL_BYTES, ) @@ -156,17 +155,14 @@ def test_pools_persist_bytes_and_a_held_pool_refuses_claims(tmp_path: Path) -> N with _running_service(pool_dir) as socket_path: client = KVCRClient(socket_path) - first = client.claim(0, [("", _BLOCK_SIZE_BYTES)], _DIGEST, _control_bind(0)) - second = client.claim(1, [("", _BLOCK_SIZE_BYTES)], _DIGEST, _control_bind(1)) + first = client.claim(0, _POOL_LAYOUTS, _DIGEST, _control_bind(0)) + second = client.claim(1, _POOL_LAYOUTS, _DIGEST, _control_bind(1)) try: - first_address = first.local_dram.pools[0][1] - assert first_address != second.local_dram.pools[0][1] - ctypes.memmove(first_address, payload, len(payload)) + assert first.local_dram.pools[0][1] != second.local_dram.pools[0][1] + ctypes.memmove(first.local_dram.pools[0][1], payload, len(payload)) first.release() - replacement = client.claim( - 0, [("", _BLOCK_SIZE_BYTES)], _DIGEST, _control_bind(0) - ) + replacement = client.claim(0, _POOL_LAYOUTS, _DIGEST, _control_bind(0)) try: assert ( ctypes.string_at(replacement.local_dram.pools[0][1], len(payload)) @@ -184,7 +180,7 @@ def test_pools_persist_bytes_and_a_held_pool_refuses_claims(tmp_path: Path) -> N controller = KVCR( KVCRConfig( nixl_agent_name="target", - pool_layouts=[("", _BLOCK_SIZE_BYTES)], + pool_layouts=_POOL_LAYOUTS, nixl_listen_port=1, ), KVCRBindings( @@ -202,15 +198,13 @@ def test_pools_persist_bytes_and_a_held_pool_refuses_claims(tmp_path: Path) -> N ) try: with pytest.raises(KVCRServiceError, match="held"): - client.claim(0, [("", _BLOCK_SIZE_BYTES)], _DIGEST, (host, port)) + client.claim(0, _POOL_LAYOUTS, _DIGEST, (host, port)) finally: controller.close() control.close() # Closing the worker released the pool: the next claim is served. - reclaimed = client.claim( - 0, [("", _BLOCK_SIZE_BYTES)], _DIGEST, (host, port) - ) + reclaimed = client.claim(0, _POOL_LAYOUTS, _DIGEST, (host, port)) reclaimed.release() finally: second.release() @@ -229,15 +223,12 @@ def test_cli_daemon_sets_geometry_and_restart_reclaims_only_unattached( # The deployed flags produce the requested pool and data geometry. pools = list(pool_dir.iterdir()) - assert len(pools) == 2, "--guard-count pools at startup" - pool_bytes = int(float(_CLI_POOL_SIZE_GB) * (1 << 30)) + assert len(pools) == 2, "--guard-count allocations at startup" + usable = _CLI_POOL_SIZE_BYTES assert all( - path.stat().st_size == _DEFAULT_JOURNAL_BYTES + pool_bytes - for path in pools + path.stat().st_size == _DEFAULT_JOURNAL_BYTES + usable for path in pools ) - pool_name, address, size_bytes = hold.local_dram.pools[0] - assert (pool_name, size_bytes) == ("", pool_bytes) - assert address > 0 + assert hold.local_dram.pools[0][2] == usable attached_pool = next(pool_dir.glob("kvcr-pool_0-*")) unclaimed_pool = next(pool_dir.glob("kvcr-pool_1-*")) diff --git a/tests/unit/test_recovery_journal.py b/tests/unit/test_recovery_journal.py index 76acd9e..d24ef76 100644 --- a/tests/unit/test_recovery_journal.py +++ b/tests/unit/test_recovery_journal.py @@ -8,10 +8,12 @@ from pathlib import Path from unittest.mock import Mock, patch +import msgspec import pytest from _kvcr_test_utils import _recovered_record from kvcr.core import _BlockRecord +from kvcr.guard_protocol import _PoolDescriptor from kvcr.memory import KVCRPoolAttachment, KVCRPoolSpec, _KVCRPoolOwner from kvcr.recovery_journal import ( _JOURNAL_HEADER_BYTES, @@ -283,12 +285,48 @@ def _write_slot(pool: KVCRPoolAttachment, terms: bytes, key: bytes, slot: int) - write_recovery_snapshot(pool, terms, frames) +def test_canonical_pool_terms_bind_ordered_geometry_and_allocation_identity() -> None: + spec = KVCRPoolSpec( + pool_id="pool_0", + path=f"/tmp/kvcr-pool_0-{_GENERATION}", + generation=_GENERATION, + device=7, + inode=11, + mapping_bytes=5 * mmap.PAGESIZE, + journal_bytes=2 * mmap.PAGESIZE, + ) + pools = ( + _PoolDescriptor("pool0", mmap.PAGESIZE, 1024, 2 * mmap.PAGESIZE), + _PoolDescriptor("pool1", 2 * mmap.PAGESIZE, 2048, 3 * mmap.PAGESIZE), + ) + + def terms_for(candidate=pools, digest=_TEST_DIGEST, allocation=spec): + return canonical_pool_terms(digest, candidate, allocation) + + terms = terms_for() + + for field, value in ( + ("name", "other"), + ("size_bytes", 8192), + ("block_size_bytes", 2048), + ("offset_bytes", 12288), + ): + changed = msgspec.structs.replace(pools[0], **{field: value}) + assert terms_for((changed, pools[1])) != terms + assert terms_for(tuple(reversed(pools))) != terms + assert terms_for(allocation=msgspec.structs.replace(spec, device=8)) != terms + + def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) -> None: """Replayed whole under its own terms, discardable when torn, gone once released.""" with _attached(tmp_path) as pool: path = Path(pool._spec.path) - pool_layouts = [("pool0", 4096)] - terms = canonical_pool_terms(_TEST_DIGEST, pool_layouts, pool._spec) + pools = ( + _PoolDescriptor( + "pool0", pool._spec.data_bytes, 4096, pool._spec.journal_bytes + ), + ) + terms = canonical_pool_terms(_TEST_DIGEST, pools, pool._spec) assert list(read_recovery_snapshot(pool, terms)) == [] records = { @@ -309,12 +347,9 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - assert mirror.take_records() == records # A slot number only means the same bytes under the same geometry. - for other in ( - canonical_pool_terms("another-digest", pool_layouts, pool._spec), - canonical_pool_terms(_TEST_DIGEST, [("pool0", 8192)], pool._spec), - ): - with pytest.raises(RecoveryJournalError, match="other terms"): - list(read_recovery_snapshot(pool, other)) + other = canonical_pool_terms("another-digest", pools, pool._spec) + with pytest.raises(RecoveryJournalError, match="other terms"): + list(read_recovery_snapshot(pool, other)) # Stopped once the replacing body has landed but before its header has. interrupted = Mock( @@ -337,7 +372,7 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - region[: _SNAPSHOT_HEADER.size] = bytes(_SNAPSHOT_HEADER.size) with pytest.raises(RecoveryJournalTornError, match="unfinished"): list(read_recovery_snapshot(pool, terms)) - assert read_handback(pool, _TEST_DIGEST, pool_layouts)._records == {} + assert read_handback(pool, _TEST_DIGEST, pools)._records == {} assert list(read_recovery_snapshot(pool, terms)) == [] # A released region is truncated away, so it replays nothing. From 1ccbdf307c7b7ded3f087f901e5cf362c23c14ee Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Wed, 9 Sep 2026 18:40:47 +0300 Subject: [PATCH 03/16] refactor(core): remove scalar block size Signed-off-by: Kapil Arya --- src/kvcr/core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 60c4506..e77dbcc 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -119,7 +119,6 @@ def __init__( self.pool_layouts = list(config.pool_layouts) _validate_pool_layouts(self.pool_layouts) self._block_sizes = dict(self.pool_layouts) - self.block_size_bytes = self.pool_layouts[0][1] if self.config.operation_timeout_ms <= 0: raise ValueError("operation_timeout_ms must be positive") if self.config.inventory_report_interval_ms < 0: @@ -223,7 +222,7 @@ def __init__( _G3( self, g3_config, - self.block_size_bytes, + self.pool_layouts[0][1], ) if g3_config is not None and local_dram_config is not None else None From 548dff27f62d4f534b9f7e1052766eaa523cf67e Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Wed, 9 Sep 2026 19:18:51 +0300 Subject: [PATCH 04/16] fix(core)!: report pressure per pool BREAKING CHANGE: capacity_needed_callback now receives an ordered list of pool name and slot count pairs. Signed-off-by: Kapil Arya --- docs/design_overview.md | 4 +- src/kvcr/api.py | 2 +- src/kvcr/core.py | 41 +++++++++++++------ src/kvcr/local_dram.py | 26 ++++++++---- src/kvcr/policy_runtime.py | 7 +++- src/kvcr/remote_fw_dram.py | 6 +-- tests/unit/_kvcr_test_utils.py | 2 + tests/unit/test_kvcr_local_dram.py | 66 +++++++++++++++++++++++++----- 8 files changed, 113 insertions(+), 41 deletions(-) diff --git a/docs/design_overview.md b/docs/design_overview.md index b6edc04..8d27e79 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -154,7 +154,7 @@ kvcr.abort(operation_handle, block_key_list=None) # cancel this kvcr.close() # synchronous teardown after framework-submitted jobs are drained # KVCR → Framework -framework.capacity_needed(num_slots) # last-resort capacity pressure signal +framework.capacity_needed([(pool_name, num_slots), ...]) # last-resort capacity pressure signal framework.request_pin(block_key_list) -> PinRequestId # enqueue a framework-owned source pin request framework.poll_pin_results() -> list[tuple[PinRequestId, PinResult]] # drain completed requests; one pin may cover the full list @@ -179,7 +179,7 @@ pool name may omit it. `deposit` also accepts a `no_evict` flag (batch-level, applies to all entries): when set, the KVCR keeps every completed slot non-evictable and returns a release handle per entry. A framework that wants guaranteed local DRAM residency behavior for selected KV blocks can get that behavior through `no_evict`, while the KVCR still handles sharing, routing visibility, transfer setup, and tiering policy. The framework calls `release` with the corresponding handle to clear the no-evict claim. -The tradeoff is backpressure: when policy cannot free enough capacity—for example, because `no_evict` claims occupy the pool or an attempted eviction does not free its source—KVCR may invoke `capacity_needed` as a last-resort pressure signal. The framework should release enough claims to free the requested slots. If sufficient capacity remains unavailable, affected committed entries complete with errors. Pool size and the free-slot threshold that triggers `capacity_needed` are deployment knobs. +The tradeoff is backpressure: when policy cannot free enough capacity—for example, because `no_evict` claims occupy a pool or an attempted eviction does not free its source—KVCR may invoke `capacity_needed` as a last-resort pressure signal. Each `(pool_name, num_slots)` request identifies the pool-local release batch the framework should satisfy. Pressure and its configured low watermark are tracked independently per pool. If sufficient capacity remains unavailable, affected committed entries complete with errors. Pool size and the free-slot threshold that triggers `capacity_needed` are deployment knobs. **Block availability query** diff --git a/src/kvcr/api.py b/src/kvcr/api.py index d120cb0..aa1599b 100644 --- a/src/kvcr/api.py +++ b/src/kvcr/api.py @@ -68,7 +68,7 @@ class KVCRBindings: inventory_sink: InventorySink | None = None # Capacity pressure, telemetry, and placement policy. - capacity_needed_callback: Callable[[int], None] | None = None + capacity_needed_callback: Callable[[list[tuple[str, int]]], None] | None = None stats_factory: Callable[[], TelemetryStats] | None = None policy: "KVCachePolicy | None" = None diff --git a/src/kvcr/core.py b/src/kvcr/core.py index e77dbcc..3c5d59b 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -162,7 +162,7 @@ def __init__( self._block_record_map: dict[BlockKey, _BlockRecord] = {} self._pending_inventory_events: list[InventoryEvent] = [] self._inventory_flush_deadline: float | None = None - self._capacity_pressure_active = False + self._capacity_pressure_pools: set[str] = set() self._closed = False self._outstanding_operations = 0 self._framework_pin_keys: dict[PinHandle, set[BlockKey]] = {} @@ -208,10 +208,17 @@ def __init__( if local_dram_config is not None else None ) - self._capacity_low_watermark_slots = ceil( - (self._local_dram._total_slots if self._local_dram is not None else 0) - * self.config.capacity_low_watermark_percent - / 100 + self._capacity_low_watermarks = { + name: ceil(count * self.config.capacity_low_watermark_percent / 100) + for name, count in ( + self._local_dram._slot_counts.items() + if self._local_dram is not None + else () + ) + } + self._capacity_pressure_enabled = ( + self._capacity_needed_callback is not None + and any(self._capacity_low_watermarks.values()) ) self._remote_fw_dram = _RemoteFWDram( self, @@ -632,18 +639,26 @@ def _flush_inventory(self, *, force: bool = False) -> None: for event in events: self._send_inventory(event) - def _update_capacity_pressure(self, reclaimable_slots: int) -> None: + def _update_capacity_pressure(self, reclaimable_slots: Mapping[str, int]) -> None: callback = self._capacity_needed_callback - if callback is None or self._capacity_low_watermark_slots == 0: + if callback is None or not self._capacity_pressure_enabled: return - if reclaimable_slots >= self._capacity_low_watermark_slots: - self._capacity_pressure_active = False - return - if self._capacity_pressure_active: + pressured = { + name + for name, watermark in self._capacity_low_watermarks.items() + if watermark and reclaimable_slots[name] < watermark + } + newly_pressured = pressured - self._capacity_pressure_pools + self._capacity_pressure_pools = pressured + if not newly_pressured: return - self._capacity_pressure_active = True + request = [ + (name, self._capacity_low_watermarks[name]) + for name, _ in self.pool_layouts + if name in newly_pressured + ] try: - callback(self._capacity_low_watermark_slots) + callback(request) except Exception: logger.warning("KVCR capacity callback failed", exc_info=True) diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index 44a38a5..f6fc4c5 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -153,6 +153,7 @@ def __init__( self._kvcr = kvcr self._backend = region.backend self._pools: dict[str, tuple[int, int, int]] = {} + self._slot_counts: dict[str, int] = {} self._free_slots: dict[str, deque[int]] = {} for (pool_name, address, length), (_, slot_size) in zip( region.pools, kvcr.pool_layouts @@ -165,8 +166,10 @@ def __init__( if not slot_count: raise ValueError("local DRAM pool must hold at least one block") self._pools[pool_name] = (address, length, slot_size) + self._slot_counts[pool_name] = slot_count self._free_slots[pool_name] = deque(range(slot_count)) self._evictable = _EvictionQueue() + self._evictable_slots: Counter[str] = Counter() self._unscored: set[BlockKey] = set() self._pending_residency_ops: dict[_OpId, _PendingResidencyOp] = {} self._pending_deliver_ops: dict[_OpId, _PendingDeliverOp] = {} @@ -973,7 +976,7 @@ def _new_public_claim( def _acquire_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: if residency.state is not _LocalDramState.READY: raise RuntimeError(f"cannot claim unready local DRAM entry {key!r}") - self._remove_evictable(key) + self._remove_evictable(key, residency) residency.claim_count += 1 def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: @@ -1035,9 +1038,7 @@ def _allocate_slots( skipped.add(key) continue size_bytes = self._size_bytes(residency.slots) - free_before = { - name: len(self._free_slots[name]) for name in required - } + free_before = {name: len(self._free_slots[name]) for name in required} decision, eviction_pending = self._kvcr._decide_eviction( self._kvcr._block_meta(key, record, size_bytes), CacheTier.LOCAL_G2, @@ -1064,7 +1065,7 @@ def _allocate_slots( return None, [], False for key, record, residency, size_bytes in victims: - self._evictable.remove(key) + self._remove_evictable(key, residency) record.local_dram = None self._residency_observer(key, record) self._kvcr._on_remove(self._kvcr._block_meta(key, record, size_bytes)) @@ -1090,11 +1091,13 @@ def _make_evictable(self, key: BlockKey) -> None: self._unscored.add(key) return self._unscored.discard(key) - self._evictable.insert(key, score, len(record.local_dram.slots)) + if self._evictable.insert(key, score, len(record.local_dram.slots)): + self._evictable_slots.update(name for name, _ in record.local_dram.slots) - def _remove_evictable(self, key: BlockKey) -> None: + def _remove_evictable(self, key: BlockKey, residency: _LocalDramResidency) -> None: self._unscored.discard(key) - self._evictable.remove(key) + if self._evictable.remove(key): + self._evictable_slots.subtract(name for name, _ in residency.slots) def _retry_unscored(self) -> None: for key in tuple(self._unscored): @@ -1132,6 +1135,11 @@ def _same_layout( ] def _update_capacity_pressure(self) -> None: + if not self._kvcr._capacity_pressure_enabled: + return self._kvcr._update_capacity_pressure( - sum(map(len, self._free_slots.values())) + self._evictable.total_weight + { + name: len(slots) + self._evictable_slots[name] + for name, slots in self._free_slots.items() + } ) diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index 4a747a6..b6eaa96 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -161,7 +161,7 @@ def __init__(self) -> None: def __len__(self) -> int: return len(self._live) - def insert(self, key: BlockKey, score: float, weight: int = 1) -> None: + def insert(self, key: BlockKey, score: float, weight: int = 1) -> bool: previous = self._live.get(key) if previous is not None: self.total_weight -= previous.weight @@ -170,11 +170,14 @@ def insert(self, key: BlockKey, score: float, weight: int = 1) -> None: self._live[key] = entry self.total_weight += weight heapq.heappush(self._heap, (entry.score, entry.sequence, key)) + return previous is None - def remove(self, key: BlockKey) -> None: + def remove(self, key: BlockKey) -> bool: entry = self._live.pop(key, None) if entry is not None: self.total_weight -= entry.weight + return True + return False def select(self, excluded: set[BlockKey]) -> BlockKey | None: skipped: list[tuple[float, int, BlockKey]] = [] diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index 9424030..d464543 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -982,9 +982,9 @@ def _submit_prepared_source_write( destination = source_pin.dst_descriptors[index] if source is None: break - if [ - (descriptor.info, descriptor.size) for descriptor in source - ] != [(descriptor.info, descriptor.size) for descriptor in destination]: + if [(descriptor.info, descriptor.size) for descriptor in source] != [ + (descriptor.info, descriptor.size) for descriptor in destination + ]: logger.warning( "KVCR start_write layout mismatch op=%d key=%r", source_pin.op_handle, diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index d173e95..051dd54 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -457,6 +457,7 @@ def _new_kvcr( local_dram: LocalDramOptions | None = None, g3: G3Options | None = None, inventory_sink=None, + capacity_needed_callback=None, policy=None, ) -> KVCR: config = replace( @@ -480,6 +481,7 @@ def _new_kvcr( framework_control=control, key_adapter=key_adapter, inventory_sink=inventory_sink, + capacity_needed_callback=capacity_needed_callback, policy=policy, stats_factory=(FakeTelemetryStats if config.enable_telemetry else None), ), diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 37ce8c4..f4d4c42 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -387,7 +387,7 @@ def test_local_claims_fetch_deliver_release_and_capacity() -> None: agent = FakeNixlAgent() policy = _RecordingFIFOPolicy() - capacity_requests: list[int] = [] + capacity_requests: list[list[tuple[str, int]]] = [] kvcr = _new_local_kvcr( agent, local, @@ -403,7 +403,7 @@ def test_local_claims_fetch_deliver_release_and_capacity() -> None: deposit = kvcr.deposit({first_key: [_mem_descriptor(primary_addr)]}, no_evict=True) _wait_until(lambda: bool(agent.transfers)) - assert capacity_requests == [1] + assert capacity_requests == [[("", 1)]] with pytest.raises(ValueError, match="expected layout"): kvcr.fetch((first_key, second_key), expected_layout=["unknown"]) fetch = kvcr.fetch((first_key,), expected_layout=[""]) @@ -467,7 +467,7 @@ def test_local_claims_fetch_deliver_release_and_capacity() -> None: replacement = kvcr.deposit( {second_key: [_mem_descriptor(primary_addr + block_size)]} ) - assert capacity_requests == [1, 1] + assert capacity_requests == [[("", 1)], [("", 1)]] assert _poll_until(kvcr, lambda results: bool(results)) == [ (replacement, _op_entries({second_key: True})) ] @@ -476,7 +476,7 @@ def test_local_claims_fetch_deliver_release_and_capacity() -> None: def test_capacity_needed_is_edge_triggered() -> None: local = ctypes.create_string_buffer(10) - capacity_requests: list[int] = [] + capacity_requests: list[list[tuple[str, int]]] = [] kvcr = _new_local_kvcr( FakeNixlAgent(), local, @@ -485,14 +485,58 @@ def test_capacity_needed_is_edge_triggered() -> None: capacity_needed_callback=capacity_requests.append, ) - kvcr._core._update_capacity_pressure(2) - kvcr._core._update_capacity_pressure(1) - kvcr._core._update_capacity_pressure(0) - assert capacity_requests == [2] + kvcr._core._update_capacity_pressure({"": 2}) + kvcr._core._update_capacity_pressure({"": 1}) + kvcr._core._update_capacity_pressure({"": 0}) + assert capacity_requests == [[("", 2)]] - kvcr._core._update_capacity_pressure(2) - kvcr._core._update_capacity_pressure(1) - assert capacity_requests == [2, 2] + kvcr._core._update_capacity_pressure({"": 2}) + kvcr._core._update_capacity_pressure({"": 1}) + assert capacity_requests == [[("", 2)], [("", 2)]] + + +def test_capacity_pressure_is_pool_local() -> None: + pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(16)] + source = ctypes.create_string_buffer(24) + capacity_requests: list[list[tuple[str, int]]] = [] + agent = FakeNixlAgent() + agent.state = "DONE" + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig( + nixl_agent_name="target", + pool_layouts=[("full", 8), ("swa", 8)], + capacity_low_watermark_percent=100, + ), + local_dram=LocalDramOptions( + [ + ("full", ctypes.addressof(pools[0]), 8), + ("swa", ctypes.addressof(pools[1]), 16), + ] + ), + capacity_needed_callback=capacity_requests.append, + ) + descriptors = [ + _mem_descriptor(ctypes.addressof(source), 8, info="full"), + _mem_descriptor(ctypes.addressof(source) + 8, 8, info="swa"), + _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"), + ] + + full = kvcr.deposit({BlockKey(b"full"): descriptors[:1]}, no_evict=True) + _poll_until(kvcr, lambda done: full in dict(done)) + assert capacity_requests == [[("full", 1)]] + + swa_key = BlockKey(b"swa") + swa = kvcr.deposit({swa_key: descriptors[1:]}) + _poll_until(kvcr, lambda done: swa in dict(done)) + assert capacity_requests == [[("full", 1)], [("swa", 2)]] + capacity_requests.clear() + + claim = kvcr.fetch((swa_key,), expected_layout=["swa", "swa"]) + _poll_until(kvcr, lambda done: claim in dict(done)) + assert capacity_requests == [[("swa", 2)]] @pytest.mark.parametrize( From 86c380725086bd6e109adac876bda58cd4301816 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Wed, 9 Sep 2026 23:58:26 +0300 Subject: [PATCH 05/16] refactor(policy): remove duplicate slot weights Signed-off-by: Kapil Arya --- src/kvcr/local_dram.py | 4 ++-- src/kvcr/policy_runtime.py | 10 ++-------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index f6fc4c5..b4560f7 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -252,7 +252,7 @@ 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": self._evictable.total_weight, + "local_g2_evictable_slots": sum(self._evictable_slots.values()), } def deposit( @@ -1091,7 +1091,7 @@ def _make_evictable(self, key: BlockKey) -> None: self._unscored.add(key) return self._unscored.discard(key) - if self._evictable.insert(key, score, len(record.local_dram.slots)): + if self._evictable.insert(key, score): self._evictable_slots.update(name for name, _ in record.local_dram.slots) def _remove_evictable(self, key: BlockKey, residency: _LocalDramResidency) -> None: diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index b6eaa96..ae1fba9 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -148,7 +148,6 @@ def on_remove(self, meta: BlockMeta) -> None: class _Entry: score: float sequence: int - weight: int class _EvictionQueue: @@ -156,26 +155,21 @@ def __init__(self) -> None: self._heap: list[tuple[float, int, BlockKey]] = [] self._live: dict[BlockKey, _Entry] = {} self._next_sequence = 0 - self.total_weight = 0 def __len__(self) -> int: return len(self._live) - def insert(self, key: BlockKey, score: float, weight: int = 1) -> bool: + def insert(self, key: BlockKey, score: float) -> bool: previous = self._live.get(key) - if previous is not None: - self.total_weight -= previous.weight - entry = _Entry(score, self._next_sequence, weight) + entry = _Entry(score, self._next_sequence) self._next_sequence += 1 self._live[key] = entry - self.total_weight += weight heapq.heappush(self._heap, (entry.score, entry.sequence, key)) return previous is None def remove(self, key: BlockKey) -> bool: entry = self._live.pop(key, None) if entry is not None: - self.total_weight -= entry.weight return True return False From 588a9e5f5d5b79b45f352adb94d78cad8086c104 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 00:07:53 +0300 Subject: [PATCH 06/16] refactor(core): simplify capacity accounting Signed-off-by: Kapil Arya --- src/kvcr/core.py | 18 ++++++++---------- src/kvcr/local_dram.py | 19 ++++++++++--------- src/kvcr/policy_runtime.py | 13 ++++--------- 3 files changed, 22 insertions(+), 28 deletions(-) diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 3c5d59b..694225d 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -209,17 +209,15 @@ def __init__( else None ) self._capacity_low_watermarks = { - name: ceil(count * self.config.capacity_low_watermark_percent / 100) - for name, count in ( - self._local_dram._slot_counts.items() - if self._local_dram is not None - else () + name: ceil( + length // block_size * self.config.capacity_low_watermark_percent / 100 + ) + for name, (_, length, block_size) in ( + self._local_dram._pools.items() if self._local_dram is not None else () ) } - self._capacity_pressure_enabled = ( - self._capacity_needed_callback is not None - and any(self._capacity_low_watermarks.values()) - ) + if not any(self._capacity_low_watermarks.values()): + self._capacity_needed_callback = None self._remote_fw_dram = _RemoteFWDram( self, backend_configs.remote_fw_dram, @@ -641,7 +639,7 @@ def _flush_inventory(self, *, force: bool = False) -> None: def _update_capacity_pressure(self, reclaimable_slots: Mapping[str, int]) -> None: callback = self._capacity_needed_callback - if callback is None or not self._capacity_pressure_enabled: + if callback is None: return pressured = { name diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index b4560f7..ad4bdd9 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -153,7 +153,6 @@ def __init__( self._kvcr = kvcr self._backend = region.backend self._pools: dict[str, tuple[int, int, int]] = {} - self._slot_counts: dict[str, int] = {} self._free_slots: dict[str, deque[int]] = {} for (pool_name, address, length), (_, slot_size) in zip( region.pools, kvcr.pool_layouts @@ -166,7 +165,6 @@ def __init__( if not slot_count: raise ValueError("local DRAM pool must hold at least one block") self._pools[pool_name] = (address, length, slot_size) - self._slot_counts[pool_name] = slot_count self._free_slots[pool_name] = deque(range(slot_count)) self._evictable = _EvictionQueue() self._evictable_slots: Counter[str] = Counter() @@ -976,7 +974,11 @@ def _new_public_claim( def _acquire_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: if residency.state is not _LocalDramState.READY: raise RuntimeError(f"cannot claim unready local DRAM entry {key!r}") - self._remove_evictable(key, residency) + if residency.claim_count == 0: + if key in self._unscored: + self._unscored.remove(key) + else: + self._remove_evictable(key, residency) residency.claim_count += 1 def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: @@ -1091,13 +1093,12 @@ def _make_evictable(self, key: BlockKey) -> None: self._unscored.add(key) return self._unscored.discard(key) - if self._evictable.insert(key, score): - self._evictable_slots.update(name for name, _ in record.local_dram.slots) + self._evictable.insert(key, score) + self._evictable_slots.update(name for name, _ in record.local_dram.slots) def _remove_evictable(self, key: BlockKey, residency: _LocalDramResidency) -> None: - self._unscored.discard(key) - if self._evictable.remove(key): - self._evictable_slots.subtract(name for name, _ in residency.slots) + self._evictable.remove(key) + self._evictable_slots.subtract(name for name, _ in residency.slots) def _retry_unscored(self) -> None: for key in tuple(self._unscored): @@ -1135,7 +1136,7 @@ def _same_layout( ] def _update_capacity_pressure(self) -> None: - if not self._kvcr._capacity_pressure_enabled: + if self._kvcr._capacity_needed_callback is None: return self._kvcr._update_capacity_pressure( { diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index ae1fba9..cea2ce1 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -159,19 +159,14 @@ def __init__(self) -> None: def __len__(self) -> int: return len(self._live) - def insert(self, key: BlockKey, score: float) -> bool: - previous = self._live.get(key) + def insert(self, key: BlockKey, score: float) -> None: entry = _Entry(score, self._next_sequence) self._next_sequence += 1 self._live[key] = entry heapq.heappush(self._heap, (entry.score, entry.sequence, key)) - return previous is None - def remove(self, key: BlockKey) -> bool: - entry = self._live.pop(key, None) - if entry is not None: - return True - return False + def remove(self, key: BlockKey) -> None: + self._live.pop(key, None) def select(self, excluded: set[BlockKey]) -> BlockKey | None: skipped: list[tuple[float, int, BlockKey]] = [] @@ -179,7 +174,7 @@ def select(self, excluded: set[BlockKey]) -> BlockKey | None: while self._heap: score, sequence, key = self._heap[0] entry = self._live.get(key) - if entry is None or (entry.score, entry.sequence) != (score, sequence): + if entry != _Entry(score, sequence): heapq.heappop(self._heap) continue if key not in excluded: From d402ccebc56e0de565864a549942ded7d41349b3 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 00:22:56 +0300 Subject: [PATCH 07/16] docs(core): clarify G3 layout limits Signed-off-by: Kapil Arya --- docs/dev-guide.md | 4 ++-- src/kvcr/core.py | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/dev-guide.md b/docs/dev-guide.md index fb6fcd9..5c264ae 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -357,8 +357,8 @@ pools do not have to agree on a block size. `KVCRConfig.pool_layouts` supplies the same ordered layouts to direct and `KVCRGuardConfig`-driven construction. Remote-transfer peers must use the same -pool names, block sizes, and order; a mismatch fails that operation. G3 remains -limited to a single-pool layout. +pool names, block sizes, and order; a mismatch fails that operation. G3 supports +one pool and exactly one block per key; repeated slots in that pool are unsupported. **A pool group's configuration is fixed by its first claim.** Every later claim on that Guard must name the same ordered pool layout and, when G3 is diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 694225d..86d200d 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -441,8 +441,6 @@ def fetch( name not in self._block_sizes for name in expected_layout ): raise ValueError("expected layout must use configured pools") - if self._g3 is not None and len(expected_layout) != 1: - raise ValueError("G3 does not support multi-block layouts") op_handle = self._next_op_handle self._next_op_handle += 1 local_dram = self._local_dram From bcf6c33b080393677b11ee7514fddc5a07e22bbb Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 00:45:45 +0300 Subject: [PATCH 08/16] refactor: trim multi-pool implementation Signed-off-by: Kapil Arya --- docs/dev-guide.md | 5 +++-- src/kvcr/core.py | 4 ++-- src/kvcr/guard_protocol.py | 2 -- src/kvcr/local_dram.py | 8 +++----- src/kvcr/recovery_journal.py | 12 +----------- src/kvcr/remote_fw_dram.py | 13 +++++-------- tests/unit/_kvcr_test_utils.py | 2 +- tests/unit/test_guard.py | 11 ++++++----- tests/unit/test_guard_integration.py | 20 ++++++-------------- tests/unit/test_guard_protocol.py | 27 --------------------------- tests/unit/test_kvcr.py | 13 ------------- tests/unit/test_kvcr_service.py | 1 + tests/unit/test_recovery_journal.py | 19 ++++--------------- 13 files changed, 32 insertions(+), 105 deletions(-) diff --git a/docs/dev-guide.md b/docs/dev-guide.md index 5c264ae..a345838 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -357,8 +357,9 @@ pools do not have to agree on a block size. `KVCRConfig.pool_layouts` supplies the same ordered layouts to direct and `KVCRGuardConfig`-driven construction. Remote-transfer peers must use the same -pool names, block sizes, and order; a mismatch fails that operation. G3 supports -one pool and exactly one block per key; repeated slots in that pool are unsupported. +pool names, block sizes, and order; a mismatch fails that operation. The current +G3 data plane stores one slot per key, so it supports one pool and one block per +key; multiple blocks in that pool remain out of scope. **A pool group's configuration is fixed by its first claim.** Every later claim on that Guard must name the same ordered pool layout and, when G3 is diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 86d200d..61af860 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -649,8 +649,8 @@ def _update_capacity_pressure(self, reclaimable_slots: Mapping[str, int]) -> Non if not newly_pressured: return request = [ - (name, self._capacity_low_watermarks[name]) - for name, _ in self.pool_layouts + (name, watermark) + for name, watermark in self._capacity_low_watermarks.items() if name in newly_pressured ] try: diff --git a/src/kvcr/guard_protocol.py b/src/kvcr/guard_protocol.py index 8520fbe..8074450 100644 --- a/src/kvcr/guard_protocol.py +++ b/src/kvcr/guard_protocol.py @@ -61,8 +61,6 @@ class _PoolDescriptor(msgspec.Struct, frozen=True, forbid_unknown_fields=True): def __post_init__(self) -> None: _compute_pool_geometry(self.size_bytes, self.block_size_bytes) - if type(self.offset_bytes) is not int or self.offset_bytes < 0: - raise ValueError("pool offset must be a non-negative integer") class _TierConfig(msgspec.Struct, frozen=True, forbid_unknown_fields=True): diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index ad4bdd9..d78db6e 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -189,10 +189,6 @@ def __init__( def memory_regions(self) -> tuple[tuple[int, int], ...]: return tuple((address, length) for address, length, _ in self._pools.values()) - @property - def _total_slots(self) -> int: - return sum(length // slot_size for _, length, slot_size in self._pools.values()) - def observe_residency( self, observer: Callable[[BlockKey, "_BlockRecord"], None] ) -> None: @@ -244,7 +240,9 @@ def rank_recovered(self, records: Mapping[BlockKey, "_BlockRecord"]) -> None: self._make_evictable(key) def telemetry_state(self) -> dict[str, int]: - total_slots = self._total_slots + total_slots = sum( + length // slot_size for _, length, slot_size in self._pools.values() + ) free_slots = sum(map(len, self._free_slots.values())) return { "local_g2_total_slots": total_slots, diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index c9b3946..20044af 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -602,17 +602,7 @@ def canonical_pool_terms( spec.device, spec.inode, ) - + msgspec.msgpack.encode( - [ - ( - pool.name, - pool.size_bytes, - pool.block_size_bytes, - pool.offset_bytes, - ) - for pool in pools - ] - ) + + msgspec.msgpack.encode(pools) ) diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index d464543..bada79a 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -16,6 +16,7 @@ from collections.abc import Collection, Iterator, Mapping from dataclasses import dataclass, field, replace from enum import Enum, auto +from itertools import chain from typing import TYPE_CHECKING, Any, cast import msgspec @@ -269,15 +270,11 @@ def progress( try: transfer_id, submitted = progress.submit_transfer( "WRITE", + tuple(chain.from_iterable(self.src_descriptors)), tuple( - descriptor - for descriptors in self.src_descriptors - for descriptor in descriptors - ), - tuple( - descriptor - for descriptors in self.dst_descriptors[: self.completed_count] - for descriptor in descriptors + chain.from_iterable( + self.dst_descriptors[: self.completed_count] + ) ), remote_side_agent=self.remote_agent, backend=backend._options.backend, diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 051dd54..aca210f 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -563,7 +563,7 @@ def decode(self, key): def _recovered_record( - *, g2: int | list[tuple[str, int]] | None = None, g3: int | None = None + *, g2: list[tuple[str, int]] | None = None, g3: int | None = None ) -> _BlockRecord: """A block record as recovery rebuilds one: settled residencies, nothing live.""" return _BlockRecord( diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 1618b8f..6f777de 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -81,6 +81,7 @@ def _fake_attachment() -> Mock: """A stand-in with the pool-tail surface a Guard reaches for.""" attachment = Mock( address=1234, + data_address=1234 + _TEST_SPEC.journal_bytes, _spec=_TEST_SPEC, ) attachment.mapped_snapshot.return_value = nullcontext(None) @@ -212,8 +213,8 @@ def test_a_serving_guard_reports_a_poll_failure_and_fences_its_core(caplog) -> N assert caplog.records[0].exc_info[1] is error core.close.assert_called_once_with() control.close.assert_not_called() - failure_callback.assert_called_once_with(guard, error) journal.invalidate.assert_called_once_with() + failure_callback.assert_called_once_with(guard, error) def test_standby_guard_failure_releases_adopted_listener() -> None: @@ -713,19 +714,19 @@ def test_a_grant_that_never_arrived_resumes_the_guard_it_stood_down() -> None: guard._promote = lambda: outcomes.append("promote") guard._release = lambda: outcomes.append("release") - lease = Mock(close=Mock(side_effect=lambda: outcomes.append("close"))) + lease = Mock() guard._pool_lease.current = lease guard._abort(lease) - assert outcomes == ["promote", "close"] + assert outcomes == ["promote"] lease.close.assert_called_once_with() assert guard._pool_lease.current is None assert guard._phase is _Phase.STANDBY guard._resumable = False - stale = Mock(close=Mock(side_effect=lambda: outcomes.append("close"))) + stale = Mock() guard._pool_lease.current = stale guard._abort(stale) - assert outcomes == ["promote", "close", "release", "close"] + assert outcomes == ["promote", "release"] assert guard._phase is _Phase.IDLE diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 68949ab..9a287ec 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -204,8 +204,6 @@ def _group_primary_child(socket_path: str, control_port: str) -> None: _DIGEST, ("127.0.0.1", int(control_port)), ) - for index, (_name, address, size_bytes) in enumerate(hold.local_dram.pools): - ctypes.memset(address, ord("A") + index, size_bytes) record = _BlockRecord( local_dram=_LocalDramResidency( [("pool0", 0), ("pool1", 0)], _LocalDramState.READY @@ -258,12 +256,13 @@ def live_service( request: pytest.FixtureRequest, ) -> Iterator[tuple[_KVCRService, Callable[..., subprocess.Popen[str]]]]: """A service on its own thread; children it spawns die with it.""" - pool_count = getattr(request, "param", 1) pool_dir = tmp_path / "pools" pool_dir.mkdir() page_size = os.sysconf("SC_PAGE_SIZE") pool_sizes = ( - (2 * page_size, page_size) if pool_count == 2 else (page_size,) * pool_count + (2 * page_size, page_size) + if getattr(request, "param", 1) == 2 + else (page_size,) ) service = _KVCRService( tmp_path / "service.sock", @@ -492,16 +491,9 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( _DIGEST, replacement._pools, ).take_records() - recovered_record = recovered[key] - assert recovered_record.local_dram is not None - assert recovered_record.local_dram.slots == [("pool0", 0), ("pool1", 0)] - for index, (_name, address, _size_bytes) in enumerate( - replacement.local_dram.pools - ): - assert ( - ctypes.string_at(address, page_size) - == bytes((ord("A") + index,)) * page_size - ) + assert recovered[key].local_dram == _LocalDramResidency( + [("pool0", 0), ("pool1", 0)], _LocalDramState.READY + ) finally: replacement.release() diff --git a/tests/unit/test_guard_protocol.py b/tests/unit/test_guard_protocol.py index 92de3f1..d756f63 100644 --- a/tests/unit/test_guard_protocol.py +++ b/tests/unit/test_guard_protocol.py @@ -219,37 +219,10 @@ def _connect_with( def test_pool_descriptor_constraints_are_part_of_the_wire_contract() -> None: with pytest.raises(ValueError, match="positive"): _PoolDescriptor("pool", 0, 1) - with pytest.raises(ValueError, match="non-negative"): - _PoolDescriptor("pool", 1, 1, -1) with pytest.raises(ValueError, match="complete KV block"): _PoolDescriptor("pool", 1023, 1024) -def test_collection_wire_shapes_are_nonempty_and_have_no_scalar_aliases() -> None: - claim_wire = msgspec.to_builtins( - _Claim(_GUARD_INDEX, _DIGEST, _TIER_CONFIG, "127.0.0.1", 5555, 1) - ) - empty_claim = {**claim_wire, "tier_config": {"pool_layouts": [], "g3": None}} - negative_claim = {**claim_wire, "guard_index": -1} - scalar_tier = {"row_stride": _POOL_LAYOUTS[0][1], "g3": None} - scalar_claim = { - **claim_wire, - "tier_config": scalar_tier, - } - scalar_claim["pool_index"] = scalar_claim.pop("guard_index") - scalar_grant = {**msgspec.to_builtins(_grant()), "tier_config": scalar_tier} - scalar_grant["pool_index"] = scalar_grant.pop("guard_index") - scalar_grant.pop("pools") - for decoder, wire in ( - (protocol_module._CLAIM_DECODER, empty_claim), - (protocol_module._CLAIM_DECODER, negative_claim), - (protocol_module._CLAIM_DECODER, scalar_claim), - (protocol_module._CLAIM_RESPONSE_DECODER, scalar_grant), - ): - with pytest.raises(msgspec.ValidationError): - decoder.decode(msgspec.msgpack.encode(wire)) - - def test_g3_config_keeps_its_intrinsic_path_checks() -> None: good = { "paths": ("/g3/a",), diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index 865bbde..f4c4281 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -454,19 +454,6 @@ def test_kvcr_rejects_no_dram_backends() -> None: ) -def test_kvcr_accepts_multi_pool_layouts() -> None: - kvcr = _new_kvcr( - FakeNixlAgent(), - FakePrimaryPinning(), - FakeBytesControl(), - KVCRConfig( - nixl_agent_name="target", - pool_layouts=[("full", 8), ("swa", 4)], - ), - ) - assert kvcr._core.pool_layouts == [("full", 8), ("swa", 4)] - - def test_kvcr_rejects_ambiguous_pool_names() -> None: bindings = KVCRBindings(Mock(), Mock(), Mock()) for pool_layouts, message in ( diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index bb760e9..ad0f239 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -1029,6 +1029,7 @@ def test_promotion_failure_fails_the_pool_and_stops_the_whole_service( _claim(registry, 0, liveness) _kill_and_wait(registry, 0, liveness) + assert server._fatal_error is failure server.shutdown.assert_called_once_with() assert registry._refusing.is_set() is True diff --git a/tests/unit/test_recovery_journal.py b/tests/unit/test_recovery_journal.py index d24ef76..64cf5a0 100644 --- a/tests/unit/test_recovery_journal.py +++ b/tests/unit/test_recovery_journal.py @@ -300,21 +300,10 @@ def test_canonical_pool_terms_bind_ordered_geometry_and_allocation_identity() -> _PoolDescriptor("pool1", 2 * mmap.PAGESIZE, 2048, 3 * mmap.PAGESIZE), ) - def terms_for(candidate=pools, digest=_TEST_DIGEST, allocation=spec): - return canonical_pool_terms(digest, candidate, allocation) - - terms = terms_for() - - for field, value in ( - ("name", "other"), - ("size_bytes", 8192), - ("block_size_bytes", 2048), - ("offset_bytes", 12288), - ): - changed = msgspec.structs.replace(pools[0], **{field: value}) - assert terms_for((changed, pools[1])) != terms - assert terms_for(tuple(reversed(pools))) != terms - assert terms_for(allocation=msgspec.structs.replace(spec, device=8)) != terms + terms = canonical_pool_terms(_TEST_DIGEST, pools, spec) + changed = msgspec.structs.replace(pools[0], block_size_bytes=2048) + assert canonical_pool_terms(_TEST_DIGEST, (changed, pools[1]), spec) != terms + assert canonical_pool_terms(_TEST_DIGEST, tuple(reversed(pools)), spec) != terms def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) -> None: From 1f6282fc0fde5fc27afacb817e08e58a14eee84b Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 01:07:43 +0300 Subject: [PATCH 09/16] fix(core): harden multi-pool boundaries Signed-off-by: Kapil Arya --- src/kvcr/local_dram.py | 8 ++++ src/kvcr/recovery_journal.py | 17 +------ tests/unit/test_kvcr.py | 14 +++++- tests/unit/test_kvcr_local_dram.py | 19 ++++++++ tests/unit/test_kvcr_remote_target.py | 64 +++++++++++++++++++++++++++ 5 files changed, 105 insertions(+), 17 deletions(-) diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index d78db6e..e50c4d4 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -166,6 +166,14 @@ def __init__( raise ValueError("local DRAM pool must hold at least one block") self._pools[pool_name] = (address, length, slot_size) self._free_slots[pool_name] = deque(range(slot_count)) + ranges = sorted( + (address, address + length) for address, length, _ in self._pools.values() + ) + if any( + left_end > right_start + for (_, left_end), (right_start, _) in zip(ranges, ranges[1:]) + ): + raise ValueError("local DRAM pools must not overlap") self._evictable = _EvictionQueue() self._evictable_slots: Counter[str] = Counter() self._unscored: set[BlockKey] = set() diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index 20044af..a680a86 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -17,15 +17,10 @@ import msgspec -from .config import ( - KVCRBackendConfigs, - KVCRConfig, - KVCRGuardConfig, - _validate_pool_layouts, -) +from .config import KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig from .core import _BlockRecord, _KVCRCore from .guard_protocol import KVCRClient, KVCRPoolHold, _PoolDescriptor -from .local_disk import _G3, _G3Residency, _validate_g3_slot_geometry +from .local_disk import _G3, _G3Residency from .local_dram import _LocalDram, _LocalDramResidency, _LocalDramState from .memory import _JOURNAL_HEADER_BYTES, KVCRPoolAttachment, KVCRPoolSpec from .types import BlockKey, RecoveryMirrorError @@ -457,14 +452,6 @@ def claim_guarded_pool( """ if backend_configs.local_dram is not None: raise ValueError("guard_config conflicts with backend_configs.local_dram") - if backend_configs.g3 is not None: - _validate_pool_layouts(config.pool_layouts) - if len(config.pool_layouts) != 1: - raise ValueError("G3 does not support multiple pools") - _validate_g3_slot_geometry( - backend_configs.g3, - config.pool_layouts[0][1], - ) # Duck-typed: what matters is whether the framework's control can hand its # endpoint over, not what class it is. framework_control = bindings.framework_control diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index f4c4281..74dcd8e 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -137,8 +137,18 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: [ ("control-absent", ValueError, "share its control endpoint", []), ("control-cannot-share", ValueError, "share its control endpoint", []), - ("g3-invalid", ValueError, "page aligned", []), - ("g3-multi-pool", ValueError, "does not support multiple pools", []), + ( + "g3-invalid", + ValueError, + "page aligned", + ["claim", "hold.release"], + ), + ( + "g3-multi-pool", + ValueError, + "does not support multiple pools", + ["claim", "hold.release"], + ), ( "handback-unreadable", RuntimeError, diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index f4d4c42..728e8cf 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -131,6 +131,25 @@ def test_local_transfer_accepts_multiple_blocks_in_one_pool() -> None: ].success +def test_local_dram_rejects_overlapping_pools() -> None: + memory = ctypes.create_string_buffer(16) + address = ctypes.addressof(memory) + + with pytest.raises(ValueError, match="overlap"): + _new_kvcr( + FakeNixlAgent(), + FakePrimaryPinning(), + FakeBytesControl(), + KVCRConfig( + nixl_agent_name="target", + pool_layouts=[("full", 8), ("swa", 8)], + ), + local_dram=LocalDramOptions( + [("full", address, 16), ("swa", address + 8, 8)] + ), + ) + + def test_multi_pool_residency_moves_and_evicts_as_one_key() -> None: full = ctypes.create_string_buffer(16) swa = ctypes.create_string_buffer(16) diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index cdce976..50a98d2 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -234,6 +234,70 @@ def test_remote_fetch_uses_local_then_framework_sources() -> None: ) +def test_remote_fetch_preserves_a_multi_pool_layout() -> None: + names = ("full", "swa") + layout = [(name, 8) for name in names] + source_primary = ctypes.create_string_buffer(16) + source_local = [ctypes.create_string_buffer(8) for _ in names] + target_local = [ctypes.create_string_buffer(8) for _ in names] + source_agent = FakeNixlAgent(metadata=b"source-md") + target_agent = FakeNixlAgent(metadata=b"target-md") + source_control = FakeBytesControl("tcp://source:1") + target_control = FakeBytesControl("tcp://target:1") + config = KVCRConfig(nixl_agent_name="unused", pool_layouts=layout) + + def dram(memories: list[ctypes.Array]) -> LocalDramOptions: + return LocalDramOptions( + [ + (name, ctypes.addressof(memory), len(memory)) + for name, memory in zip(names, memories, strict=True) + ] + ) + + source = _new_kvcr( + source_agent, + FakePrimaryPinning(), + source_control, + config, + name="source", + local_dram=dram(source_local), + ) + target = _new_kvcr( + target_agent, + FakePrimaryPinning(), + target_control, + config, + key_adapter=_ConstantHashAdapter(), + remote_options=RemoteFWDramOptions(eager_ctrl_connect=False), + local_dram=dram(target_local), + ) + key = BlockKey(b"multi-pool") + descriptors = [ + _mem_descriptor(ctypes.addressof(source_primary) + index * 8, 8, info=name) + for index, name in enumerate(names) + ] + source_agent.state = "DONE" + deposit = source.deposit({key: descriptors}) + assert dict(_poll_until(source, bool))[deposit][key].success + source_agent.state = "PROC" + + target.submit_hint(_router_hint("tcp://source:1"), request_id="req") + fetch = target.fetch((key,), "req", expected_layout=list(names)) + _wait_until(lambda: bool(target_control.sent)) + source_control.incoming.extend(message for _, message in target_control.sent) + _poll_until(source, lambda _: len(source_agent.xfers) == 2) + source_xfer = source_agent.xfers[1] + assert len(source_xfer[1]) == len(source_xfer[3]) == 2 + + notification = source_xfer[5] + source_agent.state = "DONE" + _poll_until(source, lambda _: not _has_outstanding_operations(source)) + target_agent.notifs["source"] = [notification] + result = dict(_poll_until(target, bool))[fetch][key] + assert result.success + assert [descriptor.info for descriptor in result.descriptors or ()] == list(names) + + def test_remote_staging_commits_available_prefix() -> None: block_size = 16 local = ctypes.create_string_buffer(block_size * 2) From 474bae3f2c3a54c723120a9615c9aaa5ef464264 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 01:29:39 +0300 Subject: [PATCH 10/16] fix(core): tighten multi-pool boundaries Signed-off-by: Kapil Arya --- src/kvcr/core.py | 30 ++++++++++------- src/kvcr/guard.py | 2 +- src/kvcr/local_disk.py | 10 +++++- src/kvcr/local_dram.py | 13 +++----- src/kvcr/recovery_journal.py | 46 ++++++++++----------------- tests/unit/test_g3.py | 2 ++ tests/unit/test_guard.py | 3 +- tests/unit/test_guard_integration.py | 6 +--- tests/unit/test_guard_protocol.py | 5 +++ tests/unit/test_kvcr.py | 22 ++++--------- tests/unit/test_kvcr_local_dram.py | 22 ++++++++++--- tests/unit/test_kvcr_remote_target.py | 35 +++++++++++++------- tests/unit/test_kvcr_service.py | 3 +- tests/unit/test_recovery_journal.py | 7 ++-- tests/unit/test_recovery_mirror.py | 10 +++--- 15 files changed, 116 insertions(+), 100 deletions(-) diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 61af860..04cfeb3 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -380,7 +380,7 @@ def deliver( if self._is_local_resident(key): local_blocks[key] = destination elif self._g3 is not None and self._g3.is_ready(key): - g3_blocks[key] = destination[0] + g3_blocks[key] = self._g3._single_descriptor(destination) else: remote_blocks[key] = destination @@ -437,10 +437,9 @@ def fetch( hints: object | None = None, ) -> OpHandle: expected_layout = [""] if expected_layout is None else list(expected_layout) - if not expected_layout or any( - name not in self._block_sizes for name in expected_layout - ): - raise ValueError("expected layout must use configured pools") + self._validate_block_layout( + expected_layout, "expected layout must use configured pools" + ) op_handle = self._next_op_handle self._next_op_handle += 1 local_dram = self._local_dram @@ -690,7 +689,10 @@ def _start_local_fill( if source is CacheTier.G3: started = self._g3 is not None and self._g3.start_fill( fill_handle, - {key: descriptors[0] for key, descriptors in blocks.items()}, + { + key: self._g3._single_descriptor(descriptors) + for key, descriptors in blocks.items() + }, deadline, ) elif source is CacheTier.REMOTE_G2: @@ -724,16 +726,22 @@ def _normalize_descriptors( isinstance(descriptor, MemDescriptor) for descriptor in descriptors ): raise ValueError("each block requires at least one descriptor") + self._validate_block_layout( + [descriptor.info for descriptor in descriptors], + "block descriptors must use configured pools", + ) for descriptor in descriptors: - block_size = self._block_sizes.get(descriptor.info) - if block_size is None: - raise ValueError(f"unknown descriptor pool {descriptor.info!r}") + block_size = self._block_sizes[descriptor.info] if descriptor.size != block_size: raise ValueError("block descriptor has the wrong byte count") - if self._g3 is not None and len(descriptors) != 1: - raise ValueError("G3 does not support multi-block layouts") return list(descriptors) + def _validate_block_layout(self, layout: list[str], invalid_message: str) -> None: + if not layout or any(name not in self._block_sizes for name in layout): + raise ValueError(invalid_message) + if self._g3 is not None and len(layout) != 1: + raise ValueError("G3 does not support multi-block layouts") + def _release_local_dram_sources( self, op_id: _OpId, diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index 453e7bf..2feef6b 100644 --- a/src/kvcr/guard.py +++ b/src/kvcr/guard.py @@ -318,7 +318,7 @@ def _write_handback(self, records: Mapping[BlockKey, _BlockRecord]) -> None: write_recovery_snapshot( self.attachment, canonical_pool_terms(self._compatibility_digest, self.pools, self._spec), - _recovery_frames(records, tuple(pool.name for pool in self.pools)), + _recovery_frames(records), ) def close(self) -> None: diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index d5555f9..a58bf4d 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -295,7 +295,9 @@ def resolve_eviction( if key not in sources: return (PlacementAction.KEEP, None), False try: - if not self._start_store(op_id, {key: sources[key][0]}, deadline): + if not self._start_store( + op_id, {key: self._single_descriptor(sources[key])}, deadline + ): self._recover_store_failure(key, "G3 destination unavailable") return (PlacementAction.KEEP, None), False except Exception: @@ -319,6 +321,12 @@ def start_deliver( ) -> bool: return self._start_read("deliver", op_handle, blocks, deadline) + @staticmethod + def _single_descriptor(descriptors: list[MemDescriptor]) -> MemDescriptor: + if len(descriptors) != 1: + raise ValueError("G3 requires one descriptor per block") + return descriptors[0] + def poll_main(self, items: Collection[object]) -> list[object]: unhandled: list[object] = [] for item in items: diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index e50c4d4..acf1e70 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -390,16 +390,13 @@ def fetch( to_reserve.append(key) else: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) - elif residency.state is _LocalDramState.READY: - if [name for name, _ in residency.slots] == layout: - self._kvcr._record_access((key,)) - op.results[key] = self._new_public_claim( - key, residency, include_descriptors=True - ) - else: - op.results[key] = OpEntryResult(OpEntryStatus.FAILED) elif [name for name, _ in residency.slots] != layout: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) + elif residency.state is _LocalDramState.READY: + self._kvcr._record_access((key,)) + op.results[key] = self._new_public_claim( + key, residency, include_descriptors=True + ) elif residency.state is _LocalDramState.DISCARDING: # A discarded fill still owns its slot, so this block cannot be # reserved yet. Wait for the slot instead of failing a key a diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index a680a86..3c2a2aa 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -78,7 +78,13 @@ def store_release(self, value: int) -> None: # count and pool-name lengths. class _RecoveryBlock(msgspec.Struct, frozen=True, array_like=True): # Ordered pool locations, or nothing. Pool names may repeat. - g2: list[tuple[str, int]] | None = None + g2: ( + Annotated[ + list[tuple[str, Annotated[int, msgspec.Meta(ge=0)]]], + msgspec.Meta(min_length=1), + ] + | None + ) = None g3: Annotated[int, msgspec.Meta(ge=0)] | None = None @@ -103,40 +109,25 @@ def _is_recoverable(record: _BlockRecord) -> bool: ) -def _g2_locations( - slots: list[tuple[str, int]], pool_names: tuple[str, ...] -) -> list[tuple[str, int]]: - if type(slots) is not list or not slots: - raise ValueError("G2 recovery locations must be a non-empty list") - allowed = set(pool_names) - for location in slots: - if type(location) is not tuple or len(location) != 2: - raise ValueError("G2 recovery location must be a pool and slot pair") - pool_name, pool_slot = location - if pool_name not in allowed or type(pool_slot) is not int or pool_slot < 0: - raise ValueError("G2 recovery location does not match the pool group") - return slots - - -def _project_recovery_record( - record: _BlockRecord, pool_names: tuple[str, ...] -) -> _RecoveryBlock: +def _project_recovery_record(record: _BlockRecord) -> _RecoveryBlock: g3 = record.g3.slot if record.g3 is not None else None local_dram = record.local_dram if local_dram is None or local_dram.state is not _LocalDramState.READY: return _RecoveryBlock(g3=g3) - return _RecoveryBlock(g2=_g2_locations(local_dram.slots, pool_names), g3=g3) + return _RecoveryBlock(g2=local_dram.slots, g3=g3) def _decode_recovery_record( payload: bytes, pool_names: tuple[str, ...] ) -> _BlockRecord: recovered = _RECOVERY_DECODER.decode(payload) + if recovered.g2 is not None and any( + name not in pool_names for name, _ in recovered.g2 + ): + raise ValueError("G2 recovery location does not match the pool group") return _BlockRecord( local_dram=( - _LocalDramResidency( - _g2_locations(recovered.g2, pool_names), _LocalDramState.READY - ) + _LocalDramResidency(recovered.g2, _LocalDramState.READY) if recovered.g2 is not None else None ), @@ -357,7 +348,6 @@ def adopt(self, records: dict[BlockKey, _BlockRecord]) -> None: if local_dram.state is not _LocalDramState.READY: record.local_dram = None else: - _g2_locations(local_dram.slots, self._pool_names) local_dram.claim_count = 0 local_dram.retire_on_release = False if record.g3 is not None: @@ -386,7 +376,6 @@ def take_records(self) -> dict[BlockKey, _BlockRecord]: def _attach_journal( local_dram: _LocalDram, journal: RecoveryJournal, - pool_names: tuple[str, ...], g3: _G3 | None = None, ) -> None: """Attach stable G2/G3 residency publication to one journal.""" @@ -410,7 +399,7 @@ def publish_frame(record_type: int, key: bytes, payload: bytes) -> None: def publish(key: BlockKey, record: _BlockRecord) -> None: # TODO: Publish per-tier deltas if full-record journal traffic is material. - recovered = _project_recovery_record(record, pool_names) + recovered = _project_recovery_record(record) publish_frame(_RECORD_BLOCK, bytes(key), _RECOVERY_ENCODER.encode(recovered)) if g3 is not None: @@ -519,7 +508,6 @@ def adopt_claimed_pool(core: _KVCRCore, claimed: ClaimedPool) -> None: _attach_journal( core._local_dram, RecoveryJournal(hold._attachment), - tuple(pool[0] for pool in hold.local_dram.pools), core._g3, ) install_recovery_records(core, claimed.recovered.take_records()) @@ -554,13 +542,13 @@ def _pack_frame(record_type: int, key: bytes, payload: bytes, size: int) -> byte def _recovery_frames( - records: Mapping[BlockKey, _BlockRecord], pool_names: tuple[str, ...] + records: Mapping[BlockKey, _BlockRecord], ) -> Iterator[tuple[int, bytes, bytes]]: """Every frame a returning primary needs to rebuild this state.""" for key, record in records.items(): if not _is_recoverable(record): continue - payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record, pool_names)) + payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record)) yield _RECORD_BLOCK, bytes(key), payload diff --git a/tests/unit/test_g3.py b/tests/unit/test_g3.py index 0ccd540..c046690 100644 --- a/tests/unit/test_g3.py +++ b/tests/unit/test_g3.py @@ -525,6 +525,8 @@ def test_g3_spill_deliver_and_fill_reuse_existing_progress(tmp_path) -> None: assert kvcr.query((first,)) == [(QueryStatus.FETCHABLE, CacheTier.G3)] now = 2.0 + with pytest.raises(ValueError, match="multi-block"): + kvcr.fetch((first,), expected_layout=["", ""]) fetch = kvcr.fetch((first,)) fetch_result = dict(_poll_until(kvcr, bool))[fetch][first] assert fetch_result.success and fetch_result.descriptors is not None diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index 6f777de..bc3f0e7 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -81,7 +81,6 @@ def _fake_attachment() -> Mock: """A stand-in with the pool-tail surface a Guard reaches for.""" attachment = Mock( address=1234, - data_address=1234 + _TEST_SPEC.journal_bytes, _spec=_TEST_SPEC, ) attachment.mapped_snapshot.return_value = nullcontext(None) @@ -106,7 +105,7 @@ def drain(self): def _frame(key: BlockKey, record: _BlockRecord) -> tuple[int, bytes, bytes]: """One journal frame, exactly as a primary would publish it.""" - payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record, ("",))) + payload = _RECOVERY_ENCODER.encode(_project_recovery_record(record)) return (_RECORD_BLOCK, bytes(key), payload) diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 9a287ec..281f136 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -210,11 +210,7 @@ def _group_primary_child(socket_path: str, control_port: str) -> None: ) ) journal = RecoveryJournal(hold._attachment) - journal.publish( - *next( - iter(_recovery_frames({BlockKey(b"grouped"): record}, ("pool0", "pool1"))) - ) - ) + journal.publish(*next(iter(_recovery_frames({BlockKey(b"grouped"): record})))) print("ready", flush=True) time.sleep(60) diff --git a/tests/unit/test_guard_protocol.py b/tests/unit/test_guard_protocol.py index d756f63..7509b24 100644 --- a/tests/unit/test_guard_protocol.py +++ b/tests/unit/test_guard_protocol.py @@ -245,6 +245,11 @@ def test_g3_config_keeps_its_intrinsic_path_checks() -> None: [("pool", mmap.PAGESIZE)], _G3Config(**{**good, "capacity_bytes_per_file": mmap.PAGESIZE + 1}), ) + with pytest.raises(ValueError, match="page-aligned"): + _TierConfig( + [("pool", mmap.PAGESIZE // 2)], + _G3Config(**{**good, "capacity_bytes_per_file": mmap.PAGESIZE}), + ) def test_claim_and_release_round_trip_typed_messages_and_geometry( diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index 74dcd8e..b4d3824 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -143,12 +143,6 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: "page aligned", ["claim", "hold.release"], ), - ( - "g3-multi-pool", - ValueError, - "does not support multiple pools", - ["claim", "hold.release"], - ), ( "handback-unreadable", RuntimeError, @@ -166,7 +160,6 @@ def observe(key: BlockKey, record: _BlockRecord) -> None: "control-absent", "control-cannot-share", "g3-invalid", - "g3-multi-pool", "handback-unreadable", "install-fails", ], @@ -200,7 +193,7 @@ def claim(*_args, **_kwargs) -> SimpleNamespace: control = None elif stage == "control-cannot-share": control = SimpleNamespace(control_bind_address=None, adopt_listener=None) - elif stage in ("g3-invalid", "g3-multi-pool"): + elif stage == "g3-invalid": backend_configs = KVCRBackendConfigs( g3=G3Options( paths=(tmp_path / "g3",), @@ -243,11 +236,9 @@ def is_quiescent(self) -> bool: KVCR( KVCRConfig( nixl_agent_name="target", - pool_layouts=( - [("full", 1024), ("swa", 1024)] - if stage == "g3-multi-pool" - else [("", mmap.PAGESIZE // 2 if stage == "g3-invalid" else 1024)] - ), + pool_layouts=[ + ("", mmap.PAGESIZE // 2 if stage == "g3-invalid" else 1024) + ], nixl_listen_port=1, ), KVCRBindings(Mock(), Mock(), Mock(), framework_control=control), @@ -366,11 +357,10 @@ def make_journal(pool) -> object: events.append("journal") return journal - def attach_journal(local, configured_journal, pool_name, disk) -> None: - assert (local, configured_journal, pool_name, disk) == ( + def attach_journal(local, configured_journal, disk) -> None: + assert (local, configured_journal, disk) == ( local_dram, journal, - ("",), g3, ) events.append("attach") diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 728e8cf..08b8c82 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -177,6 +177,8 @@ def test_multi_pool_residency_moves_and_evicts_as_one_key() -> None: ] first, second = BlockKey(b"first"), BlockKey(b"second") + with pytest.raises(ValueError, match="configured pools"): + kvcr.deposit({first: [_mem_descriptor(info="unknown")]}) operation = kvcr.deposit({first: descriptors}) _wait_until(lambda: bool(agent.transfers)) wrong_layout = kvcr.fetch((first,), expected_layout=["swa"]) @@ -195,6 +197,13 @@ def test_multi_pool_residency_moves_and_evicts_as_one_key() -> None: kvcr.release([result.release_handle]) assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 3 + wrong = kvcr.deliver({first: descriptors[1:]}) + assert dict(kvcr.poll_completed())[wrong][first].status is OpEntryStatus.FAILED + matching = kvcr.deliver({first: descriptors}) + assert dict(_poll_until(kvcr, lambda done: matching in dict(done)))[matching][ + first + ].success + operation = kvcr.deposit({second: descriptors}) _poll_until(kvcr, lambda done: operation in dict(done)) assert kvcr.query((first, second)) == [ @@ -240,7 +249,7 @@ def test_failed_group_reservation_does_not_evict_a_partial_group() -> None: def test_group_allocation_evicts_enough_whole_keys() -> None: - pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(16)] + pools = [ctypes.create_string_buffer(16), ctypes.create_string_buffer(16)] source = ctypes.create_string_buffer(24) agent = FakeNixlAgent() agent.state = "DONE" @@ -251,7 +260,7 @@ def test_group_allocation_evicts_enough_whole_keys() -> None: KVCRConfig(nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)]), local_dram=LocalDramOptions( [ - ("full", ctypes.addressof(pools[0]), 8), + ("full", ctypes.addressof(pools[0]), 16), ("swa", ctypes.addressof(pools[1]), 16), ] ), @@ -273,12 +282,12 @@ def test_group_allocation_evicts_enough_whole_keys() -> None: assert result[grouped].success assert kvcr.query((full, swa0, swa1, grouped)) == [ - (QueryStatus.MISS, None), + (QueryStatus.HIT, CacheTier.LOCAL_G2), (QueryStatus.MISS, None), (QueryStatus.MISS, None), (QueryStatus.HIT, CacheTier.LOCAL_G2), ] - assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 3 + assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 4 @pytest.mark.parametrize( @@ -543,6 +552,11 @@ def test_capacity_pressure_is_pool_local() -> None: _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"), ] + kvcr._core._update_capacity_pressure({"full": 0, "swa": 0}) + assert capacity_requests == [[("full", 1), ("swa", 2)]] + kvcr._core._update_capacity_pressure({"full": 1, "swa": 2}) + capacity_requests.clear() + full = kvcr.deposit({BlockKey(b"full"): descriptors[:1]}, no_evict=True) _poll_until(kvcr, lambda done: full in dict(done)) assert capacity_requests == [[("full", 1)]] diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index 50a98d2..81a022a 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -234,7 +234,13 @@ def test_remote_fetch_uses_local_then_framework_sources() -> None: ) -def test_remote_fetch_preserves_a_multi_pool_layout() -> None: +@pytest.mark.parametrize( + ("expected_layout", "success"), + [(["swa", "full"], False), (["full", "swa"], True)], +) +def test_remote_fetch_validates_a_multi_pool_layout( + expected_layout: list[str], success: bool, caplog: pytest.LogCaptureFixture +) -> None: names = ("full", "swa") layout = [(name, 8) for name in names] source_primary = ctypes.create_string_buffer(16) @@ -282,20 +288,27 @@ def dram(memories: list[ctypes.Array]) -> LocalDramOptions: source_agent.state = "PROC" target.submit_hint(_router_hint("tcp://source:1"), request_id="req") - fetch = target.fetch((key,), "req", expected_layout=list(names)) + fetch = target.fetch((key,), "req", expected_layout=expected_layout) _wait_until(lambda: bool(target_control.sent)) source_control.incoming.extend(message for _, message in target_control.sent) - _poll_until(source, lambda _: len(source_agent.xfers) == 2) - source_xfer = source_agent.xfers[1] - assert len(source_xfer[1]) == len(source_xfer[3]) == 2 - - notification = source_xfer[5] - source_agent.state = "DONE" - _poll_until(source, lambda _: not _has_outstanding_operations(source)) + if success: + _poll_until(source, lambda _: len(source_agent.xfers) == 2) + source_xfer = source_agent.xfers[1] + assert len(source_xfer[1]) == len(source_xfer[3]) == 2 + notification = source_xfer[5] + source_agent.state = "DONE" + _poll_until(source, lambda _: not _has_outstanding_operations(source)) + else: + _poll_until(source, lambda _: bool(source_agent.sent_notifs)) + notification = source_agent.sent_notifs[0][1] + assert "start_write layout mismatch" in caplog.text target_agent.notifs["source"] = [notification] result = dict(_poll_until(target, bool))[fetch][key] - assert result.success - assert [descriptor.info for descriptor in result.descriptors or ()] == list(names) + assert result.success is success + if success: + assert [descriptor.info for descriptor in result.descriptors or ()] == list( + names + ) def test_remote_staging_commits_available_prefix() -> None: diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index ad0f239..f2f6a00 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -222,7 +222,6 @@ def _stand_in_pool(spec) -> Mock: def _new_registry( tmp_path: Path, guard_count: int = 1, - pool_sizes_bytes: tuple[int, ...] = _TEST_POOL_SIZES_BYTES, ) -> _PoolRegistry: """A registry of real Guards over stand-in pool mappings.""" journal = Mock() @@ -234,7 +233,7 @@ def _new_registry( return _PoolRegistry( tmp_path, guard_count, - pool_sizes_bytes, + _TEST_POOL_SIZES_BYTES, _TEST_JOURNAL_BYTES, _TEST_DIGEST, ) diff --git a/tests/unit/test_recovery_journal.py b/tests/unit/test_recovery_journal.py index 64cf5a0..99333cb 100644 --- a/tests/unit/test_recovery_journal.py +++ b/tests/unit/test_recovery_journal.py @@ -140,7 +140,7 @@ def test_publisher_streams_mutations_until_the_journal_refuses_or_fails( journal, _ = journal_and_mapping local_dram, g3 = _Source(), _Source() key = BlockKey(b"block") - _attach_journal(local_dram, journal, ("pool0",), g3) + _attach_journal(local_dram, journal, g3) caplog.set_level("WARNING", logger="kvcr.recovery_journal") local_dram.emit(key, _recovered_record(g2=[("pool0", 2)])) @@ -174,7 +174,7 @@ def test_publisher_streams_mutations_until_the_journal_refuses_or_fails( fresh = RecoveryJournal(_attachment(fresh_mapping, _TEST_JOURNAL_BYTES)) fresh.reset() source = _Source() - _attach_journal(source, fresh, ("pool0",)) + _attach_journal(source, fresh) assert not fresh.is_invalid() with patch.object(fresh, "publish", side_effect=RuntimeError("publish failed")): source.emit( @@ -280,7 +280,6 @@ def _write_slot(pool: KVCRPoolAttachment, terms: bytes, key: bytes, slot: int) - """One-slot handback region: the smallest finished snapshot.""" frames = _recovery_frames( {BlockKey(key * 32): _recovered_record(g2=[("pool0", slot)])}, - ("pool0",), ) write_recovery_snapshot(pool, terms, frames) @@ -324,7 +323,7 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - BlockKey(b"b" * 32): _recovered_record(g2=[("pool0", 4)], g3=9), BlockKey(b"c" * 32): _recovered_record(g3=2), } - write_recovery_snapshot(pool, terms, _recovery_frames(records, ("pool0",))) + write_recovery_snapshot(pool, terms, _recovery_frames(records)) # Inside the pool file, so it has no name of its own to be found under. assert set(tmp_path.iterdir()) == {path} assert path.stat().st_size > pool._spec.mapping_bytes diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index 9c90011..11d9c40 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -23,8 +23,8 @@ _TWO_POOLS = ("full", "swa") -def _payload(record: _BlockRecord, pool_names: tuple[str, ...] = _ONE_POOL) -> bytes: - return _RECOVERY_ENCODER.encode(_project_recovery_record(record, pool_names)) +def _payload(record: _BlockRecord) -> bytes: + return _RECOVERY_ENCODER.encode(_project_recovery_record(record)) # Every live-only field set, to prove projection strips all of it. @@ -92,7 +92,7 @@ def test_recovery_wire_round_trip_keeps_only_settled_slots( recovered: _BlockRecord, ) -> None: """Only settled G2/G3 slots reach the wire; decode rebuilds fresh live state.""" - encoded = _payload(record, pool_names) + encoded = _payload(record) # The outer record stays positional; G2 locations carry their pool names. assert msgspec.msgpack.decode(encoded) == wire @@ -137,8 +137,6 @@ def test_mirror_rejects_malformed_or_unknown_wire_state(payload: bytes) -> None: with pytest.raises(RecoveryMirrorError, match="malformed"): mirror.apply(_RECORD_BLOCK, b"block", payload) - assert mirror._records == {} - def test_mirror_replaces_blocks_whole_and_hands_them_over_uncopied() -> None: """Frames replace blocks whole in _records (mirrored table); take transfers it.""" @@ -212,7 +210,7 @@ def test_mirror_adopts_exactly_what_a_handback_region_would_carry() -> None: # A kept mirror must match exactly what the handback frames carry. framed = { BlockKey(key): _decode_recovery_record(payload, _TWO_POOLS) - for _, key, payload in _recovery_frames(served, _TWO_POOLS) + for _, key, payload in _recovery_frames(served) } mirror = _RecoveryMirror(_TWO_POOLS) mirror.adopt(served) From 4632f4efb9887db07e259f88a92bd8892a29d1f8 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 04:39:14 +0300 Subject: [PATCH 11/16] fix(core): reject incompatible G3 recovery Signed-off-by: Kapil Arya --- src/kvcr/core.py | 12 +++++++----- src/kvcr/local_disk.py | 10 +--------- tests/unit/test_g3.py | 21 ++++++++++++++++++++- 3 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 04cfeb3..98d4274 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -289,6 +289,11 @@ def adopt_recovery_records(self, records: dict[BlockKey, _BlockRecord]) -> None: g3 = self._g3 if g3 is None and any(record.g3 is not None for record in records.values()): raise RecoveryMirrorError("recovered G3 residency has no configured G3") + if g3 is not None and any( + record.local_dram is not None and len(record.local_dram.slots) != 1 + for record in records.values() + ): + raise RecoveryMirrorError("G3 recovery requires one local DRAM slot") if self._block_record_map: raise RecoveryMirrorError("recovered records need a core that holds none") @@ -380,7 +385,7 @@ def deliver( if self._is_local_resident(key): local_blocks[key] = destination elif self._g3 is not None and self._g3.is_ready(key): - g3_blocks[key] = self._g3._single_descriptor(destination) + g3_blocks[key] = destination[0] else: remote_blocks[key] = destination @@ -689,10 +694,7 @@ def _start_local_fill( if source is CacheTier.G3: started = self._g3 is not None and self._g3.start_fill( fill_handle, - { - key: self._g3._single_descriptor(descriptors) - for key, descriptors in blocks.items() - }, + {key: descriptors[0] for key, descriptors in blocks.items()}, deadline, ) elif source is CacheTier.REMOTE_G2: diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index a58bf4d..d5555f9 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -295,9 +295,7 @@ def resolve_eviction( if key not in sources: return (PlacementAction.KEEP, None), False try: - if not self._start_store( - op_id, {key: self._single_descriptor(sources[key])}, deadline - ): + if not self._start_store(op_id, {key: sources[key][0]}, deadline): self._recover_store_failure(key, "G3 destination unavailable") return (PlacementAction.KEEP, None), False except Exception: @@ -321,12 +319,6 @@ def start_deliver( ) -> bool: return self._start_read("deliver", op_handle, blocks, deadline) - @staticmethod - def _single_descriptor(descriptors: list[MemDescriptor]) -> MemDescriptor: - if len(descriptors) != 1: - raise ValueError("G3 requires one descriptor per block") - return descriptors[0] - def poll_main(self, items: Collection[object]) -> list[object]: unhandled: list[object] = [] for item in items: diff --git a/tests/unit/test_g3.py b/tests/unit/test_g3.py index c046690..4f530b1 100644 --- a/tests/unit/test_g3.py +++ b/tests/unit/test_g3.py @@ -36,7 +36,7 @@ ) from kvcr.core import _BlockRecord from kvcr.local_disk import _G3Residency -from kvcr.local_dram import _LocalDramState +from kvcr.local_dram import _LocalDramResidency, _LocalDramState from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy from kvcr.recovery_journal import install_recovery_records from kvcr.types import ( @@ -45,6 +45,7 @@ InventoryEvent, PlacementAction, QueryStatus, + RecoveryMirrorError, ) @@ -484,6 +485,24 @@ def test_g3_recovery_rejects_invalid_slots(tmp_path, slots: tuple[int, int]) -> assert kvcr._core._block_record_map == {} +def test_g3_recovery_rejects_multi_block_local_residency(tmp_path) -> None: + page_size = os.sysconf("SC_PAGE_SIZE") + local = ctypes.create_string_buffer(2 * page_size) + kvcr = _new_g3_kvcr(tmp_path, local, slot_count=2) + + with pytest.raises(RecoveryMirrorError, match="one local DRAM slot"): + install_recovery_records( + kvcr._core, + { + BlockKey(b"multi"): _BlockRecord( + local_dram=_LocalDramResidency( + [("", 0), ("", 1)], _LocalDramState.READY + ) + ) + }, + ) + + def test_g3_spill_deliver_and_fill_reuse_existing_progress(tmp_path) -> None: page_size = os.sysconf("SC_PAGE_SIZE") primary = ctypes.create_string_buffer(page_size * 2) From 8f3b77923e680d5b189b7550e0d4bffa590344b0 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 05:20:20 +0300 Subject: [PATCH 12/16] refactor: simplify multi-pool bookkeeping Signed-off-by: Kapil Arya --- src/kvcr/guard.py | 3 +- src/kvcr/local_disk.py | 15 +++---- src/kvcr/local_dram.py | 61 +++++++++++++--------------- src/kvcr/policy_runtime.py | 4 +- src/kvcr/remote_fw_dram.py | 4 +- tests/unit/test_g3.py | 22 +++------- tests/unit/test_guard_integration.py | 10 ++--- tests/unit/test_kvcr.py | 6 +-- tests/unit/test_kvcr_service.py | 13 ++---- tests/unit/test_recovery_journal.py | 38 +++++------------ tests/unit/test_recovery_mirror.py | 2 - 11 files changed, 62 insertions(+), 116 deletions(-) diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index 2feef6b..a0507ae 100644 --- a/src/kvcr/guard.py +++ b/src/kvcr/guard.py @@ -31,7 +31,6 @@ from .memory import ( KVCRPoolAttachment, KVCRPoolSpec, - _compute_pool_geometry, _KVCRPoolOwner, ) from .recovery_journal import ( @@ -856,7 +855,7 @@ def reject_pin(keys: object) -> int: ( pool.name, self._recovery.attachment.address + pool.offset_bytes, - _compute_pool_geometry(pool.size_bytes, pool.block_size_bytes)[0], + pool.size_bytes, ) for pool in self._recovery.pools ], diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index d5555f9..df4c4b2 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -33,15 +33,6 @@ logger = logging.getLogger(__name__) -def _validate_g3_slot_geometry(config: G3Options, slot_size: int) -> None: - """Validate the scalar data-plane relationship Guard protocol ignores.""" - if slot_size <= 0 or slot_size % os.sysconf("SC_PAGE_SIZE"): - raise ValueError("G3 slot size must be positive and page aligned") - capacity = config.capacity_bytes_per_file - if capacity <= 0 or capacity % slot_size: - raise ValueError("G3 file capacity must contain complete slots") - - @dataclass(slots=True) class _G3Residency: slot: int @@ -134,7 +125,11 @@ def __init__(self, kvcr: "_KVCRCore", config: G3Options, slot_size: int) -> None raise ValueError("G3 requires at least one file path") if len(paths) != len(set(paths)): raise ValueError("G3 file paths must be unique") - _validate_g3_slot_geometry(config, slot_size) + if slot_size <= 0 or slot_size % os.sysconf("SC_PAGE_SIZE"): + raise ValueError("G3 slot size must be positive and page aligned") + capacity = config.capacity_bytes_per_file + if capacity <= 0 or capacity % slot_size: + raise ValueError("G3 file capacity must contain complete slots") if not config.backend: raise ValueError("G3 NIXL backend must be non-empty") if not all( diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index acf1e70..03eb359 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -411,7 +411,7 @@ def fetch( required_local=True, deadline=deadline, framework_hints=hints, - layouts={key: layout for key in to_reserve}, + layout=layout, ) op.remote_fill_keys.update(destinations) for key in eviction_pending: @@ -686,13 +686,14 @@ def reserve_fill( required_local: bool, deadline: float, framework_hints: object | None = None, - layouts: Mapping[BlockKey, list[str]], + layout: list[str], ) -> tuple[dict[BlockKey, list[MemDescriptor]], set[BlockKey]]: keys = tuple(dict.fromkeys(keys)) protected = set(keys) destinations: dict[BlockKey, list[MemDescriptor]] = {} eviction_pending: set[BlockKey] = set() evicted: list[BlockKey] = [] + size_bytes = sum(self._pools[name][2] for name in layout) for key in keys: record = self._kvcr._block_record_map.get(key) if record is None: @@ -700,11 +701,7 @@ def reserve_fill( if record.local_dram is not None: continue decision = self._kvcr._policy.decide_ingest( - self._kvcr._block_meta( - key, - record, - sum(self._pools[name][2] for name in layouts[key]), - ), + self._kvcr._block_meta(key, record, size_bytes), sources[key], required_local, framework_hints=framework_hints, @@ -712,7 +709,7 @@ def reserve_fill( if decision[0] is PlacementAction.DROP: continue locations, evicted_keys, waiting = self._allocate_slots( - layouts[key], protected, deadline + layout, protected, deadline ) evicted.extend(evicted_keys) if locations is None: @@ -977,11 +974,7 @@ def _new_public_claim( def _acquire_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: if residency.state is not _LocalDramState.READY: raise RuntimeError(f"cannot claim unready local DRAM entry {key!r}") - if residency.claim_count == 0: - if key in self._unscored: - self._unscored.remove(key) - else: - self._remove_evictable(key, residency) + self._remove_evictable(key, residency) residency.claim_count += 1 def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: @@ -1014,8 +1007,9 @@ def _allocate_slots( self, pool_names: list[str], protected: set[BlockKey], deadline: float ) -> tuple[list[tuple[str, int]] | None, list[BlockKey], bool]: required = Counter(pool_names) - available = {name: len(self._free_slots[name]) for name in required} - if all(available[name] >= count for name, count in required.items()): + if all( + len(self._free_slots[name]) >= count for name, count in required.items() + ): return ( [(name, self._free_slots[name].popleft()) for name in pool_names], [], @@ -1026,7 +1020,19 @@ def _allocate_slots( self._retry_unscored() skipped = set(protected) victims: list[tuple[BlockKey, "_BlockRecord", _LocalDramResidency, int]] = [] - while (key := self._evictable.select(skipped)) is not None: + freed: Counter[str] = Counter() + + def short() -> set[str]: + return { + name + for name, count in required.items() + if len(self._free_slots[name]) + freed[name] < count + } + + while deficient := short(): + key = self._evictable.select(skipped) + if key is None: + return None, [], False record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None if ( @@ -1036,22 +1042,16 @@ def _allocate_slots( or residency.claim_count ): raise RuntimeError(f"invalid evictable local DRAM entry {key!r}") - if not any( - name in required and available[name] < required[name] - for name, _ in residency.slots - ): + if not any(name in deficient for name, _ in residency.slots): skipped.add(key) continue size_bytes = self._size_bytes(residency.slots) - free_before = {name: len(self._free_slots[name]) for name in required} decision, eviction_pending = self._kvcr._decide_eviction( self._kvcr._block_meta(key, record, size_bytes), CacheTier.LOCAL_G2, deadline, ) - for name in required: - available[name] += len(self._free_slots[name]) - free_before[name] - if all(available[name] >= count for name, count in required.items()): + if not short(): break if eviction_pending: self._capacity_eviction_key = key @@ -1061,13 +1061,7 @@ def _allocate_slots( continue victims.append((key, record, residency, size_bytes)) skipped.add(key) - for name, _ in residency.slots: - if name in available: - available[name] += 1 - if all(available[name] >= count for name, count in required.items()): - break - else: - return None, [], False + freed.update(name for name, _ in residency.slots) for key, record, residency, size_bytes in victims: self._remove_evictable(key, residency) @@ -1100,8 +1094,9 @@ def _make_evictable(self, key: BlockKey) -> None: self._evictable_slots.update(name for name, _ in record.local_dram.slots) def _remove_evictable(self, key: BlockKey, residency: _LocalDramResidency) -> None: - self._evictable.remove(key) - self._evictable_slots.subtract(name for name, _ in residency.slots) + self._unscored.discard(key) + if self._evictable.remove(key): + self._evictable_slots.subtract(name for name, _ in residency.slots) def _retry_unscored(self) -> None: for key in tuple(self._unscored): diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index cea2ce1..0a7068f 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -165,8 +165,8 @@ def insert(self, key: BlockKey, score: float) -> None: self._live[key] = entry heapq.heappush(self._heap, (entry.score, entry.sequence, key)) - def remove(self, key: BlockKey) -> None: - self._live.pop(key, None) + def remove(self, key: BlockKey) -> bool: + return self._live.pop(key, None) is not None def select(self, excluded: set[BlockKey]) -> BlockKey | None: skipped: list[tuple[float, int, BlockKey]] = [] diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index bada79a..41e72cf 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -979,8 +979,8 @@ def _submit_prepared_source_write( destination = source_pin.dst_descriptors[index] if source is None: break - if [(descriptor.info, descriptor.size) for descriptor in source] != [ - (descriptor.info, descriptor.size) for descriptor in destination + if [descriptor.info for descriptor in source] != [ + descriptor.info for descriptor in destination ]: logger.warning( "KVCR start_write layout mismatch op=%d key=%r", diff --git a/tests/unit/test_g3.py b/tests/unit/test_g3.py index 4f530b1..9f0c9d7 100644 --- a/tests/unit/test_g3.py +++ b/tests/unit/test_g3.py @@ -18,6 +18,7 @@ _mem_descriptor, _new_kvcr, _poll_until, + _recovered_record, _router_hint, _write_done_notification, ) @@ -36,7 +37,7 @@ ) from kvcr.core import _BlockRecord from kvcr.local_disk import _G3Residency -from kvcr.local_dram import _LocalDramResidency, _LocalDramState +from kvcr.local_dram import _LocalDramState from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy from kvcr.recovery_journal import install_recovery_records from kvcr.types import ( @@ -493,13 +494,7 @@ def test_g3_recovery_rejects_multi_block_local_residency(tmp_path) -> None: with pytest.raises(RecoveryMirrorError, match="one local DRAM slot"): install_recovery_records( kvcr._core, - { - BlockKey(b"multi"): _BlockRecord( - local_dram=_LocalDramResidency( - [("", 0), ("", 1)], _LocalDramState.READY - ) - ) - }, + {BlockKey(b"multi"): _recovered_record(g2=[("", 0), ("", 1)])}, ) @@ -806,7 +801,8 @@ def test_full_g3_does_not_hide_a_synchronously_freed_local_slot(tmp_path) -> Non kvcr = _new_g3_kvcr(tmp_path, local, policy=policy, g3_slot_count=1) first, second, third = (BlockKey(bytes((index,))) for index in range(3)) - for index, key in enumerate((first, second)): + for index, key in enumerate((first, second, third)): + policy.keep_g3 = index == 2 assert _deposit( kvcr, key, @@ -814,14 +810,6 @@ def test_full_g3_does_not_hide_a_synchronously_freed_local_slot(tmp_path) -> Non page_size, ).success - policy.keep_g3 = True - assert _deposit( - kvcr, - third, - ctypes.addressof(primary) + 2 * page_size, - page_size, - ).success - assert kvcr.query((first, second, third)) == [ (QueryStatus.FETCHABLE, CacheTier.G3), (QueryStatus.MISS, None), diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 281f136..98cb757 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -22,6 +22,7 @@ _mem_descriptor, _new_kvcr, _poll_until, + _recovered_record, _router_hint, _use_nixl_agent, _wait_until, @@ -39,7 +40,6 @@ RemoteFWDramOptions, ) from kvcr.control_channels import ZmqPeerControlChannel -from kvcr.core import _BlockRecord from kvcr.guard import _Guard from kvcr.kvcr_service import _KVCRService from kvcr.local_dram import _LocalDramResidency, _LocalDramState @@ -204,11 +204,7 @@ def _group_primary_child(socket_path: str, control_port: str) -> None: _DIGEST, ("127.0.0.1", int(control_port)), ) - record = _BlockRecord( - local_dram=_LocalDramResidency( - [("pool0", 0), ("pool1", 0)], _LocalDramState.READY - ) - ) + record = _recovered_record(g2=[("pool0", 0), ("pool1", 0)]) journal = RecoveryJournal(hold._attachment) journal.publish(*next(iter(_recovery_frames({BlockKey(b"grouped"): record})))) print("ready", flush=True) @@ -467,7 +463,7 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( assert guard._core._local_dram.memory_regions == ( ( guard._recovery.attachment.address + pools[0].offset_bytes, - page_size + page_size // 2, + 2 * page_size, ), ( guard._recovery.attachment.address + pools[1].offset_bytes, diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index b4d3824..b92014f 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -358,11 +358,7 @@ def make_journal(pool) -> object: return journal def attach_journal(local, configured_journal, disk) -> None: - assert (local, configured_journal, disk) == ( - local_dram, - journal, - g3, - ) + assert (local, configured_journal, disk) == (local_dram, journal, g3) events.append("attach") monkeypatch.setattr( diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index f2f6a00..6ac865c 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -33,6 +33,7 @@ _Error, _G3Config, _Granted, + _PoolDescriptor, _Release, _Released, _TierConfig, @@ -219,10 +220,7 @@ def _stand_in_pool(spec) -> Mock: return attachment -def _new_registry( - tmp_path: Path, - guard_count: int = 1, -) -> _PoolRegistry: +def _new_registry(tmp_path: Path, guard_count: int = 1) -> _PoolRegistry: """A registry of real Guards over stand-in pool mappings.""" journal = Mock() journal.read_next.return_value = None @@ -281,11 +279,8 @@ def test_client_claims_one_grouped_allocation_with_independent_strides( spec = guard._owner.spec assert spec.mapping_bytes == _TEST_JOURNAL_BYTES + sum(pool_sizes) offsets = (_TEST_JOURNAL_BYTES, _TEST_JOURNAL_BYTES + pool_sizes[0]) - assert tuple( - (pool.name, pool.size_bytes, pool.block_size_bytes, pool.offset_bytes) - for pool in guard._recovery.pools - ) == tuple( - (name, size, block_size, offset) + assert guard._recovery.pools == tuple( + _PoolDescriptor(name, size, block_size, offset) for (name, block_size), size, offset in zip( pool_layouts, pool_sizes, offsets, strict=True ) diff --git a/tests/unit/test_recovery_journal.py b/tests/unit/test_recovery_journal.py index 99333cb..bfec1ad 100644 --- a/tests/unit/test_recovery_journal.py +++ b/tests/unit/test_recovery_journal.py @@ -284,35 +284,13 @@ def _write_slot(pool: KVCRPoolAttachment, terms: bytes, key: bytes, slot: int) - write_recovery_snapshot(pool, terms, frames) -def test_canonical_pool_terms_bind_ordered_geometry_and_allocation_identity() -> None: - spec = KVCRPoolSpec( - pool_id="pool_0", - path=f"/tmp/kvcr-pool_0-{_GENERATION}", - generation=_GENERATION, - device=7, - inode=11, - mapping_bytes=5 * mmap.PAGESIZE, - journal_bytes=2 * mmap.PAGESIZE, - ) - pools = ( - _PoolDescriptor("pool0", mmap.PAGESIZE, 1024, 2 * mmap.PAGESIZE), - _PoolDescriptor("pool1", 2 * mmap.PAGESIZE, 2048, 3 * mmap.PAGESIZE), - ) - - terms = canonical_pool_terms(_TEST_DIGEST, pools, spec) - changed = msgspec.structs.replace(pools[0], block_size_bytes=2048) - assert canonical_pool_terms(_TEST_DIGEST, (changed, pools[1]), spec) != terms - assert canonical_pool_terms(_TEST_DIGEST, tuple(reversed(pools)), spec) != terms - - def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) -> None: """Replayed whole under its own terms, discardable when torn, gone once released.""" with _attached(tmp_path) as pool: path = Path(pool._spec.path) pools = ( - _PoolDescriptor( - "pool0", pool._spec.data_bytes, 4096, pool._spec.journal_bytes - ), + _PoolDescriptor("pool0", 2048, 1024, pool._spec.journal_bytes), + _PoolDescriptor("pool1", 2048, 1024, pool._spec.journal_bytes + 2048), ) terms = canonical_pool_terms(_TEST_DIGEST, pools, pool._spec) assert list(read_recovery_snapshot(pool, terms)) == [] @@ -335,9 +313,15 @@ def test_a_handback_region_lives_and_dies_inside_the_pool_file(tmp_path: Path) - assert mirror.take_records() == records # A slot number only means the same bytes under the same geometry. - other = canonical_pool_terms("another-digest", pools, pool._spec) - with pytest.raises(RecoveryJournalError, match="other terms"): - list(read_recovery_snapshot(pool, other)) + changed = msgspec.structs.replace(pools[0], block_size_bytes=2048) + for digest, layout in ( + ("another-digest", pools), + (_TEST_DIGEST, (changed, pools[1])), + (_TEST_DIGEST, tuple(reversed(pools))), + ): + other = canonical_pool_terms(digest, layout, pool._spec) + with pytest.raises(RecoveryJournalError, match="other terms"): + list(read_recovery_snapshot(pool, other)) # Stopped once the replacing body has landed but before its header has. interrupted = Mock( diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index 11d9c40..3a453df 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -125,8 +125,6 @@ class _RecoveryBlockV2(msgspec.Struct, frozen=True, array_like=True): msgspec.msgpack.encode({"g2": {"slot": "0"}}), msgspec.msgpack.encode([[["other", 0]], None]), msgspec.msgpack.encode([[["", -1]], None]), - msgspec.msgpack.encode([0, None]), - msgspec.msgpack.encode([[[""]], None]), msgspec.msgpack.encode([[], None]), ], ) From 599fd0cfd15ec2fc6206bd13f425769954ef38a2 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 05:24:36 +0300 Subject: [PATCH 13/16] refactor: remove unrelated G3 validation churn Restore the base constructor spelling to keep the multi-pool diff focused. Signed-off-by: Kapil Arya --- src/kvcr/local_disk.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/kvcr/local_disk.py b/src/kvcr/local_disk.py index df4c4b2..3a376c3 100644 --- a/src/kvcr/local_disk.py +++ b/src/kvcr/local_disk.py @@ -120,15 +120,18 @@ class _G3: """Own bounded files and the metadata needed to use them as G3 cache.""" def __init__(self, kvcr: "_KVCRCore", config: G3Options, slot_size: int) -> None: + page_size = os.sysconf("SC_PAGE_SIZE") paths = tuple(Path(path).expanduser().resolve() for path in config.paths) if not paths: raise ValueError("G3 requires at least one file path") if len(paths) != len(set(paths)): raise ValueError("G3 file paths must be unique") - if slot_size <= 0 or slot_size % os.sysconf("SC_PAGE_SIZE"): + if slot_size <= 0 or slot_size % page_size: raise ValueError("G3 slot size must be positive and page aligned") - capacity = config.capacity_bytes_per_file - if capacity <= 0 or capacity % slot_size: + if ( + config.capacity_bytes_per_file <= 0 + or config.capacity_bytes_per_file % slot_size + ): raise ValueError("G3 file capacity must contain complete slots") if not config.backend: raise ValueError("G3 NIXL backend must be non-empty") From 24b77ed45ab89aeb74dba5ee3831c57623ad983a Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 05:35:35 +0300 Subject: [PATCH 14/16] refactor: trim multi-pool setup and layout checks Restore base snapshot-term ordering and journal formatting; retain malformed-frame coverage while sharing repeated setup. Signed-off-by: Kapil Arya --- src/kvcr/local_dram.py | 21 ++++------ src/kvcr/recovery_journal.py | 16 +++---- tests/unit/test_kvcr_local_dram.py | 67 +++++++++++------------------- tests/unit/test_recovery_mirror.py | 2 + 4 files changed, 41 insertions(+), 65 deletions(-) diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index 03eb359..2df5be6 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -45,6 +45,10 @@ class _LocalDramResidency: claim_count: int = 0 retire_on_release: bool = False + @property + def layout(self) -> list[str]: + return [name for name, _ in self.slots] + @dataclass class _PendingResidencyOp(_Op): @@ -291,7 +295,7 @@ def deposit( record = self._kvcr._block_record(key) residency = record.local_dram if residency is not None: - if not self._same_layout(residency.slots, sources): + if residency.layout != [descriptor.info for descriptor in sources]: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) elif residency.state is _LocalDramState.READY: op.results[key] = ( @@ -390,7 +394,7 @@ def fetch( to_reserve.append(key) else: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) - elif [name for name, _ in residency.slots] != layout: + elif residency.layout != layout: op.results[key] = OpEntryResult(OpEntryStatus.FAILED) elif residency.state is _LocalDramState.READY: self._kvcr._record_access((key,)) @@ -743,7 +747,8 @@ def _start_deliveries( continue elif ( residency.state is _LocalDramState.DISCARDING - or not self._same_layout(residency.slots, op.destinations[key]) + or residency.layout + != [descriptor.info for descriptor in op.destinations[key]] or now >= op.deadline ): op.results[key] = OpEntryResult(OpEntryStatus.FAILED) @@ -887,7 +892,7 @@ def _resume_capacity_waiters(self) -> None: if residency is not None: self._capacity_waiters.popleft() op.capacity_waiters.remove(waiter.key) - if [name for name, _ in residency.slots] != waiter.layout: + if residency.layout != waiter.layout: op.results[waiter.key] = OpEntryResult(OpEntryStatus.FAILED) elif residency.state is _LocalDramState.READY: op.results[waiter.key] = ( @@ -1125,14 +1130,6 @@ def _free(self, locations: Collection[tuple[str, int]]) -> None: def _size_bytes(self, locations: Collection[tuple[str, int]]) -> int: return sum(self._pools[pool_name][2] for pool_name, _ in locations) - @staticmethod - def _same_layout( - locations: Collection[tuple[str, int]], descriptors: Collection[MemDescriptor] - ) -> bool: - return [name for name, _ in locations] == [ - descriptor.info for descriptor in descriptors - ] - def _update_capacity_pressure(self) -> None: if self._kvcr._capacity_needed_callback is None: return diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index 3c2a2aa..7a6360e 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -374,9 +374,7 @@ def take_records(self) -> dict[BlockKey, _BlockRecord]: def _attach_journal( - local_dram: _LocalDram, - journal: RecoveryJournal, - g3: _G3 | None = None, + local_dram: _LocalDram, journal: RecoveryJournal, g3: _G3 | None = None ) -> None: """Attach stable G2/G3 residency publication to one journal.""" enabled = True @@ -505,11 +503,7 @@ def adopt_claimed_pool(core: _KVCRCore, claimed: ClaimedPool) -> None: hold = claimed.hold if core._local_dram is None: raise ValueError("a claimed pool must give the core its local DRAM tier") - _attach_journal( - core._local_dram, - RecoveryJournal(hold._attachment), - core._g3, - ) + _attach_journal(core._local_dram, RecoveryJournal(hold._attachment), core._g3) install_recovery_records(core, claimed.recovered.take_records()) hold.hand_listener_to(claimed.adopt_listener) @@ -557,7 +551,7 @@ def _recovery_frames( # different pool of the same shape; the digest separates finished from filling. _SNAPSHOT_HEADER = struct.Struct("<32sQ") _SNAPSHOT_DOMAIN = b"KVCR-HANDBACK\0" -_SNAPSHOT_ALLOCATION_TERMS = struct.Struct(" None: block_size = 16 primary = ctypes.create_string_buffer(block_size * 3) @@ -155,20 +174,13 @@ def test_multi_pool_residency_moves_and_evicts_as_one_key() -> None: swa = ctypes.create_string_buffer(16) source = ctypes.create_string_buffer(32) agent = FakeNixlAgent() - kvcr = _new_kvcr( + kvcr = _two_pool_kvcr( agent, - FakePrimaryPinning(), - FakeBytesControl(), + (full, swa), KVCRConfig( nixl_agent_name="target", pool_layouts=[("full", 16), ("swa", 8)], ), - local_dram=LocalDramOptions( - [ - ("full", ctypes.addressof(full), 16), - ("swa", ctypes.addressof(swa), 16), - ] - ), ) descriptors = [ _mem_descriptor(ctypes.addressof(source), 16, info="full"), @@ -217,18 +229,7 @@ def test_failed_group_reservation_does_not_evict_a_partial_group() -> None: source = ctypes.create_string_buffer(16) agent = FakeNixlAgent() agent.state = "DONE" - kvcr = _new_kvcr( - agent, - FakePrimaryPinning(), - FakeBytesControl(), - KVCRConfig(nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)]), - local_dram=LocalDramOptions( - [ - ("full", ctypes.addressof(pools[0]), 8), - ("swa", ctypes.addressof(pools[1]), 8), - ] - ), - ) + kvcr = _two_pool_kvcr(agent, pools) full, swa, grouped = (BlockKey(name) for name in (b"full", b"swa", b"grouped")) descriptors = [ _mem_descriptor(ctypes.addressof(source), 8, info="full"), @@ -253,18 +254,7 @@ def test_group_allocation_evicts_enough_whole_keys() -> None: source = ctypes.create_string_buffer(24) agent = FakeNixlAgent() agent.state = "DONE" - kvcr = _new_kvcr( - agent, - FakePrimaryPinning(), - FakeBytesControl(), - KVCRConfig(nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)]), - local_dram=LocalDramOptions( - [ - ("full", ctypes.addressof(pools[0]), 16), - ("swa", ctypes.addressof(pools[1]), 16), - ] - ), - ) + kvcr = _two_pool_kvcr(agent, pools) full, swa0, swa1, grouped = ( BlockKey(name) for name in (b"full", b"swa0", b"swa1", b"grouped") ) @@ -529,21 +519,14 @@ def test_capacity_pressure_is_pool_local() -> None: capacity_requests: list[list[tuple[str, int]]] = [] agent = FakeNixlAgent() agent.state = "DONE" - kvcr = _new_kvcr( + kvcr = _two_pool_kvcr( agent, - FakePrimaryPinning(), - FakeBytesControl(), + pools, KVCRConfig( nixl_agent_name="target", pool_layouts=[("full", 8), ("swa", 8)], capacity_low_watermark_percent=100, ), - local_dram=LocalDramOptions( - [ - ("full", ctypes.addressof(pools[0]), 8), - ("swa", ctypes.addressof(pools[1]), 16), - ] - ), capacity_needed_callback=capacity_requests.append, ) descriptors = [ diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index 3a453df..a5d849e 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -123,6 +123,8 @@ class _RecoveryBlockV2(msgspec.Struct, frozen=True, array_like=True): msgspec.msgpack.encode({"g2": {"slot": 0, "state": "ready"}}), msgspec.msgpack.encode({"g3": {"slot": -1}}), msgspec.msgpack.encode({"g2": {"slot": "0"}}), + msgspec.msgpack.encode([0, None]), + msgspec.msgpack.encode([[[""]], None]), msgspec.msgpack.encode([[["other", 0]], None]), msgspec.msgpack.encode([[["", -1]], None]), msgspec.msgpack.encode([[], None]), From a2f1df61c2558558238b5d263f450fa4c0093c0d Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 07:35:08 +0300 Subject: [PATCH 15/16] test: verify multi-pool bytes through failover Cover native Guard transfers and reclaim with multiple descriptors. Consolidate overlapping transfer, malformed-input, and capacity cases. Signed-off-by: Kapil Arya --- tests/unit/test_guard_integration.py | 103 +++++++++++++++++++------- tests/unit/test_kvcr_local_dram.py | 46 ++---------- tests/unit/test_kvcr_remote_target.py | 63 +++++++++++----- tests/unit/test_recovery_mirror.py | 3 - 4 files changed, 128 insertions(+), 87 deletions(-) diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 98cb757..885e36b 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -106,12 +106,13 @@ def transfer(self, handle): def _make_kvcr( socket_path: str, - g3_path: str, + g3_path: str | None, control_port: int | str, agent_name: str, *, agent: FakeNixlAgent | None = None, framework: "ctypes.Array | None" = None, + pool_layouts: list[tuple[str, int]] | None = None, ) -> KVCR: """A claiming KVCR: a fake agent gets a MOCK G3, a real one POSIX plus NIXL-registered framework memory every descriptor it hands KVCR points into.""" @@ -121,7 +122,7 @@ def _make_kvcr( return KVCR( KVCRConfig( nixl_agent_name=agent_name, - pool_layouts=[("", page_size)], + pool_layouts=pool_layouts or [("", page_size)], nixl_listen_port=0, inventory_report_interval_ms=0, ), @@ -145,7 +146,9 @@ def _make_kvcr( page_size if agent is not None else page_size * 2 ), backend="MOCK" if agent is not None else "POSIX", - ), + ) + if g3_path is not None + else None, remote_fw_dram=RemoteFWDramOptions(eager_ctrl_connect=False), ), KVCRGuardConfig( @@ -322,25 +325,37 @@ def _real_nixl_runs_first(request: pytest.FixtureRequest) -> None: f"{request.node.name} must run first in this module; " f"{_RAN_BEFORE_REAL_NIXL[0]} already ran in this process" ) - _RAN_BEFORE_REAL_NIXL.append(request.node.name) + if "real_nixl" not in request.node.name: + _RAN_BEFORE_REAL_NIXL.append(request.node.name) +@pytest.mark.parametrize( + ("live_service", "multi_pool"), [(1, False), (2, True)], indirect=["live_service"] +) def test_a_promoted_guard_serves_real_nixl_transfers( tmp_path: Path, live_service: tuple[_KVCRService, Callable[..., subprocess.Popen[str]]], + multi_pool: bool, ) -> None: """With nothing faked, a promoted Guard serves a real UCX read then stands down.""" # Not a decorator: children import this module, and NIXL logs to their stdout. if not _real_nixl_available(): pytest.skip("no runnable NIXL agent on this machine") page_size = os.sysconf("SC_PAGE_SIZE") - second_payload = b"B" * page_size + layout = _real_nixl_layout(multi_pool) + second_payload = b"".join( + bytes([ord("B") + index]) * size for index, (_, size) in enumerate(layout) + ) g3_path = tmp_path / "g3.data" control_port = free_port() service, spawn = live_service primary = spawn( - "_real_nixl_primary_child", service.socket_path, g3_path, control_port + "_real_nixl_primary_child", + service.socket_path, + g3_path, + control_port, + multi_pool, ) _await_marker(primary, "ready", _REAL_NIXL_TIMEOUT_SECONDS) guard = service._registry._guards[0] @@ -353,12 +368,12 @@ def test_a_promoted_guard_serves_real_nixl_transfers( # A real UCX read through the Guard: the agent did not exist at write time. source_endpoint = f"tcp://127.0.0.1:{control_port}" - target_memory = ctypes.create_string_buffer(page_size) + target_memory = ctypes.create_string_buffer(len(second_payload)) target_pinning = FakePrimaryPinning() target = KVCR( KVCRConfig( nixl_agent_name="real-target", - pool_layouts=[("", page_size)], + pool_layouts=list(dict(layout).items()), nixl_listen_port=0, inventory_report_interval_ms=0, operation_timeout_ms=_REAL_NIXL_TIMEOUT_SECONDS * 1000, @@ -382,7 +397,11 @@ def test_a_promoted_guard_serves_real_nixl_transfers( served_key = BlockKey(b"resident-b") target.submit_hint(_router_hint(source_endpoint), request_id="from-guard") operation = target.deliver( - {served_key: [_mem_descriptor(ctypes.addressof(target_memory), page_size)]}, + { + served_key: _real_nixl_descriptors( + ctypes.addressof(target_memory), layout + ) + }, request_id="from-guard", ) deadline = time.monotonic() + _REAL_NIXL_TIMEOUT_SECONDS @@ -392,10 +411,7 @@ def test_a_promoted_guard_serves_real_nixl_transfers( time.sleep(0.01) assert results, "the Guard never answered the target" assert results[operation][served_key].success - assert ( - ctypes.string_at(ctypes.addressof(target_memory), page_size) - == second_payload - ) + assert target_memory.raw == second_payload # _serving stands for "answering peers": a served read must not end it. assert guard._serving is True finally: @@ -404,23 +420,24 @@ def test_a_promoted_guard_serves_real_nixl_transfers( # Only then does a replacement take the pool back and stand the Guard down. # Inline, not a child: adoption only claims and reads, and this process # already runs real agents beside the Guard's own. - framework = ctypes.create_string_buffer(page_size * 2) + framework = ctypes.create_string_buffer(len(second_payload) * 2) replacement = _make_kvcr( str(service.socket_path), - str(g3_path), + None if multi_pool else str(g3_path), control_port, "real-replacement", framework=framework, + pool_layouts=list(dict(layout).items()), ) try: - destination = ctypes.addressof(framework) + page_size - for key, payload in ( - (BlockKey(b"resident-a"), b"A" * page_size), - (BlockKey(b"resident-b"), second_payload), - ): + destination = ctypes.addressof(framework) + len(second_payload) + recovered = [(BlockKey(b"resident-b"), second_payload)] + if not multi_pool: + recovered.insert(0, (BlockKey(b"resident-a"), b"A" * page_size)) + for key, payload in recovered: ctypes.memset(destination, 0, len(payload)) operation = replacement.deliver( - {key: [_mem_descriptor(destination, len(payload))]} + {key: _real_nixl_descriptors(destination, layout)} ) result = dict( _poll_until(replacement, bool, timeout=_REAL_NIXL_TIMEOUT_SECONDS) @@ -759,13 +776,49 @@ def _real_nixl_available() -> bool: return True -def _real_nixl_primary_child(socket_path: str, g3_path: str, control_port: str) -> None: +def _real_nixl_layout(multi_pool: bool) -> list[tuple[str, int]]: + page = os.sysconf("SC_PAGE_SIZE") + return ( + [("full", page + page // 2), ("swa", page // 2), ("swa", page // 2)] + if multi_pool + else [("", page)] + ) + + +def _real_nixl_descriptors(address: int, layout: list[tuple[str, int]]): + descriptors = [] + for name, size in layout: + descriptors.append(_mem_descriptor(address, size, info=name)) + address += size + return descriptors + + +def _real_nixl_primary_child( + socket_path: str, g3_path: str, control_port: str, multi_pool: str +) -> None: """Fill the pool through a real agent, then hold the claim until killed.""" page_size = os.sysconf("SC_PAGE_SIZE") - framework = ctypes.create_string_buffer(page_size * 2) + layout = _real_nixl_layout(multi_pool == "True") + framework = ctypes.create_string_buffer(sum(size for _, size in layout) * 2) kvcr = _make_kvcr( - socket_path, g3_path, control_port, "real-primary", framework=framework + socket_path, + None if multi_pool == "True" else g3_path, + control_port, + "real-primary", + framework=framework, + pool_layouts=list(dict(layout).items()), ) - _deposit_two_blocks(kvcr, ctypes.addressof(framework), page_size) + if multi_pool == "True": + payload = b"".join( + bytes([ord("B") + index]) * size for index, (_, size) in enumerate(layout) + ) + ctypes.memmove(ctypes.addressof(framework), payload, len(payload)) + key = BlockKey(b"resident-b") + operation = kvcr.deposit( + {key: _real_nixl_descriptors(ctypes.addressof(framework), layout)} + ) + assert dict(_poll_until(kvcr, bool))[operation][key].success + else: + _deposit_two_blocks(kvcr, ctypes.addressof(framework), page_size) print("ready", flush=True) time.sleep(60) diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 49c28aa..05dd915 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -130,26 +130,6 @@ def test_local_deposit_deduplicates_and_evicts_fifo() -> None: ] -def test_local_transfer_accepts_multiple_blocks_in_one_pool() -> None: - agent = FakeNixlAgent() - agent.state = "DONE" - source = ctypes.create_string_buffer(32) - kvcr = _new_local_kvcr(agent, ctypes.create_string_buffer(32), 2) - - operation = kvcr.deposit( - { - BlockKey(b"key"): [ - _mem_descriptor(ctypes.addressof(source)), - _mem_descriptor(ctypes.addressof(source) + 16), - ] - } - ) - - assert dict(_poll_until(kvcr, lambda results: bool(results)))[operation][ - BlockKey(b"key") - ].success - - def test_local_dram_rejects_overlapping_pools() -> None: memory = ctypes.create_string_buffer(16) address = ctypes.addressof(memory) @@ -492,27 +472,6 @@ def test_local_claims_fetch_deliver_release_and_capacity() -> None: assert local.raw == b"b" * block_size -def test_capacity_needed_is_edge_triggered() -> None: - local = ctypes.create_string_buffer(10) - capacity_requests: list[list[tuple[str, int]]] = [] - kvcr = _new_local_kvcr( - FakeNixlAgent(), - local, - 10, - capacity_low_watermark_percent=20, - capacity_needed_callback=capacity_requests.append, - ) - - kvcr._core._update_capacity_pressure({"": 2}) - kvcr._core._update_capacity_pressure({"": 1}) - kvcr._core._update_capacity_pressure({"": 0}) - assert capacity_requests == [[("", 2)]] - - kvcr._core._update_capacity_pressure({"": 2}) - kvcr._core._update_capacity_pressure({"": 1}) - assert capacity_requests == [[("", 2)], [("", 2)]] - - def test_capacity_pressure_is_pool_local() -> None: pools = [ctypes.create_string_buffer(8), ctypes.create_string_buffer(16)] source = ctypes.create_string_buffer(24) @@ -535,9 +494,14 @@ def test_capacity_pressure_is_pool_local() -> None: _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"), ] + kvcr._core._update_capacity_pressure({"full": 1, "swa": 2}) + kvcr._core._update_capacity_pressure({"full": 0, "swa": 1}) kvcr._core._update_capacity_pressure({"full": 0, "swa": 0}) assert capacity_requests == [[("full", 1), ("swa", 2)]] kvcr._core._update_capacity_pressure({"full": 1, "swa": 2}) + kvcr._core._update_capacity_pressure({"full": 0, "swa": 1}) + assert capacity_requests == [[("full", 1), ("swa", 2)]] * 2 + kvcr._core._update_capacity_pressure({"full": 1, "swa": 2}) capacity_requests.clear() full = kvcr.deposit({BlockKey(b"full"): descriptors[:1]}, no_evict=True) diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index 81a022a..0706a4f 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -235,28 +235,53 @@ def test_remote_fetch_uses_local_then_framework_sources() -> None: @pytest.mark.parametrize( - ("expected_layout", "success"), - [(["swa", "full"], False), (["full", "swa"], True)], + ("layout", "expected_layout", "success"), + [ + ([("full", 16), ("swa", 8)], ["swa", "full"], False), + ([("full", 16), ("swa", 8)], ["full", "swa"], True), + ([("", 16), ("", 16)], ["", ""], True), + ], ) -def test_remote_fetch_validates_a_multi_pool_layout( - expected_layout: list[str], success: bool, caplog: pytest.LogCaptureFixture +def test_remote_fetch_preserves_block_layout_and_bytes( + layout: list[tuple[str, int]], + expected_layout: list[str], + success: bool, + caplog: pytest.LogCaptureFixture, ) -> None: - names = ("full", "swa") - layout = [(name, 8) for name in names] - source_primary = ctypes.create_string_buffer(16) - source_local = [ctypes.create_string_buffer(8) for _ in names] - target_local = [ctypes.create_string_buffer(8) for _ in names] - source_agent = FakeNixlAgent(metadata=b"source-md") + class CopyingWriteAgent(FakeNixlAgent): + def transfer(self, handle): + self.transfers.append(handle) + _, sources, _, destinations, _, _ = self.xfers[handle - 1] + for (src, src_size, _), (dst, dst_size, _) in zip( + sources, destinations, strict=True + ): + ctypes.memmove(dst, src, min(src_size, dst_size)) + return "PROC" + + payloads = [bytes([index + 1]) * size for index, (_, size) in enumerate(layout)] + source_primary = [ctypes.create_string_buffer(data, len(data)) for data in payloads] + pool_data = dict.fromkeys(dict(layout), b"") + for (name, _), data in zip(layout, payloads, strict=True): + pool_data[name] += data + source_local = { + name: ctypes.create_string_buffer(len(data)) for name, data in pool_data.items() + } + target_local = { + name: ctypes.create_string_buffer(len(data)) for name, data in pool_data.items() + } + source_agent = CopyingWriteAgent(metadata=b"source-md") target_agent = FakeNixlAgent(metadata=b"target-md") source_control = FakeBytesControl("tcp://source:1") target_control = FakeBytesControl("tcp://target:1") - config = KVCRConfig(nixl_agent_name="unused", pool_layouts=layout) + config = KVCRConfig( + nixl_agent_name="unused", pool_layouts=list(dict(layout).items()) + ) - def dram(memories: list[ctypes.Array]) -> LocalDramOptions: + def dram(memories: dict[str, ctypes.Array]) -> LocalDramOptions: return LocalDramOptions( [ (name, ctypes.addressof(memory), len(memory)) - for name, memory in zip(names, memories, strict=True) + for name, memory in memories.items() ] ) @@ -279,8 +304,8 @@ def dram(memories: list[ctypes.Array]) -> LocalDramOptions: ) key = BlockKey(b"multi-pool") descriptors = [ - _mem_descriptor(ctypes.addressof(source_primary) + index * 8, 8, info=name) - for index, name in enumerate(names) + _mem_descriptor(ctypes.addressof(memory), size, info=name) + for (name, size), memory in zip(layout, source_primary, strict=True) ] source_agent.state = "DONE" deposit = source.deposit({key: descriptors}) @@ -306,9 +331,11 @@ def dram(memories: list[ctypes.Array]) -> LocalDramOptions: result = dict(_poll_until(target, bool))[fetch][key] assert result.success is success if success: - assert [descriptor.info for descriptor in result.descriptors or ()] == list( - names - ) + assert [ + (descriptor.info, descriptor.size) + for descriptor in result.descriptors or () + ] == layout + assert {name: memory.raw for name, memory in target_local.items()} == pool_data def test_remote_staging_commits_available_prefix() -> None: diff --git a/tests/unit/test_recovery_mirror.py b/tests/unit/test_recovery_mirror.py index a5d849e..fd51a82 100644 --- a/tests/unit/test_recovery_mirror.py +++ b/tests/unit/test_recovery_mirror.py @@ -120,9 +120,6 @@ class _RecoveryBlockV2(msgspec.Struct, frozen=True, array_like=True): [ b"", msgspec.msgpack.encode({"g4": {"slot": 0}}), - msgspec.msgpack.encode({"g2": {"slot": 0, "state": "ready"}}), - msgspec.msgpack.encode({"g3": {"slot": -1}}), - msgspec.msgpack.encode({"g2": {"slot": "0"}}), msgspec.msgpack.encode([0, None]), msgspec.msgpack.encode([[[""]], None]), msgspec.msgpack.encode([[["other", 0]], None]), From 32ef8badef8df1d8363dbdceeede98df2eb5c214 Mon Sep 17 00:00:00 2001 From: Kapil Arya Date: Thu, 10 Sep 2026 07:51:47 +0300 Subject: [PATCH 16/16] test: allow slower native NIXL startup in CI Signed-off-by: Kapil Arya --- tests/unit/test_guard_integration.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 885e36b..fadc017 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -336,8 +336,13 @@ def test_a_promoted_guard_serves_real_nixl_transfers( tmp_path: Path, live_service: tuple[_KVCRService, Callable[..., subprocess.Popen[str]]], multi_pool: bool, + monkeypatch: pytest.MonkeyPatch, ) -> None: """With nothing faked, a promoted Guard serves a real UCX read then stands down.""" + # Native startup on CI can exceed the production thread timeout. + monkeypatch.setattr( + kvcr_progress, "_JOIN_TIMEOUT_SECONDS", _REAL_NIXL_TIMEOUT_SECONDS + ) # Not a decorator: children import this module, and NIXL logs to their stdout. if not _real_nixl_available(): pytest.skip("no runnable NIXL agent on this machine") @@ -797,6 +802,7 @@ def _real_nixl_primary_child( socket_path: str, g3_path: str, control_port: str, multi_pool: str ) -> None: """Fill the pool through a real agent, then hold the claim until killed.""" + kvcr_progress._JOIN_TIMEOUT_SECONDS = _REAL_NIXL_TIMEOUT_SECONDS page_size = os.sysconf("SC_PAGE_SIZE") layout = _real_nixl_layout(multi_pool == "True") framework = ctypes.create_string_buffer(sum(size for _, size in layout) * 2)