diff --git a/docs/source/zh_archive/kunpeng_ub_transport.md b/docs/source/zh_archive/kunpeng_ub_transport.md new file mode 100644 index 0000000000..6f41b1ee1b --- /dev/null +++ b/docs/source/zh_archive/kunpeng_ub_transport.md @@ -0,0 +1,90 @@ +# Kunpeng UB Transport +Kunpeng UbTransport源代码路径为Mooncake/mooncake-transfer-engine/src/transport/kunpneg_transport，该路径下有UB协议的Transport对接代码和实现逻辑。 + +## 概述 +UB（Unified Bus，统一总线） 是与RDMA、CXL、NVLink 和TCP处于同一抽象层的传输协议，属于可在应用层灵活选择的传输方案。目前 UB 协议有两个开源实现：URMA（远程内存访问语义）和 OBMM（Load/Store 语义）。 + +URMA（Unified Remote Memory Access，统一远程内存访问）是UB协议为上层应用提供的统一编程抽象与核心语义层。它基于 UB 协议低延迟、高带宽的底层特性，为远程共享内存的访问与操作提供统一的 API 和语义接口。 + +URMA 开源代码仓库：https://atomgit.com/openeuler/umdk + +OBMM (Ownership Based Memory Management) 是面向超节点环境的内核内存管理系统，支持跨节点的物理内存共享。该系统通过内核模块 (obmm.ko) 和用户态库 (libobmm.so) 提供高效的远程内存访问能力。 + +OBMM 开源代码仓库：https://atomgit.com/openeuler/obmm + +## 新增依赖 +Kunpeng UbTransport在Mooncake本身依赖的基础上，新增了一部分URMA和OBMM的依赖： + +- **硬件平台**: 支持原生UB互联架构的鲲鹏950 CPU +- **OS版本**: openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3) +- **URMA依赖**: UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建。 +- **协议优势**: URMA 提供类似 RDMA 的内存语义，针对鲲鹏芯片片上互联进行了优化 + +--- + +## 构建与编译 + +**前置条件** + +- openEuler 24.03 (LTS-SP3) [下载链接](https://www.openeuler.openatom.cn/zh/download/#openEuler%2024.03%20LTS%20SP3) +- 已安装 UMDK: `yum install umdk-urma-devel` 或从[源码](https://atomgit.com/openeuler/umdk)构建 + +**CMake 配置** + +```bash +# 克隆 Mooncake 仓库 +git clone https://github.com/kvcache-ai/Mooncake.git +cd Mooncake + +# 启用 UB 传输层进行配置 +mkdir build && cd build +cmake .. -DUSE_UB=ON \ + -DURMA_INCLUDE_DIR=/usr/include \ + -DURMA_LIBRARY=/usr/lib64/liburma.so + +# 编译 +make -j$(nproc) +``` + +**验证** + +```bash +# 检查 UB 传输层是否已注册 +./mooncake_server --list-transports +# 预期输出: rdma, tcp, nvlink, ub +``` + +--- + +## 运行与测试 + +**单节点基准测试** + +```bash +# 终端 1: 目标端（Target） +./transfer_engine_bench \ + --mode=target \ + --protocol=ub \ + --device_name=urma0 \ + --local_server_name=127.0.0.1 \ + --metadata_server=P2PHANDSHAKE + +# 终端 2: 发起端（Initiator） +./transfer_engine_bench \ + --mode=initiator \ + --protocol=ub \ + --device_name=urma0 \ + --metadata_server=P2PHANDSHAKE \ + --segment_size=8388608 \ + --batch_size=1\ + --segment_id=127.0.0.1:$PORT +``` + +**多设备基准测试** + +```bash +# 自动发现多个 URMA 设备 +./transfer_engine_bench \ + --protocol=ub \ + --device_name=urma0,urma1,urma2,urma3 +``` diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 79bd7be49b..743d13e429 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -73,7 +73,12 @@ option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF) option(USE_CXL "option for using CXL protocol" OFF) option(USE_EFA "option for using AWS EFA transport" OFF) +option(USE_UB "option for using UB protocol transport" OFF) +if (USE_UB) + add_compile_definitions(USE_UB) + message(STATUS "ub transport is enabled") +endif() if (USE_EFA) # Find libfabric headers and library; default to AWS EFA installer path find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 9643d70580..ddcbb726d1 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -231,6 +231,9 @@ static void freeMemoryPool(void *addr, size_t size) { #endif } #else + if (FLAGS_protocol == "ub") { + munmap(addr, size); // for urma + } numa_free(addr, size); #endif } @@ -459,6 +462,9 @@ static Transport *installTransportFromFlags(TransferEngine *engine) { args.get()[0] = const_cast(nic_priority_matrix.c_str()); args.get()[1] = nullptr; xport = engine->installTransport(FLAGS_protocol.c_str(), args.get()); + } else if (FLAGS_protocol == "ub") { + engine->getLocalTopology()->discover({FLAGS_device_name}); + xport = engine->installTransport(FLAGS_protocol, nullptr); } else if (FLAGS_protocol == "efa") { // EFA needs topology discovery to find devices, but auto_discovery // would auto-install RDMA transport. Manually discover instead. diff --git a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp index e29cdee993..adbf87dde5 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp @@ -313,6 +313,9 @@ int initiator() { args[0] = (void *)nic_priority_matrix.c_str(); args[1] = nullptr; xport = engine->installTransport("rdma", args); + } else if (FLAGS_protocol == "ub") { + engine->getLocalTopology()->discover({FLAGS_device_name}); + xport = engine->installTransport(FLAGS_protocol, nullptr); } else if (FLAGS_protocol == "tcp") { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { @@ -412,6 +415,9 @@ int target() { args[0] = (void *)nic_priority_matrix.c_str(); args[1] = nullptr; engine->installTransport("rdma", args); + } else if (FLAGS_protocol == "ub") { + engine->getLocalTopology()->discover({FLAGS_device_name}); + engine->installTransport(FLAGS_protocol, nullptr); } else if (FLAGS_protocol == "tcp") { engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { diff --git a/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp b/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp index 169cde43bb..3944b94608 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp @@ -386,6 +386,9 @@ int target() { args[0] = (void *)nic_priority_matrix.c_str(); args[1] = nullptr; engine->installTransport("rdma", args); + } else if (FLAGS_protocol == "ub") { + engine->getLocalTopology()->discover({FLAGS_device_name}); + xport = engine->installTransport(FLAGS_protocol, nullptr); } else if (FLAGS_protocol == "tcp") { engine->installTransport("tcp", nullptr); } else { diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index 5bc5a1cfcd..df8812706a 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -65,6 +65,13 @@ struct GlobalConfig { int ib_pci_relaxed_ordering_mode = 0; bool ascend_use_fabric_mem = false; bool ascend_agent_mode = false; + // ub config parameters + size_t num_jfc_per_ctx = 2; + size_t num_jfce_per_ctx = 2; + int eid_index = 0; + 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; }; struct RpcCommunicatorConfig { diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index df6e748535..2cf0d6bf70 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -46,16 +46,19 @@ class TransferMetadata { std::string name; uint16_t lid; std::string gid; + std::string eid; // for ub }; struct BufferDesc { std::string name; uint64_t addr; uint64_t length; - std::vector lkey; // for rdma - std::vector rkey; // for rdma - std::string shm_name; // for nvlink and hip - uint64_t offset; // for cxl + std::vector lkey; // for rdma + std::vector rkey; // for rdma + std::string shm_name; // for nvlink and hip + uint64_t offset; // for cxl + std::vector tseg; // for ub/urma + std::vector l_seg_index; // for ub/urma }; struct NVMeoFBufferDesc { @@ -82,7 +85,7 @@ class TransferMetadata { struct SegmentDesc { std::string name; std::string protocol; - // this is for rdma/shm + // this is for rdma/shm/urma std::vector devices; Topology topology; std::vector buffers; @@ -113,6 +116,9 @@ class TransferMetadata { struct HandShakeDesc { std::string local_nic_path; std::string peer_nic_path; +#ifdef USE_UB + std::vector jetty_num; // for ub/urma +#endif #ifdef USE_BAREX uint16_t barex_port; #endif diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h new file mode 100644 index 0000000000..7240ff7bbd --- /dev/null +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -0,0 +1,357 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef UB_CONTEXT_H +#define UB_CONTEXT_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "config.h" +#include "ub_transport.h" + +namespace mooncake { +class UbContext; + +// define worker pool class +class UbWorkerPool { + public: + UbWorkerPool(UbContext& context, int numa_socket_id = 0); + + ~UbWorkerPool(); + + // Add slices to queue, called by Transport + int submitPostSend(const std::vector& slice_list); + + private: + void performPostSend(int thread_id); + + void performPoll(int thread_id); + + void redispatch(std::vector& slice_list, int thread_id); + + void transferWorker(int thread_id); + + void monitorWorker(); + + int doProcessContextEvents(); + + private: + UbContext& context_; + const int numa_socket_id_; + std::vector worker_thread_; + std::atomic workers_running_; + std::atomic suspended_flag_; + std::atomic redispatch_counter_; + std::mutex cond_mutex_; + std::condition_variable cond_var_; + using SliceList = std::vector; + const static int kShardCount = 8; + std::unordered_map slice_queue_[kShardCount]; + std::atomic slice_queue_count_[kShardCount]; + TicketLock slice_queue_lock_[kShardCount]; + std::vector> + collective_slice_queue_; + std::atomic submitted_slice_count_; + std::atomic processed_slice_count_; + uint64_t success_nr_polls = 0, failed_nr_polls = 0; +}; + +class UbEndpointStore { + public: + virtual ~UbEndpointStore() = default; + virtual std::shared_ptr getEndpoint( + const std::string& peer_nic_path) = 0; + virtual std::shared_ptr insertEndpoint( + const std::string& peer_nic_path, UbContext* context) = 0; + virtual int deleteEndpoint(const std::string& peer_nic_path) = 0; + virtual void evictEndpoint() = 0; + virtual void reclaimEndpoint() = 0; + virtual size_t getSize() = 0; + + virtual int destroy() = 0; + virtual int disconnect() = 0; +}; + +// NSDI 24, similar to clock with quick demotion +class UbSIEVEEndpointStore : public UbEndpointStore { + public: + UbSIEVEEndpointStore(size_t max_size) + : waiting_list_len_(0), max_size_(max_size) {} + + std::shared_ptr getEndpoint( + const std::string& peer_nic_path) override; + std::shared_ptr insertEndpoint(const std::string& peer_nic_path, + UbContext* context) override; + int deleteEndpoint(const std::string& peer_nic_path) override; + void evictEndpoint() override; + void reclaimEndpoint() override; + size_t getSize() override; + + int destroy() override; + int disconnect() override; + + private: + RWSpinlock endpoint_map_lock_; + // The bool represents visited + std::unordered_map, std::atomic_bool>> + endpoint_map_; + std::unordered_map::iterator> fifo_map_; + std::list fifo_list_; + + std::optional::iterator> hand_; + + std::unordered_set> waiting_list_; + std::atomic waiting_list_len_; + + size_t max_size_; +}; + +// UbContext class +class UbContext { + public: + UbContext(UbTransport& engine, std::string device_name, int max_endpoints) + : device_name_(std::move(device_name)), + engine_(engine), + max_endpoints_(max_endpoints), + worker_pool_(nullptr), + active_(true), + show_work_request_flushed_error_(false) {} + + virtual ~UbContext() = default; + + int doConstruct(GlobalConfig& config) { + show_work_request_flushed_error_ = globalConfig().trace; + if (construct(config)) { + LOG(INFO) << "failed construct context " << toString(); + return 1; + } + LOG(INFO) << "finish construct context " << toString(); + endpoint_store_ = + std::make_shared(max_endpoints_); + if (endpoint_store_ == nullptr) { + LOG(INFO) << "failed create endpoint store."; + return 1; + } + LOG(INFO) << "finish create endpoint store."; + return 0; + } + + virtual int buildLocalBufferDesc(uint64_t addr, + UbTransport::BufferDesc& buffer_desc) = 0; + + virtual void* localSegWithIndex(unsigned value) = 0; + + virtual std::shared_ptr makeEndpoint() = 0; + + private: + virtual int construct(GlobalConfig& config) = 0; + + virtual int deconstruct() = 0; + + public: + virtual int registerMemoryRegion(uint64_t va, size_t length) = 0; + + virtual int unregisterMemoryRegion(uint64_t va) = 0; + + virtual int doProcessContextEvents() = 0; + + virtual void* retrieveRemoteSeg(const std::string& value) = 0; + + virtual int openDevice(const std::string& device_name, uint8_t port, + int& eid_index) = 0; + + virtual int poll(int num_entries, Transport::Slice** cr, + int jfc_index = 0) = 0; + + virtual volatile int* outstandingCount(int jfc_index) = 0; + + virtual int jfcCount() = 0; + + virtual int submitPostSend( + const std::vector& slice_list) = 0; + + virtual int getAsyncFd() = 0; + + virtual std::string getEid() = 0; + + virtual std::string toString() = 0; + + bool active() const { return active_; } + + void set_active(bool flag) { active_ = flag; } + + // EndPoint Management + std::shared_ptr endpoint() { + return endpoint("LOCAL_SEGMENT_ID"); + } + + std::shared_ptr endpoint(const std::string& peer_nic_path) { + if (!active_) { + LOG(ERROR) << "Context is not active: " << deviceName(); + return nullptr; + } + + if (peer_nic_path.empty()) { + LOG(ERROR) << "Invalid peer NIC path: " << deviceName(); + return nullptr; + } + auto endpoint = endpoint_store_->getEndpoint(peer_nic_path); + if (endpoint) { + return endpoint; + } + + endpoint = endpoint_store_->insertEndpoint(peer_nic_path, this); + endpoint_store_->reclaimEndpoint(); + return endpoint; + } + + int deleteEndpoint(const std::string& peer_nic_path) { + return endpoint_store_->deleteEndpoint(peer_nic_path); + } + + int disconnectAllEndpoints() { return endpoint_store_->disconnect(); } + + // Device name, such as `mlx5_3` + std::string deviceName() const { return device_name_; } + + // NIC Path, such as `192.168.3.76@mlx5_3` + std::string nicPath() const { + return MakeNicPath(engine_.local_server_name_, device_name_); + } + + UbTransport& engine() const { return engine_; } + + uint8_t portNum() const { return port_; } + + int activeSpeed() const { return active_speed_; } + + int eventFd() const { return event_fd_; } + + int socketId() { + std::string path = + "/sys/class/infiniband/" + device_name_ + "/device/numa_node"; + std::ifstream file(path); + if (file.is_open()) { + int socket_id; + file >> socket_id; + file.close(); + return socket_id; + } else { + return 0; + } + } + + static int hexCharToValue(char c) { + if (c >= '0' && c <= '9') return c - '0'; + if (c >= 'A' && c <= 'F') return 10 + c - 'A'; + if (c >= 'a' && c <= 'f') return 10 + c - 'a'; + throw std::invalid_argument("Invalid hexadecimal character"); + } + + static std::string serializeBinaryData(const void* data, size_t length) { + if (!data) { + throw std::invalid_argument("Data pointer cannot be null"); + } + + std::string hexString; + hexString.reserve(length * 2); + + const unsigned char* byteData = static_cast(data); + for (size_t i = 0; i < length; ++i) { + hexString.push_back("0123456789ABCDEF"[(byteData[i] >> 4) & 0x0F]); + hexString.push_back("0123456789ABCDEF"[byteData[i] & 0x0F]); + } + + return hexString; + } + + static void deserializeBinaryData(const std::string& hexString, + std::vector& buffer) { + if (hexString.length() % 2 != 0) { + throw std::invalid_argument("Input string length must be even"); + } + + buffer.clear(); + buffer.reserve(hexString.length() / 2); + + for (size_t i = 0; i < hexString.length(); i += 2) { + int high = hexCharToValue(hexString[i]); + int low = hexCharToValue(hexString[i + 1]); + buffer.push_back(static_cast((high << 4) | low)); + } + } + + protected: + static int joinNonblockingPollList(int& event_fd, int data_fd) { + event_fd = epoll_create1(0); + if (event_fd < 0) { + PLOG(ERROR) << "Failed to create epoll"; + return ERR_CONTEXT; + } + epoll_event event{}; + memset(&event, 0, sizeof(epoll_event)); + + int flags = fcntl(data_fd, F_GETFL, 0); + if (flags == -1) { + PLOG(ERROR) << "Failed to get file descriptor flags"; + return ERR_CONTEXT; + } + if (fcntl(data_fd, F_SETFL, flags | O_NONBLOCK) == -1) { + PLOG(ERROR) << "Failed to set file descriptor nonblocking"; + return ERR_CONTEXT; + } + + event.events = EPOLLIN | EPOLLET; + event.data.fd = data_fd; + if (epoll_ctl(event_fd, EPOLL_CTL_ADD, event.data.fd, &event)) { + PLOG(ERROR) << "Failed to register file descriptor to epoll"; + return ERR_CONTEXT; + } + return 0; + } + + protected: + const std::string device_name_; + UbTransport& engine_; + int max_endpoints_; + + uint8_t port_ = 0; + + std::shared_ptr endpoint_store_; + + int active_speed_ = -1; + + int event_fd_ = -1; + + std::vector background_thread_; + std::atomic threads_running_; + + std::shared_ptr worker_pool_; + + volatile bool active_; + + bool show_work_request_flushed_error_; +}; +} // namespace mooncake + +#endif // UB_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h new file mode 100644 index 0000000000..e95bd25825 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h @@ -0,0 +1,94 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef UB_ENDPOINT_H +#define UB_ENDPOINT_H +#include "common.h" +#include "config.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { +// define the UbEndpoint class +class UbEndPoint { + public: + enum Status { + INITIALIZING, + UNCONNECTED, + CONNECTED, + }; + + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + UbEndPoint() : status_(INITIALIZING), active_(true) {} + + virtual int construct(GlobalConfig& config) = 0; + + virtual int deconstruct() = 0; + + virtual void setPeerNicPath(const std::string& peer_nic_path) = 0; + + virtual int setupConnectionsByActive() = 0; + + virtual int setupConnectionsByPassive(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) = 0; + + virtual bool hasOutstandingSlice() const = 0; + + virtual int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) = 0; + + virtual const std::string toString() const = 0; + + int setupConnectionsByActive(const std::string& peer_nic_path) { + setPeerNicPath(peer_nic_path); + return setupConnectionsByActive(); + } + + bool active() const { return active_; } + + void set_active(bool flag) { + RWSpinlock::WriteGuard guard(lock_); + active_ = flag; + if (!flag) inactive_time_ = getCurrentTimeInNano(); + } + + double inactiveTime() { + if (active_) return 0.0; + return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0; + } + + bool connected() const { + return status_.load(std::memory_order_relaxed) == CONNECTED; + } + + void disconnect() { + RWSpinlock::WriteGuard guard(lock_); + disconnectUnlocked(); + } + + private: + virtual void disconnectUnlocked() = 0; + + protected: + volatile uint64_t inactive_time_{}; + volatile bool active_{}; + std::atomic status_; + RWSpinlock lock_; + std::string peer_nic_path_; +}; +} // namespace mooncake + +#endif // UB_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h new file mode 100644 index 0000000000..223567c439 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h @@ -0,0 +1,123 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef UB_TRANSPORT_H +#define UB_TRANSPORT_H +#include +#include +#include +#include "topology.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { +class UbContext; +class UbEndPoint; +class TransferMetadata; +class UbWorkerpool; + +enum UB_ENDPOINT_TYPE { URMA_ENDPOINT = 0, OBMM_ENDPOINT = 1 }; + +// ub transport supports integration of endpoints that use UB protocol software. +// Currently, two types of endpoints are supported: urma and obmm. +// urma link : https://atomgit.com/openeuler/umdk +// obmm link : https://atomgit.com/openeuler/obmm +class UbTransport : public Transport { + friend class UbContext; + friend class UbEndPoint; + friend class UbWorkerPool; + + public: + using BufferDesc = TransferMetadata::BufferDesc; + using SegmentDesc = TransferMetadata::SegmentDesc; + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + public: + UbTransport(UB_ENDPOINT_TYPE endpoint_type = URMA_ENDPOINT); + + ~UbTransport(); + + int install(std::string& local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) override; + + int registerLocalMemory(void* addr, size_t length, + const std::string& location, bool remote_accessible, + bool update_metadata = true) override; + + int unregisterLocalMemory(void* addr, bool update_metadata = true) override; + + int registerLocalMemoryBatch(const std::vector& buffer_list, + const std::string& location) override; + + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override; + + const char* getName() const override { return "ub"; } + + // TRANSFER + + Status submitTransfer(BatchID batch_id, + const std::vector& entries) override; + + Status submitTransferTask( + const std::vector& task_list) override; + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) override; + + SegmentID getSegmentID(const std::string& segment_name); + + private: + int allocateLocalSegmentID(); + + public: + int onSetupConnections(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc); + + int sendHandshake(const std::string& peer_server_name, + const HandShakeDesc& local_desc, + HandShakeDesc& peer_desc) { + return metadata_->sendHandshake(peer_server_name, local_desc, + peer_desc); + } + + private: + static int init(UbTransport* transport); + + static void uninit(UbTransport* transport); + + static int initializeUbResources(UbTransport* transport); + + static std::shared_ptr buildContext( + UbTransport* transport, const std::string& device_name, + int max_endpoints); + + int startHandshakeDaemon(std::string& local_server_name); + + public: + static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + int& buffer_id, int& device_id, int retry_cnt = 0); + static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + std::string_view hint, int& buffer_id, + int& device_id, int retry_cnt = 0); + + private: + std::vector> context_list_; + std::shared_ptr local_topology_; + UB_ENDPOINT_TYPE endpoint_type_; +}; +} // namespace mooncake + +#endif // UB_TRANSPORT_H diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h new file mode 100644 index 0000000000..5c7f9b9006 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h @@ -0,0 +1,198 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef URMA_ENDPOINT_H +#define URMA_ENDPOINT_H +#include +#include +#include +#include +#include +#include "common.h" +#include "config.h" +#include "ub/umdk/urma/urma_api.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" + +namespace mooncake { +struct UrmaJFC { + UrmaJFC() : native(nullptr), outstanding(0) {} + + urma_jfc_t* native; + volatile int outstanding; +}; + +struct UrmaJFR { + UrmaJFR() : native(nullptr), outstanding(0) {} + + urma_jfr_t* native; + volatile int outstanding; +}; + +static urma_import_seg_flag_t import_flag = { + .bs = {.cacheable = URMA_NON_CACHEABLE, + .access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC, + .mapping = URMA_SEG_NOMAP, + .reserved = 0}}; + +// define the UrmaContext class +class UrmaContext : public UbContext { + friend class UrmaEndpoint; + + public: + UrmaContext(UbTransport& engine, std::string device_name, + int max_endpoints); + ~UrmaContext(); + int registerMemoryRegion(uint64_t va, size_t length) override; + int unregisterMemoryRegion(uint64_t va) override; + int doProcessContextEvents() override; + void* retrieveRemoteSeg(const std::string& value) override; + int poll(int num_entries, Transport::Slice** cr, int jfc_index) override; + volatile int* outstandingCount(int jfc_index) override; + int submitPostSend( + const std::vector& slice_list) override; + int buildLocalBufferDesc(uint64_t addr, + UbTransport::BufferDesc& buffer_desc) override; + void* localSegWithIndex(unsigned value) override; + int jfcCount() override; + int getAsyncFd() override; + std::string getEid() override; + std::string toString() override; + std::shared_ptr makeEndpoint() override; + std::string eid() const; + std::string eid(urma_eid_t eid); + bool transEidFromString(const std::string& eid_str, urma_eid_t& eid); + urma_jfc_t* jfc(); + urma_jfr_t* jfr(); + urma_jfce_t* JFCE(); + static bool uninit(); + static bool init(); + + private: + int construct(GlobalConfig& config) override; + int deconstruct() override; + int openDevice(const std::string& device_name, uint8_t port, + int& eid_index) override; + + urma_target_seg_t* seg(uint64_t addr); + + std::vector& remote_seg_list() { return remote_seg_list_; } + + std::vector& imported_seg_list() { + return imported_seg_list_; + } + + std::vector& local_tseg_list() { + return local_tseg_list_; + } + + void updateUrmaGlobalConfig(urma_device_attr_t& device_attr) { + auto& config = globalConfig(); + if (config.max_ep_per_ctx * config.num_jetty_per_ep > + (size_t)device_attr.dev_cap.max_jetty) { + config.max_ep_per_ctx = + device_attr.dev_cap.max_jetty / config.num_jetty_per_ep; + } + if (config.num_jfc_per_ctx > (size_t)device_attr.dev_cap.max_jfc) { + config.num_jfc_per_ctx = device_attr.dev_cap.max_jfc; + } + } + + private: + std::vector jfc_list_; + urma_token_t urma_token = {.token = 0xACFE}; + urma_context_t* urma_context_ = nullptr; + // ibv_pd *pd_ = nullptr; + uint64_t max_seg_size{}; + urma_mtu active_mtu_; + urma_eid_t eid_{}; + urma_device_attr_t dev_attr_{}; + + int eid_index_ = -1; + int active_speed_ = -1; + + RWSpinlock seg_region_lock_; + std::vector> seg_region_list_; + std::vector local_tseg_list_; + std::vector remote_seg_list_; + std::vector imported_seg_list_; + + std::vector jfr_list_; + + size_t num_JFCE_ = 0; + urma_jfce_t** jfce_ = nullptr; + + std::vector background_thread_; + std::atomic threads_running_; + + std::atomic next_jfce_index_; + std::atomic next_jfce_vector_index_; + std::atomic next_jfc_list_index_; + std::atomic next_jfr_list_index_; + std::vector jfc_r_list_; + + urma_import_seg_flag_t import_flag_ = mooncake::import_flag; + std::unordered_map import_tseg_map; +}; + +// define the UrmaEndpoint class +class UrmaEndpoint : public UbEndPoint { + public: + UrmaEndpoint(UrmaContext* context) + : context_(context), jfc_outstanding_(nullptr) {} + + int construct(GlobalConfig& config) override; + + int deconstruct() override; + + void setPeerNicPath(const std::string& peer_nic_path) override; + + int setupConnectionsByActive() override; + + int setupConnectionsByPassive(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) override; + + bool hasOutstandingSlice() const override; + + int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) override; + + const std::string toString() const override; + + private: + void disconnectUnlocked() override; + + private: + std::vector JettyNum() const; + + int doSetupConnection(const std::string& peer_eid, + std::vector peer_jetty_num_list, + std::string* reply_msg = nullptr); + + int doSetupConnection(int qp_index, const std::string& peer_eid, + uint32_t peer_jetty_num, + std::string* reply_msg = nullptr); + + private: + UrmaContext* context_; + urma_token_t urma_token = {.token = 0xACFE}; + std::vector jetty_list_; + volatile int* wr_depth_list_; + int max_wr_depth_; + volatile int* jfc_outstanding_; + std::unordered_map imported_jetty_map_; +}; +} // namespace mooncake +#endif // URMA_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 704f4c9c78..e3a22c7ef5 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -125,6 +125,14 @@ class Transport { uint32_t retry_cnt; uint32_t max_retry_cnt; } rdma; + struct { + uint64_t dest_addr; + volatile int *jetty_depth; + uint32_t retry_cnt; + uint32_t max_retry_cnt; + void *r_seg; + void *l_seg; + } ub; struct { void *dest_addr; } local; diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index f94ac11e84..6e41ae7f3b 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -125,3 +125,7 @@ if(USE_EFA) message(STATUS "Enabled USE_EFA (AWS Elastic Fabric Adapter) support") target_link_libraries(transfer_engine PUBLIC fabric efa_transport) endif() +if(USE_UB) + message(STATUS "Enabled USE_UB protocol support") + target_link_libraries(transfer_engine PUBLIC ub_transport) +endif() \ No newline at end of file diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index 3c6d42d0b6..5734ad2dcb 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -25,9 +25,12 @@ void loadGlobalConfig(GlobalConfig &config) { const char *num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX"); if (num_cq_per_ctx_env) { int val = atoi(num_cq_per_ctx_env); - if (val > 0 && val < 256) + if (val > 0 && val < 256) { config.num_cq_per_ctx = val; - else + // In URMA, JFC and JFCE are bound one-to-one. + config.num_jfc_per_ctx = val; + config.num_jfce_per_ctx = val; + } else LOG(WARNING) << "Ignore value from environment variable MC_NUM_CQ_PER_CTX"; } diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index f1445ded1f..8ad5fb8bf8 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -55,6 +55,9 @@ #ifdef USE_EFA #include "transport/efa_transport/efa_transport.h" #endif +#ifdef USE_UB +#include "transport/kunpeng_transport/ub_transport.h" +#endif #include @@ -228,6 +231,11 @@ Transport *MultiTransport::installTransport(const std::string &proto, if (std::string(proto) == "rdma") { transport = new RdmaTransport(); } +#ifdef USE_UB + else if (std::string(proto) == "ub") { + transport = new UbTransport(); + } +#endif #ifdef USE_BAREX else if (std::string(proto) == "barex") { transport = new BarexTransport(); diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index 5ddeb80d97..abe026ed03 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -27,12 +27,16 @@ #include #include #include +#include #include #include #include "cuda_alike.h" #include "memory_location.h" #include "topology.h" +#ifdef USE_UB +#include +#endif namespace mooncake { @@ -192,6 +196,112 @@ static std::vector listInfiniBandDevices( return devices; } +#ifdef USE_UB +struct UBDevice { + std::string name; + std::string pci_bus_id; + int numa_node; +}; + +static std::vector listUBDevices( + const std::vector &filter) { + int num_devices = 0; + std::vector devices; + + urma_init_attr_t init_attr = { + .uasid = 0, + }; + if (urma_init(&init_attr) != URMA_SUCCESS) { + LOG(WARNING) << "Failed to urma init"; + return {}; + } + LOG(INFO) << "URMA module init success"; + urma_device_t **device_list = urma_get_device_list(&num_devices); + if (!device_list) { + LOG(WARNING) << "No UB devices found, check your device installation"; + urma_uninit(); + return {}; + } + if (device_list && num_devices <= 0) { + LOG(WARNING) << "No UB devices found, check your device installation"; + urma_free_device_list(device_list); + return {}; + } + + for (int i = 0; i < num_devices; ++i) { + std::string device_name = device_list[i]->name; + if (!filter.empty() && std::find(filter.begin(), filter.end(), + device_name) == filter.end()) + continue; + char path[PATH_MAX + 32]; + char resolved_path[PATH_MAX]; + // Get the PCI bus id for the infiniband device. Note that + snprintf(path, sizeof(path), "/sys/class/ubcore/%s", + device_list[i]->name); + LOG(INFO) << "listUBDevices: path " << path; + if (realpath(path, resolved_path) == NULL) { + LOG(ERROR) << "listUBDevices: realpath " << resolved_path + << " failed"; + continue; + } + LOG(INFO) << "listUBDevices: realpath " << resolved_path; + std::string pci_bus_id = basename(resolved_path); + + int numa_node = -1; + snprintf(path, sizeof(path), "%s/numa", + dirname(dirname(resolved_path))); + LOG(INFO) << "listUBDevices: numanodepath " << path; + std::ifstream(path) >> numa_node; + LOG(INFO) << "UBDevices : performation node ----" + << device_list[i]->name << " : " << numa_node; + + devices.push_back(UBDevice{.name = std::move(device_name), + .pci_bus_id = std::move(pci_bus_id), + .numa_node = numa_node}); + } + urma_free_device_list(device_list); + urma_uninit(); + return devices; +} + +static std::vector discoverCpuTopology( + const std::vector &all_hca) { + DIR *dir = opendir("/sys/devices/system/node"); + struct dirent *entry; + std::vector topology; + + if (dir == NULL) { + PLOG(WARNING) + << "discoverCpuTopology: open /sys/devices/system/node failed"; + return {}; + } + while ((entry = readdir(dir))) { + const char *prefix = "node"; + if (entry->d_type != DT_DIR || + strncmp(entry->d_name, prefix, strlen(prefix)) != 0) { + continue; + } + int node_id = atoi(entry->d_name + strlen(prefix)); + std::vector preferred_hca; + std::vector avail_hca; + // an HCA connected to the same cpu NUMA node is preferred + for (const auto &hca : all_hca) { + if (hca.numa_node == node_id) { + preferred_hca.push_back(hca.name); + } else { + avail_hca.push_back(hca.name); + } + } + topology.push_back( + TopologyEntry{.name = "cpu:" + std::to_string(node_id), + .preferred_hca = std::move(preferred_hca), + .avail_hca = std::move(avail_hca)}); + } + (void)closedir(dir); + return topology; +} +#endif + static std::vector discoverCpuTopology( const std::vector &all_hca) { DIR *dir = opendir("/sys/devices/system/node"); @@ -364,10 +474,17 @@ void Topology::clear() { int Topology::discover(const std::vector &filter) { matrix_.clear(); +#ifdef USE_UB + auto all_hca = listUBDevices(filter); + for (auto &ent : discoverCpuTopology(all_hca)) { + matrix_[ent.name] = ent; + } +#else auto all_hca = listInfiniBandDevices(filter); for (auto &ent : discoverCpuTopology(all_hca)) { matrix_[ent.name] = ent; } +#endif #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) for (auto &ent : discoverCudaTopology(all_hca)) { matrix_[ent.name] = ent; diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index f7c179afb3..c7e2b823fe 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -242,6 +242,15 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, LOG(INFO) << "Topology discovery complete. Found " << local_topology_->getHcaList().size() << " HCAs."; +#ifdef USE_UB + Transport* ub_transport = + multi_transports_->installTransport("ub", local_topology_); + if (!ub_transport) { + LOG(ERROR) << "Failed to install ub transport"; + return -1; + } +#endif + #ifdef USE_ASCEND_HETEROGENEOUS Transport* ascend_transport = multi_transports_->installTransport("ascend", local_topology_); diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index c9d3da8b99..d241c7ff76 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -65,6 +65,15 @@ struct TransferHandshakeUtil { #ifdef USE_EFA root["efa_addr"] = desc.efa_addr; // EFA endpoint address #endif + +#ifdef USE_UB + Json::Value jettyNums(Json::arrayValue); + for (const auto &jetty : desc.jetty_num) jettyNums.append(jetty); + root["jetty_num"] = jettyNums; + LOG(INFO) << "Encode: local_nic_path is " << desc.local_nic_path + << " peer_nic_path is " << desc.peer_nic_path + << " jetty_num size is " << desc.jetty_num.size(); +#endif return root; } @@ -80,6 +89,15 @@ struct TransferHandshakeUtil { #ifdef USE_EFA desc.efa_addr = root["efa_addr"].asString(); // EFA endpoint address #endif + +#ifdef USE_UB + for (const auto &jetty : root["jetty_num"]) { + desc.jetty_num.push_back(jetty.asUInt()); + } + LOG(INFO) << "Decode: remote_nic_path is " << desc.local_nic_path + << " peer_nic_path is " << desc.peer_nic_path + << " jetty_num size is " << desc.jetty_num.size(); +#endif return 0; } }; @@ -198,6 +216,29 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, } segmentJSON["buffers"] = buffersJSON; segmentJSON["priority_matrix"] = desc.topology.toJson(); + } else if (segmentJSON["protocol"] == "ub") { + Json::Value devicesJSON(Json::arrayValue); + for (const auto &device : desc.devices) { + Json::Value deviceJSON; + deviceJSON["name"] = device.name; + deviceJSON["eid"] = device.eid; + devicesJSON.append(deviceJSON); + } + segmentJSON["devices"] = devicesJSON; + + Json::Value buffersJSON(Json::arrayValue); + for (const auto &buffer : desc.buffers) { + Json::Value bufferJSON; + bufferJSON["name"] = buffer.name; + bufferJSON["addr"] = static_cast(buffer.addr); + bufferJSON["length"] = static_cast(buffer.length); + Json::Value tsegJSON(Json::arrayValue); + for (auto &entry : buffer.tseg) tsegJSON.append(entry); + bufferJSON["tseg"] = tsegJSON; + buffersJSON.append(bufferJSON); + } + segmentJSON["buffers"] = buffersJSON; + segmentJSON["priority_matrix"] = desc.topology.toJson(); } else if (segmentJSON["protocol"] == "tcp") { Json::Value buffersJSON(Json::arrayValue); for (const auto &buffer : desc.buffers) { @@ -375,6 +416,42 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, desc->buffers.push_back(buffer); } + int ret = desc->topology.parse( + segmentJSON["priority_matrix"].toStyledString()); + if (ret) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + } + } else if (desc->protocol == "ub") { + for (const auto &deviceJSON : segmentJSON["devices"]) { + DeviceDesc device; + device.name = deviceJSON["name"].asString(); + device.eid = deviceJSON["eid"].asString(); + if (device.name.empty() || device.eid.empty()) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + return nullptr; + } + desc->devices.push_back(device); + } + + for (const auto &bufferJSON : segmentJSON["buffers"]) { + BufferDesc buffer; + buffer.name = bufferJSON["name"].asString(); + buffer.addr = bufferJSON["addr"].asUInt64(); + buffer.length = bufferJSON["length"].asUInt64(); + for (const auto &tsegJSON : bufferJSON["tseg"]) { + buffer.tseg.push_back(tsegJSON.asString()); + } + if (buffer.name.empty() || !buffer.addr || !buffer.length || + buffer.tseg.empty()) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + return nullptr; + } + desc->buffers.push_back(buffer); + } + int ret = desc->topology.parse( segmentJSON["priority_matrix"].toStyledString()); if (ret) { diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 1b72a72159..951c7049f5 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -5,6 +5,11 @@ add_subdirectory(rpc_communicator) add_library(transport OBJECT ${XPORT_SOURCES} $ $) target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread) +if(USE_UB) + add_subdirectory(kunpeng_transport) + target_sources(transport PUBLIC $) +endif() + if (USE_TCP) add_subdirectory(tcp_transport) target_sources(transport PUBLIC $) diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt new file mode 100644 index 0000000000..3514560164 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt @@ -0,0 +1,18 @@ +file(GLOB UB_SOURCES "*.cpp") + +add_library(ub_transport OBJECT ${UB_SOURCES}) + + +target_include_directories(ub_transport + PUBLIC + /usr/include/umdk +) + +target_link_libraries(ub_transport + PUBLIC + /usr/lib64/liburma.so + PRIVATE + JsonCpp::JsonCpp + glog::glog + pthread +) \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp new file mode 100644 index 0000000000..4814a183f8 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -0,0 +1,525 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include "config.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" + +namespace mooncake { +std::shared_ptr UbSIEVEEndpointStore::getEndpoint( + const std::string& peer_nic_path) { + RWSpinlock::ReadGuard guard(endpoint_map_lock_); + auto iter = endpoint_map_.find(peer_nic_path); + if (iter != endpoint_map_.end()) { + iter->second.second.store( + true, std::memory_order_relaxed); // This is safe within read lock + // because of idempotence + return iter->second.first; + } + return nullptr; +} + +std::shared_ptr UbSIEVEEndpointStore::insertEndpoint( + const std::string& peer_nic_path, UbContext* context) { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + if (endpoint_map_.find(peer_nic_path) != endpoint_map_.end()) { + LOG(INFO) << "Endpoint " << peer_nic_path + << " already exists in SIEVEEndpointStore"; + return endpoint_map_[peer_nic_path].first; + } + auto endpoint = context->makeEndpoint(); + if (!endpoint) { + LOG(ERROR) << "Failed to allocate memory for UbEndPoint"; + return nullptr; + } + auto& config = globalConfig(); + int ret = endpoint->construct(config); + if (ret) return nullptr; + + while (this->getSize() >= max_size_) evictEndpoint(); + + endpoint->setPeerNicPath(peer_nic_path); + endpoint_map_[peer_nic_path] = std::make_pair(endpoint, false); + fifo_list_.push_front(peer_nic_path); + fifo_map_[peer_nic_path] = fifo_list_.begin(); + return endpoint; +} + +int UbSIEVEEndpointStore::deleteEndpoint(const std::string& peer_nic_path) { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + auto iter = endpoint_map_.find(peer_nic_path); + if (iter != endpoint_map_.end()) { + waiting_list_len_++; + waiting_list_.insert(iter->second.first); + endpoint_map_.erase(iter); + auto fifo_iter = fifo_map_[peer_nic_path]; + if (hand_.has_value() && hand_.value() == fifo_iter) { + fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt + : hand_ = std::prev(fifo_iter); + } + fifo_list_.erase(fifo_iter); + fifo_map_.erase(peer_nic_path); + } + return 0; +} + +void UbSIEVEEndpointStore::evictEndpoint() { + if (fifo_list_.empty()) { + return; + } + auto o = hand_.has_value() ? hand_.value() : --fifo_list_.end(); + std::string victim; + while (true) { + victim = *o; + if (endpoint_map_[victim].second.load(std::memory_order_relaxed)) { + endpoint_map_[victim].second.store(false, + std::memory_order_relaxed); + o = (o == fifo_list_.begin() ? --fifo_list_.end() : std::prev(o)); + } else { + break; + } + } + hand_ = (o == fifo_list_.begin() ? --fifo_list_.end() : std::prev(o)); + fifo_list_.erase(o); + fifo_map_.erase(victim); + auto victim_instance = endpoint_map_[victim].first; + victim_instance->set_active(false); + waiting_list_len_++; + waiting_list_.insert(victim_instance); + endpoint_map_.erase(victim); + return; +} + +void UbSIEVEEndpointStore::reclaimEndpoint() { + if (waiting_list_len_.load(std::memory_order_relaxed) == 0) return; + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + std::vector> to_delete; + for (auto& endpoint : waiting_list_) + if (!endpoint->hasOutstandingSlice()) to_delete.push_back(endpoint); + for (auto& endpoint : to_delete) waiting_list_.erase(endpoint); + waiting_list_len_ -= to_delete.size(); +} + +int UbSIEVEEndpointStore::destroy() { + for (auto& endpoint : waiting_list_) endpoint->deconstruct(); + for (auto& kv : endpoint_map_) kv.second.first->deconstruct(); + return 0; +} + +int UbSIEVEEndpointStore::disconnect() { + for (auto& endpoint : waiting_list_) endpoint->disconnect(); + for (auto& kv : endpoint_map_) kv.second.first->disconnect(); + return 0; +} + +size_t UbSIEVEEndpointStore::getSize() { return endpoint_map_.size(); } + +const static int kTransferWorkerCount = globalConfig().workers_per_ctx; + +UbWorkerPool::UbWorkerPool(UbContext& context, int numa_socket_id) + : context_(context), + numa_socket_id_(numa_socket_id), + workers_running_(true), + suspended_flag_(0), + redispatch_counter_(0), + submitted_slice_count_(0), + processed_slice_count_(0) { + for (auto& i : slice_queue_count_) i.store(0, std::memory_order_relaxed); + collective_slice_queue_.resize(kTransferWorkerCount); + for (int i = 0; i < kTransferWorkerCount; ++i) { + worker_thread_.emplace_back([this, i] { transferWorker(i); }); + } + worker_thread_.emplace_back([this] { monitorWorker(); }); +} + +UbWorkerPool::~UbWorkerPool() { + if (workers_running_) { + cond_var_.notify_all(); + workers_running_.store(false); + for (auto& entry : worker_thread_) entry.join(); + } +} + +int UbWorkerPool::submitPostSend( + const std::vector& slice_list) { +#ifdef CONFIG_CACHE_SEGMENT_DESC + thread_local uint64_t tl_last_cache_ts = getCurrentTimeInNano(); + thread_local std::unordered_map> + segment_desc_map; + uint64_t current_ts = getCurrentTimeInNano(); + + if (current_ts - tl_last_cache_ts > 1000000000) { + segment_desc_map.clear(); + tl_last_cache_ts = current_ts; + } + + for (auto& slice : slice_list) { + auto target_id = slice->target_id; + if (!segment_desc_map.count(target_id)) { + segment_desc_map[target_id] = + context_.engine().meta()->getSegmentDescByID(target_id); + if (!segment_desc_map[target_id]) { + segment_desc_map.clear(); + LOG(ERROR) << "Cannot get target segment description #" + << target_id; + return ERR_INVALID_ARGUMENT; + } + } + } +#else + std::unordered_map> + segment_desc_map; + for (auto& slice : slice_list) { + auto target_id = slice->target_id; + if (!segment_desc_map.count(target_id)) + segment_desc_map[target_id] = + context_.engine().meta()->getSegmentDescByID(target_id); + } +#endif // CONFIG_CACHE_SEGMENT_DESC + SliceList slice_list_map[kShardCount]; + uint64_t submitted_slice_count = 0; + thread_local std::unordered_map failed_target_ids; + for (auto& slice : slice_list) { + if (failed_target_ids.count(slice->target_id)) { + auto ts = failed_target_ids[slice->target_id]; + if (getCurrentTimeInNano() - ts < 100000000ull) { + slice->markFailed(); + continue; + } else { + failed_target_ids.erase(slice->target_id); + } + } + auto& peer_segment_desc = segment_desc_map[slice->target_id]; + int buffer_id, device_id; + auto hint = globalConfig().enable_dest_device_affinity + ? context_.deviceName() + : ""; + if (UbTransport::selectDevice(peer_segment_desc.get(), + slice->ub.dest_addr, slice->length, hint, + buffer_id, device_id)) { + peer_segment_desc = context_.engine().meta()->getSegmentDescByID( + slice->target_id, true); + if (!peer_segment_desc) { + LOG(ERROR) << "Cannot reload target segment #" + << slice->target_id; + slice->markFailed(); + failed_target_ids[slice->target_id] = getCurrentTimeInNano(); + continue; + } + + if (UbTransport::selectDevice(peer_segment_desc.get(), + slice->ub.dest_addr, slice->length, + hint, buffer_id, device_id)) { + slice->markFailed(); + for (const auto& dev_desc : peer_segment_desc.get()->devices) { + LOG(ERROR) << "peer device : " << dev_desc.name; + } + context_.engine().meta()->dumpMetadataContent( + peer_segment_desc->name, slice->ub.dest_addr, + slice->length); + continue; + } + } + if (!peer_segment_desc) { + slice->markFailed(); + continue; + } + auto targetSegment = + peer_segment_desc->buffers[buffer_id].tseg[device_id]; + slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment); + auto peer_nic_path = + MakeNicPath(peer_segment_desc->name, + peer_segment_desc->devices[device_id].name); + slice->peer_nic_path = peer_nic_path; + int shard_id = (slice->target_id * 10007 + device_id) % kShardCount; + slice_list_map[shard_id].push_back(slice); + submitted_slice_count++; + } + for (int shard_id = 0; shard_id < kShardCount; ++shard_id) { + if (slice_list_map[shard_id].empty()) continue; + slice_queue_lock_[shard_id].lock(); + for (auto& slice : slice_list_map[shard_id]) + slice_queue_[shard_id][slice->peer_nic_path].push_back(slice); + slice_queue_count_[shard_id].fetch_add(slice_list_map[shard_id].size(), + std::memory_order_relaxed); + slice_queue_lock_[shard_id].unlock(); + } + submitted_slice_count_.fetch_add(submitted_slice_count, + std::memory_order_relaxed); + if (suspended_flag_.load(std::memory_order_relaxed)) cond_var_.notify_all(); + return 0; +} + +void UbWorkerPool::performPostSend(int thread_id) { + auto& local_slice_queue = collective_slice_queue_[thread_id]; + for (int shard_id = thread_id; shard_id < kShardCount; + shard_id += kTransferWorkerCount) { + if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) == 0) + continue; + + slice_queue_lock_[shard_id].lock(); + for (auto& entry : slice_queue_[shard_id]) { + for (auto& slice : entry.second) + local_slice_queue[entry.first].push_back(slice); + entry.second.clear(); + } + slice_queue_count_[shard_id].store(0, std::memory_order_relaxed); + slice_queue_lock_[shard_id].unlock(); + } + + // Redispatch slices to other endpoints, for temporary failures + thread_local int tl_redispatch_counter = 0; + if (tl_redispatch_counter < + redispatch_counter_.load(std::memory_order_relaxed)) { + tl_redispatch_counter = + redispatch_counter_.load(std::memory_order_relaxed); + auto local_slice_queue_clone = local_slice_queue; + local_slice_queue.clear(); + for (auto& entry : local_slice_queue_clone) + redispatch(entry.second, thread_id); + return; + } + +#ifdef CONFIG_CACHE_ENDPOINT + thread_local uint64_t tl_last_cache_ts = getCurrentTimeInNano(); + thread_local std::unordered_map> + endpoint_map; + uint64_t current_ts = getCurrentTimeInNano(); + if (current_ts - tl_last_cache_ts > 1000000000) { + endpoint_map.clear(); + tl_last_cache_ts = current_ts; + } +#endif + + SliceList failed_slice_list; + for (auto& entry : local_slice_queue) { + if (entry.second.empty()) continue; + +#ifdef USE_FAKE_POST_SEND + for (auto& slice : entry.second) slice->markSuccess(); + processed_slice_count_.fetch_add(entry.second.size()); + entry.second.clear(); +#else +#ifdef CONFIG_CACHE_ENDPOINT + auto& endpoint = endpoint_map[entry.first]; + if (endpoint == nullptr || !endpoint->active()) + endpoint = context_.endpoint(entry.first); +#else + auto endpoint = context_.endpoint(entry.first); +#endif + if (!endpoint) { + for (auto& slice : entry.second) failed_slice_list.push_back(slice); + entry.second.clear(); + continue; + } + if (!endpoint->active()) { + if (endpoint->inactiveTime() > 1.0) + context_.deleteEndpoint(entry.first); + // enable for re-establishation + for (auto& slice : entry.second) failed_slice_list.push_back(slice); + entry.second.clear(); + continue; + } + if (!endpoint->connected() && endpoint->setupConnectionsByActive()) { + LOG(ERROR) << "Worker: Cannot make connection for endpoint: " + << entry.first << ", mark it inactive"; + for (auto& slice : entry.second) failed_slice_list.push_back(slice); + endpoint->set_active(false); + failed_nr_polls++; + if (context_.active() && failed_nr_polls > 32 && + !success_nr_polls) { + LOG(WARNING) + << "Failed to establish peer endpoints in local RNIC " + << context_.nicPath() << ", mark it inactive"; + context_.set_active(false); + } + entry.second.clear(); + continue; + } + endpoint->submitPostSend(entry.second, failed_slice_list); +#endif + } + + if (!failed_slice_list.empty()) { + for (auto& slice : failed_slice_list) slice->ub.retry_cnt++; + redispatch(failed_slice_list, thread_id); + } +} + +void UbWorkerPool::performPoll(int thread_id) { + int processed_slice_count = 0; + const static size_t kPollCount = 64; + std::unordered_map jetty_depth_set; + for (int jfc_index = thread_id; jfc_index < context_.jfcCount(); + jfc_index += kTransferWorkerCount) { + UbTransport::Slice* cr[kPollCount]; + int nr_poll = context_.poll(kPollCount, cr, jfc_index); + if (nr_poll < 0) { + LOG(ERROR) << "Worker: Failed to poll jetty for complete"; + continue; + } + for (int i = 0; i < nr_poll; ++i) { + UbTransport::Slice* slice = cr[i]; + assert(slice); + if (jetty_depth_set.count(slice->ub.jetty_depth)) + jetty_depth_set[slice->ub.jetty_depth]++; + else + jetty_depth_set[slice->ub.jetty_depth] = 1; + if (cr[i]->status != Transport::Slice::SUCCESS) { + failed_nr_polls++; + if (context_.active() && failed_nr_polls > 32 && + !success_nr_polls) { + LOG(WARNING) << "Too many errors found in local RNIC " + << context_.nicPath() << ", mark it inactive"; + context_.set_active(false); + } + context_.deleteEndpoint(slice->peer_nic_path); + slice->ub.retry_cnt++; + if (slice->ub.retry_cnt >= slice->ub.max_retry_cnt) { + slice->markFailed(); + processed_slice_count_++; + } else { + collective_slice_queue_[thread_id][slice->peer_nic_path] + .push_back(slice); + redispatch_counter_++; + } + } else { + // slice->markSuccess(); + processed_slice_count++; + success_nr_polls++; + } + } + if (nr_poll) + __sync_fetch_and_sub(context_.outstandingCount(jfc_index), nr_poll); + } + + for (auto& entry : jetty_depth_set) + __sync_fetch_and_sub(entry.first, entry.second); + + if (processed_slice_count) + processed_slice_count_.fetch_add(processed_slice_count); +} + +void UbWorkerPool::redispatch(std::vector& slice_list, + int thread_id) { + std::unordered_map> + segment_desc_map; + for (auto& slice : slice_list) { + auto target_id = slice->target_id; + if (!segment_desc_map.count(target_id)) { + segment_desc_map[target_id] = + context_.engine().meta()->getSegmentDescByID(target_id, true); + } + } + for (auto& slice : slice_list) { + if (slice->ub.retry_cnt >= slice->ub.max_retry_cnt) { + slice->markFailed(); + processed_slice_count_++; + } else { + auto& peer_segment_desc = segment_desc_map[slice->target_id]; + int buffer_id, device_id; + if (!peer_segment_desc || + UbTransport::selectDevice( + peer_segment_desc.get(), slice->ub.dest_addr, slice->length, + buffer_id, device_id, slice->ub.retry_cnt)) { + slice->markFailed(); + processed_slice_count_++; + continue; + } + auto targetSegment = + peer_segment_desc->buffers[buffer_id].tseg[device_id]; + slice->ub.r_seg = context_.retrieveRemoteSeg(targetSegment); + auto peer_nic_path = + MakeNicPath(peer_segment_desc->name, + peer_segment_desc->devices[device_id].name); + slice->peer_nic_path = peer_nic_path; + collective_slice_queue_[thread_id][peer_nic_path].push_back(slice); + } + } +} + +void UbWorkerPool::transferWorker(int thread_id) { + thread_local struct { + uint64_t post_send_count = 0; + uint64_t post_send_total_ns = 0; + std::array post_send_buckets{}; + + uint64_t poll_jfc_count = 0; + uint64_t poll_jfc_total_ns = 0; + std::array poll_jfc_buckets{}; + } stats; + bindToSocket(numa_socket_id_); + const static uint64_t kWaitPeriodInNano = 100000000; // 100ms + uint64_t last_wait_ts = getCurrentTimeInNano(); + while (workers_running_.load(std::memory_order_relaxed)) { + auto processed_slice_count = + processed_slice_count_.load(std::memory_order_relaxed); + auto submitted_slice_count = + submitted_slice_count_.load(std::memory_order_relaxed); + if (processed_slice_count == submitted_slice_count) { + uint64_t curr_wait_ts = getCurrentTimeInNano(); + if (curr_wait_ts - last_wait_ts > kWaitPeriodInNano) { + std::unique_lock lock(cond_mutex_); + suspended_flag_.fetch_add(1); + cond_var_.wait_for(lock, std::chrono::seconds(1)); + suspended_flag_.fetch_sub(1); + last_wait_ts = curr_wait_ts; + } + continue; + } + performPostSend(thread_id); +#ifndef USE_FAKE_POST_SEND + performPoll(thread_id); +#endif + } +} + +void UbWorkerPool::monitorWorker() { + bindToSocket(numa_socket_id_); + auto last_reset_ts = getCurrentTimeInNano(); + while (workers_running_) { + auto current_ts = getCurrentTimeInNano(); + if (current_ts - last_reset_ts > 1000000000ll) { + context_.set_active(true); + last_reset_ts = current_ts; + } + struct epoll_event event{}; + int num_events = epoll_wait(context_.eventFd(), &event, 1, 100); + if (num_events < 0) { + if (errno != EWOULDBLOCK && errno != EINTR) + PLOG(ERROR) << "Worker: epoll_wait()"; + continue; + } + + if (num_events == 0) continue; + + if (!(event.events & EPOLLIN)) continue; + + if (event.data.fd == context_.getAsyncFd()) doProcessContextEvents(); + } +} + +int UbWorkerPool::doProcessContextEvents() { + return context_.doProcessContextEvents(); +} +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp new file mode 100644 index 0000000000..d3e9e8b9d9 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -0,0 +1,502 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include "config.h" +#include "memory_location.h" +#include +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_transport.h" +#include "transport/kunpeng_transport/ub_endpoint.h" +#include "transport/kunpeng_transport/urma_endpoint.h" + +namespace mooncake { +UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) + : endpoint_type_(endpoint_type) {} + +UbTransport::~UbTransport() { +#ifdef CONFIG_USE_BATCH_DESC_SET + batch_desc_set_.clear(); +#endif + metadata_->removeSegmentDesc(local_server_name_); + batch_desc_set_.clear(); + context_list_.clear(); +} + +int UbTransport::install(std::string& local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) { + if (topo == nullptr) { + LOG(ERROR) << "UbTransport: missing topology"; + return ERR_INVALID_ARGUMENT; + } + metadata_ = meta; + local_server_name_ = local_server_name; + local_topology_ = topo; + auto ret = initializeUbResources(this); + if (ret) { + LOG(ERROR) << "UbTransport: cannot initialize Ub resources"; + uninit(this); + return ret; + } + LOG(INFO) << "UbTransport: initialize Ub resources done"; + + ret = allocateLocalSegmentID(); + if (ret) { + LOG(ERROR) << "Transfer engine cannot be initialized: cannot " + "allocate local segment"; + uninit(this); + return ret; + } + LOG(INFO) << "Transfer engine allocate local segment done"; + + ret = startHandshakeDaemon(local_server_name); + if (ret) { + LOG(ERROR) << "UbTransport: cannot start handshake daemon"; + uninit(this); + return ret; + } + LOG(INFO) << "UbTransport: start handshake daemon done"; + + ret = metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(ERROR) << "UbTransport: cannot publish segments"; + uninit(this); + return ret; + } + LOG(INFO) << "UbTransport: publish segments done"; + + return 0; +} + +int UbTransport::registerLocalMemory(void* addr, size_t length, + const std::string& name, + bool remote_accessible, + bool update_metadata) { + (void)remote_accessible; + BufferDesc buffer_desc; + for (auto& context : context_list_) { + int ret = context->registerMemoryRegion((uint64_t)addr, length); + if (ret) { + LOG(ERROR) << "UbTransport: cannot register LocalMemory"; + return ret; + } + ret = context->buildLocalBufferDesc((uint64_t)addr, buffer_desc); + if (ret) { + LOG(ERROR) << "UbTransport: build buffer description failed"; + return ret; + } + } + + // Get the memory location automatically after registered MR(pinned), + // when the name is kWildcardLocation("*"). + if (name == kWildcardLocation) { + bool only_first_page = true; + const std::vector entries = + getMemoryLocation(addr, length, only_first_page); + if (entries.empty()) return -1; + buffer_desc.name = entries[0].location; + buffer_desc.addr = (uint64_t)addr; + buffer_desc.length = length; + 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 rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); + if (rc) return rc; + } + + return 0; +} + +int UbTransport::unregisterLocalMemory(void* addr, bool update_metadata) { + int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata); + if (rc) return rc; + for (auto& context : context_list_) + context->unregisterMemoryRegion((uint64_t)addr); + return 0; +} + +int UbTransport::registerLocalMemoryBatch( + const std::vector& buffer_list, const std::string& location) { + std::vector> results; + results.reserve(buffer_list.size()); + for (auto& buffer : buffer_list) { + results.emplace_back( + std::async(std::launch::async, [this, buffer, location]() -> int { + return registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + })); + } + + for (size_t i = 0; i < buffer_list.size(); ++i) { + if (results[i].get()) { + LOG(WARNING) << "UbTransport: Failed to register memory: addr " + << buffer_list[i].addr << " length " + << buffer_list[i].length; + } + } + + return metadata_->updateLocalSegmentDesc(); +} + +int UbTransport::unregisterLocalMemoryBatch( + const std::vector& addr_list) { + std::vector> results; + results.reserve(addr_list.size()); + for (auto& addr : addr_list) { + results.emplace_back( + std::async(std::launch::async, [this, addr]() -> int { + return unregisterLocalMemory(addr, false); + })); + } + + for (size_t i = 0; i < addr_list.size(); ++i) { + if (results[i].get()) + LOG(WARNING) << "UbTransport: Failed to unregister memory: addr " + << addr_list[i]; + } + + return metadata_->updateLocalSegmentDesc(); +} + +Status UbTransport::submitTransfer( + BatchID batch_id, const std::vector& entries) { + auto& batch_desc = *((BatchDesc*)(batch_id)); + if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { + LOG(ERROR) << "UbTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "UbTransport: Exceed the limitation of capacity, batch id: " + + std::to_string(batch_id)); + } + size_t task_id = batch_desc.task_list.size(); + batch_desc.task_list.resize(task_id + entries.size()); + std::vector task_list; + task_list.reserve(batch_desc.task_list.size()); + for (auto& task : batch_desc.task_list) task_list.push_back(&task); + return submitTransferTask(task_list); +} + +Status UbTransport::submitTransferTask( + const std::vector& task_list) { + std::unordered_map, std::vector> + 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; + assert(task.request); + auto& request = *task.request; + auto request_buffer_id = -1, request_device_id = -1; + + if (selectDevice(local_segment_desc.get(), (uint64_t)request.source, + request.length, request_buffer_id, + request_device_id)) { + request_buffer_id = -1; + request_device_id = -1; + } + + 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 = + 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; + bool found_device = false; + if (request_buffer_id >= 0 && request_device_id >= 0) { + found_device = true; + buffer_id = request_buffer_id; + device_id = request_device_id; + } + 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++)) + continue; + assert(device_id >= 0 && + static_cast(device_id) < context_list_.size()); + auto& context = context_list_[device_id]; + assert(context.get()); + if (!context->active()) continue; + assert(buffer_id >= 0 && + static_cast(buffer_id) < + local_segment_desc->buffers.size()); + assert(local_segment_desc->buffers[buffer_id].tseg.size() == + context_list_.size()); + 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); + LOG(ERROR) + << "UbTransport: Address not registered by any device(s) " + << source_addr; + return Status::AddressNotRegistered( + "UbTransport: not registered by any device(s), " + "address: " + + std::to_string(reinterpret_cast(source_addr))); + } + // start to submit batch request task + auto& context = context_list_[device_id]; + if (!context->active()) { + LOG(ERROR) << "Device " << device_id << " is not active"; + 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; + } + } + } + for (auto& entry : slices_to_post) + entry.first->submitPostSend(entry.second); + return Status::OK(); +} + +Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) { + auto& batch_desc = *((BatchDesc*)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + if (task_id >= task_count) { + return Status::InvalidArgument( + "UbTransport::getTransportStatus invalid argument, batch id: " + + std::to_string(batch_id)); + } + auto& task = batch_desc.task_list[task_id]; + status.transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) + status.s = FAILED; + else + status.s = COMPLETED; + task.is_finished = true; + } else { + status.s = WAITING; + } + return Status::OK(); +} + +Transport::SegmentID UbTransport::getSegmentID( + const std::string& segment_name) { + return metadata_->getSegmentID(segment_name); +} + +int UbTransport::allocateLocalSegmentID() { + auto desc = std::make_shared(); + if (!desc) return ERR_MEMORY; + desc->name = local_server_name_; + desc->protocol = "ub"; + for (auto& context : context_list_) { + TransferMetadata::DeviceDesc device_desc; + device_desc.name = context->deviceName(); + device_desc.eid = context->getEid(); + desc->devices.push_back(device_desc); + } + desc->topology = *(local_topology_); + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return 0; +} + +int UbTransport::onSetupConnections(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) { + auto local_nic_name = getNicNameFromNicPath(peer_desc.peer_nic_path); + if (local_nic_name.empty()) return ERR_INVALID_ARGUMENT; + + std::shared_ptr context; + int index = 0; + for (auto& entry : local_topology_->getHcaList()) { + if (entry == local_nic_name) { + context = context_list_[index]; + break; + } + index++; + } + if (!context) return ERR_INVALID_ARGUMENT; + +#ifdef CONFIG_ERDMA + if (context->deleteEndpoint(peer_desc.local_nic_path)) return ERR_ENDPOINT; +#endif + auto endpoint = context->endpoint(peer_desc.local_nic_path); + if (!endpoint) return ERR_ENDPOINT; + return endpoint->setupConnectionsByPassive(peer_desc, local_desc); +} + +int UbTransport::startHandshakeDaemon(std::string& local_server_name) { + return metadata_->startHandshakeDaemon( + std::bind(&UbTransport::onSetupConnections, this, std::placeholders::_1, + std::placeholders::_2), + metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd); +} + +int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + int& buffer_id, int& device_id, int retry_cnt) { + return selectDevice(desc, offset, length, "", buffer_id, device_id, + retry_cnt); +} + +int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + std::string_view hint, int& buffer_id, + int& device_id, int retry_cnt) { + if (desc == nullptr) { + LOG(ERROR) << "UbTransport Get Segment Desc failed"; + return ERR_ADDRESS_NOT_REGISTERED; + } + + const auto& buffers = desc->buffers; + for (buffer_id = 0; buffer_id < static_cast(buffers.size()); + ++buffer_id) { + const auto& buffer = buffers[buffer_id]; + // Check if offset is within buffer range + if (offset < buffer.addr || length > buffer.length || + offset - buffer.addr > buffer.length - length) { + continue; + } + + device_id = + hint.empty() + ? desc->topology.selectDevice(buffer.name, retry_cnt) + : desc->topology.selectDevice(buffer.name, hint, retry_cnt); + if (device_id >= 0) return 0; + device_id = hint.empty() ? desc->topology.selectDevice( + kWildcardLocation, retry_cnt) + : desc->topology.selectDevice( + kWildcardLocation, hint, retry_cnt); + if (device_id >= 0) return 0; + } + return ERR_ADDRESS_NOT_REGISTERED; +} + +int UbTransport::initializeUbResources(UbTransport* t) { + auto ret = init(t); + if (ret != 0) { + LOG(ERROR) << "Failed to init, ret = " << ret; + return -1; + } + auto hca_list = t->local_topology_->getHcaList(); + for (auto& device_name : hca_list) { + auto& config = globalConfig(); + auto max_endpoints = config.max_ep_per_ctx; + auto context = buildContext(t, device_name, max_endpoints); + ret = context->doConstruct(config); + if (ret) { + t->local_topology_->disableDevice(device_name); + LOG(WARNING) << "Disable device " << device_name; + } else { + t->context_list_.push_back(context); + LOG(INFO) << "device " << context->deviceName() << " add to list"; + } + } + if (t->local_topology_->empty()) { + LOG(ERROR) << "UbTransport: No available RNIC"; + return ERR_DEVICE_NOT_FOUND; + } + LOG(INFO) << "ub resources init success"; + return 0; +} + +int UbTransport::init(UbTransport* transport) { + if (transport->endpoint_type_ == URMA_ENDPOINT) { + if (!UrmaContext::init()) { + LOG(ERROR) << "UrmaContext init failed"; + return -1; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmContext not support now."; + return -1; + } else { + LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; + return -1; + } + return 0; +} + +void UbTransport::uninit(UbTransport* transport) { + if (transport->endpoint_type_ == URMA_ENDPOINT) { + if (!UrmaContext::uninit()) { + LOG(ERROR) << "UrmaContext uninit failed"; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmContext not support now."; + } else { + LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; + } +} + +std::shared_ptr UbTransport::buildContext( + UbTransport* t, const std::string& device_name, int max_endpoints) { + if (t->endpoint_type_ == URMA_ENDPOINT) { + auto context = + std::make_shared(*t, device_name, max_endpoints); + if (!context) { + LOG(ERROR) << "UrmaContext build failed"; + return nullptr; + } + return context; + } else if (t->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmContext not support now."; + return nullptr; + } else { + LOG(ERROR) << "invalid endpoint type : " << t->endpoint_type_; + return nullptr; + } +} +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp new file mode 100644 index 0000000000..1d4af10763 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -0,0 +1,959 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include "config.h" +#include "transport/kunpeng_transport/urma_endpoint.h" + +namespace mooncake { +static int isNullEid(urma_eid_t* eid) { + for (int i = 0; i < URMA_EID_SIZE; ++i) { + if (eid->raw[i] != 0) return 0; + } + return 1; +} + +UrmaContext::UrmaContext(UbTransport& engine, std::string device_name, + int max_endpoints) + : UbContext(engine, std::move(device_name), max_endpoints), + next_jfce_index_(0), + next_jfce_vector_index_(0), + next_jfc_list_index_(0), + next_jfr_list_index_(0) {} + +UrmaContext::~UrmaContext() { + auto thisString = toString(); + worker_pool_.reset(); + LOG(INFO) << "destroy worker pool done."; + endpoint_store_->destroy(); + LOG(INFO) << "destroy endpoint store done."; + if (urma_context_) deconstruct(); + LOG(WARNING) << "finished destroy context : " << thisString; +} + +std::string UrmaContext::toString() { + std::ostringstream ss; + ss << "UrmaContext:{device_name : " << device_name_ + << " ,max_endpoints : " << max_endpoints_ + << " ,async_fd : " << getAsyncFd() << " }"; + return ss.str(); +} + +int UrmaContext::getAsyncFd() { return urma_context_->async_fd; } + +int UrmaContext::submitPostSend( + const std::vector& slice_list) { + return worker_pool_->submitPostSend(slice_list); +} + +int UrmaContext::construct(GlobalConfig& config) { + size_t num_jfc_list = config.num_jfc_per_ctx; + size_t num_jfces = config.num_jfce_per_ctx; + int eid_index = config.eid_index; + size_t max_jfc_e = config.max_jfc_e; + // urma: Here, num_jfc_list and num_jfces use the same value, + // meaning one JFC is bound to one JFCE. + // max_jfc_e uses the default value DEFAULT_DEPTH. + if (openDevice(device_name_, port_, eid_index)) { + LOG(ERROR) << "Failed to open device : " << device_name_ + << " with EID index : " << eid_index; + return ERR_CONTEXT; + } + + num_JFCE_ = num_jfces; + jfce_ = new urma_jfce_t*[num_JFCE_]; + for (size_t i = 0; i < num_jfces; ++i) { + jfce_[i] = urma_create_jfce(urma_context_); + if (!jfce_[i]) { + PLOG(ERROR) << "Failed to create jetty for completion events queue " + "on device" + << device_name_; + return ERR_CONTEXT; + } + } + LOG(INFO) << "create jfce done"; + if (joinNonblockingPollList(event_fd_, urma_context_->async_fd)) { + LOG(ERROR) << "Failed to register context async fd to epoll"; + close(event_fd_); + return ERR_CONTEXT; + } + + LOG(INFO) << "join blocking list done"; + jfc_list_.resize(num_jfc_list); + urma_jfc_cfg_t jfc_s_cfg[num_jfc_list] = {}; + for (size_t i = 0; i < num_jfc_list; ++i) { + jfc_s_cfg[i].depth = max_jfc_e; + jfc_s_cfg[i].jfce = NULL; + jfc_s_cfg[i].user_ctx = (uint64_t)&jfc_list_[i].outstanding; + auto jfc = urma_create_jfc(urma_context_, &jfc_s_cfg[i]); + if (!jfc) { + PLOG(ERROR) << "Failed to create jetty for completion queue(jfs)"; + close(event_fd_); + return ERR_CONTEXT; + } + jfc_list_[i].native = jfc; + LOG(INFO) << "create jfc(send) done, jfc id : " << jfc->jfc_id.id; + } + urma_jfc_cfg_t jfc_r_cfg = {}; + jfc_r_cfg.depth = max_jfc_e; + jfc_r_cfg.jfce = NULL; + jfc_r_cfg.user_ctx = 0; + jfc_r_list_.resize(num_jfc_list); + for (size_t i = 0; i < num_jfc_list; ++i) { + auto jfc = urma_create_jfc(urma_context_, &jfc_r_cfg); + if (!jfc) { + PLOG(ERROR) << "Failed to create jetty for completion queue(jfr)"; + close(event_fd_); + return ERR_CONTEXT; + } + jfc_r_list_[i] = jfc; + LOG(INFO) << "create jfc(send) done, jfc id : " << jfc->jfc_id.id; + } + jfr_list_.resize(num_jfc_list); + urma_jfr_cfg_t jfr_cfg[num_jfc_list] = {}; + /* one-side write/read, jfr no used */ + for (size_t i = 0; i < num_jfc_list; ++i) { + jfr_cfg[i].depth = 2048; + jfr_cfg[i].flag.bs.tag_matching = URMA_NO_TAG_MATCHING; + jfr_cfg[i].flag.bs.lock_free = 0; + jfr_cfg[i].trans_mode = URMA_TM_RC; + jfr_cfg[i].min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + jfr_cfg[i].token_value = urma_token; + jfr_cfg[i].id = 0; + jfr_cfg[i].max_sge = 2; + jfr_cfg[i].jfc = jfc_r_list_[i]; + auto jfr = urma_create_jfr(urma_context_, &jfr_cfg[i]); + if (!jfr) { + PLOG(ERROR) << "Failed to create jetty for receive queue"; + close(event_fd_); + return ERR_CONTEXT; + } + jfr_list_[i].native = jfr; + } + LOG(INFO) << "create jfr done"; + LOG(INFO) << "URMA device: " << urma_context_->dev->name + << ", EID: (EID_Index " << eid_index_ << ") " << eid(); + + LOG(INFO) << "context_ == NULL ? " + << (urma_context_ == nullptr ? "TRUE" : "FALSE"); + worker_pool_ = std::make_shared(*this, socketId()); + LOG(INFO) << "create workerpool done"; + return 0; +} + +int UrmaContext::deconstruct() { + for (auto& entry : seg_region_list_) { + int ret = urma_unregister_seg(entry.first); + if (ret) { + PLOG(ERROR) << "Failed to unregister segment"; + } + } + seg_region_list_.clear(); + + for (size_t i = 0; i < jfr_list_.size(); i++) { + if (!jfr_list_[i].native) continue; + + int ret = urma_delete_jfr(jfr_list_[i].native); + if (ret) { + PLOG(ERROR) << "Failed to destroy jetty for receive queue"; + } + } + jfr_list_.clear(); + + for (size_t i = 0; i < jfc_list_.size(); ++i) { + if (!jfc_list_[i].native) continue; + + int ret = urma_delete_jfc(jfc_list_[i].native); + if (ret) { + PLOG(ERROR) << "Failed to destroy jetty for completion queue"; + } + } + jfc_list_.clear(); + + for (size_t i = 0; i < jfc_r_list_.size(); i++) { + if (!jfc_r_list_[i]) continue; + + int ret = urma_delete_jfc(jfc_r_list_[i]); + if (ret) { + PLOG(ERROR) << "Failed to destroy jetty for completion queue"; + } + } + jfc_r_list_.clear(); + + if (event_fd_ >= 0) { + if (close(event_fd_)) LOG(ERROR) << "Failed to close epoll fd"; + event_fd_ = -1; + } + + if (jfce_) { + for (size_t i = 0; i < num_JFCE_; ++i) + if (jfce_[i]) + if (urma_delete_jfce(jfce_[i])) + LOG(ERROR) + << "Failed to destroy jetty for completion event queue"; + delete[] jfce_; + jfce_ = nullptr; + } + + if (urma_context_) { + if (urma_delete_context(urma_context_)) + PLOG(ERROR) << "Failed to close device context"; + urma_context_ = nullptr; + } + + urma_uninit(); + return 0; +} + +urma_target_seg_t* UrmaContext::seg(uint64_t addr) { + RWSpinlock::ReadGuard guard(seg_region_lock_); + for (auto iter = seg_region_list_.begin(); iter != seg_region_list_.end(); + ++iter) + if ((*iter).first->seg.ubva.va <= addr && + addr < (*iter).first->seg.ubva.va + (*iter).second) + return (*iter).first; + + LOG(ERROR) << "Address " << addr << " seg not found for " << deviceName(); + return 0; +} + +int UrmaContext::buildLocalBufferDesc(uint64_t addr, + UbTransport::BufferDesc& buffer_desc) { + auto str = serializeBinaryData(&seg(addr)->seg, sizeof(urma_seg_t)); + auto index = local_tseg_list().size() - 1; + buffer_desc.tseg.push_back(str); + buffer_desc.l_seg_index.push_back(index); + return 0; +} + +void* UrmaContext::localSegWithIndex(unsigned value) { + return local_tseg_list_.at(value); +} + +int UrmaContext::registerMemoryRegion(uint64_t va, size_t length) { + if (length > (size_t)globalConfig().max_seg_size) { + PLOG(WARNING) << "The buffer length exceeds device max_seg_size, " + << "shrink it to " << globalConfig().max_seg_size; + length = (size_t)globalConfig().max_seg_size; + } + LOG(INFO) << "Register memory region " << va << " length " << length; + urma_reg_seg_flag_t flag = {}; + flag.bs.token_policy = URMA_TOKEN_NONE; + flag.bs.cacheable = URMA_NON_CACHEABLE; + flag.bs.access = URMA_ACCESS_READ | URMA_ACCESS_WRITE | URMA_ACCESS_ATOMIC; + flag.bs.token_id_valid = 0; + flag.bs.reserved = 0; + + urma_seg_cfg_t seg_cfg = { + .va = va, + .len = length, + .token_id = NULL, + .token_value = urma_token, + .flag = flag, + .user_ctx = (uintptr_t)NULL, + .iova = 0, + }; + urma_target_seg_t* seg = urma_register_seg(urma_context_, &seg_cfg); + if (!seg) { + PLOG(ERROR) << "Failed to register segment " << seg_cfg.va; + return ERR_CONTEXT; + } + LOG(INFO) << "Local seg token id : " << seg->seg.token_id; + local_tseg_list_.push_back(seg); + + RWSpinlock::WriteGuard guard(seg_region_lock_); + seg_region_list_.emplace_back(seg, length); + return 0; +} + +int UrmaContext::unregisterMemoryRegion(uint64_t addr) { + RWSpinlock::WriteGuard guard(seg_region_lock_); + bool has_removed; + do { + has_removed = false; + for (auto iter = seg_region_list_.begin(); + iter != seg_region_list_.end(); ++iter) { + if ((*iter).first->seg.ubva.va <= addr && + addr < (*iter).first->seg.ubva.va + (*iter).second) { + if (urma_unregister_seg((*iter).first)) { + LOG(ERROR) << "Failed to unregister memory " + << (*iter).first->seg.ubva.va; + return ERR_CONTEXT; + } + seg_region_list_.erase(iter); + has_removed = true; + break; + } + } + } while (has_removed); + return 0; +} + +std::string UrmaContext::getEid() { return eid(); } + +int UrmaContext::doProcessContextEvents() { + urma_async_event_t event; + if (urma_get_async_event(urma_context_, &event) < 0) return ERR_CONTEXT; + LOG(WARNING) << "Worker: Received context async event " << event.event_type + << " for context " << device_name_; + if (event.event_type == URMA_EVENT_JETTY_ERR || + event.event_type == URMA_EVENT_JETTY_LIMIT) { + LOG(WARNING) << "JETTY ERR OR LIMIT" << event.event_type + << device_name_; + } else if (event.event_type == URMA_EVENT_DEV_FATAL || + event.event_type == URMA_EVENT_JFC_ERR || + event.event_type == URMA_EVENT_PORT_DOWN || + event.event_type == URMA_EVENT_EID_CHANGE) { + set_active(false); + disconnectAllEndpoints(); + LOG(INFO) << "Worker: Context " << device_name_ << " is now inactive"; + } else if (event.event_type == URMA_EVENT_PORT_ACTIVE) { + set_active(true); + LOG(INFO) << "Worker: Context " << device_name_ << " is now active"; + } + urma_ack_async_event(&event); + return 0; +} + +void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) { + auto ret = import_tseg_map.find(remoteSegmentStr); + if (ret != import_tseg_map.end()) return ret->second; + std::vector output_buffer; + deserializeBinaryData(remoteSegmentStr, output_buffer); + urma_seg_t* handle; + handle = (urma_seg_t*)malloc(sizeof(urma_seg_t)); + memcpy(handle, output_buffer.data(), sizeof(urma_seg_t)); + remote_seg_list_.push_back(handle); + auto import_tseg = + urma_import_seg(urma_context_, handle, &urma_token, 0, import_flag_); + if (import_tseg == NULL) { + LOG(ERROR) << "Import segment Failed With " << remoteSegmentStr; + free(handle); + return nullptr; + } + imported_seg_list_.push_back(import_tseg); + import_tseg_map[remoteSegmentStr] = import_tseg; + return import_tseg; +} + +int UrmaContext::openDevice(const std::string& device_name, uint8_t port, + int& eid_index) { + int num_devices = 0; + urma_context_t* context = nullptr; + urma_device_t** devices = urma_get_device_list(&num_devices); + urma_eid_info_t* eid_list; + int ret; + uint32_t eid_cnt; + if (!devices) { + LOG(ERROR) << "urma_get_device_list failed"; + return ERR_DEVICE_NOT_FOUND; + } + if (devices && num_devices <= 0) { + LOG(ERROR) << "urma_get_device_list failed"; + urma_free_device_list(devices); + return ERR_DEVICE_NOT_FOUND; + } + LOG(INFO) << "found " << num_devices << " devices."; + for (int i = 0; i < num_devices; ++i) { + if (device_name != devices[i]->name) continue; + + eid_list = urma_get_eid_list(devices[i], &eid_cnt); + if (eid_list == NULL) { + PLOG(ERROR) << "Failed to get eid list, device = " << device_name; + urma_free_device_list(devices); + return ERR_CONTEXT; + } + for (uint32_t j = 0; eid_list != NULL && j < eid_cnt; j++) { + LOG(INFO) << "device_name : " << device_name + << " EID : " << eid(eid_list[j].eid); + } + if (eid_cnt > 0) { + eid_index = eid_list[0].eid_index; + } + eid_ = eid_list[0].eid; + urma_free_eid_list(eid_list); + + context = urma_create_context(devices[i], eid_index); + if (!context) { + LOG(ERROR) << "Urma_create_device context(" << device_name + << ") failed failed"; + urma_free_device_list(devices); + return ERR_CONTEXT; + } + + ret = urma_query_device(devices[i], &dev_attr_); + if (ret) { + PLOG(ERROR) << "Failed to query dev attr( " << device_name << " ) "; + if (urma_delete_context(context)) { + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; + } + urma_free_device_list(devices); + return ERR_CONTEXT; + } + if (dev_attr_.port_cnt != 0 && + dev_attr_.port_attr[port].state != URMA_PORT_ACTIVE) { + LOG(WARNING) << "Device " << device_name << " port( " << port + << " ) not active"; + if (urma_delete_context(context)) { + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; + } + urma_free_device_list(devices); + return ERR_CONTEXT; + } + + updateUrmaGlobalConfig(dev_attr_); + +#ifndef CONFIG_SKIP_NULL_GID_CHECK + if (isNullEid(&eid_)) { + LOG(WARNING) << "GID is NULL, please check your EID index by " + "specifying MC_EID_INDEX"; + if (urma_delete_context(context)) { + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; + } + urma_free_device_list(devices); + return ERR_CONTEXT; + } +#endif // CONFIG_SKIP_NULL_GID_CHECK + + urma_context_ = context; + eid_index_ = eid_index; + port_ = port; + if (dev_attr_.port_cnt != 0) { + active_mtu_ = dev_attr_.port_attr[port].active_mtu; + active_speed_ = dev_attr_.port_attr[port].active_speed; + } else { + active_mtu_ = URMA_MTU_4096; // default mtu and speed + active_speed_ = URMA_SP_100G; + } + + urma_free_device_list(devices); + return 0; + } + + urma_free_device_list(devices); + LOG(ERROR) << "No matched device found: " << device_name; + return ERR_DEVICE_NOT_FOUND; +} + +std::string UrmaContext::eid() const { + std::string eid_str; + char buf[16] = {0}; + const static size_t kEidLength = URMA_EID_SIZE; + for (size_t i = 0; i < kEidLength; ++i) { + sprintf(buf, "%02x", eid_.raw[i]); + eid_str += i == 0 ? buf : std::string(":") + buf; + } + + return eid_str; +} + +std::string UrmaContext::eid(urma_eid_t eid) { + std::string eid_str; + char buf[16] = {0}; + const static size_t kEidLength = URMA_EID_SIZE; + for (size_t i = 0; i < kEidLength; ++i) { + sprintf(buf, "%02x", eid.raw[i]); + eid_str += i == 0 ? buf : std::string(":") + buf; + } + + return eid_str; +} + +bool UrmaContext::transEidFromString(const std::string& eid_str, + urma_eid_t& eid) { + std::stringstream ss(eid_str); + std::string byte_str; + size_t index = 0; + + while (std::getline(ss, byte_str, ':') && index < URMA_EID_SIZE) { + try { + int byte_val = std::stoi(byte_str, nullptr, 16); + eid.raw[index++] = static_cast(byte_val); + } catch (const std::exception&) { + return false; + } + } + + return index == URMA_EID_SIZE; +} + +int UrmaContext::poll(int num_entries, Transport::Slice** slices, + int jfc_index) { + urma_cr_t cr[num_entries]; + int nr_poll = urma_poll_jfc(jfc_list_[jfc_index].native, num_entries, cr); + if (nr_poll < 0) { + LOG(ERROR) << "Failed to poll JFC " << jfc_index << " of device " + << device_name_; + return ERR_CONTEXT; + } + Transport::Slice s[nr_poll]; + for (int i = 0; i < nr_poll; ++i) { + auto slice = (Transport::Slice*)cr[i].user_ctx; + if (cr[i].status == URMA_CR_SUCCESS) { + slice->markSuccess(); + slices[i] = slice; + continue; + } + if (cr[i].status != URMA_CR_WR_FLUSH_ERR || + show_work_request_flushed_error_) + LOG(ERROR) << "Worker: Process failed for slice (opcode: " + << slice->opcode + << ", source_addr: " << slice->source_addr + << ", length: " << slice->length + << ", dest_addr: " << (void*)slice->ub.dest_addr + << ", local_nic: " << deviceName() + << ", peer_nic: " << slice->peer_nic_path + << ", dest_seg_tokenid: " + << static_cast(slice->ub.r_seg) + ->seg.token_id + << ", retry_cnt: " << slice->ub.retry_cnt + << "): " << cr[i].status << ", jfc idx : " << jfc_index + << ", comp_events_acked: " + << jfc_list_[jfc_index].native->comp_events_acked << " " + << jfc_list_[jfc_index].native->async_events_acked; + } + return nr_poll; +} + +volatile int* UrmaContext::outstandingCount(int jfc_index) { + return &jfc_list_[jfc_index].outstanding; +} + +urma_jfc_t* UrmaContext::jfc() { + int index = (next_jfc_list_index_++) % jfc_list_.size(); + return jfc_list_[index].native; +} + +urma_jfr_t* UrmaContext::jfr() { + int index = (next_jfr_list_index_++) % jfr_list_.size(); + return jfr_list_[index].native; +} + +urma_jfce_t* UrmaContext::JFCE() { + int index = (next_jfce_index_++) % num_JFCE_; + return jfce_[index]; +} + +int UrmaContext::jfcCount() { return jfc_list_.size(); } + +bool UrmaContext::uninit() { + urma_uninit(); + return true; +} + +bool UrmaContext::init() { + urma_init_attr_t init_attr = { + .uasid = 0, + }; + auto ret = urma_init(&init_attr); + if (ret != URMA_SUCCESS && ret != URMA_EEXIST) { + LOG(ERROR) << "Failed to urma init, ret = " << ret; + return false; + } + LOG(INFO) << "URMA module init success"; + return true; +} + +// start define the UrmaEndpot method +int UrmaEndpoint::construct(GlobalConfig& config) { + size_t num_jetty_list = config.num_jetty_per_ep; + size_t max_wr_depth = config.max_wr; + if (status_.load(std::memory_order_relaxed) != INITIALIZING) { + LOG(ERROR) << "Endpoint has already been constructed"; + return ERR_ENDPOINT; + } + + jetty_list_.resize(num_jetty_list); + auto* jfc = context_->jfc(); + jfc_outstanding_ = (volatile int*)jfc->jfc_cfg.user_ctx; + + max_wr_depth_ = (int)max_wr_depth; // work request + wr_depth_list_ = new volatile int[num_jetty_list]; + if (!wr_depth_list_) { + LOG(ERROR) << "Failed to allocate memory for work request depth list"; + return ERR_MEMORY; + } + urma_jfs_cfg_t jfs_cfg = { + .depth = 2048, // DEFAULT_DEPTH (512) + .trans_mode = URMA_TM_RC, /* Reliable connection */ + .priority = 15, // URMA_MAX_PRIORITY 15 + .max_sge = 5, // SGE_NUM_MAX 5 + .rnr_retry = 7, // URMA_TYPICAL_RNR_RETRY 7 + .err_timeout = 17, // URMA_TYPICAL_ERR_TIMEOUT 17 + .user_ctx = 0, + }; + urma_jetty_flag_t jetty_flag = {}; + urma_jetty_cfg_t attr; + memset(&attr, 0, sizeof(attr)); + jetty_flag.bs.share_jfr = 1; + attr.flag = jetty_flag; + attr.jfs_cfg = jfs_cfg; + for (size_t i = 0; i < num_jetty_list; ++i) { + wr_depth_list_[i] = 0; + attr.jfs_cfg.jfc = jfc; + // attr.shared.jfc = jfc; + attr.shared.jfr = context_->jfr(); + jetty_list_[i] = urma_create_jetty(context_->urma_context_, &attr); + if (!jetty_list_[i]) { + PLOG(ERROR) << "Failed to create jetty"; + return ERR_ENDPOINT; + } + LOG(INFO) << "Create jetty success, jetty id = " + << jetty_list_[i]->jetty_id.id << " ,jetty jfc id = " + << jetty_list_[i]->jetty_cfg.jfs_cfg.jfc->jfc_id.id << " : " + << jfc->jfc_id.id; + } + + status_.store(UNCONNECTED, std::memory_order_relaxed); + return 0; +} + +int UrmaEndpoint::deconstruct() { + int ret = 0; + for (size_t i = 0; i < jetty_list_.size(); ++i) { + auto imported_it = imported_jetty_map_.find(jetty_list_[i]); + auto imported_jetty = (imported_it != imported_jetty_map_.end()) + ? imported_it->second + : nullptr; + ret = urma_unbind_jetty(jetty_list_[i]); + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; + if (imported_jetty != nullptr) { + ret = urma_unimport_jetty(imported_jetty); + if (ret) PLOG(ERROR) << "Failed to unimport jetty"; + } + // After destroying QP, the wr_depth_list_ won't change + bool displayed = false; + if (wr_depth_list_[i] != 0) { + if (!displayed) { + LOG(WARNING) << "Outstanding work requests found, CQ will not " + "be generated"; + displayed = true; + } + __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); + wr_depth_list_[i] = 0; + } + } + jetty_list_.clear(); + delete[] wr_depth_list_; + return 0; +} + +void UrmaEndpoint::setPeerNicPath(const std::string& peer_nic_path) { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(WARNING) << "Previous connection will be discarded"; + disconnectUnlocked(); + } + peer_nic_path_ = peer_nic_path; +} + +const std::string UrmaEndpoint::toString() const { + auto status = status_.load(std::memory_order_relaxed); + if (status == CONNECTED) + return "EndPoint: local " + context_->nicPath() + ", peer " + + peer_nic_path_; + else + return "EndPoint: local " + context_->nicPath() + " (unconnected)"; +} + +int UrmaEndpoint::setupConnectionsByActive() { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(INFO) << "Connection has been established"; + return 0; + } + + if (context_->nicPath() == peer_nic_path_) { + auto segment_desc = + context_->engine().meta()->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (segment_desc) { + for (auto& nic : segment_desc->devices) { + if (nic.name == context_->deviceName()) { + return doSetupConnection(nic.eid, JettyNum()); + } + } + } + LOG(ERROR) << "Peer NIC " << context_->deviceName() + << " not found in localhost"; + return ERR_DEVICE_NOT_FOUND; + } + + HandShakeDesc local_desc, peer_desc; + local_desc.local_nic_path = context_->nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.jetty_num = JettyNum(); + + auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); + auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); + if (peer_server_name.empty() || peer_nic_name.empty()) { + LOG(ERROR) << "Parse peer nic path failed: " << peer_nic_path_; + return ERR_INVALID_ARGUMENT; + } + + int rc = context_->engine().sendHandshake(peer_server_name, local_desc, + peer_desc); + if (rc) return rc; + if (!peer_desc.reply_msg.empty()) { + LOG(ERROR) << "Reject the handshake request by peer " + << local_desc.peer_nic_path; + return ERR_REJECT_HANDSHAKE; + } + + if (peer_desc.local_nic_path != peer_nic_path_ || + peer_desc.peer_nic_path != local_desc.local_nic_path) { + LOG(ERROR) << "Invalid argument: received packet mismatch" + << ", local.local_nic_path: " << local_desc.local_nic_path + << ", local.peer_nic_path: " << local_desc.peer_nic_path + << ", peer.local_nic_path: " << peer_desc.local_nic_path + << ", peer.peer_nic_path: " << peer_desc.peer_nic_path; + return ERR_REJECT_HANDSHAKE; + } + + auto segment_desc = + context_->engine().meta()->getSegmentDescByName(peer_server_name); + if (segment_desc) { + for (auto& nic : segment_desc->devices) { + if (nic.name == peer_nic_name) { + return doSetupConnection(nic.eid, peer_desc.jetty_num); + } + } + } + LOG(ERROR) << "Peer NIC " << peer_nic_name << " not found in " + << peer_server_name; + return ERR_DEVICE_NOT_FOUND; +} + +void UrmaEndpoint::disconnectUnlocked() { + urma_jetty_attr_t attr; + memset(&attr, 0, sizeof(attr)); + attr.state = URMA_JETTY_STATE_RESET; + + for (size_t i = 0; i < jetty_list_.size(); ++i) { + int ret = urma_modify_jetty(jetty_list_[i], &attr); + if (ret) PLOG(ERROR) << "Failed to modify jetty to RESET"; + auto imported_jetty = imported_jetty_map_[jetty_list_[i]]; + ret = urma_unbind_jetty(jetty_list_[i]); + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; + ret = urma_unimport_jetty(imported_jetty); + if (ret) PLOG(ERROR) << "Failed to unimport jetty"; + // After resetting QP, the wr_depth_list_ won't change + bool displayed = false; + if (wr_depth_list_[i] != 0) { + if (!displayed) { + LOG(WARNING) << "Outstanding work requests found, JFC will not " + "be generated"; + displayed = true; + } + __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); + wr_depth_list_[i] = 0; + } + } + status_.store(UNCONNECTED, std::memory_order_release); +} + +int UrmaEndpoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) { + RWSpinlock::WriteGuard guard(lock_); + if (connected()) { + LOG(WARNING) << "Re-establish connection: " << toString(); + disconnectUnlocked(); + } + + if (peer_desc.peer_nic_path != context_->nicPath() || + peer_desc.local_nic_path != peer_nic_path_) { + local_desc.reply_msg = + "Invalid argument: peer nic path inconsistency, expect " + + context_->nicPath() + " + " + peer_nic_path_ + ", while got " + + peer_desc.peer_nic_path + " + " + peer_desc.local_nic_path; + + LOG(ERROR) << local_desc.reply_msg; + return ERR_REJECT_HANDSHAKE; + } + + auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); + auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); + if (peer_server_name.empty() || peer_nic_name.empty()) { + local_desc.reply_msg = "Parse peer nic path failed: " + peer_nic_path_; + LOG(ERROR) << local_desc.reply_msg; + return ERR_INVALID_ARGUMENT; + } + + local_desc.local_nic_path = context_->nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.jetty_num = JettyNum(); + + auto segment_desc = + context_->engine().meta()->getSegmentDescByName(peer_server_name); + if (segment_desc) { + for (auto& nic : segment_desc->devices) + if (nic.name == peer_nic_name) + return doSetupConnection(nic.eid, peer_desc.jetty_num, + &local_desc.reply_msg); + } + local_desc.reply_msg = + "Peer nic not found in that server: " + peer_nic_path_; + LOG(ERROR) << local_desc.reply_msg; + return ERR_DEVICE_NOT_FOUND; +} + +bool UrmaEndpoint::hasOutstandingSlice() const { + if (active_) return true; + for (size_t i = 0; i < jetty_list_.size(); i++) + if (wr_depth_list_[i] != 0) return true; + return false; +} + +int UrmaEndpoint::submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) { + RWSpinlock::WriteGuard guard(lock_); + if (!active_) return 0; + int jetty_index = SimpleRandom::Get().next(jetty_list_.size()); + int wr_count = std::min(max_wr_depth_ - wr_depth_list_[jetty_index], + (int)slice_list.size()); + wr_count = + 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; + 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]; + 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.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; + wr.tjetty = imported_jetty_map_[jetty_list_[jetty_index]]; + slice->ts = getCurrentTimeInNano(); + slice->status = Transport::Slice::POSTED; + slice->ub.jetty_depth = &wr_depth_list_[jetty_index]; + } + __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; + } + } + slice_list.erase(slice_list.begin(), slice_list.begin() + wr_count); + return 0; +} + +std::vector UrmaEndpoint::JettyNum() const { + std::vector ret; + for (int jetty_index = 0; jetty_index < (int)jetty_list_.size(); + ++jetty_index) + ret.push_back(jetty_list_[jetty_index]->jetty_id.id); + return ret; +} + +int UrmaEndpoint::doSetupConnection(const std::string& peer_eid, + std::vector peer_jetty_num_list, + std::string* reply_msg) { + if (jetty_list_.size() != peer_jetty_num_list.size()) { + std::string message = + "jetty count mismatch in peer and local endpoints, check " + "MC_MAX_EP_PER_CTX"; + LOG(ERROR) << "[Handshake] " << message; + if (reply_msg) *reply_msg = message; + return ERR_INVALID_ARGUMENT; + } + + for (int jetty_index = 0; jetty_index < (int)jetty_list_.size(); + ++jetty_index) { + int ret = doSetupConnection( + jetty_index, peer_eid, peer_jetty_num_list[jetty_index], reply_msg); + if (ret) return ret; + } + + status_.store(CONNECTED, std::memory_order_relaxed); + return 0; +} + +int UrmaEndpoint::doSetupConnection(int jetty_index, + const std::string& peer_eid, + uint32_t peer_jetty_num, + std::string* reply_msg) { + if (jetty_index < 0 || jetty_index >= (int)jetty_list_.size()) + return ERR_INVALID_ARGUMENT; + auto& jetty = jetty_list_[jetty_index]; + urma_eid_t eid; + bool trans_ret = context_->transEidFromString(peer_eid, eid); + if (!trans_ret) { + PLOG(ERROR) << "Invalid peer eid: " << peer_eid; + return ERR_INVALID_ARGUMENT; + } + urma_rjetty_t rjetty = {}; + rjetty.jetty_id.id = peer_jetty_num; + rjetty.jetty_id.eid = eid; + rjetty.trans_mode = URMA_TM_RC; + rjetty.type = URMA_JETTY; + LOG(INFO) << "Peer jetty id = " << peer_jetty_num; + urma_target_jetty_t* imported_jetty = + urma_import_jetty(context_->urma_context_, &rjetty, &urma_token); + urma_status_t ret = urma_bind_jetty(jetty, imported_jetty); + if (ret != URMA_SUCCESS && ret != URMA_EEXIST) { + std::string message = "Failed to bind jetty"; + PLOG(ERROR) << "[Handshake] " << message; + if (reply_msg) *reply_msg = message + ": " + strerror(errno); + urma_unimport_jetty(imported_jetty); + return ERR_ENDPOINT; + } + imported_jetty_map_[jetty] = imported_jetty; + LOG(INFO) << "Bind jetty success, local jetty id:" << jetty->jetty_id.id + << ", remote jetty id:" << peer_jetty_num; + + return 0; +} + +std::shared_ptr UrmaContext::makeEndpoint() { + return std::make_shared(this); +} +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index a8dab57146..47ecbafec9 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -65,6 +65,12 @@ if (USE_EFA) add_test(NAME efa_transport_test COMMAND efa_transport_test) endif() +if (USE_UB) + add_executable(ub_transport_test ${WORKSPACE}/ub_transport_test.cpp) + target_link_libraries(ub_transport_test PUBLIC transfer_engine gtest gtest_main) + add_test(NAME ub_transport_test COMMAND ub_transport_test) +endif() + add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp) target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME transfer_metadata_test COMMAND transfer_metadata_test) diff --git a/mooncake-transfer-engine/tests/ub_transport_test.cpp b/mooncake-transfer-engine/tests/ub_transport_test.cpp new file mode 100644 index 0000000000..29e3474c0b --- /dev/null +++ b/mooncake-transfer-engine/tests/ub_transport_test.cpp @@ -0,0 +1,331 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// How to run: +// etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls +// http://127.0.0.1:2379 +// ./ub_transport_test --mode=target --metadata_server=127.0.0.1:2379 +// --local_server_name=127.0.0.2:12345 --device_name=bonding_dev_0 +// ./ub_transport_test --metadata_server=127.0.0.1:2379 +// --segment_id=127.0.0.2:12345 --local_server_name=127.0.0.3:12346 +// --device_name=bonding_dev_0 + +#include +#include +#include + +#include +#include +#include +#include + +#include "transfer_engine.h" +#include "transport/transport.h" +#include "common.h" + +#include "cuda_alike.h" +#if defined(USE_CUDA) && defined(USE_NVMEOF) +#include +#endif + +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) + +#include + +static void checkCudaError(cudaError_t result, const char* message) { + if (result != cudaSuccess) { + LOG(ERROR) << message << " (Error code: " << result << " - " + << cudaGetErrorString(result) << ")" << std::endl; + exit(EXIT_FAILURE); + } +} +#endif + +#define NR_SOCKETS (1) + +DEFINE_string(local_server_name, mooncake::getHostname(), + "Local server name for segment discovery"); +DEFINE_string(metadata_server, "192.168.3.77:2379", "etcd server host address"); +DEFINE_string(mode, "initiator", + "Running mode: initiator or target. Initiator node read/write " + "data blocks from target node"); +DEFINE_string(operation, "read", "Operation type: read or write"); + +DEFINE_string(protocol, "ub", "Transfer protocol: ub|tcp"); + +DEFINE_string(device_name, "bonding_dev_0", + "Device name to use, valid if protocol=ub"); +DEFINE_string(nic_priority_matrix, "", + "Path to UB NIC priority matrix file (Advanced)"); + +DEFINE_string(segment_id, "192.168.3.76", "Segment ID to access data"); + +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) +DEFINE_bool(use_vram, true, "Allocate memory from GPU VRAM"); +DEFINE_int32(gpu_id, 0, "GPU ID to use"); +#endif + +using namespace mooncake; + +static void* allocateMemoryPool(size_t size, int socket_id, + bool from_vram = false) { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) + if (from_vram) { + int gpu_id = FLAGS_gpu_id; + void* d_buf; + checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); + checkCudaError(cudaMalloc(&d_buf, size), + "Failed to allocate device memory"); + return d_buf; + } +#endif + return numa_alloc_onnode(size, socket_id); +} + +static void freeMemoryPool(void* addr, size_t size) { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) + // check pointer on GPU + cudaPointerAttributes attributes; + checkCudaError(cudaPointerGetAttributes(&attributes, addr), + "Failed to get pointer attributes"); + + if (attributes.type == cudaMemoryTypeDevice) { + cudaFree(addr); + } else if (attributes.type == cudaMemoryTypeHost) { + numa_free(addr, size); + } else { + LOG(ERROR) << "Unknown memory type"; + } +#else + numa_free(addr, size); +#endif +} + +int initiatorWorker(TransferEngine* engine, SegmentID segment_id, int thread_id, + void* addr) { + bindToSocket(0); + auto segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id); + auto remote_base = (uint64_t)segment_desc->buffers[0].addr; + const size_t kDataLength = 4096000; + { + LOG(INFO) << "Stage 1: Write Data"; + for (size_t offset = 0; offset < kDataLength; ++offset) + *((char*)(addr) + offset) = 'a' + lrand48() % 26; + + LOG(INFO) << "Write Data: " << std::string((char*)(addr), 16) << "..."; + + auto batch_id = engine->allocateBatchID(1); + Status s; + + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t*)(addr); + entry.target_id = segment_id; + entry.target_offset = remote_base; + s = engine->submitTransfer(batch_id, {entry}); + LOG_ASSERT(s.ok()); + bool completed = false; + TransferStatus status; + while (!completed) { + Status s = engine->getTransferStatus(batch_id, 0, status); + LOG_ASSERT(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) + completed = true; + else if (status.s == TransferStatusEnum::FAILED) { + LOG(INFO) << "FAILED"; + completed = true; + } + } + s = engine->freeBatchID(batch_id); + LOG_ASSERT(s.ok()); + } + + { + LOG(INFO) << "Stage 2: Read Data"; + auto batch_id = engine->allocateBatchID(1); + Status s; + + TransferRequest entry; + entry.opcode = TransferRequest::READ; + entry.length = kDataLength; + entry.source = (uint8_t*)(addr) + kDataLength; + entry.target_id = segment_id; + entry.target_offset = remote_base; + s = engine->submitTransfer(batch_id, {entry}); + LOG_ASSERT(s.ok()); + bool completed = false; + TransferStatus status; + while (!completed) { + Status s = engine->getTransferStatus(batch_id, 0, status); + LOG_ASSERT(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) + completed = true; + else if (status.s == TransferStatusEnum::FAILED) { + LOG(INFO) << "FAILED"; + completed = true; + } + } + s = engine->freeBatchID(batch_id); + LOG_ASSERT(s.ok()); + } + + int ret = + memcmp((uint8_t*)(addr), (uint8_t*)(addr) + kDataLength, kDataLength); + LOG(INFO) << "Read Data: " << std::string((char*)(addr) + kDataLength, 16) + << "..."; + LOG(INFO) << "Compare: " << (ret == 0 ? "OK" : "FAILED"); + + return 0; +} + +std::string formatDeviceNames(const std::string& device_names) { + std::stringstream ss(device_names); + std::string item; + std::vector tokens; + while (getline(ss, item, ',')) { + tokens.push_back(item); + } + + std::string formatted; + for (size_t i = 0; i < tokens.size(); ++i) { + formatted += "\"" + tokens[i] + "\""; + if (i < tokens.size() - 1) { + formatted += ","; + } + } + return formatted; +} + +std::string loadNicPriorityMatrix() { + if (!FLAGS_nic_priority_matrix.empty()) { + std::ifstream file(FLAGS_nic_priority_matrix); + if (file.is_open()) { + std::string content((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + file.close(); + return content; + } + } + // Build JSON Data + auto device_names = formatDeviceNames(FLAGS_device_name); + return "{\"cpu:0\": [[" + device_names + + "], []], " + " \"cpu:1\": [[" + + device_names + + "], []], " + " \"cuda:0\": [[" + + device_names + + "], []], " + " \"musa:0\": [[" + + device_names + "], []]}"; +} + +int initiator() { + const size_t ram_buffer_size = 1ull << 30; + // disable topology auto discovery for testing. + auto filters = std::vector({FLAGS_device_name}); + auto engine = std::make_unique(true, filters); + + auto hostname_port = parseHostNameWithPort(FLAGS_local_server_name); + engine->init(FLAGS_metadata_server, FLAGS_local_server_name, + hostname_port.first, hostname_port.second); + + Transport* xport = nullptr; + if (FLAGS_protocol == "rdma") { + auto nic_priority_matrix = loadNicPriorityMatrix(); + void** args = (void**)malloc(2 * sizeof(void*)); + args[0] = (void*)nic_priority_matrix.c_str(); + args[1] = nullptr; + xport = engine->installTransport("rdma", args); + } else if (FLAGS_protocol == "ub") { + xport = engine->installTransport("ub", nullptr); + } else if (FLAGS_protocol == "tcp") { + xport = engine->installTransport("tcp", nullptr); + } else { + LOG(ERROR) << "Unsupported protocol"; + } + + LOG_ASSERT(xport); + + void* addr = nullptr; +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) + addr = allocateMemoryPool(ram_buffer_size, 0, FLAGS_use_vram); + std::string name_prefix = FLAGS_use_vram ? GPU_PREFIX : "cpu:"; + int name_suffix = FLAGS_use_vram ? FLAGS_gpu_id : 0; + int rc = engine->registerLocalMemory( + addr, ram_buffer_size, name_prefix + std::to_string(name_suffix)); + LOG_ASSERT(!rc); +#else + addr = allocateMemoryPool(ram_buffer_size, 0, false); + int rc = + engine->registerLocalMemory(addr, ram_buffer_size, kWildcardLocation); + LOG_ASSERT(!rc); +#endif + + auto segment_id = engine->openSegment(FLAGS_segment_id.c_str()); + std::thread workers(initiatorWorker, engine.get(), segment_id, 0, addr); + workers.join(); + engine->unregisterLocalMemory(addr); + freeMemoryPool(addr, ram_buffer_size); + return 0; +} + +int target() { + const size_t ram_buffer_size = 1ull << 30; + // disable topology auto discovery for testing. + auto filters = std::vector({FLAGS_device_name}); + auto engine = std::make_unique(true, filters); + + auto hostname_port = parseHostNameWithPort(FLAGS_local_server_name); + engine->init(FLAGS_metadata_server, FLAGS_local_server_name, + hostname_port.first, hostname_port.second); + + if (FLAGS_protocol == "rdma") { + auto nic_priority_matrix = loadNicPriorityMatrix(); + void** args = (void**)malloc(2 * sizeof(void*)); + args[0] = (void*)nic_priority_matrix.c_str(); + args[1] = nullptr; + engine->installTransport("rdma", args); + } else if (FLAGS_protocol == "ub") { + engine->installTransport("ub", nullptr); + } else if (FLAGS_protocol == "tcp") { + engine->installTransport("tcp", nullptr); + } else { + LOG(ERROR) << "Unsupported protocol"; + } + + void* addr = nullptr; + addr = allocateMemoryPool(ram_buffer_size, 0); + int rc = engine->registerLocalMemory(addr, ram_buffer_size, "cpu:0"); + LOG_ASSERT(!rc); + + while (true) sleep(1); + + engine->unregisterLocalMemory(addr); + freeMemoryPool(addr, ram_buffer_size); + return 0; +} + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + + if (FLAGS_mode == "initiator") + return initiator(); + else if (FLAGS_mode == "target") + return target(); + + LOG(ERROR) << "Unsupported mode: must be 'initiator' or 'target'"; + exit(EXIT_FAILURE); +} \ No newline at end of file