Prevent remote writes from hanging after source death

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,33 @@ 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. Framework pins are released as soon as their dependent work finishes or is abandoned under the policy below, minimizing interference with framework scheduling. Backing allocations and NIXL registrations remain valid until native transfers quiesce. 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.
+
+### 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.
+At timeout `T`, the destination probes the source, which blocks unsubmitted work
+and attempts cancellation. At its fixed `2T` deadline, the source releases content
+pins without waiting for the destination, but continues native-transfer cleanup.
+At `2T`, the destination abandons any unresolved write, releases its memory,
+and sends a final cleanup probe. Nonterminal replies do not extend this deadline.
+
+A per-operation tombstone retains destination descriptors, not memory pins, so
+new work can continue even if both the source and its Guard die. A terminal reply
+from the original source clears it; Guard confirmation of that source
+incarnation's death starts one additional `T` of retention. Without either, it
+remains for the destination's lifetime. Guard takeover fails old operations
+rather than replaying them; a reused agent name alone is not death proof.
+
+This is bounded, best-effort handling, not a transport fence: neither silence nor
+metadata removal stops an already-posted write. Unresolved source cancellation
+and observed late writes report a `TransferError` through `KVCRBindings.on_error`
+during `poll_completed()`, defaulting to an error log. Source reports identify
+keys and local buffers; destination reports identify regions, not their current
+keys. A custom handler may raise to the caller without stopping progress. Writes
+without a notification, or after tombstone expiry, cannot be diagnosed this way.
 
 ---
 
@@ -253,7 +279,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.
 
 ---
 

src/kvcr/api.py

@@ -38,6 +38,7 @@
     QueryStatus,
     ReleaseHandle,
     ReleaseResult,
+    TransferError,
 )
 
 if TYPE_CHECKING:
@@ -72,6 +73,9 @@ class KVCRBindings:
     stats_factory: Callable[[], TelemetryStats] | None = None
     policy: "KVCachePolicy | None" = None
 
+    # Called by poll_completed; None logs at ERROR. A handler may raise to its caller.
+    on_error: Callable[[TransferError], None] | None = None
+
 
 class KVCR:
     """Framework-facing KV Cache Runner."""

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
@@ -38,6 +39,7 @@
     RecoveryMirrorError,
     ReleaseHandle,
     ReleaseResult,
+    TransferError,
 )
 
 if TYPE_CHECKING:
@@ -171,6 +173,7 @@ def __init__(
         ] = {}
 
         self._completion_queue: list[OpResult] = []
+        self._transfer_errors: deque[TransferError] = deque()
         self._joined_completions: dict[
             OpHandle, tuple[set[BlockKey], dict[BlockKey, OpEntryResult]]
         ] = {}
@@ -201,8 +204,12 @@ def __init__(
         )
 
         # Import lazily to keep the concrete backend private to KVCR setup.
+        from .dangling_ops import _log_transfer_error
         from .remote_fw_dram import _RemoteFWDram
 
+        self._on_error_callback = (
+            _log_transfer_error if bindings.on_error is None else bindings.on_error
+        )
         self._local_dram = (
             _LocalDram(self, local_dram_config)
             if local_dram_config is not None
@@ -500,6 +507,11 @@ def poll_completed(self) -> Iterable[OpResult]:
             progress_items = self._local_dram.poll_main(progress_items)
         self._remote_fw_dram.poll_main(progress_items)
         self._flush_inventory()
+        # Apply the whole batch before invoking user code; a raising handler must
+        # not lose completions or stop the progress thread.
+        while self._transfer_errors:
+            error = self._transfer_errors.popleft()
+            self._on_error_callback(error)
         completed = self._completion_queue
         self._completion_queue = []
         return completed

src/kvcr/dangling_ops.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Bounded handling of dangling operations and late-completion diagnostics."""
+
+import heapq
+import time
+import uuid
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any, 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_transfer_error(error: TransferError) -> None:
+    logger.error("%s", error)
+
+
+@dataclass
+class _SourceWriteStatus:
+    submitted: bool = False
+    cancel_requested: bool = False
+    cancel_deadline: float | None = None
+    abandoned: bool = False
+
+
+@dataclass
+class _Tombstone:
+    operation: "_TargetPullOp"
+    expires_at: float | None = None
+
+
+class _DanglingOps:
+    """Progress-thread state for probes, stalled sources, and tombstones."""
+
+    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] = {}
+        # TODO: Add a retention limit if unresolved tombstones accumulate.
+        self.tombstones: dict[OpHandle, _Tombstone] = {}
+        self._expirations: list[tuple[float, OpHandle]] = []
+        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:
+            # Stay disabled so queued pre-stall requests cannot get fresh deadlines.
+            self._source_stalled = True
+            logger.error(
+                "KVCR source progress stalled for %.1f ms (timeout: %d ms); "
+                "new source writes disabled until restart",
+                gap_ms,
+                timeout_ms,
+            )
+        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)]
+        first_attempt = cancelling and status.cancel_deadline is None
+        if first_attempt:
+            status.cancel_deadline = (
+                op.deadline + self._backend._kvcr.config.operation_timeout_ms / 1000
+            )
+        result = progress.poll_transfer(
+            cast(int, op.transfer_id), cancellation_requested=cancelling
+        )
+        if result is None and cancelling:
+            if first_attempt:
+                self._backend._send_write_done(
+                    progress, op.remote_agent, op.op_handle, False, terminal=False
+                )
+            if not status.abandoned and self._backend._kvcr._clock() >= cast(
+                float, status.cancel_deadline
+            ):
+                status.abandoned = True
+                # At 2T release content pins, not the registration or native work.
+                # The CANCEL_PENDING snapshot is a release request, not terminal proof.
+                error = TransferError(
+                    "KVCR source cancellation timed out",
+                    op.op_handle,
+                    source_blocks={
+                        key: list(descriptors)
+                        for key, descriptors in zip(op.ordered_keys, op.src_descriptors)
+                    },
+                )
+                self._backend._progress_outbound.extend([replace(op), error])
+        return result
+
+    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 abandon(self, progress: "_KVCRProgress", op: "_TargetPullOp") -> None:
+        # Best effort, not a transport fence. The caller may reuse these addresses.
+        self.tombstones[op.op_id[1]] = _Tombstone(op)
+        self._backend._invalidate_control_peer(op.remote_ctrl_ep)
+        self.probe(progress, op)  # Cleanup only: never extends the operation deadline.
+
+    def notification(
+        self, progress: "_KVCRProgress", op_handle: OpHandle, payload: dict[str, Any]
+    ) -> None:
+        tombstone = self.tombstones.get(op_handle)
+        if tombstone is None:
+            return
+        if (
+            payload.get("success") is True
+            and ("target", op_handle) not in progress._in_flight_ops
+        ):
+            self._backend._progress_outbound.append(
+                TransferError(
+                    f"KVCR late remote write from {payload.get('source_agent')!r}",
+                    op_handle,
+                    destination_regions=[
+                        descriptor
+                        for descriptors in tombstone.operation.dst_descriptors
+                        for descriptor in descriptors
+                    ],
+                )
+            )
+        self.tombstones.pop(op_handle)
+
+    def expire(self) -> None:
+        now = self._backend._kvcr._clock()
+        while self._expirations and self._expirations[0][0] <= now:
+            _, handle = heapq.heappop(self._expirations)
+            self.tombstones.pop(handle, None)
+
+    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 status is not None and expected in (None, self.incarnation):
+            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
+        active = cast(
+            "_TargetPullOp | None", progress._in_flight_ops.get(("target", handle))
+        )
+        tombstone = self.tombstones.get(handle)
+        op = tombstone.operation if tombstone is not None else active
+        if op is None or op.remote_ctrl_ep != endpoint:
+            return
+        expected = op.source_incarnation
+        if expected and payload.get("dead_incarnation") == expected:
+            if tombstone is None:
+                tombstone = self.tombstones[handle] = _Tombstone(op)
+            if tombstone.expires_at is None:
+                # One metadata-only grace period after confirmed process death.
+                # Expiry ends late-write diagnostics; it does not establish a fence.
+                tombstone.expires_at = (
+                    self._backend._kvcr._clock()
+                    + self._backend._kvcr.config.operation_timeout_ms / 1000
+                )
+                heapq.heappush(self._expirations, (tombstone.expires_at, handle))
+            self.sources[endpoint] = incarnation
+            if active is not None:
+                self._backend._refused_writes[("target", handle)] = {"success": False}
+        elif expected == incarnation and terminal:
+            self.tombstones.pop(handle, None)
+            if active is not None:
+                self._backend._refused_writes[("target", handle)] = {"success": False}

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
@@ -622,6 +623,8 @@ 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:
+                self._dead_incarnations.add(lease.incarnation)
             self._promote_for(lease)
         except BaseException as error:  # noqa: BLE001 - service-fatal
             self._fail(error)
@@ -883,6 +886,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.