[https://nvbugs/6601575][fix] Cherry-picked the already-reviewed 0b95818fa2: take mLock per NIXL call via a…

cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp

@@ -748,12 +748,59 @@ std::vector<uint8_t> 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<std::mutex>& lock, std::set<std::string>& 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<std::mutex>& mLock;
+    std::set<std::string>& mLoadingAgents;
+    std::condition_variable& mCv;
+    std::string const& mName;
+};
+
+} // namespace
+
 AgentConnection* AgentConnectionManager::connect(std::string const& remoteAgentName, std::string const& connectionInfo,
     std::optional<std::string> 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);
         }
     }

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 <condition_variable>
 #include <map>
+#include <set>
 
 namespace tensorrt_llm::executor::kv_cache
 {
@@ -333,6 +335,10 @@ class AgentConnectionManager : public ConnectionManager
 private:
     std::map<std::string, std::shared_ptr<AgentConnection>> 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<std::string> mLoadingAgents;
+    std::condition_variable mLoadingAgentsCv;
     /// Connection info for dynamically discovered agents that are not listed in mCommState.
     std::map<std::string, std::string> mRemoteConnectionInfo;
     std::mutex mRemoteConnectionInfoMutex;

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<std::shared_mutex> 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<std::shared_mutex> 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));
         }
     }