From cef45766e75cf2ec16e0291fae5861c9c251985d Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 4 Sep 2026 10:11:04 -0700
Subject: [PATCH 01/15] feat(ucx-exchange): Add transport payload metrics

---
 .../ucx-exchange/UcxExchangeClient.cpp        | 36 ++++++++++++++++++-
 .../ucx-exchange/UcxExchangeClient.h          |  3 ++
 .../ucx-exchange/UcxExchangeSource.cpp        | 29 ++++++++++++---
 .../ucx-exchange/UcxExchangeSource.h          | 13 +++++++
 4 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp b/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
index bf5cb2c6c05..393bde807c9 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
+++ b/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;
 }
 
diff --git a/velox/experimental/ucx-exchange/UcxExchangeClient.h b/velox/experimental/ucx-exchange/UcxExchangeClient.h
index f43e31f0533..5ccfeb6b528 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeClient.h
+++ b/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};
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
index 1235fae64e1..236ff60524c 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
@@ -266,16 +266,36 @@ folly::F14FastMap<std::string, int64_t> UcxExchangeSource::stats() const {
 
 folly::F14FastMap<std::string, RuntimeMetric> UcxExchangeSource::metrics()
     const {
+  std::lock_guard<std::mutex> lock(metricsMutex_);
   folly::F14FastMap<std::string, RuntimeMetric> map;
 
   // these metrics will be aggregated over all exchange sources of the same
   // exchange client.
   map["ucxExchangeSource.numPackedColumns"] = metrics_.numPackedColumns_;
   map["ucxExchangeSource.totalBytes"] = metrics_.totalBytes_;
+  map["ucxExchangeSource.intraNodePackedColumns"] =
+      metrics_.intraNodePackedColumns_;
+  map["ucxExchangeSource.intraNodeBytes"] = metrics_.intraNodeBytes_;
+  map["ucxExchangeSource.remotePackedColumns"] = metrics_.remotePackedColumns_;
+  map["ucxExchangeSource.remoteBytes"] = metrics_.remoteBytes_;
   map["ucxExchangeSource.rttPerRequest"] = metrics_.rttPerRequest_;
   return map;
 }
 
+void UcxExchangeSource::recordPayloadMetrics(uint64_t dataBytes) {
+  const auto bytes = saturateCast(dataBytes);
+  std::lock_guard<std::mutex> lock(metricsMutex_);
+  metrics_.numPackedColumns_.addValue(1);
+  metrics_.totalBytes_.addValue(bytes);
+  if (isIntraNodeTransfer_) {
+    metrics_.intraNodePackedColumns_.addValue(1);
+    metrics_.intraNodeBytes_.addValue(bytes);
+  } else {
+    metrics_.remotePackedColumns_.addValue(1);
+    metrics_.remoteBytes_.addValue(bytes);
+  }
+}
+
 // private methods ---
 PartitionKey UcxExchangeSource::extractTaskAndDestinationId(
     std::string_view path) {
@@ -602,8 +622,7 @@ void UcxExchangeSource::onData(ucs_status_t status, std::shared_ptr<void> arg) {
     std::shared_ptr<DataAndMetadata> ptr =
         std::static_pointer_cast<DataAndMetadata>(arg);
 
-    metrics_.numPackedColumns_.addValue(1);
-    metrics_.totalBytes_.addValue(ptr->metadata.dataSizeBytes);
+    recordPayloadMetrics(ptr->metadata.dataSizeBytes);
 
     // Create packed_columns from the received metadata and data buffer
     cudf::packed_columns packedCols(
@@ -765,12 +784,12 @@ void UcxExchangeSource::onIntraNodeData(
     return;
   }
 
+  const auto dataSize = data->gpu_data->size();
   VLOG(3) << toString()
           << " Intra-node transfer: received data for seq=" << sequenceNumber_
-          << " size=" << data->gpu_data->size();
+          << " size=" << dataSize;
 
-  metrics_.numPackedColumns_.addValue(1);
-  metrics_.totalBytes_.addValue(data->gpu_data->size());
+  recordPayloadMetrics(dataSize);
 
   // Convert packed_columns to PackedTableWithStream for the queue.
   // Create packed_columns from the shared data.
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index d22513965a8..138e86f1da0 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -36,15 +36,25 @@
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
+#include <mutex>
+
 namespace facebook::velox::ucx_exchange {
 
 struct UcxExchangeMetrics {
   UcxExchangeMetrics()
       : numPackedColumns_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
         totalBytes_(RuntimeCounter::Unit::kBytes),
+        intraNodePackedColumns_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
+        intraNodeBytes_(RuntimeCounter::Unit::kBytes),
+        remotePackedColumns_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
+        remoteBytes_(RuntimeCounter::Unit::kBytes),
         rttPerRequest_(RuntimeMetric(RuntimeCounter::Unit::kNanos)) {}
   RuntimeMetric numPackedColumns_; // total number of packed columns received.
   RuntimeMetric totalBytes_; // total number of bytes received
+  RuntimeMetric intraNodePackedColumns_;
+  RuntimeMetric intraNodeBytes_;
+  RuntimeMetric remotePackedColumns_;
+  RuntimeMetric remoteBytes_;
   RuntimeMetric rttPerRequest_;
 };
 
@@ -246,6 +256,8 @@ class UcxExchangeSource
   /// @return Returns true if state was changed, false otherwise.
   bool setStateIf(ReceiverState expected, ReceiverState desired);
 
+  void recordPayloadMetrics(uint64_t dataBytes);
+
   // The connection parameters
   const std::string host_;
   uint16_t port_;
@@ -289,6 +301,7 @@ class UcxExchangeSource
   std::atomic<bool> backpressureActive_{false};
 
   // Some metrics/counters:
+  mutable std::mutex metricsMutex_;
   UcxExchangeMetrics metrics_;
 
   // The outstanding request - there can only be one outstanding request

From 97c24cd72f6726ca116809f9e41693752d8a24b5 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 4 Sep 2026 10:13:32 -0700
Subject: [PATCH 02/15] test(ucx-exchange): Cover partitioned placeholder
 routing

---
 .../ucx-exchange/tests/SinkDriverMock.h       |  4 +
 .../ucx-exchange/tests/UcxExchangeTest.cpp    | 87 +++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/velox/experimental/ucx-exchange/tests/SinkDriverMock.h b/velox/experimental/ucx-exchange/tests/SinkDriverMock.h
index f255663c3f7..143f7d92f73 100644
--- a/velox/experimental/ucx-exchange/tests/SinkDriverMock.h
+++ b/velox/experimental/ucx-exchange/tests/SinkDriverMock.h
@@ -75,6 +75,10 @@ class SinkDriverMock {
     return numChunksReceived_.load();
   }
 
+  folly::F14FastMap<std::string, RuntimeMetric> exchangeClientStats() const {
+    return exchangeClient_->stats();
+  }
+
   bool dataIsValid() {
     return dataValidFlag_;
   }
diff --git a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
index d66e37730bc..108e3159770 100644
--- a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
@@ -146,6 +146,9 @@ struct ExchangeTestParamsPrinter {
 
 class UcxExchangeTest : public testing::TestWithParam<ExchangeTestParams> {
  protected:
+  // Tests that mutate CudfConfig::intraNodeExchange must use process-level
+  // parallelism only: it is mutable process-global state read by the
+  // communicator's Acceptor thread.
   // Chosen per process in SetUpTestCase() rather than hardcoded. The
   // communicator opens a listener on this port without address reuse, so two
   // runs of this binary in quick succession fail the second with
@@ -1083,6 +1086,90 @@ TEST_P(UcxExchangeTest, intraNodeTaskRemovalLivelock) {
   // If we get here, the source correctly detected the cancelled task.
 }
 
+// A same-process consumer can connect before the producer task has initialized
+// its output queue. The handshake must wait for the real output kind instead of
+// permanently falling back to remote UCX based on the placeholder queue.
+TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
+  // This test doesn't use parameters — run only for the first param set.
+  {
+    ExchangeTestParams p = GetParam();
+    if (p.numSrcDrivers != 1 || p.numDstDrivers != 1 || p.numPartitions != 1 ||
+        p.numChunks != 100 || p.numUpstreamTasks != 1 ||
+        p.tableType != TableType::NARROW) {
+      GTEST_SKIP() << "partitionedPlaceholderUsesIntraNodePath: runs only once";
+    }
+  }
+
+  auto& config = cudf_velox::CudfConfig::getInstance();
+  const bool origIntraNode = config.intraNodeExchange;
+  config.intraNodeExchange = true;
+  SCOPE_EXIT {
+    config.intraNodeExchange = origIntraNode;
+  };
+
+  auto metricSum = [](
+                       const folly::F14FastMap<std::string, RuntimeMetric>&
+                           stats,
+                       std::string_view name) {
+    auto it = stats.find(std::string{name});
+    return it == stats.end() ? 0 : it->second.sum;
+  };
+
+  const std::string taskPrefix = getUniqueTaskPrefix();
+  const std::string srcTaskId = taskPrefix + "partitionedPlaceholderSrc";
+  const std::string sinkTaskId = taskPrefix + "partitionedPlaceholderSink";
+  constexpr int numChunks = 2;
+  constexpr int numRowsPerChunk = 1000;
+
+  core::PlanNodeId exchangeNodeId;
+  auto sinkTask = createExchangeTask(
+      sinkTaskId, UcxTestData::kTestRowType, /*partitionId=*/0, exchangeNodeId);
+  auto sinkDriver =
+      std::make_shared<SinkDriverMock>(sinkTask, /*numDrivers=*/1);
+  std::vector<exec::Split> splits;
+  splits.emplace_back(remoteSplit(srcTaskId, /*partitionId=*/0));
+  sinkDriver->addSplits(splits);
+
+  // Start the consumer first and give its handshake time to reach the server.
+  // At this point the producer is represented only by an uninitialized task ID.
+  sinkDriver->run();
+  std::this_thread::sleep_for(std::chrono::seconds(1));
+  EXPECT_FALSE(queueManager_->canUseIntraNode(srcTaskId));
+
+  auto srcTask = createSourceTask(srcTaskId, pool_, UcxTestData::kTestRowType);
+  queueManager_->initializeTask(
+      srcTask,
+      core::PartitionedOutputNode::Kind::kPartitioned,
+      /*numDestinations=*/1,
+      /*numDrivers=*/1);
+  EXPECT_TRUE(queueManager_->canUseIntraNode(srcTaskId));
+
+  auto sourceMock = std::make_shared<UcxPartitionedOutputMock>(
+      srcTaskId,
+      /*numDrivers=*/1,
+      /*numPartitions=*/1,
+      numChunks,
+      numRowsPerChunk);
+  sourceMock->run();
+
+  sourceMock->joinThreads();
+  sinkDriver->joinThreads();
+
+  EXPECT_EQ(
+      sinkDriver->numRows(),
+      static_cast<uint64_t>(numChunks * numRowsPerChunk));
+
+  const auto sinkStats = sinkDriver->exchangeClientStats();
+  EXPECT_EQ(
+      metricSum(sinkStats, "ucxExchangeSource.intraNodePackedColumns"),
+      numChunks);
+  EXPECT_GT(metricSum(sinkStats, "ucxExchangeSource.intraNodeBytes"), 0);
+  EXPECT_EQ(metricSum(sinkStats, "ucxExchangeSource.remotePackedColumns"), 0);
+  EXPECT_EQ(metricSum(sinkStats, "ucxExchangeSource.remoteBytes"), 0);
+
+  queueManager_->removeTask(srcTaskId);
+}
+
 // Regression test for broadcast + intra-node SIGSEGV.
 // Before the fix in Acceptor.cpp, broadcast tasks using intra-node transfer
 // would crash because the intra-node source destructively moves gpu_data from

From f2a40c02e79bc6cab2c1e49767d4ed47f61a150f Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 8 Sep 2026 16:46:23 -0700
Subject: [PATCH 03/15] fix(ucx-exchange): Defer intra-node handshakes until
 task initialization

---
 velox/experimental/ucx-exchange/Acceptor.cpp  | 160 ++++++++++++------
 velox/experimental/ucx-exchange/Acceptor.h    |   9 +
 .../ucx-exchange/Communicator.cpp             |  44 ++++-
 .../experimental/ucx-exchange/Communicator.h  |  20 +++
 .../ucx-exchange/UcxOutputQueueManager.cpp    |  36 ++++
 .../ucx-exchange/UcxOutputQueueManager.h      |   7 +
 velox/experimental/ucx-exchange/UcxQueues.cpp |  75 ++++++--
 velox/experimental/ucx-exchange/UcxQueues.h   |  17 ++
 .../tests/UcxOutputQueueManagerTest.cpp       |  64 +++++++
 9 files changed, 354 insertions(+), 78 deletions(-)

diff --git a/velox/experimental/ucx-exchange/Acceptor.cpp b/velox/experimental/ucx-exchange/Acceptor.cpp
index 36fc57d40da..edf930de454 100644
--- a/velox/experimental/ucx-exchange/Acceptor.cpp
+++ b/velox/experimental/ucx-exchange/Acceptor.cpp
@@ -21,8 +21,88 @@
 #include "velox/experimental/ucx-exchange/UcxExchangeServer.h"
 #include "velox/experimental/ucx-exchange/UcxOutputQueueManager.h"
 
+#include <atomic>
+
 namespace facebook::velox::ucx_exchange {
 
+namespace {
+
+std::atomic<uint64_t> gNumIntraNodeHandshakes{0};
+std::atomic<uint64_t> gNumRemoteHandshakes{0};
+
+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);
+
+  if (isIntraNodeTransfer) {
+    gNumIntraNodeHandshakes.fetch_add(1, std::memory_order_relaxed);
+  } else {
+    gNumRemoteHandshakes.fetch_add(1, std::memory_order_relaxed);
+  }
+
+  // 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
+
+uint64_t Acceptor::numIntraNodeHandshakes() {
+  return gNumIntraNodeHandshakes.load(std::memory_order_relaxed);
+}
+
+uint64_t Acceptor::numRemoteHandshakes() {
+  return gNumRemoteHandshakes.load(std::memory_order_relaxed);
+}
+
 /*static*/
 void Acceptor::cStyleAMCallback(
     std::shared_ptr<ucxx::Request> request,
@@ -63,62 +143,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.
diff --git a/velox/experimental/ucx-exchange/Acceptor.h b/velox/experimental/ucx-exchange/Acceptor.h
index 063532b836c..a44f35ef806 100644
--- a/velox/experimental/ucx-exchange/Acceptor.h
+++ b/velox/experimental/ucx-exchange/Acceptor.h
@@ -16,6 +16,7 @@
 #pragma once
 
 #include <ucxx/api.h>
+#include <cstdint>
 #include <map>
 #include "velox/experimental/ucx-exchange/EndpointRef.h"
 
@@ -31,6 +32,14 @@ struct Acceptor {
       std::shared_ptr<ucxx::Request> request,
       ucp_ep_h ep);
 
+  /// Number of completed handshakes routed through the process-local path.
+  /// Exposed for diagnostics and regression tests.
+  static uint64_t numIntraNodeHandshakes();
+
+  /// Number of completed handshakes routed through the remote UCX path.
+  /// Exposed for diagnostics and regression tests.
+  static uint64_t numRemoteHandshakes();
+
   /// @brief Adds the endpoint reference to the handleToEndpointRef_ map such
   /// that endpoint handles can be resolved
   void registerEndpointRef(std::shared_ptr<EndpointRef> endpointRef);
diff --git a/velox/experimental/ucx-exchange/Communicator.cpp b/velox/experimental/ucx-exchange/Communicator.cpp
index 0b385407da7..d11d4f37239 100644
--- a/velox/experimental/ucx-exchange/Communicator.cpp
+++ b/velox/experimental/ucx-exchange/Communicator.cpp
@@ -146,6 +146,11 @@ 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 deferred handshakes, endpoint cleanup, and heartbeat
+  // work at the top of this loop cannot be starved.
+  constexpr size_t kMaxWorkItemsPerLoop = 256;
   while (running_) {
     try {
       // Periodic heartbeat for diagnostic logging.
@@ -185,6 +190,7 @@ void Communicator::run() {
                 << ")"
                 << " endpoints=" << numEndpoints
                 << " deferredCleanup=" << deferredEndpointCleanup_.size()
+                << " deferredActions=" << deferredActions_.size()
                 << " deferredRequests=" << deferredRequests_.size()
                 << " workItemsProcessed=" << workItemsProcessed_
                 << " GPU=" << gpuUsedMB << "/" << gpuTotalMB << "MB"
@@ -206,11 +212,24 @@ 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.
+      while (auto deferredAction = deferredActions_.pop()) {
+        deferredAction->action();
+      }
+
+      // 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.
+      size_t workItemsThisLoop = 0;
+      while (workItemsThisLoop < kMaxWorkItemsPerLoop) {
+        auto comms = workQueue_.pop();
+        if (!comms) {
+          break;
+        }
         comms->process();
         ++workItemsProcessed_;
+        ++workItemsThisLoop;
         // Progress after each work item to allow UCXX to advance
         // its internal state (complete sends/receives, fire callbacks).
         // Use non-blocking progress here to avoid blocking between
@@ -230,11 +249,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 +302,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) {
diff --git a/velox/experimental/ucx-exchange/Communicator.h b/velox/experimental/ucx-exchange/Communicator.h
index 6c8f3ffc0ac..c039b2435d3 100644
--- a/velox/experimental/ucx-exchange/Communicator.h
+++ b/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.
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
index 05ebbab5836..bef50a1c346 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
@@ -151,6 +151,42 @@ bool UcxOutputQueueManager::canUseIntraNode(std::string_view taskId) {
       queue->kind() != core::PartitionedOutputNode::Kind::kBroadcast;
 }
 
+void UcxOutputQueueManager::notifyOnIntraNodeEligibility(
+    std::string_view taskId,
+    UcxIntraNodeEligibilityCallback callback) {
+  VELOX_CHECK(callback, "Intra-node eligibility callback must be set");
+
+  const std::string taskIdStr{taskId};
+  std::shared_ptr<UcxOutputQueue> outputQueue;
+  bool taskRemoved = false;
+  queues_.withLock([&](auto& queues) {
+    auto it = queues.find(taskIdStr);
+    if (it != queues.end()) {
+      outputQueue = it->second;
+      return;
+    }
+
+    taskRemoved = removedTasks_.withLock(
+        [&](auto& removed) { return removed.count(taskIdStr) > 0; });
+    if (!taskRemoved) {
+      // A local source can connect before the producer task is initialized.
+      // Preserve the handshake until initializeTask() publishes the real kind.
+      outputQueue = std::make_shared<UcxOutputQueue>(
+          nullptr,
+          /*numDestinations=*/0,
+          /*numDrivers=*/0,
+          core::PartitionedOutputNode::Kind::kPartitioned);
+      queues[taskIdStr] = outputQueue;
+    }
+  });
+
+  if (taskRemoved) {
+    callback(false);
+    return;
+  }
+  outputQueue->notifyOnIntraNodeEligibility(std::move(callback));
+}
+
 void UcxOutputQueueManager::removeTask(const std::string& taskId) {
   std::string taskIdStr{taskId};
   auto queue =
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
index 09f059d38ff..075619cd72f 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
@@ -115,6 +115,13 @@ class UcxOutputQueueManager : public exec::OutputBufferManager {
   /// source's destructive move would corrupt data for other servers).
   bool canUseIntraNode(std::string_view taskId);
 
+  /// Resolves whether a task can use process-local exchange after its real
+  /// output kind is known. If the task is still represented by a placeholder,
+  /// callback invocation is deferred until initializeTask() or removeTask().
+  void notifyOnIntraNodeEligibility(
+      std::string_view taskId,
+      UcxIntraNodeEligibilityCallback callback);
+
   /// @brief Removes the queue for the given task from the queue manager.
   /// Calls "terminate" on the queue to awake waiting producers.
   void removeTask(const std::string& taskId) override;
diff --git a/velox/experimental/ucx-exchange/UcxQueues.cpp b/velox/experimental/ucx-exchange/UcxQueues.cpp
index c0fb5fa65d6..d4e21b17809 100644
--- a/velox/experimental/ucx-exchange/UcxQueues.cpp
+++ b/velox/experimental/ucx-exchange/UcxQueues.cpp
@@ -166,28 +166,61 @@ bool UcxOutputQueue::initialize(
     uint32_t numDestinations,
     uint32_t numDrivers,
     core::PartitionedOutputNode::Kind kind) {
-  std::lock_guard<std::mutex> l(mutex_);
-  if (task_) {
-    // already initialized!
-    return false;
+  std::vector<UcxIntraNodeEligibilityCallback> eligibilityCallbacks;
+  {
+    std::lock_guard<std::mutex> l(mutex_);
+    if (task_) {
+      // already initialized!
+      return false;
+    }
+    kind_ = kind;
+    numDrivers_ = numDrivers;
+    task_ = task;
+    maxSize_ = task_->queryCtx()->queryConfig().maxOutputBufferSize();
+    continueSize_ = (maxSize_ * kContinuePct) / 100;
+    // Publish task metadata before destination queue expansion. Acceptor only
+    // needs task/kind to choose the intra-node path; getData() takes mutex_ and
+    // waits for any queue expansion in this function to finish.
+    initialized_.store(true, std::memory_order_release);
+    // create additional queues if there are more destinations.
+    for (int i = queues_.size(); i < numDestinations; ++i) {
+      // create the destination queues inside the vector using emplace_back.
+      queues_.emplace_back(std::make_unique<UcxDestinationQueue>());
+    }
+    eligibilityCallbacks = std::move(intraNodeEligibilityCallbacks_);
   }
-  kind_ = kind;
-  numDrivers_ = numDrivers;
-  task_ = task;
-  maxSize_ = task_->queryCtx()->queryConfig().maxOutputBufferSize();
-  continueSize_ = (maxSize_ * kContinuePct) / 100;
-  // Publish task metadata before destination queue expansion. Acceptor only
-  // needs task/kind to choose the intra-node path; getData() takes mutex_ and
-  // waits for any queue expansion in this function to finish.
-  initialized_.store(true, std::memory_order_release);
-  // create additional queues if there are more destinations.
-  for (int i = queues_.size(); i < numDestinations; ++i) {
-    // create the destination queues inside the vector using emplace_back.
-    queues_.emplace_back(std::make_unique<UcxDestinationQueue>());
+
+  const bool canUseIntraNode =
+      kind != core::PartitionedOutputNode::Kind::kBroadcast;
+  for (auto& callback : eligibilityCallbacks) {
+    callback(canUseIntraNode);
   }
   return true;
 }
 
+void UcxOutputQueue::notifyOnIntraNodeEligibility(
+    UcxIntraNodeEligibilityCallback callback) {
+  VELOX_CHECK(callback, "Intra-node eligibility callback must be set");
+
+  bool invokeNow = false;
+  bool canUseIntraNode = false;
+  {
+    std::lock_guard<std::mutex> l(mutex_);
+    if (terminated_) {
+      invokeNow = true;
+    } else if (initialized_.load(std::memory_order_acquire)) {
+      invokeNow = true;
+      canUseIntraNode = kind_ != core::PartitionedOutputNode::Kind::kBroadcast;
+    } else {
+      intraNodeEligibilityCallbacks_.push_back(std::move(callback));
+    }
+  }
+
+  if (invokeNow) {
+    callback(canUseIntraNode);
+  }
+}
+
 void UcxOutputQueue::updateNumDrivers(uint32_t newNumDrivers) {
   bool isNoMoreDrivers{false};
   {
@@ -528,9 +561,11 @@ void UcxOutputQueue::deleteResults(int destination) {
 
 void UcxOutputQueue::terminate() {
   std::vector<UcxDataAvailable> pendingCallbacks;
+  std::vector<UcxIntraNodeEligibilityCallback> eligibilityCallbacks;
   std::vector<ContinuePromise> promises;
   {
     std::lock_guard<std::mutex> l(mutex_);
+    terminated_ = true;
     if (task_ && task_->isRunning()) {
       LOG(WARNING) << "UcxOutputQueue::terminate() called while task "
                    << task_->taskId() << " is still running";
@@ -546,12 +581,18 @@ void UcxOutputQueue::terminate() {
     }
     // Release any outstanding producer-side promises (blocked on queue-full).
     promises = std::move(promises_);
+    eligibilityCallbacks = std::move(intraNodeEligibilityCallbacks_);
   }
 
   // Fire callbacks outside of mutex to avoid potential deadlocks.
   for (auto& callback : pendingCallbacks) {
     callback.notify();
   }
+  // A task removed before initialization cannot use the process-local path.
+  // Resolving these callbacks lets waiting handshakes terminate cleanly.
+  for (auto& callback : eligibilityCallbacks) {
+    callback(false);
+  }
   // Unblock any blocked producers.
   for (auto& promise : promises) {
     promise.setValue();
diff --git a/velox/experimental/ucx-exchange/UcxQueues.h b/velox/experimental/ucx-exchange/UcxQueues.h
index f6715b7fb1d..3871cef92ea 100644
--- a/velox/experimental/ucx-exchange/UcxQueues.h
+++ b/velox/experimental/ucx-exchange/UcxQueues.h
@@ -42,6 +42,11 @@ using UcxDataAvailableCallback = std::function<void(
     vector_size_t numRows,
     std::vector<int64_t> remainingBytes)>;
 
+/// Called once output-queue initialization determines whether the task can use
+/// the process-local exchange path. False means the task is broadcast or was
+/// terminated before initialization.
+using UcxIntraNodeEligibilityCallback = std::function<void(bool)>;
+
 struct UcxDataAvailable {
   UcxDataAvailableCallback callback{nullptr};
   std::shared_ptr<cudf::packed_columns> data;
@@ -187,6 +192,12 @@ class UcxOutputQueue : public std::enable_shared_from_this<UcxOutputQueue> {
     return initialized_.load(std::memory_order_acquire);
   }
 
+  /// Invokes 'callback' once task metadata is initialized and its output kind
+  /// is known. If initialization already completed, invokes it immediately.
+  /// A callback receives false if this queue is broadcast or is terminated
+  /// before initialization.
+  void notifyOnIntraNodeEligibility(UcxIntraNodeEligibilityCallback callback);
+
   /// @brief When we understand the final number of split groups (for grouped
   /// execution only), we need to update the number of producing drivers here.
   void updateNumDrivers(uint32_t newNumDrivers);
@@ -318,6 +329,12 @@ class UcxOutputQueue : public std::enable_shared_from_this<UcxOutputQueue> {
   // written before the store are visible.
   std::atomic<bool> initialized_{false};
 
+  // Handshakes that arrived before initializeTask() wait here instead of
+  // permanently selecting the remote UCX path based on a placeholder queue.
+  std::vector<UcxIntraNodeEligibilityCallback> intraNodeEligibilityCallbacks_;
+
+  bool terminated_{false};
+
   // For broadcast: stores data for late-arriving destinations that need
   // backfill. Cleared once noMoreQueues_ is set.
   // Paired with the row count for the same reason QueuedPage is.
diff --git a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
index 0622382b0a7..5b99e5a6e0b 100644
--- a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
@@ -24,6 +24,7 @@
 #include <rmm/device_buffer.hpp>
 #include <memory>
 #include <type_traits>
+#include <optional>
 #include <vector>
 #include "velox/common/memory/MemoryPool.h"
 #include "velox/core/PlanNode.h"
@@ -641,6 +642,69 @@ TEST_F(UcxOutputQueueManagerTest, callbackFiredOnTerminateAfterInit) {
   EXPECT_TRUE(callback1Nullptr);
 }
 
+TEST_F(
+    UcxOutputQueueManagerTest,
+    intraNodeEligibilityWaitsForPartitionedTaskInitialization) {
+  const std::string taskId = "partitionedEligibility";
+  queueManager_->removeTask(taskId);
+
+  std::optional<bool> eligibility;
+  queueManager_->notifyOnIntraNodeEligibility(
+      taskId, [&](bool value) { eligibility = value; });
+  EXPECT_FALSE(eligibility.has_value());
+
+  auto task = initializeTask(
+      taskId,
+      /*numDestinations=*/2,
+      /*numDrivers=*/1,
+      /*cleanup=*/false,
+      core::PartitionedOutputNode::Kind::kPartitioned);
+
+  ASSERT_TRUE(eligibility.has_value());
+  EXPECT_TRUE(*eligibility);
+  queueManager_->removeTask(taskId);
+}
+
+TEST_F(
+    UcxOutputQueueManagerTest,
+    intraNodeEligibilityWaitsForBroadcastTaskInitialization) {
+  const std::string taskId = "broadcastEligibility";
+  queueManager_->removeTask(taskId);
+
+  std::optional<bool> eligibility;
+  queueManager_->notifyOnIntraNodeEligibility(
+      taskId, [&](bool value) { eligibility = value; });
+  EXPECT_FALSE(eligibility.has_value());
+
+  auto task = initializeTask(
+      taskId,
+      /*numDestinations=*/2,
+      /*numDrivers=*/1,
+      /*cleanup=*/false,
+      core::PartitionedOutputNode::Kind::kBroadcast);
+
+  ASSERT_TRUE(eligibility.has_value());
+  EXPECT_FALSE(*eligibility);
+  queueManager_->removeTask(taskId);
+}
+
+TEST_F(
+    UcxOutputQueueManagerTest,
+    intraNodeEligibilityResolvesWhenPlaceholderIsRemoved) {
+  const std::string taskId = "removedEligibility";
+  queueManager_->removeTask(taskId);
+
+  std::optional<bool> eligibility;
+  queueManager_->notifyOnIntraNodeEligibility(
+      taskId, [&](bool value) { eligibility = value; });
+  EXPECT_FALSE(eligibility.has_value());
+
+  queueManager_->removeTask(taskId);
+
+  ASSERT_TRUE(eligibility.has_value());
+  EXPECT_FALSE(*eligibility);
+}
+
 // --- Broadcast tests ---
 
 // Basic broadcast: enqueue data, all destinations receive the same data.

From 2a1b00f49c2b687aabfe563455f84fca2da7bd62 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Thu, 3 Sep 2026 14:47:12 -0700
Subject: [PATCH 04/15] fix(ucx-exchange): Clear stale cancellation before
 intra-node callbacks

---
 .../ucx-exchange/UcxOutputQueueManager.cpp    | 10 +++++-
 velox/experimental/ucx-exchange/UcxQueues.cpp | 13 +++++---
 velox/experimental/ucx-exchange/UcxQueues.h   |  4 ++-
 .../tests/UcxOutputQueueManagerTest.cpp       | 32 +++++++++++++++++++
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
index bef50a1c346..bf11f6ae93d 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
@@ -40,13 +40,15 @@ void UcxOutputQueueManager::initializeTask(
     int numDrivers,
     const std::string& /*transportOptions*/) {
   const auto& taskId = task->taskId();
+  std::vector<UcxIntraNodeEligibilityCallback> eligibilityCallbacks;
   queues_.withLock([&](auto& queues) {
     auto it = queues.find(taskId);
     if (it == queues.end()) {
       queues[taskId] = std::make_shared<UcxOutputQueue>(
           std::move(task), numDestinations, numDrivers, kind);
     } else {
-      if (!it->second->initialize(task, numDestinations, numDrivers, kind)) {
+      if (!it->second->initialize(
+              task, numDestinations, numDrivers, kind, &eligibilityCallbacks)) {
         VELOX_FAIL(
             "Registering a cudf output queue for pre-existing taskId {}",
             taskId);
@@ -59,6 +61,12 @@ void UcxOutputQueueManager::initializeTask(
   // Clear any stale "cancelled" state in the intra-node registry so
   // that the cancelledTasks_ set doesn't grow unboundedly across queries.
   IntraNodeTransferRegistry::getInstance()->clearCancelledTask(taskId);
+
+  const bool canUseIntraNode =
+      kind != core::PartitionedOutputNode::Kind::kBroadcast;
+  for (auto& callback : eligibilityCallbacks) {
+    callback(canUseIntraNode);
+  }
 }
 
 bool UcxOutputQueueManager::updateOutputBuffers(
diff --git a/velox/experimental/ucx-exchange/UcxQueues.cpp b/velox/experimental/ucx-exchange/UcxQueues.cpp
index d4e21b17809..50d78d105fb 100644
--- a/velox/experimental/ucx-exchange/UcxQueues.cpp
+++ b/velox/experimental/ucx-exchange/UcxQueues.cpp
@@ -165,8 +165,9 @@ bool UcxOutputQueue::initialize(
     std::shared_ptr<exec::Task> task,
     uint32_t numDestinations,
     uint32_t numDrivers,
-    core::PartitionedOutputNode::Kind kind) {
-  std::vector<UcxIntraNodeEligibilityCallback> eligibilityCallbacks;
+    core::PartitionedOutputNode::Kind kind,
+    std::vector<UcxIntraNodeEligibilityCallback>* eligibilityCallbacks) {
+  std::vector<UcxIntraNodeEligibilityCallback> localEligibilityCallbacks;
   {
     std::lock_guard<std::mutex> l(mutex_);
     if (task_) {
@@ -187,12 +188,16 @@ bool UcxOutputQueue::initialize(
       // create the destination queues inside the vector using emplace_back.
       queues_.emplace_back(std::make_unique<UcxDestinationQueue>());
     }
-    eligibilityCallbacks = std::move(intraNodeEligibilityCallbacks_);
+    if (eligibilityCallbacks) {
+      *eligibilityCallbacks = std::move(intraNodeEligibilityCallbacks_);
+    } else {
+      localEligibilityCallbacks = std::move(intraNodeEligibilityCallbacks_);
+    }
   }
 
   const bool canUseIntraNode =
       kind != core::PartitionedOutputNode::Kind::kBroadcast;
-  for (auto& callback : eligibilityCallbacks) {
+  for (auto& callback : localEligibilityCallbacks) {
     callback(canUseIntraNode);
   }
   return true;
diff --git a/velox/experimental/ucx-exchange/UcxQueues.h b/velox/experimental/ucx-exchange/UcxQueues.h
index 3871cef92ea..ba4e26dea0e 100644
--- a/velox/experimental/ucx-exchange/UcxQueues.h
+++ b/velox/experimental/ucx-exchange/UcxQueues.h
@@ -178,7 +178,9 @@ class UcxOutputQueue : public std::enable_shared_from_this<UcxOutputQueue> {
       uint32_t numDestinations,
       uint32_t numDrivers,
       core::PartitionedOutputNode::Kind kind =
-          core::PartitionedOutputNode::Kind::kPartitioned);
+          core::PartitionedOutputNode::Kind::kPartitioned,
+      std::vector<UcxIntraNodeEligibilityCallback>* eligibilityCallbacks =
+          nullptr);
 
   core::PartitionedOutputNode::Kind kind() const {
     return kind_;
diff --git a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
index 5b99e5a6e0b..c7f8cc8f06e 100644
--- a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
@@ -30,6 +30,7 @@
 #include "velox/core/PlanNode.h"
 #include "velox/exec/OutputTransportRegistry.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
+#include "velox/experimental/ucx-exchange/IntraNodeTransferRegistry.h"
 #include "velox/experimental/ucx-exchange/tests/UcxTestHelpers.h"
 
 using namespace facebook::velox::ucx_exchange;
@@ -705,6 +706,37 @@ TEST_F(
   EXPECT_FALSE(*eligibility);
 }
 
+TEST_F(
+    UcxOutputQueueManagerTest,
+    intraNodeEligibilityClearsStaleCancellationBeforeCallback) {
+  const std::string taskId = "reusedEligibility";
+  queueManager_->removeTask(taskId);
+
+  std::optional<bool> eligibility;
+  std::optional<IntraNodeTransferResult> registryPollResult;
+  queueManager_->notifyOnIntraNodeEligibility(
+      taskId,
+      [&](bool value) {
+        eligibility = value;
+        registryPollResult = IntraNodeTransferRegistry::getInstance()->poll(
+            IntraNodeTransferKey{
+                taskId, /*destination=*/0, /*sequenceNumber=*/0});
+      });
+  EXPECT_FALSE(eligibility.has_value());
+
+  auto task = initializeTask(
+      taskId,
+      /*numDestinations=*/1,
+      /*numDrivers=*/1,
+      /*cleanup=*/false,
+      core::PartitionedOutputNode::Kind::kPartitioned);
+
+  ASSERT_TRUE(eligibility.has_value());
+  EXPECT_TRUE(*eligibility);
+  EXPECT_FALSE(registryPollResult.has_value());
+  queueManager_->removeTask(taskId);
+}
+
 // --- Broadcast tests ---
 
 // Basic broadcast: enqueue data, all destinations receive the same data.

From 34b4396a2e3593d5f33dc4a2ece274d25579955f Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Thu, 3 Sep 2026 14:48:32 -0700
Subject: [PATCH 05/15] perf(ucx-exchange): Bound deferred communicator actions

---
 velox/experimental/ucx-exchange/Communicator.cpp | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/velox/experimental/ucx-exchange/Communicator.cpp b/velox/experimental/ucx-exchange/Communicator.cpp
index d11d4f37239..c1efbac01ea 100644
--- a/velox/experimental/ucx-exchange/Communicator.cpp
+++ b/velox/experimental/ucx-exchange/Communicator.cpp
@@ -150,6 +150,7 @@ void Communicator::run() {
   // Therefore workQueue_ can remain non-empty indefinitely. Process it in
   // bounded batches so deferred handshakes, endpoint cleanup, and heartbeat
   // work at the top of this loop cannot be starved.
+  constexpr size_t kMaxDeferredActionsPerLoop = 256;
   constexpr size_t kMaxWorkItemsPerLoop = 256;
   while (running_) {
     try {
@@ -214,22 +215,29 @@ void Communicator::run() {
 
       // Run actions submitted by other threads. These must execute here
       // because they may issue UCXX operations on this worker.
-      while (auto deferredAction = deferredActions_.pop()) {
+      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.
-      size_t workItemsThisLoop = 0;
-      while (workItemsThisLoop < kMaxWorkItemsPerLoop) {
+      for (size_t workItemsThisLoop = 0;
+           workItemsThisLoop < kMaxWorkItemsPerLoop;
+           ++workItemsThisLoop) {
         auto comms = workQueue_.pop();
         if (!comms) {
           break;
         }
         comms->process();
         ++workItemsProcessed_;
-        ++workItemsThisLoop;
         // Progress after each work item to allow UCXX to advance
         // its internal state (complete sends/receives, fire callbacks).
         // Use non-blocking progress here to avoid blocking between

From ad6fdd2bf5640381ec180ac0eae435327feabc73 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 8 Sep 2026 17:08:46 -0700
Subject: [PATCH 06/15] feat(ucx-exchange): Expose transport decisions in
 source metrics

---
 velox/experimental/ucx-exchange/Acceptor.cpp  | 19 -----
 velox/experimental/ucx-exchange/Acceptor.h    |  9 ---
 .../ucx-exchange/UcxExchangeSource.cpp        | 12 ++++
 .../ucx-exchange/UcxExchangeSource.h          |  6 ++
 .../ucx-exchange/tests/UcxExchangeTest.cpp    | 70 +++++++++++++++++++
 5 files changed, 88 insertions(+), 28 deletions(-)

diff --git a/velox/experimental/ucx-exchange/Acceptor.cpp b/velox/experimental/ucx-exchange/Acceptor.cpp
index edf930de454..7f848d0b2c3 100644
--- a/velox/experimental/ucx-exchange/Acceptor.cpp
+++ b/velox/experimental/ucx-exchange/Acceptor.cpp
@@ -21,15 +21,10 @@
 #include "velox/experimental/ucx-exchange/UcxExchangeServer.h"
 #include "velox/experimental/ucx-exchange/UcxOutputQueueManager.h"
 
-#include <atomic>
-
 namespace facebook::velox::ucx_exchange {
 
 namespace {
 
-std::atomic<uint64_t> gNumIntraNodeHandshakes{0};
-std::atomic<uint64_t> gNumRemoteHandshakes{0};
-
 void completeHandshake(
     const std::shared_ptr<Communicator>& communicator,
     const std::shared_ptr<EndpointRef>& epRef,
@@ -47,12 +42,6 @@ void completeHandshake(
   auto exchangeServer =
       UcxExchangeServer::create(communicator, epRef, key, isIntraNodeTransfer);
 
-  if (isIntraNodeTransfer) {
-    gNumIntraNodeHandshakes.fetch_add(1, std::memory_order_relaxed);
-  } else {
-    gNumRemoteHandshakes.fetch_add(1, std::memory_order_relaxed);
-  }
-
   // Add this exchangeServer to the endpoint reference.
   epRef->addCommElem(exchangeServer);
 
@@ -95,14 +84,6 @@ void completeHandshake(
 
 } // namespace
 
-uint64_t Acceptor::numIntraNodeHandshakes() {
-  return gNumIntraNodeHandshakes.load(std::memory_order_relaxed);
-}
-
-uint64_t Acceptor::numRemoteHandshakes() {
-  return gNumRemoteHandshakes.load(std::memory_order_relaxed);
-}
-
 /*static*/
 void Acceptor::cStyleAMCallback(
     std::shared_ptr<ucxx::Request> request,
diff --git a/velox/experimental/ucx-exchange/Acceptor.h b/velox/experimental/ucx-exchange/Acceptor.h
index a44f35ef806..063532b836c 100644
--- a/velox/experimental/ucx-exchange/Acceptor.h
+++ b/velox/experimental/ucx-exchange/Acceptor.h
@@ -16,7 +16,6 @@
 #pragma once
 
 #include <ucxx/api.h>
-#include <cstdint>
 #include <map>
 #include "velox/experimental/ucx-exchange/EndpointRef.h"
 
@@ -32,14 +31,6 @@ struct Acceptor {
       std::shared_ptr<ucxx::Request> request,
       ucp_ep_h ep);
 
-  /// Number of completed handshakes routed through the process-local path.
-  /// Exposed for diagnostics and regression tests.
-  static uint64_t numIntraNodeHandshakes();
-
-  /// Number of completed handshakes routed through the remote UCX path.
-  /// Exposed for diagnostics and regression tests.
-  static uint64_t numRemoteHandshakes();
-
   /// @brief Adds the endpoint reference to the handleToEndpointRef_ map such
   /// that endpoint handles can be resolved
   void registerEndpointRef(std::shared_ptr<EndpointRef> endpointRef);
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
index 236ff60524c..e3a1a37e283 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
@@ -278,10 +278,21 @@ folly::F14FastMap<std::string, RuntimeMetric> UcxExchangeSource::metrics()
   map["ucxExchangeSource.intraNodeBytes"] = metrics_.intraNodeBytes_;
   map["ucxExchangeSource.remotePackedColumns"] = metrics_.remotePackedColumns_;
   map["ucxExchangeSource.remoteBytes"] = metrics_.remoteBytes_;
+  map["ucxExchangeSource.intraNodeSources"] = metrics_.intraNodeSources_;
+  map["ucxExchangeSource.remoteSources"] = metrics_.remoteSources_;
   map["ucxExchangeSource.rttPerRequest"] = metrics_.rttPerRequest_;
   return map;
 }
 
+void UcxExchangeSource::recordTransportDecisionMetric() {
+  std::lock_guard<std::mutex> lock(metricsMutex_);
+  if (isIntraNodeTransfer_) {
+    metrics_.intraNodeSources_.addValue(1);
+  } else {
+    metrics_.remoteSources_.addValue(1);
+  }
+}
+
 void UcxExchangeSource::recordPayloadMetrics(uint64_t dataBytes) {
   const auto bytes = saturateCast(dataBytes);
   std::lock_guard<std::mutex> lock(metricsMutex_);
@@ -708,6 +719,7 @@ void UcxExchangeSource::onHandshakeResponse(
       std::static_pointer_cast<HandshakeResponse>(arg);
 
   isIntraNodeTransfer_ = response->isIntraNodeTransfer;
+  recordTransportDecisionMetric();
 
   VLOG(3) << toString() << " + onHandshakeResponse isIntraNodeTransfer="
           << isIntraNodeTransfer_;
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index 138e86f1da0..fb9f5fd4790 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -48,6 +48,8 @@ struct UcxExchangeMetrics {
         intraNodeBytes_(RuntimeCounter::Unit::kBytes),
         remotePackedColumns_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
         remoteBytes_(RuntimeCounter::Unit::kBytes),
+        intraNodeSources_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
+        remoteSources_(RuntimeMetric(RuntimeCounter::Unit::kNone)),
         rttPerRequest_(RuntimeMetric(RuntimeCounter::Unit::kNanos)) {}
   RuntimeMetric numPackedColumns_; // total number of packed columns received.
   RuntimeMetric totalBytes_; // total number of bytes received
@@ -55,6 +57,8 @@ struct UcxExchangeMetrics {
   RuntimeMetric intraNodeBytes_;
   RuntimeMetric remotePackedColumns_;
   RuntimeMetric remoteBytes_;
+  RuntimeMetric intraNodeSources_;
+  RuntimeMetric remoteSources_;
   RuntimeMetric rttPerRequest_;
 };
 
@@ -256,6 +260,8 @@ class UcxExchangeSource
   /// @return Returns true if state was changed, false otherwise.
   bool setStateIf(ReceiverState expected, ReceiverState desired);
 
+  void recordTransportDecisionMetric();
+
   void recordPayloadMetrics(uint64_t dataBytes);
 
   // The connection parameters
diff --git a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
index 108e3159770..ef1e614c3ea 100644
--- a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
@@ -1116,11 +1116,52 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
   };
 
   const std::string taskPrefix = getUniqueTaskPrefix();
+  const std::string pollingSrcTaskId = taskPrefix + "pollingSrc";
+  const std::string pollingSinkTaskId = taskPrefix + "pollingSink";
   const std::string srcTaskId = taskPrefix + "partitionedPlaceholderSrc";
   const std::string sinkTaskId = taskPrefix + "partitionedPlaceholderSink";
   constexpr int numChunks = 2;
   constexpr int numRowsPerChunk = 1000;
 
+  auto waitForMetric = [&metricSum](
+                           const std::shared_ptr<SinkDriverMock>& sinkDriver,
+                           std::string_view name,
+                           int64_t target) {
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::seconds(5);
+    while (metricSum(sinkDriver->exchangeClientStats(), name) < target &&
+           std::chrono::steady_clock::now() < deadline) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }
+    return metricSum(sinkDriver->exchangeClientStats(), name) >= target;
+  };
+
+  // Establish one intra-node source with no data. It continuously re-enqueues
+  // itself while polling the transfer registry, keeping the communicator work
+  // queue non-empty. A later deferred handshake must still make progress.
+  auto pollingSrcTask =
+      createSourceTask(pollingSrcTaskId, pool_, UcxTestData::kTestRowType);
+  queueManager_->initializeTask(
+      pollingSrcTask,
+      core::PartitionedOutputNode::Kind::kPartitioned,
+      /*numDestinations=*/1,
+      /*numDrivers=*/1);
+
+  core::PlanNodeId pollingExchangeNodeId;
+  auto pollingSinkTask = createExchangeTask(
+      pollingSinkTaskId,
+      UcxTestData::kTestRowType,
+      /*partitionId=*/0,
+      pollingExchangeNodeId);
+  auto pollingSinkDriver =
+      std::make_shared<SinkDriverMock>(pollingSinkTask, /*numDrivers=*/1);
+  std::vector<exec::Split> pollingSplits;
+  pollingSplits.emplace_back(remoteSplit(pollingSrcTaskId, /*partitionId=*/0));
+  pollingSinkDriver->addSplits(pollingSplits);
+  pollingSinkDriver->run();
+  EXPECT_TRUE(waitForMetric(
+      pollingSinkDriver, "ucxExchangeSource.intraNodeSources", 1));
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
   core::PlanNodeId exchangeNodeId;
   auto sinkTask = createExchangeTask(
       sinkTaskId, UcxTestData::kTestRowType, /*partitionId=*/0, exchangeNodeId);
@@ -1144,21 +1185,39 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
       /*numDrivers=*/1);
   EXPECT_TRUE(queueManager_->canUseIntraNode(srcTaskId));
 
+  // This specifically catches starvation in Communicator::run(): the polling
+  // source above keeps workQueue_ non-empty while this handshake is submitted
+  // through deferredActions_.
+  EXPECT_TRUE(
+      waitForMetric(sinkDriver, "ucxExchangeSource.intraNodeSources", 1));
   auto sourceMock = std::make_shared<UcxPartitionedOutputMock>(
       srcTaskId,
       /*numDrivers=*/1,
       /*numPartitions=*/1,
       numChunks,
       numRowsPerChunk);
+  auto pollingSourceMock = std::make_shared<UcxPartitionedOutputMock>(
+      pollingSrcTaskId,
+      /*numDrivers=*/1,
+      /*numPartitions=*/1,
+      numChunks,
+      numRowsPerChunk);
   sourceMock->run();
+  pollingSourceMock->run();
 
   sourceMock->joinThreads();
+  pollingSourceMock->joinThreads();
   sinkDriver->joinThreads();
+  pollingSinkDriver->joinThreads();
 
   EXPECT_EQ(
       sinkDriver->numRows(),
       static_cast<uint64_t>(numChunks * numRowsPerChunk));
 
+  EXPECT_EQ(
+      pollingSinkDriver->numRows(),
+      static_cast<uint64_t>(numChunks * numRowsPerChunk));
+
   const auto sinkStats = sinkDriver->exchangeClientStats();
   EXPECT_EQ(
       metricSum(sinkStats, "ucxExchangeSource.intraNodePackedColumns"),
@@ -1167,7 +1226,18 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
   EXPECT_EQ(metricSum(sinkStats, "ucxExchangeSource.remotePackedColumns"), 0);
   EXPECT_EQ(metricSum(sinkStats, "ucxExchangeSource.remoteBytes"), 0);
 
+  const auto pollingSinkStats = pollingSinkDriver->exchangeClientStats();
+  EXPECT_EQ(
+      metricSum(pollingSinkStats, "ucxExchangeSource.intraNodePackedColumns"),
+      numChunks);
+  EXPECT_GT(
+      metricSum(pollingSinkStats, "ucxExchangeSource.intraNodeBytes"), 0);
+  EXPECT_EQ(
+      metricSum(pollingSinkStats, "ucxExchangeSource.remotePackedColumns"), 0);
+  EXPECT_EQ(metricSum(pollingSinkStats, "ucxExchangeSource.remoteBytes"), 0);
+
   queueManager_->removeTask(srcTaskId);
+  queueManager_->removeTask(pollingSrcTaskId);
 }
 
 // Regression test for broadcast + intra-node SIGSEGV.

From 9e886183a7e85969851578e086ce8ca297b9130b Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 15:53:09 -0700
Subject: [PATCH 07/15] fix(ucx-exchange): Bound cleanup and stabilize
 regression

---
 .../ucx-exchange/Communicator.cpp             | 12 +++--
 .../ucx-exchange/UcxExchangeSource.h          |  2 +
 .../ucx-exchange/tests/UcxExchangeTest.cpp    | 46 +++++++++++++++----
 3 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/velox/experimental/ucx-exchange/Communicator.cpp b/velox/experimental/ucx-exchange/Communicator.cpp
index c1efbac01ea..6d5cbbed3a1 100644
--- a/velox/experimental/ucx-exchange/Communicator.cpp
+++ b/velox/experimental/ucx-exchange/Communicator.cpp
@@ -148,8 +148,8 @@ void Communicator::run() {
   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 deferred handshakes, endpoint cleanup, and heartbeat
-  // work at the top of this loop cannot be starved.
+  // 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_) {
@@ -204,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
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index fb9f5fd4790..642ad57d1b1 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -57,6 +57,8 @@ struct UcxExchangeMetrics {
   RuntimeMetric intraNodeBytes_;
   RuntimeMetric remotePackedColumns_;
   RuntimeMetric remoteBytes_;
+  // Completed source handshakes by selected transport, not a live-source
+  // gauge.
   RuntimeMetric intraNodeSources_;
   RuntimeMetric remoteSources_;
   RuntimeMetric rttPerRequest_;
diff --git a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
index ef1e614c3ea..1996bf3465f 100644
--- a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
@@ -1089,14 +1089,15 @@ TEST_P(UcxExchangeTest, intraNodeTaskRemovalLivelock) {
 // A same-process consumer can connect before the producer task has initialized
 // its output queue. The handshake must wait for the real output kind instead of
 // permanently falling back to remote UCX based on the placeholder queue.
-TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
+TEST_P(UcxExchangeTest, partitionedPlaceholderDefersHandshakeWithBusyWorker) {
   // This test doesn't use parameters — run only for the first param set.
   {
     ExchangeTestParams p = GetParam();
     if (p.numSrcDrivers != 1 || p.numDstDrivers != 1 || p.numPartitions != 1 ||
         p.numChunks != 100 || p.numUpstreamTasks != 1 ||
         p.tableType != TableType::NARROW) {
-      GTEST_SKIP() << "partitionedPlaceholderUsesIntraNodePath: runs only once";
+      GTEST_SKIP()
+          << "partitionedPlaceholderDefersHandshakeWithBusyWorker: runs only once";
     }
   }
 
@@ -1135,6 +1136,15 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
     }
     return metricSum(sinkDriver->exchangeClientStats(), name) >= target;
   };
+  auto waitForPlaceholder = [this](const std::string& taskId) {
+    const auto deadline =
+        std::chrono::steady_clock::now() + std::chrono::seconds(5);
+    while (!queueManager_->stats(taskId).has_value() &&
+           std::chrono::steady_clock::now() < deadline) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    }
+    return queueManager_->stats(taskId).has_value();
+  };
 
   // Establish one intra-node source with no data. It continuously re-enqueues
   // itself while polling the transfer registry, keeping the communicator work
@@ -1159,8 +1169,13 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
   pollingSplits.emplace_back(remoteSplit(pollingSrcTaskId, /*partitionId=*/0));
   pollingSinkDriver->addSplits(pollingSplits);
   pollingSinkDriver->run();
-  EXPECT_TRUE(waitForMetric(
-      pollingSinkDriver, "ucxExchangeSource.intraNodeSources", 1));
+  if (!waitForMetric(
+          pollingSinkDriver, "ucxExchangeSource.intraNodeSources", 1)) {
+    pollingSinkTask->requestAbort();
+    pollingSinkDriver->joinThreads();
+    FAIL() << "Polling source did not select intra-node transfer";
+    return;
+  }
   std::this_thread::sleep_for(std::chrono::milliseconds(100));
   core::PlanNodeId exchangeNodeId;
   auto sinkTask = createExchangeTask(
@@ -1171,10 +1186,17 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
   splits.emplace_back(remoteSplit(srcTaskId, /*partitionId=*/0));
   sinkDriver->addSplits(splits);
 
-  // Start the consumer first and give its handshake time to reach the server.
-  // At this point the producer is represented only by an uninitialized task ID.
+  // Start the consumer first. Waiting for the placeholder proves that its
+  // handshake has reached the output queue while the producer is uninitialized.
   sinkDriver->run();
-  std::this_thread::sleep_for(std::chrono::seconds(1));
+  if (!waitForPlaceholder(srcTaskId)) {
+    sinkTask->requestAbort();
+    pollingSinkTask->requestAbort();
+    sinkDriver->joinThreads();
+    pollingSinkDriver->joinThreads();
+    FAIL() << "Consumer handshake did not create the producer placeholder";
+    return;
+  }
   EXPECT_FALSE(queueManager_->canUseIntraNode(srcTaskId));
 
   auto srcTask = createSourceTask(srcTaskId, pool_, UcxTestData::kTestRowType);
@@ -1188,8 +1210,14 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderUsesIntraNodePath) {
   // This specifically catches starvation in Communicator::run(): the polling
   // source above keeps workQueue_ non-empty while this handshake is submitted
   // through deferredActions_.
-  EXPECT_TRUE(
-      waitForMetric(sinkDriver, "ucxExchangeSource.intraNodeSources", 1));
+  if (!waitForMetric(sinkDriver, "ucxExchangeSource.intraNodeSources", 1)) {
+    sinkTask->requestAbort();
+    pollingSinkTask->requestAbort();
+    sinkDriver->joinThreads();
+    pollingSinkDriver->joinThreads();
+    FAIL() << "Deferred handshake did not select intra-node transfer";
+    return;
+  }
   auto sourceMock = std::make_shared<UcxPartitionedOutputMock>(
       srcTaskId,
       /*numDrivers=*/1,

From cd99641e9915af03ace7f199184d10d4320edbc6 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 16:23:03 -0700
Subject: [PATCH 08/15] refactor(ucx-exchange): Remove unreachable stats
 fallback

---
 .../experimental/ucx-exchange/UcxExchangeClient.cpp | 13 +++----------
 .../experimental/ucx-exchange/UcxExchangeSource.cpp |  4 ----
 velox/experimental/ucx-exchange/UcxExchangeSource.h |  9 ---------
 3 files changed, 3 insertions(+), 23 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp b/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
index 393bde807c9..027c7871b47 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
@@ -89,16 +89,9 @@ 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);
-      }
+    for (const auto& [name, value] : source->metrics()) {
+      auto [iter, inserted] = stats.try_emplace(name, value.unit);
+      iter->second.merge(value);
     }
   }
 
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
index e3a1a37e283..525f9493e08 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
@@ -260,10 +260,6 @@ void UcxExchangeSource::resumeFromBackpressure() {
   }
 }
 
-folly::F14FastMap<std::string, int64_t> UcxExchangeSource::stats() const {
-  VELOX_UNREACHABLE();
-}
-
 folly::F14FastMap<std::string, RuntimeMetric> UcxExchangeSource::metrics()
     const {
   std::lock_guard<std::mutex> lock(metricsMutex_);
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index 642ad57d1b1..327a56e7634 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -101,10 +101,6 @@ class UcxExchangeSource
       std::string_view url,
       const std::shared_ptr<UcxExchangeQueue>& queue);
 
-  bool supportsMetrics() const {
-    return true;
-  }
-
   /// @brief Advances the UCXX communication that was started by issuing
   /// "request"
   void process() override;
@@ -132,11 +128,6 @@ class UcxExchangeSource
   static constexpr int32_t kBackpressureHighWaterMark = 32;
   static constexpr int32_t kBackpressureLowWaterMark = 16;
 
-  // Returns runtime statistics. ExchangeSource is expected to report
-  // background CPU time by including a runtime metric named
-  // ExchangeClient::kBackgroundCpuTimeMs.
-  folly::F14FastMap<std::string, int64_t> stats() const;
-
   /// Returns runtime statistics. ExchangeSource is expected to report
   /// Specify units of individual counters in ExchangeSource.
   /// for an example: 'totalBytes ：count: 9, sum: 11.17GB, max: 1.39GB,

From 6aa7238db4a722a1bb7f02e39487b146660b1610 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 17:34:05 -0700
Subject: [PATCH 09/15] test(ucx-exchange): Scope transport configuration
 regressions

---
 .../ucx-exchange/tests/UcxExchangeTest.cpp    | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
index 1996bf3465f..e7aeac95f48 100644
--- a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
@@ -1027,6 +1027,13 @@ TEST_P(UcxExchangeTest, intraNodeTaskRemovalLivelock) {
     }
   }
 
+  auto& config = cudf_velox::CudfConfig::getInstance();
+  const bool origIntraNode = config.intraNodeExchange;
+  config.intraNodeExchange = true;
+  SCOPE_EXIT {
+    config.intraNodeExchange = origIntraNode;
+  };
+
   const std::string taskPrefix = getUniqueTaskPrefix();
   const std::string srcTaskId = taskPrefix + "srcProducerNeverSends";
   const std::string sinkTaskId = taskPrefix + "sinkConsumer";
@@ -1041,6 +1048,7 @@ TEST_P(UcxExchangeTest, intraNodeTaskRemovalLivelock) {
       core::PartitionedOutputNode::Kind::kPartitioned,
       numPartitions,
       /*numDrivers=*/1);
+  EXPECT_TRUE(queueManager_->canUseIntraNode(srcTaskId));
 
   // 2. Create sink task with exchange plan node.
   core::PlanNodeId exchangeNodeId;
@@ -1290,6 +1298,9 @@ TEST_P(UcxExchangeTest, broadcastIntraNodeFallback) {
   auto& config = cudf_velox::CudfConfig::getInstance();
   const bool origIntraNode = config.intraNodeExchange;
   config.intraNodeExchange = true;
+  SCOPE_EXIT {
+    config.intraNodeExchange = origIntraNode;
+  };
 
   const std::string taskPrefix = getUniqueTaskPrefix();
   const std::string srcTaskId = taskPrefix + "broadcastSrc";
@@ -1305,6 +1316,7 @@ TEST_P(UcxExchangeTest, broadcastIntraNodeFallback) {
       core::PartitionedOutputNode::Kind::kBroadcast,
       numDestinations,
       numDrivers);
+  EXPECT_FALSE(queueManager_->canUseIntraNode(srcTaskId));
   // Finalize destinations for broadcast.
   queueManager_->updateOutputBuffers(srcTaskId, numDestinations, true);
 
@@ -1353,7 +1365,6 @@ TEST_P(UcxExchangeTest, broadcastIntraNodeFallback) {
 
   // Cleanup.
   queueManager_->removeTask(srcTaskId);
-  config.intraNodeExchange = origIntraNode;
 }
 
 // Regression test for broadcast + intra-node placeholder race condition.
@@ -1379,6 +1390,9 @@ TEST_P(UcxExchangeTest, broadcastIntraNodePlaceholderRace) {
   auto& config = cudf_velox::CudfConfig::getInstance();
   const bool origIntraNode = config.intraNodeExchange;
   config.intraNodeExchange = true;
+  SCOPE_EXIT {
+    config.intraNodeExchange = origIntraNode;
+  };
 
   const std::string taskPrefix = getUniqueTaskPrefix();
   const std::string srcTaskId = taskPrefix + "broadcastPlaceholderSrc";
@@ -1423,6 +1437,7 @@ TEST_P(UcxExchangeTest, broadcastIntraNodePlaceholderRace) {
       core::PartitionedOutputNode::Kind::kBroadcast,
       numDestinations,
       numDrivers);
+  EXPECT_FALSE(queueManager_->canUseIntraNode(srcTaskId));
 
   // Step 4: Finalize destinations for broadcast.
   queueManager_->updateOutputBuffers(srcTaskId, numDestinations, true);
@@ -1448,7 +1463,6 @@ TEST_P(UcxExchangeTest, broadcastIntraNodePlaceholderRace) {
 
   // Cleanup.
   queueManager_->removeTask(srcTaskId);
-  config.intraNodeExchange = origIntraNode;
 }
 
 // Test that UcxPartitionedOutput's batch accumulation correctly merges many
@@ -1771,6 +1785,9 @@ TEST_P(UcxExchangeTest, deferredRequestCleanupOnTaskAbort) {
   auto& config = cudf_velox::CudfConfig::getInstance();
   const bool origIntraNode = config.intraNodeExchange;
   config.intraNodeExchange = false;
+  SCOPE_EXIT {
+    config.intraNodeExchange = origIntraNode;
+  };
 
   const std::string taskPrefix = getUniqueTaskPrefix();
   const std::string srcTaskId = taskPrefix + "srcActiveTransfer";
@@ -1841,8 +1858,6 @@ TEST_P(UcxExchangeTest, deferredRequestCleanupOnTaskAbort) {
 
   // If we reach here without crashing, the deferred cleanup is working.
   VLOG(0) << "deferredRequestCleanupOnTaskAbort: completed without crash";
-
-  config.intraNodeExchange = origIntraNode;
 }
 
 std::shared_ptr<UcxOutputQueueManager> UcxExchangeTest::queueManager_;

From 1f47c40c6a4afe50833aa53a5f9007f457ce859a Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 18:02:11 -0700
Subject: [PATCH 10/15] docs(ucx-exchange): Clarify transport decision metrics

---
 velox/experimental/ucx-exchange/UcxExchangeSource.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index 327a56e7634..7ae53e85a64 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -57,8 +57,8 @@ struct UcxExchangeMetrics {
   RuntimeMetric intraNodeBytes_;
   RuntimeMetric remotePackedColumns_;
   RuntimeMetric remoteBytes_;
-  // Completed source handshakes by selected transport, not a live-source
-  // gauge.
+  // Cumulative count of sources that completed a handshake for each transport.
+  // Not a live-source gauge.
   RuntimeMetric intraNodeSources_;
   RuntimeMetric remoteSources_;
   RuntimeMetric rttPerRequest_;

From 56a95a7b944034a2262124645fd863ac29b0be46 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 18:04:39 -0700
Subject: [PATCH 11/15] Revert "refactor(ucx-exchange): Remove unreachable
 stats fallback"

This reverts commit cd99641e9915af03ace7f199184d10d4320edbc6.
---
 .../experimental/ucx-exchange/UcxExchangeClient.cpp | 13 ++++++++++---
 .../experimental/ucx-exchange/UcxExchangeSource.cpp |  4 ++++
 velox/experimental/ucx-exchange/UcxExchangeSource.h |  9 +++++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp b/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
index 027c7871b47..393bde807c9 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeClient.cpp
@@ -89,9 +89,16 @@ UcxExchangeClient::collectStatsLocked() const {
   folly::F14FastMap<std::string, RuntimeMetric> stats;
 
   for (const auto& source : sources_) {
-    for (const auto& [name, value] : source->metrics()) {
-      auto [iter, inserted] = stats.try_emplace(name, value.unit);
-      iter->second.merge(value);
+    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);
+      }
     }
   }
 
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
index 525f9493e08..e3a1a37e283 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.cpp
@@ -260,6 +260,10 @@ void UcxExchangeSource::resumeFromBackpressure() {
   }
 }
 
+folly::F14FastMap<std::string, int64_t> UcxExchangeSource::stats() const {
+  VELOX_UNREACHABLE();
+}
+
 folly::F14FastMap<std::string, RuntimeMetric> UcxExchangeSource::metrics()
     const {
   std::lock_guard<std::mutex> lock(metricsMutex_);
diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index 7ae53e85a64..8a023c8bb32 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -101,6 +101,10 @@ class UcxExchangeSource
       std::string_view url,
       const std::shared_ptr<UcxExchangeQueue>& queue);
 
+  bool supportsMetrics() const {
+    return true;
+  }
+
   /// @brief Advances the UCXX communication that was started by issuing
   /// "request"
   void process() override;
@@ -128,6 +132,11 @@ class UcxExchangeSource
   static constexpr int32_t kBackpressureHighWaterMark = 32;
   static constexpr int32_t kBackpressureLowWaterMark = 16;
 
+  // Returns runtime statistics. ExchangeSource is expected to report
+  // background CPU time by including a runtime metric named
+  // ExchangeClient::kBackgroundCpuTimeMs.
+  folly::F14FastMap<std::string, int64_t> stats() const;
+
   /// Returns runtime statistics. ExchangeSource is expected to report
   /// Specify units of individual counters in ExchangeSource.
   /// for an example: 'totalBytes ：count: 9, sum: 11.17GB, max: 1.39GB,

From 408f063557489a2947852c0c845c8a57a2a84e6d Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 18:11:08 -0700
Subject: [PATCH 12/15] fix(ucx-exchange): Serialize task lifecycle transitions

---
 .../ucx-exchange/UcxExchangeSource.h          |  6 +-
 .../ucx-exchange/UcxOutputQueueManager.cpp    | 91 +++++++++++--------
 .../ucx-exchange/UcxOutputQueueManager.h      |  6 ++
 velox/experimental/ucx-exchange/UcxQueues.h   |  2 +
 4 files changed, 62 insertions(+), 43 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeSource.h b/velox/experimental/ucx-exchange/UcxExchangeSource.h
index 8a023c8bb32..efd5ec2bfa6 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeSource.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeSource.h
@@ -296,10 +296,10 @@ class UcxExchangeSource
   /// while deliverEndMarker() reads it on the Communicator thread.
   std::atomic<bool> registered_{false};
 
-  /// True if the server detected that this source is on the same node.
+  /// True if the server detected that this source uses the same Communicator
+  /// instance (and therefore the same process).
   /// Set from the isIntraNodeTransfer flag in HandshakeResponse, which the
-  /// server determines by comparing this source's listener address (sent in
-  /// HandshakeMsg) with its own Communicator's listener address.
+  /// server determines by comparing the source and server worker IDs.
   /// When true, intra-node transfer optimizations bypass UCXX transfers.
   bool isIntraNodeTransfer_{false};
 
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
index bf11f6ae93d..faafcbbe699 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
@@ -41,26 +41,33 @@ void UcxOutputQueueManager::initializeTask(
     const std::string& /*transportOptions*/) {
   const auto& taskId = task->taskId();
   std::vector<UcxIntraNodeEligibilityCallback> eligibilityCallbacks;
-  queues_.withLock([&](auto& queues) {
-    auto it = queues.find(taskId);
-    if (it == queues.end()) {
-      queues[taskId] = std::make_shared<UcxOutputQueue>(
-          std::move(task), numDestinations, numDrivers, kind);
-    } else {
-      if (!it->second->initialize(
-              task, numDestinations, numDrivers, kind, &eligibilityCallbacks)) {
-        VELOX_FAIL(
-            "Registering a cudf output queue for pre-existing taskId {}",
-            taskId);
+  {
+    std::lock_guard<std::mutex> lifecycleLock(taskLifecycleMutex_);
+    // A reused task ID must clear the previous incarnation's cancellation
+    // before publishing its initialized queue to handshake callbacks.
+    IntraNodeTransferRegistry::getInstance()->clearCancelledTask(taskId);
+    queues_.withLock([&](auto& queues) {
+      // Keep queues_ before removedTasks_ lock ordering, matching removeTask()
+      // and getData().
+      removedTasks_.withLock([&](auto& removed) { removed.erase(taskId); });
+      auto it = queues.find(taskId);
+      if (it == queues.end()) {
+        queues[taskId] = std::make_shared<UcxOutputQueue>(
+            std::move(task), numDestinations, numDrivers, kind);
+      } else {
+        if (!it->second->initialize(
+                task,
+                numDestinations,
+                numDrivers,
+                kind,
+                &eligibilityCallbacks)) {
+          VELOX_FAIL(
+              "Registering a cudf output queue for pre-existing taskId {}",
+              taskId);
+        }
       }
-    }
-  });
-  // Clear any stale "removed" state so that getData() calls after this
-  // initializeTask() create proper placeholder queues if needed.
-  removedTasks_.withLock([&](auto& removed) { removed.erase(taskId); });
-  // Clear any stale "cancelled" state in the intra-node registry so
-  // that the cancelledTasks_ set doesn't grow unboundedly across queries.
-  IntraNodeTransferRegistry::getInstance()->clearCancelledTask(taskId);
+    });
+  }
 
   const bool canUseIntraNode =
       kind != core::PartitionedOutputNode::Kind::kBroadcast;
@@ -197,33 +204,37 @@ void UcxOutputQueueManager::notifyOnIntraNodeEligibility(
 
 void UcxOutputQueueManager::removeTask(const std::string& taskId) {
   std::string taskIdStr{taskId};
-  auto queue =
-      queues_.withLock([&](auto& queues) -> std::shared_ptr<UcxOutputQueue> {
-        auto it = queues.find(taskIdStr);
-        if (it == queues.end()) {
-          // Already removed. Clear any stale "removed" state so the task ID
-          // can be reused.
+  std::shared_ptr<UcxOutputQueue> queue;
+  {
+    std::lock_guard<std::mutex> lifecycleLock(taskLifecycleMutex_);
+    queue = queues_.withLock(
+        [&](auto& queues) -> std::shared_ptr<UcxOutputQueue> {
+          auto it = queues.find(taskIdStr);
+          if (it == queues.end()) {
+            // Already removed. Clear any stale "removed" state so the task ID
+            // can be reused.
+            removedTasks_.withLock(
+                [&](auto& removed) { removed.erase(taskIdStr); });
+            return nullptr;
+          }
+          auto taskQueue = it->second;
+          queues.erase(it);
+          // Insert into removedTasks_ while still holding the queues_ lock
+          // to prevent getData() from seeing a gap between erase and insert,
+          // which would cause it to create a zombie placeholder queue.
           removedTasks_.withLock(
-              [&](auto& removed) { removed.erase(taskIdStr); });
-          return nullptr;
-        }
-        auto taskQueue = it->second;
-        queues.erase(it);
-        // Insert into removedTasks_ while still holding the queues_ lock
-        // to prevent getData() from seeing a gap between erase and insert,
-        // which would cause it to create a zombie placeholder queue.
-        removedTasks_.withLock(
-            [&](auto& removed) { removed.insert(taskIdStr); });
-        return taskQueue;
-      });
+              [&](auto& removed) { removed.insert(taskIdStr); });
+          return taskQueue;
+        });
+    // Publish cancellation before a subsequent initialization is allowed to
+    // clear it for a new task incarnation.
+    IntraNodeTransferRegistry::getInstance()->cancelTask(taskId);
+  }
   VLOG(2) << "[QUEUE-MGR] removeTask=" << taskId
           << " queueExists=" << (queue != nullptr);
   if (queue != nullptr) {
     queue->terminate();
   }
-  // Notify the intra-node registry so that any sources polling for this
-  // task get an atEnd result instead of spinning forever.
-  IntraNodeTransferRegistry::getInstance()->cancelTask(taskId);
 }
 
 std::shared_ptr<UcxOutputQueue> UcxOutputQueueManager::getQueueIfExists(
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
index 075619cd72f..d704b7d6f77 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
@@ -19,6 +19,7 @@
 #include <velox/exec/OutputBufferManager.h>
 #include <velox/exec/Task.h>
 #include <functional>
+#include <mutex>
 #include <string_view>
 #include <unordered_set>
 #include "velox/experimental/ucx-exchange/UcxQueues.h"
@@ -159,6 +160,11 @@ class UcxOutputQueueManager : public exec::OutputBufferManager {
   // that exceed the placeholder's undersized queues_ vector.
   folly::Synchronized<std::unordered_set<std::string>, std::mutex>
       removedTasks_;
+
+  // Serializes task initialization and removal with registry cancellation.
+  // A newly initialized queue must never be visible while a cancellation from
+  // its previous task incarnation is still active.
+  std::mutex taskLifecycleMutex_;
 };
 
 } // namespace facebook::velox::ucx_exchange
diff --git a/velox/experimental/ucx-exchange/UcxQueues.h b/velox/experimental/ucx-exchange/UcxQueues.h
index ba4e26dea0e..0f7104a2bff 100644
--- a/velox/experimental/ucx-exchange/UcxQueues.h
+++ b/velox/experimental/ucx-exchange/UcxQueues.h
@@ -173,6 +173,8 @@ class UcxOutputQueue : public std::enable_shared_from_this<UcxOutputQueue> {
   /// later properly, and eventually the callback fires.
   /// @return True, if initialization was successful, i.e. the queue wasn't
   /// already initialized.
+  /// @param eligibilityCallbacks Receives pending callbacks so the manager can
+  /// invoke them after clearing stale cancellation state for a reused task ID.
   bool initialize(
       std::shared_ptr<exec::Task> task,
       uint32_t numDestinations,

From fd6b53069b5328ca64a3f5f6761bdf251ef05a5e Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 14 Sep 2026 18:13:59 -0700
Subject: [PATCH 13/15] test(ucx-exchange): Cover concurrent task lifecycle
 race

---
 .../ucx-exchange/UcxOutputQueueManager.cpp    |  5 ++
 .../tests/UcxOutputQueueManagerTest.cpp       | 69 +++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
index faafcbbe699..af7ab0defe5 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
@@ -20,6 +20,7 @@
 #include <cudf/table/table.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
+#include "velox/common/testutil/TestValue.h"
 #include "velox/experimental/ucx-exchange/IntraNodeTransferRegistry.h"
 
 namespace facebook::velox::ucx_exchange {
@@ -67,6 +68,10 @@ void UcxOutputQueueManager::initializeTask(
         }
       }
     });
+    // Test hook for the initializeTask()/removeTask() lifecycle race.
+    common::testutil::TestValue::adjust(
+        "facebook::velox::ucx_exchange::UcxOutputQueueManager::initializeTask::queuePublished",
+        nullptr);
   }
 
   const bool canUseIntraNode =
diff --git a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
index c7f8cc8f06e..d8851d037a6 100644
--- a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
@@ -22,11 +22,15 @@
 #include <folly/synchronization/EventCount.h>
 #include <gtest/gtest.h>
 #include <rmm/device_buffer.hpp>
+#include <atomic>
+#include <chrono>
 #include <memory>
+#include <thread>
 #include <type_traits>
 #include <optional>
 #include <vector>
 #include "velox/common/memory/MemoryPool.h"
+#include "velox/common/testutil/TestValue.h"
 #include "velox/core/PlanNode.h"
 #include "velox/exec/OutputTransportRegistry.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
@@ -37,6 +41,7 @@ using namespace facebook::velox::ucx_exchange;
 using namespace facebook::velox;
 using namespace facebook::velox::exec;
 using namespace facebook::velox::core;
+using namespace facebook::velox::common::testutil;
 
 class UcxOutputQueueManagerTest : public testing::Test {
  protected:
@@ -737,6 +742,70 @@ TEST_F(
   queueManager_->removeTask(taskId);
 }
 
+DEBUG_ONLY_TEST_F(
+    UcxOutputQueueManagerTest,
+    initializationDoesNotClearConcurrentRemovalCancellation) {
+  const std::string taskId = "concurrentLifecycle";
+  queueManager_->removeTask(taskId);
+
+  auto task = createSourceTask(taskId, pool_, UcxTestData::kTestRowType);
+  std::atomic<bool> queuePublished{false};
+  std::atomic<bool> removeStarted{false};
+
+  TestValue::enable();
+  SCOPE_EXIT {
+    TestValue::disable();
+  };
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::ucx_exchange::UcxOutputQueueManager::initializeTask::queuePublished",
+      std::function<void(void*)>([&](void*) {
+        queuePublished.store(true, std::memory_order_release);
+        const auto deadline =
+            std::chrono::steady_clock::now() + std::chrono::seconds(5);
+        while (!removeStarted.load(std::memory_order_acquire) &&
+               std::chrono::steady_clock::now() < deadline) {
+          std::this_thread::sleep_for(std::chrono::milliseconds(1));
+        }
+        VELOX_CHECK(
+            removeStarted.load(std::memory_order_acquire),
+            "Concurrent removal did not start");
+      }));
+
+  std::thread initializer([&] {
+    queueManager_->initializeTask(
+        task,
+        core::PartitionedOutputNode::Kind::kPartitioned,
+        /*numDestinations=*/1,
+        /*numDrivers=*/1);
+  });
+
+  const auto deadline =
+      std::chrono::steady_clock::now() + std::chrono::seconds(5);
+  while (!queuePublished.load(std::memory_order_acquire) &&
+         std::chrono::steady_clock::now() < deadline) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(1));
+  }
+  if (!queuePublished.load(std::memory_order_acquire)) {
+    removeStarted.store(true, std::memory_order_release);
+    initializer.join();
+    FAIL() << "Initialization did not publish its queue";
+    return;
+  }
+
+  std::thread remover([&] {
+    removeStarted.store(true, std::memory_order_release);
+    queueManager_->removeTask(taskId);
+  });
+  initializer.join();
+  remover.join();
+
+  const auto result = IntraNodeTransferRegistry::getInstance()->poll(
+      IntraNodeTransferKey{
+          taskId, /*destination=*/0, /*sequenceNumber=*/0});
+  ASSERT_TRUE(result.has_value());
+  EXPECT_TRUE(result->atEnd);
+}
+
 // --- Broadcast tests ---
 
 // Basic broadcast: enqueue data, all destinations receive the same data.

From 17df2fbc5a346dde54681ce034da676f42e718a3 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 21 Sep 2026 09:22:44 -0700
Subject: [PATCH 14/15] fix(ucx-exchange): Delete results from queue
 incarnation

---
 .../ucx-exchange/UcxExchangeServer.cpp        |  8 ++-
 .../ucx-exchange/UcxExchangeServer.h          |  5 ++
 .../ucx-exchange/UcxOutputQueueManager.cpp    | 17 +++--
 .../ucx-exchange/UcxOutputQueueManager.h      | 10 ++-
 .../tests/UcxOutputQueueManagerTest.cpp       | 62 +++++++++++++++----
 5 files changed, 83 insertions(+), 19 deletions(-)

diff --git a/velox/experimental/ucx-exchange/UcxExchangeServer.cpp b/velox/experimental/ucx-exchange/UcxExchangeServer.cpp
index 0f5fff0b860..c2e2014d17c 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeServer.cpp
+++ b/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());
     }
diff --git a/velox/experimental/ucx-exchange/UcxExchangeServer.h b/velox/experimental/ucx-exchange/UcxExchangeServer.h
index a9db5db3051..17433bced3c 100644
--- a/velox/experimental/ucx-exchange/UcxExchangeServer.h
+++ b/velox/experimental/ucx-exchange/UcxExchangeServer.h
@@ -144,6 +144,11 @@ class UcxExchangeServer
   std::shared_ptr<ucxx::Request> metaRequest_{nullptr};
   std::shared_ptr<ucxx::Request> dataRequest_{nullptr};
 
+  // Queue incarnation returned by getData(). Final cleanup must target this
+  // exact queue rather than looking up by task ID, because task IDs can be
+  // reused while an old server is still draining its terminal callback.
+  std::shared_ptr<UcxOutputQueue> outputQueue_;
+
   // Completed UCXX requests are kept alive here to prevent use-after-free.
   // UCP's ucp_wireup_replay_pending_requests can fire callbacks on already-
   // completed requests; if the ucxx::Request has been freed, the callback
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
index af7ab0defe5..7849ca9cf3f 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.cpp
@@ -122,7 +122,15 @@ void UcxOutputQueueManager::deleteResults(
   }
 }
 
-void UcxOutputQueueManager::getData(
+void UcxOutputQueueManager::deleteResultsForQueue(
+    const std::shared_ptr<UcxOutputQueue>& queue,
+    int destination) {
+  if (queue != nullptr) {
+    queue->deleteResults(destination);
+  }
+}
+
+std::shared_ptr<UcxOutputQueue> UcxOutputQueueManager::getData(
     std::string_view taskId,
     int destination,
     UcxDataAvailableCallback notify) {
@@ -158,11 +166,12 @@ void UcxOutputQueueManager::getData(
   if (taskRemoved) {
     // Fire callback immediately with nullptr to signal end-of-stream.
     notify(nullptr, /*numRows=*/0, {});
-    return;
+    return nullptr;
   }
   // outside of lock. Queue must exist.
   // get the data or install the notify callback.
   outputQueue->getData(destination, notify);
+  return outputQueue;
 }
 
 bool UcxOutputQueueManager::canUseIntraNode(std::string_view taskId) {
@@ -212,8 +221,8 @@ void UcxOutputQueueManager::removeTask(const std::string& taskId) {
   std::shared_ptr<UcxOutputQueue> queue;
   {
     std::lock_guard<std::mutex> lifecycleLock(taskLifecycleMutex_);
-    queue = queues_.withLock(
-        [&](auto& queues) -> std::shared_ptr<UcxOutputQueue> {
+    queue =
+        queues_.withLock([&](auto& queues) -> std::shared_ptr<UcxOutputQueue> {
           auto it = queues.find(taskIdStr);
           if (it == queues.end()) {
             // Already removed. Clear any stale "removed" state so the task ID
diff --git a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
index d704b7d6f77..625243a8452 100644
--- a/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
+++ b/velox/experimental/ucx-exchange/UcxOutputQueueManager.h
@@ -95,6 +95,14 @@ class UcxOutputQueueManager : public exec::OutputBufferManager {
   /// @brief
   void deleteResults(std::string_view taskId, int destination);
 
+  /// Deletes results from a specific queue incarnation. This is used by
+  /// exchange servers that registered a getData() callback against a queue and
+  /// must not mutate a newer queue if the task ID is reused before the stale
+  /// server finishes.
+  void deleteResultsForQueue(
+      const std::shared_ptr<UcxOutputQueue>& queue,
+      int destination);
+
   /// @brief Asynchronously returns the head of the queue. If data is available,
   /// the callback function is triggered immediately and true is returned.
   /// Otherwise, the callback function is registered and called once data is
@@ -104,7 +112,7 @@ class UcxOutputQueueManager : public exec::OutputBufferManager {
   /// @param taskId The unique taskId.
   /// @param destination The destination.
   /// @param notify The callback function.
-  void getData(
+  std::shared_ptr<UcxOutputQueue> getData(
       std::string_view taskId,
       int destination,
       UcxDataAvailableCallback notify);
diff --git a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
index d8851d037a6..070b7183b5f 100644
--- a/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxOutputQueueManagerTest.cpp
@@ -25,9 +25,9 @@
 #include <atomic>
 #include <chrono>
 #include <memory>
+#include <optional>
 #include <thread>
 #include <type_traits>
-#include <optional>
 #include <vector>
 #include "velox/common/memory/MemoryPool.h"
 #include "velox/common/testutil/TestValue.h"
@@ -648,6 +648,50 @@ TEST_F(UcxOutputQueueManagerTest, callbackFiredOnTerminateAfterInit) {
   EXPECT_TRUE(callback1Nullptr);
 }
 
+TEST_F(UcxOutputQueueManagerTest, staleServerCleanupDoesNotDeleteReusedTask) {
+  const std::string taskId = "reusedTaskCleanup";
+  queueManager_->removeTask(taskId);
+
+  auto oldTask = initializeTask(
+      taskId,
+      /*numDestinations=*/1,
+      /*numDrivers=*/1,
+      /*cleanup=*/false);
+
+  bool oldServerReceivedEndMarker = false;
+  auto oldQueue = queueManager_->getData(
+      taskId, 0, receiveEndMarker(0, oldServerReceivedEndMarker));
+  ASSERT_NE(oldQueue, nullptr);
+  EXPECT_FALSE(oldServerReceivedEndMarker);
+
+  oldTask->requestAbort().wait();
+  queueManager_->removeTask(taskId);
+  EXPECT_TRUE(oldServerReceivedEndMarker);
+
+  auto newTask = initializeTask(
+      taskId,
+      /*numDestinations=*/1,
+      /*numDrivers=*/1,
+      /*cleanup=*/false);
+
+  bool newServerReceivedData = false;
+  auto newQueue =
+      queueManager_->getData(taskId, 0, receiveData(0, newServerReceivedData));
+  ASSERT_NE(newQueue, nullptr);
+  ASSERT_NE(newQueue, oldQueue);
+  EXPECT_FALSE(newServerReceivedData);
+
+  queueManager_->deleteResultsForQueue(oldQueue, 0);
+  EXPECT_FALSE(newServerReceivedData);
+
+  enqueue(taskId, 0, /*size=*/10);
+  EXPECT_TRUE(newServerReceivedData);
+
+  queueManager_->deleteResultsForQueue(newQueue, 0);
+  newTask->requestAbort().wait();
+  queueManager_->removeTask(taskId);
+}
+
 TEST_F(
     UcxOutputQueueManagerTest,
     intraNodeEligibilityWaitsForPartitionedTaskInitialization) {
@@ -719,14 +763,11 @@ TEST_F(
 
   std::optional<bool> eligibility;
   std::optional<IntraNodeTransferResult> registryPollResult;
-  queueManager_->notifyOnIntraNodeEligibility(
-      taskId,
-      [&](bool value) {
-        eligibility = value;
-        registryPollResult = IntraNodeTransferRegistry::getInstance()->poll(
-            IntraNodeTransferKey{
-                taskId, /*destination=*/0, /*sequenceNumber=*/0});
-      });
+  queueManager_->notifyOnIntraNodeEligibility(taskId, [&](bool value) {
+    eligibility = value;
+    registryPollResult = IntraNodeTransferRegistry::getInstance()->poll(
+        IntraNodeTransferKey{taskId, /*destination=*/0, /*sequenceNumber=*/0});
+  });
   EXPECT_FALSE(eligibility.has_value());
 
   auto task = initializeTask(
@@ -800,8 +841,7 @@ DEBUG_ONLY_TEST_F(
   remover.join();
 
   const auto result = IntraNodeTransferRegistry::getInstance()->poll(
-      IntraNodeTransferKey{
-          taskId, /*destination=*/0, /*sequenceNumber=*/0});
+      IntraNodeTransferKey{taskId, /*destination=*/0, /*sequenceNumber=*/0});
   ASSERT_TRUE(result.has_value());
   EXPECT_TRUE(result->atEnd);
 }

From 2acc113bef0dfe5a8b1bdf7cc41be31fafe180f2 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Mon, 21 Sep 2026 17:57:41 -0700
Subject: [PATCH 15/15] style(ucx-exchange): Apply clang-format to regression
 test

---
 .../ucx-exchange/tests/UcxExchangeTest.cpp       | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
index e7aeac95f48..3be3a005494 100644
--- a/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
+++ b/velox/experimental/ucx-exchange/tests/UcxExchangeTest.cpp
@@ -1116,13 +1116,12 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderDefersHandshakeWithBusyWorker) {
     config.intraNodeExchange = origIntraNode;
   };
 
-  auto metricSum = [](
-                       const folly::F14FastMap<std::string, RuntimeMetric>&
-                           stats,
-                       std::string_view name) {
-    auto it = stats.find(std::string{name});
-    return it == stats.end() ? 0 : it->second.sum;
-  };
+  auto metricSum =
+      [](const folly::F14FastMap<std::string, RuntimeMetric>& stats,
+         std::string_view name) {
+        auto it = stats.find(std::string{name});
+        return it == stats.end() ? 0 : it->second.sum;
+      };
 
   const std::string taskPrefix = getUniqueTaskPrefix();
   const std::string pollingSrcTaskId = taskPrefix + "pollingSrc";
@@ -1266,8 +1265,7 @@ TEST_P(UcxExchangeTest, partitionedPlaceholderDefersHandshakeWithBusyWorker) {
   EXPECT_EQ(
       metricSum(pollingSinkStats, "ucxExchangeSource.intraNodePackedColumns"),
       numChunks);
-  EXPECT_GT(
-      metricSum(pollingSinkStats, "ucxExchangeSource.intraNodeBytes"), 0);
+  EXPECT_GT(metricSum(pollingSinkStats, "ucxExchangeSource.intraNodeBytes"), 0);
   EXPECT_EQ(
       metricSum(pollingSinkStats, "ucxExchangeSource.remotePackedColumns"), 0);
   EXPECT_EQ(metricSum(pollingSinkStats, "ucxExchangeSource.remoteBytes"), 0);