[TE/TENT] Enable UB Transport in TENT on Kunpeng SuperNode (Phase 3)

mooncake-common/FindUrma.cmake

@@ -1,20 +1,30 @@
 include(FetchContent)
 
-# UMDK 头文件库
-FetchContent_Declare(
-        urma
-        GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
-        GIT_TAG        v25.12.0.B081
-)
+# Allow callers to supply headers without downloading UMDK, e.g.: cmake
+# -DURMA_INCLUDE_DIR=/usr/include ... cmake
+# -DFETCHCONTENT_SOURCE_DIR_URMA=/path/to/umdk ...
+if(DEFINED URMA_INCLUDE_DIR AND URMA_INCLUDE_DIR)
+  set(urma_INCLUDE_DIR ${URMA_INCLUDE_DIR})
+  message(STATUS "Using provided URMA_INCLUDE_DIR: ${urma_INCLUDE_DIR}")
+elseif(DEFINED FETCHCONTENT_SOURCE_DIR_URMA AND FETCHCONTENT_SOURCE_DIR_URMA)
+  set(urma_SOURCE_DIR ${FETCHCONTENT_SOURCE_DIR_URMA})
+  set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
+  message(STATUS "Using FETCHCONTENT_SOURCE_DIR_URMA: ${urma_SOURCE_DIR}")
+else()
+  FetchContent_Declare(
+    urma
+    GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git
+    GIT_TAG v25.12.0.B081
+    GIT_SHALLOW TRUE)
 
-FetchContent_MakeAvailable(urma)
+  FetchContent_GetProperties(urma)
+  if(NOT urma_POPULATED)
+    FetchContent_Populate(urma)
+  endif()
 
-# 输出实际路径，确认位置
-message(STATUS "URMA source dir: ${urma_SOURCE_DIR}")
-message(STATUS "URMA binary dir: ${urma_BINARY_DIR}")
+  set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
+  message(STATUS "URMA source dir: ${urma_SOURCE_DIR}")
+  message(STATUS "URMA binary dir: ${urma_BINARY_DIR}")
+endif()
 
-# 假设 UMDK 头文件在其 include 目录下
-set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include)
-
-# 添加到需要的目标
-message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}")
+message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}")

mooncake-transfer-engine/example/transfer_engine_bench.cpp

@@ -90,10 +90,11 @@ DEFINE_string(operation, "read", "Operation type: read or write");
 
 DEFINE_string(protocol, "rdma",
               "Transfer protocol: "
-              "rdma|barex|tcp|efa|nvlink|nvlink_intra|hip|sunrise_link");
+              "rdma|barex|tcp|efa|ub|nvlink|nvlink_intra|hip|ubshmem|"
+              "sunrise_link");
 
 DEFINE_string(device_name, "mlx5_2",
-              "Device name to use, valid if protocol=rdma");
+              "Device name to use, valid if protocol=rdma|ub");
 DEFINE_string(nic_priority_matrix, "",
               "Path to RDMA NIC priority matrix file (Advanced)");
 
@@ -673,6 +674,17 @@ std::shared_ptr<mooncake::tent::Config> createTentConfig() {
     config->set("local_segment_name", FLAGS_local_server_name);
     config->set("verbose", true);
 
+    // Propagate --device_name to the UB/RDMA transport so that
+    // UbTentTransport::install() can filter which device to use
+    // (important for bonded-device setups like bonding_dev_0).
+    if (!FLAGS_device_name.empty()) {
+        if (FLAGS_protocol == "rdma") {
+            config->set("transports/rdma/device_name", FLAGS_device_name);
+        } else if (FLAGS_protocol == "ub") {
+            config->set("transports/ub/device_name", FLAGS_device_name);
+        }
+    }
+
     return config;
 }
 

mooncake-transfer-engine/include/transfer_metadata.h

@@ -162,6 +162,7 @@ class TransferMetadata {
         std::string peer_nic_path;
 #ifdef USE_UB
         std::vector<uint32_t> jetty_num;  // for ub/urma
+        std::string local_eid;            // for ub/urma
 #endif
 #ifdef USE_BAREX
         uint16_t barex_port;
@@ -188,13 +189,13 @@ class TransferMetadata {
    public:
     TransferMetadata(const std::string &conn_string);
 
-    ~TransferMetadata();
+    virtual ~TransferMetadata();
 
-    std::shared_ptr<SegmentDesc> getSegmentDescByName(
+    virtual std::shared_ptr<SegmentDesc> getSegmentDescByName(
         const std::string &segment_name, bool force_update = false);
 
-    std::shared_ptr<SegmentDesc> getSegmentDescByID(SegmentID segment_id,
-                                                    bool force_update = false);
+    virtual std::shared_ptr<SegmentDesc> getSegmentDescByID(
+        SegmentID segment_id, bool force_update = false);
 
     int updateLocalSegmentDesc(SegmentID segment_id = LOCAL_SEGMENT_ID);
 
@@ -204,7 +205,7 @@ class TransferMetadata {
     std::shared_ptr<SegmentDesc> getSegmentDesc(
         const std::string &segment_name);
 
-    SegmentID getSegmentID(const std::string &segment_name);
+    virtual SegmentID getSegmentID(const std::string &segment_name);
 
     int syncSegmentCache(const std::string &segment_name);
 
@@ -234,12 +235,12 @@ class TransferMetadata {
 
     using OnReceiveHandShake = std::function<int(const HandShakeDesc &peer_desc,
                                                  HandShakeDesc &local_desc)>;
-    int startHandshakeDaemon(OnReceiveHandShake on_receive_handshake,
-                             uint16_t listen_port, int sockfd);
+    virtual int startHandshakeDaemon(OnReceiveHandShake on_receive_handshake,
+                                     uint16_t listen_port, int sockfd);
 
-    int sendHandshake(const std::string &peer_server_name,
-                      const HandShakeDesc &local_desc,
-                      HandShakeDesc &peer_desc);
+    virtual int sendHandshake(const std::string &peer_server_name,
+                              const HandShakeDesc &local_desc,
+                              HandShakeDesc &peer_desc);
 
     int sendNotify(const std::string &peer_server_name,
                    const NotifyDesc &local_desc, NotifyDesc &peer_desc);

mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h

@@ -174,6 +174,19 @@ class UbContext {
    public:
     virtual int registerMemoryRegion(uint64_t va, size_t length) = 0;
 
+    // Returns the most recently registered local segment handle on this
+    // context (the one produced by the last successful registerMemoryRegion),
+    // or nullptr if none.  Used to share a single host-global URMA segment
+    // across all contexts (see UbTransport::registerLocalMemory).
+    virtual void* lastRegisteredSeg() = 0;
+
+    // Adopts a segment that was registered on another context for the same
+    // host virtual address.  URMA registers memory into a host-global ubva
+    // space, so a given host VA must be registered with the driver exactly
+    // once; the other contexts reference that single segment instead of
+    // re-registering it (which the driver rejects as a duplicate).
+    virtual int adoptLocalSeg(uint64_t va, size_t length, void* seg) = 0;
+
     virtual int unregisterMemoryRegion(uint64_t va) = 0;
 
     virtual int doProcessContextEvents() = 0;

mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h

@@ -18,6 +18,7 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <unordered_set>
 #include <utility>
 #include "common.h"
 #include "config.h"
@@ -55,6 +56,8 @@ class UrmaContext : public UbContext {
                 int max_endpoints);
     ~UrmaContext();
     int registerMemoryRegion(uint64_t va, size_t length) override;
+    void* lastRegisteredSeg() override;
+    int adoptLocalSeg(uint64_t va, size_t length, void* seg) override;
     int unregisterMemoryRegion(uint64_t va) override;
     int doProcessContextEvents() override;
     void* retrieveRemoteSeg(const std::string& value) override;
@@ -127,6 +130,10 @@ class UrmaContext : public UbContext {
     RWSpinlock seg_region_lock_;
     std::vector<std::pair<urma_target_seg_t*, uint64_t>> seg_region_list_;
     std::vector<urma_target_seg_t*> local_tseg_list_;
+    // Local segments actually registered with the URMA driver by THIS context
+    // (as opposed to segments adopted from another context for the same
+    // host-global VA).  Only owned segments are passed to urma_unregister_seg.
+    std::unordered_set<urma_target_seg_t*> owned_segs_;
     std::vector<urma_seg_t*> remote_seg_list_;
     std::vector<urma_target_seg_t*> imported_seg_list_;
 

mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp

@@ -86,15 +86,49 @@ int UbTransport::registerLocalMemory(void* addr, size_t length,
                                      bool update_metadata) {
     (void)remote_accessible;
     BufferDesc buffer_desc;
-    for (auto& context : context_list_) {
-        int ret = context->registerMemoryRegion((uint64_t)addr, length);
+    if (context_list_.empty()) {
+        LOG(ERROR) << "UbTransport: no available context to register memory";
+        return ERR_DEVICE_NOT_FOUND;
+    }
+
+    // URMA registers host memory into a host-global ubva space, so a given host
+    // virtual address may be registered with the driver exactly once per host
+    // (re-registering the same VA on another context is rejected by the driver
+    // as a duplicate, e.g. "registered twice within 500ms").  Register the
+    // buffer on the primary context, then share the resulting segment with the
+    // remaining contexts so every device's data path still has a valid local
+    // segment handle for this buffer.
+    int ret = context_list_[0]->registerMemoryRegion((uint64_t)addr, length);
+    if (ret) {
+        LOG(ERROR) << "UbTransport: cannot register LocalMemory on primary "
+                      "context";
+        return ret;
+    }
+    void* shared_seg = context_list_[0]->lastRegisteredSeg();
+    if (!shared_seg) {
+        LOG(ERROR) << "UbTransport: primary context returned null segment";
+        context_list_[0]->unregisterMemoryRegion((uint64_t)addr);
+        return ERR_CONTEXT;
+    }
+
+    for (size_t i = 1; i < context_list_.size(); ++i) {
+        ret =
+            context_list_[i]->adoptLocalSeg((uint64_t)addr, length, shared_seg);
         if (ret) {
-            LOG(ERROR) << "UbTransport: cannot register LocalMemory";
+            LOG(ERROR) << "UbTransport: cannot share segment to context " << i;
+            for (size_t j = 0; j < i; ++j)
+                context_list_[j]->unregisterMemoryRegion((uint64_t)addr);
             return ret;
         }
-        ret = context->buildLocalBufferDesc((uint64_t)addr, buffer_desc);
+    }
+
+    for (size_t i = 0; i < context_list_.size(); ++i) {
+        ret =
+            context_list_[i]->buildLocalBufferDesc((uint64_t)addr, buffer_desc);
         if (ret) {
             LOG(ERROR) << "UbTransport: build buffer description failed";
+            for (size_t j = 0; j < context_list_.size(); ++j)
+                context_list_[j]->unregisterMemoryRegion((uint64_t)addr);
             return ret;
         }
     }
@@ -203,19 +237,29 @@ Status UbTransport::submitTransfer(
 
 Status UbTransport::submitTransferTask(
     const std::vector<TransferTask*>& task_list) {
-    std::unordered_map<std::shared_ptr<UbContext>, std::vector<Slice*>>
-        slices_to_post;
     auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID);
     const size_t kBlockSize = globalConfig().slice_size;
     const int kMaxRetryCount = globalConfig().retry_cnt;
     const size_t kFragmentSize = globalConfig().fragment_limit;
     const size_t kSubmitWatermark =
         globalConfig().max_wr * globalConfig().num_qp_per_ep;
-    uint64_t nr_slices;
-    for (size_t index = 0; index < task_list.size(); ++index) {
-        assert(task_list[index]);
-        auto& task = *task_list[index];
-        nr_slices = 0;
+
+    struct PlannedSlice {
+        TransferTask* task;
+        std::shared_ptr<UbContext> context;
+        uint64_t offset;
+        size_t length;
+        int buffer_id;
+        int device_id;
+    };
+    std::vector<PlannedSlice> plan;
+
+    // Validate and route the complete submission before mutating TransferTask
+    // state or posting any work. This guarantees that a non-OK return has no
+    // partial submission for the caller to roll back.
+    for (auto* task_ptr : task_list) {
+        assert(task_ptr);
+        auto& task = *task_ptr;
         assert(task.request);
         auto& request = *task.request;
         auto request_buffer_id = -1, request_device_id = -1;
@@ -229,27 +273,11 @@ Status UbTransport::submitTransferTask(
 
         for (uint64_t offset = 0; offset < request.length;
              offset += kBlockSize) {
-            Slice* slice = getSliceCache().allocate();
-            assert(slice);
-            if (!slice->from_cache) {
-                nr_slices++;
-            }
             bool merge_final_slice =
                 request.length - offset <= kBlockSize + kFragmentSize;
-            slice->source_addr = (char*)request.source + offset;
-            slice->length =
+            auto* source_addr = static_cast<char*>(request.source) + offset;
+            size_t slice_length =
                 merge_final_slice ? request.length - offset : kBlockSize;
-            slice->opcode = request.opcode;
-            // LOG(INFO) << "target_offset : " << request.target_offset << ",
-            // offset : " << offset;
-            slice->ub.dest_addr = request.target_offset + offset;
-            slice->ub.retry_cnt = 0;
-            slice->ub.max_retry_cnt = kMaxRetryCount;
-            slice->task = &task;
-            slice->target_id = request.target_id;
-            slice->ts = 0;
-            slice->status = Slice::PENDING;
-            task.slice_list.push_back(slice);
 
             int buffer_id = -1, device_id = -1,
                 retry_cnt = request.advise_retry_cnt;
@@ -261,8 +289,9 @@ Status UbTransport::submitTransferTask(
             }
             while (retry_cnt < kMaxRetryCount && !found_device) {
                 if (selectDevice(local_segment_desc.get(),
-                                 (uint64_t)slice->source_addr, slice->length,
-                                 buffer_id, device_id, retry_cnt++))
+                                 reinterpret_cast<uint64_t>(source_addr),
+                                 slice_length, buffer_id, device_id,
+                                 retry_cnt++))
                     continue;
                 assert(device_id >= 0 &&
                        static_cast<size_t>(device_id) < context_list_.size());
@@ -277,10 +306,7 @@ Status UbTransport::submitTransferTask(
                 found_device = true;
                 break;
             }
-            if (device_id < 0) {
-                auto source_addr = slice->source_addr;
-                for (auto& entry : slices_to_post)
-                    for (auto s : entry.second) getSliceCache().deallocate(s);
+            if (!found_device || device_id < 0) {
                 LOG(ERROR)
                     << "UbTransport: Address not registered by any device(s) "
                     << source_addr;
@@ -296,22 +322,54 @@ Status UbTransport::submitTransferTask(
                 return Status::InvalidArgument(
                     "Device " + std::to_string(device_id) + " is not active");
             }
-            auto local_tseg_index =
-                local_segment_desc->buffers[buffer_id].l_seg_index[device_id];
-            slice->ub.l_seg = context->localSegWithIndex(local_tseg_index);
-            slices_to_post[context].push_back(slice);
-            task.total_bytes += slice->length;
-            __sync_fetch_and_add(&task.slice_count, 1);
-            if (nr_slices >= kSubmitWatermark) {
-                for (auto& entry : slices_to_post)
-                    entry.first->submitPostSend(entry.second);
-                slices_to_post.clear();
-                nr_slices = 0;
-            }
 
-            if (merge_final_slice) {
-                break;
-            }
+            plan.push_back(
+                {&task, context, offset, slice_length, buffer_id, device_id});
+
+            if (merge_final_slice) break;
+        }
+    }
+
+    std::unordered_map<std::shared_ptr<UbContext>, std::vector<Slice*>>
+        slices_to_post;
+    TransferTask* current_task = nullptr;
+    uint64_t nr_slices = 0;
+    for (const auto& planned : plan) {
+        auto& task = *planned.task;
+        auto& request = *task.request;
+        if (current_task != &task) {
+            current_task = &task;
+            nr_slices = 0;
+        }
+
+        Slice* slice = getSliceCache().allocate();
+        assert(slice);
+        if (!slice->from_cache) ++nr_slices;
+        slice->source_addr =
+            static_cast<char*>(request.source) + planned.offset;
+        slice->length = planned.length;
+        slice->opcode = request.opcode;
+        slice->ub.dest_addr = request.target_offset + planned.offset;
+        slice->ub.retry_cnt = 0;
+        slice->ub.max_retry_cnt = kMaxRetryCount;
+        slice->task = &task;
+        slice->target_id = request.target_id;
+        slice->ts = 0;
+        slice->status = Slice::PENDING;
+        task.slice_list.push_back(slice);
+
+        auto& context = planned.context;
+        auto local_tseg_index = local_segment_desc->buffers[planned.buffer_id]
+                                    .l_seg_index[planned.device_id];
+        slice->ub.l_seg = context->localSegWithIndex(local_tseg_index);
+        slices_to_post[context].push_back(slice);
+        task.total_bytes += slice->length;
+        __sync_fetch_and_add(&task.slice_count, 1);
+        if (nr_slices >= kSubmitWatermark) {
+            for (auto& entry : slices_to_post)
+                entry.first->submitPostSend(entry.second);
+            slices_to_post.clear();
+            nr_slices = 0;
         }
     }
     for (auto& entry : slices_to_post)