UCX/BACKEND: Add worker_id selection support (#938)

Signed-off-by: Michal Shalev <[email protected]>

Author: michal-shalev

src/api/cpp/backend/backend_engine.h

@@ -177,7 +177,9 @@ class nixlBackendEngine {
 
         // Initialize a signal for GPU transfer using memory handle from descriptor
         virtual nixl_status_t
-        prepGpuSignal(const nixlBackendMD &meta, void *signal) const {
+        prepGpuSignal(const nixlBackendMD &meta,
+                      void *signal,
+                      const nixl_opt_b_args_t *opt_args = nullptr) const {
             return NIXL_ERR_NOT_SUPPORTED;
         }
 

src/plugins/ucx/ucx_backend.cpp

@@ -1388,6 +1388,35 @@ nixlUcxEngine::getWorkerId() const {
     return it->second;
 }
 
+std::optional<size_t>
+nixlUcxEngine::getWorkerIdFromOptArgs(const nixl_opt_b_args_t *opt_args) const noexcept {
+    if (!opt_args || opt_args->customParam.empty()) {
+        return std::nullopt;
+    }
+
+    constexpr std::string_view worker_id_key = "worker_id=";
+    size_t pos = opt_args->customParam.find(worker_id_key);
+    if (pos == std::string::npos) {
+        return std::nullopt;
+    }
+
+    try {
+        size_t worker_id = std::stoull(opt_args->customParam.substr(pos + worker_id_key.length()));
+
+        if (worker_id >= getSharedWorkersSize()) {
+            NIXL_WARN << "Invalid worker_id " << worker_id << " (must be < "
+                      << getSharedWorkersSize() << ")";
+            return std::nullopt;
+        }
+
+        return worker_id;
+    }
+    catch (const std::exception &e) {
+        NIXL_WARN << "Failed to parse worker_id from customParam: " << e.what();
+        return std::nullopt;
+    }
+}
+
 nixl_status_t nixlUcxEngine::prepXfer (const nixl_xfer_op_t &operation,
                                        const nixl_meta_dlist_t &local,
                                        const nixl_meta_dlist_t &remote,
@@ -1401,7 +1430,8 @@ nixl_status_t nixlUcxEngine::prepXfer (const nixl_xfer_op_t &operation,
     }
 
     /* TODO: try to get from a pool first */
-    size_t worker_id = getWorkerId();
+    const auto opt_worker_id = getWorkerIdFromOptArgs(opt_args);
+    size_t worker_id = opt_worker_id.value_or(getWorkerId());
     auto *ucx_handle = new nixlUcxBackendH(getWorker(worker_id).get(), worker_id);
 
     handle = ucx_handle;
@@ -1659,6 +1689,8 @@ nixlUcxEngine::createGpuXferReq(const nixlBackendReqH &req_hndl,
 
     try {
         gpu_req_hndl = nixl::ucx::createGpuXferReq(*ep, local_mems, remote_rkeys, remote_addrs);
+        NIXL_TRACE << "Created device memory list: ep=" << ep->getEp() << " handle=" << gpu_req_hndl
+                   << " worker_id=" << workerId << " num_elements=" << local_mems.size();
         return NIXL_SUCCESS;
     }
     catch (const std::exception &e) {
@@ -1690,10 +1722,19 @@ nixlUcxEngine::getGpuSignalSize(size_t &signal_size) const {
 }
 
 nixl_status_t
-nixlUcxEngine::prepGpuSignal(const nixlBackendMD &meta, void *signal) const {
+nixlUcxEngine::prepGpuSignal(const nixlBackendMD &meta,
+                             void *signal,
+                             const nixl_opt_b_args_t *opt_args) const {
     try {
         auto *ucx_meta = static_cast<const nixlUcxPrivateMetadata *>(&meta);
-        getWorker(getWorkerId())->prepGpuSignal(ucx_meta->mem, signal);
+
+        const auto opt_worker_id = getWorkerIdFromOptArgs(opt_args);
+        if (opt_worker_id) {
+            getWorker(*opt_worker_id)->prepGpuSignal(ucx_meta->mem, signal);
+        } else {
+            getWorker(getWorkerId())->prepGpuSignal(ucx_meta->mem, signal);
+        }
+
         return NIXL_SUCCESS;
     }
     catch (const std::exception &e) {

src/plugins/ucx/ucx_backend.h

@@ -204,7 +204,9 @@ class nixlUcxEngine : public nixlBackendEngine {
     getGpuSignalSize(size_t &signal_size) const override;
 
     nixl_status_t
-    prepGpuSignal(const nixlBackendMD &meta, void *signal) const override;
+    prepGpuSignal(const nixlBackendMD &meta,
+                  void *signal,
+                  const nixl_opt_b_args_t *opt_args = nullptr) const override;
 
     int
     progress();
@@ -218,6 +220,11 @@ class nixlUcxEngine : public nixlBackendEngine {
     nixl_status_t
     checkConn(const std::string &remote_agent);
 
+private:
+    // Helper to extract worker_id from opt_args->customParam or nullopt if not found
+    [[nodiscard]] std::optional<size_t>
+    getWorkerIdFromOptArgs(const nixl_opt_b_args_t *opt_args) const noexcept;
+
 protected:
     const std::vector<std::unique_ptr<nixlUcxWorker>> &
     getWorkers() const {

src/utils/ucx/gpu_xfer_req_h.cpp

@@ -81,7 +81,8 @@ createGpuXferReq(const nixlUcxEp &ep,
                                  ucs_status_string(ucs_status));
     }
 
-    NIXL_DEBUG << "Created device memory list handle with " << local_mems.size() << " elements";
+    NIXL_DEBUG << "Created device memory list: ep=" << ep.getEp() << " handle=" << ucx_handle
+               << " num_elements=" << local_mems.size();
     return reinterpret_cast<nixlGpuXferReqH>(ucx_handle);
 }
 

test/gtest/device_api/single_write_test.cu

@@ -149,7 +149,7 @@ protected:
         nixl_b_params_t params;
 
         if (getBackendName() == "UCX") {
-            params["num_workers"] = "2";
+            params["num_workers"] = std::to_string(numWorkers);
         }
 
         return params;
@@ -194,21 +194,50 @@ protected:
         agent.registerMem(reg_list);
     }
 
+    // TODO: remove this function once a blocking CreateGpuXferReq is implemented
     void
-    completeWireup(size_t from_agent, size_t to_agent) {
-        nixl_notifs_t notifs;
-        nixl_status_t status = getAgent(from_agent).genNotif(getAgentName(to_agent), NOTIF_MSG);
-        ASSERT_EQ(status, NIXL_SUCCESS) << "Failed to complete wireup";
-
-        do {
-            nixl_status_t ret = getAgent(to_agent).getNotifs(notifs);
-            ASSERT_EQ(ret, NIXL_SUCCESS) << "Failed to get notifications during wireup";
-            std::this_thread::sleep_for(std::chrono::milliseconds(10));
-        } while (notifs.size() == 0);
+    completeWireup(size_t from_agent, size_t to_agent,
+                   const std::vector<MemBuffer> &wireup_src,
+                   const std::vector<MemBuffer> &wireup_dst) {
+        nixl_opt_args_t wireup_params;
+
+        for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+            wireup_params.customParam = "worker_id=" + std::to_string(worker_id);
+
+            nixlXferReqH *wireup_req;
+            nixl_status_t status = getAgent(from_agent)
+                                       .createXferReq(NIXL_WRITE,
+                                                      makeDescList<nixlBasicDesc>(wireup_src, VRAM_SEG),
+                                                      makeDescList<nixlBasicDesc>(wireup_dst, VRAM_SEG),
+                                                      getAgentName(to_agent),
+                                                      wireup_req,
+                                                      &wireup_params);
+
+            ASSERT_EQ(status, NIXL_SUCCESS) << "Failed to create wireup request for worker " << worker_id;
+
+            status = getAgent(from_agent).postXferReq(wireup_req);
+            ASSERT_TRUE(status == NIXL_SUCCESS || status == NIXL_IN_PROG)
+                << "Failed to post wireup for worker " << worker_id;
+
+            nixl_status_t xfer_status;
+            do {
+                xfer_status = getAgent(from_agent).getXferStatus(wireup_req);
+                std::this_thread::sleep_for(std::chrono::milliseconds(1));
+            } while (xfer_status == NIXL_IN_PROG);
+
+            ASSERT_EQ(xfer_status, NIXL_SUCCESS) << "Warmup failed for worker " << worker_id;
+
+            status = getAgent(from_agent).releaseXferReq(wireup_req);
+            ASSERT_EQ(status, NIXL_SUCCESS);
+        }
     }
 
     void
     exchangeMD(size_t from_agent, size_t to_agent) {
+        std::vector<MemBuffer> wireup_src, wireup_dst;
+        createRegisteredMem(getAgent(from_agent), 64, 1, VRAM_SEG, wireup_src);
+        createRegisteredMem(getAgent(to_agent), 64, 1, VRAM_SEG, wireup_dst);
+
         for (size_t i = 0; i < agents.size(); i++) {
             nixl_blob_t md;
             nixl_status_t status = agents[i]->getLocalMD(md);
@@ -223,7 +252,7 @@ protected:
             }
         }
 
-        completeWireup(from_agent, to_agent);
+        completeWireup(from_agent, to_agent, wireup_src, wireup_dst);
     }
 
     void
@@ -316,6 +345,7 @@ protected:
 protected:
     static constexpr size_t SENDER_AGENT = 0;
     static constexpr size_t RECEIVER_AGENT = 1;
+    static constexpr size_t numWorkers = 32;
 
 private:
     static constexpr uint64_t DEV_ID = 0;
@@ -572,6 +602,92 @@ TEST_P(SingleWriteTest, VariableSizeTest) {
     }
 }
 
+TEST_P(SingleWriteTest, MultipleWorkersTest) {
+    constexpr size_t size = 4 * 1024;
+    constexpr size_t num_iters = 100;
+    constexpr unsigned index = 0;
+    constexpr bool is_no_delay = true;
+    constexpr nixl_mem_t mem_type = VRAM_SEG;
+    constexpr size_t num_threads = 32;
+
+    std::vector<std::vector<MemBuffer>> src_buffers(numWorkers);
+    std::vector<std::vector<MemBuffer>> dst_buffers(numWorkers);
+    std::vector<std::vector<uint32_t>> patterns(numWorkers);
+
+    for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+        createRegisteredMem(getAgent(SENDER_AGENT), size, 1, mem_type, src_buffers[worker_id]);
+        createRegisteredMem(getAgent(RECEIVER_AGENT), size, 1, mem_type, dst_buffers[worker_id]);
+
+        constexpr size_t num_elements = size / sizeof(uint32_t);
+        patterns[worker_id].resize(num_elements);
+        for (size_t i = 0; i < num_elements; i++) {
+            patterns[worker_id][i] = 0xDEAD0000 | worker_id;
+        }
+        cudaMemcpy(static_cast<void *>(src_buffers[worker_id][0]), patterns[worker_id].data(),
+                   size, cudaMemcpyHostToDevice);
+    }
+
+    exchangeMD(SENDER_AGENT, RECEIVER_AGENT);
+
+    nixl_opt_args_t extra_params = {};
+    extra_params.hasNotif = true;
+    extra_params.notifMsg = NOTIF_MSG;
+
+    std::vector<nixlXferReqH *> xfer_reqs(numWorkers);
+    std::vector<nixlGpuXferReqH> gpu_req_hndls(numWorkers);
+
+    for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+        extra_params.customParam = "worker_id=" + std::to_string(worker_id);
+
+        nixl_status_t status = getAgent(SENDER_AGENT)
+                                   .createXferReq(NIXL_WRITE,
+                                                  makeDescList<nixlBasicDesc>(src_buffers[worker_id], mem_type),
+                                                  makeDescList<nixlBasicDesc>(dst_buffers[worker_id], mem_type),
+                                                  getAgentName(RECEIVER_AGENT),
+                                                  xfer_reqs[worker_id],
+                                                  &extra_params);
+
+        ASSERT_EQ(status, NIXL_SUCCESS) << "Failed to create xfer request for worker " << worker_id;
+
+        status = getAgent(SENDER_AGENT).createGpuXferReq(*xfer_reqs[worker_id], gpu_req_hndls[worker_id]);
+        ASSERT_EQ(status, NIXL_SUCCESS) << "Failed to create GPU xfer request for worker " << worker_id;
+    }
+
+    unsigned long long *start_time_ptr;
+    unsigned long long *end_time_ptr;
+    initTimingPublic(&start_time_ptr, &end_time_ptr);
+
+    for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+        nixl_status_t status = dispatchLaunchSingleWriteTest(GetParam(), num_threads,
+                                                             gpu_req_hndls[worker_id], index,
+                                                             0, 0, size, num_iters, is_no_delay,
+                                                             start_time_ptr, end_time_ptr);
+        ASSERT_EQ(status, NIXL_SUCCESS) << "Kernel launch failed for worker " << worker_id;
+    }
+
+    for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+        std::vector<uint32_t> received(size / sizeof(uint32_t));
+        cudaMemcpy(received.data(), static_cast<void *>(dst_buffers[worker_id][0]),
+                   size, cudaMemcpyDeviceToHost);
+
+        EXPECT_EQ(received, patterns[worker_id])
+            << "Worker " << worker_id << " full buffer verification failed";
+    }
+
+    Logger() << "MultipleWorkers test: " << numWorkers << " workers with explicit selection verified";
+
+    cudaFree(start_time_ptr);
+    cudaFree(end_time_ptr);
+
+    for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) {
+        getAgent(SENDER_AGENT).releaseGpuXferReq(gpu_req_hndls[worker_id]);
+        nixl_status_t status = getAgent(SENDER_AGENT).releaseXferReq(xfer_reqs[worker_id]);
+        EXPECT_EQ(status, NIXL_SUCCESS);
+    }
+
+    invalidateMD();
+}
+
 } // namespace gtest::nixl::gpu::single_write
 
 using gtest::nixl::gpu::single_write::SingleWriteTest;