diff --git a/mooncake-common/FindUrma.cmake b/mooncake-common/FindUrma.cmake index 3966437a31..4dd439d80b 100644 --- a/mooncake-common/FindUrma.cmake +++ b/mooncake-common/FindUrma.cmake @@ -7,6 +7,11 @@ find_path( NAMES urma_api.h PATHS /usr/include /usr/local/include PATH_SUFFIXES urma umdk src/urma/lib/urma/core/include) +find_path( + URMA_BOND_SYSTEM_INCLUDE_DIR + NAMES urma_ubagg.h + PATHS /usr/include /usr/local/include + PATH_SUFFIXES ub/umdk/urma urma umdk src/urma/lib/urma/bond/include) find_library( URMA_LIBRARY NAMES urma @@ -14,6 +19,7 @@ find_library( if(URMA_SYSTEM_INCLUDE_DIR) set(urma_INCLUDE_DIR "${URMA_SYSTEM_INCLUDE_DIR}") + set(urma_bond_INCLUDE_DIR "${URMA_BOND_SYSTEM_INCLUDE_DIR}") else() # The source fallback supplies headers only. Production TENT UB remains # disabled at runtime when no real liburma is present; tests inject their own @@ -24,12 +30,18 @@ else() GIT_TAG v25.12.0.B081) FetchContent_MakeAvailable(urma) set(urma_INCLUDE_DIR "${urma_SOURCE_DIR}/src/urma/lib/urma/core/include") + set(urma_bond_INCLUDE_DIR + "${urma_SOURCE_DIR}/src/urma/lib/urma/bond/include") endif() if(NOT TARGET Urma::urma) add_library(Urma::urma INTERFACE IMPORTED GLOBAL) + set(urma_interface_include_dirs "${urma_INCLUDE_DIR}") + if(urma_bond_INCLUDE_DIR) + list(APPEND urma_interface_include_dirs "${urma_bond_INCLUDE_DIR}") + endif() set_property(TARGET Urma::urma PROPERTY INTERFACE_INCLUDE_DIRECTORIES - "${urma_INCLUDE_DIR}") + "${urma_interface_include_dirs}") if(URMA_LIBRARY) set_property(TARGET Urma::urma PROPERTY INTERFACE_LINK_LIBRARIES "${URMA_LIBRARY}") diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index cdb156303c..d8e287c9c8 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -662,6 +662,10 @@ class Client { // Return sorted NUMA node IDs that have at least one RDMA NIC. [[nodiscard]] std::vector GetNicNumaNodes() const; + // Return the total number of NUMA nodes (counted from the local topology, + // which has one cpu:N entry per node regardless of NIC presence). + [[nodiscard]] int GetNumaNodeCount() const; + tl::expected GetPreferredReplica( const std::vector& replica_list); diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 77291f2913..6d73811324 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -13,7 +13,6 @@ #include #include #include -#include #include #include #include @@ -3432,6 +3431,20 @@ std::vector Client::GetNicNumaNodes() const { return {nodes.begin(), nodes.end()}; } +int Client::GetNumaNodeCount() const { + if (!transfer_engine_) return 0; + auto topo = transfer_engine_->getLocalTopology(); + if (!topo) return 0; + // discoverCpuTopology emits one "cpu:N" entry per NUMA node (regardless of + // whether it hosts a NIC), so counting them gives the NUMA node count. + int count = 0; + for (auto& [name, entry] : topo->getMatrix()) { + (void)entry; + if (name.rfind("cpu:", 0) == 0) ++count; + } + return count; +} + tl::expected Client::MountSegment( const void* buffer, size_t size, const std::string& protocol, const std::string& location) { diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 5a93169618..b4f4f05754 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -40,6 +40,9 @@ #ifdef USE_NOF #include "spdk/spdk_wrapper.h" #endif +#ifdef USE_UB +#include "ub_allocator.h" +#endif #ifdef USE_ASCEND_DIRECT #include "acl/acl_rt.h" #include "transport/ascend_transport/ascend_direct_transport/context_manager.h" @@ -951,7 +954,8 @@ tl::expected RealClient::setup_internal( #endif // For RDMA, auto-discover NUMA nodes with NICs and distribute - // global_segment across them for full NIC utilization. + // global_segment across them for full NIC utilization. RDMA keeps the + // legacy single-segment-with-multi-region ("segments:...") behavior. std::vector seg_numa_nodes; if (protocol == "rdma") { seg_numa_nodes = client_->GetNicNumaNodes(); @@ -968,6 +972,32 @@ tl::expected RealClient::setup_internal( } } + // For UB, mount ONE segment per NUMA node (each bound to its node, + // location="cpu:N"). Memory is spread across ALL NUMA nodes by count, + // not just NIC-bearing ones: this works even with a single bonded + // device (e.g. bonding_dev_0, NUMA=-1) where NIC-NUMA discovery would + // be empty, and it relies only on the (decimal) NUMA count, so it is + // unaffected by the hex "numa" attribute. Each segment then drives + // selectDevice (cpu:N -> local NIC, else any) and chip affinity + // (cpu:N -> chip via numaNodeToChipId). Automatic whenever UB has more + // than one NUMA node; independent of both MC_UB_NUMA_AFFINITY_ENABLE + // and MC_URMA_BONDING_MULTIPATH_ENABLE. +#ifdef USE_UB + std::vector ub_numa_nodes; + if (protocol == "ub") { + int numa_count = client_->GetNumaNodeCount(); + if (numa_count > 1) { + std::string nodes_str; + for (int i = 0; i < numa_count; ++i) { + ub_numa_nodes.push_back(i); + if (i) nodes_str += ","; + nodes_str += std::to_string(i); + } + LOG(INFO) << "UB per-NUMA mode: NUMA node count=" << numa_count + << ", nodes=[" << nodes_str << "]"; + } + } +#endif // USE_UB const bool parallel_hugetlb_population = protocol == "rdma" && should_use_hugepage; @@ -981,6 +1011,57 @@ tl::expected RealClient::setup_internal( } global_segment_size -= segment_size; + // UB NUMA affinity: split this chunk into one segment per NIC-NUMA + // node, each physically bound to its node and registered with + // location "cpu:N" (so selectDevice picks the NUMA-local NIC). +#ifdef USE_UB + if (!ub_numa_nodes.empty()) { + size_t page_sz = should_use_hugepage + ? get_hugepage_size_from_env() + : static_cast(getpagesize()); + size_t n = ub_numa_nodes.size(); + size_t per_node_size = align_up(segment_size / n, page_sz); + if (per_node_size == 0) { + LOG(ERROR) << "UB per-NUMA: per_node_size is 0, segment " + "too small for " + << n << " NUMA nodes"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + for (int node : ub_numa_nodes) { + // Use UB's own allocator bound to this node: + // numa_alloc_onnode via libnuma, registered in the + // store-memory table, so URMA can register it. (A raw + // mmap+mbind buffer cannot be registered by + // urma_register_seg -- it fails with error 2048 because the + // VMA has no backing pages at reg time.) + void *ptr = mooncake::ub_allocate_memory_onnode( + /*alignment=*/page_sz, per_node_size, node); + if (!ptr) { + LOG(ERROR) << "UB per-NUMA: failed to allocate " + "segment for node " + << node; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + // numa_alloc-backed => free via ub_free_memory/numa_free, + // NOT munmap. Track with UbSegmentDeleter accordingly. + ub_segment_ptrs_.emplace_back( + ptr, UbSegmentDeleter{per_node_size}); + + std::string loc = genCpuNodeName(node); // "cpu:" + LOG(INFO) << "Mounting UB per-NUMA segment: node=" << node + << ", size=" << per_node_size << ", loc=" << loc; + auto mr = client_->MountSegment(ptr, per_node_size, + protocol, loc); + if (!mr.has_value()) { + LOG(ERROR) << "Failed to mount UB per-NUMA segment: " + << toString(mr.error()); + return tl::unexpected(mr.error()); + } + } + continue; // this chunk fully mounted across NUMA nodes + } +#endif // USE_UB + size_t mapped_size = segment_size; void *ptr = nullptr; std::string seg_location = kWildcardLocation; diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index fa09744749..7662ed9c22 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -167,6 +167,14 @@ struct GlobalConfig { uint64_t max_seg_size = 0x10000000000; size_t max_jfc_e = 4096; // urma is temporarily using this default value. size_t num_jetty_per_ep = 1; + // Enable URMA bonding multipath mode. Default is off; override via + // MC_URMA_BONDING_MULTIPATH_ENABLE. + bool urma_bonding_multipath = false; + // Enable UB NUMA affinity: store splits the global segment into one + // segment per NIC-NUMA node, and transfers pin src/dst chip by NUMA. + // Independent from urma_bonding_multipath; default off; override via + // MC_UB_NUMA_AFFINITY_ENABLE. + bool ub_numa_affinity = false; }; struct RpcCommunicatorConfig { diff --git a/mooncake-transfer-engine/include/memory_location.h b/mooncake-transfer-engine/include/memory_location.h index 6e47dc986b..692260273c 100644 --- a/mooncake-transfer-engine/include/memory_location.h +++ b/mooncake-transfer-engine/include/memory_location.h @@ -31,6 +31,19 @@ struct MemoryLocationEntry { std::string location; }; +const static uint8_t INVALID_CHIP_ID = 0xFF; + +// "cpu:3" -> 3; "*" 或 非cpu串返回-1 +int parseCpuNumaNode(const std::string &location); + +// NUMA 节点 -> chip id:优先按 sysfs 真实拓扑(physical_package_id)映射; +// sysfs 不可读时回退「前半 chip1 / 后半 chip2」启发式;失败返回 +// INVALID_CHIP_ID。 +uint8_t numaNodeToChipId(int numa_node, size_t numa_count = 0); + +// NUMA 节点 -> location 字符串:"cpu:N"(node>=0)或 "*"(node<0)。 +std::string genCpuNodeName(int node); + // If only_first_page is true, only the location of the first page will be // returned. Scan all pages may take a long time, so set only_first_page if only // the location of the first page is needed. diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 228003d984..99a4158e85 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -75,6 +75,11 @@ class TransferMetadata { uint64_t offset; // for cxl std::vector tseg; // for ub/urma std::vector l_seg_index; // for ub/urma + // for ub: NUMA->chip id of this buffer, computed once by the owner at + // registration and published. -1 = not provided (consumer falls back + // to resolving from `name`). Only meaningful for single-NUMA ("cpu:N") + // buffers; multi-NUMA "segments:..." buffers leave it -1. + int chip_id = -1; bool operator==(const BufferDesc &other) const = default; }; diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h index 56767c9dee..1a51eb45ef 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -138,7 +138,9 @@ class UbContext { max_endpoints_(max_endpoints), worker_pool_(nullptr), active_(true), - show_work_request_flushed_error_(false) {} + show_work_request_flushed_error_(false), + multipath_(globalConfig().urma_bonding_multipath), + numa_affinity_(globalConfig().ub_numa_affinity) {} virtual ~UbContext() = default; @@ -214,6 +216,10 @@ class UbContext { void set_active(bool flag) { active_ = flag; } + bool multipath() const { return multipath_; } + + bool numa_affinity() const { return numa_affinity_; } + // EndPoint Management std::shared_ptr endpoint() { return endpoint("LOCAL_SEGMENT_ID"); @@ -369,6 +375,10 @@ class UbContext { volatile bool active_; bool show_work_request_flushed_error_; + + bool multipath_ = false; + + bool numa_affinity_ = false; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 5ece720f95..3a50389c34 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -161,6 +161,8 @@ class Transport { uint32_t max_retry_cnt; void *r_seg; void *l_seg; + uint8_t src_chip_id; + uint8_t dst_chip_id; void *endpoint; } ub; struct { diff --git a/mooncake-transfer-engine/include/ub_allocator.h b/mooncake-transfer-engine/include/ub_allocator.h index a753f42165..4a9375e8b2 100644 --- a/mooncake-transfer-engine/include/ub_allocator.h +++ b/mooncake-transfer-engine/include/ub_allocator.h @@ -4,6 +4,12 @@ namespace mooncake { void* ub_allocate_memory(size_t alignment, size_t total_size); +// Same as ub_allocate_memory but binds the allocation to a specific NUMA node +// (numa_node < 0 falls back to node-local). Still allocated via libnuma and +// registered in the store-memory range table, so URMA can register it. +void* ub_allocate_memory_onnode(size_t alignment, size_t total_size, + int numa_node); + void ub_free_memory(void* ptr); bool ub_is_store_memory(void* addr, size_t length); diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index ceb23f8658..7fc8a88f78 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -728,6 +728,31 @@ void loadGlobalConfig(GlobalConfig& config) { } } + const char* urma_bonding_multipath_enable = + std::getenv("MC_URMA_BONDING_MULTIPATH_ENABLE"); + if (urma_bonding_multipath_enable && *urma_bonding_multipath_enable) { + std::string val(urma_bonding_multipath_enable); + if (val == "true" || val == "1" || val == "on") { + config.urma_bonding_multipath = true; + LOG(WARNING) << "MC_URMA_BONDING_MULTIPATH_ENABLE is " << val; + } else + LOG(WARNING) + << "Ignore value from environment variable " + "MC_URMA_BONDING_MULTIPATH_ENABLE, it should be true|1|on"; + } + + const char* ub_numa_affinity_enable = + std::getenv("MC_UB_NUMA_AFFINITY_ENABLE"); + if (ub_numa_affinity_enable && *ub_numa_affinity_enable) { + std::string val(ub_numa_affinity_enable); + if (val == "true" || val == "1" || val == "on") { + config.ub_numa_affinity = true; + LOG(WARNING) << "MC_UB_NUMA_AFFINITY_ENABLE is " << val; + } else + LOG(WARNING) + << "Ignore value from environment variable " + "MC_UB_NUMA_AFFINITY_ENABLE, it should be true|1|on"; + } const char* mlx5_qp_lag_port_balance_env = std::getenv("MC_MLX5_QP_LAG_PORT_BALANCE"); if (mlx5_qp_lag_port_balance_env && *mlx5_qp_lag_port_balance_env) { diff --git a/mooncake-transfer-engine/src/memory_location.cpp b/mooncake-transfer-engine/src/memory_location.cpp index 32b488b05e..98e7da0212 100644 --- a/mooncake-transfer-engine/src/memory_location.cpp +++ b/mooncake-transfer-engine/src/memory_location.cpp @@ -12,6 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + #include "memory_location.h" #include "cuda_alike.h" @@ -20,6 +31,146 @@ namespace mooncake { uintptr_t alignPage(uintptr_t address) { return address & ~(pagesize - 1); } +int parseCpuNumaNode(const std::string &location) { + const std::string prefix = "cpu:"; + if (location.rfind(prefix, 0) != 0) return -1; + try { + return std::stoi(location.substr(prefix.size())); + } catch (const std::exception &) { + return -1; + } +} + +// 数 /sys/devices/system/node 下的 NUMA 节点数 +static size_t getNumaNodeCount() { + int count = 0; + DIR *dir = opendir("/sys/devices/system/node"); + if (!dir) return 0; + for (dirent *e = readdir(dir); e; e = readdir(dir)) + if (strncmp(e->d_name, "node", 4) == 0 && + isdigit((unsigned char)e->d_name[4])) + count++; + closedir(dir); + return (size_t)count; +} + +namespace { + +// 读取 NUMA 节点 nodeN 下第一个 CPU 的物理 package(chip) 编号。 +// 路径:/sys/devices/system/node/nodeN/cpulist -> 取首个 cpu -> +// /sys/devices/system/cpu/cpuX/topology/physical_package_id。失败返回 +// -1。 +int readNodePackageId(int node) { + char path[256]; + snprintf(path, sizeof(path), "/sys/devices/system/node/node%d/cpulist", + node); + std::ifstream cpulist(path); + if (!cpulist) return -1; + std::string s; + std::getline(cpulist, s); // 形如 "0-23" 或 "0-7,16-23" + int cpu = -1; + try { + cpu = std::stoi(s); // 取列表首个 CPU + } catch (const std::exception &) { + return -1; + } + if (cpu < 0) return -1; + + snprintf(path, sizeof(path), + "/sys/devices/system/cpu/cpu%d/topology/physical_package_id", cpu); + std::ifstream pkg(path); + if (!pkg) return -1; + int package_id = -1; + pkg >> package_id; + return package_id; +} + +// 进程内只构建一次:从服务器拓扑(sysfs)读取 NUMA节点 -> chip(物理package) +// 映射。 chip id 采用「package 排序后 1-based」编号,兼容旧约定(双 chip 即 +// 1/2), 同时正确处理节点与 package 非顺序对应的拓扑。 +class NumaChipMap { + public: + static const NumaChipMap &Instance() { + static const NumaChipMap inst; + return inst; + } + + // 返回该 NUMA 节点所属 chip id;未知返回 INVALID_CHIP_ID。 + uint8_t ChipId(int numa_node) const { + auto it = node_to_chip_.find(numa_node); + return it == node_to_chip_.end() ? INVALID_CHIP_ID : it->second; + } + + bool Empty() const { return node_to_chip_.empty(); } + + private: + NumaChipMap() { Build(); } + + void Build() { + std::map node_to_pkg; // numa node -> physical package id + std::set packages; + + DIR *dir = opendir("/sys/devices/system/node"); + if (!dir) return; + for (dirent *e = readdir(dir); e; e = readdir(dir)) { + if (strncmp(e->d_name, "node", 4) != 0 || + !isdigit((unsigned char)e->d_name[4])) + continue; + int node = atoi(e->d_name + 4); + int pkg = readNodePackageId(node); + if (pkg < 0) continue; + node_to_pkg[node] = pkg; + packages.insert(pkg); + } + closedir(dir); + + // package id 排序去重 -> 1-based chip id + std::map pkg_to_chip; + uint8_t chip = 1; + for (int p : packages) pkg_to_chip[p] = chip++; + + for (const auto &[node, pkg] : node_to_pkg) + node_to_chip_[node] = pkg_to_chip[pkg]; + + std::string mapping; + for (const auto &[node, pkg] : node_to_pkg) { + if (!mapping.empty()) mapping += ", "; + mapping += "numa:" + std::to_string(node) + + " package:" + std::to_string(pkg) + + " chip:" + std::to_string(pkg_to_chip[pkg]); + } + LOG(INFO) << "[numa_affinity] numa_chip_map nodes=" + << node_to_chip_.size() << " chips=" << packages.size() + << " mapping={" << mapping << "}"; + } + + std::unordered_map node_to_chip_; +}; + +} // namespace + +// NUMA 节点 -> chip id:优先按 sysfs 真实拓扑(physical_package_id)映射; +// 若 sysfs 不可读则回退到旧的「前半 chip1 / 后半 chip2」启发式。 +uint8_t numaNodeToChipId(int numa_node, size_t numa_count) { + if (numa_node < 0) return INVALID_CHIP_ID; // 对应 INVALID_NUMA_ID + + const auto &chip_map = NumaChipMap::Instance(); + if (!chip_map.Empty()) { + return chip_map.ChipId(numa_node); + } + + // Fallback:sysfs + // 读取失败(如受限容器),退回原启发式,保证不破坏既有行为。 + constexpr uint8_t chipId1 = 1, chipId2 = 2; + if (numa_count == 0) { + numa_count = getNumaNodeCount(); + if (numa_count == 0) return INVALID_CHIP_ID; + } + if ((size_t)numa_node >= numa_count) return INVALID_CHIP_ID; + const size_t firstHalfCount = (numa_count + 1) / 2; + return ((size_t)numa_node < firstHalfCount) ? chipId1 : chipId2; +} + std::string genCpuNodeName(int node) { if (node >= 0) return "cpu:" + std::to_string(node); return kWildcardLocation; diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index 5383f1e3d7..59fb9918f5 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -399,7 +399,20 @@ static std::vector listUBDevices( snprintf(path, sizeof(path), "%s/numa", dirname(dirname(resolved_path))); LOG(INFO) << "listUBDevices: numanodepath " << path; - std::ifstream(path) >> numa_node; + // The ubcore "numa" attribute is written in hex (e.g. "0x01"), unlike + // the infiniband "numa_node" which is decimal. A plain `>> int` parses + // "0x01" as 0 (stops at 'x'), so read as string and pick the base by + // the 0x prefix. + std::string numa_str; + std::ifstream(path) >> numa_str; + if (!numa_str.empty()) { + int base = + (numa_str.rfind("0x", 0) == 0 || numa_str.rfind("0X", 0) == 0) + ? 16 + : 10; + numa_node = + static_cast(strtol(numa_str.c_str(), nullptr, base)); + } LOG(INFO) << "UBDevices : performation node ----" << device_list[i]->name << " : " << numa_node; diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index 3148861208..88cab7b976 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -479,6 +479,7 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, Json::Value tsegJSON(Json::arrayValue); for (auto &entry : buffer.tseg) tsegJSON.append(entry); bufferJSON["tseg"] = tsegJSON; + bufferJSON["chip_id"] = buffer.chip_id; buffersJSON.append(bufferJSON); } segmentJSON["buffers"] = buffersJSON; @@ -902,6 +903,10 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, for (const auto &tsegJSON : bufferJSON["tseg"]) { buffer.tseg.push_back(tsegJSON.asString()); } + // Backward compatible: old peers don't publish chip_id -> keep -1. + buffer.chip_id = bufferJSON.isMember("chip_id") + ? bufferJSON["chip_id"].asInt() + : -1; if (buffer.name.empty() || !buffer.addr || !buffer.length || buffer.tseg.empty()) { LOG(WARNING) << "Corrupted segment descriptor, name " diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp index 609ecf9843..88ce30d3e4 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_allocator.cpp @@ -33,15 +33,22 @@ size_t remove_store_memory_range(void* ptr) { return sz; } -void* ub_allocate_memory(size_t alignment, size_t total_size) { - void* ptr = numa_alloc_local(total_size); +void* ub_allocate_memory_onnode(size_t alignment, size_t total_size, + int numa_node) { + // numa_node < 0 keeps the original node-local behavior; otherwise bind to + // the requested node. Both go through libnuma (same family as + // numa_alloc_local), so URMA can register the result -- unlike a raw + // mmap+mbind buffer. + void* ptr = (numa_node < 0) ? numa_alloc_local(total_size) + : numa_alloc_onnode(total_size, numa_node); if (!ptr) { LOG(ERROR) << "failed for UB protocol, size=" << total_size - << ", alignment : " << alignment; + << ", node=" << numa_node << ", alignment : " << alignment; return nullptr; } LOG(INFO) << "UB: allocated total size : " << total_size - << ", alignment : " << alignment << " addr at " << ptr; + << ", node : " << numa_node << ", alignment : " << alignment + << " addr at " << ptr; std::lock_guard store_lock(g_ub_store_mem_mutex); g_ub_store_mem_ranges.push_back({ptr, total_size}); @@ -49,6 +56,10 @@ void* ub_allocate_memory(size_t alignment, size_t total_size) { return ptr; } +void* ub_allocate_memory(size_t alignment, size_t total_size) { + return ub_allocate_memory_onnode(alignment, total_size, -1); +} + void ub_free_memory(void* ptr) { if (!ptr) { return; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index 46a3d1e467..2254af0dc8 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -21,8 +21,12 @@ #include "config.h" #include "transport/kunpeng_transport/ub_context.h" #include "transport/kunpeng_transport/ub_endpoint.h" +#include "memory_location.h" namespace mooncake { +namespace { +constexpr uint64_t kNumaAffinitySampleInterval = 10000; +} // namespace std::shared_ptr UbSIEVEEndpointStore::getEndpoint( const std::string& peer_nic_path) { RWSpinlock::ReadGuard guard(endpoint_map_lock_); @@ -269,6 +273,30 @@ int UbWorkerPool::submitPostSend( auto targetSegment = peer_segment_desc->buffers[buffer_id].tseg[device_id]; slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment); + if (context_.numa_affinity()) { + const auto& peer_buf = peer_segment_desc->buffers[buffer_id]; + int data_numa = parseCpuNumaNode(peer_buf.name); + if (peer_buf.chip_id >= 0) { + slice->ub.dst_chip_id = (uint8_t)peer_buf.chip_id; + } else { + slice->ub.dst_chip_id = numaNodeToChipId(data_numa); + } + static std::atomic numa_log_counter{0}; + if (VLOG_IS_ON(2) && + numa_log_counter.fetch_add(1, std::memory_order_relaxed) % + kNumaAffinitySampleInterval == + 0) { + VLOG(2) << "[numa_affinity] remote_sample batch_id=" + << slice->task->batch_id + << " target_id=" << slice->target_id << " opcode=" + << (slice->opcode == Transport::TransferRequest::READ + ? "READ" + : "WRITE") + << " remote_data_numa=" << data_numa + << " dst_chip=" << (int)slice->ub.dst_chip_id + << " remote_name=" << peer_buf.name; + } + } if (!slice->ub.r_seg) { LOG(ERROR) << "[UB] retrieveRemoteSeg failed for target_id=" << slice->target_id << " buffer_id=" << buffer_id @@ -492,6 +520,31 @@ void UbWorkerPool::redispatch(std::vector& slice_list, auto targetSegment = peer_segment_desc->buffers[buffer_id].tseg[device_id]; slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment); + if (context_.numa_affinity()) { + const auto& peer_buf = peer_segment_desc->buffers[buffer_id]; + int data_numa = parseCpuNumaNode(peer_buf.name); + if (peer_buf.chip_id >= 0) { + slice->ub.dst_chip_id = (uint8_t)peer_buf.chip_id; + } else { + slice->ub.dst_chip_id = numaNodeToChipId(data_numa); + } + static std::atomic numa_log_counter{0}; + if (VLOG_IS_ON(2) && + numa_log_counter.fetch_add(1, std::memory_order_relaxed) % + kNumaAffinitySampleInterval == + 0) { + VLOG(2) + << "[numa_affinity] remote_redispatch_sample batch_id=" + << slice->task->batch_id + << " target_id=" << slice->target_id << " opcode=" + << (slice->opcode == Transport::TransferRequest::READ + ? "READ" + : "WRITE") + << " remote_data_numa=" << data_numa + << " dst_chip=" << (int)slice->ub.dst_chip_id + << " remote_name=" << peer_buf.name; + } + } auto peer_nic_path = MakeNicPath(peer_segment_desc->nicPathServerName(), peer_segment_desc->devices[device_id].name); @@ -565,4 +618,4 @@ void UbWorkerPool::monitorWorker() { int UbWorkerPool::doProcessContextEvents() { return context_.doProcessContextEvents(); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index cf443a0ba5..c6c0941d54 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include "config.h" #include "memory_location.h" @@ -22,6 +23,10 @@ #include "transport/kunpeng_transport/urma/urma_endpoint.h" namespace mooncake { +namespace { +constexpr uint64_t kNumaAffinitySampleInterval = 10000; +} // namespace + UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) : endpoint_type_(endpoint_type) {} @@ -109,12 +114,18 @@ int UbTransport::registerLocalMemory(void* addr, size_t length, buffer_desc.name = entries[0].location; buffer_desc.addr = (uint64_t)addr; buffer_desc.length = length; + // Precompute chip_id for single-NUMA ("cpu:N") buffers so peers read it + // directly instead of resolving per-slice. -1 stays for non-cpu names. + int node = parseCpuNumaNode(buffer_desc.name); + if (node >= 0) buffer_desc.chip_id = numaNodeToChipId(node); int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); if (rc) return rc; } else { buffer_desc.name = name; buffer_desc.addr = (uint64_t)addr; buffer_desc.length = length; + int node = parseCpuNumaNode(buffer_desc.name); + if (node >= 0) buffer_desc.chip_id = numaNodeToChipId(node); int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); if (rc) return rc; } @@ -257,6 +268,8 @@ Status UbTransport::submitTransferTask( slice->target_id = request.target_id; slice->ts = 0; slice->status = Slice::PENDING; + slice->ub.src_chip_id = INVALID_CHIP_ID; + slice->ub.dst_chip_id = INVALID_CHIP_ID; task.slice_list.push_back(slice); int buffer_id = -1, device_id = -1, @@ -312,6 +325,34 @@ Status UbTransport::submitTransferTask( auto local_tseg_index = local_segment_desc->buffers[buffer_id].l_seg_index[device_id]; slice->ub.l_seg = context->localSegWithIndex(local_tseg_index); + if (context->numa_affinity()) { + const auto& local_buf = local_segment_desc->buffers[buffer_id]; + int data_numa = parseCpuNumaNode(local_buf.name); + if (local_buf.chip_id >= 0) { + // Prefer the chip id published at registration. + slice->ub.src_chip_id = (uint8_t)local_buf.chip_id; + } else { + // Each UB buffer belongs to one NUMA node and is named + // "cpu:N"; no offset-based segment lookup is required. + slice->ub.src_chip_id = numaNodeToChipId(data_numa); + } + static std::atomic numa_log_counter{0}; + if (VLOG_IS_ON(2) && + numa_log_counter.fetch_add(1, std::memory_order_relaxed) % + kNumaAffinitySampleInterval == + 0) { + VLOG(2) + << "[numa_affinity] local_sample batch_id=" + << slice->task->batch_id + << " target_id=" << slice->target_id << " opcode=" + << (slice->opcode == Transport::TransferRequest::READ + ? "READ" + : "WRITE") + << " local_data_numa=" << data_numa + << " src_chip=" << (int)slice->ub.src_chip_id + << " local_name=" << local_buf.name; + } + } slices_to_post[context].push_back(slice); task.total_bytes += slice->length; __sync_fetch_and_add(&task.slice_count, 1); @@ -434,10 +475,14 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, continue; } + // UB memory is allocated and mounted as one independent segment per + // NUMA node. Its BufferDesc name is already the resolved location + // ("cpu:N"), so no offset-based segments-location lookup is needed. + const std::string& location = buffer.name; device_id = hint.empty() - ? desc->topology.selectDevice(buffer.name, retry_cnt) - : desc->topology.selectDevice(buffer.name, hint, retry_cnt); + ? desc->topology.selectDevice(location, retry_cnt) + : desc->topology.selectDevice(location, hint, retry_cnt); if (device_id >= 0) return 0; device_id = hint.empty() ? desc->topology.selectDevice( kWildcardLocation, retry_cnt) diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp index c392577a6b..2919d3aaec 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/mock_urma.cpp @@ -54,6 +54,14 @@ urma_status_t urma_init(urma_init_attr_t *init_attr) { return URMA_SUCCESS; } +urma_status_t urma_user_ctl(urma_context_t *ctx, urma_user_ctl_in_t *in, + urma_user_ctl_out_t *out) { + (void)ctx; + (void)in; + (void)out; + return URMA_SUCCESS; +} + urma_status_t urma_uninit(void) { std::unique_lock lock(g_rw_mutex); initialized = false; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp index 06c65d0385..5ddbd11116 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp @@ -13,12 +13,40 @@ // limitations under the License. #include +#include #include #include #include "config.h" #include "transport/kunpeng_transport/urma/urma_endpoint.h" +#include "urma_ubagg.h" namespace mooncake { +namespace { +constexpr uint64_t kNumaAffinitySampleInterval = 10000; + +static const char* bondingModeToString(int mode) { + switch (mode) { + case BONDP_BONDING_MODE_STANDALONE: + return "STANDALONE"; + case BONDP_BONDING_MODE_BALANCE: + return "BALANCE"; + default: + return "UNKNOWN"; + } +} + +static const char* bondingLevelToString(int level) { + switch (level) { + case BONDP_BONDING_LEVEL_PORT: + return "PORT"; + case BONDP_BONDING_LEVEL_IODIE: + return "IODIE"; + default: + return "UNKNOWN"; + } +} +} // namespace + static int isNullEid(urma_eid_t* eid) { for (int i = 0; i < URMA_EID_SIZE; ++i) { if (eid->raw[i] != 0) return 0; @@ -410,7 +438,51 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, urma_free_device_list(devices); return ERR_CONTEXT; } - + if (multipath()) { + // multipath on: BALANCE mode at IODIE level (chip-aware bonding). + bondp_set_bonding_mode_in_t mode{ + .bonding_mode = BONDP_BONDING_MODE_BALANCE, + .bonding_level = BONDP_BONDING_LEVEL_IODIE}; + urma_user_ctl_in_t in{.addr = reinterpret_cast(&mode), + .len = sizeof(mode), + .opcode = BONDP_USER_CTL_SET_BONDING_MODE}; + urma_user_ctl_out_t out; + memset(&out, 0, sizeof(out)); + auto ret = urma_user_ctl(context, &in, &out); + if (ret != URMA_SUCCESS) { + LOG(ERROR) << "Failed to set bonding BALANCE/IODIE mode, ret = " + << ret; + return ERR_CONTEXT; + } + LOG(INFO) << "[multipath ON] bonding mode set on " << device_name + << " bonding_mode=" + << bondingModeToString(mode.bonding_mode) + << " bonding_level=" + << bondingLevelToString(mode.bonding_level); + } else { + // multipath off: explicitly set the bondp default (STANDALONE/PORT) + // instead of leaving it implicit. + bondp_set_bonding_mode_in_t mode{ + .bonding_mode = BONDP_BONDING_MODE_STANDALONE, + .bonding_level = BONDP_BONDING_LEVEL_PORT}; + urma_user_ctl_in_t in{.addr = reinterpret_cast(&mode), + .len = sizeof(mode), + .opcode = BONDP_USER_CTL_SET_BONDING_MODE}; + urma_user_ctl_out_t out; + memset(&out, 0, sizeof(out)); + auto ret = urma_user_ctl(context, &in, &out); + if (ret != URMA_SUCCESS) { + LOG(ERROR) + << "Failed to set bonding STANDALONE/PORT mode, ret = " + << ret; + return ERR_CONTEXT; + } + LOG(INFO) << "[multipath OFF] bonding mode set on " << device_name + << " bonding_mode=" + << bondingModeToString(mode.bonding_mode) + << " bonding_level=" + << bondingLevelToString(mode.bonding_level); + } ret = urma_query_device(devices[i], &dev_attr_); if (ret) { PLOG(ERROR) << "Failed to query dev attr( " << device_name << " ) "; @@ -644,6 +716,7 @@ int UrmaEndpoint::construct(GlobalConfig& config) { .err_timeout = 17, // URMA_TYPICAL_ERR_TIMEOUT 17 .user_ctx = 0, }; + jfs_cfg.flag.bs.multi_path = context_->multipath() ? 1 : 0; urma_jetty_flag_t jetty_flag = {}; urma_jetty_cfg_t attr; memset(&attr, 0, sizeof(attr)); @@ -880,68 +953,114 @@ int UrmaEndpoint::submitPostSend( std::min(int(globalConfig().max_jfc_e) - *jfc_outstanding_, wr_count); if (wr_count <= 0) return 0; - urma_jfs_wr_t wr_list[wr_count], *bad_wr = nullptr; + // chip 亲和(用 bondp_jfs_wr_t 带 src/dst_chip_id)由 NUMA 亲和开关控制; + // bonding 模式本身(设备级)由 multipath 在 openDevice 设置。 + const bool numa_affinity = context_->numa_affinity(); urma_sge_t l_sge_list[wr_count]; urma_sge_t r_sge_list[wr_count]; - memset(wr_list, 0, sizeof(urma_jfs_wr_t) * wr_count); - for (int i = 0; i < wr_count; ++i) { - auto slice = slice_list[i]; - auto& l_sge = l_sge_list[i]; - auto& r_sge = r_sge_list[i]; + + // 两分支共用:填一个 WR 的公共字段(SGE + // 方向、opcode、flag、tjetty、user_ctx + slice 记账) + auto fill_common = [&](urma_jfs_wr_t& wr, Transport::Slice* slice, + urma_sge_t& l_sge, urma_sge_t& r_sge) { + const bool is_read = slice->opcode == Transport::TransferRequest::READ; l_sge.addr = (uint64_t)slice->source_addr; l_sge.len = slice->length; l_sge.tseg = static_cast(slice->ub.l_seg); r_sge.addr = slice->ub.dest_addr; r_sge.len = slice->length; r_sge.tseg = static_cast(slice->ub.r_seg); - - auto& wr = wr_list[i]; wr.user_ctx = (uint64_t)slice; - wr.opcode = slice->opcode == Transport::TransferRequest::READ - ? URMA_OPC_READ - : URMA_OPC_WRITE; - wr.rw.src.sge = - slice->opcode == Transport::TransferRequest::READ ? &r_sge : &l_sge; - wr.rw.src.num_sge = 1; - wr.rw.dst.sge = - slice->opcode == Transport::TransferRequest::READ ? &l_sge : &r_sge; + wr.opcode = is_read ? URMA_OPC_READ : URMA_OPC_WRITE; + wr.rw.src.sge = is_read ? &r_sge : &l_sge; + wr.rw.src.num_sge = 1; // 按数据流向定 src/dst + wr.rw.dst.sge = is_read ? &l_sge : &r_sge; wr.rw.dst.num_sge = 1; - wr.next = (i + 1 == wr_count) ? nullptr : &wr_list[i + 1]; wr.flag.bs.complete_enable = 1; wr.flag.bs.inline_flag = 0; - // Check if the jetty is in the imported_jetty_map_ auto it = imported_jetty_map_.find(jetty_list_[jetty_index]); - if (it == imported_jetty_map_.end()) { - LOG(ERROR) << "Jetty not imported for endpoint, tjetty is nullptr" - << jetty_index << ", local_nic="; - } - if (it != imported_jetty_map_.end()) { - wr.tjetty = it->second; - } else { - // If not found, use a dummy value - wr.tjetty = nullptr; - } + wr.tjetty = (it != imported_jetty_map_.end()) ? it->second : nullptr; slice->ts = getCurrentTimeInNano(); slice->status = Transport::Slice::POSTED; slice->ub.jetty_depth = &wr_depth_list_[jetty_index]; // Set endpoint pointer for each slice before submitting slice->ub.endpoint = this; - } + }; __sync_fetch_and_add(&wr_depth_list_[jetty_index], wr_count); __sync_fetch_and_add(jfc_outstanding_, wr_count); - if (jetty_list_[jetty_index]->remote_jetty == NULL) { - } - int rc = - urma_post_jetty_send_wr(jetty_list_[jetty_index], wr_list, &bad_wr); - if (rc) { - PLOG(ERROR) << "Failed to urma_post_jetty_send_wr"; - while (bad_wr) { - int i = bad_wr - wr_list; - LOG(ERROR) << "slice (" << i << ") post send failed."; - failed_slice_list.push_back(slice_list[i]); - __sync_fetch_and_sub(&wr_depth_list_[jetty_index], 1); - __sync_fetch_and_sub(jfc_outstanding_, 1); - bad_wr = bad_wr->next; + urma_jfs_wr_t* bad_wr = nullptr; + int rc; + + if (numa_affinity) { + // —— 分支一:chip 亲和,用 bondp_jfs_wr_t 链 —— + static std::atomic numa_log_counter{0}; + if (VLOG_IS_ON(2) && + numa_log_counter.fetch_add(1, std::memory_order_relaxed) % + kNumaAffinitySampleInterval == + 0) { + VLOG(2) << "[numa_affinity] wr_sample nic=" << peer_nic_path_ + << " batch_id=" << slice_list[0]->task->batch_id + << " target_id=" << slice_list[0]->target_id << " opcode=" + << (slice_list[0]->opcode == + Transport::TransferRequest::READ + ? "READ" + : "WRITE") + << " wr_count=" << wr_count + << " src_chip=" << (int)slice_list[0]->ub.src_chip_id + << " dst_chip=" << (int)slice_list[0]->ub.dst_chip_id; + } + bondp_jfs_wr_t wr_list[wr_count]; + memset(wr_list, 0, sizeof(bondp_jfs_wr_t) * wr_count); + for (int i = 0; i < wr_count; ++i) { + fill_common(wr_list[i].base, slice_list[i], l_sge_list[i], + r_sge_list[i]); + wr_list[i].base.next = + (i + 1 == wr_count) ? nullptr : &wr_list[i + 1].base; + // src/dst 均使用远端 segment 所在 chip,避免跨 chip 传输。 + const uint8_t remote_chip_id = slice_list[i]->ub.dst_chip_id; + wr_list[i].src_chip_id = remote_chip_id; + wr_list[i].dst_chip_id = remote_chip_id; + } + rc = urma_post_jetty_send_wr(jetty_list_[jetty_index], &wr_list[0].base, + &bad_wr); + if (rc) { + PLOG(ERROR) << "Failed to urma_post_jetty_send_wr"; + while (bad_wr) { + // 用 user_ctx 直接还原失败 slice —— 不依赖 base 是不是首成员 + failed_slice_list.push_back( + (Transport::Slice*)bad_wr->user_ctx); + __sync_fetch_and_sub(&wr_depth_list_[jetty_index], 1); + __sync_fetch_and_sub(jfc_outstanding_, 1); + bad_wr = bad_wr->next; + } + } + } else { + // —— 分支二:非亲和,保持 Mooncake 现状(urma_jfs_wr_t 链)—— + static std::atomic disabled_logged{false}; + bool expected = false; + if (VLOG_IS_ON(2) && disabled_logged.compare_exchange_strong( + expected, true, std::memory_order_relaxed)) { + VLOG(2) << "[numa_affinity] disabled, plain urma_jfs_wr_t " + "will be used"; + } + urma_jfs_wr_t wr_list[wr_count]; + memset(wr_list, 0, sizeof(urma_jfs_wr_t) * wr_count); + for (int i = 0; i < wr_count; ++i) { + fill_common(wr_list[i], slice_list[i], l_sge_list[i], + r_sge_list[i]); + wr_list[i].next = (i + 1 == wr_count) ? nullptr : &wr_list[i + 1]; + } + rc = urma_post_jetty_send_wr(jetty_list_[jetty_index], &wr_list[0], + &bad_wr); + if (rc) { + PLOG(ERROR) << "Failed to urma_post_jetty_send_wr"; + while (bad_wr) { + int i = bad_wr - wr_list; + failed_slice_list.push_back(slice_list[i]); + __sync_fetch_and_sub(&wr_depth_list_[jetty_index], 1); + __sync_fetch_and_sub(jfc_outstanding_, 1); + bad_wr = bad_wr->next; + } } } slice_list.erase(slice_list.begin(), slice_list.begin() + wr_count); diff --git a/mooncake-transfer-engine/tests/memory_location_test.cpp b/mooncake-transfer-engine/tests/memory_location_test.cpp index 777813fbf3..bf2e043588 100644 --- a/mooncake-transfer-engine/tests/memory_location_test.cpp +++ b/mooncake-transfer-engine/tests/memory_location_test.cpp @@ -129,3 +129,11 @@ TEST(MemoryLocationTest, MallocMultipleNodes) { numa_free(addr, size); } + +TEST(MemoryLocationTest, ParseCpuNumaNode) { + EXPECT_EQ(mooncake::parseCpuNumaNode("cpu:3"), 3); + EXPECT_EQ(mooncake::parseCpuNumaNode("cpu:0"), 0); + EXPECT_EQ(mooncake::parseCpuNumaNode("*"), -1); + EXPECT_EQ(mooncake::parseCpuNumaNode("gpu:0"), -1); + EXPECT_EQ(mooncake::parseCpuNumaNode("cpu:not-a-node"), -1); +}