Refactoring: derive nixlUcxThreadPoolEngine from nixlUcxThreadEngine

src/plugins/ucx/ucx_backend.cpp

@@ -261,12 +261,6 @@ class nixlUcxThread {
         return tls;
     }
 
-    static bool
-    isProgressThread(const nixlUcxEngine *engine) noexcept {
-        nixlUcxThread *thread = tlsThread();
-        return thread && thread->engine_ == engine;
-    }
-
     friend std::ostream &
     operator<<(std::ostream &os, const nixlUcxThread &thread) {
         return os << "thread " << &thread << "{engine: " << thread.engine_ << ", worker_ids: ["
@@ -368,22 +362,29 @@ class nixlUcxSharedThread : public nixlUcxThread {
     std::vector<pollfd> pollFds_;
 };
 
-nixlUcxThreadEngine::nixlUcxThreadEngine(const nixlBackendInitParams &init_params)
-    : nixlUcxEngine(init_params) {
+nixlUcxThreadEngine::nixlUcxThreadEngine(const nixlBackendInitParams &init_params,
+                                         size_t num_dedicated_workers)
+    : nixlUcxEngine(init_params, num_dedicated_workers) {
+    if (!init_params.enableProgTh) {
+        return;
+    }
+
     if (!nixlUcxMtLevelIsSupported(nixl::ucx::mt_mode_t::WORKER)) {
         throw std::invalid_argument("UCX library does not support multi-threading");
     }
 
-    size_t num_workers = getWorkers().size();
-    thread_ = std::make_unique<nixlUcxSharedThread>(this, num_workers, init_params.pthrDelay);
-    for (size_t i = 0; i < num_workers; i++) {
-        thread_->addWorker(getWorkers()[i].get(), i);
+    const size_t shared_count = getSharedWorkers().size();
+    thread_ = std::make_unique<nixlUcxSharedThread>(this, shared_count, init_params.pthrDelay);
+    for (size_t i = 0; i < shared_count; i++) {
+        thread_->addWorker(getSharedWorkers()[i].get(), i);
     }
     thread_->start();
 }
 
 nixlUcxThreadEngine::~nixlUcxThreadEngine() {
-    thread_->join();
+    if (thread_) {
+        thread_->join();
+    }
 }
 
 void
@@ -398,6 +399,10 @@ nixlUcxThreadEngine::getNotifs(notif_list_t &notif_list) {
         return NIXL_ERR_INVALID_PARAM;
     }
 
+    if (!thread_) {
+        progressLoop();
+    }
+
     const std::lock_guard lock(notifMutex_);
     notifList_.swap(notif_list);
     return NIXL_SUCCESS;
@@ -643,42 +648,24 @@ class nixlUcxDedicatedThread : public nixlUcxThread {
     std::vector<nixlUcxChunkBackendReqH *> requests_;
 };
 
-nixlUcxThreadPoolEngine::nixlUcxThreadPoolEngine(const nixlBackendInitParams &init_params)
-    : nixlUcxEngine(init_params) {
-    const size_t num_threads =
-        nixl::getBackendParamDefaulted(init_params.customParams, "num_threads", 0u);
-    numSharedWorkers_ = getWorkers().size() - num_threads;
-    NIXL_ASSERT(numSharedWorkers_ > 0);
-
+nixlUcxThreadPoolEngine::nixlUcxThreadPoolEngine(const nixlBackendInitParams &init_params,
+                                                 size_t num_threads)
+    : nixlUcxThreadEngine(init_params, num_threads) {
     splitBatchSize_ =
         nixl::getBackendParamDefaulted(init_params.customParams, "split_batch_size", 1024u);
 
-    if (init_params.enableProgTh) {
-        sharedThread_ =
-            std::make_unique<nixlUcxSharedThread>(this, numSharedWorkers_, init_params.pthrDelay);
-        for (size_t i = 0; i < numSharedWorkers_; i++) {
-            sharedThread_->addWorker(getWorkers()[i].get(), i);
-        }
-        sharedThread_->start();
-    }
-
-    if (num_threads > 0) {
-        io_.reset(new asio::io_context());
-        dedicatedThreads_.reserve(num_threads);
-        for (size_t i = 0; i < num_threads; ++i) {
-            size_t worker_id = numSharedWorkers_ + i;
-            dedicatedThreads_.emplace_back(std::make_unique<nixlUcxDedicatedThread>(this, *io_));
-            dedicatedThreads_.back()->addWorker(getWorker(worker_id).get(), worker_id);
-            dedicatedThreads_.back()->start();
-        }
+    const auto dedicated_workers = getDedicatedWorkers();
+    io_.reset(new asio::io_context());
+    dedicatedThreads_.reserve(dedicated_workers.size());
+    for (size_t i = 0; i < dedicated_workers.size(); ++i) {
+        const size_t worker_id = getSharedWorkersSize() + i;
+        dedicatedThreads_.emplace_back(std::make_unique<nixlUcxDedicatedThread>(this, *io_));
+        dedicatedThreads_.back()->addWorker(dedicated_workers[i].get(), worker_id);
+        dedicatedThreads_.back()->start();
     }
 }
 
 nixlUcxThreadPoolEngine::~nixlUcxThreadPoolEngine() {
-    if (sharedThread_) {
-        sharedThread_->join();
-    }
-
     if (io_) {
         io_->stop();
         for (auto &thread : dedicatedThreads_) {
@@ -702,9 +689,9 @@ nixlUcxThreadPoolEngine::prepXfer(const nixl_xfer_op_t &operation,
     size_t chunk_size = std::max(batch_size / dedicatedThreads_.size(), splitBatchSize_);
     size_t num_chunks = (batch_size + chunk_size - 1) / chunk_size;
 
-    size_t worker_id = getWorkerId();
+    size_t worker_id = getSharedWorkerId();
     const auto comp_handle = new nixlUcxCompositeBackendReqH(
-        getWorker(worker_id).get(), worker_id, chunk_size, num_chunks);
+        getSharedWorker(worker_id).get(), worker_id, chunk_size, num_chunks);
     NIXL_TRACE << "created " << *comp_handle;
     handle = comp_handle;
     return NIXL_SUCCESS;
@@ -766,27 +753,6 @@ nixlUcxThreadPoolEngine::sendXferRange(const nixl_xfer_op_t &operation,
     return status.load();
 }
 
-void
-nixlUcxThreadPoolEngine::appendNotif(std::string &&remote_name, std::string &&msg) {
-    const std::lock_guard lock(notifMutex_);
-    notifList_.emplace_back(std::move(remote_name), std::move(msg));
-}
-
-nixl_status_t
-nixlUcxThreadPoolEngine::getNotifs(notif_list_t &notif_list) {
-    if (!notif_list.empty()) {
-        return NIXL_ERR_INVALID_PARAM;
-    }
-
-    if (!sharedThread_) {
-        progressLoop();
-    }
-
-    const std::lock_guard lock(notifMutex_);
-    notifList_.swap(notif_list);
-    return NIXL_SUCCESS;
-}
-
 /****************************************
  * Constructor/Destructor
  *****************************************/
@@ -797,7 +763,7 @@ nixlUcxEngine::create(const nixlBackendInitParams &init_params) {
     const size_t num_threads =
         nixl::getBackendParamDefaulted(init_params.customParams, "num_threads", 0u);
     if (num_threads > 0) {
-        engine = new nixlUcxThreadPoolEngine(init_params);
+        engine = new nixlUcxThreadPoolEngine(init_params, num_threads);
     } else if (init_params.enableProgTh) {
         engine = new nixlUcxThreadEngine(init_params);
     } else {
@@ -806,7 +772,7 @@ nixlUcxEngine::create(const nixlBackendInitParams &init_params) {
     return std::unique_ptr<nixlUcxEngine>(engine);
 }
 
-nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params)
+nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params, size_t num_dedicated_workers)
     : nixlBackendEngine(&init_params),
       sharedWorkerIndex_(1) {
     std::vector<std::string> devs; /* Empty vector */
@@ -817,14 +783,14 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params)
     }
 
     size_t num_workers = nixl::getBackendParamDefaulted(custom_params, "num_workers", 1u);
-    const size_t num_threads = nixl::getBackendParamDefaulted(custom_params, "num_threads", 0u);
+    if (num_workers <= num_dedicated_workers) {
+        num_workers = num_dedicated_workers + 1;
+    }
+    numSharedWorkers_ = num_workers - num_dedicated_workers;
+
     const size_t num_device_channels =
         nixl::getBackendParamDefaulted(custom_params, "ucx_num_device_channels", 4u);
 
-    if (num_workers <= num_threads) {
-        /* There must be at least one shared worker */
-        num_workers = num_threads + 1;
-    }
 
     ucp_err_handling_mode_t err_handling_mode = UCP_ERR_HANDLING_MODE_PEER;
     if (const auto opt = nixl::getBackendParamOptional<std::string>(
@@ -846,13 +812,15 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params)
 
     uc->warnAboutHardwareSupportMismatch();
 
+    workers_.reserve(num_workers);
     for (size_t i = 0; i < num_workers; i++) {
-        uws.emplace_back(std::make_unique<nixlUcxWorker>(*uc, err_handling_mode, ep_close_flags));
+        workers_.emplace_back(
+            std::make_unique<nixlUcxWorker>(*uc, err_handling_mode, ep_close_flags));
     }
 
-    auto &uw = uws.front();
-    workerAddr = uw->epAddr();
-    uw->regAmCallback(nixl::ucx::am_cb_op_t::NOTIF_STR, notifAmCb, this);
+    auto &worker = workers_.front();
+    workerAddr = worker->epAddr();
+    worker->regAmCallback(nixl::ucx::am_cb_op_t::NOTIF_STR, notifAmCb, this);
 }
 
 nixl_mem_list_t nixlUcxEngine::getSupportedMems () const {
@@ -877,10 +845,6 @@ nixlUcxEngine::~nixlUcxEngine() {
  * Connection management
 *****************************************/
 
-nixl_status_t nixlUcxEngine::checkConn(const std::string &remote_agent) {
-    return remoteConnMap.count(remote_agent) ? NIXL_SUCCESS : NIXL_ERR_NOT_FOUND;
-}
-
 nixl_status_t nixlUcxEngine::getConnInfo(std::string &str) const {
     str = workerAddr;
     return NIXL_SUCCESS;
@@ -919,12 +883,12 @@ nixl_status_t nixlUcxEngine::loadRemoteConnInfo (const std::string &remote_agent
 
     nixlSerDes::_stringToBytes(addr.data(), remote_conn_info, size);
     std::shared_ptr<nixlUcxConnection> conn = std::make_shared<nixlUcxConnection>();
-    for (auto &uw : uws) {
-        std::unique_ptr<nixlUcxEp> result = uw->connect(addr.data(), size);
-        if (!result) {
+    for (const auto &uw : workers_) {
+        std::unique_ptr<nixlUcxEp> ep = uw->connect(addr.data(), size);
+        if (!ep) {
             return NIXL_ERR_BACKEND;
         }
-        conn->eps.push_back(std::move(result));
+        conn->eps.push_back(std::move(ep));
     }
 
     remoteConnMap.insert({remote_agent, conn});
@@ -1004,7 +968,7 @@ nixlUcxEngine::internalMDHelper (const nixl_blob_t &blob,
         }
         // nixlSerDes::_stringToBytes() was used to "unpack" blob here.
         output = new nixlUcxPublicMetadata(
-            it->second, makePublicMetadataRkeys(it->second, uws.size(), blob.data()));
+            it->second, makePublicMetadataRkeys(it->second, workers_.size(), blob.data()));
         return NIXL_SUCCESS;
     }
     catch (const std::runtime_error &e) {
@@ -1043,7 +1007,7 @@ nixl_status_t nixlUcxEngine::unloadMD (nixlBackendMD* input) {
 *****************************************/
 
 size_t
-nixlUcxEngine::getWorkerId(const nixl_opt_b_args_t *opt_args) const noexcept {
+nixlUcxEngine::getSharedWorkerId(const nixl_opt_b_args_t *opt_args) const noexcept {
     if (opt_args) {
         const std::optional<size_t> worker_id = getWorkerIdFromOptArgs(*opt_args);
         if (worker_id) {
@@ -1098,9 +1062,9 @@ nixl_status_t nixlUcxEngine::prepXfer (const nixl_xfer_op_t &operation,
         return NIXL_ERR_INVALID_PARAM;
     }
 
-    const size_t worker_id = getWorkerId(opt_args);
+    const size_t worker_id = getSharedWorkerId(opt_args);
     /* TODO: try to get from a pool first */
-    handle = new nixlUcxBackendReqH(getWorker(worker_id).get(), worker_id);
+    handle = new nixlUcxBackendReqH(getSharedWorker(worker_id).get(), worker_id);
 
     return NIXL_SUCCESS;
 }
@@ -1358,7 +1322,7 @@ unsigned
 nixlUcxEngine::progress() {
     // TODO: add listen for connection handling if necessary
     unsigned ret = 0;
-    for (auto &uw : uws) {
+    for (const auto &uw : getSharedWorkers()) {
         ret += uw->progress();
     }
     return ret;
@@ -1460,7 +1424,7 @@ nixlUcxEngine::genNotif(const std::string &remote_agent, const std::string &msg)
         return NIXL_ERR_NOT_FOUND;
     }
 
-    const nixl_status_t ret = notifSendPriv(remote_agent, msg, conn->getEp(getWorkerId()));
+    const nixl_status_t ret = notifSendPriv(remote_agent, msg, conn->getEp(getSharedWorkerId()));
     if (ret == NIXL_IN_PROG) {
         return NIXL_SUCCESS;
     }
@@ -1471,9 +1435,9 @@ nixl_status_t
 nixlUcxEngine::prepMemView(const nixl_remote_meta_dlist_t &dlist,
                            nixlMemViewH &mvh,
                            const nixl_opt_b_args_t *opt_args) const {
-    const size_t worker_id = getWorkerId(opt_args);
+    const size_t worker_id = getSharedWorkerId(opt_args);
     try {
-        mvh = nixl::ucx::createMemList(dlist, worker_id, *getWorker(worker_id));
+        mvh = nixl::ucx::createMemList(dlist, worker_id, *getSharedWorker(worker_id));
         return NIXL_SUCCESS;
     }
     catch (const std::exception &e) {
@@ -1486,9 +1450,9 @@ nixl_status_t
 nixlUcxEngine::prepMemView(const nixl_meta_dlist_t &dlist,
                            nixlMemViewH &mvh,
                            const nixl_opt_b_args_t *opt_args) const {
-    const size_t worker_id = getWorkerId(opt_args);
+    const size_t worker_id = getSharedWorkerId(opt_args);
     try {
-        mvh = nixl::ucx::createMemList(dlist, *getWorker(worker_id));
+        mvh = nixl::ucx::createMemList(dlist, *getSharedWorker(worker_id));
         return NIXL_SUCCESS;
     }
     catch (const std::exception &e) {