[TRTLLM-12721][fix] Bound disagg transfer status polling

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -53,6 +53,7 @@
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include "tensorrt_llm/runtime/utils/pgUtils.h"
 #include <algorithm>
+#include <chrono>
 #include <cstddef>
 #include <numeric>
 #include <unordered_map>
@@ -66,6 +67,20 @@ std::mutex CacheTransceiver::mDllMutex;
 namespace
 {
 
+constexpr int kTransferFuturePollIntervalMs = 10;
+
+// Finite status checks are scheduler polls, not terminal deadlines. Keep each
+// slice short so the caller can yield back to scheduling and transfer progress.
+std::chrono::milliseconds getTransferFuturePollInterval(std::optional<int> const& configuredTimeoutMs)
+{
+    auto waitMs = kTransferFuturePollIntervalMs;
+    if (configuredTimeoutMs.has_value())
+    {
+        waitMs = std::max(1, std::min(configuredTimeoutMs.value(), kTransferFuturePollIntervalMs));
+    }
+    return std::chrono::milliseconds(waitMs);
+}
+
 using RequestIdType = LlmRequest::RequestIdType;
 
 enum class TransferConsensusState : std::uint64_t
@@ -700,13 +715,13 @@ void updateKVCacheTransferBW(std::shared_ptr<CacheTransceiverComm> const& mComm,
 RequestStatuses CacheTransceiver::checkContextTransferStatus(
     std::optional<int> const& atLeastRequestNum, bool markComplete)
 {
-    bool blockAll = !atLeastRequestNum.has_value();
+    bool const blockAll = !atLeastRequestNum.has_value();
     std::optional<int> senderFutureTimeoutMs = std::nullopt;
-    // If blockAll is true, we want to block and not use a timeout
-    if (!blockAll && mCacheTransceiverConfig.has_value())
+    if (mCacheTransceiverConfig.has_value())
     {
         senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs();
     }
+    auto const futurePollInterval = getTransferFuturePollInterval(senderFutureTimeoutMs);
     // Observe-only: WARN per-request when the wall-clock transfer time exceeds
     // kvTransferTimeoutMs. No cancellation, eviction, or state transition.
     std::optional<int> kvTransferTimeoutMs = std::nullopt;
@@ -770,27 +785,26 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus(
     for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();)
     {
         auto& [request, future] = *it;
+        auto const requestId = request->mRequestId;
         if (kvTransferTimeoutMs.has_value())
         {
             auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                 LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart());
             auto elapsedMs = static_cast<long>(elapsed.count());
-            if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(request->mRequestId).second)
+            if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second)
             {
                 TLLM_LOG_WARNING(
                     "Context KV cache transfer for request %ld exceeded configured timeout: "
                     "elapsed %ld ms > limit %d ms (observe-only).",
-                    request->mRequestId, elapsedMs, kvTransferTimeoutMs.value());
+                    requestId, elapsedMs, kvTransferTimeoutMs.value());
             }
         }
-        if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()))
+        if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end()))
         {
-            auto const requestId = request->mRequestId;
             try
             {
-                // Wait for up to a specified timeout
-                auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0)));
-                if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value())
+                auto const status = blockAll ? std::future_status::ready : future.wait_for(futurePollInterval);
+                if (status == std::future_status::ready)
                 {
                     future.get();
                     bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR;
@@ -800,8 +814,10 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus(
                 }
                 else if (status == std::future_status::timeout)
                 {
-                    TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.",
-                        senderFutureTimeoutMs.value());
+                    TLLM_LOG_DEBUG(
+                        "Context KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping it "
+                        "in progress.",
+                        requestId, static_cast<long>(futurePollInterval.count()));
                     ++it;
                 }
                 else
@@ -866,7 +882,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus(
 
 void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastRequestNum)
 {
-    bool blockAll = !atLeastRequestNum.has_value();
+    bool const blockAll = !atLeastRequestNum.has_value();
     std::vector<LlmRequest::RequestIdType> genTransferReadyRequestIds;
     for (auto&& [request, future] : mRequesterFutures)
     {
@@ -977,6 +993,13 @@ void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastR
             " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(),
             atLeastRequestNum.value_or(0));
     }
+    std::optional<int> requesterFutureTimeoutMs = std::nullopt;
+    if (mCacheTransceiverConfig.has_value())
+    {
+        requesterFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs();
+    }
+    auto const futurePollInterval = getTransferFuturePollInterval(requesterFutureTimeoutMs);
+
     // Observe-only: gen-side mirror of the context-side timeout WARN.
     std::optional<int> kvTransferTimeoutMs = std::nullopt;
     if (mCacheTransceiverConfig.has_value())
@@ -992,28 +1015,48 @@ void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastR
             auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                 LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart());
             auto elapsedMs = static_cast<long>(elapsed.count());
-            if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(request->mRequestId).second)
+            if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(requestId).second)
             {
                 TLLM_LOG_WARNING(
                     "Generation KV cache transfer for request %ld exceeded configured timeout: "
                     "elapsed %ld ms > limit %d ms (observe-only).",
-                    request->mRequestId, elapsedMs, kvTransferTimeoutMs.value());
+                    requestId, elapsedMs, kvTransferTimeoutMs.value());
             }
         }
         if (blockAll || toCompleteIdSet.find(requestId) != toCompleteIdSet.end())
         {
             try
             {
-                it->second.get();
-                bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR;
-                if (failed)
+                auto const status = blockAll ? std::future_status::ready : it->second.wait_for(futurePollInterval);
+                if (status == std::future_status::ready)
                 {
-                    // The receiver uses the error state as a local transfer-failed signal.
-                    // Keep that signal local until the consensus outcome commits it globally.
-                    request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS);
+                    it->second.get();
+                    bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR;
+                    if (failed)
+                    {
+                        // The receiver uses the error state as a local transfer-failed signal.
+                        // Keep that signal local until the consensus outcome commits it globally.
+                        request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS);
+                    }
+                    recordLocalTransferOutcome(requestId, request, failed, mCompletedRequesterRequestIds,
+                        mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus);
+                }
+                else if (status == std::future_status::timeout)
+                {
+                    TLLM_LOG_DEBUG(
+                        "Generation KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping "
+                        "it in progress.",
+                        requestId, static_cast<long>(futurePollInterval.count()));
+                    ++it;
+                    continue;
+                }
+                else
+                {
+                    TLLM_LOG_ERROR(
+                        "Future returned unexpected status for request %ld. Recording as failed.", requestId);
+                    recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds,
+                        mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus);
                 }
-                recordLocalTransferOutcome(requestId, request, failed, mCompletedRequesterRequestIds,
-                    mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus);
             }
             catch (std::exception const& e)
             {

tests/unittest/others/test_kv_cache_transceiver.py

@@ -26,6 +26,10 @@
 AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType
 LlmRequestType = tensorrt_llm.bindings.internal.batch_manager.LlmRequestType
 DataType = tensorrt_llm.bindings.DataType
+DEFAULT_KV_TRANSFER_TIMEOUT_S = (
+    CacheTransceiverConfig.model_fields["kv_transfer_timeout_ms"].default /
+    1000.0)
+KV_TRANSFER_COMPLETION_MARGIN_S = 10.0
 
 
 def create_kv_cache_manager(mapping, dtype):
@@ -54,6 +58,20 @@ def fill_kv_cache_buffer(kv_cache_manager):
         buffer.copy_(random_values)
 
 
+def wait_for_transfer_completion(poll_fn,
+                                 is_done_fn,
+                                 timeout_s=DEFAULT_KV_TRANSFER_TIMEOUT_S +
+                                 KV_TRANSFER_COMPLETION_MARGIN_S):
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        poll_fn()
+        if is_done_fn():
+            return
+        time.sleep(0.01)
+    poll_fn()
+    assert is_done_fn(), "Timed out waiting for KV cache transfer completion"
+
+
 @pytest.fixture(scope="function")
 def ctx_gen_kv_cache_dtype(request):
     if request.param == "ctx_fp8_gen_fp8":
@@ -155,8 +173,21 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype,
     # send gen request
     kv_cache_transceiver_gen.request_and_receive_async(gen_request)
 
-    kv_cache_transceiver_ctx.check_context_transfer_status(1)
-    kv_cache_transceiver_gen.check_gen_transfer_status(1)
+    completed_ctx_ids = set()
+
+    def poll_transfers():
+        completed, failed = kv_cache_transceiver_ctx.check_context_transfer_status(
+            1)
+        assert failed == []
+        completed_ctx_ids.update(completed)
+        kv_cache_transceiver_gen.check_gen_transfer_status(1)
+
+    def transfers_done():
+        return (ctx_request.py_request_id in completed_ctx_ids
+                and gen_request.state
+                == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE)
+
+    wait_for_transfer_completion(poll_transfers, transfers_done)
 
     assert torch.equal(
         kv_cache_manager_gen.get_buffers(0),
@@ -323,8 +354,23 @@ def test_async_transfer_keeps_llm_request_alive():
         "the requester future is in flight")
 
     # Drive the transfer to completion so the harness tears down cleanly.
-    transceiver_ctx.check_context_transfer_status(1)
-    transceiver_gen.check_gen_transfer_status(1)
+    completed_ctx_ids = set()
+
+    def poll_transfers():
+        completed, failed = transceiver_ctx.check_context_transfer_status(1)
+        assert failed == []
+        completed_ctx_ids.update(completed)
+        transceiver_gen.check_gen_transfer_status(1)
+
+    def transfers_done():
+        ctx_request = ctx_ref()
+        gen_request = gen_ref()
+        return (ctx_request is not None and gen_request is not None
+                and ctx_request.py_request_id in completed_ctx_ids
+                and gen_request.state
+                == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE)
+
+    wait_for_transfer_completion(poll_transfers, transfers_done)
 
 
 def _build_ctx_request_for_timeout_test(request_id: int) -> LlmRequest:
@@ -428,6 +474,76 @@ def test_kv_transfer_timeout_silent_when_unset(capfd):
     transceiver_ctx.cancel_request(ctx_request)
 
 
+@pytest.mark.timeout(60)
+def test_context_transfer_bounded_poll_keeps_request_in_progress():
+    """A bounded transfer poll must not make a non-ready request terminal."""
+    mapping = Mapping(world_size=1, rank=0)
+    dist = Distributed.get(mapping)
+    kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF)
+    kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF)
+
+    cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT",
+                                                      max_tokens_in_buffer=512)
+    transceiver_ctx = create_kv_cache_transceiver(mapping, dist,
+                                                  kv_cache_manager_ctx,
+                                                  AttentionTypeCpp.DEFAULT,
+                                                  cache_transceiver_config)
+    transceiver_gen = create_kv_cache_transceiver(mapping, dist,
+                                                  kv_cache_manager_gen,
+                                                  AttentionTypeCpp.DEFAULT,
+                                                  cache_transceiver_config)
+
+    fill_kv_cache_buffer(kv_cache_manager_ctx)
+
+    ctx_request = _build_ctx_request_for_timeout_test(request_id=100)
+    kv_cache_manager_ctx.impl.add_sequence_batch(
+        [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request])
+
+    transceiver_ctx.respond_and_send_async(ctx_request)
+
+    start = time.monotonic()
+    completed_request_ids, error_request_ids = (
+        transceiver_ctx.check_context_transfer_status(1))
+    elapsed = time.monotonic() - start
+
+    assert elapsed < 1.0, (
+        f"Bounded poll should yield back to the executor quickly; "
+        f"elapsed={elapsed:.3f}s")
+    assert completed_request_ids == []
+    assert error_request_ids == []
+    assert ctx_request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS
+
+    sampling_params = SamplingParams()
+    gen_request = LlmRequest(
+        request_id=100,
+        max_new_tokens=1,
+        input_tokens=list(range(256)),
+        sampling_config=tensorrt_llm.bindings.SamplingConfig(
+            sampling_params._get_sampling_config()),
+        is_streaming=False,
+        llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY,
+        context_phase_params=ctx_request.context_phase_params)
+
+    kv_cache_manager_gen.impl.add_sequence_batch(
+        [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request])
+    transceiver_gen.request_and_receive_async(gen_request)
+
+    completed_ctx_ids = set()
+
+    def poll_transfers():
+        completed, failed = transceiver_ctx.check_context_transfer_status(1)
+        assert failed == []
+        completed_ctx_ids.update(completed)
+        transceiver_gen.check_gen_transfer_status(1)
+
+    def transfers_done():
+        return (ctx_request.py_request_id in completed_ctx_ids
+                and gen_request.state
+                == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE)
+
+    wait_for_transfer_completion(poll_transfers, transfers_done)
+
+
 def create_hybrid_cache_manager(mapping,
                                 dtype,
                                 mamba_conv_dtype=torch.float16,
@@ -629,8 +745,21 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes,
 
     cache_transceiver_gen.request_and_receive_async(gen_request)
 
-    cache_transceiver_ctx.check_context_transfer_status(1)
-    cache_transceiver_gen.check_gen_transfer_status(1)
+    completed_ctx_ids = set()
+
+    def poll_transfers():
+        completed, failed = cache_transceiver_ctx.check_context_transfer_status(
+            1)
+        assert failed == []
+        completed_ctx_ids.update(completed)
+        cache_transceiver_gen.check_gen_transfer_status(1)
+
+    def transfers_done():
+        return (ctx_request.py_request_id in completed_ctx_ids
+                and gen_request.state
+                == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE)
+
+    wait_for_transfer_completion(poll_transfers, transfers_done)
 
     assert torch.equal(
         hybrid_cache_manager_gen.get_buffers(0),