[None][refactor] Harden the KvCacheTransceiver contract and add a conformance fake

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -45,7 +45,11 @@
 from tensorrt_llm._torch.disaggregation.resource.page import CacheKind
 from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool
 from tensorrt_llm._torch.distributed.communicator import Distributed
-from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver
+from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
+    CtxTransferStatus,
+    GenTransferStatus,
+    KvCacheTransceiver,
+)
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
     MambaHybridCacheManager,
@@ -243,7 +247,7 @@ def summarize(
             f"waiting_for_peer_info={len(self._wait_reqs)}"
         )
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         if getattr(self, "_shutdown", False):
             return
         self._shutdown = True
@@ -676,7 +680,7 @@ def _finalize_send(self, req: LlmRequest, session: TxSessionBase):
         self._send_reqs[rid] = req
 
     @nvtx_range("KvCacheTransceiverV2.respond_and_send_async")
-    def respond_and_send_async(self, req: LlmRequest):
+    def respond_and_send_async(self, req: LlmRequest) -> None:
         self._ever_had_send_session = True
         req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now())
         session = self._get_or_create_send_session(req)
@@ -685,7 +689,7 @@ def respond_and_send_async(self, req: LlmRequest):
         self._finalize_send(req, session)
 
     @nvtx_range("KvCacheTransceiverV2.request_and_receive_sync")
-    def request_and_receive_sync(self, req: LlmRequest):
+    def request_and_receive_sync(self, req: LlmRequest) -> None:
         rid = get_unique_rid(req)
         if rid in self._recv_sessions:
             logger.warning(
@@ -721,7 +725,7 @@ def request_and_receive_sync(self, req: LlmRequest):
             self._recv_reqs.pop(rid, None)
 
     @nvtx_range("KvCacheTransceiverV2.request_and_receive_async")
-    def request_and_receive_async(self, req: LlmRequest):
+    def request_and_receive_async(self, req: LlmRequest) -> None:
         self._ever_had_recv_session = True
         req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now())
         rid = get_unique_rid(req)
@@ -740,15 +744,15 @@ def request_and_receive_async(self, req: LlmRequest):
 
     def check_context_transfer_status(
         self, at_least_request_num: Optional[int], mark_complete: bool = False
-    ):
+    ) -> CtxTransferStatus:
         # A worker that never sends KV has nothing to reconcile here, so skip the consensus. Safe
         # because the flag flips together on every rank and never resets, so they all skip in step;
         # gating on the live session dict instead would not be, since a cancel clears it per-rank.
         # Keep the original sweep (only when tp/pp sync is on) so nothing is leaked.
         if not self._ever_had_send_session:
             if self._ctx_need_tp_sync or self._ctx_need_pp_sync:
                 self._transfer_worker.sweep_stale_req_infos()
-            return [], []
+            return CtxTransferStatus([], [])
         block_all = at_least_request_num is None
         wait_num = at_least_request_num if not block_all else 0
         need_progress = wait_num > 0
@@ -810,11 +814,11 @@ def check_context_transfer_status(
         # DP ranks (entries that will never have a TxSession created for them).
         self._transfer_worker.sweep_stale_req_infos()
 
-        return completed, failed
+        return CtxTransferStatus(completed, failed)
 
-    def check_gen_transfer_status(self, at_least_request_num: Optional[int]):
+    def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenTransferStatus:
         if not self._ever_had_recv_session and not self._gen_need_sync:
-            return [], [], []
+            return GenTransferStatus([], [], [])
         block_all = at_least_request_num is None
         wait_num = at_least_request_num if not block_all else 0
         need_progress = wait_num > 0
@@ -895,7 +899,7 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]):
             )
         self._close_failed_sessions(self._recv_sessions, self._recv_reqs, failed)
 
-        return completed, failed, cancelled_reqs
+        return GenTransferStatus(completed, failed, cancelled_reqs)
 
     def _poll_gen_sessions_for_poll_interval(self, wait_num: int) -> None:
         self._poll_sessions_for_interval(
@@ -1016,7 +1020,7 @@ def get_disaggregated_params(self) -> Dict[str, Any]:
             else None,
         }
 
-    def prepare_context_requests(self, requests: List[LlmRequest]):
+    def prepare_context_requests(self, requests: List[LlmRequest]) -> None:
         # Place new generation-first context requests into wait state, then
         # use allgather consensus to promote ready requests to CONTEXT_INIT.
         for req in requests:

tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from os import environ, getenv
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, NamedTuple, Optional
 
 import tensorrt_llm
 from tensorrt_llm import logger
@@ -142,12 +142,13 @@ def mapping_to_world_config(mapping: Mapping) -> WorldConfig:
 
 
 def create_kv_cache_transceiver(
-        mapping: Mapping,
-        dist: Distributed,
-        kv_cache_manager: KVCacheManager,
-        attention_type: AttentionTypeCpp,
-        cache_transceiver_config: CacheTransceiverConfig,
-        mamba_cache_manager: Optional[BaseMambaCacheManager] = None):
+    mapping: Mapping,
+    dist: Distributed,
+    kv_cache_manager: KVCacheManager,
+    attention_type: AttentionTypeCpp,
+    cache_transceiver_config: CacheTransceiverConfig,
+    mamba_cache_manager: Optional[BaseMambaCacheManager] = None
+) -> Optional["KvCacheTransceiver"]:
     if cache_transceiver_config is None or cache_transceiver_config.backend is None:
         logger.info("cache_transceiver is disabled")
         return None
@@ -230,44 +231,191 @@ def create_kv_cache_transceiver(
                                   mamba_cache_manager)
 
 
+class CtxTransferStatus(NamedTuple):
+    """Typed result of ``KvCacheTransceiver.check_context_transfer_status``.
+
+    Unpacks positionally as ``(completed_request_ids, error_request_ids)``
+    for backward compatibility; prefer named field access in new code.
+    """
+    # Requests whose KV send settled successfully during this call.
+    completed_request_ids: List[int]
+    # Requests whose KV send failed during this call.
+    error_request_ids: List[int]
+
+
+class GenTransferStatus(NamedTuple):
+    """Typed result of ``KvCacheTransceiver.check_gen_transfer_status``.
+
+    Unpacks positionally as ``(completed_request_ids, error_request_ids,
+    cancelled_requests)`` for backward compatibility; prefer named field
+    access in new code.
+
+    The C++ transceiver runtime reports receive outcomes exclusively through
+    request-state mutation, so its lists are always empty; only the Python
+    (V2) runtime populates them.
+    """
+    # Requests whose KV receive settled successfully during this call.
+    completed_request_ids: List[int]
+    # Requests whose KV receive failed during this call.
+    error_request_ids: List[int]
+    # Requests whose receive session was cancelled (locally via
+    # ``cancel_request`` or by a remote CANCEL message). Their sessions are
+    # closed; the caller decides the final request state, distinguishing
+    # user cancellation from remote-initiated cancellation.
+    cancelled_requests: List[LlmRequest]
+
+
 class KvCacheTransceiver(ABC):
+    """Contract for moving KV cache between disaggregated instances.
+
+    Implementations: ``BindKvCacheTransceiver`` (C++ runtime) and
+    ``KvCacheTransceiverV2`` (Python/NIXL runtime). The executor must not
+    depend on which one it holds beyond this interface.
+
+    Rank symmetry: ``check_context_transfer_status`` and
+    ``check_gen_transfer_status`` participate in intra-instance consensus
+    collectives so that all ranks agree on per-request outcomes. Every rank
+    of an instance must call them the same number of times with the same
+    arguments per iteration; divergence can deadlock the instance.
+
+    Request ids in results are the ids the transceiver tracks: the request's
+    disaggregated unique id when its disaggregated params carry one,
+    otherwise ``request_id``.
+    """
+
+    # KV-transfer timeout budget in milliseconds, taken from
+    # CacheTransceiverConfig. None means no timeout is enforced: transfers
+    # may remain in flight indefinitely and the executor skips its
+    # timeout/cancellation sweeps. Implementations must set this attribute.
+    kv_transfer_timeout_ms: Optional[int]
 
     @abstractmethod
-    def respond_and_send_async(self, req: LlmRequest):
+    def respond_and_send_async(self, req: LlmRequest) -> None:
+        """Start sending ``req``'s KV cache to the requesting instance.
+
+        Non-blocking. Postcondition: ``req.state`` is
+        ``DISAGG_CONTEXT_TRANS_IN_PROGRESS``. Completion, failure, or
+        cancellation is reported by later
+        ``check_context_transfer_status`` calls.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def request_and_receive_sync(self, req: LlmRequest):
+    def request_and_receive_sync(self, req: LlmRequest) -> None:
+        """Receive ``req``'s KV cache, blocking until the transfer settles.
+
+        Postcondition: ``req.state`` is
+        ``DISAGG_GENERATION_TRANS_COMPLETE`` on success or
+        ``DISAGG_TRANS_ERROR`` on failure.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def request_and_receive_async(self, req: LlmRequest):
+    def request_and_receive_async(self, req: LlmRequest) -> None:
+        """Start receiving ``req``'s KV cache without blocking.
+
+        Postcondition: ``req.state`` is
+        ``DISAGG_GENERATION_TRANS_IN_PROGRESS``. Completion, failure, or
+        cancellation is reported by later ``check_gen_transfer_status``
+        calls.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def check_context_transfer_status(self, at_least_request_num: int):
+    def check_context_transfer_status(
+            self,
+            at_least_request_num: Optional[int],
+            mark_complete: bool = False) -> CtxTransferStatus:
+        """Poll send-side transfers and reap the ones that settled.
+
+        Args:
+            at_least_request_num: None enters the implementation's
+                blocking mode, whose semantics are runtime-specific and
+                NOT portable:
+
+                * C++ runtime: blocks until every in-flight send
+                  completes — potentially unboundedly, since the
+                  transfer timeout is observe-only on this path — and
+                  rejects None outright while in-flight cancellation is
+                  enabled (a finite poll is required).
+                * Python (V2) runtime: a blocking wait bounded by the
+                  session transfer timeout; transfers may still be
+                  pending on return, so draining requires re-polling.
+
+                Runtime-independent callers (e.g. drain loops) must use
+                a finite int in an explicit loop instead: 0 is a
+                non-blocking sweep; N > 0 additionally waits for sends
+                to settle, bounded by
+                ``kv_transfer_sender_future_timeout_ms`` — the C++
+                runtime waits it once per still-pending selected
+                transfer (so a single call may wait it several times),
+                the Python (V2) runtime once per call — and may still
+                report fewer than N when the bound expires.
+            mark_complete: When True, transition requests whose send
+                completed to ``DISAGG_CONTEXT_COMPLETE`` before returning;
+                when False, that transition is the caller's responsibility.
+
+        Cancelled sends are not reported as a distinct category:
+        implementations close them internally and the caller tracks
+        cancellation through its own bookkeeping.
+
+        Participates in rank-consensus collectives; see the class docstring
+        for the symmetry requirement.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def check_gen_transfer_status(self, at_least_request_num: int):
+    def check_gen_transfer_status(
+            self, at_least_request_num: Optional[int]) -> GenTransferStatus:
+        """Poll receive-side transfers and reap the ones that settled.
+
+        ``at_least_request_num`` follows the same runtime-specific
+        blocking / bounded-polling semantics as
+        ``check_context_transfer_status``, except that on this receive
+        side both runtimes bound N > 0 local readiness polling by
+        ``kv_transfer_poll_interval_ms``, applied as a single deadline
+        per call (unlike the C++ context path); the rank-consensus
+        collectives that follow are outside that local wait bound.
+        Postconditions: requests whose
+        receive completed are transitioned to
+        ``DISAGG_GENERATION_TRANS_COMPLETE`` and failed ones to
+        ``DISAGG_TRANS_ERROR``; cancelled sessions are closed and their
+        requests returned in ``cancelled_requests`` with the state left for
+        the caller to decide. ``cancelled_requests`` is populated only by
+        the Python (V2) runtime — the C++ runtime reports every outcome,
+        cancellation included, through request-state mutation (see
+        ``GenTransferStatus``).
+
+        Participates in rank-consensus collectives; see the class docstring
+        for the symmetry requirement.
+        """
         raise NotImplementedError
 
     @abstractmethod
-    def check_gen_transfer_complete(self):
+    def check_gen_transfer_complete(self) -> bool:
+        """Return True when no receive-side transfer remains in flight."""
         raise NotImplementedError
 
     @abstractmethod
-    def cancel_request(self, req: LlmRequest):
+    def cancel_request(self, req: LlmRequest) -> bool:
+        """Best-effort cancellation of ``req``'s in-flight transfers.
+
+        Returns True when the transfers are cancelled and it is safe to
+        release the request's KV resources; False when a task is mid-write
+        and the caller must retry on a later iteration.
+        """
         raise NotImplementedError
 
     def supports_inflight_request_cancellation(self) -> bool:
+        """Return True when in-flight transfers can be cancelled safely."""
         return False
 
     def has_poisoned_transfer_buffer(self) -> bool:
+        """Return True when a cancelled transfer may have corrupted a shared buffer."""
         return False
 
     @abstractmethod
-    def prepare_context_requests(self, requests: List[LlmRequest]):
+    def prepare_context_requests(self, requests: List[LlmRequest]) -> None:
         """
         Prepare the context request for the cache transceiver in generation-first mode.
         This method should set the context request state to DISAGG_CONTEXT_WAIT_SCHEDULER
@@ -295,7 +443,7 @@ def get_status_dump(self) -> str:
         """Return a human-readable dump of transceiver state for debugging hangs."""
         return ""
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         """Shut down the transceiver and release registered resources."""
 
 
@@ -374,25 +522,34 @@ def __init__(self,
             cache_transceiver_config._to_pybind(), rnn_layer_num_per_pp_rank,
             indexer_layer_num_per_pp_rank)
 
-    def respond_and_send_async(self, req: LlmRequest):
+    def respond_and_send_async(self, req: LlmRequest) -> None:
         return self.impl.respond_and_send_async(req)
 
-    def request_and_receive_sync(self, req: LlmRequest):
+    def request_and_receive_sync(self, req: LlmRequest) -> None:
         return self.impl.request_and_receive_sync(req)
 
-    def request_and_receive_async(self, req: LlmRequest):
+    def request_and_receive_async(self, req: LlmRequest) -> None:
         return self.impl.request_and_receive_async(req)
 
-    def check_context_transfer_status(self, at_least_request_num: int):
-        return self.impl.check_context_transfer_status(at_least_request_num)
-
-    def check_gen_transfer_status(self, at_least_request_num: int):
-        return self.impl.check_gen_transfer_status(at_least_request_num)
-
-    def check_gen_transfer_complete(self):
+    def check_context_transfer_status(
+            self,
+            at_least_request_num: Optional[int],
+            mark_complete: bool = False) -> CtxTransferStatus:
+        completed_ids, error_ids = self.impl.check_context_transfer_status(
+            at_least_request_num, mark_complete)
+        return CtxTransferStatus(completed_ids, error_ids)
+
+    def check_gen_transfer_status(
+            self, at_least_request_num: Optional[int]) -> GenTransferStatus:
+        # The C++ runtime reports outcomes via request-state mutation only,
+        # so the returned lists are empty by design.
+        self.impl.check_gen_transfer_status(at_least_request_num)
+        return GenTransferStatus([], [], [])
+
+    def check_gen_transfer_complete(self) -> bool:
         return self.impl.check_gen_transfer_complete()
 
-    def cancel_request(self, req: LlmRequest):
+    def cancel_request(self, req: LlmRequest) -> bool:
         return self.impl.cancel_request(req)
 
     def supports_inflight_request_cancellation(self) -> bool:
@@ -406,7 +563,7 @@ def has_poisoned_transfer_buffer(self) -> bool:
     def get_status_dump(self) -> str:
         return self.impl.get_status_dump()
 
-    def prepare_context_requests(self, requests: List[LlmRequest]):
+    def prepare_context_requests(self, requests: List[LlmRequest]) -> None:
         # not implemented, an empty placeholder to allow being invoked unconditionally
         ...