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 b3bb08b57425..e6403341b15f 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -748,12 +748,59 @@ std::vector const& AgentConnectionManager::getBufferKinds() const return mBufferKinds; } +namespace +{ + +/// Publishes a remote agent as "handshake in flight", then releases the held lock for the duration +/// of that handshake and takes it back afterwards. Owning the unlock/relock as well as the mark is +/// what makes the exception path safe: if loadRemoteAgent throws, the destructor still re-acquires +/// the lock before clearing the mark, so the mark never disappears unsynchronized and a failed +/// handshake stays retryable. +class ScopedAgentLoad +{ +public: + ScopedAgentLoad(std::unique_lock& lock, std::set& loadingAgents, + std::condition_variable& cv, std::string const& name) + : mLock(lock) + , mLoadingAgents(loadingAgents) + , mCv(cv) + , mName(name) + { + mLoadingAgents.insert(mName); + mLock.unlock(); + } + + ~ScopedAgentLoad() + { + // Unconditional: the constructor unlocked, and the guarded scope never touches the lock. + mLock.lock(); + mLoadingAgents.erase(mName); + mCv.notify_all(); + } + + ScopedAgentLoad(ScopedAgentLoad const&) = delete; + ScopedAgentLoad& operator=(ScopedAgentLoad const&) = delete; + +private: + std::unique_lock& mLock; + std::set& mLoadingAgents; + std::condition_variable& mCv; + std::string const& mName; +}; + +} // namespace + AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentName, std::string const& connectionInfo, std::optional metadata, bool isSender) { TLLM_LOG_DEBUG(mRank, "mAgentName: %s connect to %s", mAgentName.c_str(), remoteAgentName.c_str()); - std::scoped_lock lock(mConnectionsMutex); + std::unique_lock lock(mConnectionsMutex); + // Observe an in-flight handshake's result instead of racing it: starting a second load for the + // same peer would insert a second AgentConnection over the first, dangling the raw pointer an + // earlier caller already returned, and would let invalidateRemoteAgent below run against a poll + // still in progress. Only same-peer callers wait; every other peer proceeds immediately. + mLoadingAgentsCv.wait(lock, [&] { return mLoadingAgents.count(remoteAgentName) == 0; }); auto it = mConnections.find(remoteAgentName); if (it != mConnections.end()) { @@ -788,6 +835,11 @@ AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentN TLLM_CHECK_WITH_INFO(!isSender, "Sender shouldn't call loadRemoteAgent"); TLLM_LOG_DEBUG(mRank, "mAgentName: %s connect to %s with loadRemoteAgent", mAgentName.c_str(), remoteAgentName.c_str()); + // This overload waits on the peer's metadata reply -- unbounded network I/O. Holding + // mConnectionsMutex across it let one silent peer stall connect() for every other peer + // too, and stall notification draining with it: the sender thread blocks here while + // holding mNotificationMutex, which updateUnhandledNotifications then can never take. + ScopedAgentLoad const loading{lock, mLoadingAgents, mLoadingAgentsCv, remoteAgentName}; m_Agent->loadRemoteAgent(remoteAgentName, connectionInfo); } } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index 83f7df2541b9..7c3839a55f4d 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -24,7 +24,9 @@ #include "tensorrt_llm/executor/cacheCommunicator.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/transferAgent.h" +#include #include +#include namespace tensorrt_llm::executor::kv_cache { @@ -333,6 +335,10 @@ class AgentConnectionManager : public ConnectionManager private: std::map> mConnections; std::mutex mConnectionsMutex; + /// Remote agents whose loadRemoteAgent handshake is in flight on some other thread, which runs + /// with mConnectionsMutex released. See connect() for why same-peer callers wait on it. + std::set mLoadingAgents; + std::condition_variable mLoadingAgentsCv; /// Connection info for dynamically discovered agents that are not listed in mCommState. std::map mRemoteConnectionInfo; std::mutex mRemoteConnectionInfoMutex; 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 adfd6d1deffe..20b8e8a18903 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -667,8 +667,6 @@ ConnectionInfoType NixlTransferAgent::getLocalConnectionInfo() void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoType const& connectionInfo) { - std::unique_lock lock(mLock); - TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::loadRemoteAgent called after shutdown"); auto const separator = connectionInfo.rfind(':'); TLLM_CHECK_WITH_INFO(separator != std::string::npos, "Invalid NIXL connection info, expected 'ip:port' or '[ipv6]:port': %s", connectionInfo.c_str()); @@ -685,22 +683,58 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoT nixl_opt_args_t md_extra_params; md_extra_params.ipAddr = ip; md_extra_params.port = std::stoi(port); - auto status = mRawAgent->fetchRemoteMD(name, &md_extra_params); - TLLM_CHECK_WITH_INFO( - status == NIXL_SUCCESS, "fetchRemoteMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); - // status = mRawAgent->sendLocalMD(&md_extra_params); - // TLLM_CHECK_WITH_INFO( - // status == NIXL_SUCCESS, "sendLocalMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); - status = NIXL_ERR_NOT_FOUND; + // The metadata wait below is unbounded network I/O. Holding mLock across it would block every + // other entry point -- exclusively for the writers (registerMemory, invalidateRemoteAgent, + // shutdown), and, because glibc's shared_mutex prefers writers, transitively for readers queued + // behind them. So take the lock per NIXL call instead of across the whole wait: writers now + // queue behind a single call rather than the entire handshake. + // + // Scoping it this way, rather than observing mRawAgent through a weak_ptr, is what preserves the + // second duty the original exclusive lock was carrying: shutdown() takes mLock exclusively + // *specifically* to drain in-flight callers before mRawAgent.reset(). A weak_ptr handle would let + // the reset land while this thread still held a strong ref, migrating ~nixlAgent (documented to + // release UCX and the progress thread synchronously) onto the waiter. Holding the shared_lock + // across the dereference keeps that release on the teardown thread where it belongs. + auto withAgent = [this, &name](auto&& fn) + { + std::shared_lock lock(mLock); + // Same invariant every other entry point relies on: shutdown() resets mRawAgent under the + // exclusive lock only after setting mShutdown, so !mShutdown here implies a live agent. + TLLM_CHECK_WITH_INFO(!mShutdown.load(), + "NixlTransferAgent shut down while awaiting metadata of remote agent '%s'", name.c_str()); + return fn(*mRawAgent); + }; + + auto fetchRemoteMD = [&] + { + auto const status = withAgent([&](nixlAgent& agent) { return agent.fetchRemoteMD(name, &md_extra_params); }); + TLLM_CHECK_WITH_INFO( + status == NIXL_SUCCESS, "fetchRemoteMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); + }; + fetchRemoteMD(); + + nixl_status_t status{NIXL_ERR_NOT_FOUND}; nixl_xfer_dlist_t descs{DRAM_SEG}; + // fetchRemoteMD is fire-and-forget with no retransmit, so if that one request is lost -- or + // reaches a listener that is not serving yet -- no reply ever arrives and polling alone waits + // forever. Re-ask periodically so the handshake recovers on its own. + auto constexpr kRefetchInterval = std::chrono::seconds(5); + auto lastFetch = std::chrono::steady_clock::now(); while (status == NIXL_ERR_NOT_FOUND) { - status = mRawAgent->checkRemoteMD(name, descs); + status = withAgent([&](nixlAgent& agent) { return agent.checkRemoteMD(name, descs); }); TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS || status == NIXL_ERR_NOT_FOUND, "checkRemoteMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); if (status == NIXL_ERR_NOT_FOUND) { + if (auto const now = std::chrono::steady_clock::now(); now - lastFetch >= kRefetchInterval) + { + lastFetch = now; + TLLM_LOG_WARNING(mRank, "Still waiting for NIXL metadata of remote agent '%s' at %s; re-requesting.", + name.c_str(), connectionInfo.c_str()); + fetchRemoteMD(); + } std::this_thread::sleep_for(std::chrono::milliseconds(1)); } }