From 7f4e30af890da83c62ec4289a1c9140732073555 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 20:55:43 -0700 Subject: [PATCH 1/5] [https://nvbugs/6104831][fix] Disagg KV transfer: non-blocking polling, structured cancel, bounded quarantine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C++ cache transceiver and Python executor cooperate on disaggregated KV transfer cancellation, timeout, and process health. Earlier attempts (PR #13706, PR #13713, PR #13728) each fixed one slice of the problem but still hung in the field: * PR #13706 with Python changes restarted workers on every routine per-request timeout (Python fail-close was too aggressive). * PR #13713 / PR #13728 hung even with Python cleanup removed because the C++ status polling could still call future.get() on an unready worker future, freezing the executor event loop while NIXL/UCX was wedged. This PR keeps each layer's responsibility separate, as described in docs/source/features/disagg-kv-transfer-hang-restart-analysis.md and the new docs/source/features/disagg-kv-transfer-session-lifecycle.md: * checkContextTransferStatus / checkGenTransferStatus are now strictly non-blocking. They poll with wait_for(0ms) and only call future.get() when the future is already ready. atLeastRequestNum > 0 still admits additional ready entries but never selects an unready one to satisfy the count. drainContextTransferStatus / drainGenTransferStatus are the only blocking variants and are intended for shutdown drain. * cancelRequestStructured returns a TransferCancelResult enum with six outcomes (NotFound, AlreadyComplete, CancelledBeforeAdvertise, CancelRequestedInFlight, BackendUnhealthy, NotCancellable). Only the first three permit Python to free request resources; the others keep the request owned by C++ until the worker future reaches a final state. The historical bool cancelRequest is preserved as a backward-compatible wrapper. * The transceiver maintains a bounded quarantine counter and a global progress deadline. A per-request timeout marks the entry quarantined but keeps the future pinned so NIXL/UCX cannot write into freed memory. If quarantined entries exceed mQuarantineBudget or no worker has reached a final state for longer than mGlobalProgressDeadlineMs, isHealthy() flips false and getHealth() surfaces the snapshot for orchestration. * PyExecutor adds _can_terminate_request_now / _inflight_cancel_- requested_ids so _do_terminate_request defers freeing resources for any request that is still in disagg transfer state or whose cancel is in flight. Per-request timeouts no longer clear active_requests, the waiting queue, or set is_shutdown — that was the PR #13706 restart-loop trigger and is explicitly forbidden by the analysis doc. * The ADP synchronized pending-response flush from PR #13112 is ported here (the base commit precedes that merge). Transfer- completion responses created in _end_transfer_and_maybe_terminate are buffered and flushed at synchronised loop points so every DP rank participates in the tp_gather collective. Test coverage: * tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py — 14 unit tests for the deferred-terminate guard, the structured- cancel decision tree, and the ADP flush symmetry. They run without GPU or MPI so they fail fast in pre-merge CI. References: * analysis: docs/source/features/disagg-kv-transfer-hang-restart-analysis.md * contract: docs/source/features/disagg-kv-transfer-session-lifecycle.md Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 179 +++++- .../batch_manager/cacheTransceiver.cpp | 518 ++++++++++++------ .../batch_manager/cacheTransceiver.cpp | 77 ++- ...isagg-kv-transfer-hang-restart-analysis.md | 281 ++++++++++ .../disagg-kv-transfer-session-lifecycle.md | 122 +++++ .../_torch/pyexecutor/kv_cache_transceiver.py | 41 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 159 +++++- .../test_kv_transfer_session_lifecycle.py | 348 ++++++++++++ 8 files changed, 1532 insertions(+), 193 deletions(-) create mode 100644 docs/source/features/disagg-kv-transfer-hang-restart-analysis.md create mode 100644 docs/source/features/disagg-kv-transfer-session-lifecycle.md create mode 100644 tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 8f8330603893..f4abcc149e4d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -200,6 +200,60 @@ struct RequestStatuses std::unordered_set errorRequestIds; }; +/// Structured outcome of a cancellation request. Replaces the historical +/// boolean return that conflated "Python may free request resources" with +/// "C++ has nothing to do." Only @c kCancelledBeforeAdvertise and +/// @c kAlreadyComplete imply that downstream cleanup may proceed +/// immediately. The other states require Python to leave the request owned +/// by C++ until a final transfer state is observed (or the transceiver +/// reports unhealthy). +enum class TransferCancelResult : uint8_t +{ + /// Request is unknown to the transceiver (already cleaned up or never + /// registered). + kNotFound = 0, + /// Request has already reached worker-final state; the future is ready + /// and any allocated buffers have been released. + kAlreadyComplete = 1, + /// Request was queued but no transfer buffer/handle was advertised to + /// a peer yet. Memory is safe to release immediately. + kCancelledBeforeAdvertise = 2, + /// Cancellation was accepted but the worker is still in flight. The + /// transceiver retains ownership of the request and its future until a + /// final state is observed; Python must not free request resources. + kCancelRequestedInFlight = 3, + /// The transceiver is unhealthy (quarantine budget exceeded or no + /// backend progress past the global deadline). Orchestration should + /// restart the worker. + kBackendUnhealthy = 4, + /// Cancellation cannot be performed at this point (e.g., the sender is + /// in the middle of advertising the buffer to a peer). Retry later. + kNotCancellable = 5, +}; + +/// Health snapshot exposed for observability. All counters are best-effort +/// approximations updated on the executor thread; they are not real-time +/// kernel-level metrics. +struct TransceiverHealth +{ + bool isHealthy{true}; + /// Number of in-flight transfers that have exceeded their per-request + /// timeout but whose worker future has not yet reached a final state. + /// These transfers are tracked but their request resources are kept + /// pinned by C++ — Python must not free them. + size_t quarantinedTransferCount{0}; + /// Maximum number of quarantined transfers permitted before the + /// transceiver is flipped to unhealthy. + size_t quarantineBudget{0}; + /// Wall-clock seconds since the last observed worker future + /// transition. Only meaningful when the transceiver has tracked at + /// least one transfer. + double secondsSinceLastProgress{0.0}; + /// Wall-clock seconds beyond which "no progress" is treated as a + /// global backend wedge. + double globalProgressDeadlineSeconds{0.0}; +}; + class BaseCacheTransceiver { public: @@ -212,16 +266,66 @@ class BaseCacheTransceiver virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; - /// Check all requests transferring context, and return the requests that have completed or encountered an error. + /// Non-blocking poll. Returns the requests that have completed or + /// encountered an error since the last call. With + /// @c atLeastRequestNum unset the call defaults to a pure poll: it + /// must @b never block on an unready future on the executor thread. + /// Callers that genuinely need to drain (shutdown only) must use + /// @ref drainContextTransferStatus instead. virtual RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) = 0; + /// Non-blocking poll. Same contract as @ref checkContextTransferStatus. virtual void checkGenTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) = 0; + /// Blocking drain — @b only safe on dedicated drain/shutdown paths. + /// The default forwards to the polling variant for backward + /// compatibility; subclasses that own worker futures override this to + /// actually wait for completion. + virtual RequestStatuses drainContextTransferStatus(bool markComplete = false) + { + return checkContextTransferStatus(std::nullopt, markComplete); + } + + /// Blocking drain — @b only safe on dedicated drain/shutdown paths. + virtual void drainGenTransferStatus() + { + checkGenTransferStatus(std::nullopt); + } + [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; + /// Structured cancellation. See @ref TransferCancelResult for the full + /// contract. Subclasses override this; the boolean wrapper below maps + /// the structured result onto the historical "is the request safe to + /// release now" bool. + virtual TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) + { + return cancelRequest(llmRequest) ? TransferCancelResult::kCancelledBeforeAdvertise + : TransferCancelResult::kNotFound; + } + + /// Backward-compatible wrapper: returns true only when Python is safe + /// to free the request's resources immediately (pre-advertise cancel + /// or already-complete worker). Callers that need to distinguish + /// in-flight cancellation from "nothing to do" should use + /// @ref cancelRequestStructured. virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + + /// Whether the transceiver is currently healthy. Flips to false when + /// quarantined transfers exceed the budget or when no backend + /// progress has been observed past the global deadline. + [[nodiscard]] virtual bool isHealthy() const + { + return true; + } + + /// Health snapshot for orchestration / metrics. + [[nodiscard]] virtual TransceiverHealth getHealth() const + { + return {}; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -265,21 +369,90 @@ class CacheTransceiver : public BaseCacheTransceiver void checkGenTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) override; + /// Blocking drain — only invoked from shutdown/teardown. + RequestStatuses drainContextTransferStatus(bool markComplete = false) override; + + /// Blocking drain — only invoked from shutdown/teardown. + void drainGenTransferStatus() override; + [[nodiscard]] bool checkGenTransferComplete() const override; + TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) override; + virtual bool cancelRequest(LlmRequest* llmRequest) override; + [[nodiscard]] bool isHealthy() const override; + + [[nodiscard]] TransceiverHealth getHealth() const override; + private: void initializeCommState(); void setContextState(LlmRequest* llmRequest); + /// Per-entry tracking for tracked worker futures. We keep the future + /// pinned until the worker reaches a final state, even after the + /// per-request deadline elapses, so that cancelling a still-running + /// worker never frees its KV/buffer resources prematurely. + struct TrackedFuture + { + LlmRequest* request; + std::future future; + std::chrono::steady_clock::time_point deadline; + /// True once the per-request timeout has fired and we have flipped + /// the request to an error state. The entry stays in the vector + /// (and the future stays pinned) until the worker actually + /// finishes — possibly forever if the backend is wedged, in which + /// case the global progress deadline raises an unhealthy signal. + bool quarantined{false}; + /// True once we have already advertised a buffer/handle to a + /// peer. Pre-advertise cancellation can release normally; post- + /// advertise cancellation must wait for worker quiescence. + bool advertised{false}; + }; + + /// Update the global "last progress" timestamp and re-evaluate the + /// transceiver health flag against the quarantine budget and the + /// global progress deadline. Caller must hold @ref mHealthMutex. + void updateHealthLocked(); + + /// Drop a TrackedFuture entry from a vector. Decrements the + /// quarantine counter if the entry was quarantined, marks progress. + void releaseTrackedFutureLocked(std::vector& vec, size_t index); + + /// Apply per-entry timeout policy: if the configured + /// kvTransferTimeoutMs has elapsed and the future is not yet ready, + /// flip the entry to quarantined and surface an error to the caller. + /// Returns true if the entry was newly quarantined. + bool maybeQuarantineLocked(TrackedFuture& entry, RequestStatuses* outStatus); + + /// Internal worker for the public polling and drain entry points. + /// @c allowBlocking is true only on shutdown drain — that is the one + /// path where waiting on a worker future is acceptable. + RequestStatuses checkContextTransferStatusImpl( + std::optional const& atLeastRequestNum, bool markComplete, bool allowBlocking); + void checkGenTransferStatusImpl(std::optional const& atLeastRequestNum, bool allowBlocking); + + /// Compute the per-entry deadline for a tracked transfer. + [[nodiscard]] std::chrono::steady_clock::time_point computeTrackedFutureDeadline( + std::chrono::steady_clock::time_point requestStart) const; + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + std::vector mSenderFutures; + std::vector mRequesterFutures; mpi::MpiComm const* mMpiWorldComm{nullptr}; + /// Health/quarantine accounting. The mutex protects only this small + /// block; the executor thread is the sole writer/reader of the future + /// vectors above. + mutable std::mutex mHealthMutex; + size_t mQuarantinedTransferCount{0}; + size_t mQuarantineBudget{16}; + std::chrono::steady_clock::time_point mLastProgressTime{std::chrono::steady_clock::now()}; + std::chrono::milliseconds mGlobalProgressDeadlineMs{60'000}; + bool mIsHealthy{true}; + std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..ac50654e21b2 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -339,7 +339,8 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } setContextState(llmRequest); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -355,7 +356,8 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest.get(), std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mSenderFutures.push_back(TrackedFuture{llmRequest.get(), std::move(future), deadline, false, false}); } } @@ -374,7 +376,7 @@ void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& pair) { return pair.first->mRequestId == llmRequest->mRequestId; }) + [llmRequest](auto const& entry) { return entry.request->mRequestId == llmRequest->mRequestId; }) != mRequesterFutures.end()) { TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); @@ -382,7 +384,8 @@ void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) } auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mRequesterFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); } @@ -480,29 +483,127 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } } -RequestStatuses CacheTransceiver::checkContextTransferStatus( - std::optional const& atLeastRequestNum, bool markComplete) +// ----- Internal helpers ---------------------------------------------------- + +std::chrono::steady_clock::time_point CacheTransceiver::computeTrackedFutureDeadline( + std::chrono::steady_clock::time_point requestStart) const +{ + // Returns time_point::max() when no per-request timeout is configured, + // so the polling path treats the entry as "never times out" — matching + // the legacy behaviour for callers that did not opt into timeouts. + if (!mCacheTransceiverConfig.has_value()) + { + return std::chrono::steady_clock::time_point::max(); + } + auto const timeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + if (!timeoutMs.has_value()) + { + return std::chrono::steady_clock::time_point::max(); + } + return requestStart + std::chrono::milliseconds(timeoutMs.value()); +} + +void CacheTransceiver::updateHealthLocked() +{ + auto const now = std::chrono::steady_clock::now(); + auto const sinceProgress + = std::chrono::duration_cast(now - mLastProgressTime).count(); + bool wedged = sinceProgress > mGlobalProgressDeadlineMs.count() && (!mSenderFutures.empty() || !mRequesterFutures.empty()); + bool overBudget = mQuarantinedTransferCount > mQuarantineBudget; + bool nextHealthy = !overBudget && !wedged; + if (mIsHealthy && !nextHealthy) + { + TLLM_LOG_WARNING( + "CacheTransceiver flipping to UNHEALTHY: quarantined=%zu budget=%zu sinceProgressMs=%lld deadlineMs=%lld", + mQuarantinedTransferCount, mQuarantineBudget, static_cast(sinceProgress), + static_cast(mGlobalProgressDeadlineMs.count())); + } + mIsHealthy = nextHealthy; +} + +void CacheTransceiver::releaseTrackedFutureLocked(std::vector& vec, size_t index) +{ + TLLM_CHECK(index < vec.size()); + { + std::scoped_lock lk(mHealthMutex); + if (vec[index].quarantined && mQuarantinedTransferCount > 0) + { + --mQuarantinedTransferCount; + } + mLastProgressTime = std::chrono::steady_clock::now(); + updateHealthLocked(); + } + vec.erase(vec.begin() + index); +} + +bool CacheTransceiver::maybeQuarantineLocked(TrackedFuture& entry, RequestStatuses* outStatus) { - bool blockAll = !atLeastRequestNum.has_value(); - std::optional senderFutureTimeoutMs = std::nullopt; - // If blockAll is true, we want to block and not use a timeout - if (!blockAll && mCacheTransceiverConfig.has_value()) + if (entry.quarantined) { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + return false; } + if (entry.deadline == std::chrono::steady_clock::time_point::max()) + { + return false; + } + auto const now = std::chrono::steady_clock::now(); + if (now < entry.deadline) + { + return false; + } + entry.quarantined = true; + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + if (outStatus != nullptr) + { + outStatus->errorRequestIds.insert(entry.request->mRequestId); + } + { + std::scoped_lock lk(mHealthMutex); + ++mQuarantinedTransferCount; + updateHealthLocked(); + } + TLLM_LOG_WARNING( + "Quarantining request %ld: per-request KV transfer deadline exceeded; future stays pinned in C++ until " + "worker reaches a final state.", + entry.request->mRequestId); + return true; +} + +// ----- Public API ---------------------------------------------------------- +RequestStatuses CacheTransceiver::checkContextTransferStatus( + std::optional const& atLeastRequestNum, bool markComplete) +{ + // Non-blocking poll. The historical "blockAll" semantics (no + // atLeastRequestNum, fall through to future.get on every entry) are + // gone: a wedged worker thread must never freeze the executor event + // loop. Use @ref drainContextTransferStatus for shutdown drain. + return checkContextTransferStatusImpl(atLeastRequestNum, markComplete, /*allowBlocking=*/false); +} + +RequestStatuses CacheTransceiver::drainContextTransferStatus(bool markComplete) +{ + // Blocking — only safe on shutdown/teardown paths. + return checkContextTransferStatusImpl(std::nullopt, markComplete, /*allowBlocking=*/true); +} + +RequestStatuses CacheTransceiver::checkContextTransferStatusImpl( + std::optional const& atLeastRequestNum, bool markComplete, bool allowBlocking) +{ auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; + + // Pass 1: which futures are READY right now? std::vector contextCompleteRequestIds; - for (auto&& [request, future] : mSenderFutures) + for (auto const& entry : mSenderFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + if (entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - contextCompleteRequestIds.push_back(request->mRequestId); + contextCompleteRequestIds.push_back(entry.request->mRequestId); } } std::unordered_map frequencyMap; - if ((syncComm) && syncComm->getSize() > 1) + if (syncComm && syncComm->getSize() > 1) { auto gatherRequestIdVec = gatherRequestIds(syncComm, contextCompleteRequestIds); for (auto&& requestId : gatherRequestIdVec) @@ -517,101 +618,150 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( frequencyMap[requestId]++; } } - std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); + std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); std::sort(freqVec.begin(), freqVec.end(), - [](std::pair const& left, - std::pair const& right) { return left.second > right.second; }); + [](auto const& left, auto const& right) { return left.second > right.second; }); + + int const expectedFreq = syncComm ? syncComm->getSize() : 1; std::unordered_set toCompleteIdSet; - for (auto&& [requestId, freq] : freqVec) + for (auto const& [requestId, freq] : freqVec) { - if (freq == ((syncComm) ? syncComm->getSize() : 1)) + if (freq == expectedFreq) { toCompleteIdSet.insert(requestId); } } - // Make sure there are at least atLeastRequestNum requests in toCompleteIdSet. - // This will preserve the order of insertion for KVCache transfer requests. - for (auto it = mSenderFutures.begin(); - atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size()) && it != mSenderFutures.end(); ++it) + // For atLeastRequestNum > 0 we may need to admit additional requests. + // CRITICAL: prefer entries whose future is already ready. Never select + // an unready entry just to satisfy the count — that is the + // "future.get on selected-but-unready entry" wedge the analysis doc + // calls out. If we run out of ready entries we simply return what we + // have; the caller polls again next iteration. + if (atLeastRequestNum.has_value()) { - auto& [request, future] = *it; - toCompleteIdSet.insert(request->mRequestId); + for (auto const& entry : mSenderFutures) + { + if (static_cast(toCompleteIdSet.size()) >= atLeastRequestNum.value()) + { + break; + } + if (entry.future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + continue; + } + toCompleteIdSet.insert(entry.request->mRequestId); + } } RequestStatuses requestsStatus{}; - - // Complete all the requests in toCompleteIdSet - for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) + bool madeProgress = false; + for (size_t i = 0; i < mSenderFutures.size();) { - auto& [request, future] = *it; - if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) + auto& entry = mSenderFutures[i]; + bool const selected = toCompleteIdSet.find(entry.request->mRequestId) != toCompleteIdSet.end(); + bool const isReady = entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; + + if (selected && isReady) { 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()) - { - future.get(); - requestsStatus.completedRequestIds.insert(request->mRequestId); - if (markComplete) - { - request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); - } - it = mSenderFutures.erase(it); - } - else if (status == std::future_status::timeout) - { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; - } - else + entry.future.get(); + requestsStatus.completedRequestIds.insert(entry.request->mRequestId); + if (markComplete) { - TLLM_LOG_ERROR( - "Future returned unexpected status for request %ld. Marking as error", request->mRequestId); - - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); - it = mSenderFutures.erase(it); + entry.request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); - it = mSenderFutures.erase(it); + TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(entry.request->mRequestId); } + releaseTrackedFutureLocked(mSenderFutures, i); + madeProgress = true; + continue; } - else + + if (allowBlocking) { - ++it; + // Drain path: shutdown is happening. Wait for the worker to + // finish so we can release resources cleanly. Note this still + // never blocks on the executor thread — only shutdown calls + // this codepath. + try + { + entry.future.get(); + requestsStatus.completedRequestIds.insert(entry.request->mRequestId); + if (markComplete) + { + entry.request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + } + catch (std::exception const& e) + { + TLLM_LOG_ERROR("Error occurred during context transfer drain for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(entry.request->mRequestId); + } + releaseTrackedFutureLocked(mSenderFutures, i); + madeProgress = true; + continue; } + + // Per-entry timeout: surface error to caller, but keep the future + // pinned. Worker may still be writing; releasing the entry now + // would let Python free request resources that NIXL/UCX may still + // touch. The future stays here until the worker finishes (or the + // global progress deadline marks the transceiver unhealthy and + // orchestration restarts). + maybeQuarantineLocked(entry, &requestsStatus); + ++i; } + if (madeProgress) + { + std::scoped_lock lk(mHealthMutex); + mLastProgressTime = std::chrono::steady_clock::now(); + updateHealthLocked(); + } + else + { + std::scoped_lock lk(mHealthMutex); + updateHealthLocked(); + } return requestsStatus; } void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { - bool blockAll = !atLeastRequestNum.has_value(); + checkGenTransferStatusImpl(atLeastRequestNum, /*allowBlocking=*/false); +} + +void CacheTransceiver::drainGenTransferStatus() +{ + checkGenTransferStatusImpl(std::nullopt, /*allowBlocking=*/true); +} + +void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLeastRequestNum, bool allowBlocking) +{ + auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + std::vector genTransferReadyRequestIds; - for (auto&& [request, future] : mRequesterFutures) + for (auto const& entry : mRequesterFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + if (entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - genTransferReadyRequestIds.push_back(request->mRequestId); + genTransferReadyRequestIds.push_back(entry.request->mRequestId); } } - std::unordered_map frequencyMap; - std::vector toBlockRequestIds; - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - if ((syncComm) && syncComm->getSize() > 1) + std::unordered_map frequencyMap; + if (syncComm && syncComm->getSize() > 1) { auto gatherRequestIdVec = gatherRequestIds(syncComm, genTransferReadyRequestIds); for (auto&& requestId : gatherRequestIdVec) @@ -628,127 +778,97 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); - std::sort(freqVec.begin(), freqVec.end(), - [](std::pair const& left, - std::pair const& right) { return left.second > right.second; }); + [](auto const& left, auto const& right) { return left.second > right.second; }); + + int const expectedFreq = syncComm ? syncComm->getSize() : 1; std::unordered_set toCompleteIdSet; - size_t idx = 0; - while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + for (auto const& [requestId, freq] : freqVec) { - if (idx >= freqVec.size()) - { - break; - } - toCompleteIdSet.insert(freqVec.at(idx).first); - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); - } - else + if (freq == expectedFreq) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); + toCompleteIdSet.insert(requestId); } - idx++; } - idx = 0; - // insert order - while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + if (atLeastRequestNum.has_value()) { - if (idx >= mRequesterFutures.size()) + for (auto const& entry : mRequesterFutures) { - break; - } - if (toCompleteIdSet.find(mRequesterFutures.at(idx).first->mRequestId) == toCompleteIdSet.end()) - { - toCompleteIdSet.insert(mRequesterFutures.at(idx).first->mRequestId); - if (useMPI()) + if (static_cast(toCompleteIdSet.size()) >= atLeastRequestNum.value()) { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + break; } - else + if (entry.future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + continue; } + toCompleteIdSet.insert(entry.request->mRequestId); } - idx++; } - for (auto&& [requestId, freq] : freqVec) - { - if (freq == ((syncComm != nullptr) ? syncComm->getSize() : 1)) - { - toCompleteIdSet.insert(requestId); - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", - requestId, freq); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", requestId, freq); - } - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } - for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) + + bool madeProgress = false; + for (size_t i = 0; i < mRequesterFutures.size();) { - if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) + auto& entry = mRequesterFutures[i]; + bool const selected = toCompleteIdSet.find(entry.request->mRequestId) != toCompleteIdSet.end(); + bool const isReady = entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; + + if (selected && isReady) { try { - it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - - // Gather the kv cache transfer time from all workers and update to leader rank + entry.future.get(); + entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + updateKVCacheTransferBW(syncComm, entry.request); } } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during generation transfer for request %ld: %s", it->first->mRequestId, e.what()); - it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + TLLM_LOG_ERROR("Error occurred during generation transfer for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } - if (useMPI()) + releaseTrackedFutureLocked(mRequesterFutures, i); + madeProgress = true; + continue; + } + + if (allowBlocking) + { + try { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + entry.future.get(); + entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + updateKVCacheTransferBW(syncComm, entry.request); + } } - else + catch (std::exception const& e) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + TLLM_LOG_ERROR("Error occurred during generation transfer drain for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } - it = mRequesterFutures.erase(it); + releaseTrackedFutureLocked(mRequesterFutures, i); + madeProgress = true; + continue; } - else + + maybeQuarantineLocked(entry, /*outStatus=*/nullptr); + ++i; + } + + { + std::scoped_lock lk(mHealthMutex); + if (madeProgress) { - ++it; + mLastProgressTime = std::chrono::steady_clock::now(); } + updateHealthLocked(); } } @@ -757,17 +877,87 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRequest) { + TLLM_CHECK(llmRequest != nullptr); + + if (!isHealthy()) + { + return TransferCancelResult::kBackendUnhealthy; + } + if (llmRequest->isContextOnlyRequest()) { - return mCacheSender->cancelRequest(*llmRequest); + // Pre-advertise: the sender's queue still owns the request. + if (mCacheSender->cancelRequest(*llmRequest)) + { + return TransferCancelResult::kCancelledBeforeAdvertise; + } + // Did the worker future already finish? + for (size_t i = 0; i < mSenderFutures.size(); ++i) + { + if (mSenderFutures[i].request->mRequestId != llmRequest->mRequestId) + { + continue; + } + auto status = mSenderFutures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) + { + return TransferCancelResult::kAlreadyComplete; + } + return TransferCancelResult::kCancelRequestedInFlight; + } + return TransferCancelResult::kNotFound; } - else if (llmRequest->isGenerationOnlyRequest()) + if (llmRequest->isGenerationOnlyRequest()) { - return mCacheReceiver->cancelRequest(*llmRequest); + if (mCacheReceiver->cancelRequest(*llmRequest)) + { + return TransferCancelResult::kCancelledBeforeAdvertise; + } + for (size_t i = 0; i < mRequesterFutures.size(); ++i) + { + if (mRequesterFutures[i].request->mRequestId != llmRequest->mRequestId) + { + continue; + } + auto status = mRequesterFutures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) + { + return TransferCancelResult::kAlreadyComplete; + } + return TransferCancelResult::kCancelRequestedInFlight; + } + return TransferCancelResult::kNotFound; } - return false; + return TransferCancelResult::kNotCancellable; +} + +bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +{ + auto result = cancelRequestStructured(llmRequest); + return result == TransferCancelResult::kCancelledBeforeAdvertise || result == TransferCancelResult::kAlreadyComplete; +} + +bool CacheTransceiver::isHealthy() const +{ + std::scoped_lock lk(mHealthMutex); + return mIsHealthy; +} + +TransceiverHealth CacheTransceiver::getHealth() const +{ + std::scoped_lock lk(mHealthMutex); + auto const sinceProgress = std::chrono::duration_cast>( + std::chrono::steady_clock::now() - mLastProgressTime) + .count(); + return TransceiverHealth{ + mIsHealthy, + mQuarantinedTransferCount, + mQuarantineBudget, + sinceProgress, + std::chrono::duration_cast>(mGlobalProgressDeadlineMs).count(), + }; } } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 61cac5df4c7e..6b2db97ca7d5 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -17,8 +17,10 @@ #include "cacheTransceiver.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/bindingUtils.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/nanobind/common/customCasters.h" @@ -51,6 +53,14 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } + void respondAndSendLayerWise( + tb::RequestVector const& /*requests*/, std::shared_ptr const& /*progress*/) override + { + // Layer-wise sender path is not exposed to Python overrides; this is + // a stub so the trampoline class is fully concrete. + TLLM_THROW("respondAndSendLayerWise is not overridable from Python"); + } + void requestAndReceiveSync(tb::LlmRequest* llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); @@ -86,31 +96,64 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver void tb::CacheTransceiverBindings::initBindings(nb::module_& m) { + nb::enum_(m, "TransferCancelResult") + .value("NotFound", tb::TransferCancelResult::kNotFound) + .value("AlreadyComplete", tb::TransferCancelResult::kAlreadyComplete) + .value("CancelledBeforeAdvertise", tb::TransferCancelResult::kCancelledBeforeAdvertise) + .value("CancelRequestedInFlight", tb::TransferCancelResult::kCancelRequestedInFlight) + .value("BackendUnhealthy", tb::TransferCancelResult::kBackendUnhealthy) + .value("NotCancellable", tb::TransferCancelResult::kNotCancellable); + + nb::class_(m, "TransceiverHealth") + .def_ro("is_healthy", &tb::TransceiverHealth::isHealthy) + .def_ro("quarantined_transfer_count", &tb::TransceiverHealth::quarantinedTransferCount) + .def_ro("quarantine_budget", &tb::TransceiverHealth::quarantineBudget) + .def_ro("seconds_since_last_progress", &tb::TransceiverHealth::secondsSinceLastProgress) + .def_ro("global_progress_deadline_seconds", &tb::TransceiverHealth::globalProgressDeadlineSeconds); + + auto checkContextStatusLambda + = [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete) + { + RequestStatuses result; + { + nb::gil_scoped_release release; + result = self.checkContextTransferStatus(atLeastRequestNum, markComplete); + } + auto completedRequestIds + = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); + auto errorRequestIds = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); + return nb::make_tuple(completedRequestIds, errorRequestIds); + }; + + auto drainContextStatusLambda = [](tb::BaseCacheTransceiver& self, bool markComplete) + { + RequestStatuses result; + { + nb::gil_scoped_release release; + result = self.drainContextTransferStatus(markComplete); + } + auto completedRequestIds + = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); + auto errorRequestIds = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); + return nb::make_tuple(completedRequestIds, errorRequestIds); + }; + nb::class_(m, "BaseCacheTransceiver") .def("respond_and_send_async", &BaseCacheTransceiver::respondAndSendAsync) .def("request_and_receive_sync", &BaseCacheTransceiver::requestAndReceiveSync) .def("request_and_receive_async", &BaseCacheTransceiver::requestAndReceiveAsync) - .def( - "check_context_transfer_status", - [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete = false) - { - RequestStatuses result; - { - nb::gil_scoped_release release; - result = self.checkContextTransferStatus(atLeastRequestNum, markComplete); - } - - auto completedRequestIds - = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); - auto errorRequestIds - = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); - return nb::make_tuple(completedRequestIds, errorRequestIds); - }, + .def("check_context_transfer_status", checkContextStatusLambda, nb::arg("at_least_request_num") = std::nullopt, nb::arg("mark_complete") = false) + .def("drain_context_transfer_status", drainContextStatusLambda, nb::arg("mark_complete") = false) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus, nb::call_guard()) + .def("drain_gen_transfer_status", &BaseCacheTransceiver::drainGenTransferStatus, + nb::call_guard()) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete) - .def("cancel_request", &BaseCacheTransceiver::cancelRequest); + .def("cancel_request", &BaseCacheTransceiver::cancelRequest) + .def("cancel_request_structured", &BaseCacheTransceiver::cancelRequestStructured) + .def("is_healthy", &BaseCacheTransceiver::isHealthy) + .def("get_health", &BaseCacheTransceiver::getHealth); nb::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) diff --git a/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md b/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md new file mode 100644 index 000000000000..a8107d6043a7 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md @@ -0,0 +1,281 @@ +# Disaggregated KV transfer restart and hang analysis + +This note summarizes the latest behavior observed across PR 13706, PR 13713, +and PR 13728, then proposes the fix direction that avoids both unsafe memory +reuse and permanent serving hangs. + +## Observed behavior + +The latest tests split the failure into two different classes: + +| Change set | `py_executor.py` changes included | Observed behavior | Immediate implication | +| --- | --- | --- | --- | +| PR 13706 | yes | frequent prefill worker restarts | The Python policy is too aggressive and turns common in-flight transfer timeouts into process-level fail-close events. | +| PR 13706 | no | no frequent restarts | The C++ side alone is not what triggers those restart loops. | +| PR 13713 | yes | hangs | The hang is not explained only by Python fail-close behavior. | +| PR 13713 | no | hangs | There is still a lower-level C++ or transport progress problem. | +| PR 13728 | yes | hangs | The added Python containment does not repair the lower-level hang. | +| PR 13728 | no | hangs | The lower-level hang persists independently of Python cleanup policy. | + +```mermaid +flowchart TD + A["PR 13706 with Python changes"] --> B["Python sees transfer timeout"] + B --> C["Python fails closed"] + C --> D["Worker restarts often"] + + E["PR 13706 without Python changes"] --> F["No Python fail-close trigger"] + F --> G["No frequent restart loop"] + + H["PR 13713 or PR 13728"] --> I["Hangs with Python changes"] + H --> J["Hangs without Python changes"] + I --> K["Hang is below Python policy"] + J --> K +``` + +The PR 13706 result tells us that per-request Python fail-close is the wrong +default policy. It can be useful as an emergency containment mechanism, but it +is too blunt when it fires for routine transfer timeouts or cancellable races. + +The PR 13713 and PR 13728 results tell us something different: removing the +Python policy is not sufficient. If those images still hang without the +`py_executor.py` changes, then at least one C++ or backend path can still stop +making progress without surfacing a final request error. + +## Relevant code paths + +The Python-side restart behavior comes from timeout and cleanup paths around: + +- `PyExecutor._check_kv_transfer_timeout()` in + `tensorrt_llm/_torch/pyexecutor/py_executor.py`: marks in-flight context or + generation transfers as timed out. +- `PyExecutor._check_disagg_ctx_cache_transfer_status()`: observes timed-out + context transfers and calls `kv_cache_transceiver.cancel_request(request)`. +- `PyExecutor._handle_responses()`: observes timed-out generation transfers and + calls `kv_cache_transceiver.cancel_request(request)`. +- The PR 13706-style Python policy then treated unknown quiescence as a fatal + executor condition, cleared active queues, and made the worker exit or become + unhealthy. + +The lower-level hang behavior must be analyzed around the C++ transfer +progress paths: + +- `CacheTransceiver::checkContextTransferStatus()` and + `CacheTransceiver::checkGenTransferStatus()` in + `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`. +- `CacheSender::Impl` and `CacheReceiver::Impl` in + `cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`. +- Transfer-buffer assignment and release in + `cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp`. +- Agent/NIXL calls reached through + `cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp`. + +If a PR hangs with Python cleanup removed, the main event loop is either still +blocking on a C++ wait, or C++ has stranded an in-flight request/future without +making it observable as ready, error, or unhealthy. + +## Why the restart-only fix is not acceptable + +A pure fail-closed policy is memory safe, but it is operationally noisy: + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant Cpp as "C++ transceiver" + participant Orch as "Orchestrator" + + Py->>Py: "KV transfer timeout" + Py->>Cpp: "cancel_request(request)" + Cpp-->>Py: "cancel accepted" + Py->>Py: "quiescence unknown" + Py->>Py: "set shutdown or unhealthy" + Orch->>Py: "restart worker" +``` + +This avoids reusing memory that a transport may still write, but it restarts on +events that are not necessarily fatal. Under bursty traffic, cancellations, or +slow backend progress, that can turn a recoverable per-request timeout into a +cluster-wide restart loop. + +## Why the C++-only fix is not sufficient + +A C++-only approach can improve lifetime and avoid some leaks, but the latest +PR 13713 and PR 13728 observations show that it can still hang. That means the +implementation still lacks at least one of these properties: + +- every transfer future eventually becomes ready, error, or explicitly + abandoned with a health signal; +- no executor-facing status function blocks on a future or backend wait unless + the caller explicitly requested blocking behavior; +- cancellation has a final state that distinguishes "safe to free" from + "cancel requested but remote writes may still exist"; +- poisoned or quarantined resources are bounded and converted into an + explicit unhealthy signal before the executor silently starves. + +## Proper fix direction + +The fix should separate three concepts that the earlier PRs mixed together: + +1. Request failure: the user request can receive a timeout or network error. +2. Resource quiescence: C++ has proof that no sender, receiver, or backend + worker can still touch the request's KV blocks or transfer buffers. +3. Process health: the worker should restart only when the backend or resource + pool is no longer capable of safe progress. + +### 1. Python should not fail-close per request + +Python should keep these responsibilities: + +- mark transfer start time; +- mark a request timed out after `kv_transfer_timeout_ms`; +- call `cancel_request()` once for a timed-out transfer; +- avoid `_terminate_request()` while the request is still in a disaggregated + transfer state; +- report a request-level error only when C++ reports a final transfer state. + +Python should not set global shutdown or clear serving queues solely because a +single in-flight transfer timed out. That is the PR 13706 restart-loop trigger. + +### 2. C++ should own transfer sessions until final state + +C++ needs a transfer-session concept, explicit or equivalent, that owns all +objects needed by the in-flight transfer: + +- a `shared_ptr`; +- receive or send buffer-index leases; +- any request-owned KV-block lease needed by direct or zero-copy paths; +- the promise/future pair; +- cancellation state and timeout deadline. + +The session can release resources only on proven-safe outcomes: + +| Outcome | Safe action | +| --- | --- | +| Completed transfer | release buffers and KV-block leases | +| Peer returned explicit not-ready before data phase | release buffers | +| Cancelled before buffer advertisement | release buffers | +| Worker future ready with structured error | release or quarantine according to where the error occurred | +| Timeout after buffer advertisement and no backend completion | quarantine or poison advertised resources | +| Backend worker wedged or future never resolves | keep resources pinned or quarantined and raise health signal | + +The important rule is that `cancel_request()` must not mean "safe to free". +It should return a structured result such as: + +- `not_found`; +- `already_complete`; +- `cancelled_before_advertise`; +- `cancel_requested_in_flight`; +- `backend_unhealthy`; +- `not_cancellable`. + +Only the first three can permit immediate Python cleanup. The in-flight cases +must remain owned by C++ until a final status arrives or until the resource is +quarantined and the process is marked unhealthy. + +### 3. C++ status checks must be non-blocking by default + +The executor-facing status functions should be polling APIs: + +- `checkContextTransferStatus(0)` and `checkGenTransferStatus(0)` must not + block on an unready future. +- `atLeastNum=1` should skip unready futures rather than calling `future.get()` + on a selected-but-unready entry. +- Timeout handling should mark the transfer session cancelled or failed, but + should not erase the tracking future until the worker reaches a final state. +- Blocking behavior should exist only for explicit drain/shutdown paths. + +If a backend call can block indefinitely inside NIXL or UCX, it must not run on +the Python event-loop thread. If it runs on a dedicated worker and wedges, the +event loop must still be able to observe "no progress" and produce a health +signal. + +### 4. Quarantine should be bounded and health-driven + +Quarantine or poisoning is still needed for memory safety after a buffer has +been advertised to a peer and quiescence is unknown. But poisoning should not +immediately restart the worker for a single request unless no safe capacity +remains. + +Recommended policy: + +- pre-advertise cancellation releases normally; +- post-advertise unknown-quiescence quarantines the specific buffer or transfer + slot; +- the transceiver refuses to reuse quarantined resources; +- if the number of quarantined slots or pinned KV blocks exceeds a small budget, + or if no transfer progress occurs for a global backend deadline, mark the + transceiver unhealthy; +- health failure should be explicit and observable, so orchestration restarts + the worker instead of the server silently hanging. + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant S as "C++ transfer session" + participant B as "Transfer buffer or KV block" + participant H as "Health monitor" + + Py->>S: "timeout, cancel requested" + alt not advertised + S->>B: "release normally" + S-->>Py: "final request error" + else advertised and backend completes + S->>B: "release after completion" + S-->>Py: "final request error" + else advertised and backend does not quiesce + S->>B: "quarantine, do not reuse" + S-->>Py: "request error if safe to report" + S->>H: "increase poisoned or pinned resource count" + H-->>Py: "unhealthy only if budget or progress deadline is exceeded" + end +``` + +This policy avoids both bad extremes: + +- it does not free or reuse memory that may still be touched by the transport; +- it does not restart the worker for every ordinary timeout. + +## Expected behavior after the proper fix + +| Scenario | Expected result | +| --- | --- | +| Ordinary successful transfer | request completes and resources release normally | +| User cancellation before transfer advertisement | request fails or cancels; resources release normally | +| Timeout before transfer advertisement | request fails; resources release normally | +| Timeout after advertisement but backend later completes | request fails; resources release after worker final state | +| Timeout after advertisement and backend never completes | resource is quarantined; server keeps serving if capacity remains | +| Repeated non-quiescing transfers exhaust quarantine budget | worker becomes unhealthy and restarts with a clear reason | +| NIXL or UCX internal deadlock | no silent hang; health fails after progress deadline unless the backend root cause is fixed | + +## Recommended implementation order + +1. Remove the per-request Python fail-close behavior from PR 13706-style + branches. Keep only the "do not terminate while transfer is in progress" + guard. +2. Make C++ cancellation/status return a structured final-state enum instead + of a boolean. +3. Keep futures and transfer sessions tracked until the worker reaches ready, + error, or quarantined final state. +4. Add bounded quarantine counters for advertised transfer buffers and any + request-owned KV-block leases. +5. Add a transceiver health state that flips only when quarantine capacity is + exhausted or backend progress is absent past a global deadline. +6. Ensure ADP response enqueueing uses the synchronized pending-response flush + path from PR 13112 when testing on rc13-derived images. +7. Add regression tests for: + - PR 13706-style timeout bursts without restart loops; + - PR 13713/13728-style C++ path without Python changes, proving status APIs + return error or unhealthy instead of hanging; + - advertised-buffer timeout followed by block reuse, proving the block is + quarantined and not reused; + - ADP response flush, proving all ranks enter response collectives together. + +## Bottom line + +The Python fail-close policy explains PR 13706's frequent restarts, but it does +not explain PR 13713 or PR 13728 hanging without Python changes. The correct +solution is therefore not "more Python fail-close" and not "C++ lifetime only". + +The correct solution is a C++ transfer-session lifecycle with structured cancel +states, non-blocking status polling, bounded quarantine for unknown-quiescence +resources, and an explicit health signal only when safe serving capacity or +backend progress is genuinely lost. diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md new file mode 100644 index 000000000000..f81f4438135f --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -0,0 +1,122 @@ +# Disaggregated KV transfer: session lifecycle and cancel contract + +> **Status:** beta — applies to disaggregated serving with the C++ cache +> transceiver (NIXL / UCX backends). + +This note describes the contract between the Python executor and the C++ +cache transceiver for cancellation, timeout, quarantine, and process +health. It supersedes the per-request fail-close behaviour explored in +PR 13706 and the C++-only lifetime hardening in PRs 13713 / 13728: +both versions still hung in the field because neither isolated the +"backend is wedged but the executor thread is blocked" failure mode. + +The fix direction is documented in +`disagg-kv-transfer-hang-restart-analysis.md`. This file describes the +mechanism that this codebase actually implements. + +## Three responsibilities, separated + +| Responsibility | Owner | Mechanism | +| ----------------- | -------------------- | ----------------------------------------- | +| Request failure | Python executor | Surface a request-level error to the user when C++ reports a final transfer state. | +| Resource quiescence | C++ transceiver | Worker future stays pinned until ready / error / quarantine. | +| Process health | C++ transceiver | Quarantine budget + global progress deadline → `isHealthy() == false`. | + +## Status polling is non-blocking by default + +`BaseCacheTransceiver::checkContextTransferStatus` and +`checkGenTransferStatus` poll worker futures with `wait_for(0ms)` and +**never** call `future.get()` on an unready entry. When +`atLeastRequestNum > 0` is requested, the polling code admits additional +ready futures but skips unready ones rather than blocking the executor +thread. + +For shutdown / drain there are explicit blocking variants: +`drainContextTransferStatus` and `drainGenTransferStatus`. These are the +only paths permitted to block on a worker future. + +## Structured cancellation + +`cancelRequest(LlmRequest*)` returns a backward-compatible bool; +`cancelRequestStructured(LlmRequest*)` returns a +`TransferCancelResult` enum: + +| Value | Caller may release request resources? | +| --------------------------- | -------------------------------------- | +| `NotFound` | yes — C++ already reached a final state | +| `AlreadyComplete` | yes — worker future is ready | +| `CancelledBeforeAdvertise` | yes — no buffer was advertised to the peer | +| `CancelRequestedInFlight` | **no** — worker may still touch buffers | +| `BackendUnhealthy` | **no** — orchestration must restart | +| `NotCancellable` | **no** — retry later | + +Only the first three states allow Python to free the request's KV / +transfer-buffer resources. The remaining states keep the request owned +by C++; the Python executor adds the request id to +`_inflight_cancel_requested_ids`, which `_can_terminate_request_now` +checks before any `_do_terminate_request` call. + +## Bounded quarantine, explicit health signal + +When a tracked transfer exceeds its per-request `kvTransferTimeoutMs` +deadline and the worker future is still not ready, the entry is flipped +to *quarantined*: an error is surfaced to the caller, but the future +stays in the tracking vector. A counter (`mQuarantinedTransferCount`) +increments. If the counter exceeds `mQuarantineBudget`, **or** no +worker has reached a final state for longer than +`mGlobalProgressDeadlineMs`, the transceiver flips +`isHealthy()` to false. Orchestration is expected to restart the +worker on this signal — neither C++ nor Python frees pinned memory +locally. + +`getHealth()` returns a `TransceiverHealth` snapshot exposing +`is_healthy`, `quarantined_transfer_count`, `quarantine_budget`, +`seconds_since_last_progress`, and `global_progress_deadline_seconds` +for metrics / observability. + +## Python policy + +The `PyExecutor` keeps these responsibilities only: + +* mark transfer start time (`py_kv_transfer_start_time`); +* mark a request timed out after `kv_transfer_timeout_ms` + (`py_kv_transfer_timed_out`); +* call `cancel_request_structured()` once for a timed-out transfer; +* defer `_terminate_request()` while the request is still in a + disaggregated transfer state — see `_can_terminate_request_now`; +* report a request-level error only when the structured cancel result + is one of the safe-to-release outcomes. + +The PR 13706 anti-pattern of clearing `active_requests`, +`waiting_queue`, `request_accumulated`, and setting `is_shutdown=True` +on a single in-flight transfer timeout is **not** present and must not +be added back. Per-request timeouts are recoverable; only an unhealthy +transceiver should drive a process restart, and that restart is +orchestration's responsibility. + +## ADP response flush + +When attention data parallelism (`enable_attention_dp=True`) is +enabled, `_enqueue_responses` runs a `tp_gather` collective. Calling it +inside `_end_transfer_and_maybe_terminate` deadlocks because only the +DP rank that owns the request reaches that point. The executor instead +buffers transfer-completion responses in `_pending_transfer_responses` +and flushes the buffer at synchronised loop points where every DP rank +participates in the collective: + +* before `_handle_canceled_requests` in the PP path; +* before `_handle_canceled_requests` in the standard executor loop; +* unconditionally in the overlap loop (so ranks with differing + `should_process_previous_batch` still participate). + +`_flush_pending_transfer_responses` always calls `_enqueue_responses` +when ADP is enabled, even if the local buffer is empty, so the +collective is symmetric across ranks. + +## Test coverage + +* `tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py` + — unit tests for the deferred-terminate guard, the structured-cancel + decision tree, and the ADP flush symmetry. These tests intentionally + do not require GPU or MPI; they replay the policy decisions as plain + Python so failures fail-fast in pre-merge CI. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index e27290b472fa..a521842672a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -18,6 +18,12 @@ CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +# Re-export the structured cancel-result enum so the rest of the Python +# executor can match on it without reaching into the bindings module +# directly. +TransferCancelResult = tensorrt_llm.bindings.internal.batch_manager.TransferCancelResult +TransceiverHealth = tensorrt_llm.bindings.internal.batch_manager.TransceiverHealth + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -122,6 +128,32 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def cancel_request_structured(self, req: LlmRequest): + """Structured cancel result. + + Default implementation maps the historical bool result onto the + enum so subclasses (e.g., the Python transceiver) keep working + without changes. The C++ binding overrides this with a real + implementation that distinguishes pre-advertise cancellation from + in-flight cancellation. + """ + cancelled = self.cancel_request(req) + return TransferCancelResult.CancelledBeforeAdvertise if cancelled else TransferCancelResult.NotFound + + def is_healthy(self) -> bool: + """Whether the transceiver is currently healthy. + + Defaults to True; subclasses with a real health signal override. + """ + return True + + def get_health(self): + """Return a TransceiverHealth snapshot. Defaults to a healthy + snapshot for transceivers that have not yet implemented health + accounting. + """ + return None + @abstractmethod def prepare_context_requests(self, requests: List[LlmRequest]): """ @@ -205,6 +237,15 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): return self.impl.cancel_request(req) + def cancel_request_structured(self, req: LlmRequest): + return self.impl.cancel_request_structured(req) + + def is_healthy(self) -> bool: + return self.impl.is_healthy() + + def get_health(self): + return self.impl.get_health() + def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 98cf521d1ff4..8d6eb1304334 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -55,7 +55,7 @@ from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits from .hang_detector import HangDetector -from .kv_cache_transceiver import KvCacheTransceiver +from .kv_cache_transceiver import KvCacheTransceiver, TransferCancelResult from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, get_draft_token_length) from .model_engine import ModelEngine @@ -414,6 +414,21 @@ def __init__( self.max_num_active_requests = model_engine.get_max_num_sequences() self.active_requests: List[LlmRequest] = [] self.expected_num_active_requests = 0 + # Buffer for transfer-completion responses created inside + # _end_transfer_and_maybe_terminate. With ADP enabled, + # _enqueue_responses runs a tp_gather collective; calling it from + # _send_kv_async causes a collective mismatch because only the + # owning DP rank reaches that point. Buffering and flushing at a + # synchronised loop point (where every DP rank participates) + # avoids the deadlock. + self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Requests whose transfer timed out and we have already issued a + # cancel_request to C++. We must not call _terminate_request on + # these until C++ reports a final transfer state, otherwise we + # would free request resources while NIXL/UCX may still be + # writing into them. Keyed by py_request_id to survive request + # state transitions. + self._inflight_cancel_requested_ids: set[int] = set() # TODO: Remove the condition on the PP size once disagg support from KVCache reuse # path is fixed. self.async_transfer_manager = AsyncTransferManager( @@ -634,7 +649,14 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): response = request.create_response(False, self.dist.rank) if response: response.result.cached_tokens = request.cached_tokens - self._enqueue_responses([(request.py_request_id, response)]) + # Buffer the response instead of enqueueing immediately. + # _enqueue_responses runs a tp_gather collective when ADP + # is enabled; calling it here would deadlock because only + # the owning DP rank reaches this point. The buffer is + # flushed later at _flush_pending_transfer_responses + # where all ranks participate. + self._pending_transfer_responses.append( + (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): self.active_requests.remove(request) self._terminate_request(request) @@ -647,6 +669,20 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): if not self.async_transfer_manager.should_store_blocks: self._terminate_request(request) + def _flush_pending_transfer_responses(self): + """Enqueue buffered transfer-completion responses. + + Must be called at a point where ALL DP ranks execute in lockstep so + the tp_gather inside _enqueue_responses cannot deadlock. + """ + responses = self._pending_transfer_responses + self._pending_transfer_responses = [] + # Even when this rank has nothing to enqueue, we must still call + # _enqueue_responses with an empty list so that other DP ranks' + # tp_gather can complete. The collective is symmetric. + if responses or self.enable_attention_dp: + self._enqueue_responses(responses) + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1700,6 +1736,7 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): if self.kv_cache_transceiver: finished_ctx_reqs = scheduled_requests.context_requests_last_chunk self._send_kv_async(finished_ctx_reqs) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2174,6 +2211,7 @@ def _executor_loop(self): self._update_requests(sample_state, self.resource_manager) self._send_kv_async(scheduled_batch.all_requests()) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2421,6 +2459,11 @@ def _executor_loop_overlap(self): self._send_kv_async( self.previous_batch.scheduled_requests.all_requests()) + # Flush outside the conditional so all DP ranks participate + # in the tp_gather collective even when + # should_process_previous_batch differs between ranks. + self._flush_pending_transfer_responses() + if self.drafter is not None and self.use_spec_decode and should_process_previous_batch: # Cleanup previous draft resources used in the draft model self.drafter.cleanup_previous_draft_resources() @@ -3331,6 +3374,11 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) + # Worker reached a final state — clear any pending in-flight cancel + # bookkeeping for those requests. + if completed_req_ids: + self._inflight_cancel_requested_ids -= completed_req_ids + for request_id in completed_req_ids: if request_id not in requests_in_transfer: @@ -3349,14 +3397,42 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - # If cancel is successful, mark as complete so it can be cleaned up - # Otherwise, try at next iteration - if is_cancelled: + # Use structured cancel so we can distinguish + # "safe-to-free" from "cancellation accepted but worker + # still in flight". For the in-flight case we keep the + # request owned by C++ until a final transfer state is + # observed; freeing it now would let NIXL/UCX overwrite + # released memory. + cancel_result = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel_result in ( + TransferCancelResult.CancelledBeforeAdvertise, + TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound): + # NotFound at this point means the C++ worker future + # has already reached a final state and was erased + # by check_context_transfer_status above; continue + # cleanup as if it were AlreadyComplete. request.py_kv_transfer_start_time = None request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - + self._inflight_cancel_requested_ids.discard(request_id) self._end_transfer_and_maybe_terminate(request) + elif cancel_result == TransferCancelResult.CancelRequestedInFlight: + # Cancellation accepted but the worker is still + # touching the buffer/KV memory. Mark the request as + # in-flight-cancel so _do_terminate_request defers + # freeing resources, and re-check next iteration. + self._inflight_cancel_requested_ids.add(request_id) + elif cancel_result == TransferCancelResult.BackendUnhealthy: + # Soft signal — the transceiver is unhealthy. We do + # not touch this request; orchestration health checks + # will restart the worker once the global deadline + # elapses. + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"context request {request_id}; deferring to " + f"health-driven restart.") + self._inflight_cancel_requested_ids.add(request_id) self._check_cache_transfer_errors("context requests") @@ -3568,6 +3644,33 @@ def _handle_errors(self, for request in failed_requests: self._terminate_request(request) + def _can_terminate_request_now(self, request: LlmRequest) -> bool: + """Return False if the request still has an in-flight C++ + transfer that may write to its KV/buffer memory. + + Calling _do_terminate_request on such a request would let + free_resources() return KV blocks to the pool while NIXL/UCX is + still writing into them. We instead leave the request owned by + C++ until cancel_request_structured reports CancelledBefore- + Advertise / AlreadyComplete, or the transceiver is flipped to + unhealthy and orchestration restarts the worker. + """ + if self.kv_cache_transceiver is None: + return True + if request.is_dummy_request: + return True + # Active disagg transfer state: the C++ worker future is still + # tracked. Don't free. + if (request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + or request.state == + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS): + return False + # We have already issued a cancel for this request and C++ has + # not yet returned a final state. Do not free. + if request.py_request_id in self._inflight_cancel_requested_ids: + return False + return True + def _terminate_request(self, request: LlmRequest): # Dummy requests don't participate in disagg KV cache transfers, # so they must bypass the PP termination handler to avoid stale @@ -3580,6 +3683,17 @@ def _terminate_request(self, request: LlmRequest): self._do_terminate_request(request) def _do_terminate_request(self, request: LlmRequest): + if not self._can_terminate_request_now(request): + # Defer: the C++ worker future is still alive. We will retry + # termination on the next iteration once the structured + # cancel result transitions to CancelledBeforeAdvertise / + # AlreadyComplete, or the request reaches a final transfer + # state via _check_disagg_ctx_cache_transfer_status / + # check_gen_transfer_status. + logger.debug( + f"Deferring _do_terminate_request for {request.py_request_id}: " + f"transfer still in flight (state={request.state})") + return self.resource_manager.free_resources(request) if self.gather_all_responses or self.dist.rank == 0: @@ -3729,11 +3843,38 @@ def _handle_responses(self): # Check if generation request needs cleanup due to KV cache transfer timeout if request.py_kv_transfer_timed_out: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - if is_cancelled: + # Use structured cancel so an "accepted but worker + # in-flight" outcome does not free request resources + # while NIXL/UCX may still write into them. + cancel_result = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel_result in ( + TransferCancelResult.CancelledBeforeAdvertise, + TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound): + # NotFound at this stage means C++ already moved the + # request to a final state (worker future ready and + # erased by check_gen_transfer_status); treat the + # same as AlreadyComplete so we surface the timeout + # error to the user and free resources. + self._inflight_cancel_requested_ids.discard(req_id) self._handle_errors( error_msg=f"Request {request.py_request_id} timed out", requests=[request]) + elif cancel_result == TransferCancelResult.CancelRequestedInFlight: + # Worker still touching buffers — keep the request + # active so we re-poll next iteration. Once C++ + # transitions the worker future to a final state we + # will see NotFound / AlreadyComplete and clean up. + self._inflight_cancel_requested_ids.add(req_id) + new_active_requests.append(request) + elif cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"generation request {req_id}; deferring to " + f"health-driven restart.") + self._inflight_cancel_requested_ids.add(req_id) + new_active_requests.append(request) continue if request.is_generation_only_request() and not request.is_finished: diff --git a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py new file mode 100644 index 000000000000..44a4f41d571d --- /dev/null +++ b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the disaggregated KV-transfer session lifecycle. + +These tests exercise the Python policy layer added on top of the C++ +transceiver and intentionally avoid requiring a real GPU, MPI runtime, or +the full PyExecutor stack. The focus is the contract that the analysis +doc requires us to keep: + +* per-request timeouts must NOT trigger a global executor shutdown; +* a CancelRequestedInFlight result must keep the request owned by C++ + (no _do_terminate_request, no removal from active queues); +* CancelledBeforeAdvertise / AlreadyComplete / NotFound (after a final + state has been observed) must release resources promptly; +* the ADP pending-response flush is symmetric across DP ranks. +""" + +from __future__ import annotations + +import enum +import types +from dataclasses import dataclass, field +from typing import List, Optional + +import pytest + + +class _FakeTransferCancelResult(enum.Enum): + """Mirror of the C++ enum exposed by the binding so the tests do not + pull in the heavyweight binding module on import.""" + NotFound = 0 + AlreadyComplete = 1 + CancelledBeforeAdvertise = 2 + CancelRequestedInFlight = 3 + BackendUnhealthy = 4 + NotCancellable = 5 + + +@dataclass +class _FakeRequest: + py_request_id: int + state: str = "DISAGG_GENERATION_TRANS_IN_PROGRESS" + py_kv_transfer_timed_out: bool = False + py_kv_transfer_start_time: Optional[float] = None + is_dummy_request: bool = False + + @property + def is_disagg_generation_transmission_in_progress(self) -> bool: + return self.state == "DISAGG_GENERATION_TRANS_IN_PROGRESS" + + +@dataclass +class _FakeTransceiver: + """Programmable stand-in for the C++-backed KvCacheTransceiver. The + test sets `next_cancel_results` to a list keyed by request id; each + cancel call pops the head.""" + next_cancel_results: dict = field(default_factory=dict) + cancel_calls: List[int] = field(default_factory=list) + healthy: bool = True + + def cancel_request_structured(self, req: _FakeRequest): + self.cancel_calls.append(req.py_request_id) + return self.next_cancel_results.get( + req.py_request_id, _FakeTransferCancelResult.CancelRequestedInFlight) + + def is_healthy(self) -> bool: + return self.healthy + + +def _make_executor_under_test(): + """Build a minimal object whose interface matches what the helper + methods need. We avoid constructing a real PyExecutor because that + drags in CUDA, MPI, and the full model engine.""" + obj = types.SimpleNamespace() + + obj.kv_cache_transceiver = _FakeTransceiver() + obj._inflight_cancel_requested_ids = set() + obj._pending_transfer_responses = [] + obj.enable_attention_dp = False + obj.active_requests = [] + obj.terminated_request_ids = [] + obj.enqueue_calls = [] + + # Bind the real implementations (copied from py_executor.py) by + # patching through closures so we do not have to import the whole + # module — the unit-under-test functions are pure Python with no + # CUDA/MPI dependencies. + + def _can_terminate_request_now(self, request) -> bool: + if self.kv_cache_transceiver is None: + return True + if request.is_dummy_request: + return True + if request.state in ("DISAGG_CONTEXT_TRANS_IN_PROGRESS", + "DISAGG_GENERATION_TRANS_IN_PROGRESS"): + return False + if request.py_request_id in self._inflight_cancel_requested_ids: + return False + return True + + def _do_terminate_request(self, request): + if not self._can_terminate_request_now(request): + return + self.terminated_request_ids.append(request.py_request_id) + + def _enqueue_responses(self, responses): + # Mimics the dist-side aggregator; we just record what was sent. + self.enqueue_calls.append(list(responses)) + + def _flush_pending_transfer_responses(self): + responses = self._pending_transfer_responses + self._pending_transfer_responses = [] + if responses or self.enable_attention_dp: + self._enqueue_responses(responses) + + obj._can_terminate_request_now = types.MethodType( + _can_terminate_request_now, obj) + obj._do_terminate_request = types.MethodType(_do_terminate_request, obj) + obj._enqueue_responses = types.MethodType(_enqueue_responses, obj) + obj._flush_pending_transfer_responses = types.MethodType( + _flush_pending_transfer_responses, obj) + return obj + + +# --------------------------------------------------------------------------- +# _can_terminate_request_now +# --------------------------------------------------------------------------- + + +def test_can_terminate_request_defers_during_active_transfer(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + assert not executor._can_terminate_request_now(req) + + +def test_can_terminate_request_defers_when_cancel_in_flight(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, state="GENERATION_COMPLETE") + executor._inflight_cancel_requested_ids.add(42) + assert not executor._can_terminate_request_now(req) + + +def test_can_terminate_request_allows_when_safe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, state="GENERATION_COMPLETE") + assert executor._can_terminate_request_now(req) + + +def test_dummy_request_is_always_terminable(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS", + is_dummy_request=True) + assert executor._can_terminate_request_now(req) + + +def test_no_transceiver_means_terminable(): + executor = _make_executor_under_test() + executor.kv_cache_transceiver = None + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + assert executor._can_terminate_request_now(req) + + +# --------------------------------------------------------------------------- +# _do_terminate_request honors the guard +# --------------------------------------------------------------------------- + + +def test_do_terminate_defers_when_unsafe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=99, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + executor._do_terminate_request(req) + assert executor.terminated_request_ids == [] + + +def test_do_terminate_runs_when_safe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=99, state="GENERATION_COMPLETE") + executor._do_terminate_request(req) + assert executor.terminated_request_ids == [99] + + +# --------------------------------------------------------------------------- +# ADP response flush +# --------------------------------------------------------------------------- + + +def test_flush_pending_transfer_responses_empties_buffer(): + executor = _make_executor_under_test() + executor._pending_transfer_responses = [(1, "r1"), (2, "r2")] + executor._flush_pending_transfer_responses() + + assert executor._pending_transfer_responses == [] + assert executor.enqueue_calls == [[(1, "r1"), (2, "r2")]] + + +def test_flush_with_adp_enabled_calls_enqueue_even_when_empty(): + """Even when the local rank has nothing to enqueue, ADP requires every + rank to participate in tp_gather. If we skipped the call we would + deadlock the other rank's collective.""" + executor = _make_executor_under_test() + executor.enable_attention_dp = True + executor._flush_pending_transfer_responses() + assert executor.enqueue_calls == [[]] + + +def test_flush_without_adp_skips_when_empty(): + executor = _make_executor_under_test() + executor.enable_attention_dp = False + executor._flush_pending_transfer_responses() + assert executor.enqueue_calls == [] + + +# --------------------------------------------------------------------------- +# Per-request timeout policy: do NOT fail-close the executor +# --------------------------------------------------------------------------- + + +def _apply_timeout_policy(executor, + request, + *, + cancel_result=_FakeTransferCancelResult. + CancelRequestedInFlight): + """Mini-replay of the timeout branch in _handle_responses for a single + request. Mirrors the production decision tree without dragging in the + full event loop.""" + request.py_kv_transfer_timed_out = True + executor.kv_cache_transceiver.next_cancel_results[ + request.py_request_id] = cancel_result + + cancel = executor.kv_cache_transceiver.cancel_request_structured(request) + safe_to_release = cancel in ( + _FakeTransferCancelResult.CancelledBeforeAdvertise, + _FakeTransferCancelResult.AlreadyComplete, + _FakeTransferCancelResult.NotFound, + ) + if safe_to_release: + executor._inflight_cancel_requested_ids.discard(request.py_request_id) + return "released" + if cancel == _FakeTransferCancelResult.CancelRequestedInFlight: + executor._inflight_cancel_requested_ids.add(request.py_request_id) + return "deferred" + if cancel == _FakeTransferCancelResult.BackendUnhealthy: + executor._inflight_cancel_requested_ids.add(request.py_request_id) + return "unhealthy-deferred" + return "ignored" + + +def test_timeout_with_in_flight_worker_does_not_clear_active_queue(): + executor = _make_executor_under_test() + other_req = _FakeRequest(py_request_id=1, state="GENERATION_IN_PROGRESS") + timed_out = _FakeRequest(py_request_id=2, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + executor.active_requests = [other_req, timed_out] + + outcome = _apply_timeout_policy( + executor, + timed_out, + cancel_result=_FakeTransferCancelResult.CancelRequestedInFlight) + + # The PR 13706 anti-pattern was to clear active_requests / shutdown + # the executor on a per-request timeout. We must NOT do that. + assert outcome == "deferred" + assert other_req in executor.active_requests + assert timed_out in executor.active_requests + assert 2 in executor._inflight_cancel_requested_ids + # _do_terminate_request must defer. + executor._do_terminate_request(timed_out) + assert executor.terminated_request_ids == [] + + +def test_timeout_pre_advertise_releases_resources(): + executor = _make_executor_under_test() + timed_out = _FakeRequest(py_request_id=7, + state="DISAGG_CONTEXT_TRANS_IN_PROGRESS") + + outcome = _apply_timeout_policy( + executor, + timed_out, + cancel_result=_FakeTransferCancelResult.CancelledBeforeAdvertise) + + assert outcome == "released" + assert 7 not in executor._inflight_cancel_requested_ids + timed_out.state = "GENERATION_COMPLETE" + executor._do_terminate_request(timed_out) + assert executor.terminated_request_ids == [7] + + +def test_timeout_followed_by_worker_completion_eventually_releases(): + """Two-iteration sequence: first iteration the worker is in flight, + second iteration C++ has erased the future. Both must be handled + without ever calling _do_terminate_request before the worker + quiesces.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=11, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + + # Iteration 1: in flight. + assert _apply_timeout_policy( + executor, + request, + cancel_result=_FakeTransferCancelResult.CancelRequestedInFlight + ) == "deferred" + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [] + + # Iteration 2: C++ has reached final state and erased the future. + request.state = "DISAGG_TRANS_ERROR" + assert _apply_timeout_policy( + executor, request, + cancel_result=_FakeTransferCancelResult.NotFound) == "released" + assert 11 not in executor._inflight_cancel_requested_ids + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [11] + + +def test_unhealthy_transceiver_keeps_request_pinned(): + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=33, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + + outcome = _apply_timeout_policy( + executor, + request, + cancel_result=_FakeTransferCancelResult.BackendUnhealthy) + + # Even when the C++ side reports unhealthy, Python must NOT free the + # request right away. Orchestration is supposed to restart the + # worker; freeing locally would leak buffers to the wedged backend. + assert outcome == "unhealthy-deferred" + assert 33 in executor._inflight_cancel_requested_ids + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [] From aa53654b35f9451b1c7a7b5e1c8bf10c7736613d Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 21:10:53 -0700 Subject: [PATCH 2/5] [https://nvbugs/6104831][fix] Pin LlmRequest with shared_ptr through the cache transceiver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the previous commit on this branch. Changes the BaseCacheTransceiver / CacheTransceiver public API and the tracked- future vectors from raw LlmRequest* to std::shared_ptr so that the C++ object outlives every worker thread access regardless of when Python's _terminate_request drops its pybind reference. This closes the historical use-after-free class on raw LlmRequest* that showed up in field traces as mRequestId == 0x5555555555555555. Specifically: * respondAndSendAsync, requestAndReceiveSync, requestAndReceiveAsync, cancelRequest, and cancelRequestStructured all take std::shared_ptr. * TrackedFuture::request is now std::shared_ptr; the shared_ptr in mSenderFutures / mRequesterFutures keeps the LlmRequest pinned for as long as the transceiver tracks the worker future. * trtGptModelInflightBatching.cpp now passes the shared_ptr directly instead of stripping it via .get(). * The nanobind trampoline (PyCacheTransceiver) signatures match. The binding relies on (already included) to provide a Python-aware shared_ptr that pins the wrapper alive while C++ holds it. The Python-side _can_terminate_request_now guard stays in place. Its job is no longer about lifetime — the shared_ptr handles that — but about resource quiescence: free_resources() releases KV blocks back to the pool, and the transport may still be writing into them after the LlmRequest object exists. Termination must still wait for the structured cancel result to signal safety. The doc at docs/source/features/disagg-kv-transfer-session-lifecycle.md is updated to describe the new "Object lifetime" section. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 36 ++++++++++++------- .../batch_manager/cacheTransceiver.cpp | 30 ++++++++-------- .../trtGptModelInflightBatching.cpp | 6 ++-- .../batch_manager/cacheTransceiver.cpp | 6 ++-- .../disagg-kv-transfer-session-lifecycle.md | 21 +++++++++++ 5 files changed, 67 insertions(+), 32 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index f4abcc149e4d..fbdc7719a0f3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -258,13 +258,20 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + + // Methods take std::shared_ptr so the transceiver and its + // worker threads pin the LlmRequest's lifetime independently of when + // Python's _terminate_request drops the pybind reference. This closes + // the use-after-free class on raw LlmRequest* observed in the field + // (mRequestId reads as 0x5555555555555555 after Python free) without + // changing the Python policy layer's contract. + virtual void respondAndSendAsync(std::shared_ptr llmRequest) = 0; virtual void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) = 0; - virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; - virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; + virtual void requestAndReceiveSync(std::shared_ptr llmRequest) = 0; + virtual void requestAndReceiveAsync(std::shared_ptr llmRequest) = 0; /// Non-blocking poll. Returns the requests that have completed or /// encountered an error since the last call. With @@ -300,7 +307,7 @@ class BaseCacheTransceiver /// contract. Subclasses override this; the boolean wrapper below maps /// the structured result onto the historical "is the request safe to /// release now" bool. - virtual TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) + virtual TransferCancelResult cancelRequestStructured(std::shared_ptr llmRequest) { return cancelRequest(llmRequest) ? TransferCancelResult::kCancelledBeforeAdvertise : TransferCancelResult::kNotFound; @@ -311,7 +318,7 @@ class BaseCacheTransceiver /// or already-complete worker). Callers that need to distinguish /// in-flight cancellation from "nothing to do" should use /// @ref cancelRequestStructured. - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; /// Whether the transceiver is currently healthy. Flips to false when /// quarantined transfers exceed the budget or when no backend @@ -356,13 +363,13 @@ class CacheTransceiver : public BaseCacheTransceiver virtual ~CacheTransceiver(); - void respondAndSendAsync(LlmRequest* llmRequest) override; + void respondAndSendAsync(std::shared_ptr llmRequest) override; void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) override; - void requestAndReceiveSync(LlmRequest* llmRequest) override; - void requestAndReceiveAsync(LlmRequest* llmRequest) override; + void requestAndReceiveSync(std::shared_ptr llmRequest) override; + void requestAndReceiveAsync(std::shared_ptr llmRequest) override; RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) override; @@ -377,9 +384,9 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) override; + TransferCancelResult cancelRequestStructured(std::shared_ptr llmRequest) override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; [[nodiscard]] bool isHealthy() const override; @@ -393,10 +400,15 @@ class CacheTransceiver : public BaseCacheTransceiver /// Per-entry tracking for tracked worker futures. We keep the future /// pinned until the worker reaches a final state, even after the /// per-request deadline elapses, so that cancelling a still-running - /// worker never frees its KV/buffer resources prematurely. + /// worker never frees its KV/buffer resources prematurely. The + /// shared_ptr also pins the LlmRequest object's lifetime — Python + /// may drop its pybind reference at any time, but the C++ worker + /// thread (and any asynchronous queue inside CacheSender / + /// CacheReceiver) can dereference the request safely as long as the + /// transceiver has not erased this entry. struct TrackedFuture { - LlmRequest* request; + std::shared_ptr request; std::future future; std::chrono::steady_clock::time_point deadline; /// True once the per-request timeout has fired and we have flipped diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index ac50654e21b2..f0473d409613 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -323,7 +323,7 @@ void CacheTransceiver::setContextState(LlmRequest* llmRequest) } } -void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) +void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS); @@ -337,7 +337,7 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); + setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); @@ -357,11 +357,11 @@ void CacheTransceiver::respondAndSendLayerWise( setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{llmRequest.get(), std::move(future), deadline, false, false}); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { @@ -371,15 +371,16 @@ void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); } -void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); + auto const reqId = llmRequest->mRequestId; if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& entry) { return entry.request->mRequestId == llmRequest->mRequestId; }) + [reqId](auto const& entry) { return entry.request->mRequestId == reqId; }) != mRequesterFutures.end()) { - TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); + TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", reqId); return; } @@ -822,7 +823,7 @@ void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLe entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - updateKVCacheTransferBW(syncComm, entry.request); + updateKVCacheTransferBW(syncComm, entry.request.get()); } } catch (std::exception const& e) @@ -844,7 +845,7 @@ void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLe entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - updateKVCacheTransferBW(syncComm, entry.request); + updateKVCacheTransferBW(syncComm, entry.request.get()); } } catch (std::exception const& e) @@ -877,7 +878,7 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRequest) +TransferCancelResult CacheTransceiver::cancelRequestStructured(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest != nullptr); @@ -886,6 +887,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe return TransferCancelResult::kBackendUnhealthy; } + auto const reqId = llmRequest->mRequestId; if (llmRequest->isContextOnlyRequest()) { // Pre-advertise: the sender's queue still owns the request. @@ -896,7 +898,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe // Did the worker future already finish? for (size_t i = 0; i < mSenderFutures.size(); ++i) { - if (mSenderFutures[i].request->mRequestId != llmRequest->mRequestId) + if (mSenderFutures[i].request->mRequestId != reqId) { continue; } @@ -917,7 +919,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe } for (size_t i = 0; i < mRequesterFutures.size(); ++i) { - if (mRequesterFutures[i].request->mRequestId != llmRequest->mRequestId) + if (mRequesterFutures[i].request->mRequestId != reqId) { continue; } @@ -933,9 +935,9 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe return TransferCancelResult::kNotCancellable; } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { - auto result = cancelRequestStructured(llmRequest); + auto result = cancelRequestStructured(std::move(llmRequest)); return result == TransferCancelResult::kCancelledBeforeAdvertise || result == TransferCancelResult::kAlreadyComplete; } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 8bb2c0e2ba88..94679a54f022 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -915,7 +915,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1596,11 +1596,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 6b2db97ca7d5..bd02e750b9a2 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -48,7 +48,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); - void respondAndSendAsync(tb::LlmRequest* llmRequest) override + void respondAndSendAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } @@ -61,12 +61,12 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver TLLM_THROW("respondAndSendLayerWise is not overridable from Python"); } - void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + void requestAndReceiveSync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); } - void requestAndReceiveAsync(tb::LlmRequest* llmRequest) override + void requestAndReceiveAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveAsync, llmRequest); } diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index f81f4438135f..ae7949cdbff3 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -22,6 +22,27 @@ mechanism that this codebase actually implements. | Resource quiescence | C++ transceiver | Worker future stays pinned until ready / error / quarantine. | | Process health | C++ transceiver | Quarantine budget + global progress deadline → `isHealthy() == false`. | +### Object lifetime + +The C++ transceiver takes ownership of the `LlmRequest` for the +duration of an in-flight transfer: + +* `BaseCacheTransceiver::respondAndSendAsync`, + `requestAndReceiveSync`, `requestAndReceiveAsync`, and `cancelRequest` + / `cancelRequestStructured` all take `std::shared_ptr`. +* The tracked-future vectors (`mSenderFutures`, `mRequesterFutures`) + store the same shared_ptr, so the C++ object outlives every worker + thread access regardless of when Python's `_terminate_request` drops + its pybind reference. + +This closes the historical "raw `LlmRequest*` use-after-free" class +(`mRequestId == 0x5555555555555555` field traces). The Python-side +`_can_terminate_request_now` guard remains in place — but its job is +now memory *quiescence*, not lifetime: even with the request object +pinned, `free_resources()` would release KV blocks back to the pool +while the transport may still be writing into them, so termination +must still wait for the structured cancel result to signal safety. + ## Status polling is non-blocking by default `BaseCacheTransceiver::checkContextTransferStatus` and From ea0c1b59026ec808d1111267b82e0f2ff40cfd9c Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 22:02:40 -0700 Subject: [PATCH 3/5] [https://nvbugs/6104831][fix] Pin LlmRequest in dataTransceiver worker queues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the prior commit. The worker queues inside CacheSender::Impl and CacheReceiver::Impl still held raw LlmRequest* pointers — meaning the executor's mSenderFutures / mRequesterFutures pinned the lifetime but the worker thread that actually dereferences the request had only a raw observer. If the executor erased its tracking entry while the worker was mid-flight, the LlmRequest could be freed under the worker. This commit closes that surface: * CacheSender::Impl::Response::mRequest → std::shared_ptr. * CacheReceiver::Impl::RequestAndPromise::mRequest → std::shared_ptr. Move/copy semantics simplified now that the field is a smart pointer. * CacheSender::sendAsync(LlmRequest&) → sendAsync(std::shared_ptr). * CacheReceiver::receiveAsync(LlmRequest&) → receiveAsync(std::shared_ptr). * CacheReceiver::Impl::requestAndReceiveAsyncMultiThreads similarly. * CacheReceiver::Impl::receiveAsync now captures the shared_ptr by value in the std::async lambda so the worker thread pins the LlmRequest independently of the caller. * CacheTransceiver::respondAndSendAsync/-LayerWise/requestAndReceive* pass the shared_ptr (no longer .get()-strip it) and move into the TrackedFuture entry where appropriate. Also fixes the eval-order UAF in CacheSender::Impl::handleAsyncSend that PR #13713 / PR #13728 called out: once Response::mRequest is a shared_ptr, the one-liner sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp)); becomes undefined behaviour because C++ argument evaluation order is unspecified — the compiler may evaluate std::move(resp) first, leaving resp.mRequest empty when reading mRequestId. Materialise the id into a local before the move. The PyCacheTransceiver trampoline's bool cancelRequest override is updated to match the new shared_ptr signature. The doc at docs/source/features/disagg-kv-transfer-session-lifecycle.md spells out the full ownership chain: executor tracking + worker queue both hold shared_ptr, and TransferSession::mRequest is an ephemeral observer used only inside the worker frame where the shared_ptr is already held. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 12 +-- .../batch_manager/dataTransceiver.cpp | 91 +++++++++---------- .../batch_manager/dataTransceiver.h | 14 +-- .../batch_manager/cacheTransceiver.cpp | 2 +- .../disagg-kv-transfer-session-lifecycle.md | 29 ++++-- 5 files changed, 80 insertions(+), 68 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f0473d409613..03726ece6a72 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -338,9 +338,9 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques return; } setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(*llmRequest); + auto future = mCacheSender->sendAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); + mSenderFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -355,7 +355,7 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(*llmRequest); + auto future = mCacheSender->sendAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } @@ -365,7 +365,7 @@ void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequ { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto future = mCacheReceiver->receiveAsync(llmRequest); future.get(); } llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); @@ -384,10 +384,10 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq return; } - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto future = mCacheReceiver->receiveAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mRequesterFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + mRequesterFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); } std::vector gatherRequestIds( diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 3ecceb9f3f2c..5a8545e17bb6 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -296,16 +296,17 @@ class CacheSender::Impl } } - [[nodiscard]] std::future sendAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future sendAsync(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); std::promise promise; auto future = promise.get_future(); - llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + auto const reqId = llmRequest->mRequestId; { { std::scoped_lock lkResp(mSenderMutex); - mReadyResponses.emplace( - llmRequest.mRequestId, Response{std::addressof(llmRequest), std::move(promise)}); + mReadyResponses.emplace(reqId, Response{std::move(llmRequest), std::move(promise)}); } std::unique_lock lkCond(mCondMutex); mAnyReady = true; @@ -477,7 +478,11 @@ class CacheSender::Impl private: struct Response { - LlmRequest* mRequest; + // shared_ptr so the LlmRequest stays alive as long as the worker + // queue holds the response, independently of the executor's + // mSenderFutures tracking. Closes the raw-pointer UAF surface in + // the worker thread. + std::shared_ptr mRequest; std::promise mPromise; }; @@ -511,7 +516,12 @@ class CacheSender::Impl resp = std::move(resource.mSendQueue.front()); resource.mSendQueue.pop_front(); } - sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp)); + // Materialize the request id before std::move(resp) — C++ + // argument evaluation order is unspecified, so reading + // resp.mRequest->mRequestId after the move would dereference + // an empty shared_ptr. + auto const reqId = resp.mRequest->mRequestId; + sendAndRemoveResponse(reqId, std::move(resp)); } } @@ -753,23 +763,29 @@ class CacheReceiver::Impl TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); } - [[nodiscard]] std::future receiveAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future receiveAsync(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); // TODO: Modify the implementation here to avoid frequent thread creation. - return std::async(std::launch::async, &CacheReceiver::Impl::requestSync, this, std::ref(llmRequest)); + // The worker captures the shared_ptr by value, pinning the LlmRequest + // for the lifetime of the request worker independently of the + // executor-side tracking. + return std::async(std::launch::async, + [this, llmRequest = std::move(llmRequest)]() mutable { this->requestSync(*llmRequest); }); } - [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(LlmRequest& llmRequest) + [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); try { auto promise = std::make_unique>(); auto future = promise->get_future(); - TLLM_CHECK(llmRequest.getDataTransceiverState().getCommState().has_value()); + TLLM_CHECK(llmRequest->getDataTransceiverState().getCommState().has_value()); std::string processInfo = kDefaultProcessInfo; if (common::getEnvRequestKVCacheConcurrent()) { - processInfo = llmRequest.getDataTransceiverState().getCommState()->toString(); + processInfo = llmRequest->getDataTransceiverState().getCommState()->toString(); } if (mInstanceToAsyncResource.find(processInfo) == mInstanceToAsyncResource.end()) { @@ -782,7 +798,7 @@ class CacheReceiver::Impl auto& asyncResource = mInstanceToAsyncResource.at(processInfo); { std::unique_lock lck(asyncResource->mMtxForQueue); - asyncResource->mRequestsQueue.emplace_back(std::addressof(llmRequest), std::move(promise)); + asyncResource->mRequestsQueue.emplace_back(std::move(llmRequest), std::move(promise)); } asyncResource->mCVforQueue.notify_all(); return future; @@ -1074,47 +1090,24 @@ class CacheReceiver::Impl struct RequestAndPromise { - LlmRequest* mRequest; + // shared_ptr so the LlmRequest stays alive while the worker queue + // holds it, independently of the executor's mRequesterFutures + // tracking. + std::shared_ptr mRequest; std::unique_ptr> mPromise; - RequestAndPromise() - : mRequest(nullptr) - , mPromise(nullptr) - { - } + RequestAndPromise() = default; - RequestAndPromise(LlmRequest* request, std::unique_ptr>&& promise) - : mRequest(request) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise) + : mRequest(std::move(request)) , mPromise(std::move(promise)) { } RequestAndPromise(RequestAndPromise const&) = delete; - - RequestAndPromise(RequestAndPromise&& other) noexcept - : mRequest(other.mRequest) - , mPromise(std::move(other.mPromise)) - { - other.mRequest = nullptr; - } - - RequestAndPromise& operator=(RequestAndPromise&& other) noexcept - { - if (this != &other) - { - mRequest = nullptr; - if (mPromise) - { - mPromise.reset(); - } - - mRequest = other.mRequest; - mPromise = std::move(other.mPromise); - - other.mRequest = nullptr; - } - return *this; - } + RequestAndPromise& operator=(RequestAndPromise const&) = delete; + RequestAndPromise(RequestAndPromise&&) noexcept = default; + RequestAndPromise& operator=(RequestAndPromise&&) noexcept = default; }; struct AsyncResource @@ -1209,9 +1202,9 @@ CacheSender::CacheSender( { } -std::future CacheSender::sendAsync(LlmRequest& llmRequest) const +std::future CacheSender::sendAsync(std::shared_ptr llmRequest) const { - return mImpl->sendAsync(llmRequest); + return mImpl->sendAsync(std::move(llmRequest)); } executor::kv_cache::CommState const& CacheSender::getCommState() const @@ -1252,9 +1245,9 @@ CacheReceiver::CacheReceiver( { } -std::future CacheReceiver::receiveAsync(LlmRequest& llmRequest) const +std::future CacheReceiver::receiveAsync(std::shared_ptr llmRequest) const { - return mImpl->requestAndReceiveAsyncMultiThreads(llmRequest); + return mImpl->requestAndReceiveAsyncMultiThreads(std::move(llmRequest)); } CacheReceiver::~CacheReceiver() = default; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 4d072a04521b..d81207824829 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -256,10 +256,11 @@ class CacheSender CacheSender() = default; /// @brief Asynchronously respond to the request and send data. - /// @param llmRequest Request object. Its data should be ready when called, and the data for this request - /// should remain valid until future synchronization. + /// @param llmRequest Request object. The transceiver pins the LlmRequest + /// for the lifetime of the worker future via the shared_ptr, so the + /// caller may drop its own strong reference once this returns. /// @return Once the data is fully sent, the future object will become valid. - [[nodiscard]] virtual std::future sendAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future sendAsync(std::shared_ptr llmRequest) const; /// @brief Return the internal communicator status. /// @return The communicator status. @@ -313,10 +314,11 @@ class CacheReceiver CacheReceiver() = default; /// @brief Asynchronously send a request to receive data. - /// @param llmRequest Request object. Its data should be in an allocated but unwritten state when called, and the - /// data for this request should remain intact only after future synchronization. + /// @param llmRequest Request object. The transceiver pins the LlmRequest + /// for the lifetime of the worker future via the shared_ptr, so the + /// caller may drop its own strong reference once this returns. /// @return Once the data is fully received, the future object will become valid. - [[nodiscard]] virtual std::future receiveAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future receiveAsync(std::shared_ptr llmRequest) const; virtual TransferSession sendRequestInfo(LlmRequest const& llmRequest); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index bd02e750b9a2..cf02471ff473 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -87,7 +87,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(checkGenTransferComplete); } - bool cancelRequest(tb::LlmRequest* llmRequest) override + bool cancelRequest(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(cancelRequest, llmRequest); } diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index ae7949cdbff3..c072d663d754 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -31,18 +31,35 @@ duration of an in-flight transfer: `requestAndReceiveSync`, `requestAndReceiveAsync`, and `cancelRequest` / `cancelRequestStructured` all take `std::shared_ptr`. * The tracked-future vectors (`mSenderFutures`, `mRequesterFutures`) - store the same shared_ptr, so the C++ object outlives every worker - thread access regardless of when Python's `_terminate_request` drops - its pybind reference. - -This closes the historical "raw `LlmRequest*` use-after-free" class -(`mRequestId == 0x5555555555555555` field traces). The Python-side + store the same shared_ptr. +* Each worker queue **also** stores a `std::shared_ptr` — + `CacheSender::Impl::Response::mRequest` and + `CacheReceiver::Impl::RequestAndPromise::mRequest`. The worker thread + thus pins the LlmRequest independently of the executor, so a wedged + worker cannot be freed by the executor erasing its tracking entry. + +Together this closes the historical "raw `LlmRequest*` use-after-free" +class (`mRequestId == 0x5555555555555555` field traces). The Python-side `_can_terminate_request_now` guard remains in place — but its job is now memory *quiescence*, not lifetime: even with the request object pinned, `free_resources()` would release KV blocks back to the pool while the transport may still be writing into them, so termination must still wait for the structured cancel result to signal safety. +`CacheSender::Impl::handleAsyncSend` materialises the request id into +a local before `std::move(resp)` to avoid a C++ argument-eval-order +bug — once `Response::mRequest` is a `shared_ptr`, the previous +one-liner `sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp))` +became undefined behaviour because the compiler may evaluate the move +first and leave `resp.mRequest` empty when reading `mRequestId`. Field +traces of PR #13713 had a SIGSEGV at exactly that argument +construction; the materialised-id pattern fixes it. + +`TransferSession::mRequest` is an ephemeral `LlmRequest const*` that is +only set inside `sendSync` and used during measurement / formatter +calls within the same worker frame; the worker still holds the +shared_ptr at that point, so the observer is safe by construction. + ## Status polling is non-blocking by default `BaseCacheTransceiver::checkContextTransferStatus` and From d7d189a108f24b6ec19c8be73905d46c8fb011e7 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 22:09:22 -0700 Subject: [PATCH 4/5] [https://nvbugs/6104831][fix] Audit fixes: cancel ordering, trampoline size, test callsites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three issues caught in self-review of the prior commits: 1. cancelRequestStructured returned BackendUnhealthy too eagerly — even for a request that was still queued pre-advertise. Pre-advertise release is always safe (no buffer was ever exposed to a peer), and freeing those resources during the unhealthy window REDUCES backend pressure rather than adds to it. Reorder: check sender / receiver queues first, return CancelledBeforeAdvertise unconditionally if found there. Only consult isHealthy() for the in-flight branch, where freeing would race with NIXL/UCX writes. 2. NB_TRAMPOLINE size was 6 but the trampoline has 7 NB_OVERRIDE_PURE entries (pre-existing under-allocation of nanobind's dispatch hash table). Bump to 7 since we touched the file. 3. cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp called mSender->sendAsync(*llmRequest) and mRequester->receiveAsync(*llmRequest) passing LlmRequest&. After the prior commit those signatures take std::shared_ptr. The test was already holding a shared_ptr (request->mLlmRequest), so just drop the dereference. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 71 ++++++++----------- .../batch_manager/cacheTransceiver.cpp | 4 +- .../multi_gpu/cacheTransceiverTest.cpp | 8 +-- 3 files changed, 38 insertions(+), 45 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 03726ece6a72..70f7d2736b85 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -882,57 +882,48 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(std::shared_ptrmRequestId; + bool const sender = llmRequest->isContextOnlyRequest(); + bool const receiver = llmRequest->isGenerationOnlyRequest(); + if (!sender && !receiver) { - return TransferCancelResult::kBackendUnhealthy; + return TransferCancelResult::kNotCancellable; } - auto const reqId = llmRequest->mRequestId; - if (llmRequest->isContextOnlyRequest()) + // Pre-advertise check first: a request that is still queued in the + // sender / receiver's pending list has not exposed any buffer to a + // peer and is always safe to release — even when the transceiver + // is otherwise unhealthy. Releasing pre-advertise requests during + // an unhealthy window reduces the resource pressure on the wedged + // backend rather than adding to it. + if (sender && mCacheSender->cancelRequest(*llmRequest)) { - // Pre-advertise: the sender's queue still owns the request. - if (mCacheSender->cancelRequest(*llmRequest)) - { - return TransferCancelResult::kCancelledBeforeAdvertise; - } - // Did the worker future already finish? - for (size_t i = 0; i < mSenderFutures.size(); ++i) - { - if (mSenderFutures[i].request->mRequestId != reqId) - { - continue; - } - auto status = mSenderFutures[i].future.wait_for(std::chrono::milliseconds(0)); - if (status == std::future_status::ready) - { - return TransferCancelResult::kAlreadyComplete; - } - return TransferCancelResult::kCancelRequestedInFlight; - } - return TransferCancelResult::kNotFound; + return TransferCancelResult::kCancelledBeforeAdvertise; } - if (llmRequest->isGenerationOnlyRequest()) + if (receiver && mCacheReceiver->cancelRequest(*llmRequest)) { - if (mCacheReceiver->cancelRequest(*llmRequest)) + return TransferCancelResult::kCancelledBeforeAdvertise; + } + + auto& futures = sender ? mSenderFutures : mRequesterFutures; + for (size_t i = 0; i < futures.size(); ++i) + { + if (futures[i].request->mRequestId != reqId) { - return TransferCancelResult::kCancelledBeforeAdvertise; + continue; } - for (size_t i = 0; i < mRequesterFutures.size(); ++i) + auto status = futures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) { - if (mRequesterFutures[i].request->mRequestId != reqId) - { - continue; - } - auto status = mRequesterFutures[i].future.wait_for(std::chrono::milliseconds(0)); - if (status == std::future_status::ready) - { - return TransferCancelResult::kAlreadyComplete; - } - return TransferCancelResult::kCancelRequestedInFlight; + return TransferCancelResult::kAlreadyComplete; } - return TransferCancelResult::kNotFound; + // Worker is mid-flight. If the transceiver is unhealthy the + // caller must defer cleanup pending an orchestration restart; + // otherwise it is just a normal in-flight cancel. + return isHealthy() ? TransferCancelResult::kCancelRequestedInFlight + : TransferCancelResult::kBackendUnhealthy; } - return TransferCancelResult::kNotCancellable; + return TransferCancelResult::kNotFound; } bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index cf02471ff473..9b5aa502b2b3 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -46,7 +46,9 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver { public: // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors - NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); + // 7 NB_OVERRIDE_PURE entries below; the historical value was 6 which + // under-allocates the dispatch table (latent pre-existing bug). + NB_TRAMPOLINE(tb::BaseCacheTransceiver, 7); void respondAndSendAsync(std::shared_ptr llmRequest) override { diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index fae2e7d3a365..8d0e11f82349 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -347,11 +347,11 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- TLLM_CUDA_CHECK(cudaMemset(it->data(), llmRequest->getPromptLen(), it->getSizeInBytes())); } } - mFutures.emplace_back(mSender->sendAsync(*llmRequest)); + mFutures.emplace_back(mSender->sendAsync(llmRequest)); } else { - auto future = mRequester->receiveAsync(*llmRequest); + auto future = mRequester->receiveAsync(llmRequest); future.get(); TLLM_CUDA_CHECK(cudaDeviceSynchronize()); auto blockRange = BlockRange::fromAllBlockIds(*mManager, llmRequest->mRequestId); @@ -970,7 +970,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsendAsync(*llmRequest); + auto future = mSender->sendAsync(llmRequest); return future; } @@ -980,7 +980,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParammLlmRequest; mManager->addSequence(llmRequest->mRequestId, llmRequest->getNumTokens(beamIdx), beamWidth, llmRequest); - return mRequester->receiveAsync(*llmRequest); + return mRequester->receiveAsync(llmRequest); } void generationVerifyKVCache(std::shared_ptr const& request) From ec5dfffb4742794bde0c7053c3380098e18f6159 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 6 May 2026 10:17:13 -0700 Subject: [PATCH 5/5] [https://nvbugs/6104831][fix] Close timeout recovery loop: deferred KV release + health-driven shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The earlier commits surfaced the user-visible "exceeded timeout" error and pinned the LlmRequest object via shared_ptr, but the recovery story still had two gaps that an internal review caught: 1. When the C++ deadline-quarantine path marked a request as DISAGG_TRANS_ERROR, _can_terminate_request_now no longer recognised the request as "still in flight" (its state was no longer DISAGG_*_TRANS_IN_PROGRESS), so _do_terminate_request happily called free_resources() — releasing KV blocks back to the cache pool while the C++ worker thread might still be writing into them via NIXL/UCX. The shared_ptr lifetime fix protected the LlmRequest object but not its KV memory. 2. C++ exposed isHealthy() / getHealth() but nothing on the Python side consumed those signals. A worker whose backend was genuinely wedged would silently accumulate quarantined transfers forever, never triggering an orchestration restart. This commit closes both: * PyExecutor adds _pending_resource_release (a list of LlmRequest whose KV cleanup is deferred) and _maybe_release_pending_resources which polls cancel_request_structured each iteration. free_resources runs only when the structured cancel result transitions to AlreadyComplete / NotFound / CancelledBeforeAdvertise — the three outcomes that prove the worker has reached a final state. * The timeout error paths in _check_disagg_ctx_cache_transfer_status and _handle_responses (gen side) call _defer_resource_release_for_inflight_transfer() before _handle_errors / _end_transfer_and_maybe_terminate, so the user-visible error response is still surfaced immediately while the actual free_resources() call defers until C++ quiesces. * PyExecutor adds _check_transceiver_health which monitors transceiver.is_healthy(). It records the first-unhealthy timestamp and, after a grace period (default 2 * global_progress_deadline_- seconds, falling back to 120s), sets self.is_shutdown so orchestration's existing restart path takes the worker out of service. Recovery from a transient unhealthy window does not trigger restart. * Both new methods plus the existing _check_kv_transfer_timeout are called at every executor-loop tick that already checks for KV transfer timeouts — five sites across PP / non-overlap / overlap loops. The doc gains a "Recovery model" section that spells out the three timescales: user-visible error (immediate), KV resource release (bounded by worker quiescence), worker restart on persistent wedge (global deadline + grace). Tests: * Five new unit tests in test_kv_transfer_session_lifecycle.py: - test_defer_resource_release_holds_until_quiesced (full sequence timeout → in-flight → still in-flight → quiesced → freed). - test_defer_resource_release_handles_not_found_as_safe. - test_defer_resource_release_no_op_when_empty. - test_check_transceiver_health_resets_on_recovery (transient unhealthy must not shut down). - test_check_transceiver_health_triggers_shutdown_after_grace (sustained unhealthy DOES shut down). * All 19 tests pass without GPU/MPI. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../disagg-kv-transfer-session-lifecycle.md | 15 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 216 +++++++++++++++--- .../test_kv_transfer_session_lifecycle.py | 158 +++++++++++++ 3 files changed, 355 insertions(+), 34 deletions(-) diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index c072d663d754..ec5faf9dfda6 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -112,6 +112,21 @@ locally. `seconds_since_last_progress`, and `global_progress_deadline_seconds` for metrics / observability. +## Recovery model + +When a per-request KV transfer exceeds `kvTransferTimeoutMs`, three +things must happen, on three different timescales: + +| Concern | Path | Timing | +|---|---|---| +| **User-visible error response** | C++ `maybeQuarantineLocked` adds the request to `errorRequestIds`; Python's `_check_disagg_ctx_cache_transfer_status` / `_handle_responses` surfaces the error via `_handle_errors` / `_end_transfer_and_maybe_terminate`. | Immediate (next polling iteration). | +| **KV / buffer resource release** | Python registers the request in `_pending_resource_release` and `_inflight_cancel_requested_ids`. Each iteration `_maybe_release_pending_resources` polls `cancelRequestStructured`; only when it returns `AlreadyComplete` / `NotFound` / `CancelledBeforeAdvertise` does Python actually call `free_resources()`. The C++ shared_ptr in `mSenderFutures` keeps the LlmRequest alive until then. | Bounded by worker quiescence — could be the rest of `kvTransferTimeoutMs`, or longer if the backend is genuinely wedged. | +| **Worker restart on persistent wedge** | C++ `getHealth().is_healthy` flips false when the quarantine budget overflows or the global progress deadline elapses. Python's `_check_transceiver_health` watches that flag and sets `is_shutdown = True` after a grace period (default `2 * mGlobalProgressDeadlineMs`). Orchestration's existing process-restart path / k8s liveness probe / `disagg_auto_scaling.py` then restarts the worker. | Bounded by the global progress deadline + grace period (default 60 s + 120 s). | + +This separation is what makes the recovery deterministic without +either freeing memory the transport may still write into or hanging +the user's request indefinitely. + ## Python policy The `PyExecutor` keeps these responsibilities only: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8d6eb1304334..e7a84ddd61cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -429,6 +429,20 @@ def __init__( # writing into them. Keyed by py_request_id to survive request # state transitions. self._inflight_cancel_requested_ids: set[int] = set() + # Requests whose timeout error response has already been surfaced + # to the user but whose KV / buffer resources cannot yet be + # released because the C++ worker has not quiesced. The request + # is no longer in active_requests, so the executor does not see + # it again — _maybe_release_pending_resources is the only place + # that polls C++ and triggers the deferred free once the worker + # reaches a final state. + self._pending_resource_release: List[LlmRequest] = [] + # Wall-clock deadline tracker for the transceiver-unhealthy + # signal. None while the transceiver is healthy; set when + # transceiver.is_healthy() first flips false. If it stays false + # longer than the configured grace period we set is_shutdown + # so orchestration restarts the worker. + self._transceiver_unhealthy_since: Optional[float] = None # TODO: Remove the condition on the PP size once disagg support from KVCache reuse # path is fixed. self.async_transfer_manager = AsyncTransferManager( @@ -683,6 +697,100 @@ def _flush_pending_transfer_responses(self): if responses or self.enable_attention_dp: self._enqueue_responses(responses) + def _defer_resource_release_for_inflight_transfer( + self, request: LlmRequest) -> None: + """Mark a timed-out transfer's resources for deferred release. + + Called from the timeout-error paths when C++ structured cancel + reports the worker is still in flight. The request has already + been removed from active_requests by _handle_errors and the user + has been told the request errored; we still must not call + free_resources() until the worker quiesces, otherwise NIXL/UCX + could write into KV blocks that have been returned to the pool. + """ + req_id = request.py_request_id + self._inflight_cancel_requested_ids.add(req_id) + if request not in self._pending_resource_release: + self._pending_resource_release.append(request) + + def _maybe_release_pending_resources(self): + """Drain _pending_resource_release for requests whose C++ worker + has finally reached a final state. + + Each iteration we ask the transceiver whether each pending + request is still in flight. AlreadyComplete / NotFound mean the + worker has quiesced (future was either ready or already erased + by check_*_transfer_status), so it is now safe to free + request-owned KV blocks. Any other result keeps the request on + the deferred list — we will retry next iteration. + """ + if (not self._pending_resource_release + or self.kv_cache_transceiver is None): + return + still_pending: List[LlmRequest] = [] + for request in self._pending_resource_release: + cancel = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel in (TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound, + TransferCancelResult.CancelledBeforeAdvertise): + req_id = request.py_request_id + self._inflight_cancel_requested_ids.discard(req_id) + # _do_terminate_request now passes the guard because the + # request is no longer in _inflight_cancel_requested_ids + # and its state is no longer DISAGG_*_TRANS_IN_PROGRESS. + self._do_terminate_request(request) + else: + still_pending.append(request) + self._pending_resource_release = still_pending + + def _check_transceiver_health(self): + """Watch the C++ transceiver's health signal and trigger an + executor shutdown if it stays unhealthy past the grace period. + + The C++ side flips `is_healthy() == False` only after the + quarantine budget has been exceeded or no worker has reached a + final state for longer than mGlobalProgressDeadlineMs. That is + already a strong signal; we wait an additional grace period + (default: 2 * deadline) before forcing shutdown so that brief + transient unhealthy windows during heavy cancellation traffic do + not cause a restart loop. + """ + if self.kv_cache_transceiver is None: + return + if self.kv_cache_transceiver.is_healthy(): + self._transceiver_unhealthy_since = None + return + now = time.time() + if self._transceiver_unhealthy_since is None: + self._transceiver_unhealthy_since = now + health = self.kv_cache_transceiver.get_health() + logger.warning( + f"KV cache transceiver flipped UNHEALTHY at {now:.3f}: " + f"quarantined={getattr(health, 'quarantined_transfer_count', '?')} " + f"budget={getattr(health, 'quarantine_budget', '?')} " + f"seconds_since_last_progress=" + f"{getattr(health, 'seconds_since_last_progress', '?')}. " + f"Will trigger executor shutdown if it stays unhealthy.") + return + elapsed = now - self._transceiver_unhealthy_since + # Grace period: 2x the global progress deadline reported by C++. + # Fall back to 120s if the snapshot is unavailable. + try: + health = self.kv_cache_transceiver.get_health() + grace_period = 2.0 * float( + getattr(health, 'global_progress_deadline_seconds', 60.0)) + except Exception: + grace_period = 120.0 + if elapsed > grace_period and not self.is_shutdown: + logger.error( + f"KV cache transceiver has been UNHEALTHY for {elapsed:.1f}s " + f"(grace period {grace_period:.1f}s). Setting is_shutdown so " + f"orchestration can restart the worker.") + self.is_shutdown = True + if self.shutdown_event is not None: + self.shutdown_event.set() + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1367,6 +1475,8 @@ def _pp_retry_until_can_schedule(self, scheduled_batch): # Let cache transceiver finish at least one cache transmission and release requests' KV cache resources self._check_disagg_ctx_cache_transfer_status(1) self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() else: raise RuntimeError( f"Reach maximum PP retry count ({self.pp_scheduler_max_retry_count}) but still cannot run first PP's schedule result. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value. Or you can set `TLLM_PP_SCHEDULER_MAX_RETRY_COUNT` to a larger value to allow more retries." @@ -1759,6 +1869,8 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() if self._disagg_pp_termination_handler is not None: self._disagg_pp_termination_handler.terminate_pending_requests() @@ -1867,6 +1979,8 @@ def _prepare_and_schedule_batch(self): self._check_disagg_ctx_schedulable_status(new_requests) self._check_disagg_gen_transfer_status() self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() iter_stats = None if self.enable_iter_perf_stats: @@ -2236,6 +2350,8 @@ def _executor_loop(self): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() self._kv_connector_terminate_requests() @@ -2529,6 +2645,8 @@ def _executor_loop_overlap(self): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() self._kv_connector_terminate_requests() @@ -3374,10 +3492,12 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) - # Worker reached a final state — clear any pending in-flight cancel - # bookkeeping for those requests. - if completed_req_ids: - self._inflight_cancel_requested_ids -= completed_req_ids + # Worker reached a final state for entries that came back as + # completed (future ready). Quarantined-but-error entries have + # NOT actually quiesced; we keep their inflight-cancel bookkeeping. + finished_set = set(finished_requests) + if finished_set: + self._inflight_cancel_requested_ids -= finished_set for request_id in completed_req_ids: @@ -3388,6 +3508,23 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request = requests_in_transfer[request_id] + # If the request errored because of timeout / quarantine + # rather than worker completion, the C++ worker may still be + # in flight (the future is pinned in mSenderFutures). Defer + # the free_resources() call by adding the request to + # _pending_resource_release before _end_transfer_and_maybe_- + # terminate hits _do_terminate_request. The deferred entry + # will be drained by _maybe_release_pending_resources once + # the structured-cancel result transitions to a safe state. + if request_id in set(error_requests): + cancel = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + self._defer_resource_release_for_inflight_transfer( + request) + self._end_transfer_and_maybe_terminate(request) # The set of requests in transfer may have changed since we terminated some requests. @@ -3417,22 +3554,26 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE self._inflight_cancel_requested_ids.discard(request_id) self._end_transfer_and_maybe_terminate(request) - elif cancel_result == TransferCancelResult.CancelRequestedInFlight: - # Cancellation accepted but the worker is still - # touching the buffer/KV memory. Mark the request as - # in-flight-cancel so _do_terminate_request defers - # freeing resources, and re-check next iteration. - self._inflight_cancel_requested_ids.add(request_id) - elif cancel_result == TransferCancelResult.BackendUnhealthy: - # Soft signal — the transceiver is unhealthy. We do - # not touch this request; orchestration health checks - # will restart the worker once the global deadline - # elapses. - logger.warning( - f"Transceiver reports unhealthy while cancelling " - f"context request {request_id}; deferring to " - f"health-driven restart.") - self._inflight_cancel_requested_ids.add(request_id) + elif cancel_result in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + # Worker is still touching buffers (or transceiver is + # unhealthy). Defer the actual free_resources() to + # _maybe_release_pending_resources by registering the + # request BEFORE the cleanup path. The user-visible + # error response is still surfaced via the normal + # transfer-error flow on a subsequent iteration once + # C++ confirms quiescence. + request.py_kv_transfer_start_time = None + request.state = LlmRequestState.DISAGG_TRANS_ERROR + self._defer_resource_release_for_inflight_transfer( + request) + if cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"context request {request_id}; deferring resource " + f"release to health-driven restart.") + self._end_transfer_and_maybe_terminate(request) self._check_cache_transfer_errors("context requests") @@ -3861,20 +4002,27 @@ def _handle_responses(self): self._handle_errors( error_msg=f"Request {request.py_request_id} timed out", requests=[request]) - elif cancel_result == TransferCancelResult.CancelRequestedInFlight: - # Worker still touching buffers — keep the request - # active so we re-poll next iteration. Once C++ - # transitions the worker future to a final state we - # will see NotFound / AlreadyComplete and clean up. - self._inflight_cancel_requested_ids.add(req_id) - new_active_requests.append(request) - elif cancel_result == TransferCancelResult.BackendUnhealthy: - logger.warning( - f"Transceiver reports unhealthy while cancelling " - f"generation request {req_id}; deferring to " - f"health-driven restart.") - self._inflight_cancel_requested_ids.add(req_id) - new_active_requests.append(request) + elif cancel_result in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + # Worker still touching buffers (or transceiver + # unhealthy). Surface the timeout error to the user + # immediately, but defer free_resources() to + # _maybe_release_pending_resources, which polls + # cancel_request_structured each iteration and frees + # only once the worker quiesces. We register the + # request BEFORE _handle_errors so _do_terminate_- + # request's guard sees req_id in the inflight set. + self._defer_resource_release_for_inflight_transfer( + request) + if cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"generation request {req_id}; deferring resource " + f"release to health-driven restart.") + self._handle_errors( + error_msg=f"Request {request.py_request_id} timed out", + requests=[request]) continue if request.is_generation_only_request() and not request.is_finished: diff --git a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py index 44a4f41d571d..3d81768d2f15 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py +++ b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py @@ -61,6 +61,15 @@ def is_disagg_generation_transmission_in_progress(self) -> bool: return self.state == "DISAGG_GENERATION_TRANS_IN_PROGRESS" +@dataclass +class _FakeHealth: + is_healthy: bool = True + quarantined_transfer_count: int = 0 + quarantine_budget: int = 16 + seconds_since_last_progress: float = 0.0 + global_progress_deadline_seconds: float = 60.0 + + @dataclass class _FakeTransceiver: """Programmable stand-in for the C++-backed KvCacheTransceiver. The @@ -69,6 +78,7 @@ class _FakeTransceiver: next_cancel_results: dict = field(default_factory=dict) cancel_calls: List[int] = field(default_factory=list) healthy: bool = True + health: _FakeHealth = field(default_factory=_FakeHealth) def cancel_request_structured(self, req: _FakeRequest): self.cancel_calls.append(req.py_request_id) @@ -78,6 +88,9 @@ def cancel_request_structured(self, req: _FakeRequest): def is_healthy(self) -> bool: return self.healthy + def get_health(self) -> _FakeHealth: + return self.health + def _make_executor_under_test(): """Build a minimal object whose interface matches what the helper @@ -88,6 +101,10 @@ def _make_executor_under_test(): obj.kv_cache_transceiver = _FakeTransceiver() obj._inflight_cancel_requested_ids = set() obj._pending_transfer_responses = [] + obj._pending_resource_release = [] + obj._transceiver_unhealthy_since = None + obj.is_shutdown = False + obj.shutdown_event = None obj.enable_attention_dp = False obj.active_requests = [] obj.terminated_request_ids = [] @@ -125,12 +142,58 @@ def _flush_pending_transfer_responses(self): if responses or self.enable_attention_dp: self._enqueue_responses(responses) + def _defer_resource_release_for_inflight_transfer(self, request): + self._inflight_cancel_requested_ids.add(request.py_request_id) + if request not in self._pending_resource_release: + self._pending_resource_release.append(request) + + def _maybe_release_pending_resources(self): + if not self._pending_resource_release or self.kv_cache_transceiver is None: + return + still_pending = [] + safe = ( + _FakeTransferCancelResult.AlreadyComplete, + _FakeTransferCancelResult.NotFound, + _FakeTransferCancelResult.CancelledBeforeAdvertise, + ) + for request in self._pending_resource_release: + cancel = self.kv_cache_transceiver.cancel_request_structured(request) + if cancel in safe: + self._inflight_cancel_requested_ids.discard(request.py_request_id) + self._do_terminate_request(request) + else: + still_pending.append(request) + self._pending_resource_release = still_pending + + def _check_transceiver_health(self): + if self.kv_cache_transceiver is None: + return + if self.kv_cache_transceiver.is_healthy(): + self._transceiver_unhealthy_since = None + return + import time as _time + now = _time.time() + if self._transceiver_unhealthy_since is None: + self._transceiver_unhealthy_since = now + return + elapsed = now - self._transceiver_unhealthy_since + health = self.kv_cache_transceiver.get_health() + grace = 2.0 * float(getattr(health, "global_progress_deadline_seconds", 60.0)) + if elapsed > grace and not self.is_shutdown: + self.is_shutdown = True + obj._can_terminate_request_now = types.MethodType( _can_terminate_request_now, obj) obj._do_terminate_request = types.MethodType(_do_terminate_request, obj) obj._enqueue_responses = types.MethodType(_enqueue_responses, obj) obj._flush_pending_transfer_responses = types.MethodType( _flush_pending_transfer_responses, obj) + obj._defer_resource_release_for_inflight_transfer = types.MethodType( + _defer_resource_release_for_inflight_transfer, obj) + obj._maybe_release_pending_resources = types.MethodType( + _maybe_release_pending_resources, obj) + obj._check_transceiver_health = types.MethodType( + _check_transceiver_health, obj) return obj @@ -346,3 +409,98 @@ def test_unhealthy_transceiver_keeps_request_pinned(): assert 33 in executor._inflight_cancel_requested_ids executor._do_terminate_request(request) assert executor.terminated_request_ids == [] + + +# --------------------------------------------------------------------------- +# Deferred-resource-release polling +# --------------------------------------------------------------------------- + + +def test_defer_resource_release_holds_until_quiesced(): + """Replay the full recovery sequence: timeout fires, error response + is surfaced, KV cleanup is deferred. Subsequent iterations poll + cancel_request_structured; when it transitions to AlreadyComplete / + NotFound, the deferred terminate runs.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=77, state="DISAGG_TRANS_ERROR") + executor._defer_resource_release_for_inflight_transfer(request) + + # Iteration 1: worker still in flight — must NOT free. + executor.kv_cache_transceiver.next_cancel_results[77] = ( + _FakeTransferCancelResult.CancelRequestedInFlight) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + assert request in executor._pending_resource_release + assert 77 in executor._inflight_cancel_requested_ids + + # Iteration 2: still in flight. + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + + # Iteration 3: worker finally quiesced. + executor.kv_cache_transceiver.next_cancel_results[77] = ( + _FakeTransferCancelResult.AlreadyComplete) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [77] + assert request not in executor._pending_resource_release + assert 77 not in executor._inflight_cancel_requested_ids + + +def test_defer_resource_release_handles_not_found_as_safe(): + """C++ may erase the worker entry between iterations; NotFound at + that point means the worker is gone and resources are safe to + free.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=88, state="DISAGG_TRANS_ERROR") + executor._defer_resource_release_for_inflight_transfer(request) + + executor.kv_cache_transceiver.next_cancel_results[88] = ( + _FakeTransferCancelResult.NotFound) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [88] + + +def test_defer_resource_release_no_op_when_empty(): + executor = _make_executor_under_test() + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + + +# --------------------------------------------------------------------------- +# Health-driven shutdown escape +# --------------------------------------------------------------------------- + + +def test_check_transceiver_health_resets_on_recovery(): + executor = _make_executor_under_test() + executor.kv_cache_transceiver.healthy = False + executor._check_transceiver_health() + assert executor._transceiver_unhealthy_since is not None + + # Recovery before grace period elapses — clear the timestamp, + # do not shut down. + executor.kv_cache_transceiver.healthy = True + executor._check_transceiver_health() + assert executor._transceiver_unhealthy_since is None + assert not executor.is_shutdown + + +def test_check_transceiver_health_triggers_shutdown_after_grace(): + executor = _make_executor_under_test() + # Set an extremely short deadline so the grace period (2x) is + # easy to exceed in a unit test. + executor.kv_cache_transceiver.health = _FakeHealth( + is_healthy=False, + global_progress_deadline_seconds=0.01, + ) + executor.kv_cache_transceiver.healthy = False + + # First call records the timestamp, does not shut down yet. + executor._check_transceiver_health() + assert not executor.is_shutdown + assert executor._transceiver_unhealthy_since is not None + + # Move the recorded timestamp into the past beyond the grace period. + executor._transceiver_unhealthy_since = executor._transceiver_unhealthy_since - 1.0 + executor._check_transceiver_health() + assert executor.is_shutdown