tfbuilder: add blocking mode for rdma operations (optional, saves 0.5 CPU core)

Author: ironMann

src/TfBuilder/TfBuilderInputUCX.cxx

@@ -191,7 +191,8 @@ bool TfBuilderInputUCX::start()
 {
   // setting configuration options
   mThreadPoolSize = std::clamp(mConfig->getUInt64Param(UcxTfBuilderThreadPoolSizeKey, UcxTfBuilderThreadPoolSizeDefault), std::size_t(1), std::size_t(256));
-  IDDLOG("TfBuilderInputUCX: Configuration loaded. thread_pool={}", mThreadPoolSize);
+  mRdmaPollingWait =mConfig->getBoolParam(UcxPollForRDMACompletionKey, UcxPollForRDMACompletionDefault);
+  IDDLOG("TfBuilderInputUCX: Configuration loaded. thread_pool={} polling={}", mThreadPoolSize, mRdmaPollingWait);
 
   auto &lConfStatus = mConfig->status();
 
@@ -628,9 +629,18 @@ void TfBuilderInputUCX::DataHandlerThread(const unsigned pThreadIdx)
           ucx::io::get(lConn->ucp_ep, lTxgUcxPtr, lStfTxg.len(), lStfTxg.start(), lConn->mRemoteKeys[lStfMeta.data_regions(lStfTxg.region()).region_rkey()], &lRmaReqSem);
         }
       }
+
       // wait for final completion
-      if (!ucx::io::ucp_wait_poll(lConn->mWorker, lRmaReqSem)) {
-        break;
+      if (mRdmaPollingWait) {
+        // polling
+        if (!lRmaReqSem.wait_poll(lConn->mWorker)) {
+          break;
+        }
+      } else {
+        // blocking
+        if (!lRmaReqSem.wait(lConn->mWorker)) {
+          break;
+        }
       }
 
       // notify StfSender we completed (use mapped scratch memory)

src/TfBuilder/TfBuilderInputUCX.h

@@ -21,6 +21,7 @@
 
 #include <SubTimeFrameDataModel.h>
 #include <ConcurrentQueue.h>
+#include <DataDistributionOptions.h>
 
 #include <UCXUtilities.h>
 #include <ucp/api/ucp.h>
@@ -171,6 +172,7 @@ class TfBuilderInputUCX
     std::unordered_map<std::string, std::size_t> mStfSenderToWorkerMap;
     std::unordered_map<std::string, std::shared_ptr<ConcurrentQueue<StfMetaRdmaInfo> > >  mStfMetaWorkerQueues;
   std::vector<std::thread> mThreadPool;
+  bool mRdmaPollingWait = UcxPollForRDMACompletionDefault;
 
   // STF postprocess threads
   ConcurrentQueue<StfMetaRdmaInfo> mStfPostprocessQueue;

src/common/discovery/DataDistributionOptions.h

@@ -81,7 +81,7 @@ static constexpr std::uint64_t UcxStfSenderThreadPoolSizeDefault = 8;
 ////////////////////////////////////////////////////////////////////////////////
 
 // Define maximum number of concurrent STF transfers
-// The value should not be much grater than UcxTfBuilderThreadPoolSize as the requests will will not be processed immediately.
+// The value should be greater than max number of FLPs.
 static constexpr std::string_view MaxNumStfTransfersKey = "MaxNumStfTransfers";
 static constexpr std::uint64_t MaxNumStfTransferDefault = 300;
 
@@ -95,14 +95,13 @@ static constexpr std::uint64_t StfSenderGrpcThreadPoolSizeDefault = 8;
 
 
 /// UCX transport
-// Size of receiver treadpool. Default 0 (number of cpu cores)
+// Size of receiver treadpool. Default 1, works best. Should not be set over 2, to avoid congestion on the receiver.
 static constexpr std::string_view UcxTfBuilderThreadPoolSizeKey = "UcxTfBuilderThreadPoolSize";
 static constexpr std::uint64_t UcxTfBuilderThreadPoolSizeDefault = 1;
 
-// Number of rma_get operation in flight, per ucx thread
-// NOTE: deprecated in DD 1.4
-// static constexpr std::string_view UcxNumConcurrentRmaGetOpsKey = "UcxNumConcurrentRmaGetOps";
-// static constexpr std::uint64_t UcxNumConcurrentRmaGetOpsDefault = 32;
+// Use polling or blocking waiting method for RDMA completion.
+static constexpr std::string_view UcxPollForRDMACompletionKey = "UcxPollForRDMACompletion";
+static constexpr bool UcxPollForRDMACompletionDefault = false;
 
 
 ////////////////////////////////////////////////////////////////////////////////

src/common/ucxtools/UCXSendRecv.h

@@ -127,6 +127,49 @@ struct dd_ucp_multi_req_v2 {
     }
     return true;
   }
+
+  inline
+  bool wait(dd_ucp_worker &pDDCtx) {
+    for (;;) {
+      // check if request is done
+      if (done()) {
+        return true;
+      } else if (ucp_worker_progress(pDDCtx.ucp_worker)) {
+        continue;
+      }
+
+      // block on the worker
+      auto status = ucp_worker_arm(pDDCtx.ucp_worker);
+
+      if (UCS_OK == status) {
+        int epoll_ret;
+        do {
+          epoll_ret = epoll_wait(pDDCtx.epoll_fd, &pDDCtx.ev, 1, 100);
+        } while ((epoll_ret == -1) && (errno == EINTR || errno == EAGAIN));
+
+        if (epoll_ret == -1) {
+          EDDLOG("Failed ucp_advance epoll. errno={}", errno);
+          return done();
+        }
+      } else if (UCS_ERR_BUSY == status) {
+        continue; // could not arm, recheck the request
+      }
+        // epoll returned or timeout, recheck the request
+    }
+    return done();
+  }
+
+  inline
+  bool wait_poll(dd_ucp_worker &pDDCtx) const {
+    for (;;) {
+      // check if request is done
+      if (done()) {
+        return true;
+      }
+      while (ucp_worker_progress(pDDCtx.ucp_worker) > 0) { }
+    }
+    return true;
+  }
 };