diff --git a/docs/design_overview.md b/docs/design_overview.md index 8d27e79..d972587 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -198,6 +198,7 @@ These statuses describe current KVCR knowledge, not a reservation or guarantee. `deposit` copies from framework-owned memory into the KVCR's pool, while `deliver` places data into a framework-provided destination. `deliver` does not name a source; source selection remains with the KVCR and router. The KVCR does not allocate or free framework memory. To serve from framework-owned memory, the KVCR acquires a pin asynchronously through `request_pin` and `poll_pin_results`, reusing covered keys and requesting only the remainder; the framework keeps it valid until the KVCR calls `release_pin`. +`release_pin` must be safely retryable: `False` or an exception leaves release pending; `True` means the framework accepts responsibility for completing release. A deployment may choose to use only framework-owned memory. In that case, it uses the pinning mechanism together with `deliver` and does not use `deposit`, `fetch`, or `release`. diff --git a/src/kvcr/control_channels.py b/src/kvcr/control_channels.py index f5fa932..3337b8f 100644 --- a/src/kvcr/control_channels.py +++ b/src/kvcr/control_channels.py @@ -295,6 +295,8 @@ def recv(self) -> list[bytes]: except zmq.Again: return messages except zmq.ZMQError: + # Warning suppression can be added if persistent ZMQ failures + # cause excessive receive logs. logger.warning("KVCR control recv failed", exc_info=True) return messages diff --git a/src/kvcr/core.py b/src/kvcr/core.py index cfc9497..346eb1d 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -773,7 +773,10 @@ def _discard_local_dram_fill(self, keys: Collection[BlockKey]) -> None: self._local_dram.discard_fill(keys) def _block_record(self, key: BlockKey) -> _BlockRecord: - return self._block_record_map.setdefault(key, _BlockRecord()) + record = self._block_record_map.get(key) + if record is None: + record = self._block_record_map[key] = _BlockRecord() + return record def _is_local_resident(self, key: BlockKey) -> bool: """Report local DRAM residency a new operation can still be served from. diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index 2df5be6..3fb73f4 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -432,14 +432,9 @@ def complete_fill(self, keys: Collection[BlockKey], *, success: bool) -> None: for key in ordered_keys: record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None - if ( - residency is None - or residency.state - not in ( - _LocalDramState.FILLING, - _LocalDramState.DISCARDING, - ) - or (success and residency.state is not _LocalDramState.FILLING) + if residency is None or residency.state not in ( + _LocalDramState.FILLING, + _LocalDramState.DISCARDING, ): raise RuntimeError(f"local DRAM fill state lost for {key!r}") slots.append(tuple(residency.slots)) @@ -605,6 +600,7 @@ def _apply_fill_result( source: CacheTier, ) -> None: committed: list[BlockKey] = [] + failed: list[BlockKey] = [] affected_residency_ops: dict[_OpId, _PendingResidencyOp] = {} affected_deliver_ops: dict[_OpId, _PendingDeliverOp] = {} deliver_keys: dict[_OpId, list[BlockKey]] = {} @@ -621,10 +617,12 @@ def _apply_fill_result( _LocalDramState.FILLING, _LocalDramState.DISCARDING, ) - or (success and residency.state is not _LocalDramState.FILLING) ): raise RuntimeError(f"local DRAM fill state lost for {key!r}") - if success: + # Main may discard a fill after progress queues its success. + # A terminal completion then frees the slot instead of committing it. + key_success = success and residency.state is _LocalDramState.FILLING + if key_success: record.last_access = now residency.state = _LocalDramState.READY self._residency_observer(key, record) @@ -637,11 +635,12 @@ def _apply_fill_result( else: record.local_dram = None self._free(residency.slots) + failed.append(key) for op_id in record.active_op_ids: residency_op = self._pending_residency_ops.get(op_id) if residency_op is not None and key in residency_op.keys: - if success and ( + if key_success and ( residency_op.op_id[0] == "deposit" or now < residency_op.deadline ): @@ -665,7 +664,7 @@ def _apply_fill_result( deliver_op = self._pending_deliver_ops.get(op_id) if deliver_op is not None and key in deliver_op.keys: - if success: + if key_success: deliver_keys.setdefault(op_id, []).append(key) else: deliver_op.results[key] = OpEntryResult(OpEntryStatus.FAILED) @@ -677,9 +676,8 @@ def _apply_fill_result( self._finish_residency_if_ready(residency_op) for op_id, deliver_op in affected_deliver_ops.items(): self._start_deliveries(deliver_op, deliver_keys.get(op_id, ())) - if not success: - for key in ordered_keys: - self._kvcr._prune_block_record(key) + for key in failed: + self._kvcr._prune_block_record(key) self._resume_capacity_waiters() def reserve_fill( diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index 0a7068f..8d4cd15 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -151,6 +151,9 @@ class _Entry: class _EvictionQueue: + # TODO: Bound stale heap growth from DRAM/G3 claim/release cycles. + # Removal only invalidates _live; select() removes stale heap entries as + # it encounters them, so repeated cache use can grow _heap without eviction. def __init__(self) -> None: self._heap: list[tuple[float, int, BlockKey]] = [] self._live: dict[BlockKey, _Entry] = {} diff --git a/src/kvcr/progress.py b/src/kvcr/progress.py index 2f31b3a..8dd79ca 100644 --- a/src/kvcr/progress.py +++ b/src/kvcr/progress.py @@ -19,8 +19,9 @@ logger = logging.getLogger(__name__) _IDLE_WAIT_SECONDS = 0.001 -_CLOSE_TIMEOUT_SECONDS = 5.0 +_OP_CLEANUP_TIMEOUT_SECONDS = 5.0 _JOIN_TIMEOUT_SECONDS = 10.0 +_STARTUP_TIMEOUT_SECONDS = 30.0 _RELEASE_LOG_INTERVAL_SECONDS = 1.0 _STOP = object() _OpId = tuple[str, Any] @@ -108,6 +109,7 @@ def __init__( ) self._failure: BaseException | None = None self._stop_requested = False + self._startup_stage = "thread startup" @property def nixl_agent(self) -> Any: @@ -256,8 +258,11 @@ def _release_transfer(self, transfer_id: int, state: _TransferState) -> bool: def start(self) -> None: self._thread.start() - if not self._ready.wait(timeout=_JOIN_TIMEOUT_SECONDS): - raise RuntimeError("KVCR progress thread did not start") + if not self._ready.wait(timeout=_STARTUP_TIMEOUT_SECONDS): + raise RuntimeError( + "KVCR progress initialization timed out after " + f"{_STARTUP_TIMEOUT_SECONDS:g}s (stage: {self._startup_stage})" + ) self.raise_if_failed() def submit(self, item: object) -> None: @@ -302,12 +307,17 @@ def close(self) -> None: def _run(self) -> None: try: + self._startup_stage = "NIXL agent initialization" self._initialize_nixl() # Let KVCR backends initialize NIXL resources before common # memory registration. + self._startup_stage = "backend initialization" self._initialize(self) + self._startup_stage = "memory registration" self._register_memory_regions() + self._startup_stage = "agent metadata capture" self._capture_agent_metadata() + self._startup_stage = "ready" self._ready.set() while not self._stop_requested: if not self._run_one_iteration(): @@ -315,6 +325,7 @@ def _run(self) -> None: except BaseException as error: self._failure = error finally: + self._startup_stage = "cleanup" try: try: self._close_progress_ops() @@ -329,7 +340,7 @@ def _run(self) -> None: self._ready.set() def _close_progress_ops(self) -> None: - deadline = time.monotonic() + _CLOSE_TIMEOUT_SECONDS + deadline = time.monotonic() + _OP_CLEANUP_TIMEOUT_SECONDS while self._in_flight_ops: for op_id, op in list(self._in_flight_ops.items()): if op.close(self): @@ -397,15 +408,14 @@ def _initialize_nixl(self) -> None: ) def _register_memory_regions(self) -> None: - if self._nixl_agent is None: + if self._nixl_agent is None or not self._memory_regions: return - for address, size in self._memory_regions: - self._memory_registrations.append( - self._nixl_agent.register_memory( - [(address, size, 0, "")], - mem_type="DRAM", - ) + self._memory_registrations.append( + self._nixl_agent.register_memory( + [(address, size, 0, "") for address, size in self._memory_regions], + mem_type="DRAM", ) + ) def _capture_agent_metadata(self) -> None: get_agent_metadata = getattr(self._nixl_agent, "get_agent_metadata", None) diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index 41e72cf..8bde839 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -411,6 +411,7 @@ def __init__( self._pending_pin_keys: dict[BlockKey, set[PinRequestId]] = {} # Pins retained while their source operations execute. self._fw_pins_by_op: dict[_OpId, set[PinHandle]] = {} + self._releasing_framework_pins: set[PinHandle] = set() # Progress-thread state: G2 control and outbound events. self._progress_outbound: list[object] = [] @@ -560,6 +561,7 @@ def fetch( ) def poll_main(self, items: Collection[object]) -> None: + self._release_framework_pins(tuple(self._releasing_framework_pins)) for item in items: if isinstance(item, _SourcePinOp): self._start_source_pin(item) @@ -603,6 +605,8 @@ def poll_main(self, items: Collection[object]) -> None: if now >= op.deadline: logger.warning("KVCR operation %r expired", op_id) self._expire_source_pin(op_id, op) + elif not op.pending_pin_ids: + self._resume_source_pin(op_id, op) def discard_hint(self, request_id: str) -> None: self._request_hints.pop(request_id, None) @@ -700,6 +704,8 @@ def poll_progress( else: raise TypeError(f"unsupported KVCR progress item: {type(item)!r}") + # Warning suppression can be added if persistent backend faults + # cause excessive polling logs. observed_work |= self._process_control_messages(progress) # A real notification outranks a refusal for the same operation. events = {**self._refused_writes, **self._poll_notifications(progress)} @@ -973,7 +979,7 @@ def _submit_prepared_source_write( and record.fw_mem is not None } sources = {} if force_failure else {**framework_sources, **local_sources} - completed_keys: tuple[BlockKey, ...] = () + completed_count = 0 for index, key in enumerate(source_pin.ordered_keys): source = sources.get(key) destination = source_pin.dst_descriptors[index] @@ -988,7 +994,8 @@ def _submit_prepared_source_write( key, ) break - completed_keys = source_pin.ordered_keys[: index + 1] + completed_count += 1 + completed_keys = source_pin.ordered_keys[:completed_count] kvcr._release_local_dram_sources( source_pin.op_id, local_sources.keys() - set(completed_keys) @@ -1083,7 +1090,7 @@ def _process_pending_pin_results(self) -> None: and request in op.pending_pin_ids ] if not ops: - self._discard_pin_result(result) + self._discard_pin_result(result, wait.keys if wait is not None else ()) continue now = kvcr._clock() @@ -1091,7 +1098,7 @@ def _process_pending_pin_results(self) -> None: active_ops = [(op_id, op) for op_id, op in ops if now < op.deadline] if not active_ops: self._record_pending_pin_wait(wait, "timeout") - self._discard_pin_result(result) + self._discard_pin_result(result, wait.keys) for op_id, op in expired_ops: op.pending_pin_ids.discard(request) self._expire_source_pin(op_id, op) @@ -1165,6 +1172,11 @@ def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: key for key in op.ordered_keys if key not in local_sources ) if unresolved_keys and not op.framework_acquire_attempted: + if any( + not kvcr._framework_pin_keys[pin].isdisjoint(unresolved_keys) + for pin in self._releasing_framework_pins + ): + return op.framework_acquire_attempted = True framework_sources = self._acquire_framework_sources(unresolved_keys) if isinstance(framework_sources, _PendingFrameworkSources): @@ -1270,7 +1282,9 @@ def _record_pending_pin_wait( if wait is not None: self._kvcr._record_duration("framework_pin_wait", wait.started_at, result) - def _discard_pin_result(self, result: PinResult) -> None: + def _discard_pin_result( + self, result: PinResult, keys: Collection[BlockKey] = () + ) -> None: if result is None: return try: @@ -1278,7 +1292,11 @@ def _discard_pin_result(self, result: PinResult) -> None: except (IndexError, TypeError): return if isinstance(pin_handle, str): - self._try_release_pin(pin_handle) + pin_keys = self._kvcr._framework_pin_keys.setdefault(pin_handle, set()) + pin_keys.update(keys) + if len(result) > 1 and isinstance(result[1], Mapping): + pin_keys.update(result[1]) + self._release_framework_pins((pin_handle,)) # Framework pin ownership. @@ -1310,7 +1328,6 @@ def _install_framework_pin( keys: Collection[BlockKey], pin_result: tuple[PinHandle, Mapping[BlockKey, list[MemDescriptor] | None]], ) -> PinHandle | None: - pin_handle: PinHandle | None = None try: pin_handle, descriptors = pin_result if not isinstance(pin_handle, str) or not isinstance(descriptors, Mapping): @@ -1340,8 +1357,7 @@ def _install_framework_pin( pin_keys.add(key) return pin_handle except Exception: - if pin_handle is not None: - self._try_release_pin(pin_handle) + self._discard_pin_result(pin_result, keys) return None def _acquire_framework_sources( @@ -1404,9 +1420,10 @@ def _release_framework_pins(self, framework_pins: Collection[PinHandle]) -> None pin_keys = kvcr._framework_pin_keys.get(pin_handle) if pin_keys is None: continue - if not self._try_release_pin(pin_handle): - continue - kvcr._framework_pin_keys.pop(pin_handle, None) + # Once release starts, its descriptors are no longer safe to reuse. + # Keep the keys until release is accepted, so overlapping pins wait. + first_attempt = pin_handle not in self._releasing_framework_pins + self._releasing_framework_pins.add(pin_handle) for key in pin_keys: record = kvcr._block_record_map.get(key) if ( @@ -1416,17 +1433,24 @@ def _release_framework_pins(self, framework_pins: Collection[PinHandle]) -> None ): record.fw_mem = None kvcr._prune_block_record(key) + if self._try_release_pin(pin_handle, warn=first_attempt): + kvcr._framework_pin_keys.pop(pin_handle, None) + self._releasing_framework_pins.discard(pin_handle) - def _try_release_pin(self, pin_handle: PinHandle) -> bool: + def _try_release_pin(self, pin_handle: PinHandle, *, warn: bool) -> bool: + # True means the framework accepts responsibility for completing release. + # False or exceptions cause retries, which must be safe after partial release. try: released = self._kvcr._release_pin_callback(pin_handle) except Exception: - logger.warning( - "KVCR release_pin failed for pin=%r", pin_handle, exc_info=True - ) + if warn: + logger.warning( + "KVCR release_pin failed for pin=%r", pin_handle, exc_info=True + ) return False - if released is False: - logger.warning("KVCR release_pin failed for pin=%r", pin_handle) + if released is not True: + if warn: + logger.warning("KVCR release_pin failed for pin=%r", pin_handle) return False return True diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 714641c..bfe3be9 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -329,6 +329,10 @@ def initialize_xfer( ): local_descs = list(local_descs) remote_descs = list(remote_descs) + if len(local_descs) != len(remote_descs) or any( + local[1] != remote[1] for local, remote in zip(local_descs, remote_descs) + ): + raise RuntimeError("NIXL rejected unaligned descriptors") self.xfers.append( ( op, @@ -348,10 +352,10 @@ def transfer(self, handle): handle - 1 ] if op == "WRITE" and remote_agent == self.name: - for local_index, remote_index in zip(local_indices, local_indices): + for local_index in local_indices: src_addr, src_size, _ = local_descs[local_index] - dst_addr, dst_size, _ = remote_descs[remote_index] - ctypes.memmove(dst_addr, src_addr, min(src_size, dst_size)) + dst_addr, _, _ = remote_descs[local_index] + ctypes.memmove(dst_addr, src_addr, src_size) return "PROC" def check_xfer_state(self, handle): diff --git a/tests/unit/test_guard.py b/tests/unit/test_guard.py index bc3f0e7..bc4c873 100644 --- a/tests/unit/test_guard.py +++ b/tests/unit/test_guard.py @@ -9,7 +9,6 @@ import queue import select import socket -import uuid from contextlib import nullcontext from unittest.mock import Mock @@ -261,7 +260,7 @@ def test_guard_lives_out_adopt_promote_and_readopt_in_ownership_order( ) closed: list[str] = [] order: list[object] = [] - constructed: list[tuple] = [] + agent_names: list[str] = [] cores: list[Mock] = [] channels: list[Mock] = [] attachment = _fake_attachment() @@ -270,7 +269,16 @@ def test_guard_lives_out_adopt_promote_and_readopt_in_ownership_order( # The seeding mechanics live on the core (adopt_recovery_records); this # test orders the Guard's calls around it, not what happens inside it. def new_core(config, bindings, backends) -> Mock: - constructed.append((config, bindings, backends)) + agent_names.append(config.nixl_agent_name) + assert config.nixl_listen_port == 0 + assert bindings.framework_control is channels[-1] + assert backends.local_dram == LocalDramOptions( + [("", 1234 + 8192, 2 * _PAGE_BLOCK_SIZE_BYTES)], + "REMOTE", + ) + assert backends.remote_fw_dram.backend == "REMOTE" + # A Guard serves G2 and holds G3 records for the returning primary. + assert backends.g3 is None core = Mock(_local_dram=Mock(), _g3=None, _block_record_map={}) def adopt(records) -> None: @@ -316,27 +324,13 @@ def new_channel() -> Mock: attach.assert_called_once_with(_PAGE_SPEC) assert journal.reset_called # Adoption only grants; a core exists once a promotion needs one. - assert constructed == [] + assert cores == [] with pytest.raises(RecoveryMirrorError, match="another tier configuration"): guard._refuse_incompatible(_tier(16)) promoted_records = guard._recovery.mirror._records guard._promote() - config, bindings, backends = constructed[0] - prefix = "KVCR-Guard-" - assert config.nixl_agent_name.startswith(prefix) - uuid.UUID(config.nixl_agent_name.removeprefix(prefix)) - assert config.nixl_listen_port == 0 - assert bindings.framework_control is channels[0] - assert backends.local_dram == LocalDramOptions( - [("", 1234 + 8192, 2 * _PAGE_BLOCK_SIZE_BYTES)], - "REMOTE", - ) - assert backends.remote_fw_dram.backend == "REMOTE" - # A Guard opens no G3: it serves the G2 half and keeps the rest for the - # primary that takes the pool back. - assert backends.g3 is None assert journal.pending == [] assert order == [ ("adopt", (first, second)), @@ -380,6 +374,7 @@ def new_channel() -> Mock: ) guard._promote() + assert len(set(agent_names)) == 2 assert set(guard._recovery._g3_records) == {fresh} assert guard._recovery._g3_records[fresh].slot == 7 assert order[3:] == [("adopt", (second, fresh)), "clear", "start"] diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index fadc017..36ff424 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -15,6 +15,7 @@ import msgspec import pytest +import zmq from _kvcr_test_utils import ( FakeNixlAgent, FakePrimaryPinning, @@ -207,6 +208,8 @@ def _group_primary_child(socket_path: str, control_port: str) -> None: _DIGEST, ("127.0.0.1", int(control_port)), ) + for index, (_, address, size) in enumerate(hold.local_dram.pools): + ctypes.memset(address, ord("A") + index, size) record = _recovered_record(g2=[("pool0", 0), ("pool1", 0)]) journal = RecoveryJournal(hold._attachment) journal.publish(*next(iter(_recovery_frames({BlockKey(b"grouped"): record})))) @@ -217,8 +220,6 @@ def _group_primary_child(socket_path: str, control_port: str) -> None: 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.""" - import zmq - context = zmq.Context() pull = context.socket(zmq.PULL) pull.setsockopt(zmq.RCVTIMEO, int(_TIMEOUT_SECONDS * 1000)) @@ -245,10 +246,21 @@ def _stale_peer_child(control_port: str, probe_port: str) -> None: return +@pytest.fixture +def _zmq_context(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + # Terminate this test's context after its service and clients close, so + # ZMQ background threads do not survive into the next native-NIXL test. + context = zmq.Context() + monkeypatch.setattr(zmq.Context, "instance", classmethod(lambda cls: context)) + yield + context.term() + + @pytest.fixture def live_service( tmp_path: Path, request: pytest.FixtureRequest, + _zmq_context: None, ) -> Iterator[tuple[_KVCRService, Callable[..., subprocess.Popen[str]]]]: """A service on its own thread; children it spawns die with it.""" pool_dir = tmp_path / "pools" @@ -308,31 +320,10 @@ def spawn(child_function: str, *args: object) -> subprocess.Popen[str]: assert not server_thread.is_alive() -_RAN_BEFORE_REAL_NIXL: list[str] = [] - - -@pytest.fixture(autouse=True) -def _real_nixl_runs_first(request: pytest.FixtureRequest) -> None: - # Enforced, not just asked for: this module's real-NIXL scenario builds - # real NIXL agents in this process, and real createXferReq starts failing - # with NIXL_ERR_INVALID_PARAM when fake-agent/ZMQ scenarios have run here - # first. Pre-existing sensitivity, reproduced without any of the - # surrounding tests' recent changes; worth its own investigation. A - # reordering (xdist, -p, a new test added above) fails loudly here instead - # of as an inscrutable NIXL error. - if "real_nixl" in request.node.name and _RAN_BEFORE_REAL_NIXL: - pytest.fail( - f"{request.node.name} must run first in this module; " - f"{_RAN_BEFORE_REAL_NIXL[0]} already ran in this process" - ) - 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( +def test_promoted_guard_serves_real_nixl_transfers( tmp_path: Path, live_service: tuple[_KVCRService, Callable[..., subprocess.Popen[str]]], multi_pool: bool, @@ -341,7 +332,7 @@ def test_a_promoted_guard_serves_real_nixl_transfers( """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 + kvcr_progress, "_STARTUP_TIMEOUT_SECONDS", _REAL_NIXL_TIMEOUT_SECONDS ) # Not a decorator: children import this module, and NIXL logs to their stdout. if not _real_nixl_available(): @@ -461,6 +452,7 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( ) -> None: """One crash moves both pools to the Guard and one claim takes both back.""" page_size = os.sysconf("SC_PAGE_SIZE") + payloads = [b"A" * (page_size + page_size // 2), b"B" * page_size] control_port = free_port() guard_agent = _FileBackedNixlAgent() guard_agent.state = "DONE" @@ -492,6 +484,10 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( page_size, ), ) + assert [ + ctypes.string_at(descriptor.addr, descriptor.size) + for descriptor in guard._core._local_dram._descriptors(record.local_dram.slots) + ] == payloads replacement = KVCRClient(service.socket_path).claim( 0, @@ -508,6 +504,12 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( assert recovered[key].local_dram == _LocalDramResidency( [("pool0", 0), ("pool1", 0)], _LocalDramState.READY ) + assert [ + ctypes.string_at(address, len(payload)) + for (_, address, _), payload in zip( + replacement.local_dram.pools, payloads, strict=True + ) + ] == payloads finally: replacement.release() @@ -672,9 +674,8 @@ def test_replacement_primary_takes_the_cache_back_from_a_guard( _wait_until(lambda: first_guard._serving, timeout=_TIMEOUT_SECONDS) assert first_guard._core._block_record_map == {} - # A stale peer's request gets a terminal refusal, not silence. In its own - # process, as a real peer is -- and because a ZMQ probe in this process - # destabilizes the real-NIXL test that follows. + # A stale peer's request gets a terminal refusal, not silence. It runs + # in its own process, as a real peer does. peer = spawn("_stale_peer_child", control_port, free_port()) _await_marker(peer, "refused") peer.wait(timeout=_TIMEOUT_SECONDS) @@ -802,7 +803,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 + kvcr_progress._STARTUP_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) diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index b92014f..49b3de8 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -275,6 +275,13 @@ def test_startup_timeout_retains_nonquiescent_resources( def create_core(*args, **kwargs): core = core_type(*args, **kwargs) cores.append(core) + ready_wait = core._progress._ready.wait + + def wait_for_ready(timeout): + assert entered.wait(timeout=1) + return ready_wait(timeout) + + monkeypatch.setattr(core._progress._ready, "wait", wait_for_ready) return core def create_agent(*_args, **_kwargs): @@ -292,12 +299,15 @@ def create_agent(*_args, **_kwargs): monkeypatch.setattr(kvcr_recovery, "_KVCRCore", create_core) monkeypatch.setattr(kvcr_recovery, "RecoveryJournal", Mock()) monkeypatch.setattr(kvcr_api, "_NONQUIESCENT_STARTUP_RESOURCES", retained) + monkeypatch.setattr(kvcr_progress, "_STARTUP_TIMEOUT_SECONDS", 0) monkeypatch.setattr(kvcr_progress, "_JOIN_TIMEOUT_SECONDS", 0) monkeypatch.setattr(kvcr_progress, "nixl_agent", create_agent) monkeypatch.setattr(kvcr_progress, "nixl_agent_config", lambda **kwargs: kwargs) try: - with pytest.raises(RuntimeError, match="progress thread did not start"): + with pytest.raises( + RuntimeError, match="timed out after 0s .*NIXL agent initialization" + ): KVCR( KVCRConfig( nixl_agent_name="target", @@ -316,7 +326,6 @@ def create_agent(*_args, **_kwargs): guard_config, ) - assert entered.wait(timeout=1) assert len(cores) == 1 core = cores[0] assert not core.is_quiescent() @@ -528,7 +537,7 @@ def deregister_memory(self, handle): kvcr = KVCR( KVCRConfig( nixl_agent_name="target", - pool_layouts=[("", 64)], + pool_layouts=[("full", 64), ("swa", 32)], nixl_listen_port=1234, ), KVCRBindings( @@ -541,7 +550,9 @@ def deregister_memory(self, handle): ), KVCRBackendConfigs( framework_dram=FrameworkDramInput(128, 256), - local_dram=LocalDramOptions([("", 384, 128)], "LOCAL"), + local_dram=LocalDramOptions( + [("full", 384, 128), ("swa", 512, 64)], "LOCAL" + ), remote_fw_dram=RemoteFWDramOptions(backend="REMOTE"), ), ) @@ -556,10 +567,9 @@ def deregister_memory(self, handle): assert len(set(lifecycle_threads)) == 1 assert lifecycle_threads[0] != main_thread assert agent.registrations == [ - ([(128, 256, 0, "")], "DRAM"), - ([(384, 128, 0, "")], "DRAM"), + ([(128, 256, 0, ""), (384, 128, 0, ""), (512, 64, 0, "")], "DRAM"), ] - assert agent.deregistered == [2, 1] + assert agent.deregistered == [1] @pytest.fixture diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index 18eb767..e99992f 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -333,11 +333,14 @@ def release_xfer_handle(self, handle): assert source_agent.telemetry_handles == [1] -def test_kvcr_pin_release_failure_is_logged_without_escaping(kvcr_caplog): +@pytest.mark.parametrize("failure", [False, None, 1, RuntimeError("release failed")]) +def test_kvcr_pin_release_failure_is_logged_and_retried(kvcr_caplog, failure): class FailingPinRelease(FakePrimaryPinning): def release_pin(self, pin_handle): self.unpins.append(pin_handle) - raise RuntimeError("release failed") + if isinstance(failure, Exception): + raise failure + return failure source_agent = FakeNixlAgent(metadata=b"source-md") pinning = FailingPinRelease() @@ -356,10 +359,77 @@ def release_pin(self, pin_handle): assert pinning.unpins == ["pin"] assert "pin" in kvcr._core._framework_pin_keys + assert not kvcr._core._block_record_map warnings = [record.getMessage() for record in kvcr_caplog.records] assert any("release_pin failed" in message for message in warnings) - # Shutdown reports the pin instead; drop it so the fixture can close. - kvcr._core._framework_pin_keys.clear() + + failure = True + kvcr.poll_completed() + assert pinning.unpins == ["pin", "pin"] + assert not kvcr._core._framework_pin_keys + + +@pytest.mark.parametrize("expires", [False, True]) +def test_framework_reacquisition_waits_for_pin_release(expires: bool) -> None: + now = 0.0 + agent = FakeNixlAgent() + agent.state = "DONE" + pinning = FakePrimaryPinning() + pinning.release_pin = Mock(return_value=False) + control = FakeBytesControl() + source = _new_kvcr(agent, pinning, control, name="source") + source._core._clock = lambda: now + backend = source._core._remote_fw_dram + key = BlockKey(b"k0") + + control.incoming.append(_start_write_message(1, key)) + _poll_until(source, lambda _: pinning.release_pin.called) + control.incoming.append(_start_write_message(2, key)) + _poll_until(source, lambda _: backend._source_pin_ops or len(agent.xfers) > 1) + assert pinning.searches == [(key,)] + assert len(agent.xfers) == 1 + + if expires: + now = 2.0 + _poll_until(source, lambda _: bool(agent.sent_notifs)) + assert _decode_notif(agent.sent_notifs[-1][1]) == { + "type": "write_done", + "op_handle": 2, + "success": False, + } + assert pinning.searches == [(key,)] + + pinning.release_pin.return_value = True + _poll_until(source, lambda _: not _has_outstanding_operations(source)) + assert pinning.searches == ([(key,)] if expires else [(key,), (key,)]) + assert len(agent.xfers) == (1 if expires else 2) + assert not source._core._framework_pin_keys + + +@pytest.mark.parametrize("late", [False, True], ids=["invalid-result", "late-result"]) +def test_discarded_framework_pin_release_is_retried(late: bool) -> None: + now = 0.0 + agent = FakeNixlAgent() + pinning = PendingPrimaryPinning() + pinning.release_pin = Mock(return_value=False) + control = FakeBytesControl() + source = _new_kvcr(agent, pinning, control, name="source") + source._core._clock = lambda: now + key = BlockKey(b"k0") + control.incoming.append(_start_write_message(1, key)) + _poll_until(source, lambda _: pinning.searches) + if late: + now = 2.0 + _poll_until(source, lambda _: pinning.cancelled) + pinning.complete(0, missing_indices=() if late else (0,)) + _poll_until(source, lambda _: pinning.release_pin.called) + assert "pin" in source._core._framework_pin_keys + assert agent.xfers == [] + + pinning.release_pin.return_value = True + source.poll_completed() + assert pinning.release_pin.call_count >= 2 + assert not source._core._framework_pin_keys def test_framework_pin_poll_failures_are_logged_without_escaping(kvcr_caplog): diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index ac5ea76..367a633 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -358,7 +358,12 @@ def test_remote_staging_commits_available_prefix() -> None: assert target.release((release_handle,)) == [(release_handle, True)] -def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal() -> None: +@pytest.mark.parametrize( + "completion_before_timeout", [False, True], ids=["late", "queued"] +) +def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal( + completion_before_timeout: bool, +) -> None: now = 0.0 block_size = 16 local = ctypes.create_string_buffer(block_size) @@ -384,21 +389,41 @@ def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal() -> None: _wait_until(lambda: bool(control.sent)) message = _decode_control_message(control.sent[0][1]) + if completion_before_timeout: + # Progress accepts success before expiry; main consumes it after expiry. + agent.notifs["source"] = [_write_done_notification(message["op_handle"])] + _wait_until(lambda: not target._core._progress._completed.empty()) + now = 0.02 assert _poll_until(target, lambda completed: bool(completed)) == [ (fetch, _op_entries({key: False})) ] - assert target.query((key,), "req") == [(QueryStatus.FETCHABLE, CacheTier.REMOTE_G2)] - assert _has_outstanding_operations(target) - blocked = target.deposit({replacement: [_mem_descriptor(size=block_size)]}) - assert list(target.poll_completed()) == [ - (blocked, _op_entries({replacement: False})) - ] - - agent.notifs["source"] = [_write_done_notification(message["op_handle"])] - assert _poll_until(target, lambda _: not _has_outstanding_operations(target)) == [] + if not completion_before_timeout: + assert target.query((key,), "req") == [ + (QueryStatus.FETCHABLE, CacheTier.REMOTE_G2) + ] + assert _has_outstanding_operations(target) + blocked = target.deposit({replacement: [_mem_descriptor(size=block_size)]}) + assert list(target.poll_completed()) == [ + (blocked, _op_entries({replacement: False})) + ] + + agent.notifs["source"] = [_write_done_notification(message["op_handle"])] + assert ( + _poll_until(target, lambda _: not _has_outstanding_operations(target)) == [] + ) + assert not _has_outstanding_operations(target) assert key not in target._core._block_record_map + # The terminal completion makes the single slot reusable. + primary = ctypes.create_string_buffer(b"a" * block_size, block_size) + agent.state = "DONE" + deposit = target.deposit( + {replacement: [_mem_descriptor(ctypes.addressof(primary), block_size)]} + ) + assert _poll_until(target, bool) == [(deposit, _op_entries({replacement: True}))] + assert local.raw == primary.raw + def test_kvcr_deliver_propagates_source_pin_miss(): target_agent = FakeNixlAgent(metadata=b"target-md")