diff --git a/docs/design_overview.md b/docs/design_overview.md index d972587..3870a1d 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -104,7 +104,7 @@ A KVCR-owned DRAM pool may be allocated by the framework and passed to KVCR, or `compatibility_manifest` identifies the framework, model, KV layout, and host representation needed to interpret cached data. A pool may contain multiple internal pools for different attention-head requirements. When `kvcr_guard_endpoint` is provided, KVCR attaches to the relevant preserved pool and verifies that the supplied manifest is compatible; initialization fails if the pool is unavailable or incompatible. -If the engine or GPU fails, KVCR-Guard fences the failed owner before activating its backup KVCR. A replacement in-process KVCR can attach to the preserved pool, recover the committed state, resynchronize inventory if needed, and assume ownership through a fenced handoff. Partial writes, in-flight operations, and framework-owned GPU or host memory are not recovered. Recovery and handoff must preserve committed-data integrity and prevent concurrent ownership. +If the engine or GPU fails, KVCR-Guard verifies that the owning process has died before activating its backup KVCR; a timeout alone is not sufficient. A replacement in-process KVCR can attach to the preserved pool, recover the committed state, resynchronize inventory if needed, and assume ownership through a fenced handoff. Partial writes, in-flight operations, and framework-owned GPU or host memory are not recovered. Recovery and handoff must preserve committed-data integrity and prevent concurrent ownership. ### State Model @@ -124,7 +124,22 @@ The event loop owns KVCR metadata mutations and never performs blocking external The active in-process KVCR owns the NIXL agent used for KVCR-owned memory and framework memory exposed through the KVCR bindings. A future integration may instead coordinate with a framework-owned agent. When resilience is enabled, the backup KVCR has its own agent but uses it only after a fenced takeover. -Each operation is bounded by a deadline. When an operation times out or is cancelled, KVCR reports caller-visible completion and begins safe release immediately. Framework pins are released as soon as their dependent work finishes, minimizing interference with framework scheduling. If NIXL may still access an underlying descriptor, physical release waits until the transfer has quiesced. Expired pins are not reused; any still-needed keys are acquired again. If physical cleanup extends beyond the deadline, it does so only for safe release, not further KVCR scheduling. +Each operation is bounded by a deadline. When an operation times out or is cancelled, KVCR begins failure handling and cleanup. Logical failure and memory reuse are separate: backing allocations and NIXL registrations remain valid until native access has quiesced. Framework pins may be handed back under the uncertainty contract below; KVCR-owned memory stays claimed while unresolved. Expired pins are not reused; any still-needed keys are acquired again. Cleanup beyond the deadline does not schedule further work for the failed operation. + +### Failed Peers and Dangling Operations + +An internal per-instance identifier distinguishes processes even when NIXL agent names are reused; it is not proof of process death. Remote `fetch` and `deliver` use NIXL writes from the source to the destination. Both deadlines are measured from operation start: `operation_timeout_ms` (`T`, default 1000) begins cancellation, and `abandon_timeout_ms` (`A`, default 5000) must be at least `2T`. At `T`, the destination probes the source; cancellation blocks unsubmitted work and sends an advisory while native cleanup continues. At `A`, unresolved operations report uncertainty before handing framework memory back. Nonterminal replies do not extend either deadline. Once cancelled, the operation stays failed; later native success only completes cleanup. + +`KVCRBindings.on_resilience_event` receives `TransferError` events during `poll_completed()` and may receive final cleanup events during `close()`: + +- `state="uncertain"` identifies regions that native work may still access. It arrives before KVCR relinquishes framework source pins or returns a completion for a framework destination. +- `state="quiesced"` carries the same `op_handle` and regions when that operation can no longer access them. It permits reclamation of that operation's hold, without restoring success or clearing other operations' overlapping holds. + +Source events identify original keys and local buffers; destination events identify regions, not their current keys. `op_handle` is local to the reporting KVCR; `source_blocks` or `destination_regions` identifies its role. Pair the two states by that local handle and role; their regions are unchanged. The framework owns quarantine and capacity policy for its memory. It may return `False` from `release_pin` to leave release pending, or `True` to accept responsibility while keeping the allocation quarantined. By default, `uncertain` logs at ERROR and `quiesced` at INFO. A custom callback may raise to its caller; pending releases and completions are retained for a later poll or close attempt. Callback failure does not stop native progress. Framework quarantine is opt-in; without a custom handler, events are only logged and framework-owned memory is not quarantined by KVCR. + +The source retains its existing native operation state until NIXL reports `DONE`; pending states, `ERR`, and state-query failures do not establish quiescence. Releasing a handle is not proof that native access stopped. The destination keeps its existing operation in `QUARANTINED` state as the tombstone, retaining the original source incarnation and local regions. After `A`, it probes every `T` until a matching terminal reply or Guard-confirmed death of that original source process resolves it. There is no separate tombstone registry or expiry deadline. Process-death resolution relies on the supported transport's process-lifetime contract; it is not a general revocation guarantee for persistent RDMA transports. Neither heartbeat timeout, a reused agent name, nor elapsed retention time resolves uncertainty. Guard takeover fails old operations rather than replaying them. + +KVCR-owned local G2 source claims and discarded destination fills remain held while unresolved, even indefinitely. Their slots cannot be evicted or reused until quiescence; framework quarantine alone cannot protect KVCR's allocator. Source progress stalls are reported once as a `RuntimeError` through the same callback. New source writes remain disabled until restart; native cleanup continues. --- @@ -162,12 +177,7 @@ framework.cancel_pin_request(pin_request_id) framework.release_pin(pin_handle) # release an acquired framework-owned source pin ``` -The list-shaped API allows a key to span multiple pools. A descriptor's `info` -can identify its pool and may be extended for other descriptor metadata. -`fetch` may receive the expected layout shared by its keys as an ordered list -of pool names so KVCR can allocate the destinations. Repeated names represent -multiple descriptors from the same pool. A single-pool caller using the empty -pool name may omit it. +The list-shaped API allows a key to span multiple pools. A descriptor's `info` can identify its pool and may be extended for other descriptor metadata. `fetch` may receive the expected layout shared by its keys as an ordered list of pool names so KVCR can allocate the destinations. Repeated names represent multiple descriptors from the same pool. A single-pool caller using the empty pool name may omit it. ### Operating flow @@ -197,8 +207,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. +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 `release_pin` accepts release. `release_pin` must be safely retryable: `False` or an exception leaves release pending; `True` means the framework accepts responsibility for completing release, including retaining uncertain allocations until the matching `quiesced` event. 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`. @@ -253,7 +262,7 @@ The router is never on the data path: KVCR instances execute transfers peer-to-p Peer transfers use a separate control channel for connection metadata, acknowledgements, and transfer control; payload bytes move directly through NIXL and never traverse the router or control channel. Peer-protocol versioning and compatibility checks may be added as needed. -The engine and router already maintain engine liveness, so loss of an in-process KVCR is covered by the engine's existing heartbeat path and does not require another KVCR heartbeat. KVCR-Guard sends its own heartbeat to the main process. If a future recovery design requires KVCR-Guard to advertise liveness or takeover directly to the router, that channel can be added then. +The engine and router already maintain engine liveness, so loss of an in-process KVCR is covered by the engine's existing heartbeat path and does not require another periodic KVCR heartbeat. If a future recovery design requires KVCR-Guard to advertise liveness or takeover directly to the router, that channel can be added then. --- diff --git a/docs/dev-guide.md b/docs/dev-guide.md index a345838..60859e0 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -657,6 +657,7 @@ one local DP rank and uses illustrative capacities and ports: "control_advertise_host": "127.0.0.1", "eager_ctrl_connect": true, "operation_timeout_ms": 1000, + "abandon_timeout_ms": 5000, "enable_telemetry": true } ] @@ -694,7 +695,8 @@ The important fields are: | `eager_ctrl_connect` | Establishes peer control earlier; disabling it moves setup onto the request path | | `local_dram_backend` | NIXL backend used for local DRAM transfers | | `remote_fw_dram_backend` | NIXL backend used for peer DRAM transfers | -| `operation_timeout_ms` | Deadline for KVCR operations; timeout begins safe cancellation and cleanup | +| `operation_timeout_ms` | Deadline for KVCR operations; timeout begins cancellation and cleanup | +| `abandon_timeout_ms` | Deadline from operation start to report unresolved memory as uncertain; default `5000`, at least twice `operation_timeout_ms` | | `enable_telemetry` | Publishes KVCR operation, transfer, and state metrics through the vLLM wrapper | For several local DP ranks, provide one `control_ports` entry per local rank in diff --git a/docs/quick-start.md b/docs/quick-start.md index c708cde..b21dec8 100644 --- a/docs/quick-start.md +++ b/docs/quick-start.md @@ -146,6 +146,7 @@ export KV_TRANSFER_CONFIG='{ "control_ports": [17771, 17772], "control_advertise_host": "127.0.0.1", "operation_timeout_ms": 5000, + "abandon_timeout_ms": 10000, "enable_telemetry": true } ] diff --git a/src/kvcr/api.py b/src/kvcr/api.py index aa1599b..366a39d 100644 --- a/src/kvcr/api.py +++ b/src/kvcr/api.py @@ -72,6 +72,9 @@ class KVCRBindings: stats_factory: Callable[[], TelemetryStats] | None = None policy: "KVCachePolicy | None" = None + # Resilience failures and transfer lifecycle events; defaults to logging. + on_resilience_event: Callable[[Exception], None] | None = None + class KVCR: """Framework-facing KV Cache Runner.""" diff --git a/src/kvcr/config.py b/src/kvcr/config.py index 1131961..c7f5362 100644 --- a/src/kvcr/config.py +++ b/src/kvcr/config.py @@ -124,6 +124,7 @@ class KVCRConfig: pool_layouts: PoolBlockLayouts enable_telemetry: bool = False operation_timeout_ms: int = 1000 + abandon_timeout_ms: int = 5000 inventory_report_interval_ms: int = 10 capacity_low_watermark_percent: float = 0 nixl_listen_port: int | None = None diff --git a/src/kvcr/core.py b/src/kvcr/core.py index 346eb1d..0ab2582 100644 --- a/src/kvcr/core.py +++ b/src/kvcr/core.py @@ -5,6 +5,7 @@ import functools import logging import time +from collections import deque from collections.abc import Callable, Collection, Iterable, Mapping from dataclasses import dataclass from math import ceil @@ -121,6 +122,10 @@ def __init__( self._block_sizes = dict(self.pool_layouts) if self.config.operation_timeout_ms <= 0: raise ValueError("operation_timeout_ms must be positive") + if self.config.abandon_timeout_ms < 2 * self.config.operation_timeout_ms: + raise ValueError( + "abandon_timeout_ms must be at least twice operation_timeout_ms" + ) if self.config.inventory_report_interval_ms < 0: raise ValueError("inventory_report_interval_ms must be non-negative") if not 0 <= self.config.capacity_low_watermark_percent <= 100: @@ -171,6 +176,8 @@ def __init__( ] = {} self._completion_queue: list[OpResult] = [] + self._resilience_errors: deque[Exception] = deque() + self._pending_progress_items: list[object] = [] self._joined_completions: dict[ OpHandle, tuple[set[BlockKey], dict[BlockKey, OpEntryResult]] ] = {} @@ -201,8 +208,14 @@ def __init__( ) # Import lazily to keep the concrete backend private to KVCR setup. + from .dangling_ops import _log_resilience_event from .remote_fw_dram import _RemoteFWDram + self._on_resilience_event_callback = ( + _log_resilience_event + if bindings.on_resilience_event is None + else bindings.on_resilience_event + ) self._local_dram = ( _LocalDram(self, local_dram_config) if local_dram_config is not None @@ -493,7 +506,10 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: # TODO: Expose individual entry completions as they become available. def poll_completed(self) -> Iterable[OpResult]: - progress_items = self._progress.take_completed() + self._progress.raise_if_failed() + self._notify_transfer_errors(self._progress.take_completed()) + progress_items = self._pending_progress_items + self._pending_progress_items = [] if self._g3 is not None: progress_items = self._g3.poll_main(progress_items) if self._local_dram is not None: @@ -504,6 +520,17 @@ def poll_completed(self) -> Iterable[OpResult]: self._completion_queue = [] return completed + def _notify_transfer_errors(self, progress_items: list[object]) -> None: + # Notify before cleanup can return a buffer to its allocator. Retain the + # batch if user code raises, consuming each notification exactly once. + for item in progress_items: + if isinstance(item, Exception): + self._resilience_errors.append(item) + else: + self._pending_progress_items.append(item) + while self._resilience_errors: + self._on_resilience_event_callback(self._resilience_errors.popleft()) + def abort( self, op_handle: OpHandle, @@ -545,6 +572,12 @@ def close(self) -> None: except BaseException as error: # noqa: BLE001 - re-raised below progress_error = error + # Shutdown can observe native completion after the last public poll. + # Deliver those lifecycle notifications before returning any memory. + self._notify_transfer_errors([]) + while progress_items := self._progress.take_completed(): + self._notify_transfer_errors(progress_items) + # Cleanup mutates backend state native operations still reference, and # a stopped thread does not prove they are done with it. The caller # unmaps the pool as soon as close() returns, so anything short of a diff --git a/src/kvcr/dangling_ops.py b/src/kvcr/dangling_ops.py new file mode 100644 index 0000000..e5ef16e --- /dev/null +++ b/src/kvcr/dangling_ops.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Cancellation, quarantine, and lifecycle events for unresolved remote writes.""" + +import time +import uuid +from dataclasses import dataclass, replace +from itertools import chain +from typing import TYPE_CHECKING, Any, Literal, cast + +from .core import logger +from .types import OpHandle, TransferError + +if TYPE_CHECKING: + from .progress import _KVCRProgress + from .remote_fw_dram import _RemoteFWDram, _SourceWriteOp, _TargetPullOp + + +def _log_resilience_event(error: Exception) -> None: + log = ( + logger.info + if isinstance(error, TransferError) and error.state == "quiesced" + else logger.error + ) + log("%s", error) + + +@dataclass +class _SourceWriteStatus: + submitted: bool = False + cancel_requested: bool = False + cancel_deadline: float | None = None + abandoned: bool = False + + +class _DanglingOps: + """Progress-thread resilience policy for operations that retain their memory.""" + + def __init__(self, backend: "_RemoteFWDram") -> None: + self._backend = backend + self.incarnation = uuid.uuid4().hex + self.dead_incarnations: set[str] = set() + self.sources: dict[str, str] = {} + self.source_writes: dict[tuple[str, OpHandle], _SourceWriteStatus] = {} + self._last_progress_at: float | None = None + self._source_stalled = False + + def begin_poll(self) -> None: + polled_at = time.monotonic() + self.check_source_progress() + self._last_progress_at = polled_at + + def check_source_progress(self) -> bool: + if self._source_stalled or self._last_progress_at is None: + return not self._source_stalled + gap_ms = (time.monotonic() - self._last_progress_at) * 1000 + timeout_ms = self._backend._kvcr.config.operation_timeout_ms + if gap_ms >= timeout_ms: + # Require restart so queued pre-stall requests cannot get fresh deadlines. + self._source_stalled = True + self._backend._progress_outbound.append( + RuntimeError( + f"KVCR source progress stalled for {gap_ms:.1f} ms " + f"(timeout: {timeout_ms} ms); " + "new source writes disabled until restart" + ) + ) + return not self._source_stalled + + def poll_source( + self, progress: "_KVCRProgress", op: "_SourceWriteOp", *, cancelling: bool + ) -> tuple[bool, Any | None] | None: + status = self.source_writes[(op.route[0], op.op_handle)] + if cancelling and status.cancel_deadline is None: + config = self._backend._kvcr.config + status.cancel_deadline = ( + op.deadline + + (config.abandon_timeout_ms - config.operation_timeout_ms) / 1000 + ) + self._backend._send_write_done( + progress, op.remote_agent, op.op_handle, False, terminal=False + ) + result = progress.poll_transfer( + cast(int, op.transfer_id), require_completion=True + ) + if ( + result is None + and cancelling + and not status.abandoned + and self._backend._kvcr._clock() >= cast(float, status.cancel_deadline) + ): + status.abandoned = True + self.report_source(op, "uncertain") + self._backend._progress_outbound.append(replace(op)) + return result + + def finish_source(self, op: "_SourceWriteOp") -> None: + status = self.source_writes.pop((op.route[0], op.op_handle), None) + if status is not None and status.abandoned: + self.report_source(op, "quiesced") + + def report_source( + self, op: "_SourceWriteOp", state: Literal["uncertain", "quiesced"] + ) -> None: + self._backend._progress_outbound.append( + TransferError( + "KVCR source write memory", + OpHandle(op.op_id[1]), + state=state, + source_blocks={ + key: list(descriptors) + for key, descriptors in zip(op.source_keys, op.src_descriptors) + }, + ) + ) + + def poll_target( + self, progress: "_KVCRProgress", op: "_TargetPullOp", now: float + ) -> bool: + first_poll = not op.uncertain + if first_poll: + op.uncertain = True + self.report_target(op, "uncertain") + self._backend._progress_outbound.append(replace(op)) + self._backend._invalidate_control_peer(op.remote_ctrl_ep) + if first_poll or now >= op.deadline: + op.deadline = now + self._backend._kvcr.config.operation_timeout_ms / 1000 + return self.probe(progress, op) + return False + + def report_target( + self, op: "_TargetPullOp", state: Literal["uncertain", "quiesced"] + ) -> None: + self._backend._progress_outbound.append( + TransferError( + "KVCR remote write memory", + op.op_id[1], + state=state, + destination_regions=list(chain.from_iterable(op.dst_descriptors)), + ) + ) + + def probe(self, progress: "_KVCRProgress", op: "_TargetPullOp") -> bool: + return self._backend._send_control( + progress, + op.remote_ctrl_ep, + { + "type": "write_probe", + "op_handle": op.op_id[1], + "source_incarnation": op.source_incarnation, + }, + ) + + def handle_probe(self, progress: "_KVCRProgress", payload: dict[str, Any]) -> None: + target_agent = payload.get("target_agent") + handle = payload.get("op_handle") + reply_to = payload.get("sender_control_endpoint") + endpoint = payload.get("source_control_endpoint") + expected = payload.get("source_incarnation") + if ( + not isinstance(target_agent, str) + or not target_agent + or type(handle) is not int + or not isinstance(reply_to, str) + or not reply_to + or not isinstance(endpoint, str) + or not endpoint + or (expected is not None and not isinstance(expected, str)) + ): + return + status = self.source_writes.get((target_agent, handle)) + if expected in (None, self.incarnation): + # A probe can overtake start_write after a control reconnect. Keep + # this small cancellation fence until restart, even if no op exists. + status = self.source_writes.setdefault( + (target_agent, handle), _SourceWriteStatus() + ) + status.cancel_requested = True + response = { + "type": "write_probe_ack", + "sender_control_endpoint": endpoint, + "op_handle": handle, + "terminal": expected in (None, self.incarnation) + and (status is None or not status.submitted), + } + if expected and expected in self.dead_incarnations: + response["dead_incarnation"] = expected + self._backend._send_control(progress, reply_to, response) + + def handle_probe_ack( + self, progress: "_KVCRProgress", payload: dict[str, Any] + ) -> None: + endpoint = payload.get("sender_control_endpoint") + handle = payload.get("op_handle") + terminal = payload.get("terminal") + incarnation = payload.get("sender_incarnation") + if ( + not isinstance(endpoint, str) + or type(handle) is not int + or type(terminal) is not bool + or not isinstance(incarnation, str) + or not incarnation + ): + return + op = cast( + "_TargetPullOp | None", progress._in_flight_ops.get(("target", handle)) + ) + if op is None or op.remote_ctrl_ep != endpoint: + return + expected = op.source_incarnation + if expected and payload.get("dead_incarnation") == expected: + # Guard verifies death of this process, not just an expired heartbeat. + self.sources[endpoint] = incarnation + self._backend._refused_writes[("target", handle)] = {"success": False} + elif expected == incarnation and terminal: + self._backend._refused_writes[("target", handle)] = {"success": False} diff --git a/src/kvcr/guard.py b/src/kvcr/guard.py index a0507ae..2ce1da5 100644 --- a/src/kvcr/guard.py +++ b/src/kvcr/guard.py @@ -379,6 +379,7 @@ def __init__( self._owner = owner self._refusing = refusing self._pool_lease = _PoolLease(guard_index) + self._dead_incarnations: set[str] = set() # Owned by the current primary. self._control: ZmqPeerControlChannel | None = None self._configured: _TierConfig | None = None @@ -410,6 +411,12 @@ def _fail(self, error: BaseException) -> None: self._phase = _Phase.FAILED self._escalate(error) + @property + def dead_incarnations(self) -> tuple[str, ...]: + """Snapshot confirmed deaths for a replacement primary's claim.""" + with self._phase_lock: + return tuple(sorted(self._dead_incarnations)) + def start(self) -> None: """Attach the pool and begin the lifecycle thread, before any claim.""" if self._started: @@ -622,6 +629,9 @@ def _observe_holder(self) -> None: # The process may still be alive: promoting could seat a # second server over a live mapping. raise OSError(f"pidfd poll returned without POLLIN: {flags:#x}") + if lease.incarnation is not None: + with self._phase_lock: + self._dead_incarnations.add(lease.incarnation) self._promote_for(lease) except BaseException as error: # noqa: BLE001 - service-fatal self._fail(error) @@ -883,6 +893,9 @@ def reject_pin(keys: object) -> int: ), ) self._core = core + core._remote_fw_dram._dangling_ops.dead_incarnations = ( + self._dead_incarnations.copy() + ) core.adopt_recovery_records(records) # A previous handover describes slots this Guard is about to move, and it is # already in the mirror. Leaving it would map keys to overwritten bytes. diff --git a/src/kvcr/guard_protocol.py b/src/kvcr/guard_protocol.py index 8074450..cee17e2 100644 --- a/src/kvcr/guard_protocol.py +++ b/src/kvcr/guard_protocol.py @@ -9,6 +9,7 @@ import os import socket import threading +import uuid from collections.abc import Callable from dataclasses import dataclass, field from typing import Annotated, Literal @@ -88,6 +89,7 @@ class _Claim(msgspec.Struct, frozen=True, tag="claim"): control_host: str control_port: Annotated[int, msgspec.Meta(ge=1, le=65535)] version: ProtocolVersion + incarnation: Annotated[str, msgspec.Meta(min_length=1)] | None = None def __post_init__(self) -> None: # A literal address, because the service binds this holding the lock @@ -115,6 +117,7 @@ class _Granted(msgspec.Struct, frozen=True, tag="granted"): tier_config: _TierConfig pools: tuple[_PoolDescriptor, ...] version: ProtocolVersion + dead_incarnations: tuple[str, ...] = () class _Released(msgspec.Struct, frozen=True, tag="released"): @@ -137,6 +140,7 @@ class PidfdLiveness: def __init__(self, pidfd: int) -> None: self._pidfd = pidfd + self.incarnation: str | None = None self._close_lock = threading.Lock() @classmethod @@ -183,6 +187,8 @@ class KVCRPoolHold: _attachment: KVCRPoolAttachment _connection: FramedConnection _control_listener_fd: int | None = None + _incarnation: str | None = None + _dead_incarnations: tuple[str, ...] = () _release_attempted: bool = field(default=False, init=False, repr=False) def hand_listener_to(self, adopt: Callable[[int], None]) -> None: @@ -256,6 +262,7 @@ def claim( "control_host": control_bind[0], "control_port": control_bind[1], "version": _PROTOCOL_VERSION, + "incarnation": uuid.uuid4().hex, }, type=_Claim, ) @@ -296,6 +303,8 @@ def claim( _attachment=attachment, _connection=connection, _control_listener_fd=listener_fd, + _incarnation=request.incarnation, + _dead_incarnations=response.dead_incarnations, ) except BaseException as error: # Release the lease only after local access has stopped, or the diff --git a/src/kvcr/kvcr_service.py b/src/kvcr/kvcr_service.py index 8111918..d12d15d 100644 --- a/src/kvcr/kvcr_service.py +++ b/src/kvcr/kvcr_service.py @@ -405,12 +405,17 @@ def dispatch( raise KVCRServiceError( "KVCR compatibility digest does not match the service" ) + liveness.incarnation = request.incarnation + # Retain the Guard across a concurrent close; claim validates the request. + guard = self.registry._guards.get(request.guard_index) spec, pools, listener_fd, lease = self.registry.claim( request.guard_index, request.tier_config, liveness, (request.control_host, request.control_port), ) + if guard is None: + raise KVCRServiceError("KVCR claim succeeded without a Guard") return ( _Granted( request.guard_index, @@ -418,6 +423,7 @@ def dispatch( request.tier_config, pools, _PROTOCOL_VERSION, + dead_incarnations=guard.dead_incarnations, ), (request.guard_index, listener_fd, lease), ) diff --git a/src/kvcr/progress.py b/src/kvcr/progress.py index 8dd79ca..8480877 100644 --- a/src/kvcr/progress.py +++ b/src/kvcr/progress.py @@ -134,6 +134,7 @@ def poll_transfer( transfer_id: int, *, cancellation_requested: bool = False, + require_completion: bool = False, ) -> tuple[bool, Any | None] | None: """Advance a transfer and return its result after releasing its handle.""" state = self._active_transfers.get(transfer_id) @@ -141,19 +142,23 @@ def poll_transfer( raise KeyError(f"unknown transfer {transfer_id}") agent = self.nixl_agent outcome = state.outcome - if outcome is None: + if outcome is None or (require_completion and not outcome): try: xfer_state = agent.check_xfer_state(state.handle) except Exception: - logger.warning( - "NIXL transfer progress failed", - exc_info=True, - ) + if state.outcome is not False: + logger.warning("NIXL transfer progress failed", exc_info=True) xfer_state = "ERR" + # Releasing a pending NIXL/UCX handle can leave DMA running and lose + # its completion signal. Remote writes retain it until actual DONE. + if require_completion and xfer_state != "DONE": + if xfer_state not in ("PROC", "PEND"): + state.outcome = False + return None pending = xfer_state in ("PROC", "PEND") if pending and not cancellation_requested: return None - outcome = xfer_state == "DONE" + outcome = xfer_state == "DONE" and state.outcome is not False if not pending: state.outcome = outcome if outcome and state.capture_telemetry: @@ -270,7 +275,6 @@ def submit(self, item: object) -> None: self._submissions.put(item) def take_completed(self) -> list[object]: - self.raise_if_failed() completed: list[object] = [] while len(completed) < self._batch_size: try: @@ -331,19 +335,29 @@ def _run(self) -> None: self._close_progress_ops() finally: try: - self._close() + self._completed_backlog.extend(self._flush()) + self._publish_completed(sys.maxsize) finally: - self._close_nixl() + try: + self._close() + finally: + self._close_nixl() except BaseException as error: if self._failure is None: self._failure = error self._ready.set() def _close_progress_ops(self) -> None: + self._stop_requested = True deadline = time.monotonic() + _OP_CLEANUP_TIMEOUT_SECONDS while self._in_flight_ops: + events, _ = self._poll(self, []) for op_id, op in list(self._in_flight_ops.items()): - if op.close(self): + closed = op.close(self) + # Pending cleanup still needs timers and peer replies. + if not closed: + closed, _ = op.progress(self, events.get(op_id)) + if closed: self._in_flight_ops.pop(op_id, None) if not self._in_flight_ops: return @@ -425,6 +439,8 @@ def _capture_agent_metadata(self) -> None: def _close_nixl(self) -> None: if self._active_transfers: raise RuntimeError("cannot close NIXL with active transfers") + if self._in_flight_ops: + raise RuntimeError("cannot close NIXL with unresolved operations") failure: BaseException | None = None pending_registrations: list[Any] = [] if self._nixl_agent is not None: diff --git a/src/kvcr/recovery_journal.py b/src/kvcr/recovery_journal.py index 7a6360e..ae5805c 100644 --- a/src/kvcr/recovery_journal.py +++ b/src/kvcr/recovery_journal.py @@ -501,6 +501,10 @@ def adopt_claimed_pool(core: _KVCRCore, claimed: ClaimedPool) -> None: which is what closes it on release. """ hold = claimed.hold + if hold._incarnation is not None: + dangling = core._remote_fw_dram._dangling_ops + dangling.incarnation = hold._incarnation + dangling.dead_incarnations.update(hold._dead_incarnations) 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) diff --git a/src/kvcr/remote_fw_dram.py b/src/kvcr/remote_fw_dram.py index e35d50d..47351ad 100644 --- a/src/kvcr/remote_fw_dram.py +++ b/src/kvcr/remote_fw_dram.py @@ -28,6 +28,7 @@ TRANSFER_BYTES_METRIC, logger, ) +from .dangling_ops import _DanglingOps, _SourceWriteStatus from .progress import _KVCRProgress, _Op, _OpId, _ProgressOp from .types import ( BlockKey, @@ -66,6 +67,7 @@ class _TargetPullState(Enum): START_WRITE = auto() WAITING_WRITE_DONE = auto() WAITING_TERMINAL = auto() + QUARANTINED = auto() FINISHED = auto() @@ -99,14 +101,21 @@ class _TargetPullOp(_RemoteOp): request_id: str | None = None success: bool = False completed_keys: set[BlockKey] = field(default_factory=set) + probe_sent: bool = False + source_incarnation: str | None = None + uncertain: bool = False def progress( self, progress: _KVCRProgress, event: object | None ) -> tuple[bool, bool]: backend = self._backend now = backend._kvcr._clock() + cancelled = isinstance(event, Mapping) and event.get("cancelled", False) scope = "remote_fetch" if self.local_fill else "remote_deliver" if self.state is _TargetPullState.START_WRITE: + self.source_incarnation = backend._dangling_ops.sources.get( + self.remote_ctrl_ep + ) if now >= self.deadline: self.success = False self.state = _TargetPullState.FINISHED @@ -119,6 +128,7 @@ def progress( "type": "start_write", "op_handle": self.op_id[1], "remaining_timeout_ms": (self.deadline - now) * 1000, + "source_incarnation": self.source_incarnation, "keys": list(self.ordered_keys), "dst_descriptors": self.dst_descriptors, }, @@ -131,16 +141,26 @@ def progress( self.state = _TargetPullState.WAITING_WRITE_DONE return False, True - if self.state in ( - _TargetPullState.WAITING_WRITE_DONE, - _TargetPullState.WAITING_TERMINAL, - ) and isinstance(event, Mapping): - success = event.get("success") is True and ( - not self.local_fill - or ( - self.state is _TargetPullState.WAITING_WRITE_DONE - and now < self.deadline - ) + if ( + self.state + in ( + _TargetPullState.WAITING_WRITE_DONE, + _TargetPullState.WAITING_TERMINAL, + _TargetPullState.QUARANTINED, + ) + and isinstance(event, Mapping) + and event.get("terminal", True) is True + ): + if self.uncertain: + backend._dangling_ops.report_target(self, "quiesced") + self.state = _TargetPullState.FINISHED + return True, True + # After cancellation, native success only permits cleanup. + success = ( + event.get("success") is True + and not cancelled + and self.state is _TargetPullState.WAITING_WRITE_DONE + and (not self.local_fill or now < self.deadline) ) try: if success: @@ -175,14 +195,26 @@ def progress( backend._record_progress_duration(scope, self.started_at, result) return True, True - if now >= self.deadline: + if self.state is _TargetPullState.QUARANTINED: + # Retain this tombstone indefinitely until quiescence is proven; + # elapsed time alone cannot make its destination safe to reuse. + return False, backend._dangling_ops.poll_target(progress, self, now) + + if now >= self.deadline or ( + cancelled and self.state is _TargetPullState.WAITING_WRITE_DONE + ): if self.state is _TargetPullState.WAITING_TERMINAL: - return False, False - # TODO: Safely retry operations that span primary-to-Guard promotion. A - # deadline says the write has not been reported, not that the source cannot - # still make it, so the destination is held rather than handed back. Until - # then, callers must submit a new request. + self.state = _TargetPullState.QUARANTINED + backend._dangling_ops.poll_target(progress, self, now) + backend._record_progress_duration(scope, self.started_at, "failed") + return False, True self.state = _TargetPullState.WAITING_TERMINAL + config = backend._kvcr.config + self.deadline += ( + config.abandon_timeout_ms - config.operation_timeout_ms + ) / 1000 + backend._invalidate_control_peer(self.remote_ctrl_ep) + self.probe_sent = backend._dangling_ops.probe(progress, self) if self.local_fill: backend._progress_outbound.append( replace( @@ -191,10 +223,17 @@ def progress( completed_keys=set(self.completed_keys), ) ) - backend._invalidate_control_peer(self.remote_ctrl_ep) return False, True + if self.state is _TargetPullState.WAITING_TERMINAL and not self.probe_sent: + # Retry failed enqueues without extending the grace period. + self.probe_sent = backend._dangling_ops.probe(progress, self) + return False, self.probe_sent return False, False + def close(self, progress: _KVCRProgress) -> bool: + # Shutdown must not release destinations still awaiting a remote write. + return self.state in (_TargetPullState.START_WRITE, _TargetPullState.FINISHED) + @dataclass class _SourcePinOp(_Op): @@ -227,6 +266,7 @@ class _SourceWriteOp(_RemoteOp): state: _SourceWriteState remote_agent: bytes op_handle: int + source_keys: tuple[BlockKey, ...] dst_descriptors: tuple[tuple[MemDescriptor, ...], ...] _backend: "_RemoteFWDram" = field(repr=False, compare=False) framework_pins: set[PinHandle] = field(default_factory=set) @@ -241,7 +281,14 @@ def progress( ) -> tuple[bool, bool]: backend = self._backend observed_work = False + write_id = (self.route[0], self.op_handle) + status = backend._dangling_ops.source_writes[write_id] if self.transfer_id is None: + if ( + not backend._dangling_ops.check_source_progress() + or status.cancel_requested + ): + self.state = _SourceWriteState.NOTIFY_FAILURE if ( self.state is _SourceWriteState.NOTIFY_FAILURE or backend._kvcr._clock() >= self.deadline @@ -254,6 +301,7 @@ def progress( backend._record_progress_duration( "source_write", self.started_at, "failed" ) + backend._dangling_ops.finish_source(self) return True, True if self.state is not _SourceWriteState.READY_TO_WRITE: raise RuntimeError(f"KVCR source operation {self.op_id!r} is not ready") @@ -276,8 +324,10 @@ def progress( backend._record_progress_duration( "source_write", self.started_at, "failed" ) + backend._dangling_ops.finish_source(self) return True, True submit_started_at = backend._kvcr._timer() + status.submitted = True try: transfer_id, submitted = progress.submit_transfer( "WRITE", @@ -325,26 +375,29 @@ def progress( "source_write", self.started_at, "failed" ) self.state = _SourceWriteState.FINISHED + backend._dangling_ops.finish_source(self) return True, True transfer_id = self.transfer_id if transfer_id is None: raise RuntimeError(f"KVCR source operation {self.op_id!r} lost transfer") - if ( - self.state is not _SourceWriteState.CANCEL_PENDING - and backend._kvcr._clock() >= self.deadline + if self.state is not _SourceWriteState.CANCEL_PENDING and ( + status.cancel_requested or backend._kvcr._clock() >= self.deadline ): self.state = _SourceWriteState.CANCEL_PENDING observed_work = True - transfer_result = progress.poll_transfer( - transfer_id, - cancellation_requested=self.state is _SourceWriteState.CANCEL_PENDING, + cancelling = self.state is _SourceWriteState.CANCEL_PENDING + transfer_result = backend._dangling_ops.poll_source( + progress, self, cancelling=cancelling ) if transfer_result is None: + if progress._active_transfers[transfer_id].outcome is False: + self.state = _SourceWriteState.CANCEL_PENDING return False, observed_work self.transfer_id = None success, telemetry = transfer_result - if success: + self.success = success and not cancelling + if self.success: backend._record_transfer_telemetry(telemetry) backend._record_progress_counter( TRANSFER_BLOCKS_METRIC, @@ -353,17 +406,24 @@ def progress( ) else: backend._send_write_done(progress, self.remote_agent, self.op_handle, False) - result = "success" if success else "failed" + result = "success" if self.success else "failed" backend._record_progress_duration("source_write", self.started_at, result) - self.success = success self.state = _SourceWriteState.FINISHED + backend._dangling_ops.finish_source(self) return True, True def close(self, progress: _KVCRProgress) -> bool: if self.transfer_id is not None: - if not progress.cancel_transfer(self.transfer_id): + if ( + progress.poll_transfer(self.transfer_id, require_completion=True) + is None + ): return False self.transfer_id = None + self._backend._send_write_done( + progress, self.remote_agent, self.op_handle, False + ) + self._backend._dangling_ops.finish_source(self) return True @@ -433,6 +493,7 @@ def __init__( self._metadata_acked_sources: set[str] = set() self._metadata_retry_after: dict[str, float] = {} self._refused_writes: dict[_OpId, dict[str, bool]] = {} + self._dangling_ops = _DanglingOps(self) self._next_source_op_id = 1 self._control = kvcr.framework_control @@ -576,15 +637,22 @@ def poll_main(self, items: Collection[object]) -> None: if isinstance(item, _SourcePinOp): self._start_source_pin(item) elif isinstance(item, _SourceWriteOp): - if item.state is _SourceWriteState.FINISHED: - self._fw_pins_by_op.pop(item.op_id, None) - self._kvcr._remove_block_dependencies(item) - if item.success: - self._kvcr._record_access( - self._kvcr._local_dram_sources_by_op.get(item.op_id, ()) - ) - self._kvcr._release_local_dram_sources(item.op_id) - self._release_framework_pins(item.framework_pins) + if item.state in ( + _SourceWriteState.FINISHED, + _SourceWriteState.CANCEL_PENDING, + ): + # Framework ownership can be handed back after uncertainty; + # KVCR-owned contents stay claimed until native quiescence. + pins = self._fw_pins_by_op.pop(item.op_id, None) + if item.state is _SourceWriteState.FINISHED: + self._kvcr._remove_block_dependencies(item) + if item.success: + self._kvcr._record_access( + self._kvcr._local_dram_sources_by_op.get(item.op_id, ()) + ) + self._kvcr._release_local_dram_sources(item.op_id) + if pins is not None: + self._release_framework_pins(pins) else: raise RuntimeError( f"KVCR source operation {item.op_id!r} returned to main " @@ -597,6 +665,12 @@ def poll_main(self, items: Collection[object]) -> None: "non-local target pull is waiting for terminal state" ) self._kvcr._discard_local_dram_fill(item.keys) + elif item.state is _TargetPullState.QUARANTINED: + self._finish_target_pull(item) + elif item.state is _TargetPullState.FINISHED and item.uncertain: + self._kvcr._remove_block_dependencies(item) + if item.local_fill: + self._kvcr._complete_local_dram_fill(item.keys, success=False) elif item.state is _TargetPullState.FINISHED: self._finish_target_pull(item) else: @@ -662,7 +736,8 @@ def close_main(self) -> None: def _finish_target_pull(self, op: _TargetPullOp) -> None: kvcr = self._kvcr - kvcr._remove_block_dependencies(op) + if op.state is not _TargetPullState.QUARANTINED: + kvcr._remove_block_dependencies(op) completed_keys = op.completed_keys if op.success else set() if not op.success: self._fail_request_hint(op.request_id) @@ -675,6 +750,8 @@ def _finish_target_pull(self, op: _TargetPullOp) -> None: missing_keys=(hint.missing_keys | op.keys) - completed_keys, ) if op.local_fill: + if op.state is _TargetPullState.QUARANTINED: + return # DISCARDING already failed callers, but still owns the slots. if completed_keys: kvcr._complete_local_dram_fill( tuple(key for key in op.ordered_keys if key in completed_keys), @@ -710,6 +787,7 @@ def initialize_progress(self, _progress: _KVCRProgress) -> None: def poll_progress( self, progress: _KVCRProgress, submissions: list[object] ) -> tuple[dict[object, object], bool]: + self._dangling_ops.begin_poll() observed_work = bool(submissions) for item in submissions: if isinstance(item, _TargetMetadataRequest): @@ -726,8 +804,11 @@ def poll_progress( # 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)} + # Only terminal notifications outrank a refusal for the same operation. + events = self._poll_notifications(progress) + for op_id, refusal in self._refused_writes.items(): + if not events.get(op_id, {}).get("terminal"): + events[op_id] = refusal self._refused_writes.clear() observed_work |= bool(events) return events, observed_work @@ -749,6 +830,11 @@ def flush_progress(self) -> list[object]: return outbound def close_progress(self) -> None: + # Accepted writes may still be waiting for main-thread pin acquisition. + for (target, handle), status in self._dangling_ops.source_writes.items(): + cached = self._remote_agents_by_target.get(target) + if not status.submitted and cached is not None: + self._send_write_done(self._kvcr._progress, cached[1], handle, False) close_control = getattr(self._control, "close", None) if close_control is not None: close_control() @@ -776,9 +862,13 @@ def _process_control_messages(self, progress: _KVCRProgress) -> bool: if message_type == "target_metadata": self._handle_target_metadata(progress, payload) elif message_type == "target_metadata_ack": - self._handle_target_metadata_ack(payload) + self._handle_target_metadata_ack(progress, payload) elif message_type == "write_refused": self._handle_write_refused(progress, payload) + elif message_type == "write_probe": + self._dangling_ops.handle_probe(progress, payload) + elif message_type == "write_probe_ack": + self._dangling_ops.handle_probe_ack(progress, payload) elif message_type == "start_write": self._handle_start_write(progress, payload) else: @@ -801,6 +891,7 @@ def _send_control( self._record_progress_duration("control_enqueue", started_at, "failed") return False payload["target_agent"] = kvcr.nixl_agent_name + payload["sender_incarnation"] = self._dangling_ops.incarnation sender_endpoint = getattr(self._control, "endpoint", None) if isinstance(sender_endpoint, str): payload.setdefault("sender_control_endpoint", sender_endpoint) @@ -866,10 +957,24 @@ def _handle_write_refused( # start_write make a source write. self._refused_writes[("target", op_handle)] = {"success": False} - def _handle_target_metadata_ack(self, payload: dict[str, Any]) -> None: + def _handle_target_metadata_ack( + self, progress: _KVCRProgress, payload: dict[str, Any] + ) -> None: source_endpoint = payload.get("sender_control_endpoint") if not isinstance(source_endpoint, str) or not source_endpoint: return + incarnation = payload.get("sender_incarnation") + if isinstance(incarnation, str) and incarnation: + self._dangling_ops.sources[source_endpoint] = incarnation + handle = payload.get("op_handle") + if type(handle) is int: + op = progress._in_flight_ops.get(("target", handle)) + if ( + isinstance(op, _TargetPullOp) + and op.remote_ctrl_ep == source_endpoint + and op.source_incarnation is None + ): + op.source_incarnation = incarnation self._metadata_acked_sources.add(source_endpoint) self._metadata_retry_after.pop(source_endpoint, None) @@ -891,6 +996,8 @@ def _ack_target_metadata( "type": "target_metadata_ack", "sender_control_endpoint": source_control_endpoint, } + if payload.get("type") == "start_write": + response["op_handle"] = payload["op_handle"] if not self._send_control(progress, target_control_endpoint, response): # Kept, route and cache both: an operation this start_write queued # still transfers over this route, and the peer re-sends its @@ -955,6 +1062,21 @@ def _handle_start_write( self._notify_start_write_failure(progress, payload, op_handle) return + write_id = (target_agent, OpHandle(op_handle)) + if status := self._dangling_ops.source_writes.get(write_id): + if status.cancel_requested and not status.submitted: + self._send_write_done(progress, remote_agent, op_handle, False) + return + expected = payload.get("source_incarnation") + if ( + progress._stop_requested + or not self._dangling_ops.check_source_progress() + or (expected is not None and expected != self._dangling_ops.incarnation) + ): + self._send_write_done(progress, remote_agent, op_handle, False) + return + self._dangling_ops.source_writes[write_id] = _SourceWriteStatus() + op_id = ("source", self._next_source_op_id) self._next_source_op_id += 1 self._progress_outbound.append( @@ -1055,6 +1177,7 @@ def _submit_prepared_source_write( route=source_pin.route, _backend=self, framework_pins=framework_pins, + source_keys=completed_keys, src_descriptors=tuple(tuple(sources[key]) for key in completed_keys), completed_indices=tuple(completed_indices), ) @@ -1526,12 +1649,14 @@ def _remote_agent( # Progress notifications, telemetry, and resource cleanup. - def _poll_notifications(self, progress: _KVCRProgress) -> dict[object, object]: + def _poll_notifications( + self, progress: _KVCRProgress + ) -> dict[_OpId, dict[str, Any]]: agent = progress.nixl_agent get_new_notifs = getattr(agent, "get_new_notifs", None) if get_new_notifs is None: return {} - events: dict[object, object] = {} + events: dict[_OpId, dict[str, Any]] = {} try: for notifs in get_new_notifs().values(): for raw in notifs: @@ -1542,7 +1667,15 @@ def _poll_notifications(self, progress: _KVCRProgress) -> dict[object, object]: op_handle = int(payload["op_handle"]) except (KeyError, TypeError, ValueError): continue - events[("target", op_handle)] = payload + op_id = ("target", op_handle) + previous = events.get(op_id, {}) + terminal = payload.get("terminal", True) is True + cancelled = previous.get("cancelled", False) or not payload.get( + "success", False + ) + terminal |= previous.get("terminal", False) + payload.update(terminal=terminal, cancelled=cancelled) + events[op_id] = payload except Exception: logger.warning("KVCR notification receive failed", exc_info=True) return {} @@ -1616,6 +1749,8 @@ def _send_write_done( remote_agent: bytes, op_handle: OpHandle, success: bool, + *, + terminal: bool = True, ) -> None: agent = progress.nixl_agent send_notif = getattr(agent, "send_notif", None) @@ -1626,7 +1761,9 @@ def _send_write_done( ) return try: - result = send_notif(remote_agent, _write_done_notif(op_handle, success)) + result = send_notif( + remote_agent, _write_done_notif(op_handle, success, terminal=terminal) + ) except Exception: logger.warning( "KVCR write_done notification failed for op=%d", @@ -1658,6 +1795,8 @@ def _write_done_notif( op_handle: OpHandle, success: bool, completed_indices: tuple[int, ...] = (), + *, + terminal: bool = True, ) -> bytes: payload: dict[str, Any] = { "type": "write_done", @@ -1666,6 +1805,8 @@ def _write_done_notif( } if success: payload["completed_indices"] = completed_indices + if not terminal: + payload["terminal"] = False return _NOTIF_PREFIX + msgspec.msgpack.encode(payload) diff --git a/src/kvcr/types.py b/src/kvcr/types.py index 1a4ace9..cbc8de8 100644 --- a/src/kvcr/types.py +++ b/src/kvcr/types.py @@ -5,7 +5,7 @@ from collections.abc import Mapping from dataclasses import dataclass from enum import Enum -from typing import Annotated, NewType +from typing import Annotated, Literal, NewType import msgspec @@ -47,6 +47,34 @@ class MemDescriptor: PinResult = tuple[PinHandle, Mapping[BlockKey, list[MemDescriptor] | None]] | None +class TransferError(RuntimeError): + """Lifecycle report for memory exposed by a failed transfer. + + Source reports identify the original keys and local buffers. Destination + reports contain only local regions. ``quiesced`` clears this operation's + hazard; it never makes the failed data valid or clears other operations. + Handles are local to this KVCR instance and report side (source/destination). + """ + + def __init__( + self, + message: str, + op_handle: OpHandle, + *, + state: Literal["uncertain", "quiesced"] = "uncertain", + source_blocks: dict[BlockKey, list[MemDescriptor]] | None = None, + destination_regions: list[MemDescriptor] | None = None, + ) -> None: + self.op_handle = op_handle + self.state = state + self.source_blocks = source_blocks + self.destination_regions = destination_regions + super().__init__( + f"{message}: state={state}, op={op_handle}, sources={source_blocks!r}, " + f"destinations={destination_regions!r}" + ) + + class OpEntryStatus(Enum): # TODO: Add specific statuses for timeout, abort, capacity, and unavailable sources. SUCCESS = "SUCCESS" diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 630ff6d..ba29312 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -423,6 +423,7 @@ def _write_done_notification( *, success: bool = True, completed_indices: tuple[int, ...] = (0,), + terminal: bool = True, ) -> bytes: payload = { "type": "write_done", @@ -431,6 +432,8 @@ def _write_done_notification( } if success: payload["completed_indices"] = completed_indices + if not terminal: + payload["terminal"] = False return b"KVCR:" + msgspec.msgpack.encode(payload) @@ -477,6 +480,7 @@ def _new_kvcr( inventory_sink=None, capacity_needed_callback=None, policy=None, + on_resilience_event=None, ) -> KVCR: config = replace( config @@ -501,6 +505,7 @@ def _new_kvcr( inventory_sink=inventory_sink, capacity_needed_callback=capacity_needed_callback, policy=policy, + on_resilience_event=on_resilience_event, stats_factory=(FakeTelemetryStats if config.enable_telemetry else None), ), KVCRBackendConfigs( diff --git a/tests/unit/test_guard_integration.py b/tests/unit/test_guard_integration.py index 36ff424..d18e1a9 100644 --- a/tests/unit/test_guard_integration.py +++ b/tests/unit/test_guard_integration.py @@ -373,6 +373,7 @@ def test_promoted_guard_serves_real_nixl_transfers( nixl_listen_port=0, inventory_report_interval_ms=0, operation_timeout_ms=_REAL_NIXL_TIMEOUT_SECONDS * 1000, + abandon_timeout_ms=_REAL_NIXL_TIMEOUT_SECONDS * 2000, ), KVCRBindings( target_pinning.request_pin, @@ -514,11 +515,15 @@ def test_two_pool_group_survives_guard_failover_and_reclaim( replacement.release() -@pytest.mark.parametrize("recovery", ["kept", "given-up"]) +@pytest.mark.parametrize( + ("recovery", "late_promotion"), + [("kept", False), ("kept", True), ("given-up", True)], +) def test_request_timeout_during_promotion_then_retry_uses_guard( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, recovery: str, + late_promotion: bool, live_service: tuple[_KVCRService, Callable[..., subprocess.Popen[str]]], ) -> None: """A retry after promotion is served warm or refused cold, never left hanging.""" @@ -569,6 +574,7 @@ def pause_promotion(guard: _Guard) -> None: nixl_agent_name="target", pool_layouts=[("", page_size)], operation_timeout_ms=5000, + abandon_timeout_ms=10_000, ), remote_options=RemoteFWDramOptions(eager_ctrl_connect=False), framework_dram=FrameworkDramInput( @@ -581,7 +587,7 @@ def pause_promotion(guard: _Guard) -> None: key = BlockKey(b"resident-b") stalled_destination = (ctypes.c_char * page_size).from_buffer(target_memory) target.submit_hint(_router_hint(source_endpoint), request_id="stalled") - target.deliver( + stalled_operation = target.deliver( {key: [_mem_descriptor(ctypes.addressof(stalled_destination), page_size)]}, request_id="stalled", ) @@ -609,11 +615,27 @@ def pause_promotion(guard: _Guard) -> None: assert list(target.poll_completed()) == [] assert _has_outstanding_operations(target) + if late_promotion: + now[0] = 10.0 + completed = _poll_until(target, bool, timeout=2) + assert completed[0][0] == stalled_operation + assert not completed[0][1][key].success continue_promotion.set() _wait_until(lambda: guard._serving, timeout=2) + assert guard._core is not None + _wait_until( + lambda: ( + target._core._remote_fw_dram._dangling_ops.sources.get(source_endpoint) + == guard._core._remote_fw_dram._dangling_ops.incarnation + ), + timeout=2, + ) + if not late_promotion: + completed = _poll_until(target, bool, timeout=2) + assert completed[0][0] == stalled_operation + assert not completed[0][1][key].success # A real core either way, answering on the endpoint it inherited. - assert guard._core is not None destination = (ctypes.c_char * page_size).from_buffer(target_memory, page_size) target.submit_hint(_router_hint(source_endpoint), request_id="retry") operation = target.deliver( @@ -669,6 +691,7 @@ def test_replacement_primary_takes_the_cache_back_from_a_guard( idle = spawn("_primary_child", service.socket_path, g3_path, control_port, "idle") _await_marker(idle, "ready") first_guard = service._registry._guards[0] + idle_incarnation = first_guard._pool_lease.current.incarnation idle.kill() idle.wait(timeout=_TIMEOUT_SECONDS) _wait_until(lambda: first_guard._serving, timeout=_TIMEOUT_SECONDS) @@ -684,6 +707,7 @@ def test_replacement_primary_takes_the_cache_back_from_a_guard( "_primary_child", service.socket_path, g3_path, control_port, "held" ) _await_marker(primary, "ready") + primary_incarnation = first_guard._pool_lease.current.incarnation primary.kill() primary.wait(timeout=_TIMEOUT_SECONDS) @@ -703,6 +727,10 @@ def test_replacement_primary_takes_the_cache_back_from_a_guard( agent=replacement_agent, ) try: + assert replacement._core._remote_fw_dram._dangling_ops.dead_incarnations == { + idle_incarnation, + primary_incarnation, + } for key, payload in ( (BlockKey(b"resident-a"), b"A" * page_size), (BlockKey(b"resident-b"), b"B" * page_size), diff --git a/tests/unit/test_guard_protocol.py b/tests/unit/test_guard_protocol.py index 7509b24..fcbd55b 100644 --- a/tests/unit/test_guard_protocol.py +++ b/tests/unit/test_guard_protocol.py @@ -257,7 +257,11 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( ) -> None: """A claim/release round-trips typed wire messages, geometry, and ownership.""" events: list[str] = [] - connection = _RecordingConnection([_grant(), _Released(1)], events) + grant = msgspec.structs.replace(_grant(), dead_incarnations=("dead-primary",)) + decoded = protocol_module._CLAIM_RESPONSE_DECODER.decode( + msgspec.msgpack.encode(grant) + ) + connection = _RecordingConnection([decoded, _Released(1)], events) attachment = _Attachment(events) attach = Mock(return_value=attachment) _connect_with(monkeypatch, connection) @@ -267,6 +271,8 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( _GUARD_INDEX, _POOL_LAYOUTS, _DIGEST, ("127.0.0.1", 5555) ) + assert hold._incarnation + assert hold._dead_incarnations == ("dead-primary",) assert msgspec.to_builtins(connection.sent[0]) == { "type": "claim", "guard_index": _GUARD_INDEX, @@ -279,8 +285,10 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( "control_host": "127.0.0.1", "control_port": 5555, "version": 1, + "incarnation": hold._incarnation, } - grant_wire = msgspec.to_builtins(_grant()) + grant_wire = msgspec.to_builtins(grant) + assert grant_wire["dead_incarnations"] == ("dead-primary",) assert grant_wire["type"] == "granted" assert grant_wire["version"] == 1 assert grant_wire["guard_index"] == _GUARD_INDEX @@ -290,7 +298,7 @@ def test_claim_and_release_round_trip_typed_messages_and_geometry( "remote_fw_dram_backend": "UCX", } assert grant_wire["pools"] == msgspec.to_builtins(_WIRE_POOLS) - attach.assert_called_once_with(_grant().spec) + attach.assert_called_once_with(grant.spec) 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 @@ -383,11 +391,10 @@ def test_a_failed_claim_is_released_without_masking_the_original( assert raised.value is original # A mismatched or undecodable grant is refused before the pool is mapped. assert attach.call_count == (1 if mapping_error else 0) - assert connection.sent == [ - _Claim(_GUARD_INDEX, _DIGEST, _TIER_CONFIG, "127.0.0.1", 5555, 1), - # Unactivated: this claim never served, so the Guard may resume. - _Release(1, activated=False), - ] + claim, release = connection.sent + assert isinstance(claim, _Claim) + # Unactivated: this claim never served, so the Guard may resume. + assert release == _Release(1, activated=False) assert connection.closed is True diff --git a/tests/unit/test_kvcr.py b/tests/unit/test_kvcr.py index 49b3de8..b75031a 100644 --- a/tests/unit/test_kvcr.py +++ b/tests/unit/test_kvcr.py @@ -50,7 +50,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 = SimpleNamespace(_incarnation=None, _dead_incarnations=(), **fields) hold.hand_listener_to = partial(KVCRPoolHold.hand_listener_to, hold) return hold @@ -450,6 +450,14 @@ def test_service_dram_rejects_explicit_local_dram_before_claim(monkeypatch) -> N client.assert_not_called() +@pytest.mark.parametrize("abandon_timeout_ms", [1999, 2000]) +def test_kvcr_validates_abandon_timeout(abandon_timeout_ms): + config = KVCRConfig("target", [("", 16)], abandon_timeout_ms=abandon_timeout_ms) + expected = pytest.raises(ValueError, match="abandon_timeout_ms") + with expected if abandon_timeout_ms == 1999 else nullcontext(): + _new_kvcr(FakeNixlAgent(), FakePrimaryPinning(), FakeBytesControl(), config) + + def test_kvcr_rejects_no_dram_backends() -> None: with pytest.raises(ValueError, match="at least one DRAM backend"): KVCR( @@ -601,6 +609,9 @@ def close(self) -> None: def is_quiescent(self) -> bool: return self._quiescent + def take_completed(self) -> list[object]: + return [] + def test_close_cleans_backends_once_when_progress_is_quiescent( monkeypatch, new_kvcr diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index 6b1159e..12df7b6 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """KVCR remote framework-DRAM source-side tests.""" +import ctypes import logging import threading import time +from dataclasses import replace from types import SimpleNamespace from unittest.mock import Mock, patch @@ -16,22 +18,38 @@ FakePrimaryPinning, FakeTelemetryStats, PendingPrimaryPinning, + _decode_control_message, _decode_notif, _has_outstanding_operations, _mem_descriptor, _new_kvcr, + _op_entries, _poll_until, _start_write_message, _wait_until, ) from kvcr import DURATION_METRIC, TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC -from kvcr.config import KVCRConfig +from kvcr.config import KVCRConfig, LocalDramOptions from kvcr.core import _BlockRecord, _KVCRCore +from kvcr.progress import _STOP from kvcr.remote_fw_dram import _FwMemResidency, _RemoteFWDram, _SourcePinOp from kvcr.types import BlockKey, PinHandle, PinRequestId +def _write_probe_message(op_handle: int, incarnation=None) -> bytes: + return msgspec.msgpack.encode( + { + "type": "write_probe", + "source_incarnation": incarnation, + "op_handle": op_handle, + "target_agent": "target", + "sender_control_endpoint": "tcp://target:1", + "source_control_endpoint": "tcp://source:1", + } + ) + + @pytest.mark.parametrize("pin_before_deadline", [True, False]) def test_kvcr_start_write_respects_framework_pin_deadline( pin_before_deadline: bool, @@ -56,6 +74,7 @@ def clock() -> float: nixl_agent_name="source", pool_layouts=[("", 16)], operation_timeout_ms=10_000, + abandon_timeout_ms=20_000, ), name="source", ) @@ -85,6 +104,89 @@ def clock() -> float: } +@pytest.mark.parametrize("incarnation", [None, "matching", "other"]) +@pytest.mark.parametrize("before_start", [False, True]) +def test_write_probe_fences_write_waiting_on_framework_pin( + incarnation, before_start +) -> None: + agent = FakeNixlAgent(metadata=b"source-md") + pinning = PendingPrimaryPinning() + control = FakeBytesControl("tcp://source:1") + source = _new_kvcr(agent, pinning, control, name="source") + start = _start_write_message(9, BlockKey(b"k0"), target_agent="target") + if not before_start: + control.incoming.append(start) + _poll_until(source, lambda _: bool(pinning.pending)) + + if incarnation == "matching": + incarnation = source._core._remote_fw_dram._dangling_ops.incarnation + control.incoming.append(_write_probe_message(9, incarnation)) + _wait_until(lambda: bool(control.sent)) + assert _decode_control_message(control.sent[-1][1])["terminal"] is ( + incarnation != "other" + ) + + if before_start: + if incarnation != "other": + for metadata in (b"target-md", b"reconnected-target-md"): + payload = msgspec.msgpack.decode(start) + payload["target_agent_metadata"] = metadata + control.incoming.append(msgspec.msgpack.encode(payload)) + _wait_until(lambda: not control.incoming) + time.sleep(0.01) + assert list(source.poll_completed()) == [] + assert pinning.searches == [] + assert agent.xfers == [] + return + control.incoming.append(start) + _poll_until(source, lambda _: bool(pinning.pending)) + pinning.complete(0) + if incarnation == "other": + assert _poll_until(source, lambda _: bool(agent.xfers)) == [] + agent.state = "DONE" + assert _poll_until(source, lambda _: bool(pinning.unpins)) == [] + assert len(agent.xfers) == (1 if incarnation == "other" else 0) + + +def test_stalled_source_refuses_queued_and_future_writes() -> None: + errors = [] + + def on_resilience_event(error): + errors.append(error) + raise error + + agent = FakeNixlAgent(metadata=b"source-md") + pinning = PendingPrimaryPinning() + control = FakeBytesControl() + source = _new_kvcr( + agent, pinning, control, name="source", on_resilience_event=on_resilience_event + ) + + stalled_for = 0.0 + with patch( + "kvcr.dangling_ops.time", + SimpleNamespace(monotonic=lambda: time.monotonic() + stalled_for), + ): + for handle in (1, 2): + control.incoming.append( + _start_write_message(handle, BlockKey(b"k0"), target_agent="target") + ) + if handle == 1: + _poll_until(source, lambda _: bool(pinning.searches)) + stalled_for = 2.0 + pinning.complete(0) + with pytest.raises(RuntimeError, match="source progress stalled"): + _poll_until(source, lambda _: bool(errors)) + assert _poll_until(source, lambda _: len(agent.sent_notifs) == handle) == [] + assert not _decode_notif(agent.sent_notifs[-1][1])["success"] + assert agent.xfers == [] + assert pinning.searches == [(BlockKey(b"k0"),)] + assert pinning.unpins == ["pin"] + assert len(errors) == 1 # A raising callback must not stop native cleanup. + assert "timeout: 1000 ms" in str(errors[0]) + assert "new source writes disabled until restart" in str(errors[0]) + + def test_kvcr_close_cleans_pending_pin_operations(): agent = FakeNixlAgent(metadata=b"source-md") pinning = PendingPrimaryPinning() @@ -101,8 +203,22 @@ def test_kvcr_close_cleans_pending_pin_operations(): ) assert source._core._remote_fw_dram._pending_pin_ops + # The second request is accepted but has not reached main-thread pinning. + control.incoming.append(_start_write_message(2, key)) + _wait_until( + lambda: len(source._core._remote_fw_dram._dangling_ops.source_writes) == 2 + ) + notifications_at_close = [] + control.close = lambda: notifications_at_close.extend( + _decode_notif(raw) for _, raw in agent.sent_notifs + ) source.close() + assert notifications_at_close == [ + {"type": "write_done", "op_handle": handle, "success": False} + for handle in (1, 2) + ] + assert agent.xfers == [] assert pinning.cancelled == [PinRequestId(0)] assert not source._core._remote_fw_dram._source_pin_ops assert not source._core._remote_fw_dram._pending_pin_ops @@ -174,7 +290,7 @@ def send_notif(self, agent_name, notif_msg): ) -@pytest.mark.parametrize("failure", ["initialize", "error", "exception"]) +@pytest.mark.parametrize("failure", ["initialize", "error", "exception", "async"]) def test_kvcr_source_transfer_error_notifies_failure_and_cleans_up( failure: str, ) -> None: @@ -190,7 +306,7 @@ def transfer(self, handle): self.transfers.append(handle) if failure == "exception": raise RuntimeError("ambiguous submission") - return "ERR" + return "PROC" if failure == "async" else "ERR" source_agent = FailingTransferAgent(metadata=b"source-md") pinning = FakePrimaryPinning() @@ -198,12 +314,22 @@ def transfer(self, handle): key = BlockKey(b"k0") control.incoming.append(_start_write_message(5, key)) kvcr = _new_kvcr(source_agent, pinning, control, name="source") + kvcr._core._clock = lambda: 0.0 # Failure must come from ERR, not a timeout. + if failure == "async": + assert _poll_until(kvcr, lambda _: bool(source_agent.xfers)) == [] + assert source_agent.sent_notifs == [] + source_agent.state = "ERR" + if failure != "initialize": + assert _poll_until(kvcr, lambda _: bool(source_agent.sent_notifs)) == [] + assert _decode_notif(source_agent.sent_notifs[0][1])["terminal"] is False + assert pinning.unpins == [] + assert source_agent.released_xfers == [] + source_agent.state = "DONE" assert _poll_until(kvcr, lambda _: pinning.unpins == ["pin"]) == [] assert source_agent.transfers == ([] if failure == "initialize" else [1]) - assert len(source_agent.sent_notifs) == 1 - agent_name, notif = source_agent.sent_notifs[0] + agent_name, notif = source_agent.sent_notifs[-1] assert agent_name == b"remote-1" assert _decode_notif(notif) == { "type": "write_done", @@ -214,40 +340,6 @@ def transfer(self, handle): assert not _has_outstanding_operations(kvcr) -def test_kvcr_source_async_transfer_error_notifies_failure(): - source_agent = FakeNixlAgent(metadata=b"source-md") - pinning = FakePrimaryPinning() - control = FakeBytesControl() - key = BlockKey(b"k0") - control.incoming.append(_start_write_message(7, key)) - kvcr = _new_kvcr(source_agent, pinning, control, name="source") - - assert _poll_until(kvcr, lambda _: bool(source_agent.xfers)) == [] - assert source_agent.sent_notifs == [] - - source_agent.state = "ERR" - assert ( - _poll_until( - kvcr, - lambda _: ( - bool(source_agent.sent_notifs) and not _has_outstanding_operations(kvcr) - ), - ) - == [] - ) - - assert len(source_agent.sent_notifs) == 1 - agent_name, notif = source_agent.sent_notifs[0] - assert agent_name == b"remote-1" - assert _decode_notif(notif) == { - "type": "write_done", - "op_handle": 7, - "success": False, - } - assert source_agent.released_xfers == [1] - assert pinning.unpins == ["pin"] - - def test_kvcr_source_ignores_malformed_control_messages(): """Malformed control payloads never crash the scheduler or cause effects.""" source_agent = FakeNixlAgent(metadata=b"source-md") @@ -267,70 +359,183 @@ def test_kvcr_source_ignores_malformed_control_messages(): @pytest.mark.parametrize( - "terminal_state", - [None, "ERR", "DONE"], - ids=["cancelled", "failed", "completed"], + ("terminal_state", "abandon", "raises"), + [ + ("DONE", False, False), + ("DONE", True, False), + ("DONE", True, True), + ("shutdown", True, False), + ], ) -def test_kvcr_source_timeout_holds_pins_until_safe_release( - terminal_state: str | None, -) -> None: - class DelayedReleaseAgent(FakeNixlAgent): - def __init__(self): - super().__init__(metadata=b"source-md") - self.release_attempts = 0 - self.allow_release = False +def test_kvcr_source_timeout_releases_pins_on_completion_or_abandonment( + terminal_state, + abandon, + raises, +): + now = 0.0 + agent, pinning, control = FakeNixlAgent(), FakePrimaryPinning(), FakeBytesControl() + errors = [] - def release_xfer_handle(self, handle): - self.release_attempts += 1 - if self.release_attempts == 1: - raise RuntimeError("busy") - if not self.allow_release: - return False - super().release_xfer_handle(handle) + def on_resilience_event(error): + if not errors: + assert pinning.unpins == [] + errors.append(error) + if raises: + raise RuntimeError("callback failed") - now = 0.0 - source_agent = DelayedReleaseAgent() - pinning = FakePrimaryPinning() - control = FakeBytesControl() - key = BlockKey(b"k0") kvcr = _new_kvcr( - source_agent, - pinning, - control, - KVCRConfig( - nixl_agent_name="source", - pool_layouts=[("", 16)], - operation_timeout_ms=1000, - enable_telemetry=True, - ), - name="source", + agent, pinning, control, name="source", on_resilience_event=on_resilience_event ) kvcr._core._clock = lambda: now - control.incoming.append(_start_write_message(12, key)) + key = BlockKey(b"k0") + try: + control.incoming.append(_start_write_message(12, key, target_agent="target")) + assert _poll_until(kvcr, lambda _: bool(agent.xfers)) == [] + source_handle = next(iter(kvcr._core._progress._in_flight_ops))[1] + now = 1.5 # A late first poll must not restart the abandonment deadline. + _wait_until(lambda: bool(agent.sent_notifs)) + assert _decode_notif(agent.sent_notifs[0][1]).get("terminal", True) is False + # Releasing a PROC handle is not evidence that the native write stopped. + assert agent.released_xfers == [] + assert pinning.unpins == [] + if abandon: + now = 5.0 + if raises: + with pytest.raises(RuntimeError, match="callback failed"): + _poll_until(kvcr, lambda _: bool(errors)) + assert pinning.unpins == [] + assert _poll_until(kvcr, lambda _: bool(pinning.unpins)) == [] + assert [error.state for error in errors] == ["uncertain"] + assert agent.released_xfers == [] + control.incoming.append(_write_probe_message(12)) + _wait_until(lambda: bool(control.sent)) + assert _decode_control_message(control.sent[-1][1])["terminal"] is False + if terminal_state == "shutdown": + kvcr._core._progress._submissions.put(_STOP) + _wait_until(lambda: kvcr._core._progress._startup_stage == "cleanup") + assert agent.released_xfers == [] + agent.state = "DONE" + if terminal_state == "shutdown": + kvcr.close() + assert not kvcr._core._framework_pin_keys + assert not kvcr._core._local_dram_sources_by_op + else: + if raises: + with pytest.raises(RuntimeError, match="callback failed"): + _poll_until(kvcr, lambda _: len(errors) == 2) + assert ( + _poll_until(kvcr, lambda _: not _has_outstanding_operations(kvcr)) == [] + ) + assert pinning.unpins == ["pin"] + assert agent.released_xfers == [1] + assert [error.state for error in errors] == ( + ["uncertain", "quiesced"] if abandon else [] + ) + for error in errors: + assert error.op_handle == source_handle + assert error.source_blocks == {key: [_mem_descriptor(addr=0)]} + assert error.destination_regions is None + finally: + agent.state = "DONE" + kvcr.close() - assert _poll_until(kvcr, lambda _: bool(source_agent.xfers)) == [] - now = 2.0 - _wait_until(lambda: source_agent.release_attempts > 0) - assert source_agent.released_xfers == [] - assert pinning.unpins == [] - assert _has_outstanding_operations(kvcr) - if terminal_state is not None: - source_agent.state = terminal_state - source_agent.allow_release = True - assert _poll_until(kvcr, lambda _: not _has_outstanding_operations(kvcr)) == [] - assert not kvcr._core._remote_fw_dram._source_pin_ops - assert pinning.unpins == ["pin"] - assert source_agent.released_xfers == [1] - if terminal_state != "DONE": - assert _decode_notif(source_agent.sent_notifs[0][1]) == { - "type": "write_done", - "op_handle": 12, - "success": False, - } - else: - assert source_agent.sent_notifs == [] - assert source_agent.telemetry_handles == [1] +def test_source_lifecycles_distinguish_targets_reusing_the_same_handle(): + now = 0.0 + agent, control, errors, done = FakeNixlAgent(), FakeBytesControl(), [], set() + agent.check_xfer_state = lambda handle: "DONE" if handle in done else "PROC" + source = _new_kvcr( + agent, + FakePrimaryPinning(), + control, + name="source", + on_resilience_event=errors.append, + ) + source._core._clock = lambda: now + key = BlockKey(b"shared") + try: + for target in ("target-a", "target-b"): + control.incoming.append(_start_write_message(12, key, target_agent=target)) + _poll_until(source, lambda _: len(agent.xfers) == 2) + now = 5.0 + _poll_until(source, lambda _: len(errors) == 2) + pending = {error.op_handle for error in errors} + assert len(pending) == 2 + assert [error.state for error in errors] == ["uncertain", "uncertain"] + assert ( + errors[0].source_blocks + == errors[1].source_blocks + == {key: [_mem_descriptor(addr=0)]} + ) + for native_handle in (1, 2): + done.add(native_handle) + _poll_until(source, lambda _: len(errors) == 2 + native_handle) + event = errors[-1] + assert event.state == "quiesced" + assert event.source_blocks == errors[0].source_blocks + pending.remove(event.op_handle) + assert len(pending) == 2 - native_handle + assert not _has_outstanding_operations(source) + finally: + done.update((1, 2)) + source.close() + + +def test_abandoned_source_keeps_local_slot_claimed_until_quiescence(): + now = 0.0 + memory = ctypes.create_string_buffer(16) + descriptor = _mem_descriptor(ctypes.addressof(memory)) + agent, control, errors = FakeNixlAgent(), FakeBytesControl(), [] + source = _new_kvcr( + agent, + FakePrimaryPinning(missing_indices=(0,)), + control, + name="source", + on_resilience_event=errors.append, + local_dram=LocalDramOptions([("", ctypes.addressof(memory), len(memory))]), + ) + source._core._clock = lambda: now + key, replacement = BlockKey(b"k0"), BlockKey(b"k1") + missing, framework_hit = BlockKey(b"missing"), BlockKey(b"framework-hit") + expected_sources = { + key: [replace(descriptor, end_point_name="source")], + framework_hit: [_mem_descriptor(addr=0)], + } + try: + agent.state = "DONE" + deposit = source.deposit({key: [descriptor]}) + assert _poll_until(source, bool) == [(deposit, _op_entries({key: True}))] + agent.state = "PROC" + payload = msgspec.msgpack.decode( + _start_write_message(12, key, target_agent="target") + ) + payload["keys"] = [key, missing, framework_hit] + payload["dst_descriptors"] = [ + [_mem_descriptor(128 + 16 * index).__dict__] for index in range(3) + ] + control.incoming.append(msgspec.msgpack.encode(payload)) + _poll_until(source, lambda _: len(agent.xfers) == 2) + now = 5.0 + assert _poll_until(source, lambda _: bool(errors)) == [] + assert [error.state for error in errors] == ["uncertain"] + assert errors[0].source_blocks == expected_sources + blocked = source.deposit({replacement: [descriptor]}) + assert list(source.poll_completed()) == [ + (blocked, _op_entries({replacement: False})) + ] + assert source._core._block_record_map[key].local_dram.claim_count == 1 + agent.state = "DONE" + assert _poll_until(source, lambda _: len(errors) == 2) == [] + assert [error.state for error in errors] == ["uncertain", "quiesced"] + assert source._core._block_record_map[key].local_dram.claim_count == 0 + assert errors[1].source_blocks == expected_sources + deposit = source.deposit({replacement: [descriptor]}) + assert _poll_until(source, bool) == [ + (deposit, _op_entries({replacement: True})) + ] + finally: + agent.state = "DONE" + source.close() @pytest.mark.parametrize("failure", [False, None, 1, RuntimeError("release failed")]) @@ -457,10 +662,12 @@ def poll_pin_results(self): assert any("framework pin result polling failed" in message for message in warnings) -def test_kvcr_source_poll_failure_is_terminal_and_logged(kvcr_caplog): +def test_kvcr_source_poll_failure_waits_for_quiescence_and_is_logged(kvcr_caplog): class RaisingAgent(FakeNixlAgent): def check_xfer_state(self, handle): - raise RuntimeError("boom") + if self.state != "DONE": + raise RuntimeError("boom") + return self.state source_agent = RaisingAgent(metadata=b"source-md") pinning = FakePrimaryPinning() @@ -469,18 +676,20 @@ def check_xfer_state(self, handle): kvcr = _new_kvcr(source_agent, pinning, control, name="source") control.incoming.append(_start_write_message(5, key)) - assert ( - _poll_until( - kvcr, - lambda _: ( - bool(source_agent.sent_notifs) and not _has_outstanding_operations(kvcr) - ), - ) - == [] + _poll_until( + kvcr, + lambda _: any( + "transfer progress failed" in rec.getMessage() + for rec in kvcr_caplog.records + ), ) + assert source_agent.released_xfers == [] + assert pinning.unpins == [] + source_agent.state = "DONE" + assert _poll_until(kvcr, lambda _: not _has_outstanding_operations(kvcr)) == [] assert source_agent.released_xfers == [1] assert pinning.unpins == ["pin"] - assert _decode_notif(source_agent.sent_notifs[0][1]) == { + assert _decode_notif(source_agent.sent_notifs[-1][1]) == { "type": "write_done", "op_handle": 5, "success": False, @@ -564,6 +773,8 @@ def test_pending_pin_waiters_share_partial_results_and_request_uncovered_keys( ) assert set(pinning.unpins) == expected_unpins assert pinning.searches == expected_searches + agent.state = "DONE" + _poll_until(source, lambda _: not _has_outstanding_operations(source)) def test_source_telemetry_precedes_release_and_is_not_duplicated() -> None: diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index cc7d3f4..0d90a5f 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -5,7 +5,8 @@ import ctypes import logging import time -from unittest.mock import Mock +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import DEFAULT, Mock import msgspec import pytest @@ -60,9 +61,10 @@ def test_submit_hint_filters_unlisted_hash(): def test_kvcr_opportunistic_query_accepts_key_outside_hint(): + agent = FakeNixlAgent() control = FakeBytesControl("tcp://target:1") target = _new_kvcr( - FakeNixlAgent(), + agent, FakePrimaryPinning(), control, key_adapter=_ConstantHashAdapter(), @@ -77,13 +79,15 @@ def test_kvcr_opportunistic_query_accepts_key_outside_hint(): assert target.query((requested_key,), "req") == [ (QueryStatus.FETCHABLE, CacheTier.REMOTE_G2) ] - target.deliver({requested_key: [_mem_descriptor()]}, request_id="req") + op_handle = target.deliver({requested_key: [_mem_descriptor()]}, request_id="req") assert list(target.poll_completed()) == [] _wait_until(lambda: bool(control.sent)) assert _decode_control_message(control.sent[0][1])["keys"] == [requested_key] target.discard_hint("req") assert target.query((requested_key,), "req") == [(QueryStatus.MISS, None)] + agent.notifs["source"] = [_write_done_notification(op_handle)] + _poll_until(target, bool) def test_remote_fetch_uses_local_then_framework_sources() -> None: @@ -399,11 +403,12 @@ def test_remote_completion_rejects_invalid_notification( @pytest.mark.parametrize( - "completion_before_timeout", [False, True], ids=["late", "queued"] + "resolution", ["queued", "late", "notification", "probe", "guard", "unknown"] ) def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal( - completion_before_timeout: bool, + resolution, ) -> None: + errors = [] now = 0.0 block_size = 16 local = ctypes.create_string_buffer(block_size) @@ -417,8 +422,10 @@ def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal( KVCRConfig( nixl_agent_name="target", pool_layouts=[("", 16)], - operation_timeout_ms=10, + operation_timeout_ms=1000, + abandon_timeout_ms=2000, ), + on_resilience_event=errors.append, key_adapter=_ConstantHashAdapter(), remote_options=RemoteFWDramOptions(eager_ctrl_connect=False), local_dram=LocalDramOptions([("", ctypes.addressof(local), len(local))]), @@ -429,31 +436,77 @@ def test_remote_fetch_timeout_keeps_slot_until_source_is_terminal( _wait_until(lambda: bool(control.sent)) message = _decode_control_message(control.sent[0][1]) - if completion_before_timeout: + progress = target._core._progress + _wait_until(lambda: ("target", message["op_handle"]) in progress._in_flight_ops) + operation = progress._in_flight_ops[("target", message["op_handle"])] + assert not operation.close(progress) + with pytest.raises(RuntimeError, match="unresolved operations"): + progress._close_nixl() + assert agent.deregistered == [] + if resolution != "unknown": + control.incoming.append( + msgspec.msgpack.encode( + { + "type": "target_metadata_ack", + "sender_control_endpoint": "tcp://source:1", + "sender_incarnation": "source", + "op_handle": message["op_handle"], + } + ) + ) + _wait_until(lambda: operation.source_incarnation == "source") + if resolution == "queued": # 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 + now = 1.5 # Past T, but still inside the destination's grace period. assert _poll_until(target, lambda completed: bool(completed)) == [ (fetch, _op_entries({key: False})) ] - if not completion_before_timeout: + if resolution != "queued": assert target.query((key,), "req") == [ (QueryStatus.FETCHABLE, CacheTier.REMOTE_G2) ] assert _has_outstanding_operations(target) + assert not operation.close(progress) + if resolution != "late": + now = 2.0 + assert _poll_until(target, lambda _: bool(errors)) == [] + assert [error.state for error in errors] == ["uncertain"] + control.incoming.append( + _probe_ack( + message["op_handle"], + sender_incarnation="new", + terminal=True, + dead_incarnation="source" if resolution == "unknown" else None, + ) + ) + _wait_until(lambda: not control.incoming) + now = 100.0 # Time and a replacement incarnation cannot free the slot. + assert list(target.poll_completed()) == [] 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"])] + if resolution in ("probe", "guard"): + fields = {"terminal": True} + if resolution == "guard": + fields.update(sender_incarnation="guard", dead_incarnation="source") + control.incoming.append(_probe_ack(message["op_handle"], **fields)) + else: + agent.notifs["source"] = [_write_done_notification(message["op_handle"])] assert ( _poll_until(target, lambda _: not _has_outstanding_operations(target)) == [] ) + if resolution not in ("queued", "late"): + assert [error.state for error in errors] == ["uncertain", "quiesced"] + assert all(error.op_handle == message["op_handle"] for error in errors) + assert errors[0].destination_regions == errors[1].destination_regions assert not _has_outstanding_operations(target) assert key not in target._core._block_record_map + assert operation.close(progress) # The terminal completion makes the single slot reusable. primary = ctypes.create_string_buffer(b"a" * block_size, block_size) @@ -522,6 +575,8 @@ def test_kvcr_deliver_propagates_source_pin_miss(): (QueryStatus.MISS, None), (QueryStatus.FETCHABLE, CacheTier.REMOTE_G2), ] + assert _poll_until(source, lambda _: not _has_outstanding_operations(source)) == [] + assert not source._core._remote_fw_dram._dangling_ops.source_writes target_control.sent = [] retry_handle = target.deliver( @@ -554,15 +609,30 @@ def test_kvcr_deliver_propagates_source_pin_miss(): assert target.query((key,), "req") == [(QueryStatus.FETCHABLE, CacheTier.REMOTE_G2)] +def _probe_ack(handle, source="tcp://source:1", **fields): + return msgspec.msgpack.encode( + { + "type": "write_probe_ack", + "sender_control_endpoint": source, + "sender_incarnation": "source", + "op_handle": handle, + "terminal": False, + **fields, + } + ) + + def _acked_deliver(control, kvcr, source, key): """Drive a deliver whose start_write carries no metadata, and return it.""" kvcr.submit_hint(_router_hint(source), request_id="metadata") _wait_until(lambda: len(control.sent) == 1) - control.incoming.append( - msgspec.msgpack.encode( - {"type": "target_metadata_ack", "sender_control_endpoint": source} - ) - ) + ack = { + "type": "target_metadata_ack", + "sender_control_endpoint": source, + "target_agent": "source", + "sender_incarnation": "source", + } + control.incoming.append(msgspec.msgpack.encode(ack)) _wait_until(lambda: not control.incoming) control.sent = [] @@ -672,10 +742,9 @@ def test_kvcr_metadata_ack_retry_lifecycle(): ] == ["target_metadata"] -@pytest.mark.parametrize("terminal_success", [False, True]) -def test_kvcr_deliver_timeout_waits_for_terminal_notification( - terminal_success: bool, -) -> None: +@pytest.mark.parametrize("source_responsive", [False, True]) +def test_kvcr_deliver_timeout_probes_source_before_finishing(source_responsive): + errors = [] now = 0.0 agent = FakeNixlAgent(metadata=b"target-md") control = FakeBytesControl() @@ -687,38 +756,221 @@ def test_kvcr_deliver_timeout_waits_for_terminal_notification( nixl_agent_name="target", pool_layouts=[("", 16)], operation_timeout_ms=1000, + abandon_timeout_ms=7000, ), + on_resilience_event=errors.append, ) kvcr._core._clock = lambda: now - key = BlockKey(b"k0") + key, source = BlockKey(b"k0"), "tcp://source:1" + handle, _ = _acked_deliver(control, kvcr, source, key) + control.sent.clear() + control.send = Mock(wraps=control.send, side_effect=[False, DEFAULT]) + now = 1.0 + _wait_until(lambda: control.send.call_count == 2) + control.send.side_effect = None + assert list(kvcr.poll_completed()) == [] + if source_responsive: + agent.notifs["source"] = [ + _write_done_notification(handle, success=False, terminal=False) + ] + control.incoming.append(_probe_ack(handle, source)) + _wait_until(lambda: not control.incoming and not agent.notifs) + assert list(kvcr.poll_completed()) == [] + now = 6.0 + time.sleep(0.01) + assert list(kvcr.poll_completed()) == [] + now = 7.0 + assert _poll_until(kvcr, bool) == [(handle, _op_entries({key: False}))] + assert control.send.call_count == 3 # Retry at T, then cleanup at abandonment. + assert [error.state for error in errors] == ["uncertain"] + assert errors[0].op_handle == handle - kvcr.submit_hint(_router_hint("tcp://source:1"), request_id="req") - op_handle = kvcr.deliver({key: [_mem_descriptor()]}, request_id="req") + # An abandoned operation must not blacklist the endpoint for fresh work. + kvcr.submit_hint(_router_hint(source), request_id="retry") + retry = kvcr.deliver({key: [_mem_descriptor()]}, request_id="retry") _wait_until( lambda: any( - _decode_control_message(message)["type"] == "start_write" - for _, message in control.sent + _decode_control_message(raw).get("op_handle") == retry + for _, raw in control.sent ) ) - _wait_until( - lambda: "tcp://source:1" in kvcr._core._remote_fw_dram._metadata_retry_after - ) + agent.notifs["source"] = [_write_done_notification(retry)] + assert _poll_until(kvcr, bool) == [(retry, _op_entries({key: True}))] - now = 2.0 - _wait_until( - lambda: "tcp://source:1" not in kvcr._core._remote_fw_dram._metadata_retry_after + control.incoming.append( + _probe_ack(handle, source, sender_incarnation="guard", terminal=True) ) + _wait_until(lambda: not control.incoming) + now = 100.0 assert list(kvcr.poll_completed()) == [] assert _has_outstanding_operations(kvcr) + assert len(errors) == 1 + control.incoming.append( + _probe_ack( + handle, source, sender_incarnation="guard", dead_incarnation="source" + ) + ) + assert _poll_until(kvcr, lambda _: len(errors) == 2) == [] + assert not _has_outstanding_operations(kvcr) + assert [error.state for error in errors] == ["uncertain", "quiesced"] + assert errors[0].destination_regions == errors[1].destination_regions - agent.notifs["source"] = [ - _write_done_notification(op_handle, success=terminal_success) - ] - assert _poll_until(kvcr, lambda completed: bool(completed)) == [ - (op_handle, _op_entries({key: terminal_success})) + +@pytest.mark.parametrize( + "outcome", + [ + "cancelled", + "cancelled_same_poll", + "failed_first", + "failed_last", + "terminal", + "log", + "raise", + "shutdown", + "shutdown_probe", + ], +) +def test_remote_write_cancellation_and_late_completion( + outcome, + caplog, +): + caplog.set_level(logging.INFO, logger="kvcr.core") + shutting_down = outcome.startswith("shutdown") + errors = [] + + def on_resilience_event(error): + errors.append(error) + raise error + + now = 0.0 + descriptor = _mem_descriptor() + agent, control = FakeNixlAgent(), FakeBytesControl() + kvcr = _new_kvcr( + agent, + FakePrimaryPinning(), + control, + KVCRConfig( + nixl_agent_name="target", pool_layouts=[("", 16)], operation_timeout_ms=1000 + ), + on_resilience_event=on_resilience_event if outcome == "raise" else None, + ) + kvcr._core._clock = lambda: now + key, source = BlockKey(b"k0"), "tcp://source:1" + handle, _ = _acked_deliver(control, kvcr, source, key) + if outcome in ("cancelled", "cancelled_same_poll", "failed_first", "failed_last"): + # The source can cancel before the target's own timeout. + now = 0.5 + cancelled = _write_done_notification( + handle, success=False, terminal=outcome in ("failed_first", "failed_last") + ) + notifications = [_write_done_notification(handle)] + if outcome == "cancelled": + agent.notifs["source"] = [cancelled] + _wait_until(lambda: not agent.notifs) + assert list(kvcr.poll_completed()) == [] + elif outcome == "failed_first": + notifications.insert(0, cancelled) + else: + notifications.append(cancelled) + agent.notifs["source"] = notifications + assert _poll_until(kvcr, bool) == [(handle, _op_entries({key: False}))] + return + now = 1.0 + _wait_until( + lambda: any( + _decode_control_message(raw).get("type") == "write_probe" + for _, raw in control.sent + ) + ) + if outcome == "terminal": + # The Guard's reply and a completed write may be observed in one poll. + # The destination is still held, so this is not a late write. + def guard_reply(): + control.recv = lambda: [] + agent.notifs["source"] = [ + _write_done_notification(handle, success=False), + _write_done_notification(handle), + _write_done_notification(handle, success=False), + ] + return [ + _probe_ack( + handle, + source, + sender_incarnation="guard", + dead_incarnation="source", + ), + ] + + control.recv = guard_reply + assert _poll_until(kvcr, bool) == [(handle, _op_entries({key: False}))] + return + now = 5.0 + if outcome == "raise": + with pytest.raises(RuntimeError): + _poll_until(kvcr, bool) + assert [error.state for error in errors] == ["uncertain"] + assert _poll_until(kvcr, bool) == [(handle, _op_entries({key: False}))] + kvcr.submit_hint(_router_hint(source), request_id="retry") + other_handle = kvcr.deliver({key: [_mem_descriptor()]}, request_id="retry") + _wait_until( + lambda: any( + _decode_control_message(raw).get("op_handle") == other_handle + for _, raw in control.sent + ) + ) + notifications = [ + _write_done_notification(handle, success=False, terminal=False), + _write_done_notification(handle), + _write_done_notification(handle, success=False, terminal=False), + _write_done_notification(other_handle), ] - assert not kvcr._core._remote_fw_dram._source_pin_ops - assert not kvcr._core._block_record_map + if shutting_down: + # Drain framework jobs, then resolve the old quarantine during close. + agent.notifs["source"] = [_write_done_notification(other_handle)] + assert _poll_until(kvcr, bool) == [(other_handle, _op_entries({key: True}))] + with ThreadPoolExecutor(max_workers=1) as executor: + closing = executor.submit(kvcr.close) + _wait_until(lambda: kvcr._core._progress._startup_stage == "cleanup") + if outcome == "shutdown_probe": + + def guard_reply(_endpoint, message): + if _decode_control_message(message)["type"] == "write_probe": + control.incoming.append( + _probe_ack( + handle, + source, + sender_incarnation="guard", + dead_incarnation="source", + ) + ) + return True + + control.send = guard_reply + now = 6.0 + else: + agent.notifs["source"] = notifications + closing.result(timeout=6) + assert kvcr._core.is_quiescent() + else: + agent.notifs["source"] = notifications + if outcome == "raise": + with pytest.raises(RuntimeError): + _poll_until(kvcr, bool) + # Even a raising handler must not kill progress or lose another completion. + if not shutting_down: + assert _poll_until(kvcr, bool) == [(other_handle, _op_entries({key: True}))] + if outcome == "log" or shutting_down: + errors = [ + record.args[0] + for record in caplog.records + if record.args and getattr(record.args[0], "op_handle", None) == handle + ] + assert [error.state for error in errors] == ["uncertain", "quiesced"] + for error in errors: + assert error.op_handle == handle + assert error.destination_regions == [descriptor] + assert error.source_blocks is None + kvcr.close() def test_kvcr_target_ignores_unknown_op_handle_notification(): @@ -826,6 +1078,11 @@ def test_kvcr_request_scoped_sources_do_not_overwrite(): "tcp://source-B:1", ] assert {("target", op_a), ("target", op_b)} <= _block_op_ids(kvcr) + agent.notifs["source"] = [ + _write_done_notification(op_a), + _write_done_notification(op_b), + ] + _poll_until(kvcr, lambda _: not _has_outstanding_operations(kvcr)) @pytest.mark.parametrize( diff --git a/tests/unit/test_kvcr_service.py b/tests/unit/test_kvcr_service.py index 6ac865c..57b7a33 100644 --- a/tests/unit/test_kvcr_service.py +++ b/tests/unit/test_kvcr_service.py @@ -76,6 +76,7 @@ class _FakeLiveness: def __init__(self) -> None: self._read, self._write = os.pipe() self.closed = False + self.incarnation = None def fileno(self) -> int: if self.closed: