diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 8f8330603893..2b8c797565ac 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -200,28 +200,139 @@ 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: 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; - /// 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; - virtual bool cancelRequest(LlmRequest* llmRequest) = 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(std::shared_ptr 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(std::shared_ptr 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 @@ -252,34 +363,111 @@ 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; 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; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + TransferCancelResult cancelRequestStructured(std::shared_ptr llmRequest) override; + + virtual bool cancelRequest(std::shared_ptr 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. 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. + /// + /// The per-request KV transfer timeout is anchored to + /// @c LlmRequest::getKvCacheActualTransferStart(), recorded by the + /// formatters once a transfer-buffer slot has been acquired. There + /// is intentionally no captured deadline on this struct: a + /// deeply-queued request without an actual-transfer-start time + /// cannot be quarantined, which prevents queue-starvation false + /// positives. + struct TrackedFuture + { + std::shared_ptr request; + std::future future; + /// 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); + 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/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index c1ecffff026b..3be16c9e608f 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -1797,6 +1797,46 @@ class GenericLlmRequest .count()); } + /// @brief Mark when the actual KV transfer (post slot acquisition, + /// post pre-data handshake) began for this request. + /// + /// This is distinct from @ref setKvCacheTransferStart: that one is + /// called at admit / worker-dequeue time and feeds the + /// `kv_cache_transfer_time_ms` perf metric (which has historically + /// included queue + slot wait). The "actual transfer start" is + /// recorded by the formatters after `assignBufferIndexFor*` returns, + /// when the worker is about to begin the data movement, and is the + /// correct anchor for the `kvTransferTimeoutMs` deadline used by + /// CacheTransceiver's quarantine logic. + /// + /// Idempotent: subsequent calls within the same transfer (e.g., for + /// MLA / hybrid models that acquire multiple pool slots) are no-ops + /// once the field is set. + void setKvCacheActualTransferStart(TimePoint time) const + { + if (mKvCacheActualTransferStart == TimePoint{}) + { + mKvCacheActualTransferStart = maybeToGlobalSteadyClock(time); + } + } + + [[nodiscard]] TimePoint getKvCacheActualTransferStart() const noexcept + { + return mKvCacheActualTransferStart; + } + + [[nodiscard]] bool hasKvCacheActualTransferStart() const noexcept + { + return mKvCacheActualTransferStart != TimePoint{}; + } + + /// Reset before re-issuing a transfer (e.g., on retry). Not used in + /// PR #13796's flow but kept symmetric with the perf-metric API. + void clearKvCacheActualTransferStart() const noexcept + { + mKvCacheActualTransferStart = TimePoint{}; + } + void updateKvCacheSize(size_t targetBufferSize) const { mPerfMetrics.timingMetrics.kvCacheSize += targetBufferSize; @@ -2118,6 +2158,17 @@ class GenericLlmRequest bool mReturnPerfMetrics{false}; mutable executor::RequestPerfMetrics mPerfMetrics; + /// Time at which the disagg KV cache transfer actually started for + /// this request — recorded by the formatters once a transfer-buffer + /// slot has been acquired (i.e., once the worker is about to begin + /// data movement). Distinct from + /// @c mPerfMetrics.timingMetrics.kvCacheTransferStart, which is set + /// at admit / worker-dequeue time and thus includes queue / slot + /// wait. Used by @c CacheTransceiver::maybeQuarantineLocked as the + /// anchor for @c kvTransferTimeoutMs so that a deeply-queued + /// request is not quarantined for queue starvation. + mutable TimePoint mKvCacheActualTransferStart{}; + // Guided decoding params. std::optional mGuidedDecodingParams{std::nullopt}; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 58092897ebbe..48e5d7149ca8 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" +#include #include namespace tensorrt_llm::batch_manager @@ -233,8 +234,41 @@ std::optional BaseTransBufferManager::assignBufferIndex( return std::nullopt; } std::unique_lock lk(resource.mBuffersMutex); - resource.mBuffersCV.wait( - lk, [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }); + + // [wedge-trace] convert the unbounded CV wait into wait_for(interval) + // so we can periodically log when a thread has been blocked waiting + // for a slot. Slot exhaustion at this layer is the downstream + // symptom of a stuck transfer (slot held by a wedged worker thread) + // — see docs/source/features/disagg-kv-transfer-debug-stuck-slot.md. + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + auto const haveSlot + = [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }; + if (heartbeatIntervalMs > 0) + { + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + while (!haveSlot()) + { + if (resource.mBuffersCV.wait_for(lk, std::chrono::milliseconds(heartbeatIntervalMs), haveSlot)) + { + break; + } + auto const now = std::chrono::steady_clock::now(); + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot " + "for %lld ms (concurrence=%d budget=%zu). The slot is most likely held by a wedged " + "worker thread; check sender/receiver heartbeats and run the stuck-slot diagnosis.", + static_cast(elapsedMs), static_cast(resource.mConcurrence), bufferCount); + lastHeartbeat = now; + (void) lastHeartbeat; + } + } + else + { + resource.mBuffersCV.wait(lk, haveSlot); + } int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index e05e8d6f76fc..6af1242c8743 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -495,6 +495,12 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + // The slot is now reserved and the worker is about to begin + // actual data movement. Anchor the per-request KV transfer + // timeout from this point so queue / slot wait does not + // consume the deadline budget. See + // CacheTransceiver::maybeQuarantineLocked. + session.getLlmRequest().setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -815,6 +821,12 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); } + // Slot reserved (dynamically or pre-assigned). Anchor + // the per-request KV transfer timeout from this point — + // before this, the worker was still queued or waiting + // for a slot, neither of which is "in transfer." See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..999757df6ccb 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,9 +337,13 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); - auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + setContextState(llmRequest.get()); + auto future = mCacheSender->sendAsync(llmRequest); + // Note: no admit-time deadline is captured. The per-request KV + // transfer timeout (kvTransferTimeoutMs) is anchored to + // LlmRequest::getKvCacheActualTransferStart(), set by the + // formatters once the worker has acquired a transfer-buffer slot. + mSenderFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -354,36 +358,64 @@ 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 future = mCacheSender->sendAsync(llmRequest); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), false, false}); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); - { - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto const reqId = llmRequest->mRequestId; + { + auto future = mCacheReceiver->receiveAsync(llmRequest); + + // [wedge-trace] turn the previously unbounded future.get() into a + // wait_for(interval) loop so the executor thread can log a + // heartbeat when this synchronous gen-side path + // (TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1) is blocked. The + // worker thread's own heartbeats will tell us whether the wedge + // is in NIXL/UCX or somewhere downstream. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + if (heartbeatIntervalMs > 0) + { + auto const startTime = std::chrono::steady_clock::now(); + while (future.wait_for(std::chrono::milliseconds(heartbeatIntervalMs)) != std::future_status::ready) + { + auto const elapsedMs = std::chrono::duration_cast( + std::chrono::steady_clock::now() - startTime) + .count(); + TLLM_LOG_WARNING( + "[wedge-trace] CacheTransceiver::requestAndReceiveSync executor thread blocked on " + "future.get for request %ld for %lld ms (TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP path). " + "Worker thread should be emitting its own [wedge-trace] heartbeat for the actual " + "wedge location.", + reqId, static_cast(elapsedMs)); + } + } future.get(); } 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& pair) { return pair.first->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; } - auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); + auto future = mCacheReceiver->receiveAsync(llmRequest); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + // Note: no admit-time deadline; see respondAndSendAsync. + mRequesterFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), false, false}); } std::vector gatherRequestIds( @@ -480,29 +512,130 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } } -RequestStatuses CacheTransceiver::checkContextTransferStatus( - std::optional const& atLeastRequestNum, bool markComplete) +// ----- Internal helpers ---------------------------------------------------- + +void CacheTransceiver::updateHealthLocked() { - 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()) + 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()); { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + 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) +{ + if (entry.quarantined) + { + return false; + } + if (!mCacheTransceiverConfig.has_value()) + { + return false; + } + auto const timeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + if (!timeoutMs.has_value()) + { + return false; + } + // Defer the deadline check until the worker has actually started + // the transfer (post slot acquisition). Requests still sitting in + // the worker queue or waiting on a slot are not "stuck" in any + // transport-relevant sense; quarantining them for queue starvation + // turned every burst into a false-positive cascade — see + // claude-swe/disagg-deadline-at-transfer-start-fix.md and the + // d224/d227 findings (TCP at ~3 Gbps, 1-slot serialization, conc + // 128) where this mis-accounting caused ~80% of requests to be + // quarantined for queue wait rather than for any actual NIXL/UCX + // wedge. + if (!entry.request->hasKvCacheActualTransferStart()) + { + return false; + } + auto const transferStart = entry.request->getKvCacheActualTransferStart(); + auto const deadline = transferStart + std::chrono::milliseconds(timeoutMs.value()); + auto const now = std::chrono::steady_clock::now(); + if (now < 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 after slot acquisition " + "(actual transfer time > %lld ms). Future stays pinned in C++ until worker reaches a final state.", + entry.request->mRequestId, static_cast(timeoutMs.value())); + 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 +650,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) + entry.future.get(); + requestsStatus.completedRequestIds.insert(entry.request->mRequestId); + if (markComplete) { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; - } - else - { - 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 +810,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()) + if (freq == expectedFreq) { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); - } - else - { - 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.get()); } } 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.get()); + } } - 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 +909,79 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +TransferCancelResult CacheTransceiver::cancelRequestStructured(std::shared_ptr llmRequest) { - if (llmRequest->isContextOnlyRequest()) + TLLM_CHECK(llmRequest != nullptr); + + auto const reqId = llmRequest->mRequestId; + bool const sender = llmRequest->isContextOnlyRequest(); + bool const receiver = llmRequest->isGenerationOnlyRequest(); + if (!sender && !receiver) { - return mCacheSender->cancelRequest(*llmRequest); + return TransferCancelResult::kNotCancellable; } - else if (llmRequest->isGenerationOnlyRequest()) + + // 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)) { - return mCacheReceiver->cancelRequest(*llmRequest); + return TransferCancelResult::kCancelledBeforeAdvertise; } - return false; + if (receiver && 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) + { + continue; + } + auto status = futures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) + { + return TransferCancelResult::kAlreadyComplete; + } + // 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::kNotFound; +} + +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) +{ + auto result = cancelRequestStructured(std::move(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/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/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..429def3f9dcf 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -254,6 +254,11 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + // MLA indexer-K-cache slot reserved. Idempotent — if the + // primary KV slot was already acquired earlier in + // CacheFormatter::format, this call is a no-op. See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -495,6 +500,9 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); } + // MLA indexer-K-cache slot reserved on the receive side. + // Idempotent (see CacheFormatter::unformat). + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto targetNum = pickUpConnections.size(); diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 1fd1cbdc253f..54a80832ae08 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -149,6 +149,12 @@ void RnnCacheFormatter::format(TransferSession& session) } auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + // RNN-state slot reserved on the send side. Anchor the per-request + // KV transfer timeout from this point. Idempotent — for hybrid + // models that also acquire a KV slot in CacheFormatter::format, + // whichever fires first wins. See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); @@ -323,6 +329,9 @@ void RnnCacheFormatter::unformat(TransferSession& session) { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); } + // RNN-state slot reserved on the recv side. Idempotent (see send + // side / CacheFormatter::unformat). + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(sourceNum), bufferSizesPerSource, bufferManager); 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/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index bf3142a160dd..9a0117879ac7 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -454,6 +454,13 @@ std::string const& getEnvKVCacheTimeOutputPath() return outputPath; } +size_t getEnvDisaggWedgeTraceIntervalMs() +{ + // Default: 5 s. Set TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=0 to disable. + static size_t const intervalMs = getUInt64Env("TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS").value_or(5000); + return intervalMs; +} + bool getEnvKVCacheTransferUseAsyncBuffer() { diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 82c56f300362..546bd9e363f6 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,6 +115,13 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +/// Periodic [wedge-trace] heartbeat interval (ms) for the disaggregated +/// KV transfer wedge points (NixlTransferStatus::wait, +/// waitForNotification, assignBufferIndex CV wait, +/// requestAndReceiveSync future.get). When set to 0 the heartbeats are +/// disabled. Default: 5000 ms. +size_t getEnvDisaggWedgeTraceIntervalMs(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index d46defdf50ad..e1403dac4f72 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -18,6 +18,7 @@ #include "connection.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" +#include #include #include #include @@ -585,6 +586,12 @@ template void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + char const* notificationKindName + = std::is_same_v ? "NotificationSyncInfo" : "ReadySignalInfo"; + while (!terminateFlag.load()) { @@ -593,6 +600,30 @@ void AgentConnectionManager::waitForNotification( return; } updateUnhandledNotifications(); + + // [wedge-trace] periodic heartbeat for the receiver-side + // getNotifs poll loop. Fires when no matching notification has + // arrived for more than the configured interval. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + if (heartbeatIntervalMs > 0) + { + auto const now = std::chrono::steady_clock::now(); + auto const sinceHeartbeatMs + = std::chrono::duration_cast(now - lastHeartbeat).count(); + if (sinceHeartbeatMs >= static_cast(heartbeatIntervalMs)) + { + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] AgentConnectionManager::waitForNotification still pending after %lld ms " + "(remoteAgent='%s' expectedTag=%d kind=%s). Receiver-side getNotifs poll is not making " + "progress; sender may have crashed or its notification was lost.", + static_cast(elapsedMs), remoteAgentName.c_str(), + static_cast(expectedInfo.mContext.getTag()), notificationKindName); + lastHeartbeat = now; + } + } + std::scoped_lock lock(mNotificationMutex); auto it = mUnhandledNotifications.begin(); while (it != mUnhandledNotifications.end()) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index bad3e184f983..99fdf55d38d9 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -484,7 +484,9 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { - auto startTime = std::chrono::steady_clock::now(); + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); while (true) { @@ -498,6 +500,29 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const return TransferState::kFAILURE; } + // [wedge-trace] periodic heartbeat for the sender-side + // getXferStatus poll loop. Fires when the same nixlXferReqH has + // been NIXL_IN_PROG for more than the configured interval. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + if (heartbeatIntervalMs > 0) + { + auto const now = std::chrono::steady_clock::now(); + auto const sinceHeartbeatMs + = std::chrono::duration_cast(now - lastHeartbeat).count(); + if (sinceHeartbeatMs >= static_cast(heartbeatIntervalMs)) + { + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after %lld ms " + "(handle=%p, requested timeout_ms=%lld). Sender-side getXferStatus poll " + "is not making progress; peer or UCX endpoint may be wedged.", + static_cast(elapsedMs), static_cast(mHandle), + static_cast(timeout_ms)); + lastHeartbeat = now; + } + } + // If timeout_ms < 0, wait indefinitely until status is not NIXL_IN_PROG if (timeout_ms < 0) { diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 61cac5df4c7e..9b5aa502b2b3 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" @@ -44,19 +46,29 @@ 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(tb::LlmRequest* llmRequest) override + void respondAndSendAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } - void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + 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(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); } @@ -77,7 +89,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); } @@ -86,31 +98,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/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) diff --git a/docs/source/features/disagg-kv-transfer-deadline-anchor.md b/docs/source/features/disagg-kv-transfer-deadline-anchor.md new file mode 100644 index 000000000000..96cb94a58bc0 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-deadline-anchor.md @@ -0,0 +1,204 @@ +# Anchoring `kv_transfer_timeout_ms` to actual transfer start + +> Companion to +> [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md). +> Documents the bookkeeping change introduced in PR #13830 (this PR), +> stacked on PRs #13796 and #13818. + +## Background — the bug we're fixing + +PR #13796 captured the per-request deadline at admit time: + +```cpp +// cacheTransceiver.cpp::requestAndReceiveAsync (pre-fix) +auto future = mCacheReceiver->receiveAsync(llmRequest); +auto const deadline = computeTrackedFutureDeadline( + /*requestStart=*/std::chrono::steady_clock::now()); +mRequesterFutures.push_back(TrackedFuture{..., deadline, ...}); +``` + +The actual transfer doesn't start until much later inside the worker +thread — it has to wait for the worker to dequeue the request, then +for `assignBufferIndexFor{Send,Recv}` to acquire a slot, then for +`session.{recv,send}` to exchange data with the peer. + +Under burst load with a small staged-buffer pool (default +`recvBufferCount=1` and `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM=1`), +the deadline burns down while the request is queued, not while it's +transferring. Empirically — see the d224/d227 testing under +`dynamo-gb200-test/docs/disagg-wedge-pr13818-findings-2026-05-06.md` +— a 60 s deadline on TCP at ~3 Gbps with a conc=128 burst caused +~80 % of requests to be quarantined for **queue starvation** rather +than for any actual NIXL/UCX wedge. They never even started +transferring; they just sat waiting for their turn. + +## What this PR changes + +`kv_transfer_timeout_ms` is now anchored to **actual transfer +start** — the moment the worker has acquired a slot and is about to +begin data movement. + +### New `LlmRequest` field + +```cpp +mutable TimePoint mKvCacheActualTransferStart{}; + +void setKvCacheActualTransferStart(TimePoint time) const; // idempotent +TimePoint getKvCacheActualTransferStart() const noexcept; +bool hasKvCacheActualTransferStart() const noexcept; +``` + +This is **distinct from** the existing +`mPerfMetrics.timingMetrics.kvCacheTransferStart` (set at +admit / worker-dequeue time and exposed as +`req.kv_cache_transfer_time_ms`). The perf metric semantics are +preserved — downstream consumers / dashboards continue to see the +same timing model they did before. + +### Formatters set the field after slot acquisition + +| Path | File:line | Slot | +|---|---|---| +| Sender, primary KV | [`cacheFormatter.cpp:497`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp#L497) | `assignBufferIndexForSend` | +| Receiver, primary KV | [`cacheFormatter.cpp:822`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp#L822) | `assignBufferIndexForRecv` (or pre-assigned id) | +| MLA indexer-K, sender | [`mlaCacheFormatter.cpp:256`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp#L256) | indexer-K send slot | +| MLA indexer-K, receiver | [`mlaCacheFormatter.cpp:496`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp#L496) | indexer-K recv slot | +| RNN/Mamba, sender | [`rnnCacheFormatter.cpp:151`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp#L151) | RNN-state send slot | +| RNN/Mamba, receiver | [`rnnCacheFormatter.cpp:324`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp#L324) | RNN-state recv slot | + +The setter is idempotent — for hybrid (KV + indexer-K, KV + RNN-state) +models that acquire multiple pool slots in sequence, the field +captures the **first** slot acquisition only. That's the point at +which the worker can be considered "in transfer." + +### `maybeQuarantineLocked` consults the new field + +```cpp +// cacheTransceiver.cpp::maybeQuarantineLocked (post-fix) +if (!entry.request->hasKvCacheActualTransferStart()) +{ + return false; // worker still queued / waiting for slot — no deadline +} +auto const transferStart = entry.request->getKvCacheActualTransferStart(); +auto const deadline = transferStart + std::chrono::milliseconds(timeoutMs.value()); +if (now < deadline) return false; +// quarantine +``` + +Requests that haven't yet reached the formatter's slot-acquire path +are not eligible for quarantine — they're not "stuck" in any +transport-relevant sense. + +### Dropped: admit-time deadline plumbing + +`TrackedFuture::deadline` is removed. `CacheTransceiver::computeTrackedFutureDeadline` +is removed. The three call sites in `respondAndSendAsync`, +`respondAndSendLayerWise`, and `requestAndReceiveAsync` no longer +capture admit-time deadlines. + +## Behavioural impact + +### Intended + +| Scenario | Before | After | +|---|---|---| +| TCP, conc=128, 60 s deadline, 1-slot serialization | ~80 % of burst quarantined for queue wait | only transfers whose actual transfer time exceeds 60 s are quarantined (rare on functional TCP) | +| Healthy transfer in queue for 30 s, transfers in 5 s | quarantined at T+60 s mid-transfer | not quarantined; user sees normal completion | +| Genuinely stuck mid-transfer (peer crash after slot acquired, NIXL hung) | quarantined at T + 60 s after admit | quarantined at T + 60 s after slot acquisition | +| Request never gets a slot (queue depth ≫ capacity, pool wedged elsewhere) | quarantined at admit + 60 s | not quarantined for queue wait. Orchestrator's `req_timeout_secs=180` is the user-visible backstop. | + +### Side effects + +- **Perf metric `kv_cache_transfer_time_ms`** is unchanged. Still + measured from admit to `kvCacheTransferEnd` via the existing + `mPerfMetrics.timingMetrics.kvCacheTransferStart` — this PR adds + a separate field, doesn't repurpose the existing one. +- **`_check_transceiver_health` (Python side)** is unchanged. It + flips `is_shutdown` on `is_healthy=false` after the grace period. + The trigger for `is_healthy=false` (no progress for > + `mGlobalProgressDeadlineMs` OR quarantine count > budget) is + unchanged at the C++ level. With this fix in place, false-positive + quarantines from queue starvation no longer drive the budget side + of that condition. +- **Wedge-trace heartbeats from PR #13818** are unchanged — they + fire inside the worker thread regardless of where the deadline + anchor is. +- **Structured cancel result enum / deferred resource release / + `_inflight_cancel_requested_ids`** are unchanged. + +### Risks and mitigations + +1. **Unbounded queue wait.** A request in the queue without a slot + never gets its deadline to start ticking. Mitigations: + - The disagg orchestrator's HTTP-level `req_timeout_secs=180` + fires a client-visible 504/408 at 180 s. + - The transceiver's no-progress detector (`mLastProgressTime` + vs `mGlobalProgressDeadlineMs`) flips `is_healthy=false` if + the queue head stops draining for > 60 s with at least one + tracked entry, triggering health-driven shutdown via the + PR #13796 mechanism. + - For deployments that want a hard ceiling on queue wait, a + follow-up could add a separate `kv_queue_wait_timeout_ms` + knob that fires at admit + N seconds if `hasKvCacheActualTransferStart()` + is still false. Not needed for the d224/d227 scenario. + +2. **Pre-existing behavior change for users who set + `kv_transfer_timeout_ms` thinking "end-to-end including queue."** + The new semantics are "actual transfer time only." For most + deployments this is the more useful interpretation. Users who + wanted the old admit-to-completion bound should consider using + the orchestrator's `req_timeout_secs` as a wrapper budget. + +3. **MLA / RNN / hybrid-pool ordering.** If a hybrid model acquires + the indexer-K slot before the primary KV slot, the deadline + anchors at the earlier acquisition. That's correct — it's the + first moment the worker is doing real transfer work. + +## Empirical validation plan + +Re-run the d224 / d227 burst suite: + +| Repro | Pre-fix | Post-fix expectation | +|---|---|---| +| d224 r1 (TCP 3 Gbps, conc=128, no canary) | ~80 % quarantined; `is_shutdown` at T+4.4 min | quarantines drop to near zero (only true mid-transfer stalls); `is_shutdown` does not fire; bursts drain over time | +| d227 r4 (Pattern C, canary on, peer crash) | both sides quarantine + canary-driven restart | mid-transfer wedges still get quarantined → still trigger restart path; verdict unchanged | +| Single-request smoke | unchanged | unchanged | + +If d224 r1 still reaches `is_shutdown` after this fix, look at the +no-progress side of `updateHealthLocked` — that's the next layer +of the diagnosis (covered separately in +`disagg-deadline-vs-quarantine-proposal.md`). + +## Test plan + +The existing 19 Python unit tests in +`tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py` +continue to pass without modification — they exercise the Python +policy layer, which is unchanged by this PR. + +C++-side coverage is via the existing multi-GPU integration tests in +`cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp`. No +adjustments needed; that suite uses the public API and the change +is internal. + +For empirical validation of the deadline-anchor behaviour itself, +the d224/d227 reruns above are the canonical signal. + +## Sequencing relative to the broader recovery work + +This PR is **independent of** the deadline-vs-quarantine proposal +in `claude-swe/disagg-deadline-vs-quarantine-proposal.md`. They're +complementary: + +- This PR (deadline anchor): "deadline only fires for actual + in-transfer time." +- The proposal (deadline-vs-quarantine): "deadline-firing should + surface a user error but not drive the unhealthy budget." + +The proposal can be applied on top of this PR if the team agrees +with the two-layer model. Either order works; this PR alone fixes +the immediate false-positive issue under the d224 scenario, +provided no_progress alone is the only path to `is_shutdown` (which +empirically it is on the prefill side under d224 — the second +UNHEALTHY flip is driven by the time-since-progress check, not by +quarantine count). 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..ec5faf9dfda6 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -0,0 +1,175 @@ +# 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`. | + +### 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. +* 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 +`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. + +## 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: + +* 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/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md b/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md new file mode 100644 index 000000000000..31536c3de98a --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md @@ -0,0 +1,82 @@ +# `[wedge-trace]` log lines — quick reference + +Companion doc to the verify-wedge runbook (in TRT-LLM PR #13796). When +the disagg KV transfer wedges, the four polling loops where the +process can stall now emit a periodic `[wedge-trace]` log at +`WARNING` level, so an operator can identify the wedge location from +log output alone — no `gdb` attach required. + +All four log sites share the same heartbeat interval, controlled by +**`TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS`** (default `5000`). Set to +`0` to disable. + +## The four heartbeats + +| Log prefix | Source | Wedge it surfaces | +|---|---|---| +| `[wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG …` | `cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp` (inside `NixlTransferStatus::wait`) | **Sender** worker stuck polling `getXferStatus` — the in-flight `ucp_request` is not completing. UCX endpoint or peer wedged. | +| `[wedge-trace] AgentConnectionManager::waitForNotification still pending …` | `cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp` (inside `waitForNotification`) | **Receiver** worker stuck polling `getNotifs` — the post-write notification from the sender never arrived. Sender crashed or notification dropped. | +| `[wedge-trace] BaseTransBufferManager::assignBufferIndex blocked …` | `cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp` (inside `assignBufferIndex`) | A new request can't acquire a slot — slot is held by a wedged worker. **Downstream effect** of one of the above wedges. | +| `[wedge-trace] CacheTransceiver::requestAndReceiveSync executor thread blocked …` | `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp` (inside `requestAndReceiveSync`) | Executor thread blocked on `future.get`. **Only on the** `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1` **path.** Pairs with one of the worker-thread heartbeats above. | + +## Sample log output during a wedge + +``` +[2026-05-06 10:42:18.131] WARN [wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after 5023 ms (handle=0x7f1c4c01a3b0, requested timeout_ms=-1). Sender-side getXferStatus poll is not making progress; peer or UCX endpoint may be wedged. +[2026-05-06 10:42:18.155] WARN [wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot for 5009 ms (concurrence=1 budget=1). The slot is most likely held by a wedged worker thread; check sender/receiver heartbeats and run the stuck-slot diagnosis. +[2026-05-06 10:42:23.131] WARN [wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after 10025 ms (handle=0x7f1c4c01a3b0, requested timeout_ms=-1). … +[2026-05-06 10:42:23.155] WARN [wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot for 10010 ms (concurrence=1 budget=1). … +``` + +Reading this: + +* The handle pointer `0x7f1c4c01a3b0` identifies which NIXL transfer + is wedged. Cross-reference with earlier `start recv bufferIdx` + debug logs (see + [`disagg-kv-transfer-debug-stuck-slot.md`](disagg-kv-transfer-debug-stuck-slot.md)) + to map the handle to a request id. +* `concurrence=1 budget=1` means the recv-buffer pool has the default + one slot and it is fully consumed; this is the classic "single + wedged transfer pins the rank" pattern documented in + [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md). + +## Mapping log line → recovery action + +| Log line you see | Most likely cause | Action | +|---|---|---| +| Only `NixlTransferStatus::wait` heartbeats fire (sender side) | Sender's `ucp_request` not completing — peer down or UCX endpoint wedge | Health-driven worker restart via PR #13796's `_check_transceiver_health` (default ~3 min total). Ask NIXL Q1.2 — should `getXferStatus` surface UCX endpoint-error as a final state? | +| Only `waitForNotification` heartbeats fire (receiver side) | Sender notification dropped or sender crashed before sending it | Same health-driven restart path. Ask NIXL Q4.1 — are notifications reliable? | +| Both sides heartbeat | Both peers wedged; usually one started it (often a downstream pod restart) | Restart the worker that's hardest to recover. | +| Only `assignBufferIndex` heartbeats fire (no NIXL-level heartbeat) | The thread holding the slot is wedged in something OTHER than NIXL — most likely CUDA, MPI, or a TRT-LLM mutex | `gdb -p` the process to find the slot owner; this is NOT a NIXL wedge. | +| Only `requestAndReceiveSync` heartbeat fires (no worker heartbeat) | Possible. Means the receive worker died without setting the promise. Or `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1` is on AND the worker is healthy but slow | Check for crashed worker thread via `cat /proc/$PID/task/*/comm`. | + +## Disabling the heartbeats + +```bash +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=0 +``` + +…or for finer control: + +```bash +# Only fire after 30 s of no progress (suppresses occasional slow transfers) +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=30000 + +# Tight loop, useful in test environments +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=1000 +``` + +The heartbeats themselves are cheap (a single `chrono::steady_clock::now()` +plus a `TLLM_LOG_WARNING` per interval per wedged thread). They do +not affect throughput on the happy path. + +## Cross-references + +* [`disagg-kv-transfer-debug-verify-wedge.md`](disagg-kv-transfer-debug-verify-wedge.md) + — the gdb / py-spy / NIXL-telemetry / UCX-stats verification recipes. +* [`disagg-kv-transfer-debug-stuck-slot.md`](disagg-kv-transfer-debug-stuck-slot.md) + — the start-recv-without-finish-recv slot-leak diagnosis. +* [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md) + — recovery model and global progress deadline. +* [`disagg-nixl-open-questions.md`](disagg-nixl-open-questions.md) — + upstream NIXL questions whose answers would shorten the recovery path. 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..e7a84ddd61cd 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,35 @@ 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() + # 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( @@ -634,7 +663,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 +683,114 @@ 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) + + 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): @@ -1331,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." @@ -1700,6 +1846,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() @@ -1722,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() @@ -1830,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: @@ -2174,6 +2325,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() @@ -2198,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() @@ -2421,6 +2575,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() @@ -2486,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() @@ -3331,6 +3492,13 @@ 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 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: if request_id not in requests_in_transfer: @@ -3340,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. @@ -3349,13 +3534,45 @@ 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 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") @@ -3568,6 +3785,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 +3824,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,8 +3984,42 @@ 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 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]) 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..3d81768d2f15 --- /dev/null +++ b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py @@ -0,0 +1,506 @@ +# 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 _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 + 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 + health: _FakeHealth = field(default_factory=_FakeHealth) + + 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 get_health(self) -> _FakeHealth: + return self.health + + +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._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 = [] + 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) + + 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 + + +# --------------------------------------------------------------------------- +# _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 == [] + + +# --------------------------------------------------------------------------- +# 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