[https://nvbugs/6621362][fix] Unblock Python disagg transfers stalled by a context-side timeout

tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py

@@ -428,6 +428,27 @@ def supports_inflight_request_cancellation(self) -> bool:
         """Return True when in-flight transfers can be cancelled safely."""
         return False
 
+    def context_transfer_is_waiting_for_peer(self, req: LlmRequest) -> bool:
+        """Whether a context send is still waiting for its peer to ask for the data.
+
+        ``respond_and_send_async`` only posts the send; nothing is written until
+        the generation peer requests it, which can be long after the context
+        finished when the peer's request queues behind a full generation batch.
+        Runtimes that can observe that boundary report it here so the transfer
+        timeout measures the transfer rather than the peer's queueing delay.
+        The default keeps the timeout unchanged for runtimes that cannot.
+        """
+        return False
+
+    def context_transfer_peer_ready_time(self,
+                                         req: LlmRequest) -> Optional[float]:
+        """``time.monotonic()`` at which every peer rank had asked for this send.
+
+        The transfer timeout is measured from this moment when it is later than
+        the send itself. None when unknown, which keeps the timeout unchanged.
+        """
+        return None
+
     def has_poisoned_transfer_buffer(self) -> bool:
         """Return True when a cancelled transfer may have corrupted a shared buffer."""
         return False

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -245,6 +245,12 @@ def fail(self, error: Exception) -> None:
     def cancel(self, by_peer: bool) -> None:
         self._end(_LogicalOutcome(SessionStatus.CANCELLED, by_peer=by_peer))
 
+    @property
+    def failed(self) -> bool:
+        """Whether the first terminal decision was a failure; a later cancel does not change it."""
+        with self._lock:
+            return self._terminal is not None and self._terminal.status is SessionStatus.ERROR
+
     def _end(self, outcome: _LogicalOutcome) -> None:
         with self._lock:
             if self._terminal is not None:
@@ -690,6 +696,10 @@ class Sender(SenderBase):
     # RecvReqInfo that never gets consumed.  Entries older than this
     # are evicted during periodic sweeps.
     _STALE_REQ_INFO_TTL_S = 120.0
+    # How long a cleared TxSession's id is remembered. A generation peer can ask for its KV
+    # long after the context side gave up on it (its request may wait in the generation queue
+    # for minutes), so this outlives the stale-request-info TTL.
+    _RETIRED_SESSION_TTL_S = 600.0
 
     def __init__(
         self,
@@ -715,6 +725,11 @@ def __init__(
         # Who asked is not recoverable from the id alone, and the two differ: a peer's cancel is a
         # transfer error, our own side's is an ordinary end.
         self._pre_cancelled_rids: dict[int, bool] = {}
+        # unique_rid -> clear time, for sessions already cleared (completed, failed, cancelled or
+        # timed out), oldest first. A REQUEST_DATA for one of them must be answered FAILED:
+        # parking it as an early generation-first request would leave the peer's receive
+        # unresolved for good. Protected by _sessions_lock.
+        self._retired_rids: dict[int, float] = {}
         self._shutdown = False
         self._shutdown_requested = False
         self._instance_rank = self._registrar.self_rank_info.instance_rank
@@ -777,11 +792,19 @@ def _remove_req_info(self, unique_rid: int):
     def sweep_stale_req_infos(self):
         """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL.
 
-        Called opportunistically from the listener thread when a new REQUEST_DATA
-        arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate
-        entries that are never consumed; this sweep prevents unbounded growth.
+        Called from the executor's context-status poll. With gen-first ADP
+        broadcast, non-assigned DP ranks accumulate entries that are never
+        consumed; this sweep prevents unbounded growth. Without attention DP an
+        evicted entry's peer is answered FAILED.
         """
         now = time.monotonic()
+        with self._sessions_lock:
+            # Oldest first (see clear_session), so only expired entries are visited.
+            while self._retired_rids:
+                rid, ts = next(iter(self._retired_rids.items()))
+                if now - ts <= self._RETIRED_SESSION_TTL_S:
+                    break
+                del self._retired_rids[rid]
         with self._peer_requests_lock:
             stale_rids = [
                 rid
@@ -790,12 +813,28 @@ def sweep_stale_req_infos(self):
             ]
         if not stale_rids:
             return
+        # Without attention DP no request is broadcast, so an entry this old is a peer still
+        # waiting on a session that is long gone or never came, and it must be told: dropping
+        # it silently leaves the peer's receive unresolved for good. Under attention DP the
+        # entries of non-assigned DP ranks are expected and must stay silent, since a FAILED
+        # from them would fail a receive the assigned ranks are still writing.
+        answer_peers = self._registrar.self_rank_info.dp_size == 1
+        swept: list[RecvReqInfo] = []
         for rid in stale_rids:
             with self._sessions_lock, self._peer_requests_lock:
                 if rid not in self._sessions and rid in self._peer_requests:
-                    self._peer_requests.pop(rid, None)
+                    infos = self._peer_requests.pop(rid, None) or {}
                     self._peer_requests_timestamps.pop(rid, None)
                     logger.debug(f"Swept stale RecvReqInfo for rid={rid}")
+                    if answer_peers:
+                        swept.extend(infos.values())
+        for info in swept:
+            logger.warning(
+                f"Answering FAILED to a REQUEST_DATA for rid={info.unique_rid} that no TxSession "
+                f"took within {self._STALE_REQ_INFO_TTL_S:.0f}s"
+            )
+            # Queue it on a worker: this runs on the executor thread, not the listener's.
+            self._send_failed_result_to_receiver(info, defer_to_worker=True)
 
     def setup_session(self, tx_session: "TxSession"):
         unique_rid = tx_session.disagg_request_id
@@ -804,6 +843,7 @@ def setup_session(self, tx_session: "TxSession"):
         with self._sessions_lock:
             if self._shutdown_requested:
                 raise RuntimeError("cannot create a TxSession after Sender shutdown")
+            self._retired_rids.pop(unique_rid, None)
             if unique_rid in self._pre_cancelled_rids:
                 pre_cancel = True
                 cancelled_by_peer = self._pre_cancelled_rids.pop(unique_rid)
@@ -844,6 +884,9 @@ def setup_session(self, tx_session: "TxSession"):
             expected_count = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks)
             if self._is_req_ready(unique_rid, expected_count):
                 with tx_session.lock:
+                    # The executor reads both without the lock, so publish the
+                    # time first: a ready session always carries its time.
+                    tx_session.peer_ready_time = time.monotonic()
                     tx_session.receiver_ready = True
         return
 
@@ -1645,8 +1688,35 @@ def _respond_with_kv(self, _send_id: bytes, message: list[bytes]):
         with self._sessions_lock:
             session = self._get_session(info.unique_rid)
             if session is None:
-                self._save_peer_req_info(info)
-                return
+                # A retried generation request carries a fresh unique_rid, so the
+                # context-side id names the retired session as well.
+                retired = (
+                    info.unique_rid in self._retired_rids
+                    or info.sender_req_id in self._retired_rids
+                )
+                # The session it names may also still be live. Serving it would
+                # need every reply keyed by the receiver's id and could write the
+                # same pieces to two attempts, so it is refused rather than parked.
+                retry_of_live_session = (
+                    info.sender_req_id != info.unique_rid
+                    and self._get_session(info.sender_req_id) is not None
+                )
+                if not (retired or retry_of_live_session):
+                    self._save_peer_req_info(info)
+                    return
+        if session is None:
+            # Nothing will ever be sent for this request; say so instead of leaving
+            # the peer's receive (and its KV pages) waiting forever. Typically the
+            # context side timed out waiting for it while the peer's request sat in
+            # its generation queue, or it is a retry of a request still in flight.
+            logger.warning(
+                f"_respond_with_kv: TxSession {info.sender_req_id} "
+                f"{'already retired' if retired else 'serves an earlier attempt'}; "
+                f"failing REQUEST_DATA for peer rid={info.unique_rid} "
+                f"rank={info.instance_rank}"
+            )
+            self._send_failed_result_to_receiver(info)
+            return
         with session.lock:
             with self._sessions_lock:
                 if self._get_session(info.unique_rid) is not session or session._closed:
@@ -1804,6 +1874,8 @@ def _save_peer_req_info(self, peer_transfer_req_info: RecvReqInfo):
         if self._is_req_ready(req_info.unique_rid, expected_transfers):
             session = self._get_session(req_info.unique_rid)
             if session is not None and not session.receiver_ready:
+                # Time before flag, as in setup_session.
+                session.peer_ready_time = time.monotonic()
                 session.receiver_ready = True
 
     def has_all_peer_req_infos(self, unique_rid: int) -> bool:
@@ -1821,6 +1893,9 @@ def clear_session(self, unique_rid: int):
         with self._sessions_lock:
             if unique_rid in self._sessions:
                 del self._sessions[unique_rid]
+                # Re-insert so the dict stays ordered by clear time.
+                self._retired_rids.pop(unique_rid, None)
+                self._retired_rids[unique_rid] = time.monotonic()
         self._remove_req_info(unique_rid)
 
     def send_cancel_to_receivers(self, unique_rid: int) -> None:
@@ -1914,6 +1989,8 @@ def __init__(
         self._aux_buffer = aux_buffer
         self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None
         self.receiver_ready: bool = False
+        # time.monotonic() when receiver_ready turned True, i.e. every peer rank had asked.
+        self.peer_ready_time: Optional[float] = None
         self.kv_tasks = []
         self.aux_task = None
         self._reported_aux_peer_ranks: set[int] = set()
@@ -3027,6 +3104,10 @@ def __init__(
         self._exception: Optional[Exception] = None
         self._closed = False
         self._terminal_status: Optional[SessionStatus] = None
+        # When this failed receive started waiting on its remaining writers, and whether that
+        # wait has run out (see failed_receive_awaits_writers).
+        self._writer_wait_start: Optional[float] = None
+        self._writer_wait_expired = False
         self.transfer_start_time = None
         self.transfer_end_time = None
         self.kv_cache_size_bytes: int = 0
@@ -3576,14 +3657,51 @@ def has_transferring_tasks(self) -> bool:
 
         cancel_request() must return False while this is True.
         """
-        return not self.resources_drained()
+        return not self.resources_drained() or self.failed_receive_awaits_writers()
+
+    # A failed receive waits at most this many transfer timeouts for its remaining writers.
+    _WRITER_WAIT_TIMEOUT_MULTIPLIER = 2
+
+    def failed_receive_awaits_writers(self) -> bool:
+        """Whether this failed receive must stay open for writers that have not stopped.
+
+        Without physical ownership the first writer's FAILED fails the receive, while the
+        other writers it asked may still be writing into its pages, so retiring it then could
+        hand those pages to another request mid-write. A live writer ends with a result or a
+        CANCEL_SESSION, at the latest when its own context-side transfer deadline expires,
+        which comes before ``_WRITER_WAIT_TIMEOUT_MULTIPLIER`` transfer timeouts. A writer
+        still silent then is presumed gone (e.g. its process died), and the receive is let go
+        rather than pinning its pages and slot on a worker that may never act on its own
+        receive timeout. The hold covers cancellation too: a client abort or a timeout that
+        cancels the receive after it failed must not free it early either (``cancel_request``
+        refuses while ``has_transferring_tasks``). A receive cancelled before it failed is
+        not held. Polled from the executor thread.
+        """
+        if self._enforce_physical_ownership or not self._logical_outcomes.failed:
+            return False
+        if self._writer_wait_expired:
+            return False
+        if not any(task.reports_outstanding for task in self._kv_tasks):
+            return False
+        if self._timeout_s is None or self._timeout_s <= 0:
+            return True
+        now = time.monotonic()
+        if self._writer_wait_start is None:
+            self._writer_wait_start = now
+        waited_s = now - self._writer_wait_start
+        if waited_s <= self._timeout_s * self._WRITER_WAIT_TIMEOUT_MULTIPLIER:
+            return True
+        self._writer_wait_expired = True
+        logger.warning(
+            f"RxSession {self.disagg_request_id}: a writer has not reported {waited_s:.0f}s "
+            "after the receive failed; retiring it without that writer's result"
+        )
+        return False
 
     def _failed_wait_result(self) -> Optional[WaitResult]:
-        return (
-            WaitResult.FAILED
-            if not self._enforce_physical_ownership or self.resources_drained()
-            else None
-        )
+        if not self._enforce_physical_ownership:
+            return None if self.failed_receive_awaits_writers() else WaitResult.FAILED
+        return WaitResult.FAILED if self.resources_drained() else None
 
     def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]:
         """Poll or block until transfer completes.

tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py

@@ -31,6 +31,13 @@
     from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests
 
 
+# How long, in multiples of kv_transfer_timeout_ms, a context send may wait for
+# its generation peer to ask for the data before its transfer clock starts
+# counting regardless. The peer's request can queue behind a full generation
+# batch for a while, but a peer that never asks must stay reclaimable.
+_CTX_PEER_WAIT_TIMEOUT_MULTIPLIER = 3
+
+
 def is_gen_only_no_context_benchmark() -> bool:
     """Whether the ``gen_only_no_context`` benchmark skips KV transfer."""
     return os.getenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") == "1"
@@ -268,13 +275,45 @@ def flag_if_timed_out(req: LlmRequest, kind: str) -> None:
 
         # Context requests start their clock on the last chunk, which is also
         # when they enter the transfer manager, so this covers the whole
-        # context side.
+        # context side. Nothing is sent before every peer rank has asked for
+        # the data, so a send still waiting for that is not flagged, and once
+        # the peers have asked, the transfer is measured from that moment
+        # rather than charged with their queueing delay.
+        now = time.monotonic()
         for req in self._transfers.requests_in_transfer().values():
+            if self._context_send_awaits_peer(req, now, timeout_ms):
+                continue
+            peer_ready_time = self._transceiver.context_transfer_peer_ready_time(req)
+            if (
+                peer_ready_time is not None
+                and req.py_kv_transfer_start_time is not None
+                and peer_ready_time > req.py_kv_transfer_start_time
+            ):
+                req.py_kv_transfer_start_time = peer_ready_time
             flag_if_timed_out(req, "context")
         for req in self._registry.active_requests():
             if req.is_disagg_generation_transmission_in_progress:
                 flag_if_timed_out(req, "generation")
 
+    def _context_send_awaits_peer(self, req: LlmRequest, now: float, timeout_ms: int) -> bool:
+        """Whether a context send is still exempt from its transfer timeout.
+
+        The send is posted when the context finishes, but nothing is sent
+        until every generation peer rank asks, and under load that request can
+        sit in the generation queue for longer than the transfer timeout.
+        Expiring then gives up on transfers that never got to run, and the peer
+        later finds no data. A peer that never asks must still be reclaimed,
+        since this deadline is the only thing that ends such a send, so the
+        wait is capped at ``_CTX_PEER_WAIT_TIMEOUT_MULTIPLIER`` times the
+        timeout, counted from the send.
+        """
+        if req.py_kv_transfer_start_time is None:
+            return False
+        waited_ms = (now - req.py_kv_transfer_start_time) * 1000
+        if waited_ms > timeout_ms * _CTX_PEER_WAIT_TIMEOUT_MULTIPLIER:
+            return False
+        return self._transceiver.context_transfer_is_waiting_for_peer(req)
+
     # -- scheduling ----------------------------------------------------------
 
     def admit(self, fitting_gen_init: List[LlmRequest]) -> Tuple[List[LlmRequest], bool]: