fix: Deferred intra-node UCX exchange handshakes

velox/experimental/ucx-exchange/Acceptor.cpp

@@ -23,6 +23,67 @@
 
 namespace facebook::velox::ucx_exchange {
 
+namespace {
+
+void completeHandshake(
+    const std::shared_ptr<Communicator>& communicator,
+    const std::shared_ptr<EndpointRef>& epRef,
+    const PartitionKey& key,
+    bool isIntraNodeTransfer) {
+  // Eligibility may be resolved after the active-message callback returns.
+  // Do not issue UCXX operations if the endpoint closed in the meantime.
+  if (!epRef->endpoint_ || !epRef->endpoint_->isAlive()) {
+    VLOG(2) << "[ACCEPTOR] Dropping resolved handshake for " << key.toString()
+            << " because its endpoint is closed";
+    return;
+  }
+
+  const std::string peerIp = epRef->getPeerIp();
+  auto exchangeServer =
+      UcxExchangeServer::create(communicator, epRef, key, isIntraNodeTransfer);
+
+  // Add this exchangeServer to the endpoint reference.
+  epRef->addCommElem(exchangeServer);
+
+  // Register exchangeServer with communicator.
+  communicator->registerCommElement(exchangeServer);
+  VLOG(2) << "[ACCEPTOR] new server: " << exchangeServer->toString()
+          << " peerIp=" << peerIp
+          << " isIntraNodeTransfer=" << isIntraNodeTransfer;
+
+  // Send HandshakeResponse back to the source to inform about intra-node
+  // transfer. This allows the source to bypass UCXX for all subsequent data
+  // transfers.
+  auto response = std::make_shared<HandshakeResponse>();
+  response->isIntraNodeTransfer = exchangeServer->isIntraNodeTransfer();
+
+  const uint32_t keyHash = fnv1a_32(key.toString());
+  const uint64_t responseTag = getHandshakeResponseTag(keyHash);
+
+  VLOG(3) << "Sending HandshakeResponse to " << key.toString()
+          << " isIntraNodeTransfer=" << response->isIntraNodeTransfer
+          << " tag=" << std::hex << responseTag;
+
+  // Fire-and-forget: we don't need to track this request completion.
+  epRef->endpoint_->tagSend(
+      response.get(),
+      sizeof(*response),
+      ucxx::Tag{responseTag},
+      false,
+      [response, keyStr = key.toString()](
+          ucs_status_t status, std::shared_ptr<void> /*arg*/) {
+        if (status == UCS_OK) {
+          VLOG(3) << "HandshakeResponse sent successfully to " << keyStr;
+        } else {
+          VLOG(0) << "Failed to send HandshakeResponse to " << keyStr << ": "
+                  << ucs_status_string(status);
+        }
+      },
+      response);
+}
+
+} // namespace
+
 /*static*/
 void Acceptor::cStyleAMCallback(
     std::shared_ptr<ucxx::Request> request,
@@ -63,62 +124,36 @@ void Acceptor::cStyleAMCallback(
       cudf_velox::CudfConfig::getInstance().intraNodeExchange &&
       (handshakePtr->workerId == communicator->getWorkerId());
 
-  // Disable intra-node when the task is not yet initialized (placeholder
-  // queue from sinks connecting before initializeTask) or when the task
-  // uses broadcast mode (all destination servers share the same
-  // packed_columns — the intra-node source's destructive move would
-  // corrupt it for other servers).
-  if (isIntraNodeTransfer) {
-    if (!UcxOutputQueueManager::getInstanceRef()->canUseIntraNode(key.taskId)) {
-      VLOG(2) << "[ACCEPTOR] Disabling intra-node for task " << key.taskId
-              << " (not initialized or broadcast)";
-      isIntraNodeTransfer = false;
-    }
+  if (!isIntraNodeTransfer) {
+    completeHandshake(communicator, epRef, key, false);
+    return;
   }
 
-  std::string peerIp = epRef->getPeerIp();
-
-  auto exchangeServer =
-      UcxExchangeServer::create(communicator, epRef, key, isIntraNodeTransfer);
-
-  // Add this exchangeServer to the endpoint reference.
-  epRef->addCommElem(exchangeServer);
-
-  // Register exchangeServer with communicator.
-  communicator->registerCommElement(exchangeServer);
-  VLOG(2) << "[ACCEPTOR] new server: " << exchangeServer->toString()
-          << " peerIp=" << peerIp
-          << " isIntraNodeTransfer=" << isIntraNodeTransfer;
-
-  // Send HandshakeResponse back to the source to inform about intra-node
-  // transfer. This allows the source to bypass UCXX for all subsequent data
-  // transfers.
-  auto response = std::make_shared<HandshakeResponse>();
-  response->isIntraNodeTransfer = exchangeServer->isIntraNodeTransfer();
-
-  uint32_t keyHash = fnv1a_32(key.toString());
-  uint64_t responseTag = getHandshakeResponseTag(keyHash);
-
-  VLOG(3) << "Sending HandshakeResponse to " << key.toString()
-          << " isIntraNodeTransfer=" << response->isIntraNodeTransfer
-          << " tag=" << std::hex << responseTag;
-
-  // Fire-and-forget: we don't need to track this request completion
-  epRef->endpoint_->tagSend(
-      response.get(),
-      sizeof(*response),
-      ucxx::Tag{responseTag},
-      false,
-      [response, keyStr = key.toString()](
-          ucs_status_t status, std::shared_ptr<void> arg) {
-        if (status == UCS_OK) {
-          VLOG(3) << "HandshakeResponse sent successfully to " << keyStr;
-        } else {
-          VLOG(0) << "Failed to send HandshakeResponse to " << keyStr << ": "
-                  << ucs_status_string(status);
+  // The source is in this process, but its producer task may not have reached
+  // initializeTask() yet. Do not permanently route it through remote UCX based
+  // on a placeholder queue. Wait until the real output kind is known, then
+  // finish the UCXX handshake on the Communicator thread. Broadcast remains on
+  // UCX because its packed_columns object is shared by multiple destinations.
+  std::weak_ptr<Communicator> weakCommunicator = communicator;
+  std::weak_ptr<EndpointRef> weakEndpoint = epRef;
+  UcxOutputQueueManager::getInstanceRef()->notifyOnIntraNodeEligibility(
+      key.taskId, [weakCommunicator, weakEndpoint, key](bool canUseIntraNode) {
+        auto communicator = weakCommunicator.lock();
+        if (!communicator) {
+          return;
         }
-      },
-      response);
+        VLOG(2) << "[ACCEPTOR] Resolved deferred handshake for "
+                << key.toString() << " isIntraNodeTransfer=" << canUseIntraNode;
+        communicator->deferAction(
+            [weakCommunicator, weakEndpoint, key, canUseIntraNode]() {
+              auto communicator = weakCommunicator.lock();
+              auto epRef = weakEndpoint.lock();
+              if (!communicator || !epRef) {
+                return;
+              }
+              completeHandshake(communicator, epRef, key, canUseIntraNode);
+            });
+      });
 }
 
 // Add endpoint reference to ucp_cp -> epRef map.

velox/experimental/ucx-exchange/Communicator.cpp

@@ -146,6 +146,12 @@ void Communicator::run() {
 
   VLOG(3) << "Communicator running.";
   const bool blockingMode = CudfConfig::getInstance().ucxxBlockingProgress;
+  // Intra-node sources poll the transfer registry by re-enqueuing themselves.
+  // Therefore workQueue_ can remain non-empty indefinitely. Process it in
+  // bounded batches so no local queue can starve the others or heartbeat work.
+  constexpr size_t kMaxDeferredEndpointCleanupsPerLoop = 256;
+  constexpr size_t kMaxDeferredActionsPerLoop = 256;
+  constexpr size_t kMaxWorkItemsPerLoop = 256;
   while (running_) {
     try {
       // Periodic heartbeat for diagnostic logging.
@@ -185,6 +191,7 @@ void Communicator::run() {
                 << ")"
                 << " endpoints=" << numEndpoints
                 << " deferredCleanup=" << deferredEndpointCleanup_.size()
+                << " deferredActions=" << deferredActions_.size()
                 << " deferredRequests=" << deferredRequests_.size()
                 << " workItemsProcessed=" << workItemsProcessed_
                 << " GPU=" << gpuUsedMB << "/" << gpuTotalMB << "MB"
@@ -197,7 +204,13 @@ void Communicator::run() {
       // UCX callbacks cannot call closeBlocking() (which progresses the
       // worker) or iterate communicators_, so they defer cleanup to
       // this main loop via deferEndpointCleanup().
-      while (auto ep = deferredEndpointCleanup_.pop()) {
+      for (size_t endpointCleanupsThisLoop = 0;
+           endpointCleanupsThisLoop < kMaxDeferredEndpointCleanupsPerLoop;
+           ++endpointCleanupsThisLoop) {
+        auto ep = deferredEndpointCleanup_.pop();
+        if (!ep) {
+          break;
+        }
         VLOG(3) << "Processing deferred endpoint cleanup";
         // First, close all communicators associated with this endpoint.
         // This must happen before removeEndpointRef() which may destroy
@@ -206,9 +219,29 @@ void Communicator::run() {
         removeEndpointRef(ep);
       }
 
-      // Process the work queue. Make sure that communication is progressed
-      // after each call to a comms element, otherwise we will deadlock.
-      while (auto comms = workQueue_.pop()) {
+      // Run actions submitted by other threads. These must execute here
+      // because they may issue UCXX operations on this worker.
+      for (size_t deferredActionsThisLoop = 0;
+           deferredActionsThisLoop < kMaxDeferredActionsPerLoop;
+           ++deferredActionsThisLoop) {
+        auto deferredAction = deferredActions_.pop();
+        if (!deferredAction) {
+          break;
+        }
+        deferredAction->action();
+        worker_->progress();
+      }
+
+      // Process a bounded batch from the work queue. Make sure that
+      // communication is progressed after each call to a comms element,
+      // otherwise we will deadlock.
+      for (size_t workItemsThisLoop = 0;
+           workItemsThisLoop < kMaxWorkItemsPerLoop;
+           ++workItemsThisLoop) {
+        auto comms = workQueue_.pop();
+        if (!comms) {
+          break;
+        }
         comms->process();
         ++workItemsProcessed_;
         // Progress after each work item to allow UCXX to advance
@@ -230,11 +263,12 @@ void Communicator::run() {
             deferredRequests_.end());
       }
 
-      // All queues are drained. Now wait for UCXX network events.
-      // In blocking mode, this will block until a UCXX event arrives
-      // or worker_->signal() is called (from addToWorkQueue,
-      // deferEndpointCleanup, or stop).
-      if (blockingMode) {
+      // Wait for UCXX network events only if there is no process-local work.
+      // A bounded batch may intentionally leave workQueue_ non-empty; blocking
+      // in that case would strand it until an unrelated network event arrives.
+      const bool hasLocalWork = !deferredEndpointCleanup_.empty() ||
+          !deferredActions_.empty() || !workQueue_.empty();
+      if (blockingMode && !hasLocalWork) {
         worker_->progressWorkerEvent(0);
       } else {
         worker_->progress();
@@ -282,6 +316,14 @@ void Communicator::addToWorkQueue(std::shared_ptr<CommElement> comms) {
   signalWorker();
 }
 
+void Communicator::deferAction(std::function<void()> action) {
+  if (!action) {
+    return;
+  }
+  deferredActions_.push(std::make_shared<DeferredAction>(std::move(action)));
+  signalWorker();
+}
+
 void Communicator::unregister(std::shared_ptr<CommElement> comms) {
   std::lock_guard<std::mutex> lock(elemMutex_);
   if (!comms) {

velox/experimental/ucx-exchange/Communicator.h

@@ -18,6 +18,7 @@
 #include <ucxx/api.h>
 #include <chrono>
 #include <cstdint>
+#include <functional>
 #include <random>
 #include <string>
 #include <string_view>
@@ -85,6 +86,14 @@ class Communicator {
   /// @param comms The element to be added to the work queue.
   void addToWorkQueue(std::shared_ptr<CommElement> comms);
 
+  /// @brief Schedules an action to run on the Communicator thread.
+  ///
+  /// UCXX endpoint operations must run on the thread that progresses the UCX
+  /// worker. This is used when another thread resolves state needed by a
+  /// communication operation, such as output-queue initialization completing
+  /// after an exchange handshake arrives.
+  void deferAction(std::function<void()> action);
+
   /// @brief Unregisters a communication element
   /// @brief comms The communication element.
   void unregister(std::shared_ptr<CommElement> comms);
@@ -198,6 +207,17 @@ class Communicator {
   // the communicator thread.
   WorkQueue<CommElement> workQueue_;
 
+  struct DeferredAction {
+    explicit DeferredAction(std::function<void()> action)
+        : action(std::move(action)) {}
+
+    std::function<void()> action;
+  };
+
+  // Actions submitted from non-Communicator threads for execution by the main
+  // UCX progress loop.
+  WorkQueue<DeferredAction> deferredActions_;
+
   // Protects endpoints_. Must be recursive because removeEndpointRef() holds
   // the lock and calls closeBlocking(), which progresses the worker and can
   // fire listenerCallback() re-entrantly — that callback also locks this mutex.

velox/experimental/ucx-exchange/UcxExchangeClient.cpp

@@ -64,6 +64,7 @@ void UcxExchangeClient::close() {
     if (closed_) {
       return;
     }
+    stats_ = collectStatsLocked();
     closed_ = true;
     sources = std::move(sources_);
   }
@@ -76,8 +77,41 @@ void UcxExchangeClient::close() {
 }
 
 folly::F14FastMap<std::string, RuntimeMetric> UcxExchangeClient::stats() const {
-  // TODO: Implement stats collection.
+  std::lock_guard<std::mutex> l(queue_->mutex());
+  if (closed_) {
+    return stats_;
+  }
+  return collectStatsLocked();
+}
+
+folly::F14FastMap<std::string, RuntimeMetric>
+UcxExchangeClient::collectStatsLocked() const {
   folly::F14FastMap<std::string, RuntimeMetric> stats;
+
+  for (const auto& source : sources_) {
+    if (source->supportsMetrics()) {
+      for (const auto& [name, value] : source->metrics()) {
+        auto [iter, inserted] = stats.try_emplace(name, value.unit);
+        iter->second.merge(value);
+      }
+    } else {
+      for (const auto& [name, value] : source->stats()) {
+        auto [iter, inserted] = stats.try_emplace(name);
+        iter->second.addValue(value);
+      }
+    }
+  }
+
+  stats.insert_or_assign(
+      "peakBytes",
+      RuntimeMetric(queue_->peakBytes(), RuntimeCounter::Unit::kBytes));
+  stats.insert_or_assign(
+      "numReceivedTables", RuntimeMetric(queue_->receivedTables()));
+  stats.insert_or_assign(
+      "averageReceivedTableBytes",
+      RuntimeMetric(
+          queue_->averageReceivedTablesBytes(), RuntimeCounter::Unit::kBytes));
+
   return stats;
 }
 

velox/experimental/ucx-exchange/UcxExchangeClient.h

@@ -88,6 +88,8 @@ class UcxExchangeClient
   }
 
  private:
+  folly::F14FastMap<std::string, RuntimeMetric> collectStatsLocked() const;
+
   // Handy for ad-hoc logging.
   const std::string taskId_;
   const int destination_;
@@ -99,6 +101,7 @@ class UcxExchangeClient
   std::unordered_set<std::string> remoteTaskIds_;
   std::vector<std::shared_ptr<UcxExchangeSource>> sources_;
   bool closed_{false};
+  mutable folly::F14FastMap<std::string, RuntimeMetric> stats_;
 
   // Total number of packed_clumns in flight.
   int64_t totalPendingColumns_{0};

velox/experimental/ucx-exchange/UcxExchangeServer.cpp

@@ -126,7 +126,7 @@ void UcxExchangeServer::process() {
       // Use weak_ptr to prevent use-after-free if close() is called during
       // callback
       std::weak_ptr<UcxExchangeServer> weakQueue = weak_from_this();
-      queueMgr_->getData(
+      outputQueue_ = queueMgr_->getData(
           partitionKey_.taskId,
           partitionKey_.destination,
           [weakQueue](
@@ -305,7 +305,8 @@ void UcxExchangeServer::sendData() {
               key, nullptr, /*numRows=*/0, /*atEnd=*/true);
       intraNodeAtEndPublished_ = true;
 
-      queueMgr_->deleteResults(partitionKey_.taskId, partitionKey_.destination);
+      queueMgr_->deleteResultsForQueue(outputQueue_, partitionKey_.destination);
+      outputQueue_.reset();
 
       // Wait for source to acknowledge atEnd before finishing
       setState(ServerState::WaitingForIntraNodeRetrieve);
@@ -442,7 +443,8 @@ void UcxExchangeServer::sendData() {
       VLOG(3) << "@" << partitionKey_.taskId
               << " Finished transferring partition for task "
               << partitionKey_.toString();
-      queueMgr_->deleteResults(partitionKey_.taskId, partitionKey_.destination);
+      queueMgr_->deleteResultsForQueue(outputQueue_, partitionKey_.destination);
+      outputQueue_.reset();
       setState(ServerState::Done);
       communicator_->addToWorkQueue(getSelfPtr());
     }