diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 700616a3f947..944a70ce479e 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -600,10 +600,22 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { bool blockAll = !atLeastRequestNum.has_value(); + TLLM_LOG_INFO( + "[disagg-debug] C++ checkGenTransferStatus enter: atLeastRequestNum=%d blockAll=%d " + "requesterFutures=%zu", + atLeastRequestNum.value_or(-1), static_cast(blockAll), mRequesterFutures.size()); std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + auto const status = future.wait_for(std::chrono::milliseconds(0)); + auto const ctxRequestId + = request->getContextPhaseParams().has_value() ? request->getContextPhaseParams().value().getReqId() : 0; + TLLM_LOG_INFO( + "[disagg-debug] C++ checkGenTransferStatus future probe: requestId=%zu ctxRequestId=%zu " + "ready=%d state=%d", + request->mRequestId, ctxRequestId, static_cast(status == std::future_status::ready), + static_cast(request->getState())); + if (status == std::future_status::ready) { genTransferReadyRequestIds.push_back(request->mRequestId); } @@ -710,13 +722,39 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } + std::ostringstream selectedIds; + bool firstSelectedId = true; + for (auto const requestId : toCompleteIdSet) + { + if (!firstSelectedId) + { + selectedIds << ","; + } + selectedIds << requestId; + firstSelectedId = false; + } + TLLM_LOG_INFO( + "[disagg-debug] C++ checkGenTransferStatus selected requests: readyLocal=%zu freqVec=%zu " + "toComplete=%zu ids=[%s]", + genTransferReadyRequestIds.size(), freqVec.size(), toCompleteIdSet.size(), selectedIds.str().c_str()); for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) { if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) { try { + auto const ctxRequestId = it->first->getContextPhaseParams().has_value() + ? it->first->getContextPhaseParams().value().getReqId() + : 0; + TLLM_LOG_INFO( + "[disagg-debug] C++ checkGenTransferStatus waiting on requester future: requestId=%zu " + "ctxRequestId=%zu blockAll=%d", + it->first->mRequestId, ctxRequestId, static_cast(blockAll)); it->second.get(); + TLLM_LOG_INFO( + "[disagg-debug] C++ checkGenTransferStatus requester future completed: requestId=%zu " + "ctxRequestId=%zu", + it->first->mRequestId, ctxRequestId); it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); // Gather the kv cache transfer time from all workers and update to leader rank diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 3ecceb9f3f2c..19476236f660 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -33,6 +33,7 @@ #include #include #include +#include #include namespace tensorrt_llm::batch_manager @@ -301,16 +302,21 @@ class CacheSender::Impl std::promise promise; auto future = promise.get_future(); llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + size_t readyResponseCount{0}; { { std::scoped_lock lkResp(mSenderMutex); mReadyResponses.emplace( llmRequest.mRequestId, Response{std::addressof(llmRequest), std::move(promise)}); + readyResponseCount = mReadyResponses.size(); } std::unique_lock lkCond(mCondMutex); mAnyReady = true; } mSenderCv.notify_all(); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheSender sendAsync queued: requestId=%zu selfIdx=%d device=%d readyResponses=%zu", + llmRequest.mRequestId, mSelfState.getCommState().value().getSelfIdx(), mDeviceId, readyResponseCount); return future; } @@ -356,6 +362,16 @@ class CacheSender::Impl auto* agentConnectionManager = dynamic_cast(mManager); bool isAgent = agentConnectionManager != nullptr; + if (isAgent) + { + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheSender recvRequestInfo wait: selfIdx=%d device=%d localAgent=%s " + "terminate=%d managerRunning=%d", + mSelfState.getCommState().value().getSelfIdx(), mDeviceId, + agentConnectionManager->getAgentName().c_str(), static_cast(mTerminate.load()), + static_cast(mManager->isRunning())); + } + TransceiverTag::Id id; RequestInfo info; auto const* connection = isAgent @@ -391,6 +407,7 @@ class CacheSender::Impl TLLM_CHECK_WITH_INFO(peerIdx < static_cast(allCounterparts.size()), "Peer rank %d not found in expected counterparts", peerSelfIdx); + size_t connectionCount = 0; { std::unique_lock lk(mMtxForMap); auto it = mRequestToSession.find(requestId); @@ -404,7 +421,13 @@ class CacheSender::Impl it = mRequestToSession.emplace(requestId, std::move(session)).first; } it->second.setConnection(peerIdx, connection); + connectionCount = it->second.getConnections().size(); } + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheSender request-info received: requestId=%zu selfIdx=%d " + "peerSelfIdx=%d peerIdx=%d counterpartCount=%zu connections=%zu", + requestId, mSelfState.getCommState().value().getSelfIdx(), peerSelfIdx, peerIdx, allCounterparts.size(), + connectionCount); return info; } @@ -467,6 +490,10 @@ class CacheSender::Impl executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); } } + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheSender ready signal sent: requestId=%zu selfIdx=%d isReady=%d " + "connections=%zu", + requestId, mSelfState.getCommState().value().getSelfIdx(), static_cast(isReady), connections.size()); } ~Impl() @@ -623,6 +650,40 @@ class CacheSender::Impl } if (!mReadyResponses.empty()) { + auto* agentConnectionManager = dynamic_cast(mManager); + std::ostringstream readyRequestIds; + std::string currentRequest{"none"}; + size_t readyResponseCount{0}; + { + std::scoped_lock lkResp(mSenderMutex); + readyResponseCount = mReadyResponses.size(); + bool first = true; + for (auto const& [requestId, response] : mReadyResponses) + { + (void) response; + if (!first) + { + readyRequestIds << ","; + } + readyRequestIds << requestId; + first = false; + } + if (mCurrentRequest.has_value()) + { + currentRequest = std::to_string(mCurrentRequest.value()); + } + } + auto const readyRequestIdsString = readyRequestIds.str(); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheSender response waiting for request-info: selfIdx=%d device=%d " + "isAgent=%d localAgent=%s readyResponses=%zu readyRequestIds=[%s] currentRequest=%s " + "anyReady=%d terminate=%d managerRunning=%d", + mSelfState.getCommState().value().getSelfIdx(), mDeviceId, + static_cast(agentConnectionManager != nullptr), + agentConnectionManager != nullptr ? agentConnectionManager->getAgentName().c_str() : "", + readyResponseCount, readyRequestIdsString.c_str(), currentRequest.c_str(), + static_cast(mAnyReady), static_cast(mTerminate.load()), + static_cast(mManager->isRunning())); auto const& requestInfo = recvRequestInfo(); if (mTerminate || !mManager->isRunning()) { @@ -795,7 +856,14 @@ class CacheReceiver::Impl void receiveSync(TransferSession& session) { + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver receiveSync begin: requestId=%zu ctxRequestId=%zu " + "connections=%zu", + session.getLlmRequest().mRequestId, session.getLlmRequest().getContextPhaseParams().value().getReqId(), + session.getConnections().size()); mCacheTransferLayer.unformat(session); + TLLM_LOG_INFO("[disagg-debug] C++ CacheReceiver receiveSync unformat done: requestId=%zu ctxRequestId=%zu", + session.getLlmRequest().mRequestId, session.getLlmRequest().getContextPhaseParams().value().getReqId()); if (!common::getEnvKVCacheTimeOutputPath().empty()) { std::unique_lock lock(mMeasuresFileMutex); @@ -817,6 +885,10 @@ class CacheReceiver::Impl auto const& commState = contextState.getCommState().value(); auto const& destCacheState = contextState.getCacheState().value(); mCacheTransferLayer.validateSupport(contextState); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestInfo begin: requestId=%zu ctxRequestId=%zu " + "selfIdx=%d contextSelfIdx=%d", + llmRequest.mRequestId, requestId, mSelfState.getCommState().value().getSelfIdx(), commState.getSelfIdx()); RequestInfo requestInfo(requestId, mSelfState); @@ -871,6 +943,29 @@ class CacheReceiver::Impl destCacheState, mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx()) .mIRanks; } + std::ostringstream allCounterpartsStream; + for (size_t i = 0; i < allCounterparts.size(); i++) + { + if (i > 0) + { + allCounterpartsStream << ","; + } + allCounterpartsStream << allCounterparts[i]; + } + std::ostringstream kvCounterpartsStream; + for (size_t i = 0; i < kvCounterParts.size(); i++) + { + if (i > 0) + { + kvCounterpartsStream << ","; + } + kvCounterpartsStream << kvCounterParts[i]; + } + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestInfo counterparts: requestId=%zu ctxRequestId=%zu " + "all=[%s] kv=[%s] rnnCount=%zu", + llmRequest.mRequestId, requestId, allCounterpartsStream.str().c_str(), kvCounterpartsStream.str().c_str(), + rnnCounterParts.size()); auto connections = mManager->getConnections(commState); std::vector allConnections; @@ -931,19 +1026,50 @@ class CacheReceiver::Impl auto* agentConnection = dynamic_cast(connection); TLLM_CHECK(agentConnection != nullptr); + size_t activeBufferCount = 0; + for (auto const& id : idsForRank) + { + if (id.has_value()) + { + activeBufferCount++; + } + } + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestAndBufferInfo begin: requestId=%zu " + "ctxRequestId=%zu counterpartRank=%d connectionIdx=%d isKv=%d isRnn=%d " + "activeBuffers=%zu", + llmRequest.mRequestId, requestId, rank, validConnectionIdx, static_cast(isKvCounterpart), + static_cast(isRnnCounterpart), activeBufferCount); const_cast(agentConnection) ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestAndBufferInfo end: requestId=%zu " + "ctxRequestId=%zu counterpartRank=%d connectionIdx=%d", + llmRequest.mRequestId, requestId, rank, validConnectionIdx); } else { + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestInfo legacy send begin: requestId=%zu " + "ctxRequestId=%zu counterpartRank=%d", + llmRequest.mRequestId, requestId, rank); sendRequestInfo(connection, requestInfo); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestInfo legacy send end: requestId=%zu " + "ctxRequestId=%zu counterpartRank=%d", + llmRequest.mRequestId, requestId, rank); } } auto const& resource = getReceiveCacheResource(llmRequest); - return TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, + auto session = TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, !common::getEnvKVCacheTimeOutputPath().empty()); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver sendRequestInfo end: requestId=%zu ctxRequestId=%zu " + "connections=%zu", + llmRequest.mRequestId, requestId, session.getConnections().size()); + return session; } std::unique_ptr const& getReceiveCacheResource(LlmRequest const& llmRequest) @@ -1010,25 +1136,54 @@ class CacheReceiver::Impl bool isReadyFinal = true; bool isReady = false; auto const& connections = session.getConnections(); + auto const& request = session.getLlmRequest(); + auto const& counterpartRanks = session.getCounterPartRanks(); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver receiveReadySignal begin: requestId=%zu ctxRequestId=%zu " + "connections=%zu", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), connections.size()); for (size_t i = 0; i < connections.size(); i++) { + int const counterpartRank = i < counterpartRanks.size() ? static_cast(counterpartRanks[i]) : -1; auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver recvReadySignal begin: requestId=%zu ctxRequestId=%zu " + "connectionIdx=%zu counterpartRank=%d", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), i, counterpartRank); isReady = agentConnection->recvReadySignal( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, mTerminate}); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver recvReadySignal end: requestId=%zu ctxRequestId=%zu " + "connectionIdx=%zu counterpartRank=%d isReady=%d", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), i, counterpartRank, + static_cast(isReady)); } else { + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver recvReadySignal legacy begin: requestId=%zu " + "ctxRequestId=%zu connectionIdx=%zu counterpartRank=%d", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), i, counterpartRank); connections.at(i)->recv( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver recvReadySignal legacy end: requestId=%zu " + "ctxRequestId=%zu connectionIdx=%zu counterpartRank=%d isReady=%d", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), i, counterpartRank, + static_cast(isReady)); } isReadyFinal &= isReady; } + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver receiveReadySignal end: requestId=%zu ctxRequestId=%zu " + "isReadyFinal=%d", + request.mRequestId, request.getContextPhaseParams().value().getReqId(), static_cast(isReadyFinal)); return isReadyFinal; } @@ -1052,11 +1207,29 @@ class CacheReceiver::Impl TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); + TLLM_LOG_INFO("[disagg-debug] C++ CacheReceiver requestSync start: requestId=%zu ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); llmRequest.setKvCacheTransferStart(std::chrono::steady_clock::now()); TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver requestSync sendRequestInfo call: requestId=%zu " + "ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); auto session = sendRequestInfo(llmRequest); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver requestSync sendRequestInfo returned: requestId=%zu " + "ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); session.setTime(TransferSession::kTimeRequestInfo); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver requestSync receiveReadySignal call: requestId=%zu " + "ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); bool isReady = receiveReadySignal(session); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver requestSync receiveReadySignal returned: requestId=%zu " + "ctxRequestId=%zu isReady=%d", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId(), static_cast(isReady)); if (!isReady) { // Reuse the error state for the cancelled request. @@ -1064,12 +1237,20 @@ class CacheReceiver::Impl llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); return; } + TLLM_LOG_INFO("[disagg-debug] C++ CacheReceiver requestSync receiveSync call: requestId=%zu ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); receiveSync(session); + TLLM_LOG_INFO( + "[disagg-debug] C++ CacheReceiver requestSync receiveSync returned: requestId=%zu " + "ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); + TLLM_LOG_INFO("[disagg-debug] C++ CacheReceiver requestSync end: requestId=%zu ctxRequestId=%zu", + llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); } struct RequestAndPromise diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index d46defdf50ad..765861bcae7f 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,14 +18,56 @@ #include "connection.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" +#include +#include #include +#include #include +#include #include #include namespace tensorrt_llm::executor::kv_cache { +namespace +{ + +std::string describeNotificationInfo(NotificationInfo const& notificationInfo) +{ + std::ostringstream os; + os << "variant=" << notificationInfo.mInfo.index(); + if (std::holds_alternative(notificationInfo.mInfo)) + { + auto const& requestAndBufferInfo = std::get(notificationInfo.mInfo); + auto const& requestInfo = requestAndBufferInfo.mRequestInfo; + auto const& commState = requestInfo.getTransState().getCommState(); + auto const peerSelfIdx = commState.has_value() ? commState->getSelfIdx() : -1; + os << " kind=RequestAndBufferInfo" + << " senderAgent=" << requestAndBufferInfo.mAgentName << " requestId=" << requestInfo.getRequestId() + << " peerSelfIdx=" << peerSelfIdx << " connectionIdx=" << requestAndBufferInfo.mValidConnectionIdx + << " bufferDescs=" << requestAndBufferInfo.mBufferDescs.size() + << " metadata=" << static_cast(requestAndBufferInfo.mMetadata.has_value()) + << " addressBytes=" << requestAndBufferInfo.mAddress.size(); + } + else if (std::holds_alternative(notificationInfo.mInfo)) + { + auto const& syncInfo = std::get(notificationInfo.mInfo); + os << " kind=NotificationSyncInfo" + << " expectedAgent=" << syncInfo.mAgentName << " tag=" << syncInfo.mContext.getTag(); + } + else if (std::holds_alternative(notificationInfo.mInfo)) + { + auto const& readySignalInfo = std::get(notificationInfo.mInfo); + os << " kind=ReadySignalInfo" + << " expectedAgent=" << readySignalInfo.mAgentName << " tag=" << readySignalInfo.mContext.getTag() + << " isReady=" << static_cast(readySignalInfo.mIsReady); + } + return os.str(); +} + +} // namespace + std::string genUniqueAgentName() { static std::atomic counter{0}; @@ -184,6 +226,7 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque } TLLM_CHECK(!activeCacheBufferIds.empty()); + auto const activeBufferCount = activeCacheBufferIds.size(); mCacheBufferIds = std::move(activeCacheBufferIds); mBufferKinds = activeKinds; @@ -209,7 +252,16 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque std::stringstream ss; NotificationInfo notificationInfo{requestAndBufferInfo}; NotificationInfo::serialize(notificationInfo, ss); - mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); + auto payload = ss.str(); + mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, payload); + auto const& commState = requestInfo.getTransState().getCommState(); + auto const peerSelfIdx = commState.has_value() ? commState->getSelfIdx() : -1; + TLLM_LOG_INFO( + "[disagg-debug] C++ AgentConnection request-info notified: localAgent=%s remoteAgent=%s " + "requestId=%zu peerSelfIdx=%d connectionIdx=%d activeBuffers=%zu bufferDescs=%zu metadata=%d " + "payloadBytes=%zu", + mAgentName.c_str(), mRemoteAgentName.c_str(), requestInfo.getRequestId(), peerSelfIdx, connectionIdx, + activeBufferCount, bufferDescs.size(), static_cast(metadataOpt.has_value()), payload.size()); } void AgentConnection::setSenderState(std::vector cacheReceiverBufferDescs, int validSegmentIdx, @@ -241,7 +293,13 @@ void AgentConnection::sendReadySignal(DataContext const& ctx, bool isReady) cons NotificationInfo notificationInfo{readySignalInfo}; std::stringstream ss; NotificationInfo::serialize(notificationInfo, ss); - mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); + auto payload = ss.str(); + mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, payload); + TLLM_LOG_INFO( + "[disagg-debug] C++ AgentConnection ready notified: localAgent=%s remoteAgent=%s readyAgent=%s " + "tag=%d isReady=%d payloadBytes=%zu", + mAgentName.c_str(), mRemoteAgentName.c_str(), readySignalInfo.mAgentName.c_str(), ctx.getTag(), + static_cast(isReady), payload.size()); } bool AgentConnection::recvReadySignal(DataContext const& ctx) const @@ -355,6 +413,11 @@ AgentConnectionManager::AgentConnectionManager( agentStates[0] = localAgentState; } mCommState = CommState(agentStates, mpi::MpiComm::session().getRank()); + TLLM_LOG_INFO( + "[disagg-debug] C++ AgentConnectionManager agent ready: localAgent=%s device=%d registeredBuffers=%zu " + "bufferKinds=%zu backend=%s sessionRank=%d sessionSize=%d worldRank=%d", + mAgentName.c_str(), mDeviceId, memDescs.size(), mBufferKinds.size(), backendType.c_str(), + mpi::MpiComm::session().getRank(), mpi::MpiComm::session().getSize(), mpi::MpiComm::world().getRank()); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " ***** AgentConnectionManager::AgentConnectionManager mCommState: %s", mCommState.toString().c_str()); } @@ -362,14 +425,42 @@ AgentConnectionManager::AgentConnectionManager( AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( batch_manager::RequestInfo& requestInfo, std::atomic const& terminateFlag) { + auto const startTime = std::chrono::steady_clock::now(); + auto nextLogTime = startTime + std::chrono::seconds(30); + uint64_t pollCount{0}; + constexpr uint64_t kWaitLogPollPeriod{4096}; while (!terminateFlag.load()) { + pollCount++; if (!mIsRunning) { return nullptr; } updateUnhandledNotifications(); - std::scoped_lock lock(mNotificationMutex); + bool shouldLogWait{false}; + size_t pendingAgentCount{0}; + size_t pendingNotificationCount{0}; + long long elapsedMs{0}; + if (pollCount % kWaitLogPollPeriod == 0) + { + auto const now = std::chrono::steady_clock::now(); + if (now >= nextLogTime) + { + shouldLogWait = true; + elapsedMs = std::chrono::duration_cast(now - startTime).count(); + nextLogTime = now + std::chrono::seconds(30); + } + } + std::unique_lock lock(mNotificationMutex); + if (shouldLogWait) + { + pendingAgentCount = mUnhandledNotifications.size(); + for (auto const& [agent, notifs] : mUnhandledNotifications) + { + (void) agent; + pendingNotificationCount += notifs.size(); + } + } auto it = mUnhandledNotifications.begin(); while (it != mUnhandledNotifications.end()) { @@ -391,6 +482,16 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( auto metadataOpt = requestAndBufferInfo.mMetadata; auto connectionIdx = requestAndBufferInfo.mValidConnectionIdx; auto remoteAgentName = requestAndBufferInfo.mAgentName; + auto const bufferDescCount = bufferDescs.size(); + auto const& commState = requestInfo.getTransState().getCommState(); + auto const peerSelfIdx = commState.has_value() ? commState->getSelfIdx() : -1; + TLLM_LOG_INFO( + "[disagg-debug] C++ recvConnectionAndRequestInfo request-info before connect: localAgent=%s " + "notificationAgent=%s remoteAgent=%s requestId=%zu peerSelfIdx=%d connectionIdx=%d " + "bufferDescs=%zu metadata=%d addressBytes=%zu", + mAgentName.c_str(), agent.c_str(), remoteAgentName.c_str(), requestInfo.getRequestId(), + peerSelfIdx, connectionIdx, bufferDescCount, static_cast(metadataOpt.has_value()), + address.size()); TLLM_LOG_DEBUG(" recv Address:%s", address.c_str()); auto connection = connect(remoteAgentName, address, metadataOpt, true); auto bufferKinds = std::move(requestAndBufferInfo.mBufferKinds); @@ -440,6 +541,13 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( } connection->setSenderState( std::move(bufferDescs), connectionIdx, std::move(offsetRatios), std::move(bufferKinds)); + TLLM_LOG_INFO( + "[disagg-debug] C++ recvConnectionAndRequestInfo matched request-info: localAgent=%s " + "notificationAgent=%s remoteAgent=%s requestId=%zu peerSelfIdx=%d connectionIdx=%d " + "bufferDescs=%zu metadata=%d addressBytes=%zu", + mAgentName.c_str(), agent.c_str(), remoteAgentName.c_str(), requestInfo.getRequestId(), + peerSelfIdx, connectionIdx, bufferDescCount, static_cast(metadataOpt.has_value()), + address.size()); notifIt = notifs.erase(notifIt); if (notifs.empty()) { @@ -462,6 +570,16 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( it++; } } + if (shouldLogWait) + { + lock.unlock(); + TLLM_LOG_INFO( + "[disagg-debug] C++ recvConnectionAndRequestInfo still waiting: localAgent=%s pendingAgents=%zu " + "pendingNotifications=%zu pollCount=%llu elapsedMs=%lld terminate=%d running=%d", + mAgentName.c_str(), pendingAgentCount, pendingNotificationCount, + static_cast(pollCount), elapsedMs, static_cast(terminateFlag.load()), + static_cast(mIsRunning.load())); + } } return nullptr; } @@ -469,14 +587,59 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( void AgentConnectionManager::updateUnhandledNotifications() { auto notifiedSyncMessages = m_Agent->getNotifiedSyncMessages(); - std::lock_guard lock(mNotificationMutex); - // Merge new notifications with existing ones + struct NotificationDebugLog + { + std::string agent; + size_t newCount; + size_t pendingForAgent; + size_t pendingTotal; + std::string firstNotification; + }; + + std::vector debugLogs; for (auto const& [agent, notifs] : notifiedSyncMessages) { - auto& existingNotifications = mUnhandledNotifications[agent]; - existingNotifications.insert(existingNotifications.end(), std::make_move_iterator(notifs.begin()), - std::make_move_iterator(notifs.end())); + if (!notifs.empty()) + { + std::stringstream ss(notifs.front()); + auto const notificationInfo = NotificationInfo::deserialize(ss); + debugLogs.push_back( + NotificationDebugLog{agent, notifs.size(), 0, 0, describeNotificationInfo(notificationInfo)}); + } + } + { + std::lock_guard lock(mNotificationMutex); + + // Merge new notifications with existing ones + for (auto const& [agent, notifs] : notifiedSyncMessages) + { + auto& existingNotifications = mUnhandledNotifications[agent]; + existingNotifications.insert(existingNotifications.end(), std::make_move_iterator(notifs.begin()), + std::make_move_iterator(notifs.end())); + } + if (!debugLogs.empty()) + { + size_t pendingTotal{0}; + for (auto const& [agent, notifs] : mUnhandledNotifications) + { + (void) agent; + pendingTotal += notifs.size(); + } + for (auto& debugLog : debugLogs) + { + debugLog.pendingForAgent = mUnhandledNotifications[debugLog.agent].size(); + debugLog.pendingTotal = pendingTotal; + } + } + } + for (auto const& debugLog : debugLogs) + { + TLLM_LOG_INFO( + "[disagg-debug] C++ AgentConnectionManager notifications received: localAgent=%s sourceAgent=%s " + "newCount=%zu pendingForSource=%zu pendingTotal=%zu first={%s}", + mAgentName.c_str(), debugLog.agent.c_str(), debugLog.newCount, debugLog.pendingForAgent, + debugLog.pendingTotal, debugLog.firstNotification.c_str()); } } @@ -585,6 +748,18 @@ template void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { + auto const startTime = std::chrono::steady_clock::now(); + auto nextLogTime = startTime + std::chrono::seconds(30); + char const* notificationType = "unknown"; + if constexpr (std::is_same_v) + { + notificationType = "NotificationSyncInfo"; + } + else if constexpr (std::is_same_v) + { + notificationType = "ReadySignalInfo"; + } + while (!terminateFlag.load()) { @@ -594,6 +769,25 @@ void AgentConnectionManager::waitForNotification( } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); + auto const now = std::chrono::steady_clock::now(); + if (now >= nextLogTime) + { + size_t pendingNotificationCount = 0; + for (auto const& [agent, notifications] : mUnhandledNotifications) + { + pendingNotificationCount += notifications.size(); + } + auto const elapsedMs = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_INFO( + "[disagg-debug] C++ waitForNotification still waiting: type=%s remoteAgent=%s " + "expectedAgent=%s tag=%llu pendingAgents=%zu pendingNotifications=%zu elapsedMs=%lld " + "terminate=%d running=%d localAgent=%s", + notificationType, remoteAgentName.c_str(), expectedInfo.mAgentName.c_str(), + static_cast(expectedInfo.mContext.getTag()), mUnhandledNotifications.size(), + pendingNotificationCount, static_cast(elapsedMs), static_cast(terminateFlag.load()), + static_cast(mIsRunning.load()), mAgentName.c_str()); + nextLogTime = now + std::chrono::seconds(30); + } auto it = mUnhandledNotifications.begin(); while (it != mUnhandledNotifications.end()) { @@ -618,6 +812,11 @@ void AgentConnectionManager::waitForNotification( && notificationData.mAgentName == expectedInfo.mAgentName) { erase = true; + TLLM_LOG_INFO( + "[disagg-debug] C++ waitForNotification matched: type=%s remoteAgent=%s " + "expectedAgent=%s tag=%llu localAgent=%s", + notificationType, remoteAgentName.c_str(), expectedInfo.mAgentName.c_str(), + static_cast(expectedInfo.mContext.getTag()), mAgentName.c_str()); notifIt = notifs.erase(notifIt); if (notifs.empty()) { @@ -638,6 +837,12 @@ void AgentConnectionManager::waitForNotification( expectedInfo.mIsReady = readySignalData.mIsReady; erase = true; + TLLM_LOG_INFO( + "[disagg-debug] C++ waitForNotification matched: type=%s remoteAgent=%s " + "expectedAgent=%s tag=%llu isReady=%d localAgent=%s", + notificationType, remoteAgentName.c_str(), expectedInfo.mAgentName.c_str(), + static_cast(expectedInfo.mContext.getTag()), + static_cast(expectedInfo.mIsReady), mAgentName.c_str()); notifIt = notifs.erase(notifIt); if (notifs.empty()) { diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index bad3e184f983..e2a6ad74c3fe 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -485,6 +486,7 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { auto startTime = std::chrono::steady_clock::now(); + auto nextLogTime = startTime + std::chrono::seconds(30); while (true) { @@ -498,6 +500,18 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const return TransferState::kFAILURE; } + auto const now = std::chrono::steady_clock::now(); + auto const elapsed = std::chrono::duration_cast(now - startTime).count(); + if (now >= nextLogTime) + { + TLLM_LOG_INFO( + "[disagg-debug] C++ NIXL transfer wait still in progress: handle=%p timeoutMs=%lld " + "elapsedMs=%lld status=%s", + static_cast(mHandle), static_cast(timeout_ms), static_cast(elapsed), + nixlEnumStrings::statusStr(status).c_str()); + nextLogTime = now + std::chrono::seconds(30); + } + // If timeout_ms < 0, wait indefinitely until status is not NIXL_IN_PROG if (timeout_ms < 0) { @@ -506,9 +520,6 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const } // Check if timeout has elapsed - auto elapsed - = std::chrono::duration_cast(std::chrono::steady_clock::now() - startTime) - .count(); if (elapsed >= timeout_ms) { return TransferState::kIN_PROGRESS; @@ -818,8 +829,10 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c { auto status = mRawAgent->genNotif(name, syncMessage); - TLLM_CHECK_WITH_INFO( - status == NIXL_SUCCESS, "genNotif failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); + auto const statusString = nixlEnumStrings::statusStr(status); + TLLM_LOG_INFO("[disagg-debug] C++ NIXL genNotif returned: selfAgent=%s remoteAgent=%s status=%s payloadBytes=%zu", + mName.c_str(), name.c_str(), statusString.c_str(), syncMessage.size()); + TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, "genNotif failed with status: %s", statusString.c_str()); } [[nodiscard]] std::unordered_map> NixlTransferAgent::getNotifiedSyncMessages() @@ -827,8 +840,29 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c nixl_notifs_t notifs; auto status = mRawAgent->getNotifs(notifs); - TLLM_CHECK_WITH_INFO( - status == NIXL_SUCCESS, "getNotifs failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); + auto const statusString = nixlEnumStrings::statusStr(status); + TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, "getNotifs failed with status: %s", statusString.c_str()); + if (!notifs.empty()) + { + size_t totalCount{0}; + std::ostringstream sourceCounts; + bool firstSource{true}; + for (auto const& [agent, messages] : notifs) + { + totalCount += messages.size(); + if (!firstSource) + { + sourceCounts << ","; + } + firstSource = false; + sourceCounts << agent << ":" << messages.size(); + } + auto const sourceCountsString = sourceCounts.str(); + TLLM_LOG_INFO( + "[disagg-debug] C++ NIXL getNotifs returned: selfAgent=%s status=%s sources=%zu total=%zu " + "sourceCounts={%s}", + mName.c_str(), statusString.c_str(), notifs.size(), totalCount, sourceCountsString.c_str()); + } return notifs; } diff --git a/examples/disaggregated/clients/disagg_client.py b/examples/disaggregated/clients/disagg_client.py index e6c85d32231d..653fd9784743 100644 --- a/examples/disaggregated/clients/disagg_client.py +++ b/examples/disaggregated/clients/disagg_client.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import asyncio import json @@ -7,18 +22,32 @@ import aiohttp import yaml -logging.basicConfig(level=logging.INFO) +logging.basicConfig(level=logging.INFO, + format="%(asctime)s %(levelname)s:%(name)s:%(message)s") +LOGGER = logging.getLogger(__name__) + + +def _prompt_summary(prompt): + if isinstance(prompt, str): + return f"type=str chars={len(prompt)}" + if isinstance(prompt, list): + return f"type=list len={len(prompt)}" + return f"type={type(prompt).__name__}" + + +def _elapsed_ms(start_time): + return (time.monotonic() - start_time) * 1000.0 async def wait_for_server(session, server_host, server_port, timeout): url = f"http://{server_host}:{server_port}/health" start_time = time.time() - logging.info("Waiting for server to start") + LOGGER.info("Waiting for server to start: url=%s timeout=%ss", url, timeout) while time.time() - start_time < timeout: try: async with session.get(url) as response: if response.status == 200: - logging.info("Server is ready.") + LOGGER.info("Server is ready.") return except aiohttp.ClientError: pass @@ -27,7 +56,8 @@ async def wait_for_server(session, server_host, server_port, timeout): async def send_request(session, server_host, server_port, model, prompt, - max_tokens, temperature, streaming, ignore_eos): + max_tokens, temperature, streaming, ignore_eos, + request_index): url = f"http://{server_host}:{server_port}/v1/completions" headers = {"Content-Type": "application/json"} data = { @@ -40,38 +70,63 @@ async def send_request(session, server_host, server_port, model, prompt, if streaming: data["stream"] = True - async with session.post(url, headers=headers, json=data) as response: - if response.status != 200: - raise Exception(f"Error: {await response.text()}") - - if streaming: - text = "" - async for line in response.content: - if line: - line = line.decode('utf-8').strip() - if line == "data: [DONE]": - break - if line.startswith("data: "): - line = line[len("data: "):] - response_json = json.loads(line) - choices = response_json.get("choices", []) - if not choices: - continue - text += choices[0].get("text", "") - logging.info(text) - return text - else: - response_json = await response.json() - choices = response_json.get("choices", []) - if not choices: - raise ValueError("Missing choices in completion response") - text = choices[0].get("text", "") - logging.info(text) - return text + start_time = time.monotonic() + LOGGER.info( + "completion request start: index=%s url=%s model=%s stream=%s " + "max_tokens=%s temperature=%s ignore_eos=%s prompt=%s", request_index, + url, model, streaming, max_tokens, temperature, ignore_eos, + _prompt_summary(prompt)) + try: + async with session.post(url, headers=headers, json=data) as response: + LOGGER.info( + "completion response headers: index=%s status=%s " + "content_type=%s elapsed_ms=%.2f", request_index, + response.status, response.headers.get("Content-Type"), + _elapsed_ms(start_time)) + if response.status != 200: + raise RuntimeError(f"Error: {await response.text()}") + + if streaming: + text = "" + chunk_count = 0 + async for line in response.content: + if line: + chunk_count += 1 + line = line.decode('utf-8').strip() + if line == "data: [DONE]": + break + if line.startswith("data: "): + line = line[len("data: "):] + response_json = json.loads(line) + choices = response_json.get("choices", []) + if not choices: + continue + text += choices[0].get("text", "") + LOGGER.info( + "completion streaming done: index=%s chunks=%s " + "text_chars=%s elapsed_ms=%.2f text=%s", request_index, + chunk_count, len(text), _elapsed_ms(start_time), text) + return text + else: + response_json = await response.json() + choices = response_json.get("choices", []) + if not choices: + raise ValueError("Missing choices in completion response") + text = choices[0].get("text", "") + LOGGER.info( + "completion done: index=%s text_chars=%s " + "elapsed_ms=%.2f text=%s", request_index, len(text), + _elapsed_ms(start_time), text) + return text + except (asyncio.TimeoutError, aiohttp.ClientError, RuntimeError, ValueError, + json.JSONDecodeError): + LOGGER.exception("completion request failed: index=%s elapsed_ms=%.2f", + request_index, _elapsed_ms(start_time)) + raise async def send_chat_request(session, server_host, server_port, model, prompt, - max_tokens, temperature, streaming): + max_tokens, temperature, streaming, request_index): url = f"http://{server_host}:{server_port}/v1/chat/completions" headers = {"Content-Type": "application/json"} data = { @@ -92,37 +147,62 @@ async def send_chat_request(session, server_host, server_port, model, prompt, if streaming: data["stream"] = True - async with session.post(url, headers=headers, json=data) as response: - if response.status != 200: - raise Exception(f"Error: {await response.text()}") - - if streaming: - text = "" - async for line in response.content: - if line: - line = line.decode('utf-8').strip() - if line == "data: [DONE]": - break - if line.startswith("data: "): - line = line[len("data: "):] - response_json = json.loads(line) - choices = response_json.get("choices", []) - if not choices: - continue - delta = choices[0].get("delta", {}) - content = delta.get("content") - if content is not None: - text += content - logging.info(text) - return text - else: - response_json = await response.json() - choices = response_json.get("choices", []) - if not choices: - raise ValueError("Missing choices in chat completion response") - text = choices[0].get("message", {}).get("content", "") - logging.info(text) - return text + start_time = time.monotonic() + LOGGER.info( + "chat request start: index=%s url=%s model=%s stream=%s " + "max_tokens=%s temperature=%s prompt=%s", request_index, url, model, + streaming, max_tokens, temperature, _prompt_summary(prompt)) + try: + async with session.post(url, headers=headers, json=data) as response: + LOGGER.info( + "chat response headers: index=%s status=%s " + "content_type=%s elapsed_ms=%.2f", request_index, + response.status, response.headers.get("Content-Type"), + _elapsed_ms(start_time)) + if response.status != 200: + raise RuntimeError(f"Error: {await response.text()}") + + if streaming: + text = "" + chunk_count = 0 + async for line in response.content: + if line: + chunk_count += 1 + line = line.decode('utf-8').strip() + if line == "data: [DONE]": + break + if line.startswith("data: "): + line = line[len("data: "):] + response_json = json.loads(line) + choices = response_json.get("choices", []) + if not choices: + continue + delta = choices[0].get("delta", {}) + content = delta.get("content") + if content is not None: + text += content + LOGGER.info( + "chat streaming done: index=%s chunks=%s " + "text_chars=%s elapsed_ms=%.2f text=%s", request_index, + chunk_count, len(text), _elapsed_ms(start_time), text) + return text + else: + response_json = await response.json() + choices = response_json.get("choices", []) + if not choices: + raise ValueError( + "Missing choices in chat completion response") + text = choices[0].get("message", {}).get("content", "") + LOGGER.info( + "chat done: index=%s text_chars=%s elapsed_ms=%.2f " + "text=%s", request_index, len(text), + _elapsed_ms(start_time), text) + return text + except (asyncio.TimeoutError, aiohttp.ClientError, RuntimeError, ValueError, + json.JSONDecodeError): + LOGGER.exception("chat request failed: index=%s elapsed_ms=%.2f", + request_index, _elapsed_ms(start_time)) + raise async def main(): @@ -169,9 +249,19 @@ async def main(): server_host = config.get('hostname', 'localhost') server_port = config.get('port', 8000) model = config.get('model', 'TinyLlama/TinyLlama-1.1B-Chat-v1.0') + LOGGER.info( + "disagg client config: config_file=%s prompts_file=%s endpoint=%s " + "streaming=%s output_file=%s server=%s:%s model=%s max_tokens=%s " + "temperature=%s ignore_eos=%s server_start_timeout=%s", + args.disagg_config_file, args.prompts_file, args.endpoint, + args.streaming, args.output_file, server_host, server_port, model, + args.max_tokens, args.temperature, args.ignore_eos, + args.server_start_timeout) with open(args.prompts_file, "r") as file: prompts = json.load(file) + LOGGER.info("loaded prompts: count=%s summaries=%s", len(prompts), + [_prompt_summary(prompt) for prompt in prompts]) async with aiohttp.ClientSession() as session: @@ -183,16 +273,24 @@ async def main(): tasks = [ send_request(session, server_host, server_port, model, prompt, args.max_tokens, args.temperature, args.streaming, - args.ignore_eos) for prompt in prompts + args.ignore_eos, i) + for i, prompt in enumerate(prompts) ] elif args.endpoint == "chat": tasks = [ send_chat_request(session, server_host, server_port, model, prompt, args.max_tokens, args.temperature, - args.streaming) for prompt in prompts + args.streaming, i) + for i, prompt in enumerate(prompts) ] + else: + raise ValueError(f"Unknown endpoint: {args.endpoint}") + LOGGER.info("awaiting client tasks: count=%s endpoint=%s", len(tasks), + args.endpoint) responses = await asyncio.gather(*tasks) + LOGGER.info("client tasks completed: count=%s endpoint=%s", + len(responses), args.endpoint) with open(args.output_file, "w") as file: json.dump(responses, file, indent=2) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index b71d9fea8921..db524e7e7c2d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time from abc import ABC, abstractmethod from os import getenv from typing import Any, Dict, List, Optional @@ -19,6 +35,52 @@ BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +def _request_summary(req: LlmRequest) -> str: + disagg_params = getattr(req, "py_disaggregated_params", None) + if disagg_params is None: + disagg_summary = "none" + else: + encoded_opaque_state = getattr(disagg_params, "encoded_opaque_state", + None) + first_gen_tokens = getattr(disagg_params, "first_gen_tokens", None) + draft_tokens = getattr(disagg_params, "draft_tokens", None) + disagg_summary = ( + f"request_type={getattr(disagg_params, 'request_type', None)!r} " + f"ctx_request_id={getattr(disagg_params, 'ctx_request_id', None)!r} " + f"disagg_request_id={getattr(disagg_params, 'disagg_request_id', None)!r} " + f"schedule_style={getattr(disagg_params, 'schedule_style', None)!r} " + f"ctx_dp_rank={getattr(disagg_params, 'ctx_dp_rank', None)!r} " + f"ctx_info_endpoint={getattr(disagg_params, 'ctx_info_endpoint', None)!r} " + f"opaque_state_bytes={len(encoded_opaque_state) if encoded_opaque_state else 0} " + f"first_gen_tokens={len(first_gen_tokens) if first_gen_tokens else 0} " + f"draft_tokens={len(draft_tokens) if draft_tokens else 0}") + + return ( + f"request_id={getattr(req, 'py_request_id', getattr(req, 'request_id', None))!r} " + f"state={getattr(req, 'state', None)!r} " + f"prompt_len={getattr(req, 'py_prompt_len', getattr(req, 'prompt_len', None))!r} " + f"max_new_tokens={getattr(req, 'py_max_new_tokens', getattr(req, 'max_new_tokens', None))!r} " + f"seq_slot={getattr(req, 'py_seq_slot', None)!r} " + f"client_id={getattr(req, 'py_client_id', None)!r} " + f"is_child={getattr(req, 'is_child', None)!r} " + f"parent_request_id={getattr(req, 'parent_request_id', None)!r} " + f"kv_transfer_start={getattr(req, 'py_kv_transfer_start_time', None)!r} " + f"kv_transfer_timed_out={getattr(req, 'py_kv_transfer_timed_out', None)!r} " + f"disagg=({disagg_summary})") + + +def _transfer_result_summary(result: Any) -> str: + if isinstance(result, tuple): + parts = [] + for item in result: + try: + parts.append(str(len(item))) + except TypeError: + parts.append(type(item).__name__) + return f"tuple_lengths={parts}" + return repr(result) + + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: return WorldConfig(tensor_parallelism=mapping.tp_size, @@ -65,9 +127,9 @@ def create_kv_cache_transceiver( "MPI CacheTransceiver is deprecated, UCX or NIXL is recommended") elif cache_transceiver_config.backend == "UCX": logger.info( - f"Using UCX kv-cache transceiver. If your devices are not in the same domain, please consider setting " - f"UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " - f"hangs or lower-than-expected performance.") + "Using UCX kv-cache transceiver. If your devices are not in the same domain, please consider setting " + "UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " + "hangs or lower-than-expected performance.") # Select transceiver implementation based on transceiver_runtime # transceiver_runtime == None or "CPP" -> use C++ transceiver (default) @@ -167,6 +229,18 @@ def __init__(self, self.kv_transfer_timeout_ms = cache_transceiver_config.kv_transfer_timeout_ms self.kv_transfer_sender_future_timeout_ms = cache_transceiver_config.kv_transfer_sender_future_timeout_ms + logger.info( + f"[disagg-debug] creating C++ KV cache transceiver: " + f"rank={mapping.rank} tp={mapping.tp_size} pp={mapping.pp_size} " + f"cp={mapping.cp_size} gpus_per_node={mapping.gpus_per_node} " + f"attention_dp={mapping.enable_attention_dp} " + f"backend={cache_transceiver_config.backend} " + f"runtime={cache_transceiver_config.transceiver_runtime} " + f"max_tokens_in_buffer={cache_transceiver_config.max_tokens_in_buffer} " + f"kv_transfer_timeout_ms={self.kv_transfer_timeout_ms} " + f"kv_transfer_sender_future_timeout_ms={self.kv_transfer_sender_future_timeout_ms} " + f"tokens_per_block={tokens_per_block} dtype={dtype} " + f"pp_layer_num_per_pp_rank={pp_layer_num_per_pp_rank}") # Get RNN state manager and layer distribution if mamba_cache_manager is provided rnn_state_manager = None @@ -186,27 +260,83 @@ def __init__(self, pp_layer_num_per_pp_rank, dtype, attention_type, cache_transceiver_config._to_pybind(), rnn_state_manager, rnn_layer_num_per_pp_rank) + logger.info("[disagg-debug] C++ KV cache transceiver created") def respond_and_send_async(self, req: LlmRequest): - return self.impl.respond_and_send_async(req) + start_time = time.monotonic() + logger.info( + f"[disagg-debug] respond_and_send_async begin: req=({_request_summary(req)})" + ) + result = self.impl.respond_and_send_async(req) + logger.info( + f"[disagg-debug] respond_and_send_async end: elapsed_s={time.monotonic() - start_time:.3f} " + f"req=({_request_summary(req)}) result={result!r}") + return result def request_and_receive_sync(self, req: LlmRequest): - return self.impl.request_and_receive_sync(req) + start_time = time.monotonic() + logger.info( + f"[disagg-debug] request_and_receive_sync begin: req=({_request_summary(req)})" + ) + result = self.impl.request_and_receive_sync(req) + logger.info( + f"[disagg-debug] request_and_receive_sync end: elapsed_s={time.monotonic() - start_time:.3f} " + f"req=({_request_summary(req)}) result={result!r}") + return result def request_and_receive_async(self, req: LlmRequest): - return self.impl.request_and_receive_async(req) + start_time = time.monotonic() + logger.info( + f"[disagg-debug] request_and_receive_async begin: req=({_request_summary(req)})" + ) + result = self.impl.request_and_receive_async(req) + logger.info( + f"[disagg-debug] request_and_receive_async end: elapsed_s={time.monotonic() - start_time:.3f} " + f"req=({_request_summary(req)}) result={result!r}") + return result def check_context_transfer_status(self, at_least_request_num: int): - return self.impl.check_context_transfer_status(at_least_request_num) + start_time = time.monotonic() + log_fn = logger.info if at_least_request_num > 0 else logger.debug + log_fn( + f"[disagg-debug] check_context_transfer_status begin: at_least_request_num={at_least_request_num}" + ) + result = self.impl.check_context_transfer_status(at_least_request_num) + log_fn( + f"[disagg-debug] check_context_transfer_status end: at_least_request_num={at_least_request_num} " + f"elapsed_s={time.monotonic() - start_time:.3f} result={_transfer_result_summary(result)}" + ) + return result def check_gen_transfer_status(self, at_least_request_num: int): - return self.impl.check_gen_transfer_status(at_least_request_num) + start_time = time.monotonic() + log_fn = logger.info if at_least_request_num > 0 else logger.debug + log_fn( + f"[disagg-debug] check_gen_transfer_status begin: at_least_request_num={at_least_request_num}" + ) + result = self.impl.check_gen_transfer_status(at_least_request_num) + log_fn( + f"[disagg-debug] check_gen_transfer_status end: at_least_request_num={at_least_request_num} " + f"elapsed_s={time.monotonic() - start_time:.3f} result={_transfer_result_summary(result)}" + ) + return result def check_gen_transfer_complete(self): - return self.impl.check_gen_transfer_complete() + result = self.impl.check_gen_transfer_complete() + logger.debug( + f"[disagg-debug] check_gen_transfer_complete result={result!r}") + return result def cancel_request(self, req: LlmRequest): - return self.impl.cancel_request(req) + start_time = time.monotonic() + logger.info( + f"[disagg-debug] cancel_request begin: req=({_request_summary(req)})" + ) + result = self.impl.cancel_request(req) + logger.info( + f"[disagg-debug] cancel_request end: elapsed_s={time.monotonic() - start_time:.3f} " + f"req=({_request_summary(req)}) result={result!r}") + return result def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 611aae6a42f7..3e52aad9a238 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import dataclasses import datetime import functools @@ -8,7 +23,7 @@ from contextlib import contextmanager from enum import IntEnum from queue import Queue -from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union import torch @@ -89,6 +104,77 @@ PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +def _safe_len(value: Any) -> int: + if value is None: + return 0 + try: + return len(value) + except TypeError: + return -1 + + +def _summarize_disagg_params(params: Any) -> str: + if params is None: + return "none" + encoded_opaque_state = getattr(params, "encoded_opaque_state", None) + first_gen_tokens = getattr(params, "first_gen_tokens", None) + draft_tokens = getattr(params, "draft_tokens", None) + return ( + f"request_type={getattr(params, 'request_type', None)!r} " + f"ctx_request_id={getattr(params, 'ctx_request_id', None)!r} " + f"disagg_request_id={getattr(params, 'disagg_request_id', None)!r} " + f"schedule_style={getattr(params, 'schedule_style', None)!r} " + f"ctx_dp_rank={getattr(params, 'ctx_dp_rank', None)!r} " + f"ctx_info_endpoint={getattr(params, 'ctx_info_endpoint', None)!r} " + f"opaque_state_bytes={len(encoded_opaque_state) if encoded_opaque_state else 0} " + f"first_gen_tokens={_safe_len(first_gen_tokens)} " + f"draft_tokens={_safe_len(draft_tokens)}") + + +def _summarize_disagg_request(req: LlmRequest) -> str: + return ( + f"request_id={getattr(req, 'py_request_id', getattr(req, 'request_id', None))!r} " + f"state={getattr(req, 'state', None)!r} " + f"prompt_len={getattr(req, 'py_prompt_len', getattr(req, 'prompt_len', None))!r} " + f"orig_prompt_len={getattr(req, 'py_orig_prompt_len', getattr(req, 'orig_prompt_len', None))!r} " + f"max_new_tokens={getattr(req, 'py_max_new_tokens', getattr(req, 'max_new_tokens', None))!r} " + f"context_current_position={getattr(req, 'context_current_position', None)!r} " + f"decoding_iter={getattr(req, 'py_decoding_iter', None)!r} " + f"seq_slot={getattr(req, 'py_seq_slot', None)!r} " + f"client_id={getattr(req, 'py_client_id', None)!r} " + f"is_child={getattr(req, 'is_child', None)!r} " + f"parent_request_id={getattr(req, 'parent_request_id', None)!r} " + f"is_context_only={getattr(req, 'is_context_only_request', None)!r} " + f"is_gen_init={getattr(req, 'is_disagg_generation_init_state', None)!r} " + f"is_gen_transfer_in_progress={getattr(req, 'is_disagg_generation_transmission_in_progress', None)!r} " + f"is_gen_transfer_complete={getattr(req, 'is_disagg_generation_transmission_complete', None)!r} " + f"kv_transfer_start={getattr(req, 'py_kv_transfer_start_time', None)!r} " + f"kv_transfer_timed_out={getattr(req, 'py_kv_transfer_timed_out', None)!r} " + f"disagg=({_summarize_disagg_params(getattr(req, 'py_disaggregated_params', None))})" + ) + + +def _summarize_disagg_requests(requests: Iterable[LlmRequest], + limit: int = 8) -> str: + requests = list(requests) + summarized = [_summarize_disagg_request(req) for req in requests[:limit]] + if len(requests) > limit: + summarized.append(f"... {len(requests) - limit} more") + return "[" + "; ".join(summarized) + "]" + + +def _summarize_transfer_result(result: Any) -> str: + if isinstance(result, tuple): + parts = [] + for item in result: + try: + parts.append(str(len(item))) + except TypeError: + parts.append(type(item).__name__) + return f"tuple_lengths={parts}" + return repr(result) + + class PPCommTag(IntEnum): """ Unique tags for pipeline parallelism communication. @@ -666,6 +752,12 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() def _end_transfer_and_maybe_terminate(self, request: LlmRequest): + logger.info( + f"[disagg-debug] ending KV transfer: " + f"request=({_summarize_disagg_request(request)}) " + f"request_in_active={request in self.active_requests} " + f"should_store_blocks={self.async_transfer_manager.should_store_blocks}" + ) if self.kv_cache_transceiver and request in self.active_requests: # Fast-transfer: KV transfer completed in the same iteration # before _handle_responses could run. Create the response now @@ -3424,10 +3516,11 @@ def _schedule(self): @nvtx_range("_check_disagg_gen_transfer_status") def _check_disagg_gen_transfer_status(self): - need_check = any([ - req.is_disagg_generation_transmission_in_progress - for req in self.active_requests - ]) + in_progress_reqs = [ + req for req in self.active_requests + if req.is_disagg_generation_transmission_in_progress + ] + need_check = bool(in_progress_reqs) non_gen_first_reqs = [ req for req in self.active_requests if req.py_disaggregated_params and req.py_disaggregated_params. @@ -3439,6 +3532,12 @@ def _check_disagg_gen_transfer_status(self): if need_check: at_least_num = 1 if need_check_one else 0 + logger.info( + f"[disagg-debug] generation transfer status poll: " + f"at_least_num={at_least_num} need_check_one={need_check_one} " + f"in_progress={_summarize_disagg_requests(in_progress_reqs)} " + f"non_gen_first={_summarize_disagg_requests(non_gen_first_reqs)}" + ) self._check_disagg_gen_cache_transfer_status(at_least_num) return @@ -3577,6 +3676,11 @@ def _pad_attention_dp_dummy_request(self): @nvtx_range("_prepare_disagg_gen_init") def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if fitting_disagg_gen_init_requests: + logger.info( + f"[disagg-debug] prepare disagg generation init: " + f"count={len(fitting_disagg_gen_init_requests)} " + f"requests={_summarize_disagg_requests(fitting_disagg_gen_init_requests)}" + ) disagg_gen_init_to_prepare = ScheduledRequests() disagg_gen_init_to_prepare.context_requests_last_chunk = fitting_disagg_gen_init_requests @@ -3587,12 +3691,21 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if (resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager. resource_managers[resource_mgr_type] is not None): + logger.info( + f"[disagg-debug] prepare resources for disagg generation init: " + f"resource_mgr_type={resource_mgr_type} " + f"request_count={len(fitting_disagg_gen_init_requests)}" + ) self.resource_manager.resource_managers[ resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) # Trigger KV cache exchange for new disagg_gen_init_requests self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) + logger.info( + f"[disagg-debug] prepare disagg generation init done: " + f"requests={_summarize_disagg_requests(fitting_disagg_gen_init_requests)}" + ) @nvtx_range("_prepare_disagg_gen_transmission_complete") def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): @@ -3601,6 +3714,11 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): if req.is_disagg_generation_transmission_complete: cache_trans_complete_requests.append(req) if len(cache_trans_complete_requests) > 0: + logger.info( + f"[disagg-debug] disagg generation transmission complete: " + f"count={len(cache_trans_complete_requests)} " + f"requests={_summarize_disagg_requests(cache_trans_complete_requests)}" + ) requests = ScheduledRequests() requests.context_requests_last_chunk = cache_trans_complete_requests self.resource_manager.resource_managers[ @@ -3668,24 +3786,43 @@ def _has_prepended_logits(self, req) -> bool: @nvtx_range("_recv_disagg_gen_cache") def _recv_disagg_gen_cache(self, new_gen_reqs): + logger.info( + f"[disagg-debug] recv disagg generation cache start: " + f"count={len(new_gen_reqs)} " + f"disable_overlap={os.getenv('TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP')} " + f"benchmark_gen_only={os.getenv('TRTLLM_DISAGG_BENCHMARK_GEN_ONLY')} " + f"requests={_summarize_disagg_requests(new_gen_reqs)}") # For gen-only benchmarking, mark new gen request as transmission complete right away if os.getenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") == "1": for req in new_gen_reqs: req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + logger.info( + f"[disagg-debug] recv disagg generation cache skipped for gen-only benchmark: " + f"requests={_summarize_disagg_requests(new_gen_reqs)}") return if os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": for req in new_gen_reqs: + logger.info( + f"[disagg-debug] requesting synchronous generation KV cache receive: " + f"req=({_summarize_disagg_request(req)})") self.kv_cache_transceiver.request_and_receive_sync(req) else: for req in new_gen_reqs: + logger.info( + f"[disagg-debug] requesting asynchronous generation KV cache receive: " + f"req=({_summarize_disagg_request(req)})") self.kv_cache_transceiver.request_and_receive_async(req) if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: for req in new_gen_reqs: if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: req.py_kv_transfer_start_time = time.time() + logger.info( + f"[disagg-debug] generation KV transfer timeout tracking started: " + f"timeout_ms={self.kv_cache_transceiver.kv_transfer_timeout_ms} " + f"req=({_summarize_disagg_request(req)})") non_gen_first_active = [ req for req in self.active_requests @@ -3695,6 +3832,11 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): block_transfer = bool(non_gen_first_active) and all( req.is_disagg_generation_transmission_in_progress for req in non_gen_first_active) + logger.info( + f"[disagg-debug] checking generation KV cache transfer after receive request: " + f"block_transfer={block_transfer} at_least_num={1 if block_transfer else 0} " + f"non_gen_first_active={_summarize_disagg_requests(non_gen_first_active)} " + f"new_gen_reqs={_summarize_disagg_requests(new_gen_reqs)}") self._check_disagg_gen_cache_transfer_status(1 if block_transfer else 0) return @@ -3721,8 +3863,13 @@ def kv_connector_request_finished(req: LlmRequest): ) and not req.is_finished_due_to_cancellation: # 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. + logger.info(f"[disagg-debug] context KV cache send start: " + f"req=({_summarize_disagg_request(req)})") self.async_transfer_manager.start_transfer(req) self.kv_cache_transceiver.respond_and_send_async(req) + logger.info( + f"[disagg-debug] context KV cache send submitted: " + f"req=({_summarize_disagg_request(req)})") if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: req.py_kv_transfer_start_time = time.time() @@ -3757,8 +3904,22 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): @nvtx_range("_check_disagg_ctx_cache_transfer_status") def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): + requests_in_transfer = self.async_transfer_manager.requests_in_transfer( + ) + if atLeastNum > 0 or requests_in_transfer: + logger.info( + f"[disagg-debug] check context KV transfer status begin: " + f"atLeastNum={atLeastNum} " + f"inflight={_summarize_disagg_requests(requests_in_transfer.values())}" + ) + start_time = time.monotonic() finished_requests, error_requests = self.kv_cache_transceiver.check_context_transfer_status( atLeastNum) + if atLeastNum > 0 or finished_requests or error_requests: + logger.info( + f"[disagg-debug] check context KV transfer status result: " + f"atLeastNum={atLeastNum} elapsed_s={time.monotonic() - start_time:.3f} " + f"finished={finished_requests} errors={error_requests}") completed_req_ids = set(finished_requests + error_requests) @@ -3796,7 +3957,24 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): @nvtx_range("_check_disagg_gen_cache_transfer_status") def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0): + in_progress_reqs = [ + req for req in self.active_requests + if req.is_disagg_generation_transmission_in_progress + ] + if atLeastNum > 0 or in_progress_reqs: + logger.info( + f"[disagg-debug] check generation KV transfer status begin: " + f"atLeastNum={atLeastNum} " + f"in_progress={_summarize_disagg_requests(in_progress_reqs)}") + start_time = time.monotonic() result = self.kv_cache_transceiver.check_gen_transfer_status(atLeastNum) + if atLeastNum > 0 or in_progress_reqs: + logger.info( + f"[disagg-debug] check generation KV transfer status result: " + f"atLeastNum={atLeastNum} elapsed_s={time.monotonic() - start_time:.3f} " + f"result={_summarize_transfer_result(result)} " + f"in_progress_after={_summarize_disagg_requests(in_progress_reqs)}" + ) if isinstance(result, tuple): _, _, cancelled_reqs = result user_canceled_set = set(self.canceled_req_ids) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 0371f2da7e03..787d1e5a3e99 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -41,6 +41,55 @@ # yapf: enable +def _summarize_disagg_params(request: UCompletionRequest) -> str: + params = getattr(request, "disaggregated_params", None) + if params is None: + return "none" + encoded_opaque_state = getattr(params, "encoded_opaque_state", None) + first_gen_tokens = getattr(params, "first_gen_tokens", None) + draft_tokens = getattr(params, "draft_tokens", None) + return ( + f"request_type={getattr(params, 'request_type', None)!r} " + f"ctx_request_id={getattr(params, 'ctx_request_id', None)!r} " + f"disagg_request_id={getattr(params, 'disagg_request_id', None)!r} " + f"schedule_style={getattr(params, 'schedule_style', None)!r} " + f"ctx_dp_rank={getattr(params, 'ctx_dp_rank', None)!r} " + f"ctx_info_endpoint={getattr(params, 'ctx_info_endpoint', None)!r} " + f"opaque_state_bytes={len(encoded_opaque_state) if encoded_opaque_state else 0} " + f"first_gen_tokens={len(first_gen_tokens) if first_gen_tokens else 0} " + f"draft_tokens={len(draft_tokens) if draft_tokens else 0}" + ) + + +def _summarize_request(request: UCompletionRequest) -> str: + if isinstance(request, CompletionRequest): + prompt = request.prompt + if isinstance(prompt, str): + prompt_summary = f"prompt_chars={len(prompt)}" + elif isinstance(prompt, list): + prompt_summary = f"prompt_list_len={len(prompt)}" + else: + prompt_summary = f"prompt_type={type(prompt).__name__}" + elif isinstance(request, ChatCompletionRequest): + messages = getattr(request, "messages", []) + prompt_token_ids = getattr(request, "prompt_token_ids", None) + prompt_summary = ( + f"messages={len(messages)} " + f"prompt_token_ids={len(prompt_token_ids) if prompt_token_ids else 0}" + ) + else: + prompt_summary = f"request_type={type(request).__name__}" + + return ( + f"model={getattr(request, 'model', None)!r} " + f"stream={getattr(request, 'stream', None)!r} " + f"max_tokens={getattr(request, 'max_tokens', None)!r} " + f"temperature={getattr(request, 'temperature', None)!r} " + f"ignore_eos={getattr(request, 'ignore_eos', None)!r} " + f"{prompt_summary} disagg=({_summarize_disagg_params(request)})" + ) + + class OpenAIClient(ABC): async def send_request( self, @@ -129,14 +178,20 @@ async def _send_request( if server is None: server, _ = await self._router.get_next_server(request) url = f"http://{server}/{endpoint}" - logger.debug( - f"Sending {self._role} request {request.disaggregated_params.ctx_request_id} to {url}" + logger.info( + f"[disagg-debug] OpenAI client send start: role={self._role} " + f"endpoint={endpoint} server={server} url={url} " + f"request=({_summarize_request(request)})" ) try: self._metrics_collector.total_requests.inc() resp_generator = self._post_with_retry(server, url, request, hooks) if request.stream: # return the response generator, the request is not done yet + logger.info( + f"[disagg-debug] OpenAI client returning streaming generator: " + f"role={self._role} url={url}" + ) return resp_generator else: # consume the generator to get the response and return it directly when it's not streaming @@ -149,10 +204,19 @@ async def _send_request( else: hooks.on_first_token(server, request) hooks.on_resp_done(server, request, response) + logger.info( + f"[disagg-debug] OpenAI client send done: role={self._role} " + f"url={url} response_type={type(response).__name__ if response else None}" + ) return response except Exception: self._metrics_collector.error_requests.inc() # finish the request upon error + logger.error( + f"[disagg-debug] OpenAI client send failed: role={self._role} " + f"url={url} request=({_summarize_request(request)})", + traceback.format_exc(), + ) await self._finish_request(request) raise @@ -171,11 +235,24 @@ async def _post_with_retry( if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = self._disagg_id_generator() json_data = request.model_dump(exclude_unset=True, mode="json") + lines_yielded = 0 try: - lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() + logger.info( + f"[disagg-debug] HTTP post start: role={self._role} " + f"attempt={attempt}/{self._max_retries} server={server} " + f"url={url} stream={is_stream} payload_keys={sorted(json_data.keys())} " + f"request=({_summarize_request(request)})" + ) async with self._session.post(url, json=json_data) as http_response: content_type = http_response.headers.get("Content-Type", "") + logger.info( + f"[disagg-debug] HTTP post response headers: role={self._role} " + f"attempt={attempt}/{self._max_retries} url={url} " + f"status={http_response.status} reason={http_response.reason!r} " + f"content_type={content_type!r} elapsed_s=" + f"{get_steady_clock_now_in_seconds() - start_time:.3f}" + ) if not is_stream and "text/event-stream" in content_type: raise ValueError( "Received an event-stream although request stream was False" @@ -200,6 +277,11 @@ async def _post_with_retry( headers=http_response.headers, ) response_dict = await http_response.json() + logger.info( + f"[disagg-debug] HTTP post JSON received: role={self._role} " + f"url={url} response_keys={sorted(response_dict.keys())} " + f"elapsed_s={get_steady_clock_now_in_seconds() - start_time:.3f}" + ) # yield here since python forbids return statements in async generators yield response_dict # finish the request after the successful response @@ -207,8 +289,12 @@ async def _post_with_retry( self._metrics_collector.complete_latency_seconds.observe( get_steady_clock_now_in_seconds() - start_time ) + logger.info( + f"[disagg-debug] HTTP post complete: role={self._role} " + f"url={url} elapsed_s={get_steady_clock_now_in_seconds() - start_time:.3f}" + ) break # break and skip retries if the whole response is processed without exception - except (aiohttp.ClientError, OSError) as e: + except (aiohttp.ClientError, OSError, asyncio.TimeoutError) as e: if lines_yielded > 0: logger.error( f"Client error to {url}: {e} - cannot retry since {lines_yielded} lines were yielded", diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index be0817a32521..c0b7b8a29780 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -50,6 +50,67 @@ from tensorrt_llm.serve.router import KvCacheAwareRouter, Router +def _summarize_disagg_params(params: Optional[DisaggregatedParams]) -> str: + if params is None: + return "none" + encoded_opaque_state = params.encoded_opaque_state + return ( + f"request_type={params.request_type!r} " + f"ctx_request_id={params.ctx_request_id!r} " + f"disagg_request_id={params.disagg_request_id!r} " + f"schedule_style={params.schedule_style!r} " + f"ctx_dp_rank={params.ctx_dp_rank!r} " + f"ctx_info_endpoint={params.ctx_info_endpoint!r} " + f"opaque_state_bytes={len(encoded_opaque_state) if encoded_opaque_state else 0} " + f"first_gen_tokens={len(params.first_gen_tokens) if params.first_gen_tokens else 0} " + f"draft_tokens={len(params.draft_tokens) if params.draft_tokens else 0}" + ) + + +def _summarize_request(request: UCompletionRequest) -> str: + if isinstance(request, CompletionRequest): + prompt = request.prompt + if isinstance(prompt, str): + prompt_summary = f"prompt_chars={len(prompt)}" + elif isinstance(prompt, list): + prompt_summary = f"prompt_list_len={len(prompt)}" + else: + prompt_summary = f"prompt_type={type(prompt).__name__}" + elif isinstance(request, ChatCompletionRequest): + messages = getattr(request, "messages", []) + prompt_token_ids = getattr(request, "prompt_token_ids", None) + prompt_summary = ( + f"messages={len(messages)} " + f"prompt_token_ids={len(prompt_token_ids) if prompt_token_ids else 0}" + ) + else: + prompt_summary = f"request_type={type(request).__name__}" + + return ( + f"model={getattr(request, 'model', None)!r} " + f"stream={getattr(request, 'stream', None)!r} " + f"max_tokens={getattr(request, 'max_tokens', None)!r} " + f"temperature={getattr(request, 'temperature', None)!r} " + f"ignore_eos={getattr(request, 'ignore_eos', None)!r} " + f"{prompt_summary} disagg=({_summarize_disagg_params(request.disaggregated_params)})" + ) + + +def _summarize_response(response: Optional[UCompletionResponse]) -> str: + if response is None: + return "none" + if not response.choices: + return "choices=0" + choice = response.choices[0] + prompt_token_ids = getattr(response, "prompt_token_ids", None) + return ( + f"choices={len(response.choices)} " + f"finish_reason={getattr(choice, 'finish_reason', None)!r} " + f"prompt_token_ids={len(prompt_token_ids) if prompt_token_ids else 0} " + f"disagg=({_summarize_disagg_params(getattr(choice, 'disaggregated_params', None))})" + ) + + class OpenAIDisaggregatedService(OpenAIService): def __init__( self, @@ -136,18 +197,48 @@ async def _send_disagg_request_ctx_first( ctx_response = None gen_req = request disagg_request_id = get_global_disagg_request_id(self._config.node_id) + logger.info( + f"[disagg-debug] ctx-first request start: disagg_request_id={disagg_request_id} " + f"request=({_summarize_request(request)})" + ) if need_ctx: ctx_req = self._get_ctx_request(request, disagg_request_id) # ctx generator is empty ctx_server, _ = await self._ctx_router.get_next_server( ctx_req, exclude_server=gen_server ) - ctx_response = await self._ctx_client.send_request( - ctx_req, server=ctx_server, hooks=hooks + logger.info( + f"[disagg-debug] ctx-first selected context server: " + f"disagg_request_id={disagg_request_id} ctx_server={ctx_server} " + f"reserved_gen_server={gen_server} ctx_request=({_summarize_request(ctx_req)})" + ) + try: + ctx_response = await self._ctx_client.send_request( + ctx_req, server=ctx_server, hooks=hooks + ) + except Exception: + logger.error( + f"[disagg-debug] ctx-first context request failed: " + f"disagg_request_id={disagg_request_id} ctx_server={ctx_server}" + ) + raise + logger.info( + f"[disagg-debug] ctx-first context response: " + f"disagg_request_id={disagg_request_id} ctx_server={ctx_server} " + f"response=({_summarize_response(ctx_response)})" ) await self._verify_ctx_response(ctx_response) gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) + logger.info( + f"[disagg-debug] ctx-first generated generation request: " + f"disagg_request_id={disagg_request_id} gen_request=({_summarize_request(gen_req)})" + ) else: + logger.info( + f"[disagg-debug] ctx-first skipping context phase: " + f"disagg_request_id={disagg_request_id} reserved_gen_server={gen_server} " + f"request=({_summarize_request(gen_req)})" + ) # Clear synthetic disaggregated_params that may have been # injected by _extract_conversation_id (e.g. from the # X-Correlation-ID header). When need_ctx=False the gen @@ -165,11 +256,34 @@ async def _send_disagg_request_ctx_first( gen_server, _ = await self._gen_router.get_next_server( gen_req, exclude_server=ctx_server ) - gen_response = await self._gen_client.send_request( - gen_req, server=gen_server, hooks=hooks + logger.info( + f"[disagg-debug] ctx-first selected generation server: " + f"disagg_request_id={disagg_request_id} gen_server={gen_server} " + f"ctx_server={ctx_server} gen_request=({_summarize_request(gen_req)})" + ) + try: + gen_response = await self._gen_client.send_request( + gen_req, server=gen_server, hooks=hooks + ) + except Exception: + logger.error( + f"[disagg-debug] ctx-first generation request failed: " + f"disagg_request_id={disagg_request_id} gen_server={gen_server} " + f"ctx_server={ctx_server}" + ) + raise + logger.info( + f"[disagg-debug] ctx-first generation response object: " + f"disagg_request_id={disagg_request_id} gen_server={gen_server} " + f"response_type={type(gen_response).__name__}" ) return self._rewrite_disagg_usage(gen_response, ctx_response) else: + logger.info( + f"[disagg-debug] ctx-first context phase completed request without generation: " + f"disagg_request_id={disagg_request_id} ctx_server={ctx_server} " + f"ctx_response=({_summarize_response(ctx_response)})" + ) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 05f1a7dbcca0..0607d440061b 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -75,6 +75,135 @@ def cleanup_output_files(): pass +_DISAGG_DEBUG_ENV_KEYS = { + "AWS_OFI_NCCL_VERSION", + "BUILD_ID", + "BUILD_URL", + "CUDA_HOME", + "CUDA_VISIBLE_DEVICES", + "LD_LIBRARY_PATH", + "LD_PRELOAD", + "LIBRARY_PATH", + "LLM_MODELS_ROOT", + "NCCL_DEBUG", + "NCCL_DEBUG_SUBSYS", + "NCCL_IB_HCA", + "NCCL_IB_DISABLE", + "NCCL_NET", + "NCCL_NET_PLUGIN", + "NCCL_P2P_DISABLE", + "NCCL_P2P_LEVEL", + "NCCL_RUNTIME_CONNECT", + "NCCL_SOCKET_IFNAME", + "NVIDIA_DRIVER_CAPABILITIES", + "NVIDIA_VISIBLE_DEVICES", + "PATH", + "PYTHONPATH", + "TLLM_LOG_LEVEL_BY_MODULE", + "TLLM_NUMA_AWARE_WORKER_AFFINITY", + "TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", + "TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", + "TRTLLM_USE_MPI_KVCACHE", + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "UCX_CUDA_IPC_ENABLE_MNNVL", + "UCX_LOG_LEVEL", + "UCX_NET_DEVICES", + "UCX_RNDV_SCHEME", + "UCX_TLS", +} +_DISAGG_DEBUG_ENV_PREFIXES = ("NCCL_", "NIXL_", "UCX_", "CUDA_", "TRTLLM_", + "FI_", "OMPI_", "PMI_", "PMIX_") +_DISAGG_DEBUG_TESTS = ("gpt_oss_120b_harmony", "gpt_oss_120b_stress") + + +def _should_print_disagg_debug(test_desc: str | None) -> bool: + return test_desc in _DISAGG_DEBUG_TESTS + + +def _print_disagg_debug(message: str) -> None: + print(f"[disagg-debug] {message}", flush=True) + + +def _env_debug_snapshot(env: dict[str, str] | None) -> dict[str, str]: + source_env = os.environ if env is None else env + snapshot = {} + for key, value in source_env.items(): + if key in _DISAGG_DEBUG_ENV_KEYS or key.startswith( + _DISAGG_DEBUG_ENV_PREFIXES): + snapshot[key] = value + return dict(sorted(snapshot.items())) + + +def _run_disagg_debug_command(cmd: list[str], + env: dict[str, str] | None = None, + cwd: str | None = None, + timeout: int = 15) -> None: + _print_disagg_debug(f"diagnostic command start: {' '.join(cmd)}") + try: + result = subprocess.run(cmd, + cwd=cwd, + env=env, + capture_output=True, + text=True, + timeout=timeout, + check=False) + except FileNotFoundError as e: + _print_disagg_debug(f"diagnostic command missing: {cmd[0]}: {e}") + return + except subprocess.TimeoutExpired as e: + _print_disagg_debug( + f"diagnostic command timeout after {timeout}s: {' '.join(cmd)}") + if e.stdout: + _print_disagg_debug(f"stdout before timeout:\n{e.stdout[-20000:]}") + if e.stderr: + _print_disagg_debug(f"stderr before timeout:\n{e.stderr[-20000:]}") + return + + _print_disagg_debug( + f"diagnostic command done: returncode={result.returncode} cmd={' '.join(cmd)}" + ) + if result.stdout: + _print_disagg_debug(f"stdout:\n{result.stdout[-20000:]}") + if result.stderr: + _print_disagg_debug(f"stderr:\n{result.stderr[-20000:]}") + + +def _print_disagg_setup_debug(test_desc: str | None, config_file: str, + model: str | None, config: dict[str, Any], + ctx_worker_config: dict[str, Any], + gen_worker_config: dict[str, Any], + env: dict[str, str] | None, cwd: str | None, + work_dir: str, server_port: int) -> None: + if not _should_print_disagg_debug(test_desc): + return + + _print_disagg_debug( + f"setup start: test_desc={test_desc} config_file={config_file} " + f"model={model} real_model_path={os.path.realpath(model) if model else None} " + f"cwd={cwd} work_dir={work_dir} server_port={server_port}") + _print_disagg_debug( + "selected environment:\n" + + json.dumps(_env_debug_snapshot(env), indent=2, sort_keys=True)) + _print_disagg_debug("base disagg config:\n" + + yaml.safe_dump(config, sort_keys=True)) + _print_disagg_debug("context worker config:\n" + + yaml.safe_dump(ctx_worker_config, sort_keys=True)) + _print_disagg_debug("generation worker config:\n" + + yaml.safe_dump(gen_worker_config, sort_keys=True)) + + for cmd in (["nvidia-smi", "-L"], [ + "nvidia-smi", + "--query-gpu=index,name,pci.bus_id,uuid,compute_mode,memory.total", + "--format=csv,noheader", + ], ["nvidia-smi", "topo", "-m"], ["ls", "-l", "/sys/class/infiniband"], + ["ls", "-l", "/usr/lib/x86_64-linux-gnu/librdmacm.so.1" + ], ["ls", "-l", "/usr/lib64/librdmacm.so.1"], + ["ibv_devinfo", "-l"], ["rdma", "link", + "show"], ["ip", "-o", "addr", "show"]): + _run_disagg_debug_command(cmd, env=env, cwd=cwd) + + def get_default_disagg_cluster_config(): """Get default disaggregated cluster configuration.""" return { @@ -350,9 +479,14 @@ def run_client_tests(example_dir, """Run client tests against the disaggregated server.""" if client_test_set is None: client_test_set = get_client_test_set(test_desc) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"client test set: test_desc={test_desc} num_iters={num_iters} " + f"prompt_file={prompt_file} server_url={server_url} " + f"client_test_set={client_test_set}") client_dir = f"{example_dir}/clients" - for _ in range(num_iters): + for iter_idx in range(num_iters): client_cmd = [ 'python3', f'{client_dir}/disagg_client.py', '-c', f'{config_file}', '-p', f'{client_dir}/{prompt_file}', '--ignore-eos', @@ -379,6 +513,10 @@ def run_client_tests(example_dir, # Run completion test (non-streaming) if client_test_set.completion: + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"running completion client: iter={iter_idx} cmd={client_cmd}" + ) check_call(client_cmd, env=env, poll_procs=poll_procs) # Streaming client run @@ -386,12 +524,20 @@ def run_client_tests(example_dir, streaming_client_cmd = client_cmd + [ '--streaming', '-o', 'output_streaming.json' ] + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"running streaming completion client: iter={iter_idx} " + f"cmd={streaming_client_cmd}") check_call(streaming_client_cmd, env=env, poll_procs=poll_procs) # Run chat completion test if client_test_set.chat: chat_output = 'output_tool_calls.json' if test_desc == "tool_calls" else 'output_chat.json' chat_client_cmd = client_cmd + ['-e', 'chat', '-o', chat_output] + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"running chat client: iter={iter_idx} cmd={chat_client_cmd}" + ) check_call(chat_client_cmd, env=env, poll_procs=poll_procs) # Run streaming chat completion test @@ -399,6 +545,10 @@ def run_client_tests(example_dir, streaming_chat_client_cmd = client_cmd + [ '-e', 'chat', '--streaming', '-o', 'output_streaming_chat.json' ] + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"running streaming chat client: iter={iter_idx} " + f"cmd={streaming_chat_client_cmd}") check_call(streaming_chat_client_cmd, env=env, poll_procs=poll_procs) @@ -471,6 +621,7 @@ def setup_disagg_cluster( cwd: str | None = None, server_start_timeout: int = 300, schedule_style: str | None = None, + test_desc: str | None = None, ) -> tuple[dict[str, Any], list[ProcessWrapper], list[ProcessWrapper], ProcessWrapper, int, str]: """Load config, launch workers + disagg server, wait for ready. @@ -481,6 +632,7 @@ 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') + test_desc: Test description, used only to gate extra diagnostics Returns: tuple: (config, ctx_workers, gen_workers, disagg_server, server_port, work_dir) @@ -525,12 +677,21 @@ def setup_disagg_cluster( import torch num_gpus = torch.cuda.device_count() + _print_disagg_setup_debug(test_desc, config_file, model, config, + ctx_worker_config, gen_worker_config, env, cwd, + work_dir, server_port) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug(f"torch.cuda.device_count={num_gpus}") try: for i in range(num_ctx_instances): device_ids = ",".join( str(d) for d in dict.fromkeys((next_device + j) % num_gpus for j in range(gpus_per_ctx))) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"launching context worker: instance={i} " + f"device_ids={device_ids} gpus_per_ctx={gpus_per_ctx}") ctx_workers.append( run_ctx_worker(model, ctx_worker_config, @@ -538,12 +699,21 @@ def setup_disagg_cluster( port=0, device=device_ids, env=env)) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"context worker launched: instance={i} " + f"pid={ctx_workers[-1].process.pid} device_ids={device_ids}" + ) next_device += gpus_per_ctx for i in range(num_gen_instances): device_ids = ",".join( str(d) for d in dict.fromkeys((next_device + j) % num_gpus for j in range(gpus_per_gen))) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"launching generation worker: instance={i} " + f"device_ids={device_ids} gpus_per_gen={gpus_per_gen}") gen_workers.append( run_gen_worker(model, gen_worker_config, @@ -551,6 +721,11 @@ def setup_disagg_cluster( port=0, device=device_ids, env=env)) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"generation worker launched: instance={i} " + f"pid={gen_workers[-1].process.pid} device_ids={device_ids}" + ) next_device += gpus_per_gen # Build minimal server config and launch @@ -579,10 +754,18 @@ def setup_disagg_cluster( server_port, env=env, cwd=cwd) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"disagg server launched: pid={disagg_server.process.pid} " + f"server_config:\n{yaml.safe_dump(server_config, sort_keys=True)}" + ) asyncio.run( wait_for_disagg_server_ready(server_port, timeout=server_start_timeout)) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"disagg server ready: server_port={server_port}") except Exception: terminate(*ctx_workers, *gen_workers, disagg_server) shutil.rmtree(work_dir, ignore_errors=True) @@ -611,7 +794,8 @@ def run_disaggregated_test(example_dir, os.path.dirname(__file__)) config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ setup_disagg_cluster(config_file, model_name=model_path, env=env, cwd=cwd, - schedule_style=disagg_schedule_style) + schedule_style=disagg_schedule_style, + test_desc=test_desc) server_host = config.get("hostname", "localhost") @@ -626,6 +810,10 @@ def run_disaggregated_test(example_dir, dir=work_dir) with os.fdopen(temp_fd, 'w') as f: yaml.dump(client_config, f) + if _should_print_disagg_debug(test_desc): + _print_disagg_debug( + f"client config written: path={client_config_file}\n" + f"{yaml.safe_dump(client_config, sort_keys=True)}") # collect all worker processes for monitoring all_worker_procs = [w.process for w in ctx_workers @@ -1772,7 +1960,6 @@ def run_disaggregated_aiperf(config_file, env: Environment variables dict cwd: Working directory """ - cleanup_output_files() run_env = env.copy() run_env["UCX_TLS"] = get_ucx_tls() diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c845ca7dc5e4..564947c8612a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -228,7 +228,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1 disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_two_mtp[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322) disaggregated/test_disaggregated.py::test_disaggregated_diff_max_tokens[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6094100) disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) -disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] SKIP (https://nvbugs/6011317) disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6094100) disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6094100) disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6094100)