diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index c4acce47593f..71bfb2a620c1 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -102,6 +102,30 @@ using RequestIdType = LlmRequest::RequestIdType; constexpr int kTransferFuturePollIntervalMs = 10; +template +void reserveForAppend(std::vector& values, std::size_t additional) +{ + TLLM_CHECK_WITH_INFO( + additional <= values.max_size() - values.size(), "Cannot reserve %zu additional transfer records", additional); + auto const requiredCapacity = values.size() + additional; + if (requiredCapacity <= values.capacity()) + { + return; + } + + auto newCapacity = std::max(1, values.capacity()); + while (newCapacity < requiredCapacity) + { + if (newCapacity > values.max_size() / 2) + { + newCapacity = requiredCapacity; + break; + } + newCapacity *= 2; + } + values.reserve(newCapacity); +} + // Finite status checks are scheduler polls, not terminal deadlines. Pure polls // use short slices; calls that ask for at least one completion keep bounded // backpressure by waiting up to the configured future timeout. @@ -777,6 +801,10 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques } return; } + // Reserve the durable status record before handing the request to the + // sender. Once sendAsync succeeds it may own request-backed KV pages, and + // the noexcept moves below must not lose the only future that can reap it. + reserveForAppend(mSenderFutures, 1); setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); @@ -785,6 +813,7 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques void CacheTransceiver::respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) { + reserveForAppend(mSenderFutures, requests.size()); for (auto const& llmRequest : requests) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); @@ -1578,6 +1607,10 @@ bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) } if (llmRequest->isContextOnlyRequest()) { + // Keep the high-level future until checkContextTransferStatus records + // this rank's terminal outcome. Queued/current classification can + // differ across model-parallel ranks, so rank-local reaping would + // remove the evidence needed by the existing transfer consensus. return mCacheSender->cancelRequest(*llmRequest); } else if (llmRequest->isGenerationOnlyRequest()) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 0f8ded65613f..98c749c6876e 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -31,12 +31,16 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include #include +#include #include #include #include #include +#include #include #include +#include +#include #include namespace tensorrt_llm::batch_manager @@ -238,6 +242,13 @@ using DataContext = tensorrt_llm::executor::kv_cache::DataContext; namespace { +// Bound sender-side terminal state during timeout storms. No-peer markers are +// eligible for cap-driven reclamation after the minimum retention interval, +// then move into the separately bounded finalized-handshake replay history. +constexpr std::size_t kMaxPendingPreHandshakeCancellations{65'536}; +constexpr std::size_t kMaxFinalizedHandshakeIds{65'536}; +constexpr auto kPreHandshakeCancellationRetention = std::chrono::seconds{30}; + int32_t tagFromRequestId(LlmRequest::RequestIdType requestId) { constexpr int32_t kDATA_TAG{43}; @@ -345,16 +356,29 @@ class CacheSender::Impl public: using RequestIdType = LlmRequest::RequestIdType; + struct ReceivedRequestInfo + { + RequestInfo requestInfo; + bool handledWithoutTransfer{false}; + }; + Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, - std::string instanceId = "") + std::string instanceId, std::size_t maxPendingPreHandshakeCancellations, + std::chrono::milliseconds preHandshakeCancellationRetention) : mManager{manager} , mSelfState{cacheLayer.getCacheState(), executor::kv_cache::CommState{manager->getCommState()}} , mCacheTransferLayer{std::move(cacheLayer)} , mBufferManager{std::make_shared()} , mInstanceId{std::move(instanceId)} + , mMaxPendingPreHandshakeCancellations{maxPendingPreHandshakeCancellations} + , mPreHandshakeCancellationRetention{preHandshakeCancellationRetention} { TLLM_CHECK(mManager); TLLM_CHECK(mManager->getCommState().getSelfIdx() == selfIndex); + TLLM_CHECK_WITH_INFO(mMaxPendingPreHandshakeCancellations > 0, + "The maximum number of pending pre-handshake cancellations must be positive"); + TLLM_CHECK_WITH_INFO( + mPreHandshakeCancellationRetention.count() > 0, "Pre-handshake cancellation retention must be positive"); TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); mResponseFuture = std::async(std::launch::async, &Impl::response, this); int asyncSendThreadNum = common::getEnvKVCacheSendMaxConcurrenceNum(); @@ -411,14 +435,6 @@ class CacheSender::Impl mSelfState.setCommState(std::move(commState)); } - [[nodiscard]] size_t getCounterpartsCount(LlmRequest::RequestIdType requestId) - { - std::unique_lock lock(mMtxForMap); - auto it = mRequestToSession.find(requestId); - TLLM_CHECK(it != mRequestToSession.end()); - return it->second.getConnections().size(); - } - void release(LlmRequest::RequestIdType requestId) { { @@ -479,16 +495,22 @@ class CacheSender::Impl } } - [[nodiscard]] std::optional recvRequestInfo() + [[nodiscard]] std::optional recvRequestInfo(bool rejectTerminalRequest) { auto* agentConnectionManager = dynamic_cast(mManager); bool isAgent = agentConnectionManager != nullptr; TransceiverTag::Id id; RequestInfo info; - auto const* connection = isAgent - ? agentConnectionManager->recvConnectionAndRequestInfo(info, mTerminate) - : mManager->recvConnect(DataContext{TransceiverTag::kID_TAG, mTerminate}, &id, sizeof(id)); + Connection const* connection = nullptr; + if (isAgent) + { + connection = agentConnectionManager->recvConnectionAndRequestInfo(info, mTerminate); + } + else + { + connection = mManager->recvConnect(DataContext{TransceiverTag::kID_TAG, mTerminate}, &id, sizeof(id)); + } if (connection == nullptr) { return std::nullopt; @@ -507,6 +529,20 @@ class CacheSender::Impl } auto requestId = info.getRequestId(); + if (rejectTerminalRequest) + { + bool isTerminal = false; + { + std::scoped_lock lock(mSenderMutex); + isTerminal = mFinalizedHandshakeIds.find(requestId) != mFinalizedHandshakeIds.end() + || mPreHandshakeCancellationDeadlines.find(requestId) != mPreHandshakeCancellationDeadlines.end(); + } + if (isTerminal) + { + notifyRejectedPeersNoThrow(requestId, {connection}); + return ReceivedRequestInfo{std::move(info), true}; + } + } mCacheTransferLayer.validateSupport(info.getTransState()); auto allCounterparts = mCacheTransferLayer.computeCounterparts( @@ -523,26 +559,58 @@ class CacheSender::Impl { cancelFlag = getOrCreateInFlightCancelFlag(requestId); } + bool rejectRequest = false; + bool rejectedRequestHasSession = false; { - std::unique_lock lk(mMtxForMap); - auto it = mRequestToSession.find(requestId); - if (it == mRequestToSession.end()) - { - auto session = cancelFlag != nullptr - ? TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, - info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, - !common::getEnvKVCacheTimeOutputPath().empty()) - : TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, - info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, - !common::getEnvKVCacheTimeOutputPath().empty()); - session.setTime(TransferSession::kTimeRequestInfo); - it = mRequestToSession.emplace(requestId, std::move(session)).first; - } - it->second.setConnection(peerIdx, connection); - } - return info; + // Publish each peer connection atomically with cancellation so a + // queued request cannot become a transfer session between the + // cancellation decision and session lookup. + std::scoped_lock lock(mSenderMutex, mMtxForMap); + if (rejectTerminalRequest + && (mFinalizedHandshakeIds.find(requestId) != mFinalizedHandshakeIds.end() + || mPreHandshakeCancellationDeadlines.find(requestId) != mPreHandshakeCancellationDeadlines.end())) + { + // A known-terminal peer request is a late replay. Never create + // an orphan session that would head-of-line block the response + // worker; reject only IDs whose terminal state is explicit. + rejectRequest = true; + rejectedRequestHasSession = mRequestToSession.find(requestId) != mRequestToSession.end(); + } + else + { + auto it = mRequestToSession.find(requestId); + if (it == mRequestToSession.end()) + { + auto session = cancelFlag != nullptr + ? TransferSession(std::vector(allCounterparts.size(), nullptr), + DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, + info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), + nullptr, !common::getEnvKVCacheTimeOutputPath().empty()) + : TransferSession(std::vector(allCounterparts.size(), nullptr), + DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, + info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), + nullptr, !common::getEnvKVCacheTimeOutputPath().empty()); + session.setTime(TransferSession::kTimeRequestInfo); + it = mRequestToSession.emplace(requestId, std::move(session)).first; + if (rejectTerminalRequest) + { + mRemainSendCount.emplace(requestId, allCounterparts.size()); + } + } + it->second.setConnection(peerIdx, connection); + } + } + + if (rejectRequest) + { + if (!rejectedRequestHasSession && cancelFlag != nullptr) + { + discardTransferState(requestId); + } + notifyRejectedPeersNoThrow(requestId, {connection}); + return ReceivedRequestInfo{std::move(info), true}; + } + return ReceivedRequestInfo{std::move(info), false}; } void sendSync(LlmRequest const& llmRequest) @@ -561,38 +629,123 @@ class CacheSender::Impl llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); } + void reclaimExpiredPreHandshakeCancellationLocked() + { + auto const now = std::chrono::steady_clock::now(); + while (mPreHandshakeCancellationDeadlines.size() >= mMaxPendingPreHandshakeCancellations + && !mPreHandshakeCancellationOrder.empty() && mPreHandshakeCancellationOrder.front().first <= now) + { + auto const [deadline, requestId] = mPreHandshakeCancellationOrder.front(); + auto const deadlineIt = mPreHandshakeCancellationDeadlines.find(requestId); + if (deadlineIt != mPreHandshakeCancellationDeadlines.end() && deadlineIt->second == deadline) + { + // Move an expired no-peer marker into the bounded replay + // history before erasing it, so a late peer remains terminal. + recordFinalizedHandshakeLocked(requestId); + mPreHandshakeCancellationDeadlines.erase(deadlineIt); + } + mPreHandshakeCancellationOrder.pop_front(); + } + } + + bool recordCancellationLocked(RequestIdType requestId, bool isPreHandshake) + { + if (isPreHandshake) + { + if (mPreHandshakeCancellationDeadlines.find(requestId) != mPreHandshakeCancellationDeadlines.end()) + { + return true; + } + if (mPreHandshakeCancellationDeadlines.size() >= mMaxPendingPreHandshakeCancellations) + { + reclaimExpiredPreHandshakeCancellationLocked(); + if (mPreHandshakeCancellationDeadlines.size() >= mMaxPendingPreHandshakeCancellations) + { + return false; + } + } + + auto const deadline = std::chrono::steady_clock::now() + mPreHandshakeCancellationRetention; + auto const [deadlineIt, inserted] = mPreHandshakeCancellationDeadlines.emplace(requestId, deadline); + TLLM_CHECK(inserted); + try + { + mPreHandshakeCancellationOrder.emplace_back(deadline, requestId); + } + catch (...) + { + mPreHandshakeCancellationDeadlines.erase(deadlineIt); + throw; + } + return true; + } + + if (mCancelledRequests.find(requestId) != mCancelledRequests.end()) + { + return true; + } + mCancelledRequests.insert(requestId); + return true; + } + bool cancelRequest(LlmRequest const& llmRequest) { bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); bool isCancelled = false; bool isCurrentRequest = false; + bool cancellationAdmissionDeclined = false; + std::optional cancelledResponse; { - std::scoped_lock lock(mSenderMutex); + // Serialize the sender queue and transfer-session lookup. This + // makes pre-handshake cancellation atomic with publication of the + // first peer RequestInfo: default-off cancellation may win before + // any session exists, but must drain once any peer has arrived. + std::scoped_lock lock(mSenderMutex, mMtxForMap); auto it = mReadyResponses.find(llmRequest.mRequestId); if (it != mReadyResponses.end()) { isCurrentRequest = mCurrentRequest.has_value() && mCurrentRequest.value() == llmRequest.mRequestId; - // The legacy path cannot interrupt a ready/active transfer, so - // preserve its false return until the opt-in is enabled. - if (!isCurrentRequest || inflightCancelEnabled) + auto sessionIt = mRequestToSession.find(llmRequest.mRequestId); + bool const hasTransferSession = sessionIt != mRequestToSession.end(); + if (!isCurrentRequest && (inflightCancelEnabled || !hasTransferSession)) { - mCancelledRequests.insert(llmRequest.mRequestId); - isCancelled = true; - if (inflightCancelEnabled && !isCurrentRequest) + bool const isPreHandshake = !hasTransferSession; + if (recordCancellationLocked(llmRequest.mRequestId, isPreHandshake)) { - // Keep only the request ID as a tombstone so a late peer - // receives ready=false without retaining the request. - failResponse(it->second, - std::make_exception_ptr( - TLLM_REQUEST_EXCEPTION(llmRequest.mRequestId, common::RequestErrorCode::kNETWORK_ERROR, - "Context KV cache request cancelled before a peer was ready for request %zu", - llmRequest.mRequestId))); + cancelledResponse.emplace(std::move(it->second)); mReadyResponses.erase(it); + isCancelled = true; + } + else + { + cancellationAdmissionDeclined = true; } } + else if (inflightCancelEnabled) + { + // The legacy path cannot interrupt a current/active transfer. The opt-in path preserves the + // response until sendResponse coordinates ready=false or the in-flight flag stops the transfer. + isCancelled = recordCancellationLocked(llmRequest.mRequestId, false); + } + } + else if (mCancelledRequests.find(llmRequest.mRequestId) != mCancelledRequests.end() + || mPreHandshakeCancellationDeadlines.find(llmRequest.mRequestId) + != mPreHandshakeCancellationDeadlines.end()) + { + // Cancellation is idempotent while its active or no-peer + // terminal marker is retained. + isCancelled = true; } } - if (inflightCancelEnabled && (!isCancelled || isCurrentRequest)) + if (cancelledResponse.has_value()) + { + failResponse(*cancelledResponse, + std::make_exception_ptr( + TLLM_REQUEST_EXCEPTION(llmRequest.mRequestId, common::RequestErrorCode::kNETWORK_ERROR, + "Context KV cache request cancelled before a peer was ready for request %zu", + llmRequest.mRequestId))); + } + if (inflightCancelEnabled && !cancellationAdmissionDeclined && (!isCancelled || isCurrentRequest)) { std::lock_guard lg(mInFlightCancelMutex); auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); @@ -604,7 +757,17 @@ class CacheSender::Impl } if (!isCancelled) { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + if (cancellationAdmissionDeclined) + { + TLLM_LOG_DEBUG( + "Cannot cancel request %zu before its peer arrives: the pending pre-handshake " + "cancellation limit of %zu was reached", + llmRequest.mRequestId, mMaxPendingPreHandshakeCancellations); + } + else + { + TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + } } else { @@ -640,6 +803,22 @@ class CacheSender::Impl } } + void sendReadySignal(Connection const* connection, LlmRequest::RequestIdType requestId, bool isReady) + { + auto* agentConnectionManager = dynamic_cast(mManager); + if (agentConnectionManager) + { + auto* agentConnection = dynamic_cast(connection); + TLLM_CHECK(agentConnection); + agentConnection->sendReadySignal(DataContext{tagFromRequestId(requestId), mTerminate}, isReady); + } + else + { + connection->send( + executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); + } + } + ~Impl() { terminate(); @@ -791,11 +970,10 @@ class CacheSender::Impl std::optional cancelledResponse; { std::scoped_lock lock(mSenderMutex); - TLLM_CHECK(mCurrentRequest.has_value() && mCurrentRequest.value() == reqId); auto responseIt = mReadyResponses.find(reqId); + auto countIt = mRemainSendCount.find(reqId); bool const isCancelled = mCancelledRequests.find(reqId) != mCancelledRequests.end(); TLLM_CHECK(responseIt != mReadyResponses.end() || isCancelled); - auto countIt = mRemainSendCount.find(reqId); TLLM_CHECK(countIt != mRemainSendCount.end()); auto const count = --countIt->second; TLLM_CHECK(count >= 0); @@ -813,6 +991,7 @@ class CacheSender::Impl mRemainSendCount.erase(countIt); isReady = !isCancelled; allCounterpartsReady = true; + mCurrentRequest = reqId; } } @@ -841,6 +1020,7 @@ class CacheSender::Impl response = std::move(it->second); mReadyResponses.erase(it); } + recordFinalizedHandshakeLocked(reqId); mCancelledRequests.erase(reqId); mCurrentRequest = std::nullopt; } @@ -913,19 +1093,19 @@ class CacheSender::Impl // Arbitrary transfers arrive without a pre-registered response; do not gate on // mReadyResponses. - auto requestInfo = recvRequestInfo(); - if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) + auto receivedRequestInfo = recvRequestInfo(/*rejectTerminalRequest=*/true); + if (!receivedRequestInfo.has_value() || mTerminate || !mManager->isRunning()) { break; } - auto const reqId = requestInfo->getRequestId(); - - if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) + if (receivedRequestInfo->handledWithoutTransfer) { - mRemainSendCount[reqId] = getCounterpartsCount(reqId); + continue; } + auto const& requestInfo = receivedRequestInfo->requestInfo; + auto const reqId = receivedRequestInfo->requestInfo.getRequestId(); - if (requestInfo->isArbitraryTransfer()) + if (requestInfo.isArbitraryTransfer()) { // No LlmRequest will ever be registered; serve from the reuse tree off-thread. { @@ -939,6 +1119,10 @@ class CacheSender::Impl { mRemainSendCount.erase(countIt); auto pinnedIds = pinReuseTreeBlocks(reqId); + { + std::scoped_lock lock(mSenderMutex); + recordFinalizedHandshakeLocked(reqId); + } if (pinnedIds.empty()) { TLLM_LOG_ERROR( @@ -974,12 +1158,13 @@ class CacheSender::Impl // The RequestInfo may race ahead of sendAsync; wait for the specific response. { std::unique_lock lock(mSenderMutex); - mCurrentRequest = reqId; mSenderCv.wait(lock, [this, reqId]() { return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() - || mCancelledRequests.find(reqId) != mCancelledRequests.end(); + || mCancelledRequests.find(reqId) != mCancelledRequests.end() + || mPreHandshakeCancellationDeadlines.find(reqId) + != mPreHandshakeCancellationDeadlines.end(); }); if (mTerminate) { @@ -1077,6 +1262,8 @@ class CacheSender::Impl pendingResponses.swap(mReadyResponses); mCurrentRequest = std::nullopt; mCancelledRequests.clear(); + mPreHandshakeCancellationDeadlines.clear(); + mPreHandshakeCancellationOrder.clear(); mRemainSendCount.clear(); } for (auto& entry : pendingResponses) @@ -1085,6 +1272,48 @@ class CacheSender::Impl } } + void notifyRejectedPeersNoThrow(RequestIdType requestId, std::vector const& connections) noexcept + { + TLLM_LOG_DEBUG("Rejecting a terminal KV cache handshake for request %zu", requestId); + for (auto const* connection : connections) + { + try + { + sendReadySignal(connection, requestId, false); + } + catch (std::exception const& error) + { + TLLM_LOG_WARNING("Failed to notify a rejected peer for request %zu: %s", requestId, error.what()); + } + catch (...) + { + TLLM_LOG_WARNING("Failed to notify a rejected peer for request %zu: unknown error", requestId); + } + } + } + + void recordFinalizedHandshakeLocked(RequestIdType requestId) + { + auto const [idIt, inserted] = mFinalizedHandshakeIds.insert(requestId); + if (inserted) + { + try + { + mFinalizedHandshakeOrder.push_back(requestId); + } + catch (...) + { + mFinalizedHandshakeIds.erase(idIt); + throw; + } + if (mFinalizedHandshakeOrder.size() > kMaxFinalizedHandshakeIds) + { + mFinalizedHandshakeIds.erase(mFinalizedHandshakeOrder.front()); + mFinalizedHandshakeOrder.pop_front(); + } + } + } + public: void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, @@ -1097,6 +1326,10 @@ class CacheSender::Impl private: std::optional mCurrentRequest; std::set mCancelledRequests; + std::unordered_map mPreHandshakeCancellationDeadlines; + std::deque> mPreHandshakeCancellationOrder; + std::deque mFinalizedHandshakeOrder; + std::unordered_set mFinalizedHandshakeIds; std::map mReadyResponses; std::mutex mSenderMutex; std::atomic mTerminate{false}; @@ -1106,7 +1339,6 @@ class CacheSender::Impl AsyncSendResource mAsyncSendResource; std::vector> mAsyncSendFutures; int mDeviceId{-1}; - executor::kv_cache::ConnectionManager* mManager; std::map mRequestToSession; executor::DataTransceiverState mSelfState; @@ -1117,6 +1349,8 @@ class CacheSender::Impl std::mutex mInFlightCancelMutex; std::unordered_map>> mInFlightCancelFlags; std::string mInstanceId; + std::size_t const mMaxPendingPreHandshakeCancellations; + std::chrono::milliseconds const mPreHandshakeCancellationRetention; }; class CacheReceiver::Impl @@ -1438,6 +1672,12 @@ class CacheReceiver::Impl bool cancelRequest(LlmRequest const& llmRequest) { + if (!common::getEnvDisaggEnableInflightCancel()) + { + TLLM_LOG_WARNING( + "Cannot cancel generation request %zu while in-flight cancellation is disabled", llmRequest.mRequestId); + return false; + } std::string processInfo = kDefaultProcessInfo; if (common::getEnvRequestKVCacheConcurrent()) @@ -1878,8 +2118,16 @@ void CacheReceiver::ImplDeleter::operator()(Impl* ptr) CacheSender::CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, std::string instanceId) - : mImpl{ - std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer), std::move(instanceId)))} + : mImpl{std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer), + std::move(instanceId), kMaxPendingPreHandshakeCancellations, kPreHandshakeCancellationRetention))} +{ +} + +CacheSender::CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, + CacheTransferLayer cacheLayer, std::size_t maxPendingPreHandshakeCancellations, + std::chrono::milliseconds preHandshakeCancellationRetention) + : mImpl{std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer), + /*instanceId=*/"", maxPendingPreHandshakeCancellations, preHandshakeCancellationRetention))} { } @@ -1907,9 +2155,15 @@ void CacheSender::sendSync(LlmRequest const& llmRequest) RequestInfo CacheSender::recvRequestInfo() { - auto requestInfo = mImpl->recvRequestInfo(); - TLLM_CHECK(requestInfo.has_value()); - return *requestInfo; + while (true) + { + auto requestInfo = mImpl->recvRequestInfo(false); + TLLM_CHECK(requestInfo.has_value()); + if (!requestInfo->handledWithoutTransfer) + { + return std::move(requestInfo->requestInfo); + } + } } bool CacheSender::cancelRequest(LlmRequest const& llmRequest) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 778e5e80c7f4..b4c6139461e7 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -16,6 +16,8 @@ */ #pragma once +#include +#include #include #include #include @@ -38,6 +40,11 @@ #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +namespace tensorrt_llm::testing +{ +class CacheSenderTestAccess; +} // namespace tensorrt_llm::testing + namespace tensorrt_llm::batch_manager { @@ -330,6 +337,11 @@ class CacheSender virtual ~CacheSender(); private: + friend class ::tensorrt_llm::testing::CacheSenderTestAccess; + + CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, + std::size_t maxPendingPreHandshakeCancellations, std::chrono::milliseconds preHandshakeCancellationRetention); + class Impl; struct ImplDeleter diff --git a/cpp/tests/unit_tests/executor/agentCommTest.cpp b/cpp/tests/unit_tests/executor/agentCommTest.cpp index 89488ba373e8..2c5b09d6ec97 100644 --- a/cpp/tests/unit_tests/executor/agentCommTest.cpp +++ b/cpp/tests/unit_tests/executor/agentCommTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,9 +15,13 @@ * limitations under the License. */ +#include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" +#include +#include #include +#include using namespace tensorrt_llm::batch_manager::kv_cache_manager; using namespace tensorrt_llm::runtime; @@ -232,5 +236,68 @@ TEST_P(AgentCommTest, AgentConnectionManagerConnect) TLLM_LOG_INFO("after finish"); } +TEST_P(AgentCommTest, CacheSenderRejectsStandaloneFinalizedReplay) +{ + namespace tbm = tensorrt_llm::batch_manager; + namespace tr = tensorrt_llm::runtime; + namespace texec = tensorrt_llm::executor; + + std::vector bufferManagers{mTransBufferManager.get()}; + auto receiverManager = std::make_unique(bufferManagers, *mCacheState, backend); + auto senderManager = std::make_unique(bufferManagers, *mCacheState, backend); + auto receiverCommState = receiverManager->getCommState(); + auto senderCommState = senderManager->getCommState(); + + auto* receiverConnection = const_cast( + dynamic_cast(receiverManager->getConnections(senderCommState).at(0))); + ASSERT_NE(receiverConnection, nullptr); + + std::vector cacheBufferManagers{mTransBufferManager.get()}; + tbm::CacheSender sender(senderManager.get(), 0, + tbm::CacheTransferLayer(*mCacheState, + tbm::kv_cache_manager::createCacheFormatter(mCacheManager.get(), cacheBufferManagers, false))); + + tbm::LlmRequest::RequestIdType constexpr requestId{29}; + tr::SamplingConfig const samplingConfig{1}; + auto const inputTokens = std::make_shared(tbm::LlmRequest::VecTokens{1}); + auto request = std::make_shared(requestId, 1, inputTokens, samplingConfig, /*isStreaming=*/false); + auto responseFuture = sender.sendAsync(request); + ASSERT_TRUE(sender.cancelRequest(*request)); + ASSERT_EQ(responseFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(responseFuture.get(), std::exception); + + texec::DataTransceiverState receiverState{*mCacheState, receiverCommState}; + tbm::RequestInfo requestInfo{requestId, receiverState}; + std::vector> cacheBufferIds{std::optional{0}}; + int constexpr validConnectionIdx{0}; + std::atomic terminate{false}; + int32_t constexpr kDataTag{43}; + int32_t constexpr dataTag = ((requestId & 0xFFF) << 8) | (kDataTag & 0xFF); + + auto sendRequestAndAwaitRejection = [&]() -> std::optional + { + receiverConnection->sendRequestAndBufferInfo(requestInfo, cacheBufferIds, validConnectionIdx); + auto readyFuture = std::async(std::launch::async, + [&]() { + return receiverConnection->recvReadySignal(DataContext{dataTag, terminate}); + }); + if (readyFuture.wait_for(std::chrono::seconds{10}) != std::future_status::ready) + { + terminate.store(true, std::memory_order_relaxed); + return std::nullopt; + } + return readyFuture.get(); + }; + + // The first handshake consumes the queued cancellation and finalizes the + // request. The replay must be rejected while the sender has no local work. + auto firstReady = sendRequestAndAwaitRejection(); + ASSERT_TRUE(firstReady.has_value()); + EXPECT_FALSE(*firstReady); + auto replayReady = sendRequestAndAwaitRejection(); + ASSERT_TRUE(replayReady.has_value()); + EXPECT_FALSE(*replayReady); +} + INSTANTIATE_TEST_SUITE_P(AvailableBackends, AgentCommTest, ::testing::ValuesIn(getAvailableBackends()), [](::testing::TestParamInfo const& info) { return info.param; }); diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 1f216578f669..039aae244bd4 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -45,12 +45,16 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" #include +#include #include #include #include #include #include +#include +#include #include +#include #include #include #include @@ -71,6 +75,23 @@ namespace texec = tensorrt_llm::executor; using testing::Return; using testing::ReturnRef; +namespace tensorrt_llm::testing +{ + +class CacheSenderTestAccess +{ +public: + static std::unique_ptr make(executor::kv_cache::ConnectionManager* manager, + runtime::SizeType32 selfIndex, batch_manager::CacheTransferLayer cacheLayer, + std::size_t maxPendingPreHandshakeCancellations, std::chrono::milliseconds preHandshakeCancellationRetention) + { + return std::unique_ptr(new batch_manager::CacheSender(manager, selfIndex, + std::move(cacheLayer), maxPendingPreHandshakeCancellations, preHandshakeCancellationRetention)); + } +}; + +} // namespace tensorrt_llm::testing + // --------------------------------------- // RequestInfoTest // --------------------------------------- @@ -91,6 +112,203 @@ T serializeDeserialize(T const& val) return T::deserialize(iss); } +class ControlledConnection final : public texec::kv_cache::Connection +{ +public: + void setSerializedRequestInfo(std::string serializedRequestInfo) + { + std::scoped_lock lock(mMutex); + mSerializedRequestInfo = std::move(serializedRequestInfo); + } + + void send(texec::kv_cache::DataContext const& ctx, void const* data, size_t size) const override + { + if (ctx.getTag() != TransceiverTag::kREADY_SIGNAL_TAG) + { + return; + } + + TLLM_CHECK(size == sizeof(bool)); + bool isReady{false}; + std::memcpy(&isReady, data, sizeof(isReady)); + + std::unique_lock lock(mMutex); + mReadySignal = isReady; + mReadySignalObserved = true; + mReadySignalsSent.push_back(isReady); + mConditionVariable.notify_all(); + mConditionVariable.wait(lock, [this] { return mReleaseReadySignal; }); + } + + void recv(texec::kv_cache::DataContext const& ctx, void* data, size_t size) const override + { + std::unique_lock lock(mMutex); + if (ctx.getTag() == TransceiverTag::kREADY_SIGNAL_TAG) + { + TLLM_CHECK(size == sizeof(bool)); + ++mReadySignalReceiveCalls; + mConditionVariable.notify_all(); + while (mReadySignalsToReceive.empty() && !ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + mConditionVariable.wait_for(lock, std::chrono::milliseconds{1}); + } + if (ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + bool const isReady = false; + std::memcpy(data, &isReady, sizeof(isReady)); + return; + } + bool const isReady = mReadySignalsToReceive.front(); + mReadySignalsToReceive.pop_front(); + std::memcpy(data, &isReady, sizeof(isReady)); + return; + } + if (ctx.getTag() == TransceiverTag::kINFO_SIZE_TAG) + { + auto const infoSize = static_cast(mSerializedRequestInfo.size()); + TLLM_CHECK(size == sizeof(infoSize)); + std::memcpy(data, &infoSize, sizeof(infoSize)); + return; + } + TLLM_CHECK(ctx.getTag() == TransceiverTag::kINFO_TAG); + TLLM_CHECK(size == mSerializedRequestInfo.size()); + std::memcpy(data, mSerializedRequestInfo.data(), size); + } + + bool waitForReadySignal(std::chrono::milliseconds timeout) const + { + std::unique_lock lock(mMutex); + return mConditionVariable.wait_for(lock, timeout, [this] { return mReadySignalObserved; }); + } + + bool getReadySignal() const + { + std::scoped_lock lock(mMutex); + TLLM_CHECK(mReadySignalObserved); + return mReadySignal; + } + + bool waitForReadySignalCount(std::size_t count, std::chrono::milliseconds timeout) const + { + std::unique_lock lock(mMutex); + return mConditionVariable.wait_for(lock, timeout, [this, count] { return mReadySignalsSent.size() >= count; }); + } + + std::deque getReadySignals() const + { + std::scoped_lock lock(mMutex); + return mReadySignalsSent; + } + + void releaseReadySignal() + { + { + std::scoped_lock lock(mMutex); + mReleaseReadySignal = true; + } + mConditionVariable.notify_all(); + } + + bool waitForReadySignalReceiveCalls(std::size_t count, std::chrono::milliseconds timeout) const + { + std::unique_lock lock(mMutex); + return mConditionVariable.wait_for(lock, timeout, [this, count] { return mReadySignalReceiveCalls >= count; }); + } + + void provideReadySignalToReceive(bool isReady) + { + { + std::scoped_lock lock(mMutex); + mReadySignalsToReceive.push_back(isReady); + } + mConditionVariable.notify_all(); + } + +private: + mutable std::mutex mMutex; + mutable std::condition_variable mConditionVariable; + mutable std::string mSerializedRequestInfo; + mutable bool mReadySignalObserved{false}; + mutable bool mReadySignal{false}; + mutable bool mReleaseReadySignal{false}; + mutable std::deque mReadySignalsSent; + mutable std::size_t mReadySignalReceiveCalls{0}; + mutable std::deque mReadySignalsToReceive; +}; + +class ControlledConnectionManager final : public texec::kv_cache::ConnectionManager +{ +public: + void publishRequestInfo(RequestInfo const& requestInfo) + { + std::ostringstream stream; + RequestInfo::serialize(requestInfo, stream); + mConnection.setSerializedRequestInfo(stream.str()); + { + std::scoped_lock lock(mMutex); + mRequestInfoAvailable = true; + } + mConditionVariable.notify_all(); + } + + texec::kv_cache::Connection const* recvConnect( + texec::kv_cache::DataContext const& ctx, void* data, size_t size) override + { + std::unique_lock lock(mMutex); + ++mRecvConnectCalls; + mConditionVariable.notify_all(); + while (!mRequestInfoAvailable && !ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + mConditionVariable.wait_for(lock, std::chrono::milliseconds{1}); + } + if (ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + return nullptr; + } + + TLLM_CHECK(ctx.getTag() == TransceiverTag::kID_TAG); + TLLM_CHECK(size == sizeof(TransceiverTag::Id)); + auto const id = TransceiverTag::Id::REQUEST_SEND; + std::memcpy(data, &id, sizeof(id)); + mRequestInfoAvailable = false; + return &mConnection; + } + + std::vector getConnections(texec::kv_cache::CommState const&) override + { + return {&mConnection}; + } + + texec::kv_cache::CommState const& getCommState() const override + { + return mCommState; + } + + bool isRunning() const override + { + return true; + } + + ControlledConnection& getConnection() + { + return mConnection; + } + + bool waitForRecvConnectCalls(std::size_t count, std::chrono::milliseconds timeout) + { + std::unique_lock lock(mMutex); + return mConditionVariable.wait_for(lock, timeout, [this, count] { return mRecvConnectCalls >= count; }); + } + +private: + ControlledConnection mConnection; + texec::kv_cache::CommState mCommState{std::vector{0}, 0}; + std::mutex mMutex; + std::condition_variable mConditionVariable; + bool mRequestInfoAvailable{false}; + std::size_t mRecvConnectCalls{0}; +}; + } // namespace TEST(ContextTransferCoordinatorTest, CommitsStaggeredSuccessAndFailureWithoutCollectivePolling) @@ -436,12 +654,23 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return std::make_unique(arbitraryId, std::move(request)); } - void addRequestAndTransportCache(std::shared_ptr const& llmRequest) + void addRequestToCache(std::shared_ptr const& llmRequest) { auto constexpr beamIdx{0}; auto constexpr beamWidth{1}; mManager->addSequenceBatch( {{{llmRequest->mRequestId, llmRequest->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*llmRequest)}); + } + + void removeRequestFromCache(std::shared_ptr const& llmRequest) + { + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest); + mManager->removeSequence(llmRequest->mRequestId, llmRequest); + } + + void addRequestAndTransportCache(std::shared_ptr const& llmRequest) + { + addRequestToCache(llmRequest); if (isSender) { auto blockRange = BlockRange::fromAllBlockIds(*mManager, llmRequest->mRequestId); @@ -478,6 +707,44 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- } } + std::unique_ptr makeControlledSender(ControlledConnectionManager& connectionManager) + { + constexpr int maxNumTokens{1024}; + mCacheTransBufferManager = std::make_unique(mManager.get(), maxNumTokens); + std::vector bufferManagers{mCacheTransBufferManager.get()}; + return std::make_unique(&connectionManager, 0, + CacheTransferLayer(*mCacheState, createCacheFormatter(mManager.get(), bufferManagers, /*isMLA=*/false))); + } + + std::unique_ptr makeControlledSender(ControlledConnectionManager& connectionManager, + std::size_t maxPendingPreHandshakeCancellations, std::chrono::milliseconds preHandshakeCancellationRetention) + { + constexpr int maxNumTokens{1024}; + mCacheTransBufferManager = std::make_unique(mManager.get(), maxNumTokens); + std::vector bufferManagers{mCacheTransBufferManager.get()}; + return tensorrt_llm::testing::CacheSenderTestAccess::make(&connectionManager, 0, + CacheTransferLayer(*mCacheState, createCacheFormatter(mManager.get(), bufferManagers, /*isMLA=*/false)), + maxPendingPreHandshakeCancellations, preHandshakeCancellationRetention); + } + + std::unique_ptr makeControlledReceiver(ControlledConnectionManager& connectionManager) + { + constexpr int maxNumTokens{1024}; + mCacheTransBufferManager = std::make_unique(mManager.get(), maxNumTokens); + std::vector bufferManagers{mCacheTransBufferManager.get()}; + return std::make_unique(&connectionManager, 0, + CacheTransferLayer(*mCacheState, createCacheFormatter(mManager.get(), bufferManagers, /*isMLA=*/false))); + } + + RequestInfo makeControlledRequestInfo( + LlmRequest::RequestIdType requestId, ControlledConnectionManager const& connectionManager) const + { + texec::DataTransceiverState state; + state.setCommState(connectionManager.getCommState()); + state.setCacheState(*mCacheState); + return RequestInfo{requestId, std::move(state)}; + } + bool isSender{false}; tensorrt_llm::mpi::MpiComm const* mComm; SizeType32 mWorldSize{0}, mlocalRank{0}; @@ -606,6 +873,401 @@ TEST_F(SymmetricalCacheTest, ArbitraryTransferTest) tensorrt_llm::mpi::MpiComm::world().barrier(); } +TEST_F(SymmetricalCacheTest, DefaultOffDoesNotCancelCurrentSender) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + auto sender = makeControlledSender(connectionManager); + std::shared_ptr request{makeLlmRequest(10)}; + addRequestToCache(request); + auto const requestInfo = makeControlledRequestInfo(request->mRequestId, connectionManager); + + auto future = sender->sendAsync(request); + connectionManager.publishRequestInfo(requestInfo); + auto const readySignalObserved = connectionManager.getConnection().waitForReadySignal(std::chrono::seconds{10}); + auto const readySignal = readySignalObserved && connectionManager.getConnection().getReadySignal(); + auto const cancellationAccepted = sender->cancelRequest(*request); + connectionManager.getConnection().releaseReadySignal(); + + EXPECT_TRUE(readySignalObserved); + EXPECT_TRUE(readySignal); + EXPECT_FALSE(cancellationAccepted); + ASSERT_EQ(future.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_NO_THROW(future.get()); + removeRequestFromCache(request); +} + +TEST_F(SymmetricalCacheTest, DefaultOffCancelsOnlyQueuedSenderAcrossRanks) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + auto sender = makeControlledSender(connectionManager); + std::shared_ptr request{makeLlmRequest(11)}; + addRequestToCache(request); + auto const requestInfo = makeControlledRequestInfo(request->mRequestId, connectionManager); + auto future = sender->sendAsync(request); + auto const worldRank = mComm->getRank(); + + bool readySignalObserved = false; + if (worldRank == 0) + { + connectionManager.publishRequestInfo(requestInfo); + readySignalObserved = connectionManager.getConnection().waitForReadySignal(std::chrono::seconds{10}); + } + mComm->barrier(); + + auto const cancellationAccepted = sender->cancelRequest(*request); + int const localAccepted = cancellationAccepted ? 1 : 0; + std::vector allAccepted(worldSize, 0); + mComm->allgather(&localAccepted, allAccepted.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); + auto const repeatedCancellationAccepted = sender->cancelRequest(*request); + int const localRepeatedAccepted = repeatedCancellationAccepted ? 1 : 0; + std::vector allRepeatedAccepted(worldSize, 0); + mComm->allgather(&localRepeatedAccepted, allRepeatedAccepted.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); + + bool futureReadyBeforeLatePeer = false; + if (worldRank == 1) + { + futureReadyBeforeLatePeer = future.wait_for(std::chrono::seconds{1}) == std::future_status::ready; + } + + bool lateReadySignalObserved = false; + bool lateReadySignal = true; + if (worldRank == 0) + { + connectionManager.getConnection().releaseReadySignal(); + } + else + { + connectionManager.publishRequestInfo(requestInfo); + lateReadySignalObserved = connectionManager.getConnection().waitForReadySignal(std::chrono::seconds{10}); + lateReadySignal = lateReadySignalObserved && connectionManager.getConnection().getReadySignal(); + connectionManager.getConnection().releaseReadySignal(); + } + + EXPECT_EQ(allAccepted, (std::vector{0, 1})); + EXPECT_EQ(allRepeatedAccepted, (std::vector{0, 1})); + if (worldRank == 0) + { + EXPECT_TRUE(readySignalObserved); + EXPECT_FALSE(cancellationAccepted); + } + else + { + EXPECT_TRUE(futureReadyBeforeLatePeer); + EXPECT_TRUE(lateReadySignalObserved); + EXPECT_FALSE(lateReadySignal); + EXPECT_TRUE(cancellationAccepted); + } + ASSERT_EQ(future.wait_for(std::chrono::seconds{10}), std::future_status::ready); + if (worldRank == 0) + { + EXPECT_NO_THROW(future.get()); + } + else + { + EXPECT_THROW(future.get(), std::exception); + } + removeRequestFromCache(request); +} + +TEST_F(SymmetricalCacheTest, DefaultOffNoPeerCancellationTombstonesAreReclaimedAtCapacity) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + constexpr std::size_t maxPendingPreHandshakeCancellations{1}; + constexpr auto preHandshakeCancellationRetention = std::chrono::seconds{2}; + auto sender = makeControlledSender( + connectionManager, maxPendingPreHandshakeCancellations, preHandshakeCancellationRetention); + std::shared_ptr firstRequest{makeLlmRequest(10)}; + std::shared_ptr secondRequest{makeLlmRequest(11)}; + + auto firstFuture = sender->sendAsync(firstRequest); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(1, std::chrono::seconds{10})); + EXPECT_TRUE(sender->cancelRequest(*firstRequest)); + EXPECT_TRUE(sender->cancelRequest(*firstRequest)); + ASSERT_EQ(firstFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(firstFuture.get(), std::exception); + + auto secondFuture = sender->sendAsync(secondRequest); + EXPECT_FALSE(sender->cancelRequest(*secondRequest)); + EXPECT_EQ(secondFuture.wait_for(std::chrono::milliseconds{0}), std::future_status::timeout); + + std::this_thread::sleep_for(preHandshakeCancellationRetention); + bool secondCancellationAccepted = false; + auto const retryDeadline = std::chrono::steady_clock::now() + std::chrono::seconds{10}; + while (!secondCancellationAccepted && std::chrono::steady_clock::now() < retryDeadline) + { + secondCancellationAccepted = sender->cancelRequest(*secondRequest); + if (!secondCancellationAccepted) + { + std::this_thread::sleep_for(std::chrono::milliseconds{50}); + } + } + ASSERT_TRUE(secondCancellationAccepted); + ASSERT_EQ(secondFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(secondFuture.get(), std::exception); + EXPECT_FALSE(sender->cancelRequest(*firstRequest)); + + connectionManager.getConnection().releaseReadySignal(); + connectionManager.publishRequestInfo(makeControlledRequestInfo(firstRequest->mRequestId, connectionManager)); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(1, std::chrono::seconds{10})); + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{false})); +} + +TEST_F(SymmetricalCacheTest, DefaultOffPreHandshakeCancellationRejectsEachLatePeer) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + connectionManager.getConnection().releaseReadySignal(); + constexpr std::size_t maxPendingPreHandshakeCancellations{1}; + constexpr auto preHandshakeCancellationRetention = std::chrono::milliseconds{200}; + auto sender = makeControlledSender( + connectionManager, maxPendingPreHandshakeCancellations, preHandshakeCancellationRetention); + std::shared_ptr cancelledRequest{makeLlmRequest(12)}; + + constexpr SizeType32 numLayers{4}; + constexpr SizeType32 numKvHeadsPerRank{1}; + constexpr SizeType32 sizePerHead{64}; + constexpr SizeType32 tokensPerBlock{8}; + constexpr SizeType32 peerTensorParallelism{2}; + texec::kv_cache::CacheState peerCacheState{numLayers, numKvHeadsPerRank, sizePerHead, tokensPerBlock, + peerTensorParallelism, /*pipelineParallelism=*/1, /*contextParallelism=*/1, std::vector{numLayers}, + tensorrt_llm::DataType::kFLOAT}; + auto makePeerRequestInfo = [&](SizeType32 selfIdx) + { + texec::DataTransceiverState peerState; + peerState.setCommState(texec::kv_cache::CommState{std::vector{0, 1}, selfIdx}); + peerState.setCacheState(peerCacheState); + return RequestInfo{cancelledRequest->mRequestId, std::move(peerState)}; + }; + + auto cancelledFuture = sender->sendAsync(cancelledRequest); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(1, std::chrono::seconds{10})); + EXPECT_TRUE(sender->cancelRequest(*cancelledRequest)); + ASSERT_EQ(cancelledFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(cancelledFuture.get(), std::exception); + + // Reclaim the expired marker under cap pressure while installing a + // replacement marker atomically. The transition must not stop the + // non-Agent receive loop between late TP peers for the finalized request. + std::this_thread::sleep_for(preHandshakeCancellationRetention * 2); + std::shared_ptr replacementRequest{makeLlmRequest(13)}; + auto replacementFuture = sender->sendAsync(replacementRequest); + EXPECT_TRUE(sender->cancelRequest(*replacementRequest)); + ASSERT_EQ(replacementFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(replacementFuture.get(), std::exception); + EXPECT_FALSE(sender->cancelRequest(*cancelledRequest)); + + connectionManager.publishRequestInfo(makePeerRequestInfo(0)); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(2, std::chrono::seconds{10})); + bool const firstPeerRejectedBeforeSecond + = connectionManager.getConnection().waitForReadySignalCount(1, std::chrono::seconds{1}); + EXPECT_TRUE(firstPeerRejectedBeforeSecond); + if (firstPeerRejectedBeforeSecond) + { + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{false})); + } + + // Always publish the second peer so a failing implementation that created + // a partial TP2 session can unwind instead of hanging during teardown. + connectionManager.publishRequestInfo(makePeerRequestInfo(1)); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(2, std::chrono::seconds{10})); + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{false, false})); + + std::shared_ptr liveRequest{makeLlmRequest(14)}; + addRequestToCache(liveRequest); + auto liveFuture = sender->sendAsync(liveRequest); + connectionManager.publishRequestInfo(makeControlledRequestInfo(liveRequest->mRequestId, connectionManager)); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(3, std::chrono::seconds{10})); + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{false, false, true})); + ASSERT_EQ(liveFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_NO_THROW(liveFuture.get()); + removeRequestFromCache(liveRequest); +} + +TEST_F(SymmetricalCacheTest, DefaultOffDoesNotCancelPartialAsymmetricHandshake) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + auto sender = makeControlledSender(connectionManager); + std::shared_ptr request{makeLlmRequest(12)}; + addRequestToCache(request); + + constexpr SizeType32 numLayers{4}; + constexpr SizeType32 numKvHeadsPerRank{1}; + constexpr SizeType32 sizePerHead{64}; + constexpr SizeType32 tokensPerBlock{8}; + constexpr SizeType32 peerTensorParallelism{2}; + texec::kv_cache::CacheState peerCacheState{numLayers, numKvHeadsPerRank, sizePerHead, tokensPerBlock, + peerTensorParallelism, /*pipelineParallelism=*/1, /*contextParallelism=*/1, std::vector{numLayers}, + tensorrt_llm::DataType::kFLOAT}; + auto makePeerRequestInfo = [&](SizeType32 selfIdx) + { + texec::DataTransceiverState peerState; + peerState.setCommState(texec::kv_cache::CommState{std::vector{0, 1}, selfIdx}); + peerState.setCacheState(peerCacheState); + return RequestInfo{request->mRequestId, std::move(peerState)}; + }; + + auto future = sender->sendAsync(request); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(1, std::chrono::seconds{10})); + auto const firstRequestInfo = makePeerRequestInfo(0); + connectionManager.publishRequestInfo(firstRequestInfo); + + // Entering recvConnect a second time proves the first peer's RequestInfo + // created a partial TP2 session. Default-off cancellation must now drain. + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(2, std::chrono::seconds{10})); + EXPECT_FALSE(sender->cancelRequest(*request)); + EXPECT_EQ(future.wait_for(std::chrono::milliseconds{0}), std::future_status::timeout); + + auto const secondRequestInfo = makePeerRequestInfo(1); + connectionManager.publishRequestInfo(secondRequestInfo); + auto const readySignalObserved = connectionManager.getConnection().waitForReadySignal(std::chrono::seconds{10}); + auto const readySignal = readySignalObserved && connectionManager.getConnection().getReadySignal(); + connectionManager.getConnection().releaseReadySignal(); + + EXPECT_TRUE(readySignalObserved); + EXPECT_TRUE(readySignal); + ASSERT_EQ(future.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_NO_THROW(future.get()); + removeRequestFromCache(request); +} + +TEST_F(SymmetricalCacheTest, LateTerminalRequestInfoIsRejectedWithoutBlockingNextRequest) +{ + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + connectionManager.getConnection().releaseReadySignal(); + auto sender = makeControlledSender(connectionManager); + std::shared_ptr firstRequest{makeLlmRequest(17)}; + addRequestToCache(firstRequest); + auto const firstRequestInfo = makeControlledRequestInfo(firstRequest->mRequestId, connectionManager); + auto firstFuture = sender->sendAsync(firstRequest); + connectionManager.publishRequestInfo(firstRequestInfo); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(1, std::chrono::seconds{10})); + ASSERT_EQ(firstFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_NO_THROW(firstFuture.get()); + removeRequestFromCache(firstRequest); + + // A queued live response wakes the sender. Finalized replay peers are + // rejected without consuming that response or poisoning the next transfer. + connectionManager.publishRequestInfo(firstRequestInfo); + std::shared_ptr nextRequest{makeLlmRequest(18)}; + addRequestToCache(nextRequest); + auto nextFuture = sender->sendAsync(nextRequest); + + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(2, std::chrono::seconds{10})); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(3, std::chrono::seconds{10})); + connectionManager.publishRequestInfo(firstRequestInfo); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(3, std::chrono::seconds{10})); + EXPECT_TRUE(connectionManager.waitForRecvConnectCalls(4, std::chrono::seconds{10})); + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{true, false, false})); + + connectionManager.publishRequestInfo(makeControlledRequestInfo(nextRequest->mRequestId, connectionManager)); + + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalCount(4, std::chrono::seconds{10})); + EXPECT_EQ(connectionManager.getConnection().getReadySignals(), (std::deque{true, false, false, true})); + ASSERT_EQ(nextFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_NO_THROW(nextFuture.get()); + removeRequestFromCache(nextRequest); +} + +TEST_F(SymmetricalCacheTest, DefaultOffReceiverDoesNotCancelQueuedRequest) +{ + if (tensorrt_llm::common::getEnvDisaggEnableInflightCancel()) + { + GTEST_SKIP() << "This test validates the default-off cancellation path."; + } + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + + ControlledConnectionManager connectionManager; + auto receiver = makeControlledReceiver(connectionManager); + std::shared_ptr firstRequest{makeLlmRequest(13)}; + std::shared_ptr queuedRequest{makeLlmRequest(14)}; + addRequestToCache(firstRequest); + addRequestToCache(queuedRequest); + + auto firstFuture = receiver->receiveAsync(firstRequest); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalReceiveCalls(1, std::chrono::seconds{10})); + auto queuedFuture = receiver->receiveAsync(queuedRequest); + + EXPECT_FALSE(receiver->cancelRequest(*queuedRequest)); + EXPECT_EQ(queuedFuture.wait_for(std::chrono::milliseconds{0}), std::future_status::timeout); + + connectionManager.getConnection().provideReadySignalToReceive(false); + EXPECT_TRUE(connectionManager.getConnection().waitForReadySignalReceiveCalls(2, std::chrono::seconds{10})); + connectionManager.getConnection().provideReadySignalToReceive(false); + + ASSERT_EQ(firstFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + ASSERT_EQ(queuedFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready); + EXPECT_THROW(firstFuture.get(), std::exception); + EXPECT_THROW(queuedFuture.get(), std::exception); + removeRequestFromCache(firstRequest); + removeRequestFromCache(queuedRequest); +} + #if ENABLE_MULTI_DEVICE using AsymmetricTestParam = std::tuple bool: def get_from_request_queue( self, - timeout: Optional[datetime.timedelta]) -> List[RequestQueueItem]: + timeout: Optional[datetime.timedelta], + *, + cap_batch_wait_to_timeout: bool = False) -> List[RequestQueueItem]: """Fetch requests from the queue with optional timeout. Args: timeout: Optional timeout for waiting on queue. + cap_batch_wait_to_timeout: Include the batching wait in ``timeout`` + instead of adding it after the initial wait. Returns: List of RequestQueueItem fetched from the queue. """ items = [] timeout_secs = timeout.total_seconds() if timeout is not None else None + total_deadline = None + if cap_batch_wait_to_timeout and timeout_secs is not None: + total_deadline = time.monotonic() + timeout_secs try: if self.request_queue.empty() and (timeout_secs is None @@ -185,6 +195,8 @@ def get_from_request_queue( return items deadline = time.monotonic() + self.batch_wait_timeout_ms / 1000.0 + if total_deadline is not None: + deadline = min(deadline, total_deadline) while len(items) < self.max_batch_size: remaining_timeout = deadline - time.monotonic() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index a686adc97528..4493dc2888ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -213,6 +213,29 @@ def cancel_request(self, req: LlmRequest): def supports_inflight_request_cancellation(self) -> bool: return False + def context_cancellation_reports_terminal_status(self) -> bool: + """Whether accepted CTX cancellation is completed by a status poll. + + Implementations returning ``True`` retain request ownership until + ``check_context_transfer_status`` reports a rank-consistent terminal + result. Implementations returning ``False`` preserve the legacy + immediate-cleanup contract after ``cancel_request`` succeeds. + """ + return False + + def generation_cancellation_reports_terminal_status(self) -> bool: + """Whether GEN cancellation must be finalized by a status poll. + + Implementations returning ``True`` retain request ownership until + ``check_gen_transfer_status`` reports a rank-consistent terminal + result. When in-flight cancellation is disabled, callers must not make + a rank-local queued-cancellation decision because another rank may + already be in the handshake. Implementations returning ``False`` + preserve the legacy immediate-cleanup contract after + ``cancel_request`` succeeds. + """ + return False + def has_poisoned_transfer_buffer(self) -> bool: return False @@ -330,6 +353,19 @@ def cancel_request(self, req: LlmRequest): def supports_inflight_request_cancellation(self) -> bool: return self._supports_inflight_request_cancellation + def context_cancellation_reports_terminal_status(self) -> bool: + # C++ status polling performs TP/CP/PP consensus. Even when only a + # subset of ranks can cancel before the handshake, every rank retains + # ownership until the mixed local outcomes become one terminal result. + return True + + def generation_cancellation_reports_terminal_status(self) -> bool: + # A queued receiver may be cancellable on one model-parallel rank while + # another is already in the handshake. Default-off callers therefore + # wait for natural terminal status instead of cancelling rank-locally; + # the opt-in path keeps ownership until C++ GEN status reaches consensus. + return True + def has_poisoned_transfer_buffer(self) -> bool: if not is_disagg_inflight_cancel_enabled(): return False diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..2652ec0dc2b4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -505,6 +505,7 @@ class PyExecutor: # If the number of micro batches is too large, the executor will spend too much host memory (No additional GPU memory is required). # 1024 in-flight micro batches can avoid synchronization in most cases and keep host memory usage low. MIN_ASYNC_MICRO_BATCH_NUM = 1024 + ASYNC_TRANSFER_POLL_INTERVAL_FALLBACK_MS = 1000 def __init__( self, @@ -822,7 +823,12 @@ def __init__( self.is_shutdown = False self._fatal_error: Optional[BaseException] = None self._error_budget = ErrorBudget() - self._disagg_timed_out_ctx_cancelled_ids: set[int] = set() + # AsyncTransferManager also owns connector transfers. Track the + # transceiver's ownership separately so an idle executor polls C++ + # sender futures even when no transfer timeout is configured. + self._async_context_transceiver_request_ids: set[int] = set() + self._disagg_ctx_cancel_requested_ids: set[int] = set() + self._disagg_gen_cancel_requested_ids: set[int] = set() self._disagg_timed_out_gen_cancelled_ids: set[int] = set() self._disagg_inflight_cancel_unsupported_logged = False self.max_batch_size = max_batch_size @@ -1107,10 +1113,9 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): self._terminate_request(request) return if self.async_transfer_manager.end_transfer(request): - if transfer_failed: - return # Skip if the PP=1 early path already terminated this request; - # under PP>1 that path is off, so terminate here on transfer-complete. + # under PP>1 that path is off, so terminate here after either a + # successful or failed transfer releases its final async owner. if not self.force_terminate_ctx_for_partial_reuse: self._terminate_request(request) @@ -4268,10 +4273,24 @@ def _handle_control_request(self): pending = self.control_requests[0] - if pending.control_requires_drain and (len(self.active_requests) != 0 - or len(self.waiting_queue) != 0): - # drain=True: keep the sentinel parked until the engine drains. - return + if pending.control_requires_drain: + async_transfer_manager = getattr(self, "async_transfer_manager", + None) + local_has_work = ( + len(self.active_requests) != 0 or len(self.waiting_queue) != 0 + or (async_transfer_manager is not None + and async_transfer_manager.has_any_inflight_requests())) + if getattr(self.dist, "world_size", 1) > 1: + has_work = bool( + self.dist.allreduce(int(local_has_work), op=ReduceOp.MAX)) + else: + has_work = local_has_work + if has_work: + # Context responses leave active_requests before their async + # KV send releases its final owner. Keep the control parked so + # sleep/update-weights cannot mutate GPU resources still read + # by a transfer on this or another model-parallel rank. + return logger.debug(f"[control_action] firing control request " f"drain={pending.control_requires_drain} " @@ -4914,16 +4933,55 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) + def _has_async_context_transfer_work(self) -> bool: + """Return whether an idle executor must keep polling context sends.""" + local_has_transfer = bool( + getattr(self, "_async_context_transceiver_request_ids", ())) + + # Rank 0 performs the queue wait, so it must also observe transfers + # owned by another attention-DP or model-parallel rank. Every rank + # enters this reduction only after rank 0 propagated the poll bit in + # RequestBroadcaster's count header, avoiding a conditional-collective + # split when local scheduler state differs. + if getattr(self.dist, "world_size", 1) > 1: + return bool( + self.dist.allreduce(int(local_has_transfer), op=ReduceOp.MAX)) + return local_has_transfer + + def _async_context_transfer_poll_timeout(self) -> datetime.timedelta: + """Return a bounded queue wait while context sends remain in flight.""" + poll_interval_ms = getattr(self.kv_cache_transceiver, + "kv_transfer_poll_interval_ms", None) + if poll_interval_ms is None: + poll_interval_ms = self.ASYNC_TRANSFER_POLL_INTERVAL_FALLBACK_MS + transfer_timeout_ms = getattr(self.kv_cache_transceiver, + "kv_transfer_timeout_ms", None) + if transfer_timeout_ms is not None: + poll_interval_ms = min(poll_interval_ms, transfer_timeout_ms) + return datetime.timedelta(milliseconds=poll_interval_ms) + def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, total_num_active_requests: int) -> None: """Fetch requests from request_queue and enqueue to waiting_queue.""" - # Block new requests while control requests are pending - if len(self.control_requests) != 0: - return + control_pending = len(self.control_requests) != 0 # Calculate timeout - idle = (total_num_active_requests == 0) and len(waiting_queue) == 0 - if idle: + scheduler_idle = (total_num_active_requests == 0 + and len(waiting_queue) == 0) + idle = not control_pending and scheduler_idle + draining_control_pending = ( + control_pending and self.control_requests[0].control_requires_drain) + poll_async_context_transfer = ( + scheduler_idle + and getattr(self, "kv_cache_transceiver", None) is not None + and (not control_pending or draining_control_pending)) + if poll_async_context_transfer: + # Context sends are removed from active_requests while their KV + # blocks remain pinned. Keep the queue wait bounded so transfer + # completion and timeout cleanup do not depend on another request + # arriving. A positive wait avoids spinning on empty iterations. + timeout = self._async_context_transfer_poll_timeout() + elif idle: # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0 # reaches the broadcast path regularly to prevent trtllm-serve timeout when idle. timeout = datetime.timedelta( @@ -4935,20 +4993,64 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, new_requests = [] if self.dist.rank == 0: # Process accumulated requests that were queued during control request handling. - if len(self.request_accumulated) != 0: + if not control_pending and len(self.request_accumulated) != 0: new_requests.extend(self.request_accumulated) self.request_accumulated.clear() # Reset timeout to 0 to avoid hanging when no new requests are available timeout = datetime.timedelta(0) with self.hang_detector.pause(): new_requests.extend( - self.executor_request_queue.get_from_request_queue(timeout)) + self.executor_request_queue.get_from_request_queue( + timeout, + cap_batch_wait_to_timeout=poll_async_context_transfer)) # Broadcast requests and handle Python objects. RequestBroadcaster probes # the request count first and can skip the heavy payload broadcast on # empty iterations. - new_requests, py_request_objects = self.request_broadcaster.broadcast( - new_requests) + new_requests, py_request_objects, poll_async_context_transfer = ( + self.request_broadcaster.broadcast(new_requests, + poll_async_context_transfer)) + + # Run this after the request-count broadcast so every model-parallel + # rank enters the transceiver's internal status consensus in the same + # order. This also covers the PP loop when no batch is executed. + if poll_async_context_transfer: + has_async_context_transfer_work = ( + self._has_async_context_transfer_work()) + if has_async_context_transfer_work: + self._check_kv_transfer_timeout() + self._check_disagg_ctx_cache_transfer_status(0) + + if control_pending: + # A draining control action must not admit ordinary work, but cancel + # sentinels still have to pass through. Otherwise a cancel queued + # behind the control cannot release the very request that the + # control is waiting to drain. Preserve all other items, including + # shutdown, in queue order until the control action completes. + deferred_requests = [] + for request_item in new_requests: + if request_item.is_canceled_request: + self.canceled_req_ids.append(request_item.id) + elif self.dist.rank == 0: + deferred_requests.append(request_item) + if self.dist.rank == 0: + self.request_accumulated.extend(deferred_requests) + + if self.canceled_req_ids: + canceled_waiting_items = waiting_queue.remove_by_ids( + set(self.canceled_req_ids)) + self._terminalize_canceled_waiting_requests( + canceled_waiting_items) + if canceled_waiting_items: + terminalized_ids = { + request_item.id + for request_item in canceled_waiting_items + } + self.canceled_req_ids[:] = [ + request_id for request_id in self.canceled_req_ids + if request_id not in terminalized_ids + ] + return # Validate and filter requests new_requests = self._handle_special_queue_items(new_requests) @@ -4961,6 +5063,23 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, waiting_queue.add_requests(new_requests) + # Terminalize queued cancellations before admission. This fetch path + # runs even when no model batch can make progress, so a waiting request + # cannot remain unresolved behind transfer or capacity pressure. + if self.canceled_req_ids: + canceled_waiting_items = waiting_queue.remove_by_ids( + set(self.canceled_req_ids)) + self._terminalize_canceled_waiting_requests(canceled_waiting_items) + if canceled_waiting_items: + terminalized_ids = { + request_item.id + for request_item in canceled_waiting_items + } + self.canceled_req_ids[:] = [ + request_id for request_id in self.canceled_req_ids + if request_id not in terminalized_ids + ] + def _pop_from_waiting_queue( self, waiting_queue: WaitingQueue, @@ -5093,8 +5212,17 @@ def _handle_special_queue_items( self.canceled_req_ids.append(req_item.id) elif req_item.is_control_request: self.control_requests.append(req_item) + deferred_requests = [] + for deferred_item in new_requests[idx + 1:]: + # Cancellation must bypass the pending control action. A + # control action may itself be waiting for this request to + # drain, so parking the marker would deadlock both. + if deferred_item.is_canceled_request: + self.canceled_req_ids.append(deferred_item.id) + elif self.dist.rank == 0: + deferred_requests.append(deferred_item) if self.dist.rank == 0: - self.request_accumulated.extend(new_requests[idx + 1:]) + self.request_accumulated.extend(deferred_requests) break else: accepted_new_requests.append(req_item) @@ -5494,6 +5622,22 @@ def _is_disagg_inflight_cancel_active(self) -> bool: self._disagg_inflight_cancel_unsupported_logged = True return False + def _context_cancellation_reports_terminal_status(self) -> bool: + reports_status = getattr( + self.kv_cache_transceiver, + "context_cancellation_reports_terminal_status", + None, + ) + return callable(reports_status) and reports_status() is True + + def _generation_cancellation_reports_terminal_status(self) -> bool: + reports_status = getattr( + self.kv_cache_transceiver, + "generation_cancellation_reports_terminal_status", + None, + ) + return callable(reports_status) and reports_status() is True + def _request_kv_transfer_cancellation(self, request: LlmRequest) -> bool: """Best-effort cancellation that leaves ownership intact on errors.""" try: @@ -5528,7 +5672,7 @@ def _cancel_timed_out_gen_transfers(self) -> None: continue elapsed_time = ((current_time - request.py_kv_transfer_start_time) * 1000) - if (elapsed_time > timeout_ms + if (elapsed_time >= timeout_ms and not request.py_kv_transfer_timed_out): logger.warning( f"Requesting cancellation for generation request " @@ -5572,6 +5716,7 @@ def _cancel_timed_out_gen_transfers(self) -> None: is_cancelled = self._request_kv_transfer_cancellation(request) if is_cancelled: + self._disagg_gen_cancel_requested_ids.add(request_id) self._disagg_timed_out_gen_cancelled_ids.add(request_id) logger.warning( f"Cancelled timed-out generation KV transfer for request " @@ -5613,18 +5758,22 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: if req.py_kv_transfer_start_time is None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 - if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: + if elapsed_time >= timeout_ms and not req.py_kv_transfer_timed_out: verb = ("Requesting cancellation for" if self._is_disagg_inflight_cancel_active() else "Observed timeout on") logger.warning( f"{verb} {type} request {req.py_request_id} due to KV " - f"cache transfer timeout: elapsed {elapsed_time:.0f}ms > " + f"cache transfer timeout: elapsed {elapsed_time:.0f}ms >= " f"kv_transfer_timeout_ms={timeout_ms}ms") req.py_kv_transfer_timed_out = True - for req in self.async_transfer_manager.requests_in_transfer().values(): - flag_if_kv_transfer_timed_out(req, "context") + transceiver_request_ids = getattr( + self, "_async_context_transceiver_request_ids", set()) + for request_id, req in self.async_transfer_manager.requests_in_transfer( + ).items(): + if request_id in transceiver_request_ids: + flag_if_kv_transfer_timed_out(req, "context") for req in self.active_requests: if req.is_disagg_generation_transmission_in_progress: @@ -6074,16 +6223,22 @@ def kv_connector_request_finished(req: LlmRequest): if req.is_context_only_request and ( req.is_context_finished or req.is_finished_due_to_length ) and not req.is_finished_due_to_cancellation: + request_id = req.py_request_id + if request_id in self._async_context_transceiver_request_ids: + # The native transceivers already treat duplicate + # publication as a no-op. Avoid acquiring a duplicate + # Python owner that would pin this request forever. + continue # Forward is done for this request — release the # IndexMapper slot so new requests can reuse it. # KV blocks stay allocated for the upcoming transfer. if hasattr(self.kv_cache_manager, 'release_index_slot'): - self.kv_cache_manager.release_index_slot( - req.py_request_id) + self.kv_cache_manager.release_index_slot(request_id) # Order is important here: we need to start the transfer before responding # to make sure the blocks are stored for reuse before they are sent. self.async_transfer_manager.start_transfer(req) self.kv_cache_transceiver.respond_and_send_async(req) + self._async_context_transceiver_request_ids.add(request_id) if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: req.py_kv_transfer_start_time = time.monotonic() @@ -6152,13 +6307,19 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in completed_req_ids: + # A terminal status releases exactly the transceiver's ownership; + # connector ownership, if any, remains in AsyncTransferManager. + self._async_context_transceiver_request_ids.discard(request_id) + self._disagg_ctx_cancel_requested_ids.discard(request_id) + if request_id not in requests_in_transfer: logger.warning( f"Request {request_id} not found in transfer manager") continue request = requests_in_transfer[request_id] - + request.py_kv_transfer_start_time = None + request.py_kv_transfer_timed_out = False self._end_transfer_and_maybe_terminate(request) # The set of requests in transfer may have changed since we terminated some requests. @@ -6167,25 +6328,30 @@ 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 (not request.py_kv_transfer_timed_out + if (request_id not in self._async_context_transceiver_request_ids + or not request.py_kv_transfer_timed_out or request_id in completed_req_ids - or request_id in self._disagg_timed_out_ctx_cancelled_ids): + or request_id in self._disagg_ctx_cancel_requested_ids): continue is_cancelled = self._request_kv_transfer_cancellation(request) if not is_cancelled: continue - if self._is_disagg_inflight_cancel_active(): - self._disagg_timed_out_ctx_cancelled_ids.add(request_id) + if (self._is_disagg_inflight_cancel_active() + or self._context_cancellation_reports_terminal_status()): + self._disagg_ctx_cancel_requested_ids.add(request_id) logger.warning(f"Cancelled timed-out context KV transfer for " f"request {request.py_request_id}; waiting for " - "C++ transfer status to report final cleanup") + "transfer status to report rank-consistent " + "final cleanup") else: # Preserve the legacy timeout behavior when in-flight # cancellation is disabled: a queued transfer that can be # cancelled is immediately released from the async manager. + self._async_context_transceiver_request_ids.discard(request_id) request.py_kv_transfer_start_time = None + request.py_kv_transfer_timed_out = False request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE self._end_transfer_and_maybe_terminate(request) @@ -6585,7 +6751,7 @@ def _terminate_request(self, request: LlmRequest): def _do_terminate_request(self, request: LlmRequest): self.resource_manager.free_resources(request) self._prefetched_request_ids.discard(request.py_request_id) - self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) + self._disagg_gen_cancel_requested_ids.discard(request.py_request_id) self._disagg_timed_out_gen_cancelled_ids.discard(request.py_request_id) if self.gather_all_responses or self.dist.rank == 0: @@ -6611,19 +6777,94 @@ def _try_cancel_request(self, request) -> bool: if (getattr(request, "is_context_only_request", False) is True and async_transfer_manager is not None and request.py_request_id in async_transfer_manager.requests_in_transfer()): - if self._is_disagg_inflight_cancel_active(): - self._request_kv_transfer_cancellation(request) + waits_for_status = ( + self._is_disagg_inflight_cancel_active() + or self._context_cancellation_reports_terminal_status()) + if (waits_for_status and request.py_request_id + not in self._disagg_ctx_cancel_requested_ids + and self._request_kv_transfer_cancellation(request)): + self._disagg_ctx_cancel_requested_ids.add(request.py_request_id) return False if not self._is_request_in_transmission(request): return True + if (request.is_generation_only_request() + and self._generation_cancellation_reports_terminal_status()): + request_id = request.py_request_id + # With default-off cancellation, one model-parallel rank may still + # be queued while another has entered the ready-signal handshake. + # Cancelling only the queued rank would prevent it from ever + # sending RequestInfo and strand the already-connected peers. + # The opt-in path can cancel every rank and finalizes through the + # C++ status consensus. + if (self._is_disagg_inflight_cancel_active() + and request_id not in self._disagg_gen_cancel_requested_ids + and self._request_kv_transfer_cancellation(request)): + self._disagg_gen_cancel_requested_ids.add(request_id) + return False + if self._is_disagg_inflight_cancel_active(): self._request_kv_transfer_cancellation(request) return False return self._request_kv_transfer_cancellation(request) + def _terminalize_canceled_waiting_requests( + self, request_items: List[RequestQueueItem]) -> None: + """Emit final cancellation responses for requests never activated.""" + if not request_items: + return + + # A terminalized queued request still counts as consumed from the + # executor input. Benchmark fill must not wait forever for a canceled + # request that can no longer be admitted. + self.num_fetch_requests += len(request_items) + + if self.enable_iter_perf_stats and self.dist.rank == 0: + # Remove queue-timing entries without charging canceled requests + # to the latency metric for requests that became active. + self.executor_request_queue.calculate_queue_latency( + request_items, time.time()) + + adp_collective_required = (self.enable_attention_dp + and self.dist.world_size != 1) + # Waiting items are replicated before attention-DP routing. Let only + # rank 0 create their response so allgather does not duplicate it. + create_response_locally = (self.dist.rank == 0 + or (self.gather_all_responses + and not adp_collective_required)) + + canceled_requests = [] + if create_response_locally: + canceled_requests = merge_requests( + request_items, + cp_config=self.dist.cp_config, + cp_rank=self.dist.cp_rank, + cp_size=self.dist.cp_size, + exclude_last_generation_logits=self. + _should_exclude_last_generation_logits()) + + canceled_responses = [] + for request in canceled_requests: + request.finish_by_reason(FinishReason.CANCELLED) + request.decoding_iter = request.py_decoding_iter + response = request.create_response(False, self.dist.rank) + if response is not None: + canceled_responses.append( + (request.py_request_id if not request.is_child else + request.parent_request_id, response)) + + if canceled_responses or adp_collective_required: + self._enqueue_responses(canceled_responses) + + # These requests never reached resource preparation, so resource + # managers must not be asked to free them. Drop the response routing + # entries only after publishing the terminal responses. + if self.gather_all_responses or self.dist.rank == 0: + for request_item in request_items: + self.result_wait_queues.pop(request_item.id, None) + @nvtx_range("_handle_canceled_requests") def _handle_canceled_requests(self): if len(self.canceled_req_ids) == 0: @@ -6632,15 +6873,26 @@ def _handle_canceled_requests(self): # Create set from list of canceled request ids to speed up canceled test canceled_req_ids_set = set(self.canceled_req_ids) - # Remove canceled requests from the waiting queue - self.waiting_queue.remove_by_ids(canceled_req_ids_set) + # Requests canceled before activation still need a terminal response; + # otherwise their result waiters and request mappings remain live. + canceled_waiting_items = self.waiting_queue.remove_by_ids( + canceled_req_ids_set) + self._terminalize_canceled_waiting_requests(canceled_waiting_items) + handled_canceled_ids = { + request_item.id + for request_item in canceled_waiting_items + } still_pending_canceled_ids = [] + still_pending_canceled_ids_set = set() + processed_request_ids = set() for request in self.active_requests: req_id = request.py_request_id if not request.is_child else request.parent_request_id if req_id not in canceled_req_ids_set: continue + handled_canceled_ids.add(req_id) + processed_request_ids.add(request.py_request_id) is_cancelled = self._try_cancel_request(request) if is_cancelled: # Mark requests as finished, then, we reuse all existing code @@ -6649,7 +6901,45 @@ def _handle_canceled_requests(self): request.finish_by_reason(FinishReason.CANCELLED) request.decoding_iter = request.py_decoding_iter else: - still_pending_canceled_ids.append(req_id) + if req_id not in still_pending_canceled_ids_set: + still_pending_canceled_ids.append(req_id) + still_pending_canceled_ids_set.add(req_id) + + # A context-only request may leave active_requests after its response + # while AsyncTransferManager still pins its KV blocks for the C++ send. + # Reconcile those manager-only requests so a later client cancellation + # can cancel a send that has not entered the peer handshake yet. C++ + # keeps ownership until its rank-consistent status reports terminal. + requests_in_transfer = self.async_transfer_manager.requests_in_transfer( + ) if getattr(self, "async_transfer_manager", None) is not None else {} + transceiver_request_ids = getattr( + self, "_async_context_transceiver_request_ids", set()) + for request_id, request in requests_in_transfer.items(): + req_id = self._request_vote_id(request) + if (req_id not in canceled_req_ids_set + or request_id not in transceiver_request_ids + or request_id in processed_request_ids): + continue + + handled_canceled_ids.add(req_id) + if self._try_cancel_request(request): + request.py_kv_transfer_timed_out = False + request.finish_by_reason(FinishReason.CANCELLED) + request.decoding_iter = request.py_decoding_iter + else: + if req_id not in still_pending_canceled_ids_set: + still_pending_canceled_ids.append(req_id) + still_pending_canceled_ids_set.add(req_id) + + # A draining control request parks ordinary input items on rank 0. Keep + # cancellation markers whose target has not appeared yet so they can + # terminalize the parked request when the control action completes. + if getattr(self, "control_requests", None): + for req_id in self.canceled_req_ids: + if (req_id not in handled_canceled_ids + and req_id not in still_pending_canceled_ids_set): + still_pending_canceled_ids.append(req_id) + still_pending_canceled_ids_set.add(req_id) # Clear list of requests marked for cancellation and add back those that failed to cancel. self.canceled_req_ids.clear() @@ -6795,6 +7085,20 @@ def _handle_responses(self, emit_first_iter: bool = True): # Check if a generation request needs cleanup due to KV cache transfer timeout. if request.py_kv_transfer_timed_out: + if (request.is_generation_only_request() and + self._generation_cancellation_reports_terminal_status() + ): + if request.is_disagg_generation_transmission_in_progress: + # Default-off timeout handling is observe-only. Do not + # cancel a queued receiver rank-locally: another rank + # may already be in the handshake. The opt-in timeout + # driver requests cancellation separately and also waits + # for C++ GEN status consensus. + new_active_requests.append(request) + else: + timed_out_requests.append(request) + continue + if self._is_disagg_inflight_cancel_active(): if (request.is_disagg_generation_transmission_in_progress or request.state diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index e7da86608153..8343894b1f3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Utility functions for request processing.""" import os @@ -510,16 +513,26 @@ def __init__(self, dist: Distributed, hang_detector: HangDetector): self.hang_detector = hang_detector self.send_requests_handler = None - def broadcast(self, new_requests: List) -> Tuple[List, Optional[Tuple]]: - """Broadcast requests and Python objects across ranks.""" - request_count = len(new_requests) if self.dist.rank == 0 else 0 + def broadcast( + self, new_requests: List, poll_context_transfers: bool = False + ) -> Tuple[List, Optional[Tuple], bool]: + """Broadcast requests, Python objects, and the rank-0 poll decision.""" + # Encode the poll bit in the existing count probe so every rank makes + # the transfer-progress decision from the same rank-0 value without + # adding another scheduler-iteration collective. + request_header = ( + (len(new_requests) << 1) | int(poll_context_transfers) if self.dist.rank == 0 else 0 + ) # Idle non-root ranks can wait here while rank 0 blocks in the # pause-wrapped request queue fetch, so keep the probe pause-wrapped too. with self.hang_detector.pause(): - request_count = self._broadcast_request_count(request_count) + request_header = self._broadcast_request_header(request_header) + + request_count = request_header >> 1 + poll_context_transfers = bool(request_header & 1) if request_count == 0: - return [], None + return [], None, poll_context_transfers if self.dist.rank == 0: py_request_objects = self._collect_py_objects(new_requests) @@ -535,31 +548,31 @@ def broadcast(self, new_requests: List) -> Tuple[List, Optional[Tuple]]: new_requests, py_request_objects ) - return new_requests, py_request_objects + return new_requests, py_request_objects, poll_context_transfers - def _broadcast_request_count(self, request_count: int) -> int: - """Broadcast rank 0's request count using the same PP route as requests.""" + def _broadcast_request_header(self, request_header: int) -> int: + """Broadcast rank 0's encoded request-count and transfer-poll bit.""" if self.dist.world_size == 1: - return request_count + return request_header if not self.dist.has_pp: - return self.dist.broadcast(request_count, root=0) + return self.dist.broadcast(request_header, root=0) if self.dist.is_first_pp_rank: - with nvtx_range("tp_broadcast_request_count"): - request_count = self.dist.tp_cp_broadcast(request_count, root=0) + with nvtx_range("tp_broadcast_request_header"): + request_header = self.dist.tp_cp_broadcast(request_header, root=0) tag = self.dist.pp_size + 1 # Avoid the heavy request payload tag. if not self.dist.is_first_pp_rank: with nvtx_range("recv_request_count_from_prev_pp"): - request_count = self.dist.recv_object(self.dist.prev_pp_rank, tag) + request_header = self.dist.recv_object(self.dist.prev_pp_rank, tag) if not self.dist.is_last_pp_rank: with nvtx_range("send_request_count_to_next_pp"): - self.dist.send_object(request_count, self.dist.next_pp_rank, tag) + self.dist.send_object(request_header, self.dist.next_pp_rank, tag) - return request_count + return request_header def _collect_py_objects(self, new_requests: List) -> Tuple: """Collect Python-only objects from requests.""" diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/waiting_queue.py b/tensorrt_llm/_torch/pyexecutor/scheduler/waiting_queue.py index 2d3583f5ab9e..a49ee5c89d4f 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/waiting_queue.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/waiting_queue.py @@ -73,8 +73,8 @@ def prepend_requests(self, requests: Iterable[RequestQueueItem]) -> None: pass @abstractmethod - def remove_by_ids(self, request_ids: set[int]) -> None: - """Remove requests with the given IDs.""" + def remove_by_ids(self, request_ids: set[int]) -> list[RequestQueueItem]: + """Remove and return requests with the given IDs.""" pass @abstractmethod @@ -139,11 +139,13 @@ def prepend_requests(self, requests: Iterable[RequestQueueItem]) -> None: """ self.extendleft(reversed(list(requests))) - def remove_by_ids(self, request_ids: set[int]) -> None: - """Remove requests with the given IDs.""" + def remove_by_ids(self, request_ids: set[int]) -> list[RequestQueueItem]: + """Remove and return requests with the given IDs.""" + removed_requests = [req for req in self if req.id in request_ids] filtered_requests = [req for req in self if req.id not in request_ids] self.clear() self.extend(filtered_requests) + return removed_requests def __bool__(self) -> bool: """Check if queue has any requests.""" @@ -243,10 +245,12 @@ def prepend_requests(self, requests: Iterable[RequestQueueItem]) -> None: for request in reversed(list(requests)): self._push_front(request) - def remove_by_ids(self, request_ids: set[int]) -> None: - """Remove requests with the given IDs.""" + def remove_by_ids(self, request_ids: set[int]) -> list[RequestQueueItem]: + """Remove and return requests with the given IDs.""" + removed_requests = [entry[2] for entry in self._heap if entry[2].id in request_ids] self._heap = [e for e in self._heap if e[2].id not in request_ids] heapq.heapify(self._heap) + return removed_requests def __bool__(self) -> bool: return len(self._heap) > 0 diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ffc560b2d6b6..9fc1d7f9c1b4 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1788,7 +1788,7 @@ def serve_embedding( "--request_timeout", type=int, default=180, - help="Request timeout", + help="End-to-end disaggregated request timeout", status="beta") @stability_option("-l", '--log_level', diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 86475af498fa..5ca726b907ac 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -173,6 +173,10 @@ def __init__( else reservation_timeout_secs ) self._reservation_tasks: dict[tuple[str, int], asyncio.Task] = {} + # A worker can be canceled while its /select request is still running. + # Remember an early /finish so a later select commit cannot recreate a + # reservation after cleanup has already passed through the coordinator. + self._finish_tombstone_tasks: dict[tuple[str, int], asyncio.Task] = {} self._ctx_client: Optional[OpenAIClient] = None self._gen_client: Optional[OpenAIClient] = None @@ -223,9 +227,22 @@ async def select( server, info, request_id = await router.get_next_server_by_key( routing_key, req_id=req_id, exclude_server=exclude_server ) - self._reservation_tasks[reservation_key] = asyncio.create_task( - self._expire_reservation(reservation_key, router) - ) + finish_tombstone = self._finish_tombstone_tasks.pop(reservation_key, None) + if finish_tombstone is not None: + finish_tombstone.cancel() + # Register the normal expiry fallback before attempting immediate + # cleanup. If router release fails, the reservation must not lose + # its last recovery path. + reservation = asyncio.create_task(self._expire_reservation(reservation_key, router)) + self._reservation_tasks[reservation_key] = reservation + await router.finish_request_by_id(req_id, False) + if self._reservation_tasks.get(reservation_key) is reservation: + self._reservation_tasks.pop(reservation_key, None) + reservation.cancel() + else: + self._reservation_tasks[reservation_key] = asyncio.create_task( + self._expire_reservation(reservation_key, router) + ) self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) return server, self._compact_route_info(info), request_id @@ -242,12 +259,25 @@ async def get_disagg_request_id(self) -> int: async def finish(self, role: str, req_id, success: bool = True) -> None: _t0 = time.monotonic() - reservation = self._reservation_tasks.pop((self._normalize_role(role), req_id), None) - if reservation is not None: - reservation.cancel() + reservation_key = (self._normalize_role(role), req_id) + reservation = self._reservation_tasks.get(reservation_key) + if reservation is None and reservation_key not in self._finish_tombstone_tasks: + self._finish_tombstone_tasks[reservation_key] = asyncio.create_task( + self._expire_finish_tombstone(reservation_key) + ) await self._router_for_role(role).finish_request_by_id(req_id, success) + if reservation is not None and self._reservation_tasks.get(reservation_key) is reservation: + self._reservation_tasks.pop(reservation_key, None) + reservation.cancel() self._api_lat(f"finish[{role}]").record(time.monotonic() - _t0) + async def _expire_finish_tombstone(self, key: tuple[str, int]) -> None: + try: + await asyncio.sleep(self._reservation_timeout_secs) + finally: + if self._finish_tombstone_tasks.get(key) is asyncio.current_task(): + self._finish_tombstone_tasks.pop(key, None) + async def _expire_reservation(self, key: tuple[str, int], router: Router) -> None: try: await asyncio.sleep(self._reservation_timeout_secs) @@ -306,12 +336,25 @@ async def start(self) -> None: await self._wait_for_all_servers_ready() async def stop(self) -> None: - reservations = list(self._reservation_tasks.values()) + reservation_items = list(self._reservation_tasks.items()) self._reservation_tasks.clear() - for reservation in reservations: - reservation.cancel() - if reservations: - await asyncio.gather(*reservations, return_exceptions=True) + finish_tombstones = list(self._finish_tombstone_tasks.values()) + self._finish_tombstone_tasks.clear() + cleanup_tasks = [task for _, task in reservation_items] + finish_tombstones + for task in cleanup_tasks: + task.cancel() + if cleanup_tasks: + await asyncio.gather(*cleanup_tasks, return_exceptions=True) + release_results = await asyncio.gather( + *( + self._router_for_role(role).finish_request_by_id(req_id, False) + for (role, req_id), _ in reservation_items + ), + return_exceptions=True, + ) + for result in release_results: + if isinstance(result, Exception): + logger.warning(f"Failed to release coordinator reservation during stop: {result}") if self._disagg_cluster_manager: await self._disagg_cluster_manager.stop() if self._metadata_server: diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 26646a54344c..9b8259034f2b 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -17,7 +17,7 @@ import os import traceback from abc import ABC, abstractmethod -from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Tuple, Type +from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Tuple, Type, cast import aiohttp @@ -58,6 +58,88 @@ _msgpack_encoder = msgspec.msgpack.Encoder() +class UpstreamRequestTimeoutError(TimeoutError): + """An orchestrator request to a disaggregated worker timed out.""" + + def __init__(self, role: ServerRole, timeout_secs: int, detail: Optional[str] = None): + self.role = role + self.timeout_secs = timeout_secs + message = ( + f"{role.name.lower()} worker request exceeded the configured " + f"{timeout_secs}-second timeout" + ) + if detail: + message = f"{message}: {detail}" + super().__init__(message) + + +class _ManagedResponseStream: + """Own router cleanup even when a response stream is never started.""" + + def __init__(self, client, request, resp_generator, req_id: Optional[int] = None): + self._client = client + self._request = request + self._resp_generator = resp_generator + self._req_id = req_id + self._cleanup_lock = asyncio.Lock() + self._closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + return await self.asend(None) + + async def asend(self, value): + if self._closed: + raise StopAsyncIteration + try: + return await self._resp_generator.asend(value) + except StopAsyncIteration: + await self._finish(success=True) + raise + except asyncio.CancelledError: + await self._finish(success=False) + raise + except Exception: + await self._finish(success=False) + raise + + async def athrow(self, *args): + if self._closed: + return None + try: + return await self._resp_generator.athrow(*args) + except (StopAsyncIteration, GeneratorExit): + await self._finish(success=False) + raise + except asyncio.CancelledError: + await self._finish(success=False) + raise + except Exception: + await self._finish(success=False) + raise + + async def aclose(self): + await self._finish(success=False) + + async def _finish(self, success: bool) -> None: + async with self._cleanup_lock: + if self._closed: + return + if not success: + self._client._metrics_collector.error_requests.inc() + try: + await self._client._cleanup_request_safely( + self._request, + self._resp_generator, + success=success, + req_id=self._req_id, + ) + finally: + self._closed = True + + class OpenAIClient(ABC): async def send_request( self, @@ -150,6 +232,7 @@ def __init__( self._max_retries = max_retries self._retry_interval_sec = retry_interval_sec self._disagg_id_generator = disagg_id_generator + self._timeout_secs = timeout_secs async def _send_request( self, @@ -160,38 +243,78 @@ async def _send_request( hooks: Optional[ResponseHooks] = None, req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: - if server is None: - if req_id is None: - server, _ = await self._router.get_next_server(request) - else: - server, _ = await self._router.get_next_server(request, req_id=req_id) - url = f"http://{server}/{endpoint}" - # disaggregated_params is None when conditional_disagg bypasses ctx. - _dp = request.disaggregated_params - _ctx_rid = _dp.ctx_request_id if _dp is not None else None - logger.debug(f"Sending {self._role} request {_ctx_rid} to {url}") + self._metrics_collector.total_requests.inc() + resp_generator = None + cleanup_deferred = False + success = False try: - self._metrics_collector.total_requests.inc() + if server is None: + if req_id is None: + server, _ = await self._router.get_next_server(request) + else: + server, _ = await self._router.get_next_server(request, req_id=req_id) + url = f"http://{server}/{endpoint}" + # disaggregated_params is None when conditional_disagg bypasses ctx. + _dp = request.disaggregated_params + _ctx_rid = _dp.ctx_request_id if _dp is not None else None + logger.debug(f"Sending {self._role} request {_ctx_rid} to {url}") resp_generator = self._post_with_retry(server, url, request, hooks, req_id) if request.stream: - # return the response generator, the request is not done yet - return resp_generator - else: - # consume the generator to get the response and return it directly when it's not streaming - response = None - async for resp_json in resp_generator: - response = response_type(**resp_json) - if hooks: - if self._role == ServerRole.CONTEXT: - hooks.on_ctx_resp(server, response) - else: - hooks.on_first_token(server, request) - hooks.on_resp_done(server, request, response) - return response + # The POST is lazy for streaming requests. Keep cleanup around + # the generator so cancellation releases the router reservation. + response_stream = cast( + UCompletionResponseOrGenerator, + _ManagedResponseStream(self, request, resp_generator, req_id), + ) + cleanup_deferred = True + return response_stream + + response = None + async for resp_json in resp_generator: + response = response_type(**resp_json) + if hooks: + if self._role == ServerRole.CONTEXT: + hooks.on_ctx_resp(server, response) + else: + hooks.on_first_token(server, request) + hooks.on_resp_done(server, request, response) + success = True + return response + except asyncio.CancelledError: + self._metrics_collector.error_requests.inc() + raise except Exception: self._metrics_collector.error_requests.inc() - # finish the request upon error - await self._finish_request(request, success=False, req_id=req_id) + raise + finally: + if not cleanup_deferred: + await self._cleanup_request_safely( + request, resp_generator, success=success, req_id=req_id + ) + + async def _cleanup_request_safely( + self, + request: UCompletionRequest, + resp_generator: Optional[AsyncGenerator[Any, None]], + success: bool, + req_id: Optional[int] = None, + ) -> None: + """Complete response and router cleanup before propagating cancellation.""" + + async def cleanup() -> None: + try: + if resp_generator is not None: + await resp_generator.aclose() + finally: + await self._finish_request(request, success=success, req_id=req_id) + + cleanup_task = asyncio.create_task(cleanup()) + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + # The parent deadline may race the aiohttp timeout. Preserve the + # router lease until cleanup finishes, then propagate cancellation. + await cleanup_task raise async def _post_with_retry( @@ -203,15 +326,19 @@ async def _post_with_retry( req_id: Optional[int] = None, ) -> AsyncGenerator[Any, None]: is_stream = request.stream - # Loop range must cover the transient-TCP extended budget (up to 5) - # so the conditional raise inside the except block can actually decide - # to keep retrying. Non-transient errors still raise on the first - # attempt that reaches self._max_retries. - _TRANSIENT_TCP_BUDGET = 5 - loop_max = max(self._max_retries, _TRANSIENT_TCP_BUDGET) + 1 + # Keep the extended budget for transient burst-start TCP races while + # timeout exceptions remain explicitly non-retryable below. + transient_tcp_budget = 5 + loop_max = max(self._max_retries, transient_tcp_budget) + 1 for attempt in range(loop_max): - # Regenerate disagg_request_id on retry to avoid ID collision on workers - if attempt > 0 and self._disagg_id_generator is not None: + # Context retries represent a new transfer attempt and need a fresh + # ID. A generation retry must keep the ID shared with its existing + # context response. + if ( + attempt > 0 + and self._role != ServerRole.GENERATION + and self._disagg_id_generator is not None + ): dp = getattr(request, "disaggregated_params", None) if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = await self._disagg_id_generator() @@ -236,6 +363,36 @@ async def _post_with_retry( headers=req_headers, ) as http_response: content_type = http_response.headers.get("Content-Type", "") + if http_response.status >= 400: + if http_response.status == 504: + try: + error_body = await http_response.text() + except ( + aiohttp.ClientError, + asyncio.TimeoutError, + OSError, + UnicodeError, + ) as e: + # The HTTP status is authoritative. A truncated + # timeout body must not turn a known 504 into a + # retryable transport error and replay work. + logger.warning( + f"Failed to read {self._role} timeout body from {url}: {e}" + ) + error_body = "" + raise UpstreamRequestTimeoutError( + self._role, + self._timeout_secs, + detail=error_body[:2048], + ) + error_body = await http_response.text() + raise aiohttp.ClientResponseError( + http_response.request_info, + http_response.history, + status=http_response.status, + message=f"{http_response.reason}: {error_body[:2048]}", + headers=http_response.headers, + ) if not is_stream and "text/event-stream" in content_type: raise ValueError( "Received an event-stream although request stream was False" @@ -250,24 +407,20 @@ async def _post_with_retry( yield line # don't finish the request here since the response generator is not done yet else: - if http_response.status >= 400: - error_body = await http_response.text() - raise aiohttp.ClientResponseError( - http_response.request_info, - http_response.history, - status=http_response.status, - message=f"{http_response.reason}: {error_body[:2048]}", - headers=http_response.headers, - ) response_dict = await http_response.json() # yield here since python forbids return statements in async generators yield response_dict - # finish the request after the successful response - await self._finish_request(request, req_id=req_id) self._metrics_collector.complete_latency_seconds.observe( get_steady_clock_now_in_seconds() - start_time ) break # break and skip retries if the whole response is processed without exception + except UpstreamRequestTimeoutError: + raise + except asyncio.TimeoutError as e: + logger.error( + f"{self._role} request to {url} timed out after {self._timeout_secs} seconds" + ) + raise UpstreamRequestTimeoutError(self._role, self._timeout_secs) from e except (aiohttp.ClientError, OSError) as e: if lines_yielded > 0: logger.error( @@ -275,22 +428,25 @@ async def _post_with_retry( traceback.format_exc(), ) raise - # Selective retry budget: ServerDisconnectedError and - # ConnectionResetError are transient TCP races (typically at - # burst start when client keepalive vs server keepalive race). - # Give them an extended retry budget while preserving the - # original fail-fast for genuine upstream errors. is_transient_tcp = isinstance( e, (aiohttp.ServerDisconnectedError, ConnectionResetError), ) + is_pre_connect_failure = isinstance(e, aiohttp.ClientConnectorError) effective_max = self._max_retries - if is_transient_tcp: - effective_max = max(self._max_retries, _TRANSIENT_TCP_BUDGET) + if self._role == ServerRole.GENERATION and not is_pre_connect_failure: + # A generation POST is not idempotent once it may have been + # accepted. It also cannot regenerate disagg_request_id: + # the context transfer remains registered under the + # original shared ID. Retry only failures which prove that + # the connection was never established. + effective_max = 0 + elif is_transient_tcp: + effective_max = max(self._max_retries, transient_tcp_budget) if attempt >= effective_max: logger.error( - f"Client error to {url}: {e} - last retry {attempt} of {effective_max}" - "failed", + f"Client error to {url}: {e} - last retry {attempt} of " + f"{effective_max} failed", traceback.format_exc(), ) raise @@ -319,7 +475,6 @@ async def _response_generator( assert "text/event-stream" in http_response.headers.get("Content-Type", ""), ( "Response is not streaming" ) - success = True try: last_token_time = start_time i = 0 @@ -349,16 +504,7 @@ async def _response_generator( except aiohttp.ClientError as e: # a client error is expected when the response stream is done if the connector has close=True logger.error(f"{self._role} client {server} error: {e}") - self._metrics_collector.error_requests.inc() - success = False raise - except Exception: - self._metrics_collector.error_requests.inc() - success = False - raise - finally: - # finish the request after streaming response is done or error is raised - await self._finish_request(request, success=success, req_id=req_id) async def _finish_request( self, diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 5369c0cb6dbf..444959d63dcc 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -15,13 +15,15 @@ # yapf: disable import asyncio +import json import signal import socket import traceback from contextlib import asynccontextmanager -from typing import Callable, Optional +from typing import Any, AsyncGenerator, Callable, Optional import aiohttp +import anyio import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError @@ -39,11 +41,12 @@ from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient, DisaggCoordinatorService) -from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient +from tensorrt_llm.serve.openai_client import (OpenAIClient, OpenAIHttpClient, + UpstreamRequestTimeoutError) from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionRequest, CompletionRequest, UCompletionRequest, + ChatCompletionRequest, CompletionRequest, ErrorResponse, UCompletionRequest, UCompletionResponse, ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, @@ -57,6 +60,34 @@ code: f"\\x{code:02x}" for code in (*range(32), 127) } +REQUEST_CLEANUP_GRACE_SECS = 10.0 + + +class DisaggregatedRequestTimeoutError(TimeoutError): + """The end-to-end disaggregated request deadline expired.""" + + def __init__(self, timeout_secs: int): + self.timeout_secs = timeout_secs + super().__init__( + f"Disaggregated request exceeded the configured " + f"{timeout_secs}-second deadline") + + +class _CleanupStreamingResponse(StreamingResponse): + """Close upstream work even if ASGI disconnects before iteration starts.""" + + def __init__(self, *args, cleanup, cleanup_runner, **kwargs): + super().__init__(*args, **kwargs) + self._cleanup = cleanup + self._cleanup_runner = cleanup_runner + + async def __call__(self, scope, receive, send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + await self._cleanup_runner(self._cleanup, + "stream response cleanup") + class RawRequestResponseHooks(ResponseHooks): def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCollector): @@ -109,6 +140,8 @@ def __init__(self, coordinator_url: Optional[str] = None): self._config = config self._req_timeout_secs = req_timeout_secs + self._cleanup_grace_secs = REQUEST_CLEANUP_GRACE_SECS + self._background_cleanup_tasks: set[asyncio.Task] = set() self._server_start_timeout_secs = server_start_timeout_secs self._metadata_server_cfg = metadata_server_cfg self._metrics_interval_secs = metrics_interval_secs @@ -165,12 +198,7 @@ async def lifespan(app) -> None: # The cluster manager (via setup) owns server preparation + monitoring. await self._service.setup() yield - await self._service.teardown() - if self._perf_metrics_collector._background_tasks: - await asyncio.gather( - *self._perf_metrics_collector._background_tasks, - return_exceptions=True, - ) + await self._shutdown() self.app = FastAPI(lifespan=lifespan) @@ -201,6 +229,58 @@ async def validation_exception_handler(request: Request, exc): self.register_routes() + async def _shutdown(self) -> None: + pending_cleanup = await self._drain_background_cleanup_before_shutdown() + try: + await self._service.teardown() + finally: + await self._observe_background_cleanup_after_shutdown( + pending_cleanup) + if self._perf_metrics_collector._background_tasks: + await asyncio.gather( + *self._perf_metrics_collector._background_tasks, + return_exceptions=True, + ) + + async def _drain_background_cleanup_before_shutdown( + self) -> tuple[asyncio.Task, ...]: + cleanup_tasks = tuple(self._background_cleanup_tasks) + if not cleanup_tasks: + return () + + done, pending = await asyncio.wait( + cleanup_tasks, timeout=self._cleanup_grace_secs) + for task in done: + self._consume_cleanup_task(task, "server shutdown cleanup") + if not pending: + return () + + logger.warning( + f"Canceling {len(pending)} cleanup task(s) that remained active " + "during server shutdown") + for task in pending: + task.cancel() + done, pending = await asyncio.wait( + pending, timeout=self._cleanup_grace_secs) + for task in done: + self._consume_cleanup_task(task, "server shutdown cancellation") + return tuple(pending) + + async def _observe_background_cleanup_after_shutdown( + self, cleanup_tasks: tuple[asyncio.Task, ...]) -> None: + if not cleanup_tasks: + return + for task in cleanup_tasks: + task.cancel() + done, pending = await asyncio.wait( + cleanup_tasks, timeout=self._cleanup_grace_secs) + for task in done: + self._consume_cleanup_task(task, "post-shutdown cleanup") + if pending: + logger.error( + f"{len(pending)} cleanup task(s) did not stop during server shutdown" + ) + def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: async def disagg_id_generator(): return await self._coordinator.get_disagg_request_id() @@ -274,6 +354,8 @@ def _wrap_entry_point(self, entry_point: Callable, request_type: type = UComplet # CompletionRequest first and 400 every chat body, so override the wrapper's # annotation with request_type (as openai_server.py does). async def wrapper(req: request_type, raw_req: Request) -> Response: + deadline = (asyncio.get_running_loop().time() + + self._req_timeout_secs) try: self._perf_metrics_collector.total_requests.inc() if req.stream: @@ -287,18 +369,233 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: raise HTTPException(status_code=400, detail=str(e)) from e self._extract_conversation_id(req, raw_req) hooks = RawRequestResponseHooks(raw_req, self._perf_metrics_collector) - response_or_generator = await entry_point(req, hooks) + response_or_generator = await self._await_response_or_disconnect( + entry_point, req, hooks, raw_req, deadline) self._perf_metrics_collector.total_responses.inc() if req.stream: - return StreamingResponse(content=response_or_generator, media_type="text/event-stream") + upstream_generator = response_or_generator + response_or_generator = self._stream_with_deadline( + upstream_generator, deadline) + return _CleanupStreamingResponse( + content=response_or_generator, + media_type="text/event-stream", + cleanup=upstream_generator.aclose, + cleanup_runner=self._run_cleanup_bounded) else: return JSONResponse(content=response_or_generator.model_dump()) + except asyncio.CancelledError: + raise except Exception as e: - self._handle_exception(e) + return self._handle_exception(e) return wrapper + async def _await_response_or_disconnect(self, entry_point: Callable, + req: UCompletionRequest, + hooks: ResponseHooks, + raw_req: Request, + deadline: float): + request_task = asyncio.create_task(entry_point(req, hooks)) + disconnect_task = asyncio.create_task( + self._wait_for_disconnect(raw_req)) + request_task_observed = False + streaming_response_ready = False + try: + remaining = max(0, deadline - asyncio.get_running_loop().time()) + done, _ = await asyncio.wait( + (request_task, disconnect_task), + timeout=remaining, + return_when=asyncio.FIRST_COMPLETED) + # Prefer a completed response if completion and disconnect race. + if request_task in done: + request_task_observed = True + response = await request_task + streaming_response_ready = getattr(req, "stream", False) + return response + + if disconnect_task in done: + await disconnect_task + raise asyncio.CancelledError() + raise DisaggregatedRequestTimeoutError(self._req_timeout_secs) + finally: + if not request_task_observed: + await self._run_cleanup_bounded( + lambda: self._cancel_and_close_unclaimed_request_task( + request_task), "unclaimed request cleanup") + if streaming_response_ready: + # StreamingResponse becomes the sole ASGI receive owner after + # this method returns. Never background the previous listener, + # or the two consumers can race and lose a disconnect event. + await self._cancel_and_wait_strict( + disconnect_task, "stream receive handoff") + else: + await self._cancel_and_wait(disconnect_task) + + @staticmethod + async def _cancel_and_close_unclaimed_request_task( + request_task: asyncio.Task) -> None: + """Cancel setup and close a stream returned while cancellation races.""" + if not request_task.done(): + request_task.cancel() + try: + response = await request_task + except asyncio.CancelledError: + return + except Exception: + logger.debug("Suppressed unclaimed request exception during cleanup: " + f"{traceback.format_exc()}") + return + + close_response = getattr(response, "aclose", None) + if close_response is not None: + await close_response() + + async def _stream_with_deadline( + self, response_generator: AsyncGenerator[Any, None], + deadline: float) -> AsyncGenerator[Any, None]: + iterator = response_generator.__aiter__() + try: + while True: + next_task = asyncio.create_task(anext(iterator)) + try: + remaining = max( + 0, deadline - asyncio.get_running_loop().time()) + done, _ = await asyncio.wait( + (next_task, ), + timeout=remaining, + return_when=asyncio.FIRST_COMPLETED) + if next_task in done: + try: + yield await next_task + except StopAsyncIteration: + return + except UpstreamRequestTimeoutError as e: + async for chunk in self._stream_timeout_response(e): + yield chunk + return + continue + + await self._cancel_and_wait(next_task) + timeout_error = DisaggregatedRequestTimeoutError( + self._req_timeout_secs) + async for chunk in self._stream_timeout_response( + timeout_error): + yield chunk + return + except asyncio.CancelledError: + # Starlette cancels this stream from an AnyIO cancel scope + # when the client disconnects. Keep this cleanup local and + # shielded so the upstream read reaches a terminal state. + with anyio.CancelScope(shield=True): + await self._cancel_and_wait(next_task) + raise + finally: + await self._run_cleanup_bounded(iterator.aclose, + "upstream stream cleanup") + + @staticmethod + async def _wait_for_disconnect(raw_req: Request) -> None: + """Own the ASGI receive channel until StreamingResponse takes over. + + FastAPI has consumed the request body before entering this handler. + This watcher exclusively owns the receive channel until it is canceled + before ``StreamingResponse`` assumes ownership. Blocking on the channel + avoids polling during that handoff. + """ + while True: + message = await raw_req.receive() + if message["type"] == "http.disconnect": + return + + async def _cancel_and_wait(self, *tasks: asyncio.Task) -> None: + for task in tasks: + if task is not None and not task.done(): + task.cancel() + with anyio.CancelScope(shield=True): + await self._wait_for_cleanup_tasks(tasks, + "canceled request cleanup") + + @staticmethod + async def _cancel_and_wait_strict(task: asyncio.Task, label: str) -> None: + if not task.done(): + task.cancel() + with anyio.CancelScope(shield=True): + try: + await task + except asyncio.CancelledError: + pass + except Exception: + logger.debug(f"Suppressed exception during {label}: " + f"{traceback.format_exc()}") + + async def _run_cleanup_bounded(self, cleanup, label: str) -> None: + cleanup_task = asyncio.create_task(cleanup()) + with anyio.CancelScope(shield=True): + await self._wait_for_cleanup_tasks((cleanup_task, ), label) + + async def _wait_for_cleanup_tasks(self, tasks, label: str) -> None: + cleanup_tasks = tuple(task for task in tasks if task is not None) + if not cleanup_tasks: + return + done, pending = await asyncio.wait( + cleanup_tasks, + timeout=getattr(self, "_cleanup_grace_secs", + REQUEST_CLEANUP_GRACE_SECS), + ) + for task in done: + self._consume_cleanup_task(task, label) + for task in pending: + self._track_background_cleanup(task, label) + if pending: + cleanup_grace_secs = getattr(self, "_cleanup_grace_secs", + REQUEST_CLEANUP_GRACE_SECS) + logger.warning( + f"{label} exceeded the {cleanup_grace_secs}-second cleanup grace; " + f"continuing cleanup in the background") + + def _track_background_cleanup(self, task: asyncio.Task, + label: str) -> None: + background_tasks = self.__dict__.setdefault( + "_background_cleanup_tasks", set()) + if task in background_tasks: + return + background_tasks.add(task) + + def cleanup_done(completed_task: asyncio.Task) -> None: + background_tasks.discard(completed_task) + self._consume_cleanup_task(completed_task, label) + + task.add_done_callback(cleanup_done) + + @staticmethod + def _consume_cleanup_task(task: asyncio.Task, label: str) -> None: + try: + task.result() + except asyncio.CancelledError: + pass + except Exception: + logger.debug(f"Suppressed exception during {label}: " + f"{traceback.format_exc()}") + + @staticmethod + async def _stream_timeout_response( + exception: TimeoutError) -> AsyncGenerator[str, None]: + error_response = ErrorResponse(message=str(exception), + type="RequestTimeoutError", + code=504) + yield f"data: {json.dumps({'error': error_response.model_dump()})}\n\n" + yield "data: [DONE]\n\n" + def _handle_exception(self, exception): - if isinstance(exception, CppExecutorError): + if isinstance(exception, (DisaggregatedRequestTimeoutError, + UpstreamRequestTimeoutError)): + self._perf_metrics_collector.http_exceptions.inc() + logger.error(f"Request timeout: {exception}") + error_response = ErrorResponse(message=str(exception), + type="RequestTimeoutError", + code=504) + return JSONResponse(content=error_response.model_dump(), + status_code=error_response.code) + elif isinstance(exception, CppExecutorError): logger.error("CppExecutorError: ", traceback.format_exc()) signal.raise_signal(signal.SIGINT) elif isinstance(exception, HTTPException): diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e2ee00257e98..176a18248a59 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -42,6 +42,70 @@ _GEN_PENDING_FINISH_REASONS = ("length", "not_finished") +class _RouterReservation: + """Track a router reservation until a client accepts its ownership.""" + + def __init__( + self, + request: Optional[UCompletionRequest] = None, + pending: bool = False, + req_id: Optional[int] = None, + ): + self.request = request + self.pending = pending + self.req_id = req_id + + def mark_pending(self, request: UCompletionRequest, req_id: Optional[int] = None) -> None: + self.request = request + self.pending = True + self.req_id = req_id + + def clear(self) -> None: + self.pending = False + + +class _QueuedResponseStream: + """Own a background response consumer even before first iteration.""" + + def __init__(self, queue: asyncio.Queue, consume_task: asyncio.Task, upstream_stream): + self._queue = queue + self._consume_task = consume_task + self._upstream_stream = upstream_stream + self._cleanup_lock = asyncio.Lock() + self._closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._closed: + raise StopAsyncIteration + try: + item = await self._queue.get() + except asyncio.CancelledError: + await self.aclose() + raise + if item is None: + await self.aclose() + raise StopAsyncIteration + if isinstance(item, Exception): + await self.aclose() + raise item + return item + + async def aclose(self) -> None: + async with self._cleanup_lock: + if self._closed: + return + try: + if not self._consume_task.done(): + self._consume_task.cancel() + await asyncio.gather(self._consume_task, return_exceptions=True) + await self._upstream_stream.aclose() + finally: + self._closed = True + + class OpenAIDisaggregatedService(OpenAIService): def __init__( self, @@ -115,6 +179,32 @@ async def openai_chat_completion( async def _send_disagg_request_ctx_first( self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None + ) -> UCompletionResponseOrGenerator: + ctx_reservation = _RouterReservation() + gen_reservation = _RouterReservation() + try: + return await self._send_disagg_request_ctx_first_impl( + request, hooks, ctx_reservation, gen_reservation + ) + except asyncio.CancelledError: + await self._release_router_reservations( + (self._ctx_router, ctx_reservation), + (self._gen_router, gen_reservation), + ) + raise + except Exception: + await self._release_router_reservations( + (self._ctx_router, ctx_reservation), + (self._gen_router, gen_reservation), + ) + raise + + async def _send_disagg_request_ctx_first_impl( + self, + request: UCompletionRequest, + hooks: Optional[ResponseHooks], + ctx_reservation: _RouterReservation, + gen_reservation: _RouterReservation, ) -> UCompletionResponseOrGenerator: # ctx_response contains a http response with ContextPhaseParams attached after prefill compute is done @@ -124,7 +214,13 @@ async def _send_disagg_request_ctx_first( ctx_server = None disagg_request_id = await self._coordinator.get_disagg_request_id() # reserve a gen_server if conditional disagg is needed - gen_server, need_ctx = await self._check_conditional_disagg(request, disagg_request_id) + gen_server, need_ctx = await self._check_conditional_disagg( + request, disagg_request_id, gen_reservation + ) + if gen_server: + gen_reservation.mark_pending(request, req_id=disagg_request_id) + else: + gen_reservation.clear() # Context retries may replace disagg_request_id for the KV-transfer # handshake. Keep the ID used to reserve the generation server separate # so its coordinator-side load is released under the original key. @@ -140,9 +236,11 @@ async def _send_disagg_request_ctx_first( hooks.on_ctx_dispatch(request) ctx_req = self._get_ctx_request(request, disagg_request_id) # ctx generator is empty + ctx_reservation.mark_pending(ctx_req, req_id=disagg_request_id) ctx_server, _ = await self._ctx_router.get_next_server( ctx_req, exclude_server=gen_server, req_id=disagg_request_id ) + ctx_reservation.clear() ctx_response = await self._ctx_client.send_request( ctx_req, server=ctx_server, hooks=hooks, req_id=disagg_request_id ) @@ -152,10 +250,9 @@ async def _send_disagg_request_ctx_first( disagg_request_id = ctx_response_disagg_params.disagg_request_id gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) except Exception: - if gen_server: - await self._gen_router.finish_request( - request, success=False, req_id=gen_reservation_id - ) + await self._release_router_reservations( + (self._gen_router, gen_reservation), + ) raise else: # When need_ctx=False the gen server handles full generation and @@ -169,17 +266,18 @@ async def _send_disagg_request_ctx_first( gen_req.disaggregated_params = None if ctx_response is None or self._need_gen(ctx_response): if not gen_server: + gen_reservation.mark_pending(gen_req, req_id=disagg_request_id) gen_server, _ = await self._gen_router.get_next_server( gen_req, exclude_server=ctx_server, req_id=disagg_request_id ) gen_reservation_id = disagg_request_id + gen_reservation.clear() gen_response = await self._gen_client.send_request( gen_req, server=gen_server, hooks=hooks, req_id=gen_reservation_id ) return gen_response else: - if gen_server: - await self._gen_router.finish_request(request, req_id=gen_reservation_id) + await self._release_router_reservation(self._gen_router, gen_reservation) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator @@ -278,7 +376,12 @@ def _get_gen_request( request.disaggregated_params.disagg_request_id = disagg_request_id return request - async def _check_conditional_disagg(self, request: UCompletionRequest, req_id: int) -> bool: + async def _check_conditional_disagg( + self, + request: UCompletionRequest, + req_id: int, + reservation: Optional[_RouterReservation] = None, + ) -> tuple[Optional[str], bool]: if self.conditional_disagg_config: local_gen_router = ( self._gen_router._local @@ -291,6 +394,8 @@ async def _check_conditional_disagg(self, request: UCompletionRequest, req_id: i ) # Query kv cache status and select a best gen_server. # The server is reserved for generation request + if reservation is not None: + reservation.mark_pending(request, req_id=req_id) gen_server, info = await self._gen_router.get_next_server(request, req_id=req_id) match_length = info["match_length"] total_length = info["num_tokens"] @@ -348,9 +453,20 @@ async def setup(self) -> None: await self._coordinator.start() async def teardown(self) -> None: - await self._ctx_client.shutdown() - await self._gen_client.shutdown() - await self._coordinator.stop() + results = await asyncio.gather( + self._ctx_client.shutdown(), + self._gen_client.shutdown(), + return_exceptions=True, + ) + try: + await self._coordinator.stop() + except Exception as error: + results.append(error) + errors = [result for result in results if isinstance(result, BaseException)] + for error in errors: + logger.warning(f"Disaggregated service teardown failed: {error}") + if errors: + raise errors[0] async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None: if ctx_response: @@ -381,6 +497,26 @@ async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None: async def _send_disagg_request_gen_first( self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None + ) -> UCompletionResponse: + ctx_reservation = _RouterReservation() + try: + return await self._send_disagg_request_gen_first_impl(request, hooks, ctx_reservation) + except asyncio.CancelledError: + await self._release_router_reservations( + (self._ctx_router, ctx_reservation), + ) + raise + except Exception: + await self._release_router_reservations( + (self._ctx_router, ctx_reservation), + ) + raise + + async def _send_disagg_request_gen_first_impl( + self, + request: UCompletionRequest, + hooks: Optional[ResponseHooks], + ctx_reservation: _RouterReservation, ) -> UCompletionResponse: if hooks: hooks.on_req_begin(request) @@ -396,10 +532,11 @@ async def _send_disagg_request_gen_first( # arrival->here = pre-ctx wait in the orchestrator/fleet. if hooks: hooks.on_ctx_dispatch(request) + ctx_req = self._get_ctx_request(request, disagg_request_id) + ctx_reservation.mark_pending(ctx_req, req_id=disagg_request_id) ctx_server, ctx_server_info = await self._ctx_router.get_next_server( - request, req_id=disagg_request_id + ctx_req, req_id=disagg_request_id ) - ctx_req = self._get_ctx_request(request, disagg_request_id) gen_req = self._get_gen_request( request, ctx_response=None, @@ -432,56 +569,40 @@ async def _consume_gen(): await queue.put(None) # sentinel consume_task: asyncio.Task = asyncio.create_task(_consume_gen()) + response_stream = _QueuedResponseStream(queue, consume_task, gen_response) # Now send ctx request — gen server has received its request try: + ctx_reservation.clear() await self._ctx_client.send_request( ctx_req, server=ctx_server, hooks=hooks, req_id=disagg_request_id, ) + except asyncio.CancelledError: + await response_stream.aclose() + raise except Exception: - consume_task.cancel() - try: - await consume_task - except (asyncio.CancelledError, Exception): - pass + await response_stream.aclose() raise - - async def _yield_from_queue(): - try: - while True: - item = await queue.get() - if item is None: - break - if isinstance(item, Exception): - raise item - yield item - finally: - if not consume_task.done(): - consume_task.cancel() - try: - await consume_task - except asyncio.CancelledError: - pass - - return _yield_from_queue() + return response_stream else: # Non-streaming or no ctx needed: both HTTP POSTs fire eagerly # through generator consumption, so asyncio.gather works fine. tasks = [] if need_ctx: - tasks.append( - asyncio.create_task( - self._ctx_client.send_request( - ctx_req, - server=ctx_server, - hooks=hooks, - req_id=disagg_request_id, - ) + + async def _send_ctx_request(): + ctx_reservation.clear() + return await self._ctx_client.send_request( + ctx_req, + server=ctx_server, + hooks=hooks, + req_id=disagg_request_id, ) - ) + + tasks.append(asyncio.create_task(_send_ctx_request())) tasks.append( asyncio.create_task( self._gen_client.send_request( @@ -492,5 +613,45 @@ async def _yield_from_queue(): ) ) ) - responses = await asyncio.gather(*tasks) + try: + responses = await asyncio.gather(*tasks) + except asyncio.CancelledError: + await self._cancel_tasks(*tasks) + raise + except Exception: + await self._cancel_tasks(*tasks) + raise return responses[-1] + + @staticmethod + async def _release_router_reservation(router: Router, reservation: _RouterReservation) -> None: + if reservation.pending and reservation.request is not None: + try: + await router.finish_request( + reservation.request, + success=False, + req_id=reservation.req_id, + ) + finally: + reservation.clear() + + @classmethod + async def _release_router_reservations(cls, *reservations) -> None: + results = await asyncio.gather( + *( + cls._release_router_reservation(router, reservation) + for router, reservation in reservations + ), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + logger.warning(f"Failed to release router reservation: {result}") + + @staticmethod + async def _cancel_tasks(*tasks: asyncio.Task) -> None: + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tests/integration/defs/disaggregated/disagg_test_utils.py b/tests/integration/defs/disaggregated/disagg_test_utils.py index 320f08c4a278..90a33232ac24 100644 --- a/tests/integration/defs/disaggregated/disagg_test_utils.py +++ b/tests/integration/defs/disaggregated/disagg_test_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -208,7 +208,9 @@ def run_gen_worker( ) -def run_disagg_server(disagg_cluster_config, work_dir, port=0, save_log=False, env=None, cwd=None): +def run_disagg_server( + disagg_cluster_config, work_dir, port=0, save_log=False, env=None, cwd=None, request_timeout=180 +): """Launch the disaggregated server. Args: @@ -217,6 +219,8 @@ def run_disagg_server(disagg_cluster_config, work_dir, port=0, save_log=False, e port: Port number save_log: Whether to save logs to file env: Environment variables for the subprocess + cwd: Working directory for the subprocess + request_timeout: End-to-end request timeout in seconds Returns: ProcessWrapper: Wrapped subprocess @@ -225,7 +229,14 @@ def run_disagg_server(disagg_cluster_config, work_dir, port=0, save_log=False, e disagg_cluster_config["port"] = port with open(disagg_server_config_path, "w+") as f: yaml.dump(disagg_cluster_config, f) - cmds = ["trtllm-serve", "disaggregated", "-c", disagg_server_config_path] + cmds = [ + "trtllm-serve", + "disaggregated", + "-c", + disagg_server_config_path, + "--request_timeout", + str(request_timeout), + ] log_file = None log_path = None # See WAR rationale in _run_worker above (nvbugs/5821433). diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 9347bc090c0b..170e2836189b 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -22,7 +22,7 @@ import subprocess import tempfile import time -from collections import namedtuple +from collections import Counter, namedtuple from dataclasses import dataclass from typing import Any, Optional @@ -58,6 +58,11 @@ class TestConfig: cancellation_rate: Optional[int] = None cancellation_delay: Optional[float] = None concurrency: int = 512 + accuracy_concurrency: Optional[int] = None + request_timeout: int = 180 + heartbeat_interval_sec: Optional[int] = None + inactive_timeout_sec: Optional[int] = None + post_stress_recovery_timeout: Optional[int] = None def __str__(self): return self.test_desc @@ -100,6 +105,458 @@ def cleanup_output_files(): "out of memory", ) +_ACCURACY_CLIENT_TIMEOUT_GRACE_SECS = 60 +_ACCURACY_PROCESS_TIMEOUT_SECS = 1_200 +# The proxy can spend up to 10 seconds in bounded request cleanup after its +# request deadline; reserve another 5 seconds before AIPerf cancels the client +# request, then 5 seconds for recording the terminal result. +_AIPERF_CLIENT_TIMEOUT_GRACE_SECS = 15 +_AIPERF_TERMINAL_OBSERVATION_GRACE_SECS = 5 +_RECOVERY_PROBE_INTERVAL_SECS = 5 +_RECOVERY_PROBE_MIN_SUCCESSES = 6 +_RECOVERY_PROBE_ROUNDS_PER_WORKER = 3 +_RECOVERY_PROBE_OUTPUT_TOKENS = 32 +_RECOVERY_PROBE_PROMPT = "Assess service recovery readiness. " * 64 + + +def _raise_if_process_stopped(processes, phase: str) -> None: + stopped_processes = [ + proc for proc in processes if proc.process.poll() is not None + ] + if not stopped_processes: + return + + details = [ + f"{proc.log_path or 'process'} (rc={proc.process.poll()})" + for proc in stopped_processes + ] + raise RuntimeError(f"Disaggregated process exited during {phase}: " + f"{details}") + + +def wait_for_disagg_recovery(server_url: str, model_path: str, timeout: int, + request_timeout: int, required_successes: int, + workers, disagg_server) -> None: + """Wait until fresh requests can traverse the disaggregated path. + + A client-side 504 is terminal for the caller, but with in-flight KV + cancellation disabled a worker may still own that request until its + transfer finishes naturally. Starting the full accuracy phase immediately + can therefore turn its requests into another timeout burst. + + This is deliberately a behavioral readiness gate, not an assertion that + every internal queue is empty. Existing worker gauges omit asynchronous KV + transfers and can report zero during precisely this drain period. Several + consecutive cache-distinct probes, with a target scaled to the worker + count, demonstrate that nontrivial new work can again traverse context + transfer and generation before GSM8K starts. This does not prove that each + worker is idle or selected; the accuracy phase remains the authoritative + recovery check. + """ + import requests as http_requests + + if timeout <= 0: + raise ValueError("post_stress_recovery_timeout must be positive") + if request_timeout <= 0: + raise ValueError("recovery probe request_timeout must be positive") + if required_successes <= 0: + raise ValueError("recovery probe required_successes must be positive") + + deadline = time.monotonic() + timeout + start_time = time.monotonic() + consecutive_successes = 0 + attempt = 0 + last_outcome = "probe not attempted" + probe_url = f"{server_url}/v1/completions" + probe_run_id = time.monotonic_ns() + monitored_processes = [*workers, disagg_server] + + while time.monotonic() < deadline: + _raise_if_process_stopped(monitored_processes, "recovery") + + attempt += 1 + probe_prompt = (f"Recovery probe {probe_run_id}-{attempt}. " + + _RECOVERY_PROBE_PROMPT) + probe_body = { + "model": model_path, + "prompt": probe_prompt, + "max_tokens": _RECOVERY_PROBE_OUTPUT_TOKENS, + "temperature": 0.0, + "stream": False, + "ignore_eos": True, + } + remaining = deadline - time.monotonic() + if remaining <= 0: + break + per_attempt_timeout = min(request_timeout, remaining) + probe_start = time.monotonic() + try: + response = http_requests.post(probe_url, + json=probe_body, + timeout=per_attempt_timeout) + probe_latency = time.monotonic() - probe_start + if time.monotonic() >= deadline: + last_outcome = ( + f"probe completed after the {timeout}s recovery deadline") + break + if response.status_code == 200: + payload = response.json() + choices = payload.get("choices") if isinstance(payload, + dict) else None + if not isinstance(choices, list) or not choices: + raise ValueError( + "successful recovery probe returned no choices") + first_choice = choices[0] + if (not isinstance(first_choice, dict) + or not isinstance(first_choice.get("text"), str) + or not first_choice["text"].strip()): + raise ValueError( + "successful recovery probe returned no generated text") + usage = payload.get("usage") + if (not isinstance(usage, dict) + or usage.get("completion_tokens") + != _RECOVERY_PROBE_OUTPUT_TOKENS): + raise ValueError( + "successful recovery probe did not generate exactly " + f"{_RECOVERY_PROBE_OUTPUT_TOKENS} tokens") + consecutive_successes += 1 + last_outcome = (f"HTTP 200 in {probe_latency:.1f}s, " + f"streak={consecutive_successes}/" + f"{required_successes}") + else: + consecutive_successes = 0 + body = response.text.replace("\n", " ")[:300] + last_outcome = (f"HTTP {response.status_code} in " + f"{probe_latency:.1f}s: {body}") + except (http_requests.RequestException, ValueError) as error: + consecutive_successes = 0 + probe_latency = time.monotonic() - probe_start + last_outcome = (f"{type(error).__name__} after " + f"{probe_latency:.1f}s: {error}") + + elapsed = time.monotonic() - start_time + logger.info(f"Post-stress recovery probe {attempt}: " + f"elapsed={elapsed:.1f}s, {last_outcome}") + _raise_if_process_stopped(monitored_processes, "recovery") + if consecutive_successes >= required_successes: + logger.info("Disaggregated service recovered after " + f"{elapsed:.1f}s and {attempt} probe(s)") + return + + remaining = deadline - time.monotonic() + if remaining > 0: + time.sleep(min(_RECOVERY_PROBE_INTERVAL_SECS, remaining)) + + raise TimeoutError( + f"Disaggregated service did not recover within {timeout}s after " + f"{attempt} probe(s); last outcome: {last_outcome}") + + +def remove_stale_aiperf_terminal_records(artifact_dir: str) -> None: + records_path = os.path.join(artifact_dir, "profile_export.jsonl") + try: + os.remove(records_path) + except FileNotFoundError: + pass + + +def validate_aiperf_terminal_records(artifact_dir: str, expected_count: int, + client_timeout: int) -> None: + """Require one terminal AIPerf record for every profiled request. + + The stress phase intentionally allows overload errors and injected client + cancellations, so this validates terminal accounting rather than imposing + a success-rate SLO. The follow-on accuracy phase validates recovery and + semantic correctness. + """ + if expected_count <= 0: + raise ValueError("expected AIPerf record count must be positive") + if client_timeout <= 0: + raise ValueError("AIPerf client timeout must be positive") + + records_path = os.path.join(artifact_dir, "profile_export.jsonl") + if not os.path.isfile(records_path): + raise AssertionError( + f"AIPerf terminal-record export is missing: {records_path}") + + status_counts = Counter() + request_ids = set() + cancelled_count = 0 + record_count = 0 + max_terminal_latency = 0.0 + max_terminal_latency_allowed = (client_timeout + + _AIPERF_TERMINAL_OBSERVATION_GRACE_SECS) + with open(records_path, "r", encoding="utf-8") as records_file: + for line_number, line in enumerate(records_file, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise AssertionError( + f"Invalid AIPerf record at {records_path}:{line_number}: " + f"{error}") from error + if not isinstance(record, dict): + raise AssertionError( + f"AIPerf record {line_number} is not an object") + + metadata = record.get("metadata") + if not isinstance(metadata, dict): + raise AssertionError( + f"AIPerf record {line_number} has no metadata") + request_id = metadata.get("x_request_id") + if not isinstance(request_id, str) or not request_id: + raise AssertionError( + f"AIPerf record {line_number} has no request ID") + request_end_ns = metadata.get("request_end_ns") + if not isinstance(request_end_ns, int) or request_end_ns <= 0: + raise AssertionError( + f"AIPerf record {line_number} is not terminal") + request_start_ns = metadata.get("request_start_ns") + if not isinstance(request_start_ns, int) or request_start_ns <= 0: + raise AssertionError( + f"AIPerf record {line_number} has no request start time") + terminal_latency = (request_end_ns - request_start_ns) / 1e9 + if (terminal_latency < 0 + or terminal_latency > max_terminal_latency_allowed): + raise AssertionError( + f"AIPerf record {line_number} terminalized in " + f"{terminal_latency:.3f}s, outside the allowed range " + f"[0, {max_terminal_latency_allowed}]s") + max_terminal_latency = max(max_terminal_latency, terminal_latency) + if request_id in request_ids: + raise AssertionError( + f"AIPerf request ID {request_id} appears more than once") + request_ids.add(request_id) + + error = record.get("error") + if error is None: + status_counts["success"] += 1 + elif isinstance(error, dict): + status_counts[f"error:{error.get('code', 'unknown')}"] += 1 + else: + status_counts["error:malformed"] += 1 + if metadata.get("was_cancelled") is True: + cancelled_count += 1 + record_count += 1 + + logger.info("AIPerf terminal accounting: " + f"records={record_count}/{expected_count}, " + f"cancelled={cancelled_count}, " + f"max_latency={max_terminal_latency:.3f}s, " + f"statuses={dict(status_counts)}") + if record_count != expected_count: + raise AssertionError( + "AIPerf did not produce exactly one terminal record per profiled " + f"request: expected={expected_count}, actual={record_count}, " + f"statuses={dict(status_counts)}") + + +def test_validate_aiperf_terminal_records(tmp_path): + records_path = tmp_path / "profile_export.jsonl" + + with pytest.raises(ValueError, match="record count"): + validate_aiperf_terminal_records(str(tmp_path), 0, client_timeout=195) + with pytest.raises(ValueError, match="client timeout"): + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=0) + + def make_record(request_id, + start_ns=1_000_000_000, + end_ns=2_000_000_000, + error=None, + cancelled=False): + return { + "metadata": { + "x_request_id": request_id, + "request_start_ns": start_ns, + "request_end_ns": end_ns, + "was_cancelled": cancelled, + }, + "error": error, + } + + def write_records(records): + records_path.write_text("".join(f"{json.dumps(record)}\n" + for record in records), + encoding="utf-8") + + valid_records = [ + make_record("success"), + make_record("cancelled", error={"code": 499}, cancelled=True), + ] + write_records(valid_records) + validate_aiperf_terminal_records(str(tmp_path), 2, client_timeout=195) + + write_records([ + make_record("bounded-cleanup", end_ns=(1 + 190) * 1_000_000_000), + ]) + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=195) + + with pytest.raises(AssertionError, match="expected=3"): + validate_aiperf_terminal_records(str(tmp_path), 3, client_timeout=195) + + write_records([valid_records[0], valid_records[0]]) + with pytest.raises(AssertionError, match="appears more than once"): + validate_aiperf_terminal_records(str(tmp_path), 2, client_timeout=195) + + write_records([make_record("nonterminal", end_ns=None)]) + with pytest.raises(AssertionError, match="is not terminal"): + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=195) + + write_records([ + make_record("late", end_ns=(1 + 201) * 1_000_000_000), + ]) + with pytest.raises(AssertionError, match="outside the allowed range"): + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=195) + + records_path.write_text("{invalid json\n", encoding="utf-8") + with pytest.raises(AssertionError, match="Invalid AIPerf record"): + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=195) + + remove_stale_aiperf_terminal_records(str(tmp_path)) + assert not records_path.exists() + with pytest.raises(AssertionError, match="export is missing"): + validate_aiperf_terminal_records(str(tmp_path), 1, client_timeout=195) + + +def test_wait_for_disagg_recovery(monkeypatch): + import requests as http_requests + + class FakeProcess: + + def __init__(self, return_code=None): + self.return_code = return_code + + def poll(self): + return self.return_code + + class FakeWrapper: + + def __init__(self, return_code=None): + self.process = FakeProcess(return_code) + self.log_path = "fake.log" + + class FakeResponse: + + def __init__(self, status_code, completion_tokens=32): + self.status_code = status_code + self.text = "error" if status_code != 200 else "" + self.completion_tokens = completion_tokens + + def json(self): + return { + "choices": [{ + "text": "generated output" + }], + "usage": { + "completion_tokens": self.completion_tokens, + }, + } + + responses = [ + FakeResponse(504), + FakeResponse(200, completion_tokens=31), + FakeResponse(200), + FakeResponse(200), + ] + observed_prompts = [] + + def post(_url, json, timeout): + assert timeout > 0 + observed_prompts.append(json["prompt"]) + return responses.pop(0) + + monkeypatch.setattr(http_requests, "post", post) + monkeypatch.setattr(time, "sleep", lambda _seconds: None) + workers = [FakeWrapper(), FakeWrapper()] + disagg_server = FakeWrapper() + with pytest.raises(ValueError, match="post_stress_recovery_timeout"): + wait_for_disagg_recovery("http://server", + "model", + timeout=0, + request_timeout=1, + required_successes=1, + workers=workers, + disagg_server=disagg_server) + with pytest.raises(ValueError, match="request_timeout"): + wait_for_disagg_recovery("http://server", + "model", + timeout=1, + request_timeout=0, + required_successes=1, + workers=workers, + disagg_server=disagg_server) + with pytest.raises(ValueError, match="required_successes"): + wait_for_disagg_recovery("http://server", + "model", + timeout=1, + request_timeout=1, + required_successes=0, + workers=workers, + disagg_server=disagg_server) + + wait_for_disagg_recovery("http://server", + "model", + timeout=2, + request_timeout=1, + required_successes=2, + workers=workers, + disagg_server=disagg_server) + assert len(observed_prompts) == 4 + assert len(set(observed_prompts)) == 4 + + with pytest.raises(RuntimeError, match="exited during recovery"): + wait_for_disagg_recovery("http://server", + "model", + timeout=1, + request_timeout=1, + required_successes=1, + workers=[FakeWrapper(return_code=1)], + disagg_server=disagg_server) + + +def test_wait_for_disagg_recovery_rejects_post_deadline_response(monkeypatch): + import requests as http_requests + + class FakeProcess: + + def poll(self): + return None + + class FakeWrapper: + process = FakeProcess() + log_path = "fake.log" + + class FakeResponse: + status_code = 200 + text = "" + + @staticmethod + def json(): + return { + "choices": [{ + "text": "generated output" + }], + "usage": { + "completion_tokens": _RECOVERY_PROBE_OUTPUT_TOKENS, + }, + } + + def slow_post(_url, json, timeout): + del json, timeout + time.sleep(0.02) + return FakeResponse() + + monkeypatch.setattr(http_requests, "post", slow_post) + with pytest.raises(TimeoutError, match="deadline"): + wait_for_disagg_recovery("http://server", + "model", + timeout=0.005, + request_timeout=1, + required_successes=1, + workers=[FakeWrapper()], + disagg_server=FakeWrapper()) + def scan_logs_for_fatal_errors(processes): """Scan saved process logs for fatal disagg/worker error patterns. @@ -141,7 +598,7 @@ def build_worker_diag(workers, disagg_server): string if all workers are healthy. """ all_procs = list(workers) + [disagg_server] - crashed = _crashed_workers(workers) + crashed = _crashed_workers(all_procs) fatal = scan_logs_for_fatal_errors(all_procs) diag = "" if crashed: @@ -654,6 +1111,9 @@ def setup_disagg_cluster( save_log: bool = False, startup_callback=None, startup_tick: int = 30, + request_timeout: int = 180, + heartbeat_interval_sec: int | None = None, + inactive_timeout_sec: int | None = None, ) -> tuple[dict[str, Any], list[ProcessWrapper], list[ProcessWrapper], ProcessWrapper, int, str]: """Load config, launch workers + disagg server, wait for ready. @@ -664,6 +1124,12 @@ def setup_disagg_cluster( env: Environment variables to pass to subprocess (workers and disagg server) server_start_timeout: Timeout in seconds for server to become ready schedule_style: Disagg schedule style ('context_first' or 'generation_first') + save_log: Whether to save worker and server logs + startup_callback: Optional callback invoked while startup is in progress + startup_tick: Interval in seconds between startup callback invocations + request_timeout: End-to-end disaggregated request timeout in seconds + heartbeat_interval_sec: Optional worker heartbeat interval override + inactive_timeout_sec: Optional worker inactivity timeout override Returns: tuple: (config, ctx_workers, gen_workers, disagg_server, server_port, work_dir) @@ -679,6 +1145,20 @@ def setup_disagg_cluster( speculative_model) disagg_cluster = get_default_disagg_cluster_config() + if ((heartbeat_interval_sec is None) != (inactive_timeout_sec is None)): + raise ValueError("heartbeat_interval_sec and inactive_timeout_sec " + "must be overridden together") + if heartbeat_interval_sec is not None and inactive_timeout_sec is not None: + if not 0 < heartbeat_interval_sec < inactive_timeout_sec: + raise ValueError("Expected 0 < heartbeat_interval_sec < " + "inactive_timeout_sec") + disagg_cluster["heartbeat_interval_sec"] = heartbeat_interval_sec + disagg_cluster["inactive_timeout_sec"] = inactive_timeout_sec + logger.info("Disagg cluster liveness config: " + f"heartbeat_interval_sec=" + f"{disagg_cluster['heartbeat_interval_sec']}, " + f"inactive_timeout_sec=" + f"{disagg_cluster['inactive_timeout_sec']}") server_host = config.get("hostname", "localhost") server_port = get_free_port() if save_log: @@ -800,7 +1280,8 @@ def setup_disagg_cluster( server_port, save_log=save_log, env=server_env, - cwd=cwd) + cwd=cwd, + request_timeout=request_timeout) all_workers = ctx_workers + gen_workers @@ -1115,7 +1596,9 @@ def test_disaggregated_benchmark_gen_only_insufficient_kv( try: client = openai.OpenAI(api_key="tensorrt_llm", - base_url=f"http://localhost:{server_port}/v1") + base_url=f"http://localhost:{server_port}/v1", + timeout=60.0, + max_retries=0) # Send 64 concurrent requests to trigger the benchmark fill loop # and the insufficient KV cache error. @@ -1137,13 +1620,29 @@ def send_request(): except Exception as e: return e - with concurrent.futures.ThreadPoolExecutor(max_workers=64) as pool: + pool = concurrent.futures.ThreadPoolExecutor(max_workers=64) + try: futures = [pool.submit(send_request) for _ in range(64)] - results = [f.result(timeout=120) for f in futures] + _, not_done = concurrent.futures.wait(futures, timeout=120) + if not_done: + pool.shutdown(wait=False, cancel_futures=True) + pool = None + pytest.fail( + f"{len(not_done)} requests did not reach a terminal result") + results = [future.result() for future in futures] + finally: + if pool is not None: + pool.shutdown(wait=True) errors = [r for r in results if isinstance(r, Exception)] - assert len(errors) > 0, \ + assert errors, \ "Expected at least one error due to insufficient KV cache" + expected_error = "Insufficient KV cache for gen-only benchmark mode" + unrelated_errors = [ + error for error in errors if expected_error not in str(error) + ] + assert not unrelated_errors, \ + f"Received unrelated errors: {unrelated_errors!r}" finally: terminate(*ctx_workers, *gen_workers, disagg_server) shutil.rmtree(work_dir, ignore_errors=True) @@ -2315,25 +2814,31 @@ def get_config_for_benchmark(model_root, backend): return serve_config -def run_disaggregated_aiperf(config_file, - model_path, - server_start_timeout=1200, - input_tokens=128, - output_tokens=100, - input_tokens_stddev=0, - output_tokens_stddev=0, - concurrency=1, - endpoint_type='chat', - request_count=None, - warmup_request_count=10, - streaming=True, - random_seed=100, - accuracy_test=False, - threshold=0.8, - cancellation_rate=None, - cancellation_delay=None, - env=None, - cwd=None): +def run_disaggregated_aiperf( + config_file, + model_path, + server_start_timeout=1200, + input_tokens=128, + output_tokens=100, + input_tokens_stddev=0, + output_tokens_stddev=0, + concurrency=1, + endpoint_type='chat', + request_count=None, + warmup_request_count=10, + streaming=True, + random_seed=100, + accuracy_test=False, + threshold=0.8, + cancellation_rate=None, + cancellation_delay=None, + env=None, + cwd=None, + request_timeout=180, + heartbeat_interval_sec: int | None = None, + inactive_timeout_sec: int | None = None, + accuracy_concurrency: Optional[int] = None, + post_stress_recovery_timeout: Optional[int] = None): """Run disaggregated test with genai-perf for performance/stress testing. Args: @@ -2350,21 +2855,43 @@ def run_disaggregated_aiperf(config_file, random_seed: Random seed for reproducibility accuracy_test: Whether to run accuracy test threshold: Threshold for accuracy test + cancellation_rate: Percentage of requests to cancel + cancellation_delay: Delay before cancellation in seconds env: Environment variables dict cwd: Working directory + request_timeout: End-to-end disaggregated request timeout in seconds + heartbeat_interval_sec: Optional worker heartbeat interval override + inactive_timeout_sec: Optional worker inactivity timeout override + accuracy_concurrency: Optional concurrency override for the accuracy phase + post_stress_recovery_timeout: Optional bounded wait for fresh requests + to succeed before the accuracy phase """ cleanup_output_files() + if accuracy_concurrency is not None and accuracy_concurrency <= 0: + raise ValueError("accuracy_concurrency must be positive") + if request_timeout <= 0: + raise ValueError("request_timeout must be positive") + if (post_stress_recovery_timeout is not None + and post_stress_recovery_timeout <= 0): + raise ValueError("post_stress_recovery_timeout must be positive") + run_env = env.copy() run_env["UCX_TLS"] = get_ucx_tls() run_env["UCX_MM_ERROR_HANDLING"] = "y" + artifact_dir = os.path.join(cwd or ".", "benchmark-results") + aiperf_client_timeout = (request_timeout + + _AIPERF_CLIENT_TIMEOUT_GRACE_SECS) + remove_stale_aiperf_terminal_records(artifact_dir) config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd, server_start_timeout=server_start_timeout, - save_log=True) + save_log=True, + request_timeout=request_timeout, + heartbeat_interval_sec=heartbeat_interval_sec, + inactive_timeout_sec=inactive_timeout_sec) server_host = config.get("hostname", "localhost") - artifact_dir = os.path.join(cwd or ".", "benchmark-results") try: # Wait for server to be ready @@ -2392,6 +2919,8 @@ def run_disaggregated_aiperf(config_file, aiperf_cmd.extend([ '--url', f'{server_host}:{server_port}', + '--request-timeout-seconds', + str(aiperf_client_timeout), '--synthetic-input-tokens-mean', str(input_tokens), '--synthetic-input-tokens-stddev', @@ -2446,6 +2975,15 @@ def run_disaggregated_aiperf(config_file, env=env, poll_procs=all_worker_procs + [disagg_server.process]) + if request_count is not None: + validate_aiperf_terminal_records(artifact_dir, request_count, + aiperf_client_timeout) + + # This phase intentionally overloads the service and injects client + # cancellations. Its contract is bounded client terminality rather than + # a request-success SLO; the accuracy phase below verifies that the + # service recovers and remains semantically correct afterward. + # Catch cases where aiperf finished but the disagg cluster was unhealthy # during the run (e.g. context-side hangs, KV transfer timeouts) which # would otherwise be swallowed because aiperf records 500s as completed. @@ -2461,15 +2999,49 @@ def run_disaggregated_aiperf(config_file, f"logs:\n{summary}") if accuracy_test: + accuracy_server_url = f"http://{server_host}:{server_port}" + if post_stress_recovery_timeout is not None: + required_probe_successes = (max( + _RECOVERY_PROBE_MIN_SUCCESSES, + _RECOVERY_PROBE_ROUNDS_PER_WORKER * + max(len(ctx_workers), len(gen_workers)))) + wait_for_disagg_recovery( + accuracy_server_url, + model_path=model_path, + timeout=post_stress_recovery_timeout, + request_timeout=(request_timeout + + _ACCURACY_CLIENT_TIMEOUT_GRACE_SECS), + required_successes=required_probe_successes, + workers=[*ctx_workers, *gen_workers], + disagg_server=disagg_server) + # Keep lm-eval retries from recreating the stress burst while + # uncancelled AIPerf backend work finishes naturally. + resolved_accuracy_concurrency = (concurrency if accuracy_concurrency + is None else accuracy_concurrency) + # Give the proxy time to return its bounded timeout and finish + # cleanup before lm-eval abandons an individual HTTP attempt. The + # complete evaluation retains a separate, longer global safety + # cutoff; it is not a worst-case bound for every possible retry. accuracy_test_result, accuracy_value = run_accuracy_test( model_path=model_path, - server_url=f"http://{server_host}:{server_port}", - concurrency=concurrency, + server_url=accuracy_server_url, + concurrency=resolved_accuracy_concurrency, max_retries=3, - timeout=1200, + request_timeout=(request_timeout + + _ACCURACY_CLIENT_TIMEOUT_GRACE_SECS), + process_timeout=_ACCURACY_PROCESS_TIMEOUT_SECS, max_gen_toks=256, max_length=4096) + _raise_if_process_stopped( + [*ctx_workers, *gen_workers, disagg_server], "accuracy") + cluster_diag = build_worker_diag([*ctx_workers, *gen_workers], + disagg_server) + if cluster_diag: + raise AssertionError( + "Disaggregated cluster became unhealthy during the " + f"accuracy phase: {cluster_diag}") + if not accuracy_test_result: raise AssertionError( "Accuracy test failed to complete (likely worker hang or " @@ -2517,17 +3089,19 @@ def run_disaggregated_aiperf(config_file, def run_accuracy_test(model_path: str, server_url: str, concurrency: int, - max_retries: int, timeout: int, max_gen_toks: int, + max_retries: int, request_timeout: int, + process_timeout: int, max_gen_toks: int, max_length: int) -> tuple[bool, float]: """ Run accuracy test using lm_eval with GSM8K dataset Args: model_path: Path of the model being tested - server_config: Server configuration containing URL and port + server_url: Base URL of the disaggregated server concurrency: Concurrency for accuracy tests max_retries: Max retries for accuracy tests - timeout: Timeout for accuracy tests + request_timeout: Timeout for one lm-eval HTTP request + process_timeout: Timeout for the complete lm-eval subprocess max_gen_toks: Max generation tokens for accuracy tests max_length: Max length for accuracy tests @@ -2567,7 +3141,7 @@ def run_accuracy_test(model_path: str, server_url: str, concurrency: int, f"num_concurrent={concurrency}," f"max_retries={max_retries}," f"tokenized_requests=False," - f"timeout={timeout}," + f"timeout={request_timeout}," f"max_gen_toks={max_gen_toks}," f"max_length={max_length}", ] @@ -2583,7 +3157,7 @@ def run_accuracy_test(model_path: str, server_url: str, concurrency: int, result = subprocess.run(lm_eval_cmd, capture_output=True, text=True, - timeout=timeout) + timeout=process_timeout) print_info(f"Accuracy test result is: {result}") @@ -2620,8 +3194,18 @@ def run_accuracy_test(model_path: str, server_url: str, concurrency: int, logger.warning(f"stderr: {result.stderr}") return False, accuracy_value - except subprocess.TimeoutExpired: - logger.warning(f"Accuracy test timed out after {timeout} seconds") + except subprocess.TimeoutExpired as error: + duration = int(time.time() - test_start_time) + logger.warning("Accuracy test subprocess timed out after " + f"{duration} seconds (watchdog={process_timeout}, " + f"per_request={request_timeout})") + for stream_name, partial_output in (("stdout", error.stdout), + ("stderr", error.stderr)): + if partial_output: + if isinstance(partial_output, bytes): + partial_output = partial_output.decode(errors="replace") + logger.warning(f"Partial lm_eval {stream_name}:\n" + f"{partial_output[-20_000:]}") return False, accuracy_value except Exception as e: logger.warning(f"Error during accuracy test: {str(e)}") @@ -2809,7 +3393,11 @@ def test_disaggregated_qwen3_32b_fp8(disaggregated_test_root, request_count=3000, accuracy_threshold=0.72, cancellation_rate=10, - cancellation_delay=0.5), + cancellation_delay=0.5, + accuracy_concurrency=32, + heartbeat_interval_sec=5, + inactive_timeout_sec=10, + post_stress_recovery_timeout=1_800), marks=(pytest.mark.skip_less_device(2), skip_no_hopper)), pytest.param(TestConfig(model_path='GLM-5-NVFP4', test_desc='glm5_nvfp4_tp4_ep4_dp_stress', @@ -2825,7 +3413,11 @@ def test_disaggregated_qwen3_32b_fp8(disaggregated_test_root, accuracy_threshold=0.42, speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3', cancellation_rate=10, - cancellation_delay=0.5), + cancellation_delay=0.5, + accuracy_concurrency=32, + heartbeat_interval_sec=5, + inactive_timeout_sec=10, + post_stress_recovery_timeout=1_800), marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), ], ids=lambda x: x.test_desc) @@ -2847,13 +3439,15 @@ def test_disaggregated_stress_test(disaggregated_test_root, config_file = get_test_config(test_desc, disaggregated_example_root, os.path.dirname(__file__)) - # Resolve speculative_model to an absolute path for worker processes. + # Apply per-test worker overrides without changing other tests that share + # the same checked-in server configuration. if test_config.speculative_model_path is not None: + with open(config_file, 'r') as f: + patched_config = yaml.safe_load(f) + spec_model_dir = f"{llm_models_root()}/{test_config.speculative_model_path}" setup_model_symlink(llm_venv, spec_model_dir, test_config.speculative_model_path) - with open(config_file, 'r') as f: - patched_config = yaml.safe_load(f) patched_sections = [] # Check top-level speculative_config first (current YAML layout), then # fall back to per-server blocks for older config shapes. @@ -2871,30 +3465,37 @@ def test_disaggregated_stress_test(disaggregated_test_root, raise AssertionError( f"{test_desc} sets speculative_model_path, but no " "speculative_config.speculative_model field was patched") + patched_path = os.path.join(llm_venv.get_working_directory(), f"{test_desc}_patched.yaml") with open(patched_path, 'w') as f: yaml.safe_dump(patched_config, f) config_file = patched_path - run_disaggregated_aiperf(config_file=config_file, - model_path=model_dir, - server_start_timeout=7200, - input_tokens=input_tokens, - output_tokens=output_tokens, - input_tokens_stddev=0, - output_tokens_stddev=output_tokens // 10, - concurrency=concurrency, - endpoint_type='completions', - request_count=test_config.request_count, - warmup_request_count=10, - streaming=False, - accuracy_test=True, - threshold=test_config.accuracy_threshold, - cancellation_rate=test_config.cancellation_rate, - cancellation_delay=test_config.cancellation_delay, - env=llm_venv._new_env, - cwd=llm_venv.get_working_directory()) + run_disaggregated_aiperf( + config_file=config_file, + model_path=model_dir, + server_start_timeout=7200, + input_tokens=input_tokens, + output_tokens=output_tokens, + input_tokens_stddev=0, + output_tokens_stddev=output_tokens // 10, + concurrency=concurrency, + endpoint_type='completions', + request_count=test_config.request_count, + warmup_request_count=10, + streaming=False, + accuracy_test=True, + threshold=test_config.accuracy_threshold, + cancellation_rate=test_config.cancellation_rate, + cancellation_delay=test_config.cancellation_delay, + accuracy_concurrency=test_config.accuracy_concurrency, + env=llm_venv._new_env, + cwd=llm_venv.get_working_directory(), + request_timeout=test_config.request_timeout, + heartbeat_interval_sec=(test_config.heartbeat_interval_sec), + inactive_timeout_sec=(test_config.inactive_timeout_sec), + post_stress_recovery_timeout=(test_config.post_stress_recovery_timeout)) def run_cancel_stress_test(server_url: str, diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 6a0fbd26c7ac..393b5a1d0a0f 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -371,6 +371,10 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, save_async): model_fn, scheduler, worker = model_with_connector + # A transceiver-backed executor periodically wakes while idle. Keep the + # connector mock's return value valid while both workers are initialized. + worker.get_finished.return_value = [], [] + prefill_worker = model_fn( disable_overlap_scheduler=True, cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT")) diff --git a/tests/integration/test_lists/qa/llm_function_stress.txt b/tests/integration/test_lists/qa/llm_function_stress.txt index c56465dfd160..f0da958de802 100644 --- a/tests/integration/test_lists/qa/llm_function_stress.txt +++ b/tests/integration/test_lists/qa/llm_function_stress.txt @@ -1,3 +1,6 @@ +disaggregated/test_disaggregated.py::test_validate_aiperf_terminal_records +disaggregated/test_disaggregated.py::test_wait_for_disagg_recovery +disaggregated/test_disaggregated.py::test_wait_for_disagg_recovery_rejects_post_deadline_response disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-deepseek_r1_v2_fp4_stress] disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-deepseek_r1_v2_fp4_mtp_stress] disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_trtllm_stress] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6dc4776bb633..4faa8a56c604 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -154,7 +154,6 @@ full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_si full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-deepseek_r1_v2_fp4_stress] SKIP (https://nvbugs/6472256) -full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6472256) full:B200/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6475623) full:B200/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6475621) full:B200/test_e2e.py::test_multi_nodes_eval[Qwen3/Qwen3-235B-A22B-tp16-mmlu] SKIP (https://nvbugs/6424188) @@ -228,7 +227,6 @@ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_27B_VL::test_ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_35B_A3B_VL::test_auto_dtype SKIP (https://nvbugs/6442073) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_mixed_stress_test[req10k-conc512-qwen3_32b_fp8_mixed_stress] SKIP (https://nvbugs/6440089) -full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6312828) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_5_4b_fp8_stress] SKIP (https://nvbugs/6479324) full:H100/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:H100_PCIe/unittest/llmapi/test_llm_pytorch.py::test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache SKIP (https://nvbugs/5682551) diff --git a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py index 9af481b2f686..31378de11600 100644 --- a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py +++ b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -74,6 +74,7 @@ def __init__(self, kv_cache_manager, async_transfer_manager, kv_cache_transceive self.async_transfer_manager = async_transfer_manager self.kv_cache_transceiver = kv_cache_transceiver self.kv_connector_manager = None + self._async_context_transceiver_request_ids = set() self.disable_overlap_scheduler = True self.previous_batch = None @@ -104,6 +105,7 @@ def test_send_kv_async_calls_release_index_slot(self): PyExecutor._send_kv_async(executor, [request]) kv_cache_manager.release_index_slot.assert_called_once_with(42) + assert executor._async_context_transceiver_request_ids == {42} def test_send_kv_async_skips_release_for_v1_manager(self): """Verify _send_kv_async does not crash when kv_cache_manager lacks @@ -136,6 +138,54 @@ def test_release_called_once_per_request(self): kv_cache_manager.release_index_slot.assert_called_once_with(42) + def test_revisited_request_does_not_duplicate_transfer_owner(self): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = 100 + executor, transfer_manager = self._build(kv_cache_manager) + request = create_mock_request(42) + + PyExecutor._send_kv_async(executor, [request]) + PyExecutor._send_kv_async(executor, [request]) + + assert transfer_manager._request_transfer_metadata[42].counter == 1 + executor.kv_cache_transceiver.respond_and_send_async.assert_called_once_with(request) + kv_cache_manager.release_index_slot.assert_called_once_with(42) + + +def test_context_status_releases_only_transceiver_owner(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = 100 + transfer_manager = AsyncTransferManager( + create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + ) + request = create_mock_request(42) + transfer_manager.start_transfer(request) + transfer_manager.start_transfer(request) + + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = MagicMock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([42], []) + executor.async_transfer_manager = transfer_manager + executor._async_context_transceiver_request_ids = {42} + executor._disagg_ctx_cancel_requested_ids = set() + executor.active_requests = [] + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = MagicMock() + executor._check_cache_transfer_errors = MagicMock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + assert transfer_manager._request_transfer_metadata[42].counter == 1 + assert 42 in transfer_manager.requests_in_transfer() + kv_cache_manager.unpin_blocks_by_id.assert_not_called() + executor._terminate_request.assert_not_called() + + PyExecutor._end_transfer_and_maybe_terminate(executor, request) + + assert 42 not in transfer_manager.requests_in_transfer() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with(100) + executor._terminate_request.assert_called_once_with(request) + class TestIndexMapperSlotReuse: """Test IndexMapper slot recycling after early release.""" diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..05c53261e664 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -13,9 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import sys +from contextlib import nullcontext from types import SimpleNamespace -from unittest.mock import Mock, call +from unittest.mock import MagicMock, Mock, call import pytest @@ -42,6 +44,7 @@ def _reset_inflight_cancel_env_cache(monkeypatch): def _make_timeout_request(request_id=7, in_progress=False): return SimpleNamespace( is_attention_dp_dummy=False, + is_generation_only_request=Mock(return_value=True), py_kv_transfer_timed_out=True, py_request_id=request_id, is_disagg_generation_transmission_in_progress=in_progress, @@ -56,6 +59,11 @@ def _make_response_handler_stub(active_requests, tp_allgather_result): executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.cancel_request.return_value = True executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.kv_cache_transceiver.generation_cancellation_reports_terminal_status.return_value = ( + False + ) + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_gen_cancel_requested_ids = set() executor._disagg_inflight_cancel_unsupported_logged = False executor._pending_timed_out_requests = [] executor.enable_attention_dp = True @@ -73,6 +81,28 @@ def _make_response_handler_stub(active_requests, tp_allgather_result): return executor +def _make_generation_timeout_driver(request, cancel_results, tp_size=1): + executor = object.__new__(PyExecutor) + executor.active_requests = [request] + executor.canceled_req_ids = [] + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.kv_transfer_timeout_ms = 100 + executor.kv_cache_transceiver.cancel_request.side_effect = cancel_results + executor._disagg_gen_cancel_requested_ids = set() + executor._disagg_timed_out_gen_cancelled_ids = set() + executor.dist = SimpleNamespace(tp_size=tp_size) + return executor + + +def _make_generation_transfer_request(request_id=7): + return SimpleNamespace( + py_request_id=request_id, + py_kv_transfer_start_time=0.0, + py_kv_transfer_timed_out=False, + is_disagg_generation_transmission_in_progress=True, + ) + + def test_flag_unset_short_circuits_before_capability_query(monkeypatch): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() @@ -83,6 +113,13 @@ def test_flag_unset_short_circuits_before_capability_query(monkeypatch): executor.kv_cache_transceiver.supports_inflight_request_cancellation.assert_not_called() +def test_cpp_cancellation_is_finalized_by_status_consensus(): + transceiver = object.__new__(BindKvCacheTransceiver) + + assert transceiver.context_cancellation_reports_terminal_status() + assert transceiver.generation_cancellation_reports_terminal_status() + + def test_unsupported_transceiver_warns_once(monkeypatch): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() @@ -159,6 +196,40 @@ def test_flag_unset_generation_timeout_keeps_uncancellable_request_active(): executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) +def test_default_off_cpp_generation_timeout_waits_for_natural_terminal_status(): + request = _make_timeout_request(in_progress=True) + request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + executor = _make_response_handler_stub([request], [False, False]) + executor.kv_cache_transceiver.generation_cancellation_reports_terminal_status.return_value = ( + True + ) + + PyExecutor._handle_responses(executor) + + assert executor.active_requests == [request] + assert executor._pending_timed_out_requests == [] + assert executor._disagg_gen_cancel_requested_ids == set() + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + # Default-off C++ cancellation is observe-only. Repeated cleanup passes + # must continue to preserve the transfer and its KV allocation. + PyExecutor._handle_responses(executor) + + assert executor.active_requests == [request] + assert executor._pending_timed_out_requests == [] + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + # check_gen_transfer_status changes every rank only after its internal + # consensus. Python can now move the terminal request to error cleanup. + request.is_disagg_generation_transmission_in_progress = False + request.state = LlmRequestState.DISAGG_TRANS_ERROR + PyExecutor._handle_responses(executor) + + assert executor.active_requests == [] + assert executor._pending_timed_out_requests == [request] + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + def test_enabled_generation_timeout_waits_for_inflight_terminal_state(monkeypatch): request = _make_timeout_request(in_progress=True) request.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS @@ -172,6 +243,53 @@ def test_enabled_generation_timeout_waits_for_inflight_terminal_state(monkeypatc executor.kv_cache_transceiver.cancel_request.assert_not_called() +def test_timed_out_generation_cancel_is_recorded_and_not_resubmitted(monkeypatch): + request = _make_generation_transfer_request() + executor = _make_generation_timeout_driver(request, [True]) + monkeypatch.setattr(executor_module.time, "monotonic", lambda: 1.0) + + PyExecutor._cancel_timed_out_gen_transfers(executor) + PyExecutor._cancel_timed_out_gen_transfers(executor) + + assert request.py_kv_transfer_timed_out + assert executor._disagg_gen_cancel_requested_ids == {7} + assert executor._disagg_timed_out_gen_cancelled_ids == {7} + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + +def test_timed_out_generation_retries_rejected_cancel(monkeypatch): + request = _make_generation_transfer_request() + executor = _make_generation_timeout_driver(request, [False, True]) + monkeypatch.setattr(executor_module.time, "monotonic", lambda: 1.0) + + PyExecutor._cancel_timed_out_gen_transfers(executor) + + assert executor._disagg_gen_cancel_requested_ids == set() + assert executor._disagg_timed_out_gen_cancelled_ids == set() + + PyExecutor._cancel_timed_out_gen_transfers(executor) + + assert executor._disagg_gen_cancel_requested_ids == {7} + assert executor._disagg_timed_out_gen_cancelled_ids == {7} + assert executor.kv_cache_transceiver.cancel_request.call_count == 2 + + +def test_peer_timeout_is_mirrored_before_local_generation_cancel(monkeypatch): + request = _make_generation_transfer_request() + executor = _make_generation_timeout_driver(request, [True], tp_size=2) + executor.dist.tp_allreduce = Mock(return_value=1) + executor.dist.tp_allgather = Mock(return_value=[[], [7]]) + monkeypatch.setattr(executor_module.time, "monotonic", lambda: 0.01) + + PyExecutor._cancel_timed_out_gen_transfers(executor) + + assert request.py_kv_transfer_timed_out + assert executor._disagg_gen_cancel_requested_ids == {7} + assert executor._disagg_timed_out_gen_cancelled_ids == {7} + executor.dist.tp_allgather.assert_called_once_with([]) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + def test_enabled_generation_timeout_fails_transfer_that_completed_late(monkeypatch): request = _make_timeout_request(in_progress=False) request.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE @@ -197,7 +315,8 @@ def test_flag_unset_context_timeout_preserves_legacy_cleanup(): executor.async_transfer_manager.requests_in_transfer.return_value = { request.py_request_id: request } - executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True executor._disagg_inflight_cancel_unsupported_logged = False executor._end_transfer_and_maybe_terminate = Mock() @@ -209,7 +328,53 @@ def test_flag_unset_context_timeout_preserves_legacy_cleanup(): assert request.py_kv_transfer_start_time is None assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) - assert request.py_request_id not in executor._disagg_timed_out_ctx_cancelled_ids + assert request.py_request_id not in executor._async_context_transceiver_request_ids + assert request.py_request_id not in executor._disagg_ctx_cancel_requested_ids + + +def test_cpp_context_timeout_waits_for_rank_consistent_terminal_status(): + request = _make_timeout_request() + request.py_kv_transfer_start_time = 1.0 + request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_cache_transceiver.context_cancellation_reports_terminal_status.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + executor._end_transfer_and_maybe_terminate = Mock() + executor._check_cache_transfer_errors = Mock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + assert executor._disagg_ctx_cancel_requested_ids == {request.py_request_id} + executor._end_transfer_and_maybe_terminate.assert_not_called() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor._disagg_ctx_cancel_requested_ids == {request.py_request_id} + executor._end_transfer_and_maybe_terminate.assert_not_called() + + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ( + [], + [request.py_request_id], + ) + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + assert request.py_kv_transfer_start_time is None + executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert request.py_request_id not in executor._async_context_transceiver_request_ids + assert request.py_request_id not in executor._disagg_ctx_cancel_requested_ids def test_enabled_context_timeout_defers_cleanup_until_cpp_terminal_state(monkeypatch): @@ -221,11 +386,13 @@ def test_enabled_context_timeout_defers_cleanup_until_cpp_terminal_state(monkeyp executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) executor.kv_cache_transceiver.cancel_request.return_value = True executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.kv_cache_transceiver.context_cancellation_reports_terminal_status.return_value = True executor.async_transfer_manager = Mock() executor.async_transfer_manager.requests_in_transfer.return_value = { request.py_request_id: request } - executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() executor._disagg_inflight_cancel_unsupported_logged = False executor._end_transfer_and_maybe_terminate = Mock() executor._check_cache_transfer_errors = Mock() @@ -235,10 +402,49 @@ def test_enabled_context_timeout_defers_cleanup_until_cpp_terminal_state(monkeyp executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) assert request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - assert request.py_request_id in executor._disagg_timed_out_ctx_cancelled_ids + assert request.py_request_id in executor._disagg_ctx_cancel_requested_ids executor._end_transfer_and_maybe_terminate.assert_not_called() +def test_context_terminal_status_does_not_cancel_remaining_connector_owner(): + request = _make_timeout_request() + request.py_kv_transfer_start_time = 1.0 + request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ( + [request.py_request_id], + [], + ) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + # Model a second connector owner that remains after transceiver status. + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + executor._end_transfer_and_maybe_terminate = Mock() + executor._check_cache_transfer_errors = Mock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + assert request.py_kv_transfer_start_time is None + assert request.py_kv_transfer_timed_out is False + assert executor._async_context_transceiver_request_ids == set() + executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) + + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ( + [], + [], + ) + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + executor.kv_cache_transceiver.cancel_request.assert_not_called() + executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) + + def test_context_transfer_error_keeps_request_active_until_all_owners_release(): request = SimpleNamespace( state=LlmRequestState.DISAGG_TRANS_ERROR, @@ -258,6 +464,44 @@ def test_context_transfer_error_keeps_request_active_until_all_owners_release(): executor._terminate_request.assert_not_called() +def test_context_transfer_error_terminates_manager_only_request(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_TRANS_ERROR, + py_request_id=7, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.active_requests = [] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.end_transfer.return_value = True + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() + + PyExecutor._end_transfer_and_maybe_terminate(executor, request) + + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor._terminate_request.assert_called_once_with(request) + + +def test_early_request_termination_preserves_transceiver_poll_ownership(): + request = SimpleNamespace(py_request_id=7) + executor = object.__new__(PyExecutor) + executor.resource_manager = Mock() + executor._prefetched_request_ids = set() + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_gen_cancel_requested_ids = {request.py_request_id} + executor._disagg_timed_out_gen_cancelled_ids = set() + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0) + executor.result_wait_queues = {} + + PyExecutor._do_terminate_request(executor, request) + + assert executor._async_context_transceiver_request_ids == {request.py_request_id} + assert executor._disagg_gen_cancel_requested_ids == set() + + def test_context_transfer_error_cleanup_waits_for_async_owners(): request = SimpleNamespace( state=LlmRequestState.DISAGG_TRANS_ERROR, @@ -279,6 +523,157 @@ def test_context_transfer_error_cleanup_waits_for_async_owners(): assert PyExecutor._get_disagg_reqs_in_error_state(executor) == [request] +def test_user_cancelled_cpp_context_waits_for_status_before_terminalizing(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=7, + is_context_only_request=True, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.context_cancellation_reports_terminal_status.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor._async_context_transceiver_request_ids = {7} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor._disagg_ctx_cancel_requested_ids == {7} + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + executor.async_transfer_manager.requests_in_transfer.return_value = {} + request.state = LlmRequestState.DISAGG_TRANS_ERROR + assert PyExecutor._try_cancel_request(executor, request) is True + + +def test_opt_in_user_cancelled_cpp_context_waits_for_status_before_terminalizing( + monkeypatch, +): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=7, + is_context_only_request=True, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.context_cancellation_reports_terminal_status.return_value = True + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor._disagg_ctx_cancel_requested_ids == {7} + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + executor.async_transfer_manager.requests_in_transfer.return_value = {} + request.state = LlmRequestState.DISAGG_TRANS_ERROR + assert PyExecutor._try_cancel_request(executor, request) is True + + +def test_default_off_user_cancelled_cpp_generation_waits_for_natural_terminal_status(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + py_request_id=7, + is_context_only_request=False, + is_generation_only_request=Mock(return_value=True), + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.generation_cancellation_reports_terminal_status.return_value = ( + True + ) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor._disagg_gen_cancel_requested_ids = set() + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_not_called() + assert executor._disagg_gen_cancel_requested_ids == set() + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + request.state = LlmRequestState.DISAGG_TRANS_ERROR + assert PyExecutor._try_cancel_request(executor, request) is True + + +def test_opt_in_user_cancelled_cpp_generation_waits_for_status_before_terminalizing( + monkeypatch, +): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + py_request_id=7, + is_context_only_request=False, + is_generation_only_request=Mock(return_value=True), + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.generation_cancellation_reports_terminal_status.return_value = ( + True + ) + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor._disagg_gen_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) + + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor._disagg_gen_cancel_requested_ids == {7} + + # Once C++ accepts cancellation, Python retains ownership and does not + # submit it again while waiting for rank-consistent terminal status. + assert PyExecutor._try_cancel_request(executor, request) is False + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + request.state = LlmRequestState.DISAGG_TRANS_ERROR + assert PyExecutor._try_cancel_request(executor, request) is True + + +def test_opt_in_cpp_generation_retries_unaccepted_cancellation(monkeypatch): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + py_request_id=7, + is_context_only_request=False, + is_generation_only_request=Mock(return_value=True), + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.generation_cancellation_reports_terminal_status.return_value = ( + True + ) + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.kv_cache_transceiver.cancel_request.side_effect = [False, True] + executor.async_transfer_manager = Mock() + executor._disagg_gen_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) + + assert PyExecutor._try_cancel_request(executor, request) is False + assert executor._disagg_gen_cancel_requested_ids == set() + + assert PyExecutor._try_cancel_request(executor, request) is False + assert executor._disagg_gen_cancel_requested_ids == {7} + assert executor.kv_cache_transceiver.cancel_request.call_count == 2 + + assert PyExecutor._try_cancel_request(executor, request) is False + assert executor.kv_cache_transceiver.cancel_request.call_count == 2 + + def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): request = SimpleNamespace( state=LlmRequestState.DISAGG_TRANS_ERROR, @@ -293,14 +688,17 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): executor.active_requests = [request] executor.canceled_req_ids = [request.py_request_id] executor.waiting_queue = Mock() + executor.waiting_queue.remove_by_ids.return_value = [] executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.cancel_request.return_value = True executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor._disagg_ctx_cancel_requested_ids = set() executor._disagg_inflight_cancel_unsupported_logged = False executor.async_transfer_manager = Mock() executor.async_transfer_manager.requests_in_transfer.return_value = { request.py_request_id: request } + executor._async_context_transceiver_request_ids = {request.py_request_id} monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) PyExecutor._handle_canceled_requests(executor) @@ -317,6 +715,557 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): request.finish_by_reason.assert_called_once() +def test_user_cancel_finds_manager_only_cpp_context_transfer(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=7, + py_kv_transfer_timed_out=False, + py_decoding_iter=0, + is_child=False, + is_context_only_request=True, + finish_by_reason=Mock(), + ) + executor = object.__new__(PyExecutor) + executor.active_requests = [] + executor.canceled_req_ids = [request.py_request_id] + executor.waiting_queue = Mock() + executor.waiting_queue.remove_by_ids.return_value = [] + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_cache_transceiver.context_cancellation_reports_terminal_status.return_value = True + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor._async_context_transceiver_request_ids = {request.py_request_id} + executor._disagg_ctx_cancel_requested_ids = set() + executor._disagg_inflight_cancel_unsupported_logged = False + + PyExecutor._handle_canceled_requests(executor) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor.canceled_req_ids == [request.py_request_id] + assert executor._disagg_ctx_cancel_requested_ids == {request.py_request_id} + request.finish_by_reason.assert_not_called() + + +def test_parent_cancel_reaches_active_and_manager_only_child_contexts(): + def make_child(request_id): + return SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=request_id, + parent_request_id=7, + py_kv_transfer_timed_out=False, + py_decoding_iter=0, + is_child=True, + is_context_only_request=True, + finish_by_reason=Mock(), + ) + + active_child = make_child(71) + manager_only_child = make_child(72) + executor = object.__new__(PyExecutor) + executor.active_requests = [active_child] + executor.canceled_req_ids = [7] + executor.waiting_queue = Mock() + executor.waiting_queue.remove_by_ids.return_value = [] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + 71: active_child, + 72: manager_only_child, + } + executor._async_context_transceiver_request_ids = {71, 72} + executor._try_cancel_request = Mock(return_value=False) + + PyExecutor._handle_canceled_requests(executor) + + assert executor._try_cancel_request.call_args_list == [ + call(active_child), + call(manager_only_child), + ] + assert executor.canceled_req_ids == [7] + active_child.finish_by_reason.assert_not_called() + manager_only_child.finish_by_reason.assert_not_called() + + +def test_waiting_cancellation_emits_terminal_response_before_cleanup(monkeypatch): + request_item = SimpleNamespace(id=7) + response = object() + request = SimpleNamespace( + py_request_id=7, + is_child=False, + py_decoding_iter=0, + finish_by_reason=Mock(), + create_response=Mock(return_value=response), + ) + merge = Mock(return_value=[request]) + monkeypatch.setattr(executor_module, "merge_requests", merge) + + executor = object.__new__(PyExecutor) + executor.enable_iter_perf_stats = True + executor.num_fetch_requests = 0 + executor.enable_attention_dp = False + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0, world_size=1, cp_config={}, cp_rank=0, cp_size=1) + executor.executor_request_queue = Mock() + executor._should_exclude_last_generation_logits = Mock(return_value=False) + executor._enqueue_responses = Mock() + executor.result_wait_queues = Mock() + order = Mock() + order.attach_mock(executor._enqueue_responses, "respond") + order.attach_mock(executor.result_wait_queues.pop, "cleanup") + + PyExecutor._terminalize_canceled_waiting_requests(executor, [request_item]) + + request.finish_by_reason.assert_called_once_with(executor_module.FinishReason.CANCELLED) + request.create_response.assert_called_once_with(False, 0) + executor.executor_request_queue.calculate_queue_latency.assert_called_once() + executor._enqueue_responses.assert_called_once_with([(7, response)]) + assert executor.num_fetch_requests == 1 + assert order.mock_calls == [call.respond([(7, response)]), call.cleanup(7, None)] + + +@pytest.mark.parametrize("gather_all_responses", [False, True]) +def test_waiting_cancellation_enters_adp_response_collective_on_nonleader( + monkeypatch, gather_all_responses +): + merge = Mock() + monkeypatch.setattr(executor_module, "merge_requests", merge) + + executor = object.__new__(PyExecutor) + executor.enable_iter_perf_stats = False + executor.num_fetch_requests = 0 + executor.enable_attention_dp = True + executor.gather_all_responses = gather_all_responses + executor.dist = SimpleNamespace(rank=1, world_size=2, cp_config={}, cp_rank=0, cp_size=1) + executor._enqueue_responses = Mock() + waiter = object() + executor.result_wait_queues = {7: waiter} + + PyExecutor._terminalize_canceled_waiting_requests(executor, [SimpleNamespace(id=7)]) + + merge.assert_not_called() + executor._enqueue_responses.assert_called_once_with([]) + if gather_all_responses: + assert 7 not in executor.result_wait_queues + else: + assert executor.result_wait_queues[7] is waiter + + +def test_handle_canceled_requests_terminalizes_removed_waiting_items(): + request_item = SimpleNamespace(id=7) + executor = object.__new__(PyExecutor) + executor.canceled_req_ids = [7] + executor.active_requests = [] + executor.waiting_queue = Mock() + executor.waiting_queue.remove_by_ids.return_value = [request_item] + executor._terminalize_canceled_waiting_requests = Mock() + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} + + PyExecutor._handle_canceled_requests(executor) + + executor._terminalize_canceled_waiting_requests.assert_called_once_with([request_item]) + assert executor.canceled_req_ids == [] + + +def test_handle_canceled_requests_retains_marker_for_control_deferred_item(): + executor = object.__new__(PyExecutor) + executor.canceled_req_ids = [7] + executor.control_requests = [SimpleNamespace(control_requires_drain=True)] + executor.active_requests = [] + executor.waiting_queue = Mock() + executor.waiting_queue.remove_by_ids.return_value = [] + executor._terminalize_canceled_waiting_requests = Mock() + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} + + PyExecutor._handle_canceled_requests(executor) + + assert executor.canceled_req_ids == [7] + + +def test_fetch_terminalizes_waiting_cancellation_before_admission(): + cancel_marker = object() + canceled_item = SimpleNamespace(id=7) + executor = object.__new__(PyExecutor) + executor.control_requests = [] + executor.canceled_req_ids = [] + executor.num_fetch_requests = 0 + executor._disable_mpi = False + executor.dist = SimpleNamespace(rank=0, tp_size=1, has_pp=False, cp_size=1) + executor.request_accumulated = [] + executor.hang_detector = SimpleNamespace(pause=lambda: nullcontext()) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_from_request_queue.return_value = [cancel_marker] + executor.request_broadcaster = Mock() + executor.request_broadcaster.broadcast.return_value = ([cancel_marker], [], False) + + def handle_special_items(_items): + executor.canceled_req_ids.append(7) + return [] + + executor._handle_special_queue_items = Mock(side_effect=handle_special_items) + executor._terminalize_canceled_waiting_requests = Mock() + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 1 + waiting_queue.remove_by_ids.return_value = [canceled_item] + + order = Mock() + order.attach_mock(waiting_queue.add_requests, "enqueue") + order.attach_mock(waiting_queue.remove_by_ids, "remove") + order.attach_mock(executor._terminalize_canceled_waiting_requests, "terminalize") + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + assert order.mock_calls == [ + call.enqueue([]), + call.remove({7}), + call.terminalize([canceled_item]), + ] + assert executor.canceled_req_ids == [] + + +def _make_async_context_poll_executor( + *, + local_inflight=True, + enable_attention_dp=False, + peer_inflight=False, + poll_interval_ms=5000, + transfer_timeout_ms=None, + world_size=None, +): + if world_size is None: + world_size = 2 if enable_attention_dp else 1 + executor = object.__new__(PyExecutor) + executor.control_requests = [] + executor.canceled_req_ids = [] + executor._disable_mpi = False + executor.enable_attention_dp = enable_attention_dp + executor.dist = SimpleNamespace( + rank=0, + world_size=world_size, + tp_size=2 if enable_attention_dp else 1, + has_pp=False, + cp_size=1, + allreduce=Mock(return_value=int(local_inflight or peer_inflight)), + ) + executor.request_accumulated = [] + executor.hang_detector = SimpleNamespace(pause=lambda: nullcontext()) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_from_request_queue.return_value = [] + executor.request_broadcaster = Mock() + executor.request_broadcaster.broadcast.side_effect = ( + lambda requests, poll_context_transfers=False: (requests, None, poll_context_transfers) + ) + executor._handle_special_queue_items = Mock(return_value=[]) + executor._terminalize_canceled_waiting_requests = Mock() + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor._async_context_transceiver_request_ids = {7} if local_inflight else set() + executor.kv_cache_transceiver = SimpleNamespace( + kv_transfer_poll_interval_ms=poll_interval_ms, + kv_transfer_timeout_ms=transfer_timeout_ms, + ) + executor._check_disagg_ctx_cache_transfer_status = Mock() + executor._check_kv_transfer_timeout = Mock() + return executor + + +def test_transfer_only_idle_fetch_uses_bounded_wait_and_polls_after_broadcast(): + executor = _make_async_context_poll_executor() + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + order = Mock() + order.attach_mock(executor.request_broadcaster.broadcast, "broadcast") + order.attach_mock(executor._check_disagg_ctx_cache_transfer_status, "poll") + order.attach_mock(executor._check_kv_transfer_timeout, "timeout") + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=5), cap_batch_wait_to_timeout=True + ) + assert order.mock_calls == [ + call.broadcast([], True), + call.timeout(), + call.poll(0), + ] + + +def test_transfer_only_idle_fetch_uses_fallback_poll_interval(): + executor = _make_async_context_poll_executor(poll_interval_ms=None) + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=1), cap_batch_wait_to_timeout=True + ) + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_kv_transfer_timeout.assert_called_once_with() + + +def test_transfer_poll_wait_is_capped_by_transfer_timeout(): + executor = _make_async_context_poll_executor( + poll_interval_ms=5000, + transfer_timeout_ms=1000, + ) + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=1), cap_batch_wait_to_timeout=True + ) + + +@pytest.mark.parametrize("enable_attention_dp", [False, True]) +def test_peer_rank_transfer_keeps_rank_zero_polling(enable_attention_dp): + executor = _make_async_context_poll_executor( + local_inflight=False, + enable_attention_dp=enable_attention_dp, + peer_inflight=True, + world_size=2, + ) + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.dist.allreduce.assert_called_once_with(0, op=executor_module.ReduceOp.MAX) + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=5), cap_batch_wait_to_timeout=True + ) + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_kv_transfer_timeout.assert_called_once_with() + + +def test_non_root_rank_enters_transfer_poll_after_broadcast(): + executor = _make_async_context_poll_executor(world_size=2) + executor.dist.rank = 1 + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.dist.allreduce.assert_called_once_with(1, op=executor_module.ReduceOp.MAX) + executor.executor_request_queue.get_from_request_queue.assert_not_called() + executor.request_broadcaster.broadcast.assert_called_once_with([], True) + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_kv_transfer_timeout.assert_called_once_with() + + +def test_true_disagg_idle_fetch_uses_bounded_heartbeat_without_status_poll(): + executor = _make_async_context_poll_executor(local_inflight=False, world_size=2) + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=5), cap_batch_wait_to_timeout=True + ) + executor.dist.allreduce.assert_called_once_with(0, op=executor_module.ReduceOp.MAX) + executor._check_disagg_ctx_cache_transfer_status.assert_not_called() + executor._check_kv_transfer_timeout.assert_not_called() + + +def test_draining_control_keeps_polling_manager_only_context_transfer(): + executor = _make_async_context_poll_executor() + executor.control_requests = [SimpleNamespace(control_requires_drain=True)] + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(seconds=5), cap_batch_wait_to_timeout=True + ) + executor.request_broadcaster.broadcast.assert_called_once_with([], True) + executor._check_kv_transfer_timeout.assert_called_once_with() + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + + +def test_draining_control_waits_for_async_owner_on_peer_rank(monkeypatch): + control_request = SimpleNamespace(control_requires_drain=True) + executor = object.__new__(PyExecutor) + executor.control_requests = [control_request] + executor.active_requests = [] + executor.waiting_queue = [] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_any_inflight_requests.return_value = False + executor.dist = SimpleNamespace( + world_size=2, + allreduce=Mock(return_value=1), + ) + synchronize = Mock() + monkeypatch.setattr(executor_module.torch.cuda, "synchronize", synchronize) + + PyExecutor._handle_control_request(executor) + + executor.dist.allreduce.assert_called_once_with(0, op=executor_module.ReduceOp.MAX) + assert executor.control_requests == [control_request] + synchronize.assert_not_called() + + +def test_active_executor_does_not_add_transfer_only_poll(): + executor = _make_async_context_poll_executor(enable_attention_dp=True) + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 1) + + executor.executor_request_queue.get_from_request_queue.assert_called_once_with( + datetime.timedelta(0), cap_batch_wait_to_timeout=False + ) + executor.dist.allreduce.assert_not_called() + executor._check_disagg_ctx_cache_transfer_status.assert_not_called() + executor._check_kv_transfer_timeout.assert_not_called() + + +def test_fetch_processes_cancellation_while_control_waits_for_drain(): + cancel_marker = SimpleNamespace( + id=7, + is_shutdown_request=False, + is_canceled_request=True, + ) + deferred_request = SimpleNamespace( + id=8, + is_shutdown_request=False, + is_canceled_request=False, + ) + canceled_item = SimpleNamespace(id=7) + executor = object.__new__(PyExecutor) + executor.control_requests = [SimpleNamespace(control_requires_drain=True)] + executor.canceled_req_ids = [] + executor._disable_mpi = False + executor.dist = SimpleNamespace(rank=0) + executor.request_accumulated = [] + executor.hang_detector = SimpleNamespace(pause=lambda: nullcontext()) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_from_request_queue.return_value = [ + deferred_request, + cancel_marker, + ] + executor.request_broadcaster = Mock() + executor.request_broadcaster.broadcast.return_value = ( + [deferred_request, cancel_marker], + [], + False, + ) + executor._terminalize_canceled_waiting_requests = Mock() + executor.is_shutdown = False + waiting_queue = MagicMock() + waiting_queue.remove_by_ids.return_value = [canceled_item] + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 1) + + assert executor.request_accumulated == [deferred_request] + assert executor.canceled_req_ids == [] + waiting_queue.add_requests.assert_not_called() + waiting_queue.remove_by_ids.assert_called_once_with({7}) + executor._terminalize_canceled_waiting_requests.assert_called_once_with([canceled_item]) + + +def test_fetch_processes_same_batch_cancellation_behind_control(): + target_item = SimpleNamespace( + id=7, + is_shutdown_request=False, + is_canceled_request=False, + is_control_request=False, + ) + control_marker = SimpleNamespace( + id=8, + is_shutdown_request=False, + is_canceled_request=False, + is_control_request=True, + ) + cancel_marker = SimpleNamespace( + id=7, + is_shutdown_request=False, + is_canceled_request=True, + is_control_request=False, + ) + executor = object.__new__(PyExecutor) + executor.control_requests = [] + executor.canceled_req_ids = [] + executor._disable_mpi = False + executor.dist = SimpleNamespace(rank=0, tp_size=1, has_pp=False, cp_size=1) + executor.request_accumulated = [] + executor.hang_detector = SimpleNamespace(pause=lambda: nullcontext()) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_from_request_queue.return_value = [ + target_item, + control_marker, + cancel_marker, + ] + executor.request_broadcaster = Mock() + executor.request_broadcaster.broadcast.return_value = ( + [target_item, control_marker, cancel_marker], + [], + False, + ) + executor._terminalize_canceled_waiting_requests = Mock() + executor.is_shutdown = False + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + waiting_queue.remove_by_ids.return_value = [target_item] + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + assert executor.control_requests == [control_marker] + assert executor.request_accumulated == [] + assert executor.canceled_req_ids == [] + waiting_queue.add_requests.assert_called_once_with([target_item]) + waiting_queue.remove_by_ids.assert_called_once_with({7}) + executor._terminalize_canceled_waiting_requests.assert_called_once_with([target_item]) + + +def test_fetch_defers_shutdown_until_pending_control_completes(): + shutdown_marker = SimpleNamespace( + id=9, + is_shutdown_request=True, + is_canceled_request=False, + is_control_request=False, + ) + executor = object.__new__(PyExecutor) + executor.control_requests = [SimpleNamespace(control_requires_drain=True)] + executor.canceled_req_ids = [] + executor._disable_mpi = False + executor.dist = SimpleNamespace(rank=0, tp_size=1, has_pp=False, cp_size=1) + executor.request_accumulated = [] + executor.hang_detector = SimpleNamespace(pause=lambda: nullcontext()) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_from_request_queue.side_effect = [ + [shutdown_marker], + [], + ] + executor.request_broadcaster = Mock() + executor.request_broadcaster.broadcast.side_effect = ( + lambda requests, poll_context_transfers=False: (requests, [], poll_context_transfers) + ) + executor._terminalize_canceled_waiting_requests = Mock() + executor.is_shutdown = False + waiting_queue = MagicMock() + waiting_queue.__len__.return_value = 0 + + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + assert not executor.is_shutdown + assert executor.request_accumulated == [shutdown_marker] + + executor.control_requests.clear() + PyExecutor._fetch_and_enqueue_requests(executor, waiting_queue, 0) + + assert executor.is_shutdown + assert executor.request_accumulated == [] + + def test_flag_unset_generation_driver_skips_cancel_pipeline(): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() diff --git a/tests/unittest/_torch/executor/test_executor_request_queue.py b/tests/unittest/_torch/executor/test_executor_request_queue.py index 6874d578f6fb..b2357c4c6b1a 100644 --- a/tests/unittest/_torch/executor/test_executor_request_queue.py +++ b/tests/unittest/_torch/executor/test_executor_request_queue.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Tests for ExecutorRequestQueue class. This module tests the ExecutorRequestQueue class functionality including: @@ -11,7 +13,7 @@ import queue import threading import time -from unittest.mock import Mock, patch +from unittest.mock import Mock, call, patch import pytest @@ -212,6 +214,31 @@ def test_get_from_request_queue_with_timeout(executor_queue): assert elapsed < 0.2 # Should finish within timeout +def test_transfer_poll_timeout_caps_batch_wait(executor_queue): + executor_queue.batch_wait_timeout_ms = 1000 + timeout = datetime.timedelta(milliseconds=50) + item = RequestQueueItem(1, Mock()) + + with ( + patch.object(executor_queue.request_queue, + "empty", + return_value=True), + patch.object(executor_queue.request_queue, + "get", + side_effect=[item, queue.Empty]) as get, + patch( + "tensorrt_llm._torch.pyexecutor.executor_request_queue.time.monotonic", + side_effect=[100.0, 100.04, 100.04], + ), + ): + items = executor_queue.get_from_request_queue( + timeout, cap_batch_wait_to_timeout=True) + + assert items == [item] + assert get.call_args_list[0] == call(timeout=0.05) + assert get.call_args_list[1].kwargs["timeout"] == pytest.approx(0.01) + + def test_get_from_request_queue_async_behavior(executor_queue): """Test asynchronous behavior where requests arrive over time.""" import threading diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index f42375c5e892..47dcfe29cdc5 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Tests for request_utils.py functions. This module tests: @@ -6,6 +9,7 @@ """ +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import Mock, patch @@ -95,6 +99,92 @@ def test_request_broadcaster_requires_conversation_params_attr(): RequestBroadcaster._collect_py_objects(None, source_items) +def test_request_broadcaster_propagates_rank_zero_transfer_poll_bit(): + dist = SimpleNamespace(rank=0, world_size=1) + hang_detector = Mock() + hang_detector.pause.return_value = nullcontext() + broadcaster = RequestBroadcaster(dist, hang_detector) + + assert broadcaster.broadcast([], poll_context_transfers=True) == ( + [], + None, + True, + ) + + +def test_request_broadcaster_non_root_uses_rank_zero_transfer_poll_bit(): + dist = SimpleNamespace( + rank=1, + world_size=2, + has_pp=False, + broadcast=Mock(return_value=1), + ) + hang_detector = Mock() + hang_detector.pause.return_value = nullcontext() + broadcaster = RequestBroadcaster(dist, hang_detector) + + assert broadcaster.broadcast([], poll_context_transfers=False) == ( + [], + None, + True, + ) + dist.broadcast.assert_called_once_with(0, root=0) + + +def test_request_broadcaster_preserves_poll_bit_with_nonempty_batch(): + request = object() + dist = SimpleNamespace(rank=0, world_size=1) + hang_detector = Mock() + hang_detector.pause.return_value = nullcontext() + broadcaster = RequestBroadcaster(dist, hang_detector) + broadcaster._collect_py_objects = Mock(return_value=(("key", {}),)) + broadcaster._broadcast_requests = Mock(return_value=([request], (("key", {}),))) + + assert broadcaster.broadcast([request], poll_context_transfers=True) == ( + [request], + (("key", {}),), + True, + ) + broadcaster._broadcast_requests.assert_called_once_with([request], (("key", {}),)) + + +@pytest.mark.parametrize( + "is_first,is_last", + [ + (True, False), + (False, False), + (False, True), + ], + ids=["first", "middle", "last"], +) +def test_request_broadcaster_propagates_header_through_pp_chain(is_first, is_last): + dist = SimpleNamespace( + world_size=3, + has_pp=True, + pp_size=3, + is_first_pp_rank=is_first, + is_last_pp_rank=is_last, + prev_pp_rank=0, + next_pp_rank=2, + tp_cp_broadcast=Mock(side_effect=lambda value, root: value), + recv_object=Mock(return_value=5), + send_object=Mock(), + ) + broadcaster = RequestBroadcaster(dist, Mock()) + + assert broadcaster._broadcast_request_header(5) == 5 + if is_first: + dist.tp_cp_broadcast.assert_called_once_with(5, root=0) + dist.recv_object.assert_not_called() + else: + dist.tp_cp_broadcast.assert_not_called() + dist.recv_object.assert_called_once_with(0, 4) + if is_last: + dist.send_object.assert_not_called() + else: + dist.send_object.assert_called_once_with(5, 2, 4) + + def test_merge_helix_requests_with_padding(): """Test merge_helix_requests with basic valid input.""" diff --git a/tests/unittest/_torch/executor/test_waiting_queue.py b/tests/unittest/_torch/executor/test_waiting_queue.py index b1f33f235ac2..af741fbae4c8 100644 --- a/tests/unittest/_torch/executor/test_waiting_queue.py +++ b/tests/unittest/_torch/executor/test_waiting_queue.py @@ -128,9 +128,10 @@ def test_remove_by_ids(self): queue.add_requests(items) # Remove items 1 and 3 - queue.remove_by_ids({1, 3}) + removed = queue.remove_by_ids({1, 3}) assert len(queue) == 3 + assert [item.id for item in removed] == [1, 3] remaining_ids = [item.id for item in queue] assert remaining_ids == [0, 2, 4] @@ -141,8 +142,9 @@ def test_remove_nonexistent_ids(self): queue.add_requests(items) # Remove IDs that don't exist - queue.remove_by_ids({10, 20}) + removed = queue.remove_by_ids({10, 20}) + assert removed == [] assert len(queue) == 3 def test_bool_empty_queue(self): @@ -447,9 +449,10 @@ def test_remove_by_ids_removes_correct_items(self): for i, p in enumerate([0.9, 0.7, 0.5, 0.3, 0.1]): q.add_request(create_priority_request_item(i, priority=p)) - q.remove_by_ids({1, 3}) + removed = q.remove_by_ids({1, 3}) assert len(q) == 3 + assert {item.id for item in removed} == {1, 3} remaining = [q.pop_request().id for _ in range(3)] assert remaining == [0, 2, 4] @@ -459,8 +462,9 @@ def test_remove_by_ids_nonexistent_is_noop(self): q.add_request(create_priority_request_item(1, priority=0.8)) q.add_request(create_priority_request_item(2, priority=0.4)) - q.remove_by_ids({99, 100}) + removed = q.remove_by_ids({99, 100}) + assert removed == [] assert len(q) == 2 def test_remove_all_ids_leaves_empty_queue(self): diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 39ae11bc33a4..79a0539b19d3 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -268,6 +268,121 @@ async def test_coordinator_expires_stale_reservation(): assert coordinator._reservation_tasks == {} +@pytest.mark.asyncio +async def test_coordinator_early_finish_releases_later_select(): + config = _make_config([], ["gen:8000"], "round_robin", "conversation") + coordinator = DisaggCoordinatorService(config, _client_factory) + router = coordinator.gen_router + select_started = asyncio.Event() + allow_select = asyncio.Event() + original_select = router.get_next_server_by_key + + async def blocked_select(*args, **kwargs): + select_started.set() + await allow_select.wait() + return await original_select(*args, **kwargs) + + router.get_next_server_by_key = blocked_select + select_task = asyncio.create_task(coordinator.select("generation", "conversation", 123, None)) + await asyncio.wait_for(select_started.wait(), timeout=1) + + await coordinator.finish("generation", 123, False) + allow_select.set() + await asyncio.wait_for(select_task, timeout=1) + + assert router._server_content_load["gen:8000"] == 0 + assert coordinator._reservation_tasks == {} + assert coordinator._finish_tombstone_tasks == {} + + +@pytest.mark.asyncio +async def test_coordinator_finish_failure_retains_expiry_fallback(): + config = _make_config([], ["gen:8000"], "round_robin", "conversation") + coordinator = DisaggCoordinatorService( + config, + _client_factory, + reservation_timeout_secs=0.01, + ) + router = coordinator.gen_router + original_finish = router.finish_request_by_id + finish_attempts = 0 + + async def fail_once(*args, **kwargs): + nonlocal finish_attempts + finish_attempts += 1 + if finish_attempts == 1: + raise RuntimeError("release failed") + return await original_finish(*args, **kwargs) + + router.finish_request_by_id = fail_once + await coordinator.select("generation", "conversation", 123, None) + + with pytest.raises(RuntimeError, match="release failed"): + await coordinator.finish("generation", 123, False) + + assert ("generation", 123) in coordinator._reservation_tasks + await asyncio.sleep(0.02) + assert router._server_content_load["gen:8000"] == 0 + assert coordinator._reservation_tasks == {} + + +@pytest.mark.asyncio +async def test_coordinator_late_select_failure_retains_expiry_fallback(): + config = _make_config([], ["gen:8000"], "round_robin", "conversation") + coordinator = DisaggCoordinatorService( + config, + _client_factory, + reservation_timeout_secs=0.01, + ) + router = coordinator.gen_router + select_started = asyncio.Event() + allow_select = asyncio.Event() + original_select = router.get_next_server_by_key + original_finish = router.finish_request_by_id + finish_attempts = 0 + + async def blocked_select(*args, **kwargs): + select_started.set() + await allow_select.wait() + return await original_select(*args, **kwargs) + + async def fail_second_finish(*args, **kwargs): + nonlocal finish_attempts + finish_attempts += 1 + if finish_attempts == 2: + raise RuntimeError("late release failed") + return await original_finish(*args, **kwargs) + + router.get_next_server_by_key = blocked_select + router.finish_request_by_id = fail_second_finish + select_task = asyncio.create_task(coordinator.select("generation", "conversation", 123, None)) + await asyncio.wait_for(select_started.wait(), timeout=1) + await coordinator.finish("generation", 123, False) + + allow_select.set() + with pytest.raises(RuntimeError, match="late release failed"): + await asyncio.wait_for(select_task, timeout=1) + + assert ("generation", 123) in coordinator._reservation_tasks + await asyncio.sleep(0.02) + assert router._server_content_load["gen:8000"] == 0 + assert coordinator._reservation_tasks == {} + + +@pytest.mark.asyncio +async def test_coordinator_stop_releases_live_reservations(): + config = _make_config([], ["gen:8000"], "round_robin", "conversation") + coordinator = DisaggCoordinatorService(config, _client_factory) + + await coordinator.select("generation", "conversation", 123, None) + assert coordinator.gen_router._server_content_load["gen:8000"] == 1 + + await coordinator.stop() + + assert coordinator.gen_router._server_content_load["gen:8000"] == 0 + assert coordinator._reservation_tasks == {} + + def test_coordinator_compacts_route_info(): compact = DisaggCoordinatorService._compact_route_info( { diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index cead306ae3ff..fa976d1da86e 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio from unittest.mock import AsyncMock, MagicMock, Mock, patch import aiohttp import pytest from tensorrt_llm.llmapi.disagg_utils import ServerRole -from tensorrt_llm.serve.openai_client import OpenAIHttpClient +from tensorrt_llm.serve.openai_client import OpenAIHttpClient, UpstreamRequestTimeoutError from tensorrt_llm.serve.openai_protocol import ( CompletionRequest, CompletionResponse, @@ -230,6 +231,135 @@ async def test_request_error_handling( completion_request, mock_session, success=False ) + @pytest.mark.asyncio + async def test_timeout_is_not_retried( + self, openai_client, completion_request, mock_session, mock_router + ): + """A completed timeout budget must not replay upstream work.""" + mock_session.post.side_effect = asyncio.TimeoutError() + + with pytest.raises(UpstreamRequestTimeoutError, match="180-second timeout"): + await openai_client.send_request(completion_request) + + mock_session.post.assert_called_once() + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=False + ) + + @pytest.mark.asyncio + async def test_upstream_504_is_not_retried( + self, openai_client, completion_request, mock_session, mock_router + ): + """A worker timeout response must retain timeout semantics.""" + mock_http_response = AsyncMock() + mock_http_response.status = 504 + mock_http_response.reason = "Gateway Timeout" + mock_http_response.headers = {"Content-Type": "application/json"} + mock_http_response.text = AsyncMock(return_value="worker deadline expired") + mock_http_response.request_info = MagicMock() + mock_http_response.history = () + mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) + mock_http_response.__aexit__ = AsyncMock(return_value=False) + mock_session.post.return_value = mock_http_response + + with pytest.raises(UpstreamRequestTimeoutError, match="worker deadline expired"): + await openai_client.send_request(completion_request) + + mock_session.post.assert_called_once() + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=False + ) + + @pytest.mark.asyncio + async def test_truncated_upstream_504_body_is_not_retried( + self, openai_client, completion_request, mock_session, mock_router + ): + """A broken timeout body must not make a known 504 retryable.""" + mock_http_response = AsyncMock() + mock_http_response.status = 504 + mock_http_response.reason = "Gateway Timeout" + mock_http_response.headers = {"Content-Type": "application/json"} + mock_http_response.text = AsyncMock( + side_effect=aiohttp.ClientPayloadError("truncated body") + ) + mock_http_response.request_info = MagicMock() + mock_http_response.history = () + mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) + mock_http_response.__aexit__ = AsyncMock(return_value=False) + mock_session.post.return_value = mock_http_response + + with pytest.raises(UpstreamRequestTimeoutError, match="180-second timeout"): + await openai_client.send_request(completion_request) + + mock_session.post.assert_called_once() + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=False + ) + + @pytest.mark.asyncio + async def test_cancelled_request_releases_router_once( + self, openai_client, completion_request, mock_session, mock_router + ): + """Task cancellation must not bypass router cleanup.""" + mock_session.post.side_effect = asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await openai_client.send_request(completion_request) + + mock_session.post.assert_called_once() + mock_router.finish_request.assert_called_once_with( + completion_request, mock_session, success=False + ) + + @pytest.mark.asyncio + async def test_closed_stream_is_not_recorded_as_success( + self, + openai_client, + streaming_completion_request, + mock_session, + mock_router, + ): + """Closing a partially consumed stream records failed cleanup once.""" + mock_http_response = AsyncMock() + mock_http_response.status = 200 + mock_http_response.headers = {"Content-Type": "text/event-stream"} + + async def mock_iter_any(): + yield b'data: "first"\n\n' + await asyncio.Event().wait() + + mock_http_response.content = AsyncMock() + mock_http_response.content.iter_any = mock_iter_any + mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) + mock_http_response.__aexit__ = AsyncMock() + mock_session.post.return_value = mock_http_response + + response_generator = await openai_client.send_request(streaming_completion_request) + assert await response_generator.__anext__() == b'data: "first"\n\n' + await response_generator.aclose() + + mock_router.finish_request.assert_called_once_with( + streaming_completion_request, mock_session, success=False + ) + + @pytest.mark.asyncio + async def test_unstarted_stream_releases_router_once( + self, + openai_client, + streaming_completion_request, + mock_session, + mock_router, + ): + """Closing before the first chunk must still release the reservation.""" + response_generator = await openai_client.send_request(streaming_completion_request) + await response_generator.aclose() + assert await response_generator.athrow(RuntimeError("closed")) is None + + mock_session.post.assert_not_called() + mock_router.finish_request.assert_called_once_with( + streaming_completion_request, mock_session, success=False + ) + @pytest.mark.asyncio async def test_request_with_retry( self, openai_client, completion_request, mock_session, mock_router @@ -428,10 +558,9 @@ async def test_no_generator_keeps_original_id(self): class TestSelectiveTransientTcpRetry: """Selective retry budget for transient TCP race symptoms. - ServerDisconnectedError and ConnectionResetError (which include - aiohttp.ClientConnectionResetError via MRO) get an extended retry budget - of up to 5 attempts; all other client errors keep the original - max_retries fail-fast behaviour. + Context requests retain the extended retry budget. Generation requests + retry only failures which prove the connection was not established, while + preserving the transfer ID shared with the one-shot context response. """ def _ok_response(self): @@ -445,15 +574,15 @@ def _ok_response(self): ) def _mock_http_ok(self, body): - r = AsyncMock() - r.status = 200 - r.headers = {"Content-Type": "application/json"} - r.json = AsyncMock(return_value=body.model_dump()) - r.__aenter__ = AsyncMock(return_value=r) - r.__aexit__ = AsyncMock() - return r - - def _make_client(self, session, max_retries=1): + response = AsyncMock() + response.status = 200 + response.headers = {"Content-Type": "application/json"} + response.json = AsyncMock(return_value=body.model_dump()) + response.__aenter__ = AsyncMock(return_value=response) + response.__aexit__ = AsyncMock() + return response + + def _make_client(self, session, max_retries=1, role=ServerRole.CONTEXT, **kwargs): from prometheus_client.registry import REGISTRY REGISTRY._names_to_collectors = {} @@ -464,30 +593,99 @@ def _make_client(self, session, max_retries=1): router.finish_request = AsyncMock() return OpenAIHttpClient( router=router, - role=ServerRole.CONTEXT, + role=role, timeout_secs=10, max_retries=max_retries, retry_interval_sec=0, session=session, + **kwargs, ) - def _make_request(self): + def _make_request(self, request_type="context_only"): + if request_type == "generation_only": + disaggregated_params = DisaggregatedParams( + request_type=request_type, + first_gen_tokens=[123], + ctx_request_id=23, + disagg_request_id=23, + ) + else: + disaggregated_params = DisaggregatedParams( + request_type=request_type, disagg_request_id=1 + ) return CompletionRequest( model="m", prompt="hi", stream=False, - disaggregated_params=DisaggregatedParams( - request_type="context_only", disagg_request_id=1 - ), + disaggregated_params=disaggregated_params, ) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error", + [ + aiohttp.ServerDisconnectedError(), + ConnectionResetError("connection reset"), + aiohttp.ClientError("ambiguous client failure"), + ], + ids=[ + "server-disconnected", + "connection-reset", + "generic-client-error", + ], + ) + async def test_generation_client_failure_is_not_retried(self, error): + session = AsyncMock(spec=aiohttp.ClientSession) + id_generator = AsyncMock(return_value=24) + client = self._make_client( + session, + max_retries=5, + role=ServerRole.GENERATION, + disagg_id_generator=id_generator, + ) + request = self._make_request("generation_only") + session.post.side_effect = [error, self._mock_http_ok(self._ok_response())] + + with patch("asyncio.sleep", new_callable=AsyncMock) as sleep: + with pytest.raises(type(error)): + await client.send_request(request) + + assert session.post.call_count == 1 + sleep.assert_not_awaited() + id_generator.assert_not_awaited() + assert request.disaggregated_params.ctx_request_id == 23 + assert request.disaggregated_params.disagg_request_id == 23 + + @pytest.mark.asyncio + async def test_generation_pre_connect_failure_retries_with_original_id(self): + session = AsyncMock(spec=aiohttp.ClientSession) + id_generator = AsyncMock(return_value=24) + client = self._make_client( + session, + max_retries=5, + role=ServerRole.GENERATION, + disagg_id_generator=id_generator, + ) + request = self._make_request("generation_only") + session.post.side_effect = [ + aiohttp.ClientConnectorError(MagicMock(), ConnectionRefusedError("connection refused")), + self._mock_http_ok(self._ok_response()), + ] + + with patch("asyncio.sleep", new_callable=AsyncMock) as sleep: + response = await client.send_request(request) + + assert isinstance(response, CompletionResponse) + assert session.post.call_count == 2 + sleep.assert_awaited_once() + id_generator.assert_not_awaited() + assert request.disaggregated_params.ctx_request_id == 23 + assert request.disaggregated_params.disagg_request_id == 23 + @pytest.mark.asyncio async def test_server_disconnected_gets_extra_retries(self): - """ServerDisconnectedError: even with max_retries=1, retry up to 5.""" session = AsyncMock(spec=aiohttp.ClientSession) client = self._make_client(session, max_retries=1) - - # 4 disconnect failures then success on the 5th attempt session.post.side_effect = [ aiohttp.ServerDisconnectedError(), aiohttp.ServerDisconnectedError(), @@ -499,15 +697,12 @@ async def test_server_disconnected_gets_extra_retries(self): with patch("asyncio.sleep", new_callable=AsyncMock): await client.send_request(self._make_request()) - # 1 original + 4 retries = 5 total attempts (extra budget kicked in) assert session.post.call_count == 5 @pytest.mark.asyncio async def test_connection_reset_gets_extra_retries(self): - """ConnectionResetError: same extra budget as ServerDisconnectedError.""" session = AsyncMock(spec=aiohttp.ClientSession) client = self._make_client(session, max_retries=1) - session.post.side_effect = [ ConnectionResetError(), ConnectionResetError(), @@ -517,30 +712,24 @@ async def test_connection_reset_gets_extra_retries(self): with patch("asyncio.sleep", new_callable=AsyncMock): await client.send_request(self._make_request()) - # 1 original + 2 retries = 3 total attempts (within extra budget) assert session.post.call_count == 3 @pytest.mark.asyncio async def test_other_client_error_keeps_fail_fast(self): - """Generic aiohttp.ClientError still respects max_retries (=1).""" session = AsyncMock(spec=aiohttp.ClientSession) client = self._make_client(session, max_retries=1) - session.post.side_effect = aiohttp.ClientError("transient non-tcp") with patch("asyncio.sleep", new_callable=AsyncMock): with pytest.raises(aiohttp.ClientError): await client.send_request(self._make_request()) - # Original + 1 retry = 2 attempts, NOT promoted to 5 assert session.post.call_count == 2 @pytest.mark.asyncio async def test_max_retries_zero_still_gets_transient_tcp_budget(self): - """Even when max_retries=0, transient TCP races still retry up to 5.""" session = AsyncMock(spec=aiohttp.ClientSession) client = self._make_client(session, max_retries=0) - session.post.side_effect = [ aiohttp.ServerDisconnectedError(), self._mock_http_ok(self._ok_response()), @@ -553,16 +742,12 @@ async def test_max_retries_zero_still_gets_transient_tcp_budget(self): @pytest.mark.asyncio async def test_transient_tcp_capped_at_5_when_max_retries_smaller(self): - """If transient TCP keeps failing, give up after the extended budget.""" session = AsyncMock(spec=aiohttp.ClientSession) client = self._make_client(session, max_retries=1) - - # Always raise — must give up after extended (1 + 5) = 6 attempts session.post.side_effect = aiohttp.ServerDisconnectedError() with patch("asyncio.sleep", new_callable=AsyncMock): with pytest.raises(aiohttp.ServerDisconnectedError): await client.send_request(self._make_request()) - # 1 original + 5 retries assert session.post.call_count == 6 diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 5b865027c2bf..ffd2f76d3fdf 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -11,15 +11,23 @@ # 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. +import asyncio +import json from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Request from starlette.datastructures import Headers - -from tensorrt_llm.llmapi.disagg_utils import extract_disagg_cfg -from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer +from starlette.responses import StreamingResponse + +from tensorrt_llm.llmapi.disagg_utils import ServerRole, extract_disagg_cfg +from tensorrt_llm.serve.openai_client import UpstreamRequestTimeoutError +from tensorrt_llm.serve.openai_disagg_server import ( + DisaggregatedRequestTimeoutError, + OpenAIDisaggServer, + _CleanupStreamingResponse, +) from tensorrt_llm.serve.openai_protocol import ( CompletionRequest, ConversationParams, @@ -31,6 +39,41 @@ def _raw_request(headers: dict[str, str]): return SimpleNamespace(headers=Headers(headers=headers)) +@pytest.mark.asyncio +async def test_shutdown_drains_request_cleanup_before_service_teardown(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._cleanup_grace_secs = 0.5 + server._background_cleanup_tasks = set() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + order = [] + + async def cleanup(): + cleanup_started.set() + await release_cleanup.wait() + order.append("cleanup") + + cleanup_task = asyncio.create_task(cleanup()) + server._track_background_cleanup(cleanup_task, "test cleanup") + + async def teardown(): + order.append("teardown") + + server._service = SimpleNamespace(teardown=AsyncMock(side_effect=teardown)) + server._perf_metrics_collector = SimpleNamespace(_background_tasks=set()) + + shutdown_task = asyncio.create_task(server._shutdown()) + await cleanup_started.wait() + await asyncio.sleep(0) + server._service.teardown.assert_not_awaited() + + release_cleanup.set() + await asyncio.wait_for(shutdown_task, timeout=1) + + assert order == ["cleanup", "teardown"] + assert server._background_cleanup_tasks == set() + + @pytest.mark.asyncio async def test_http_cluster_storage_request_is_proxied_to_coordinator(): payload = b'{"key":"worker","value":"ready"}' @@ -174,3 +217,412 @@ def test_disagg_config_rejects_non_bool_request_chat_template_opt_in(value): generation_servers={"num_instances": 0}, allow_request_chat_template=value, ) + + +@pytest.mark.asyncio +async def test_request_deadline_cancels_upstream_work(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 1 + work_started = asyncio.Event() + work_cancelled = asyncio.Event() + + async def blocking_entry_point(_request, _hooks): + work_started.set() + try: + await asyncio.Event().wait() + finally: + work_cancelled.set() + + async def receive(): + await asyncio.Event().wait() + + raw_request = SimpleNamespace(receive=receive) + deadline = asyncio.get_running_loop().time() + 0.01 + + with pytest.raises(DisaggregatedRequestTimeoutError): + await server._await_response_or_disconnect( + blocking_entry_point, object(), object(), raw_request, deadline + ) + + assert work_started.is_set() + assert work_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_request_deadline_does_not_wait_forever_for_cleanup(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 1 + server._cleanup_grace_secs = 0.01 + server._background_cleanup_tasks = set() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + + async def blocking_entry_point(_request, _hooks): + try: + await asyncio.Event().wait() + finally: + cleanup_started.set() + await release_cleanup.wait() + + async def receive(): + await asyncio.Event().wait() + + raw_request = SimpleNamespace(receive=receive) + deadline = asyncio.get_running_loop().time() + 0.01 + + with pytest.raises(DisaggregatedRequestTimeoutError): + await asyncio.wait_for( + server._await_response_or_disconnect( + blocking_entry_point, object(), object(), raw_request, deadline + ), + timeout=0.5, + ) + + assert cleanup_started.is_set() + assert len(server._background_cleanup_tasks) == 1 + cleanup_tasks = tuple(server._background_cleanup_tasks) + release_cleanup.set() + await asyncio.gather(*cleanup_tasks, return_exceptions=True) + await asyncio.sleep(0) + assert server._background_cleanup_tasks == set() + + +@pytest.mark.asyncio +async def test_client_disconnect_cancels_upstream_work(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + work_started = asyncio.Event() + work_cancelled = asyncio.Event() + + async def blocking_entry_point(_request, _hooks): + work_started.set() + try: + await asyncio.Event().wait() + finally: + work_cancelled.set() + + async def receive(): + await work_started.wait() + return {"type": "http.disconnect"} + + raw_request = Request({"type": "http"}, receive=receive) + deadline = asyncio.get_running_loop().time() + 10 + + with pytest.raises(asyncio.CancelledError): + await server._await_response_or_disconnect( + blocking_entry_point, object(), object(), raw_request, deadline + ) + + assert work_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_completed_stream_setup_cancels_disconnect_listener(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + receive_started = asyncio.Event() + receive_cancelled = asyncio.Event() + release_receive = asyncio.Event() + response_stream = object() + + async def receive(): + receive_started.set() + try: + await release_receive.wait() + return {"type": "http.disconnect"} + finally: + receive_cancelled.set() + + async def entry_point(_request, _hooks): + await receive_started.wait() + return response_stream + + raw_request = Request({"type": "http"}, receive=receive) + deadline = asyncio.get_running_loop().time() + 10 + + request_task = asyncio.create_task( + server._await_response_or_disconnect( + entry_point, SimpleNamespace(stream=True), object(), raw_request, deadline + ) + ) + done, _ = await asyncio.wait((request_task,), timeout=1) + completed_without_release = request_task in done + if not completed_without_release: + release_receive.set() + await asyncio.wait((request_task,), timeout=1) + + assert completed_without_release + assert request_task.result() is response_stream + assert receive_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_stream_setup_waits_for_exclusive_receive_handoff(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + server._cleanup_grace_secs = 0.01 + server._background_cleanup_tasks = set() + receive_cancelled = asyncio.Event() + release_receive = asyncio.Event() + response_stream = object() + + async def receive(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + receive_cancelled.set() + await release_receive.wait() + raise + + async def entry_point(_request, _hooks): + return response_stream + + raw_request = Request({"type": "http"}, receive=receive) + deadline = asyncio.get_running_loop().time() + 10 + request_task = asyncio.create_task( + server._await_response_or_disconnect( + entry_point, SimpleNamespace(stream=True), object(), raw_request, deadline + ) + ) + + await asyncio.wait_for(receive_cancelled.wait(), timeout=1) + await asyncio.sleep(server._cleanup_grace_secs * 2) + assert not request_task.done() + + release_receive.set() + assert await asyncio.wait_for(request_task, timeout=1) is response_stream + + +@pytest.mark.asyncio +async def test_canceled_stream_setup_closes_completed_unclaimed_stream(monkeypatch): + class UnclaimedStream: + def __init__(self): + self.closed = False + + async def aclose(self): + self.closed = True + + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + server._cleanup_grace_secs = 0.1 + server._background_cleanup_tasks = set() + response_stream = UnclaimedStream() + + async def receive(): + await asyncio.Event().wait() + + async def entry_point(_request, _hooks): + return response_stream + + real_wait = asyncio.wait + wait_calls = 0 + + async def cancel_first_wait(awaitables, *args, **kwargs): + nonlocal wait_calls + wait_calls += 1 + if wait_calls == 1: + request_task, _ = tuple(awaitables) + await request_task + raise asyncio.CancelledError() + return await real_wait(awaitables, *args, **kwargs) + + monkeypatch.setattr(asyncio, "wait", cancel_first_wait) + raw_request = SimpleNamespace(receive=receive) + deadline = asyncio.get_running_loop().time() + 10 + + with pytest.raises(asyncio.CancelledError): + await server._await_response_or_disconnect( + entry_point, object(), object(), raw_request, deadline + ) + + assert response_stream.closed + + +@pytest.mark.asyncio +async def test_timed_out_stream_setup_closes_stream_returned_on_cancel(): + class UnclaimedStream: + def __init__(self): + self.closed = False + + async def aclose(self): + self.closed = True + + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 1 + server._cleanup_grace_secs = 0.1 + server._background_cleanup_tasks = set() + response_stream = UnclaimedStream() + + async def entry_point(_request, _hooks): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + return response_stream + + async def receive(): + await asyncio.Event().wait() + + raw_request = SimpleNamespace(receive=receive) + deadline = asyncio.get_running_loop().time() + 0.01 + + with pytest.raises(DisaggregatedRequestTimeoutError): + await server._await_response_or_disconnect( + entry_point, object(), object(), raw_request, deadline + ) + + assert response_stream.closed + assert server._background_cleanup_tasks == set() + + +@pytest.mark.parametrize( + "exception, expected_message", + [ + (DisaggregatedRequestTimeoutError(900), "900-second deadline"), + (UpstreamRequestTimeoutError(ServerRole.CONTEXT, 180), "180-second timeout"), + ], +) +def test_request_timeout_maps_to_structured_504(exception, expected_message): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._perf_metrics_collector = SimpleNamespace(http_exceptions=MagicMock()) + + response = server._handle_exception(exception) + body = json.loads(response.body) + + assert response.status_code == 504 + assert body["object"] == "error" + assert body["type"] == "RequestTimeoutError" + assert body["code"] == 504 + assert expected_message in body["message"] + server._perf_metrics_collector.http_exceptions.inc.assert_called_once() + + +@pytest.mark.asyncio +async def test_stream_timeout_emits_error_then_done(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 900 + stream_closed = asyncio.Event() + + async def blocking_stream(): + try: + await asyncio.Event().wait() + yield b"unreachable" + finally: + stream_closed.set() + + deadline = asyncio.get_running_loop().time() + 0.01 + chunks = [chunk async for chunk in server._stream_with_deadline(blocking_stream(), deadline)] + + assert len(chunks) == 2 + error_event = json.loads(chunks[0].removeprefix("data: ").strip()) + assert error_event["error"]["object"] == "error" + assert error_event["error"]["type"] == "RequestTimeoutError" + assert error_event["error"]["code"] == 504 + assert chunks[1] == "data: [DONE]\n\n" + assert stream_closed.is_set() + + +@pytest.mark.asyncio +async def test_stream_preserves_chunks(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + + async def stream(): + for chunk in (b"first", b"second", b"third"): + yield chunk + await asyncio.sleep(0) + + deadline = asyncio.get_running_loop().time() + 10 + chunks = [chunk async for chunk in server._stream_with_deadline(stream(), deadline)] + + assert chunks == [b"first", b"second", b"third"] + + +@pytest.mark.asyncio +async def test_streaming_response_disconnect_closes_upstream_work(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._req_timeout_secs = 180 + stream_started = asyncio.Event() + stream_closed = asyncio.Event() + + async def blocking_stream(): + stream_started.set() + try: + await asyncio.Event().wait() + yield b"unreachable" + finally: + # Exercise cleanup after Starlette has canceled the stream inside + # its AnyIO cancel scope. + await asyncio.sleep(0) + stream_closed.set() + + async def receive(): + await stream_started.wait() + return {"type": "http.disconnect"} + + async def send(_message): + pass + + scope = { + "type": "http", + "asgi": { + "version": "3.0", + "spec_version": "2.3", + }, + } + deadline = asyncio.get_running_loop().time() + 10 + response = StreamingResponse( + server._stream_with_deadline(blocking_stream(), deadline), + media_type="text/event-stream", + ) + + await asyncio.wait_for(response(scope, receive, send), timeout=1) + + assert stream_closed.is_set() + + +@pytest.mark.asyncio +async def test_streaming_response_disconnect_closes_unstarted_upstream_work(): + class UnstartedStream: + def __init__(self): + self.started = False + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started = True + await asyncio.Event().wait() + + async def aclose(self): + self.closed = True + + stream = UnstartedStream() + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._cleanup_grace_secs = 0.1 + server._background_cleanup_tasks = set() + + async def receive(): + return {"type": "http.disconnect"} + + async def send(_message): + await asyncio.Event().wait() + + scope = { + "type": "http", + "asgi": { + "version": "3.0", + "spec_version": "2.3", + }, + } + response = _CleanupStreamingResponse( + stream, + media_type="text/event-stream", + cleanup=stream.aclose, + cleanup_runner=server._run_cleanup_bounded, + ) + + await asyncio.wait_for(response(scope, receive, send), timeout=1) + + assert not stream.started + assert stream.closed diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 04defdd1cd37..2087c3e98343 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -102,6 +102,16 @@ async def test_conditional_disagg_uses_selected_server_match_length(): assert need_context is False +def _make_conditional_context_first_service() -> OpenAIDisaggregatedService: + service = _make_service("context_first") + service._config.conditional_disagg_config = ConditionalDisaggConfig(max_local_prefill_length=0) + gen_router = AsyncMock(spec=KvCacheAwareRouter) + service._gen_router = gen_router + service._coordinator._gen_router = gen_router + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + return service + + def _make_completion_response( text: str, finish_reason: str, @@ -463,6 +473,403 @@ async def _ctx_response(request, *_args, **_kwargs): assert gen_call.args[0].disaggregated_params.ctx_request_id == 202 +@pytest.mark.asyncio +async def test_gen_first_streaming_cancellation_closes_gen_consumer(): + service = _make_service("generation_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + + ctx_started = asyncio.Event() + gen_started = asyncio.Event() + gen_release = asyncio.Event() + gen_closed = asyncio.Event() + + async def _ctx_response(*_args, **_kwargs): + ctx_started.set() + await asyncio.Event().wait() + + async def _gen_response(*_args, **_kwargs): + async def _stream(): + try: + gen_started.set() + await gen_release.wait() + yield b"data: chunk\n\n" + finally: + gen_closed.set() + + return _stream() + + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_response) + service._gen_client.send_request = AsyncMock(side_effect=_gen_response) + + request = CompletionRequest(model="test-model", prompt="hello", stream=True) + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(ctx_started.wait(), timeout=1) + await asyncio.wait_for(gen_started.wait(), timeout=1) + + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + closed_by_service = gen_closed.is_set() + if not closed_by_service: + gen_release.set() + await asyncio.wait_for(gen_closed.wait(), timeout=1) + assert closed_by_service + + +@pytest.mark.asyncio +async def test_gen_first_streaming_cancellation_closes_unstarted_gen_consumer(): + service = _make_service("generation_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + + class LazyStream: + def __init__(self): + self.started = False + self.close_count = 0 + + def __aiter__(self): + self.started = True + return self + + async def __anext__(self): + await asyncio.Event().wait() + + async def aclose(self): + self.close_count += 1 + + gen_response = LazyStream() + service._gen_client.send_request = AsyncMock(return_value=gen_response) + service._ctx_client.send_request = AsyncMock(side_effect=asyncio.CancelledError) + + request = CompletionRequest(model="test-model", prompt="hello", stream=True) + with pytest.raises(asyncio.CancelledError): + await service._send_disagg_request(request) + + assert not gen_response.started + assert gen_response.close_count == 1 + + +@pytest.mark.asyncio +async def test_gen_first_releases_ctx_reservation_before_ctx_dispatch(): + service = _make_service("generation_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + gen_started = asyncio.Event() + + async def _gen_response(*_args, **_kwargs): + gen_started.set() + await asyncio.Event().wait() + + service._gen_client.send_request = AsyncMock(side_effect=_gen_response) + request = CompletionRequest(model="test-model", prompt="hello", stream=True) + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(gen_started.wait(), timeout=1) + + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + service._ctx_client.send_request.assert_not_awaited() + route_call = service._ctx_router.get_next_server.await_args + routed_ctx_request = route_call.args[0] + service._ctx_router.finish_request.assert_awaited_once_with( + routed_ctx_request, + success=False, + req_id=route_call.kwargs["req_id"], + ) + + +@pytest.mark.asyncio +async def test_ctx_first_releases_ctx_reservation_during_placement(): + service = _make_service("context_first") + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + service._check_conditional_disagg = AsyncMock(return_value=(None, True)) + service._check_gen_only_disagg = AsyncMock(return_value=False) + placement_started = asyncio.Event() + + async def _place_ctx(*_args, **_kwargs): + placement_started.set() + await asyncio.Event().wait() + + service._ctx_router.get_next_server = AsyncMock(side_effect=_place_ctx) + request = CompletionRequest(model="test-model", prompt="hello") + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(placement_started.wait(), timeout=1) + + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + route_call = service._ctx_router.get_next_server.await_args + service._ctx_router.finish_request.assert_awaited_once_with( + route_call.args[0], + success=False, + req_id=route_call.kwargs["req_id"], + ) + + +@pytest.mark.asyncio +async def test_ctx_first_releases_gen_reservation_during_placement(): + service = _make_service("context_first") + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + service._check_conditional_disagg = AsyncMock(return_value=(None, True)) + service._check_gen_only_disagg = AsyncMock(return_value=False) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + service._ctx_client = AsyncMock() + service._ctx_client.send_request = AsyncMock( + return_value=_make_completion_response("", finish_reason="length") + ) + placement_started = asyncio.Event() + + async def _place_gen(*_args, **_kwargs): + placement_started.set() + await asyncio.Event().wait() + + service._gen_router.get_next_server = AsyncMock(side_effect=_place_gen) + request = CompletionRequest(model="test-model", prompt="hello") + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(placement_started.wait(), timeout=1) + + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + route_call = service._gen_router.get_next_server.await_args + service._gen_router.finish_request.assert_awaited_once_with( + route_call.args[0], + success=False, + req_id=route_call.kwargs["req_id"], + ) + + +@pytest.mark.asyncio +async def test_ctx_first_attempts_all_reservation_cleanup_on_cancellation(): + service = _make_service("context_first") + request = CompletionRequest(model="test-model", prompt="hello") + + async def canceled_impl(_request, _hooks, ctx_reservation, gen_reservation): + ctx_reservation.mark_pending(request, req_id=101) + gen_reservation.mark_pending(request, req_id=101) + raise asyncio.CancelledError + + service._send_disagg_request_ctx_first_impl = AsyncMock(side_effect=canceled_impl) + service._ctx_router.finish_request.side_effect = RuntimeError("ctx cleanup failed") + + with pytest.raises(asyncio.CancelledError): + await service._send_disagg_request_ctx_first(request) + + service._ctx_router.finish_request.assert_awaited_once_with(request, success=False, req_id=101) + service._gen_router.finish_request.assert_awaited_once_with(request, success=False, req_id=101) + + +@pytest.mark.asyncio +async def test_teardown_attempts_all_components_after_client_failure(): + service = _make_service("context_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_client.shutdown.side_effect = RuntimeError("ctx shutdown failed") + service._coordinator.stop = AsyncMock() + + with pytest.raises(RuntimeError, match="ctx shutdown failed"): + await service.teardown() + + service._ctx_client.shutdown.assert_awaited_once() + service._gen_client.shutdown.assert_awaited_once() + service._coordinator.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_gen_first_nonstream_error_cancels_and_awaits_sibling(): + service = _make_service("generation_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + + gen_started = asyncio.Event() + gen_release = asyncio.Event() + gen_cancelled = asyncio.Event() + gen_finished = asyncio.Event() + + async def _ctx_response(*_args, **_kwargs): + await gen_started.wait() + raise RuntimeError("context failed") + + async def _gen_response(*_args, **_kwargs): + gen_started.set() + try: + await gen_release.wait() + except asyncio.CancelledError: + gen_cancelled.set() + raise + finally: + gen_finished.set() + + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_response) + service._gen_client.send_request = AsyncMock(side_effect=_gen_response) + + request = CompletionRequest(model="test-model", prompt="hello", stream=False) + with pytest.raises(RuntimeError, match="context failed"): + await service._send_disagg_request(request) + + cancelled_by_service = gen_cancelled.is_set() + if not gen_finished.is_set(): + gen_release.set() + await asyncio.wait_for(gen_finished.wait(), timeout=1) + assert cancelled_by_service + route_call = service._ctx_router.get_next_server.await_args + ctx_call = service._ctx_client.send_request.await_args + assert route_call.args[0] is ctx_call.args[0] + assert route_call.args[0] is not request + assert route_call.kwargs["req_id"] == ctx_call.kwargs["req_id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("termination", ["error", "cancellation"]) +async def test_conditional_ctx_first_releases_gen_reservation_before_dispatch( + termination, +): + service = _make_conditional_context_first_service() + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._gen_router.get_next_server = AsyncMock( + return_value=( + "gen:9001", + {"match_length": 0, "num_tokens": 3}, + ) + ) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + + ctx_started = asyncio.Event() + + async def _ctx_response(*_args, **_kwargs): + ctx_started.set() + if termination == "error": + raise RuntimeError("context failed") + await asyncio.Event().wait() + + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_response) + request = CompletionRequest(model="test-model", prompt="hello", stream=False) + + if termination == "error": + with pytest.raises(RuntimeError, match="context failed"): + await service._send_disagg_request(request) + else: + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(ctx_started.wait(), timeout=1) + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + service._gen_client.send_request.assert_not_awaited() + service._gen_router.finish_request.assert_awaited_once_with( + request, + success=False, + req_id=101, + ) + + +@pytest.mark.asyncio +async def test_conditional_ctx_first_does_not_release_after_gen_dispatch(): + service = _make_conditional_context_first_service() + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._gen_router.get_next_server = AsyncMock( + return_value=( + "gen:9001", + {"match_length": 0, "num_tokens": 3}, + ) + ) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + service._ctx_client.send_request = AsyncMock( + return_value=_make_completion_response("", finish_reason="length") + ) + gen_response = _make_completion_response("done", finish_reason="stop", context_only=False) + service._gen_client.send_request = AsyncMock(return_value=gen_response) + + request = CompletionRequest(model="test-model", prompt="hello", stream=False) + result = await service._send_disagg_request(request) + + assert result is gen_response + service._gen_client.send_request.assert_awaited_once() + service._gen_router.finish_request.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("termination", ["error", "cancellation"]) +async def test_conditional_ctx_first_transfers_failed_gen_dispatch_cleanup( + termination, +): + service = _make_conditional_context_first_service() + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._gen_router.get_next_server = AsyncMock( + return_value=( + "gen:9001", + {"match_length": 0, "num_tokens": 3}, + ) + ) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + service._ctx_client.send_request = AsyncMock( + return_value=_make_completion_response("", finish_reason="length") + ) + + gen_started = asyncio.Event() + + async def _gen_response(*_args, **_kwargs): + gen_started.set() + if termination == "error": + raise RuntimeError("generation failed") + await asyncio.Event().wait() + + service._gen_client.send_request = AsyncMock(side_effect=_gen_response) + request = CompletionRequest(model="test-model", prompt="hello", stream=False) + + if termination == "error": + with pytest.raises(RuntimeError, match="generation failed"): + await service._send_disagg_request(request) + else: + request_task = asyncio.create_task(service._send_disagg_request(request)) + await asyncio.wait_for(gen_started.wait(), timeout=1) + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + + service._gen_client.send_request.assert_awaited_once() + service._gen_router.finish_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_conditional_ctx_first_releases_gen_reservation_on_ctx_early_stop(): + service = _make_conditional_context_first_service() + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._gen_router.get_next_server = AsyncMock( + return_value=( + "gen:9001", + {"match_length": 0, "num_tokens": 3}, + ) + ) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + ctx_response = _make_completion_response("done", finish_reason="stop") + service._ctx_client.send_request = AsyncMock(return_value=ctx_response) + + request = CompletionRequest(model="test-model", prompt="hello", stream=False) + result = await service._send_disagg_request(request) + + assert result is ctx_response + service._gen_client.send_request.assert_not_awaited() + service._gen_router.finish_request.assert_awaited_once_with( + request, + success=False, + req_id=101, + ) + + def test_generation_postprocessor_rewrites_usage_from_disaggregated_params(): ctx_usage = UsageInfo( prompt_tokens=128, diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 241a4aee0d09..cef8d09061fb 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -393,52 +393,71 @@ def test_cancel_request_in_transmission(attention_type): kv_cache_transceiver_gen = create_kv_cache_transceiver( mapping, dist, kv_cache_manager_gen, attention_type, cache_transceiver_config) + try: + fill_kv_cache_buffer(kv_cache_manager_ctx) - fill_kv_cache_buffer(kv_cache_manager_ctx) - - # init ctx request - sampling_params = SamplingParams() - ctx_request = LlmRequest( - request_id=0, - max_new_tokens=1, - input_tokens=list(range(256)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) - - kv_cache_manager_ctx.impl.add_sequence_batch( - [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) - # send ctx request - kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) - - # wait for ctx request to be sent - time.sleep(2) + sampling_params = SamplingParams() + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) - # cancel ctx request - is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) - assert is_cancelled + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], + [ctx_request]) + ctx_ref = weakref.ref(ctx_request) + baseline_ctx_refcount = sys.getrefcount(ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1 + + assert kv_cache_transceiver_ctx.cancel_request(ctx_request) + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1, ( + "Accepted cancellation released ownership before the status " + "poll recorded a terminal rank outcome") + + completed_ids, error_ids = ( + kv_cache_transceiver_ctx.check_context_transfer_status(0)) + assert completed_ids == [] + assert error_ids == [ctx_request.py_request_id] + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + + # Construct the generation request before dropping the context request. + # The nanobind context_phase_params getter returns a reference into the + # context request, while the LlmRequest constructor copies the params. + # Keeping the getter result in a Python local would therefore keep the + # context request alive and invalidate the ownership check below. + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) - # init gen request - gen_request = LlmRequest( - request_id=0, - max_new_tokens=1, - input_tokens=list(range(256)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, - context_phase_params=ctx_request.context_phase_params) + del ctx_request + gc.collect() + assert ctx_ref() is None, ( + "Terminal context cancellation retained LlmRequest ownership") - kv_cache_manager_gen.impl.add_sequence_batch( - [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) - # send gen request - kv_cache_transceiver_gen.request_and_receive_async(gen_request) + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + kv_cache_transceiver_gen.request_and_receive_async(gen_request) - # Block the main thread due to the async operation - time.sleep(2) - kv_cache_transceiver_gen.check_gen_transfer_status(0) - assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + wait_for_transfer_completion( + lambda: kv_cache_transceiver_gen.check_gen_transfer_status(0), + lambda: gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, + timeout_s=10, + ) + finally: + shutdown_transceivers(kv_cache_transceiver_gen, + kv_cache_transceiver_ctx) @pytest.mark.timeout(120)