From 5966dabd0a758334bd33e3e8e1be0dfc6eedb7c9 Mon Sep 17 00:00:00 2001 From: zchuango Date: Mon, 16 Mar 2026 22:02:27 +0800 Subject: [PATCH 01/13] add the support for ub transport --- mooncake-common/common.cmake | 5 + .../example/transfer_engine_bench.cpp | 9 +- .../transfer_engine_bench_with_notify.cpp | 6 + .../transfer_engine_bench_with_retry.cpp | 3 + mooncake-transfer-engine/include/config.h | 7 + .../include/transfer_metadata.h | 6 +- .../include/transport/transport.h | 8 + .../transport/ub_transport/ub_context.h | 362 +++++++ .../transport/ub_transport/ub_endpoint.h | 97 ++ .../transport/ub_transport/ub_transport.h | 127 +++ .../transport/ub_transport/urma_endpoint.h | 204 ++++ mooncake-transfer-engine/src/CMakeLists.txt | 4 + mooncake-transfer-engine/src/config.cpp | 5 +- .../src/multi_transport.cpp | 8 + mooncake-transfer-engine/src/topology.cpp | 114 ++ .../src/transfer_engine_impl.cpp | 9 + .../src/transfer_metadata.cpp | 81 +- .../src/transport/CMakeLists.txt | 5 + .../src/transport/ub_transport/CMakeLists.txt | 18 + .../src/transport/ub_transport/ub_context.cpp | 536 ++++++++++ .../transport/ub_transport/ub_transport.cpp | 531 ++++++++++ .../transport/ub_transport/urma_endpoint.cpp | 982 ++++++++++++++++++ mooncake-transfer-engine/tests/CMakeLists.txt | 6 + .../tests/ub_transport_test.cpp | 331 ++++++ 24 files changed, 3459 insertions(+), 5 deletions(-) create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_context.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp create mode 100644 mooncake-transfer-engine/tests/ub_transport_test.cpp diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 318b246474..0b112614b8 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -69,7 +69,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..820101c23f 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -231,7 +231,11 @@ static void freeMemoryPool(void *addr, size_t size) { #endif } #else - numa_free(addr, size); + if (FLAGS_protocol == "ub") { + munmap(addr, size); // for urma + } else { + numa_free(addr, size); + } #endif } @@ -459,6 +463,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 657c8ca3bb..32f1e18600 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -64,6 +64,13 @@ struct GlobalConfig { // ib_pci_relaxed_ordering_mode: 0: off, 1: on if supported, 2: auto int ib_pci_relaxed_ordering_mode = 0; bool ascend_use_fabric_mem = 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暂时使用此默认值 + 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 dae0775757..7cf7cd572d 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -46,6 +46,7 @@ class TransferMetadata { std::string name; uint16_t lid; std::string gid; + std::string eid; // for ub }; struct BufferDesc { @@ -56,6 +57,8 @@ class TransferMetadata { 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 { @@ -81,7 +84,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; @@ -112,6 +115,7 @@ class TransferMetadata { struct HandShakeDesc { std::string local_nic_path; std::string peer_nic_path; + std::vector jetty_num; // for ub/urma #ifdef USE_BAREX uint16_t barex_port; #endif diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index afd0a621cf..01873156fd 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/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h new file mode 100644 index 0000000000..b361f6f828 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h @@ -0,0 +1,362 @@ +// 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::string, std::pair, 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) { + } + + ~UbContext() { + } + + 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< + UbSIEVEEndpointStore>(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_; +}; +} + +#endif // UB_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h new file mode 100644 index 0000000000..e96da7a740 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h @@ -0,0 +1,97 @@ +// 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/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h new file mode 100644 index 0000000000..956f8b4e12 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h @@ -0,0 +1,127 @@ +// 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_; +}; +} // mooncake + +#endif // UB_TRANSPORT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h new file mode 100644 index 0000000000..0bc7894df0 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h @@ -0,0 +1,204 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_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_; +}; +} +#endif // URMA_ENDPOINT_H diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index f05401fc8f..d1d6dc6e10 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -112,3 +112,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 3fa056f765..e70b240c7e 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -25,8 +25,11 @@ 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; + config.num_jfc_per_ctx = val; + config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 + } 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..be8aab48ec 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/ub_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..6c825b1399 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,108 @@ 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"; + 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 : perfer 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 +470,18 @@ void Topology::clear() { int Topology::discover(const std::vector &filter) { matrix_.clear(); +#ifdef USE_UB + auto all_hca = listUBDevices(filter); + // 此处protocol带不进来,使用rdma时重新编译，注释本行，走上面分支 + 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 28c7b824bb..8b0bb2361e 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,7 +216,30 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, } segmentJSON["buffers"] = buffersJSON; segmentJSON["priority_matrix"] = desc.topology.toJson(); - } else if (segmentJSON["protocol"] == "tcp") { + } 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) { Json::Value bufferJSON; @@ -376,7 +417,43 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "tcp") { + } 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) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + } + } else if (desc->protocol == "tcp") { for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 1b72a72159..abd3b23773 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(ub_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/ub_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt new file mode 100644 index 0000000000..3514560164 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_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/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp new file mode 100644 index 0000000000..e0a067da45 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp @@ -0,0 +1,536 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_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; + } + // LOG(INFO) << "Endpoint " << peer_nic_path << " not found in + // SIEVEEndpointStore"; + 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); + // remove endpoint but leaving it status unchanged + // in case it is setting up connection or submitting slice + 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); + // LOG(INFO) << victim << " evicted"; + 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();; +} +} \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp new file mode 100644 index 0000000000..b6fb57c3d6 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -0,0 +1,531 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_transport/ub_transport.h" +#include "transport/ub_transport/ub_endpoint.h" +#include "transport/ub_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 + for (auto &entry : batch_desc_set_) delete entry.second; + 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) << "UrmaTransport: 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) << "UrmaTransport: cannot initialize URMA resources"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: initialize URMA 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) << "UrmaTransport: cannot start handshake daemon"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: start handshake daemon done"; + + ret = metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(ERROR) << "UrmaTransport: cannot publish segments"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: 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) << "UrmaTransport: 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) << "UrmaTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "UrmaTransport: 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) delete s; + LOG(ERROR) + << "UrmaTransport: Address not registered by any device(s) " + << source_addr; + return Status::AddressNotRegistered( + "UrmaTransport: not registered by any device(s), " + "address: " + + std::to_string(reinterpret_cast(source_addr))); + } + // start to submit batch reqeust 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( + "UrmaTransport::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 != URMA_SUCCESS && ret != URMA_EEXIST) { + 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) << "UrmaEndpoint init failed"; + return ERR_INVALID_ARGUMENT; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmEndpoint not support now."; + return ERR_INVALID_ARGUMENT; + } else { + LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; + return ERR_INVALID_ARGUMENT; + } + return 0; +} + +void UbTransport::uninit(UbTransport* transport) { + if (transport->endpoint_type_ == URMA_ENDPOINT) { + if (!UrmaContext::uninit()) { + LOG(ERROR) << "UrmaEndpoint uninit failed"; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmEndpoint 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) << "ObmmEndpoint not support now."; + return nullptr; + } else { + LOG(ERROR) << "invalid endpoint type : " << t->endpoint_type_; + return nullptr; + } +} +} + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp new file mode 100644 index 0000000000..67d2da421a --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp @@ -0,0 +1,982 @@ +// 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/ub_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:此处num_jfc_list和num_jfces使用相同数值，即一个jfc绑定一个jfce，max_jfc_e使用默认值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; + delete 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_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"; + + if (urma_delete_jetty(jetty_list_[i])) { + PLOG(ERROR) << "Failed to destroy jetty"; + return ERR_ENDPOINT; + } + // 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); +} +} diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 2dc89be597..31c232f50c 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -60,6 +60,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..a6b6f9b295 --- /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 From 21e9633c9ef079febece59dc87c75d7f8b93b0e6 Mon Sep 17 00:00:00 2001 From: zchuango Date: Mon, 23 Mar 2026 10:34:38 +0800 Subject: [PATCH 02/13] refine the code format --- .../example/transfer_engine_bench.cpp | 2 +- mooncake-transfer-engine/include/config.h | 2 +- .../include/transfer_metadata.h | 16 +- .../transport/ub_transport/ub_context.h | 61 ++--- .../transport/ub_transport/ub_endpoint.h | 19 +- .../transport/ub_transport/ub_transport.h | 32 +-- .../transport/ub_transport/urma_endpoint.h | 58 ++--- mooncake-transfer-engine/src/config.cpp | 5 +- mooncake-transfer-engine/src/topology.cpp | 19 +- .../src/transfer_metadata.cpp | 66 ++--- .../src/transport/ub_transport/ub_context.cpp | 79 +++--- .../transport/ub_transport/ub_transport.cpp | 114 +++----- .../transport/ub_transport/urma_endpoint.cpp | 245 ++++++++---------- .../tests/ub_transport_test.cpp | 6 +- 14 files changed, 325 insertions(+), 399 deletions(-) diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 820101c23f..76c5781aad 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -232,7 +232,7 @@ static void freeMemoryPool(void *addr, size_t size) { } #else if (FLAGS_protocol == "ub") { - munmap(addr, size); // for urma + munmap(addr, size); // for urma } else { numa_free(addr, size); } diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index 32f1e18600..f24c140beb 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -69,7 +69,7 @@ struct GlobalConfig { size_t num_jfce_per_ctx = 2; int eid_index = 0; uint64_t max_seg_size = 0x10000000000; - size_t max_jfc_e = 4096;// urma暂时使用此默认值 + size_t max_jfc_e = 4096; // urma暂时使用此默认值 size_t num_jetty_per_ep = 1; }; diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 7cf7cd572d..4ff43c62d7 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -46,19 +46,19 @@ class TransferMetadata { std::string name; uint16_t lid; std::string gid; - std::string eid; // for ub + 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 tseg; // for ub/urma - std::vector l_seg_index; // for ub/urma + 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 { @@ -115,7 +115,7 @@ class TransferMetadata { struct HandShakeDesc { std::string local_nic_path; std::string peer_nic_path; - std::vector jetty_num; // for ub/urma + std::vector jetty_num; // for ub/urma #ifdef USE_BAREX uint16_t barex_port; #endif diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h index b361f6f828..f42d93d640 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h @@ -33,7 +33,7 @@ class UbContext; // define worker pool class class UbWorkerPool { -public: + public: UbWorkerPool(UbContext& context, int numa_socket_id = 0); ~UbWorkerPool(); @@ -41,7 +41,7 @@ class UbWorkerPool { // Add slices to queue, called by Transport int submitPostSend(const std::vector& slice_list); -private: + private: void performPostSend(int thread_id); void performPoll(int thread_id); @@ -54,7 +54,7 @@ class UbWorkerPool { int doProcessContextEvents(); -private: + private: UbContext& context_; const int numa_socket_id_; std::vector worker_thread_; @@ -69,14 +69,14 @@ class UbWorkerPool { std::atomic slice_queue_count_[kShardCount]; TicketLock slice_queue_lock_[kShardCount]; std::vector> - collective_slice_queue_; + 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: + public: virtual ~UbEndpointStore() = default; virtual std::shared_ptr getEndpoint( const std::string& peer_nic_path) = 0; @@ -93,15 +93,14 @@ class UbEndpointStore { // NSDI 24, similar to clock with quick demotion class UbSIEVEEndpointStore : public UbEndpointStore { -public: + public: UbSIEVEEndpointStore(size_t max_size) - : waiting_list_len_(0), max_size_(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; + 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; @@ -110,12 +109,12 @@ class UbSIEVEEndpointStore : public UbEndpointStore { int destroy() override; int disconnect() override; -private: + private: RWSpinlock endpoint_map_lock_; // The bool represents visited - std::unordered_map< - std::string, std::pair, std::atomic_bool>> - endpoint_map_; + std::unordered_map, std::atomic_bool>> + endpoint_map_; std::unordered_map::iterator> fifo_map_; std::list fifo_list_; @@ -129,18 +128,16 @@ class UbSIEVEEndpointStore : public UbEndpointStore { // UbContext class class UbContext { -public: + 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) { - } + show_work_request_flushed_error_(false) {} - ~UbContext() { - } + ~UbContext() {} int doConstruct(GlobalConfig& config) { show_work_request_flushed_error_ = globalConfig().trace; @@ -149,8 +146,8 @@ class UbContext { return 1; } LOG(INFO) << "finish construct context " << toString(); - endpoint_store_ = std::make_shared< - UbSIEVEEndpointStore>(max_endpoints_); + endpoint_store_ = + std::make_shared(max_endpoints_); if (endpoint_store_ == nullptr) { LOG(INFO) << "failed create endpoint store."; return 1; @@ -166,12 +163,12 @@ class UbContext { virtual std::shared_ptr makeEndpoint() = 0; -private: + private: virtual int construct(GlobalConfig& config) = 0; virtual int deconstruct() = 0; -public: + public: virtual int registerMemoryRegion(uint64_t va, size_t length) = 0; virtual int unregisterMemoryRegion(uint64_t va) = 0; @@ -190,8 +187,8 @@ class UbContext { virtual int jfcCount() = 0; - virtual int submitPostSend(const std::vector& slice_list) - = 0; + virtual int submitPostSend( + const std::vector& slice_list) = 0; virtual int getAsyncFd() = 0; @@ -232,9 +229,7 @@ class UbContext { return endpoint_store_->deleteEndpoint(peer_nic_path); } - int disconnectAllEndpoints() { - return endpoint_store_->disconnect(); - } + int disconnectAllEndpoints() { return endpoint_store_->disconnect(); } // Device name, such as `mlx5_3` std::string deviceName() const { return device_name_; } @@ -253,8 +248,8 @@ class UbContext { int eventFd() const { return event_fd_; } int socketId() { - std::string path = "/sys/class/infiniband/" + - device_name_ + "/device/numa_node"; + std::string path = + "/sys/class/infiniband/" + device_name_ + "/device/numa_node"; std::ifstream file(path); if (file.is_open()) { int socket_id; @@ -306,7 +301,7 @@ class UbContext { } } -protected: + protected: static int joinNonblockingPollList(int& event_fd, int data_fd) { event_fd = epoll_create1(0); if (event_fd < 0) { @@ -335,7 +330,7 @@ class UbContext { return 0; } -protected: + protected: const std::string device_name_; UbTransport& engine_; int max_endpoints_; @@ -357,6 +352,6 @@ class UbContext { bool show_work_request_flushed_error_; }; -} +} // namespace mooncake #endif // UB_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h index e96da7a740..e95bd25825 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h @@ -22,7 +22,7 @@ namespace mooncake { // define the UbEndpoint class class UbEndPoint { -public: + public: enum Status { INITIALIZING, UNCONNECTED, @@ -31,10 +31,7 @@ class UbEndPoint { using HandShakeDesc = TransferMetadata::HandShakeDesc; - UbEndPoint() - : status_(INITIALIZING), - active_(true) { - } + UbEndPoint() : status_(INITIALIZING), active_(true) {} virtual int construct(GlobalConfig& config) = 0; @@ -49,9 +46,9 @@ class UbEndPoint { virtual bool hasOutstandingSlice() const = 0; - virtual int submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) = 0; + virtual int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) = 0; virtual const std::string toString() const = 0; @@ -82,16 +79,16 @@ class UbEndPoint { disconnectUnlocked(); } -private: + private: virtual void disconnectUnlocked() = 0; -protected: + protected: volatile uint64_t inactive_time_{}; volatile bool active_{}; std::atomic status_; RWSpinlock lock_; std::string peer_nic_path_; }; -} // namespace mooncake +} // namespace mooncake #endif // UB_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h index 956f8b4e12..223567c439 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h @@ -27,10 +27,7 @@ class UbEndPoint; class TransferMetadata; class UbWorkerpool; -enum UB_ENDPOINT_TYPE { - URMA_ENDPOINT = 0, - OBMM_ENDPOINT = 1 -}; +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. @@ -41,12 +38,12 @@ class UbTransport : public Transport { friend class UbEndPoint; friend class UbWorkerPool; -public: + public: using BufferDesc = TransferMetadata::BufferDesc; using SegmentDesc = TransferMetadata::SegmentDesc; using HandShakeDesc = TransferMetadata::HandShakeDesc; -public: + public: UbTransport(UB_ENDPOINT_TYPE endpoint_type = URMA_ENDPOINT); ~UbTransport(); @@ -82,46 +79,45 @@ class UbTransport : public Transport { SegmentID getSegmentID(const std::string& segment_name); -private: + private: int allocateLocalSegmentID(); -public: + 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); + return metadata_->sendHandshake(peer_server_name, local_desc, + peer_desc); } -private: + 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); + static std::shared_ptr buildContext( + UbTransport* transport, const std::string& device_name, + int max_endpoints); int startHandshakeDaemon(std::string& local_server_name); -public: + 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: + private: std::vector> context_list_; std::shared_ptr local_topology_; UB_ENDPOINT_TYPE endpoint_type_; }; -} // mooncake +} // namespace mooncake #endif // UB_TRANSPORT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h index 0bc7894df0..a815c1d549 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h @@ -27,35 +27,30 @@ namespace mooncake { struct UrmaJFC { - UrmaJFC() : native(nullptr), outstanding(0) { - } + UrmaJFC() : native(nullptr), outstanding(0) {} urma_jfc_t* native; volatile int outstanding; }; struct UrmaJFR { - UrmaJFR() : native(nullptr), outstanding(0) { - } + 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 - } -}; + .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: + public: UrmaContext(UbTransport& engine, std::string device_name, int max_endpoints); ~UrmaContext(); @@ -84,7 +79,7 @@ class UrmaContext : public UbContext { static bool uninit(); static bool init(); -private: + private: int construct(GlobalConfig& config) override; int deconstruct() override; int openDevice(const std::string& device_name, uint8_t port, @@ -102,22 +97,23 @@ class UrmaContext : public UbContext { 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; + 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: + private: std::vector jfc_list_; urma_token_t urma_token = {.token = 0xACFE}; urma_context_t* urma_context_ = nullptr; - //ibv_pd *pd_ = nullptr; + // ibv_pd *pd_ = nullptr; uint64_t max_seg_size{}; urma_mtu active_mtu_; urma_eid_t eid_{}; @@ -152,11 +148,9 @@ class UrmaContext : public UbContext { // define the UrmaEndpoint class class UrmaEndpoint : public UbEndPoint { -public: + public: UrmaEndpoint(UrmaContext* context) - : context_(context), - jfc_outstanding_(nullptr) { - } + : context_(context), jfc_outstanding_(nullptr) {} int construct(GlobalConfig& config) override; @@ -171,16 +165,16 @@ class UrmaEndpoint : public UbEndPoint { bool hasOutstandingSlice() const override; - int submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) override; + int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) override; const std::string toString() const override; -private: + private: void disconnectUnlocked() override; -private: + private: std::vector JettyNum() const; int doSetupConnection(const std::string& peer_eid, @@ -191,7 +185,7 @@ class UrmaEndpoint : public UbEndPoint { uint32_t peer_jetty_num, std::string* reply_msg = nullptr); -private: + private: UrmaContext* context_; urma_token_t urma_token = {.token = 0xACFE}; std::vector jetty_list_; @@ -200,5 +194,5 @@ class UrmaEndpoint : public UbEndPoint { volatile int* jfc_outstanding_; std::unordered_map imported_jetty_map_; }; -} -#endif // URMA_ENDPOINT_H +} // namespace mooncake +#endif // URMA_ENDPOINT_H diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index e70b240c7e..53a2acc652 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -28,9 +28,8 @@ void loadGlobalConfig(GlobalConfig &config) { if (val > 0 && val < 256) { config.num_cq_per_ctx = val; config.num_jfc_per_ctx = val; - config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 - } - else + config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 + } else LOG(WARNING) << "Ignore value from environment variable MC_NUM_CQ_PER_CTX"; } diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index 6c825b1399..f0f1f39dd0 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -235,25 +235,28 @@ static std::vector listUBDevices( 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; + 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"; + << " 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))); + snprintf(path, sizeof(path), "%s/numa", + dirname(dirname(resolved_path))); LOG(INFO) << "listUBDevices: numanodepath " << path; std::ifstream(path) >> numa_node; - LOG(INFO) << "UBDevices : perfer node ----" << device_list[i]->name <<" : " << 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}); + .pci_bus_id = std::move(pci_bus_id), + .numa_node = numa_node}); } urma_free_device_list(device_list); urma_uninit(); @@ -276,7 +279,7 @@ static std::vector discoverCpuTopology( 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; diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index 8b0bb2361e..a8c19b1bc2 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -70,9 +70,9 @@ struct TransferHandshakeUtil { 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(); + 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; } @@ -94,9 +94,9 @@ struct TransferHandshakeUtil { 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(); + 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; } @@ -216,30 +216,30 @@ 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") { + } 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) { Json::Value bufferJSON; @@ -417,7 +417,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "ub") { + } else if (desc->protocol == "ub") { for (const auto &deviceJSON : segmentJSON["devices"]) { DeviceDesc device; device.name = deviceJSON["name"].asString(); @@ -443,7 +443,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; return nullptr; - } + } desc->buffers.push_back(buffer); } @@ -453,7 +453,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "tcp") { + } else if (desc->protocol == "tcp") { for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp index e0a067da45..885914553d 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp @@ -29,7 +29,7 @@ std::shared_ptr UbSIEVEEndpointStore::getEndpoint( 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 + true, std::memory_order_relaxed); // This is safe within read lock // because of idempotence return iter->second.first; } @@ -43,7 +43,7 @@ std::shared_ptr UbSIEVEEndpointStore::insertEndpoint( 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"; + << " already exists in SIEVEEndpointStore"; return endpoint_map_[peer_nic_path].first; } auto endpoint = context->makeEndpoint(); @@ -75,9 +75,8 @@ int UbSIEVEEndpointStore::deleteEndpoint(const std::string& peer_nic_path) { 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_iter == fifo_list_.begin() ? hand_ = std::nullopt + : hand_ = std::prev(fifo_iter); } fifo_list_.erase(fifo_iter); fifo_map_.erase(peer_nic_path); @@ -147,15 +146,12 @@ UbWorkerPool::UbWorkerPool(UbContext& context, int numa_socket_id) redispatch_counter_(0), submitted_slice_count_(0), processed_slice_count_(0) { - for (auto& i : slice_queue_count_) - i.store(0, std::memory_order_relaxed); + 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, i] { transferWorker(i); }); } - worker_thread_.emplace_back( - [this] { monitorWorker(); }); + worker_thread_.emplace_back([this] { monitorWorker(); }); } UbWorkerPool::~UbWorkerPool() { @@ -180,7 +176,7 @@ int UbWorkerPool::submitPostSend( tl_last_cache_ts = current_ts; } - for (auto &slice : slice_list) { + for (auto& slice : slice_list) { auto target_id = slice->target_id; if (!segment_desc_map.count(target_id)) { segment_desc_map[target_id] = @@ -194,8 +190,8 @@ int UbWorkerPool::submitPostSend( } } #else - std::unordered_map> + std::unordered_map> segment_desc_map; for (auto& slice : slice_list) { auto target_id = slice->target_id; @@ -222,22 +218,22 @@ int UbWorkerPool::submitPostSend( 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)) { + 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->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)) { + 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; @@ -252,8 +248,8 @@ int UbWorkerPool::submitPostSend( slice->markFailed(); continue; } - auto targetSegment = peer_segment_desc->buffers[buffer_id].tseg[ - device_id]; + 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, @@ -324,12 +320,12 @@ void UbWorkerPool::performPostSend(int thread_id) { if (entry.second.empty()) continue; #ifdef USE_FAKE_POST_SEND - for (auto &slice : entry.second) slice->markSuccess(); + 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]; + auto& endpoint = endpoint_map[entry.first]; if (endpoint == nullptr || !endpoint->active()) endpoint = context_.endpoint(entry.first); #else @@ -350,7 +346,7 @@ void UbWorkerPool::performPostSend(int thread_id) { } if (!endpoint->connected() && endpoint->setupConnectionsByActive()) { LOG(ERROR) << "Worker: Cannot make connection for endpoint: " - << entry.first << ", mark it inactive"; + << entry.first << ", mark it inactive"; for (auto& slice : entry.second) failed_slice_list.push_back(slice); endpoint->set_active(false); failed_nr_polls++; @@ -395,10 +391,10 @@ void UbWorkerPool::performPoll(int thread_id) { 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) { + 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_.nicPath() << ", mark it inactive"; context_.set_active(false); } context_.deleteEndpoint(slice->peer_nic_path); @@ -430,8 +426,8 @@ void UbWorkerPool::performPoll(int thread_id) { void UbWorkerPool::redispatch(std::vector& slice_list, int thread_id) { - std::unordered_map> + std::unordered_map> segment_desc_map; for (auto& slice : slice_list) { auto target_id = slice->target_id; @@ -448,16 +444,15 @@ void UbWorkerPool::redispatch(std::vector& slice_list, 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)) { + 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]; + 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, @@ -479,7 +474,7 @@ void UbWorkerPool::transferWorker(int thread_id) { std::array poll_jfc_buckets{}; } stats; bindToSocket(numa_socket_id_); - const static uint64_t kWaitPeriodInNano = 100000000; // 100ms + 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 = @@ -525,12 +520,12 @@ void UbWorkerPool::monitorWorker() { if (!(event.events & EPOLLIN)) continue; - if (event.data.fd == context_.getAsyncFd()) - doProcessContextEvents(); + if (event.data.fd == context_.getAsyncFd()) doProcessContextEvents(); } } int UbWorkerPool::doProcessContextEvents() { - return context_.doProcessContextEvents();; + return context_.doProcessContextEvents(); + ; } -} \ No newline at end of file +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp index b6fb57c3d6..53128deecb 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -22,13 +22,12 @@ #include "transport/ub_transport/urma_endpoint.h" namespace mooncake { -UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type): endpoint_type_( - endpoint_type) { -} +UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) + : endpoint_type_(endpoint_type) {} UbTransport::~UbTransport() { #ifdef CONFIG_USE_BATCH_DESC_SET - for (auto &entry : batch_desc_set_) delete entry.second; + for (auto& entry : batch_desc_set_) delete entry.second; batch_desc_set_.clear(); #endif metadata_->removeSegmentDesc(local_server_name_); @@ -57,7 +56,7 @@ int UbTransport::install(std::string& local_server_name, ret = allocateLocalSegmentID(); if (ret) { LOG(ERROR) << "Transfer engine cannot be initialized: cannot " - "allocate local segment"; + "allocate local segment"; uninit(this); return ret; } @@ -128,8 +127,7 @@ 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); + context->unregisterMemoryRegion((uint64_t)addr); return 0; } @@ -148,8 +146,8 @@ int UbTransport::registerLocalMemoryBatch( for (size_t i = 0; i < buffer_list.size(); ++i) { if (results[i].get()) { LOG(WARNING) << "UrmaTransport: Failed to register memory: addr " - << buffer_list[i].addr << " length " - << buffer_list[i].length; + << buffer_list[i].addr << " length " + << buffer_list[i].length; } } @@ -170,19 +168,18 @@ int UbTransport::unregisterLocalMemoryBatch( 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]; + << addr_list[i]; } return metadata_->updateLocalSegmentDesc(); } -Status UbTransport::submitTransfer(BatchID batch_id, - const std::vector& - entries) { +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) << "UrmaTransport: Exceed the limitation of current batch's " - "capacity"; + "capacity"; return Status::InvalidArgument( "UrmaTransport: Exceed the limitation of capacity, batch id: " + std::to_string(batch_id)); @@ -231,11 +228,11 @@ Status UbTransport::submitTransferTask( 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->length = + merge_final_slice ? request.length - offset : kBlockSize; slice->opcode = request.opcode; - // LOG(INFO) << "target_offset : " << request.target_offset << ", offset : " << offset; + // 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; @@ -259,15 +256,15 @@ Status UbTransport::submitTransferTask( buffer_id, device_id, retry_cnt++)) continue; assert(device_id >= 0 && - static_cast(device_id) < context_list_.size()); + 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()); + static_cast(buffer_id) < + local_segment_desc->buffers.size()); assert(local_segment_desc->buffers[buffer_id].tseg.size() == - context_list_.size()); + context_list_.size()); found_device = true; break; } @@ -283,16 +280,15 @@ Status UbTransport::submitTransferTask( "address: " + std::to_string(reinterpret_cast(source_addr))); } - // start to submit batch reqeust task + // 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"); + 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]; + 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; @@ -339,8 +335,8 @@ Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, return Status::OK(); } -Transport::SegmentID -UbTransport::getSegmentID(const std::string& segment_name) { +Transport::SegmentID UbTransport::getSegmentID( + const std::string& segment_name) { return metadata_->getSegmentID(segment_name); } @@ -387,8 +383,8 @@ int UbTransport::onSetupConnections(const HandShakeDesc& peer_desc, int UbTransport::startHandshakeDaemon(std::string& local_server_name) { return metadata_->startHandshakeDaemon( - std::bind(&UbTransport::onSetupConnections, this, - std::placeholders::_1, std::placeholders::_2), + std::bind(&UbTransport::onSetupConnections, this, std::placeholders::_1, + std::placeholders::_2), metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd); } @@ -407,8 +403,8 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, } const auto& buffers = desc->buffers; - for (buffer_id = 0; buffer_id < static_cast(buffers.size()); ++ - buffer_id) { + 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 || @@ -416,16 +412,15 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, continue; } - device_id = hint.empty() - ? desc->topology.selectDevice(buffer.name, retry_cnt) - : desc->topology.selectDevice( - buffer.name, hint, retry_cnt); + 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); + 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; @@ -487,11 +482,11 @@ void UbTransport::uninit(UbTransport* transport) { } } -std::shared_ptr UbTransport::buildContext(UbTransport* t, - const std::string& device_name, int max_endpoints) { +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); + auto context = + std::make_shared(*t, device_name, max_endpoints); if (!context) { LOG(ERROR) << "UrmaContext build failed"; return nullptr; @@ -505,27 +500,4 @@ std::shared_ptr UbTransport::buildContext(UbTransport* t, return nullptr; } } -} - - - - - - - - - - - - - - - - - - - - - - - +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp index 67d2da421a..2a79d13f10 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp @@ -32,8 +32,7 @@ UrmaContext::UrmaContext(UbTransport& engine, std::string device_name, next_jfce_index_(0), next_jfce_vector_index_(0), next_jfc_list_index_(0), - next_jfr_list_index_(0) { -} + next_jfr_list_index_(0) {} UrmaContext::~UrmaContext() { auto thisString = toString(); @@ -48,15 +47,12 @@ UrmaContext::~UrmaContext() { std::string UrmaContext::toString() { std::ostringstream ss; ss << "UrmaContext:{device_name : " << device_name_ - << " ,max_endpoints : " << max_endpoints_ - << " ,async_fd : " << getAsyncFd() - << " }"; + << " ,max_endpoints : " << max_endpoints_ + << " ,async_fd : " << getAsyncFd() << " }"; return ss.str(); } -int UrmaContext::getAsyncFd() { - return urma_context_->async_fd; -} +int UrmaContext::getAsyncFd() { return urma_context_->async_fd; } int UrmaContext::submitPostSend( const std::vector& slice_list) { @@ -70,8 +66,8 @@ int UrmaContext::construct(GlobalConfig& config) { size_t max_jfc_e = config.max_jfc_e; // urma:此处num_jfc_list和num_jfces使用相同数值，即一个jfc绑定一个jfce，max_jfc_e使用默认值DEFAULT_DEPTH if (openDevice(device_name_, port_, eid_index)) { - LOG(ERROR) << "Failed to open device : " << device_name_ << - " with EID index : " << eid_index; + LOG(ERROR) << "Failed to open device : " << device_name_ + << " with EID index : " << eid_index; return ERR_CONTEXT; } @@ -80,9 +76,9 @@ int UrmaContext::construct(GlobalConfig& config) { 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_; + PLOG(ERROR) << "Failed to create jetty for completion events queue " + "on device" + << device_name_; return ERR_CONTEXT; } } @@ -147,11 +143,10 @@ int UrmaContext::construct(GlobalConfig& config) { } LOG(INFO) << "create jfr done"; LOG(INFO) << "URMA device: " << urma_context_->dev->name - << ", EID: (EID_Index " << eid_index_ << ") " << eid(); + << ", EID: (EID_Index " << eid_index_ << ") " << eid(); - LOG(INFO) << "context_ == NULL ? " << (urma_context_ == nullptr - ? "TRUE" - : "FALSE"); + LOG(INFO) << "context_ == NULL ? " + << (urma_context_ == nullptr ? "TRUE" : "FALSE"); worker_pool_ = std::make_shared(*this, socketId()); LOG(INFO) << "create workerpool done"; return 0; @@ -197,8 +192,7 @@ int UrmaContext::deconstruct() { jfc_r_list_.clear(); if (event_fd_ >= 0) { - if (close(event_fd_)) - LOG(ERROR) << "Failed to close epoll fd"; + if (close(event_fd_)) LOG(ERROR) << "Failed to close epoll fd"; event_fd_ = -1; } @@ -206,8 +200,8 @@ int UrmaContext::deconstruct() { 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"; + LOG(ERROR) + << "Failed to destroy jetty for completion event queue"; delete[] jfce_; jfce_ = nullptr; } @@ -224,8 +218,8 @@ int UrmaContext::deconstruct() { 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) + 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; @@ -250,7 +244,7 @@ void* UrmaContext::localSegWithIndex(unsigned 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; + << "shrink it to " << globalConfig().max_seg_size; length = (size_t)globalConfig().max_seg_size; } LOG(INFO) << "Register memory region " << va << " length " << length; @@ -293,8 +287,8 @@ int UrmaContext::unregisterMemoryRegion(uint64_t addr) { 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; + LOG(ERROR) << "Failed to unregister memory " + << (*iter).first->seg.ubva.va; return ERR_CONTEXT; } seg_region_list_.erase(iter); @@ -306,32 +300,27 @@ int UrmaContext::unregisterMemoryRegion(uint64_t addr) { return 0; } -std::string UrmaContext::getEid() { - return eid(); -} +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_; + 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"; + 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"; + LOG(INFO) << "Worker: Context " << device_name_ << " is now active"; } urma_ack_async_event(&event); return 0; @@ -346,8 +335,8 @@ void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) { 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_); + 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; delete handle; @@ -386,8 +375,8 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, 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); + LOG(INFO) << "device_name : " << device_name + << " EID : " << eid(eid_list[j].eid); } if (eid_cnt > 0) { eid_index = eid_list[0].eid_index; @@ -397,8 +386,8 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, context = urma_create_context(devices[i], eid_index); if (!context) { - LOG(ERROR) << "Urma_create_device context(" << device_name << - ") failed failed"; + LOG(ERROR) << "Urma_create_device context(" << device_name + << ") failed failed"; urma_free_device_list(devices); return ERR_CONTEXT; } @@ -407,19 +396,19 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, if (ret) { PLOG(ERROR) << "Failed to query dev attr( " << device_name << " ) "; if (urma_delete_context(context)) { - PLOG(ERROR) << "urma_delete_context(" << device_name << - ") failed"; + 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 (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"; + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; } urma_free_device_list(devices); return ERR_CONTEXT; @@ -430,10 +419,10 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, #ifndef CONFIG_SKIP_NULL_GID_CHECK if (isNullEid(&eid_)) { LOG(WARNING) << "GID is NULL, please check your EID index by " - "specifying MC_EID_INDEX"; + "specifying MC_EID_INDEX"; if (urma_delete_context(context)) { - PLOG(ERROR) << "urma_delete_context(" << device_name << - ") failed"; + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; } urma_free_device_list(devices); return ERR_CONTEXT; @@ -447,7 +436,7 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, 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_mtu_ = URMA_MTU_4096; // default mtu and speed active_speed_ = URMA_SP_100G; } @@ -502,14 +491,13 @@ bool UrmaContext::transEidFromString(const std::string& eid_str, 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_; + << device_name_; return ERR_CONTEXT; } Transport::Slice s[nr_poll]; @@ -522,22 +510,21 @@ int UrmaContext::poll(int num_entries, Transport::Slice** slices, } 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; + 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; } @@ -561,10 +548,7 @@ urma_jfce_t* UrmaContext::JFCE() { return jfce_[index]; } -int UrmaContext::jfcCount() { - return jfc_list_.size(); -} - +int UrmaContext::jfcCount() { return jfc_list_.size(); } bool UrmaContext::uninit() { urma_uninit(); @@ -597,19 +581,19 @@ int UrmaEndpoint::construct(GlobalConfig& config) { auto* jfc = context_->jfc(); jfc_outstanding_ = (volatile int*)jfc->jfc_cfg.user_ctx; - max_wr_depth_ = (int)max_wr_depth; // work request + 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) + .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 + .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 = {}; @@ -621,33 +605,31 @@ int UrmaEndpoint::construct(GlobalConfig& config) { 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.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; + 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_jetty = imported_jetty_map_[jetty_list_[i]]; ret = urma_unbind_jetty(jetty_list_[i]); - if (ret) - PLOG(ERROR) << "Failed to unbind jetty"; + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; ret = urma_unimport_jetty(imported_jetty); - if (ret) - PLOG(ERROR) << "Failed to unimport jetty"; + if (ret) PLOG(ERROR) << "Failed to unimport jetty"; if (urma_delete_jetty(jetty_list_[i])) { PLOG(ERROR) << "Failed to destroy jetty"; @@ -657,8 +639,8 @@ int UrmaEndpoint::deconstruct() { bool displayed = false; if (wr_depth_list_[i] != 0) { if (!displayed) { - LOG(WARNING) - << "Outstanding work requests found, CQ will not be generated"; + LOG(WARNING) << "Outstanding work requests found, CQ will not " + "be generated"; displayed = true; } __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); @@ -688,7 +670,6 @@ const std::string UrmaEndpoint::toString() const { return "EndPoint: local " + context_->nicPath() + " (unconnected)"; } - int UrmaEndpoint::setupConnectionsByActive() { RWSpinlock::WriteGuard guard(lock_); if (connected()) { @@ -706,8 +687,8 @@ int UrmaEndpoint::setupConnectionsByActive() { } } } - LOG(ERROR) << "Peer NIC " << context_->deviceName() << - " not found in localhost"; + LOG(ERROR) << "Peer NIC " << context_->deviceName() + << " not found in localhost"; return ERR_DEVICE_NOT_FOUND; } @@ -728,17 +709,17 @@ int UrmaEndpoint::setupConnectionsByActive() { if (rc) return rc; if (!peer_desc.reply_msg.empty()) { LOG(ERROR) << "Reject the handshake request by peer " - << local_desc.peer_nic_path; + << 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; + << ", 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; } @@ -752,7 +733,7 @@ int UrmaEndpoint::setupConnectionsByActive() { } } LOG(ERROR) << "Peer NIC " << peer_nic_name << " not found in " - << peer_server_name; + << peer_server_name; return ERR_DEVICE_NOT_FOUND; } @@ -763,21 +744,18 @@ void UrmaEndpoint::disconnectUnlocked() { 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"; + 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"; + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; ret = urma_unimport_jetty(imported_jetty); - if (ret) - PLOG(ERROR) << "Failed to unimport 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"; + LOG(WARNING) << "Outstanding work requests found, JFC will not " + "be generated"; displayed = true; } __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); @@ -839,9 +817,9 @@ bool UrmaEndpoint::hasOutstandingSlice() const { return false; } -int UrmaEndpoint::submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) { +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()); @@ -871,13 +849,11 @@ int UrmaEndpoint::submitPostSend(std::vector& slice_list, 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.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.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; @@ -910,8 +886,8 @@ int UrmaEndpoint::submitPostSend(std::vector& slice_list, std::vector UrmaEndpoint::JettyNum() const { std::vector ret; - for (int jetty_index = 0; jetty_index < (int)jetty_list_.size(); ++ - jetty_index) + 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; } @@ -928,11 +904,10 @@ int UrmaEndpoint::doSetupConnection(const std::string& peer_eid, 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); + 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; } @@ -959,8 +934,8 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, 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_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"; @@ -970,8 +945,8 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, 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; + LOG(INFO) << "Bind jetty success, local jetty id:" << jetty->jetty_id.id + << ", remote jetty id:" << peer_jetty_num; return 0; } @@ -979,4 +954,4 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, std::shared_ptr UrmaContext::makeEndpoint() { return std::make_shared(this); } -} +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/ub_transport_test.cpp b/mooncake-transfer-engine/tests/ub_transport_test.cpp index a6b6f9b295..29e3474c0b 100644 --- a/mooncake-transfer-engine/tests/ub_transport_test.cpp +++ b/mooncake-transfer-engine/tests/ub_transport_test.cpp @@ -43,7 +43,7 @@ #include -static void checkCudaError(cudaError_t result, const char *message) { +static void checkCudaError(cudaError_t result, const char* message) { if (result != cudaSuccess) { LOG(ERROR) << message << " (Error code: " << result << " - " << cudaGetErrorString(result) << ")" << std::endl; @@ -83,7 +83,7 @@ static void* allocateMemoryPool(size_t size, int socket_id, #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) if (from_vram) { int gpu_id = FLAGS_gpu_id; - void *d_buf; + void* d_buf; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); checkCudaError(cudaMalloc(&d_buf, size), "Failed to allocate device memory"); @@ -184,7 +184,7 @@ int initiatorWorker(TransferEngine* engine, SegmentID segment_id, int thread_id, 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; From 7d8dfb4a705b81388e1460c4085e198d59d24f69 Mon Sep 17 00:00:00 2001 From: zchuango Date: Wed, 25 Mar 2026 17:46:59 +0800 Subject: [PATCH 03/13] change some log for ub transport --- .../transport/ub_transport/ub_transport.cpp | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp index 53128deecb..5feaad8fa9 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -39,7 +39,7 @@ int UbTransport::install(std::string& local_server_name, std::shared_ptr meta, std::shared_ptr topo) { if (topo == nullptr) { - LOG(ERROR) << "UrmaTransport: missing topology"; + LOG(ERROR) << "UbTransport: missing topology"; return ERR_INVALID_ARGUMENT; } metadata_ = meta; @@ -47,11 +47,11 @@ int UbTransport::install(std::string& local_server_name, local_topology_ = topo; auto ret = initializeUbResources(this); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot initialize URMA resources"; + LOG(ERROR) << "UbTransport: cannot initialize Ub resources"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: initialize URMA resources done"; + LOG(INFO) << "UbTransport: initialize Ub resources done"; ret = allocateLocalSegmentID(); if (ret) { @@ -64,19 +64,19 @@ int UbTransport::install(std::string& local_server_name, ret = startHandshakeDaemon(local_server_name); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot start handshake daemon"; + LOG(ERROR) << "UbTransport: cannot start handshake daemon"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: start handshake daemon done"; + LOG(INFO) << "UbTransport: start handshake daemon done"; ret = metadata_->updateLocalSegmentDesc(); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot publish segments"; + LOG(ERROR) << "UbTransport: cannot publish segments"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: publish segments done"; + LOG(INFO) << "UbTransport: publish segments done"; return 0; } @@ -145,7 +145,7 @@ int UbTransport::registerLocalMemoryBatch( for (size_t i = 0; i < buffer_list.size(); ++i) { if (results[i].get()) { - LOG(WARNING) << "UrmaTransport: Failed to register memory: addr " + LOG(WARNING) << "UbTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; } @@ -178,10 +178,10 @@ 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) << "UrmaTransport: Exceed the limitation of current batch's " + LOG(ERROR) << "UbTransport: Exceed the limitation of current batch's " "capacity"; return Status::InvalidArgument( - "UrmaTransport: Exceed the limitation of capacity, batch id: " + + "UbTransport: Exceed the limitation of capacity, batch id: " + std::to_string(batch_id)); } size_t task_id = batch_desc.task_list.size(); @@ -273,10 +273,10 @@ Status UbTransport::submitTransferTask( for (auto& entry : slices_to_post) for (auto s : entry.second) delete s; LOG(ERROR) - << "UrmaTransport: Address not registered by any device(s) " + << "UbTransport: Address not registered by any device(s) " << source_addr; return Status::AddressNotRegistered( - "UrmaTransport: not registered by any device(s), " + "UbTransport: not registered by any device(s), " "address: " + std::to_string(reinterpret_cast(source_addr))); } @@ -316,7 +316,7 @@ Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, const size_t task_count = batch_desc.task_list.size(); if (task_id >= task_count) { return Status::InvalidArgument( - "UrmaTransport::getTransportStatus invalid argument, batch id: " + + "UbTransport::getTransportStatus invalid argument, batch id: " + std::to_string(batch_id)); } auto& task = batch_desc.task_list[task_id]; @@ -428,7 +428,7 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, int UbTransport::initializeUbResources(UbTransport* t) { auto ret = init(t); - if (ret != URMA_SUCCESS && ret != URMA_EEXIST) { + if (ret != 0) { LOG(ERROR) << "Failed to init, ret = " << ret; return -1; } @@ -457,15 +457,15 @@ int UbTransport::initializeUbResources(UbTransport* t) { int UbTransport::init(UbTransport* transport) { if (transport->endpoint_type_ == URMA_ENDPOINT) { if (!UrmaContext::init()) { - LOG(ERROR) << "UrmaEndpoint init failed"; - return ERR_INVALID_ARGUMENT; + LOG(ERROR) << "UrmaContext init failed"; + return -1; } } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; - return ERR_INVALID_ARGUMENT; + LOG(ERROR) << "ObmmContext not support now."; + return -1; } else { LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; - return ERR_INVALID_ARGUMENT; + return -1; } return 0; } @@ -473,10 +473,10 @@ int UbTransport::init(UbTransport* transport) { void UbTransport::uninit(UbTransport* transport) { if (transport->endpoint_type_ == URMA_ENDPOINT) { if (!UrmaContext::uninit()) { - LOG(ERROR) << "UrmaEndpoint uninit failed"; + LOG(ERROR) << "UrmaContext uninit failed"; } } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; + LOG(ERROR) << "ObmmContext not support now."; } else { LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; } @@ -493,7 +493,7 @@ std::shared_ptr UbTransport::buildContext( } return context; } else if (t->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; + LOG(ERROR) << "ObmmContext not support now."; return nullptr; } else { LOG(ERROR) << "invalid endpoint type : " << t->endpoint_type_; From 7650d7d195f3df943e4df5ccf63f984b6de03534 Mon Sep 17 00:00:00 2001 From: zchuango Date: Fri, 27 Mar 2026 19:51:59 +0800 Subject: [PATCH 04/13] change ub_transport to kunpeng_transport --- .../{ub_transport => kunpeng_transport}/ub_context.h | 0 .../{ub_transport => kunpeng_transport}/ub_endpoint.h | 0 .../{ub_transport => kunpeng_transport}/ub_transport.h | 0 .../{ub_transport => kunpeng_transport}/urma_endpoint.h | 4 ++-- mooncake-transfer-engine/src/multi_transport.cpp | 2 +- mooncake-transfer-engine/src/transport/CMakeLists.txt | 2 +- .../{ub_transport => kunpeng_transport}/CMakeLists.txt | 0 .../{ub_transport => kunpeng_transport}/ub_context.cpp | 4 ++-- .../{ub_transport => kunpeng_transport}/ub_transport.cpp | 8 ++++---- .../{ub_transport => kunpeng_transport}/urma_endpoint.cpp | 2 +- 10 files changed, 11 insertions(+), 11 deletions(-) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_context.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_endpoint.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_transport.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/urma_endpoint.h (98%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/CMakeLists.txt (100%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/ub_context.cpp (99%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/ub_transport.cpp (98%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/urma_endpoint.cpp (99%) diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_context.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h similarity index 98% rename from mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h index a815c1d549..e06f3f7927 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h @@ -22,8 +22,8 @@ #include "common.h" #include "config.h" #include "ub/umdk/urma/urma_api.h" -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_endpoint.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" namespace mooncake { struct UrmaJFC { diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index be8aab48ec..8ad5fb8bf8 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -56,7 +56,7 @@ #include "transport/efa_transport/efa_transport.h" #endif #ifdef USE_UB -#include "transport/ub_transport/ub_transport.h" +#include "transport/kunpeng_transport/ub_transport.h" #endif #include diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index abd3b23773..951c7049f5 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -6,7 +6,7 @@ add_library(transport OBJECT ${XPORT_SOURCES} $ $ target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread) if(USE_UB) - add_subdirectory(ub_transport) + add_subdirectory(kunpeng_transport) target_sources(transport PUBLIC $) endif() diff --git a/mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt similarity index 100% rename from mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt rename to mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp similarity index 99% rename from mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index 885914553d..d3eb0cd8c9 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -19,8 +19,8 @@ #include #include #include "config.h" -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_endpoint.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" namespace mooncake { std::shared_ptr UbSIEVEEndpointStore::getEndpoint( diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp similarity index 98% rename from mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 5feaad8fa9..e206c8688f 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -16,10 +16,10 @@ #include "config.h" #include "memory_location.h" #include -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_transport.h" -#include "transport/ub_transport/ub_endpoint.h" -#include "transport/ub_transport/urma_endpoint.h" +#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) diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp similarity index 99% rename from mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index 2a79d13f10..b128587f8b 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -16,7 +16,7 @@ #include #include #include "config.h" -#include "transport/ub_transport/urma_endpoint.h" +#include "transport/kunpeng_transport/urma_endpoint.h" namespace mooncake { static int isNullEid(urma_eid_t* eid) { From 7914cfc07a416bae04b9717a44cd4b874a08c1a7 Mon Sep 17 00:00:00 2001 From: zchuango Date: Mon, 16 Mar 2026 22:02:27 +0800 Subject: [PATCH 05/13] add the support for ub transport --- mooncake-common/common.cmake | 5 + .../example/transfer_engine_bench.cpp | 9 +- .../transfer_engine_bench_with_notify.cpp | 6 + .../transfer_engine_bench_with_retry.cpp | 3 + mooncake-transfer-engine/include/config.h | 7 + .../include/transfer_metadata.h | 6 +- .../include/transport/transport.h | 8 + .../transport/ub_transport/ub_context.h | 362 +++++++ .../transport/ub_transport/ub_endpoint.h | 97 ++ .../transport/ub_transport/ub_transport.h | 127 +++ .../transport/ub_transport/urma_endpoint.h | 204 ++++ mooncake-transfer-engine/src/CMakeLists.txt | 4 + mooncake-transfer-engine/src/config.cpp | 5 +- .../src/multi_transport.cpp | 8 + mooncake-transfer-engine/src/topology.cpp | 114 ++ .../src/transfer_engine_impl.cpp | 9 + .../src/transfer_metadata.cpp | 81 +- .../src/transport/CMakeLists.txt | 5 + .../src/transport/ub_transport/CMakeLists.txt | 18 + .../src/transport/ub_transport/ub_context.cpp | 536 ++++++++++ .../transport/ub_transport/ub_transport.cpp | 531 ++++++++++ .../transport/ub_transport/urma_endpoint.cpp | 982 ++++++++++++++++++ mooncake-transfer-engine/tests/CMakeLists.txt | 6 + .../tests/ub_transport_test.cpp | 331 ++++++ 24 files changed, 3459 insertions(+), 5 deletions(-) create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_context.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h create mode 100644 mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp create mode 100644 mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp create mode 100644 mooncake-transfer-engine/tests/ub_transport_test.cpp 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..820101c23f 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -231,7 +231,11 @@ static void freeMemoryPool(void *addr, size_t size) { #endif } #else - numa_free(addr, size); + if (FLAGS_protocol == "ub") { + munmap(addr, size); // for urma + } else { + numa_free(addr, size); + } #endif } @@ -459,6 +463,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..d7c02d2483 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暂时使用此默认值 + 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..a152583243 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -46,6 +46,7 @@ class TransferMetadata { std::string name; uint16_t lid; std::string gid; + std::string eid; // for ub }; struct BufferDesc { @@ -56,6 +57,8 @@ class TransferMetadata { 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,7 @@ class TransferMetadata { struct HandShakeDesc { std::string local_nic_path; std::string peer_nic_path; + std::vector jetty_num; // for ub/urma #ifdef USE_BAREX uint16_t barex_port; #endif 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/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h new file mode 100644 index 0000000000..b361f6f828 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h @@ -0,0 +1,362 @@ +// 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::string, std::pair, 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) { + } + + ~UbContext() { + } + + 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< + UbSIEVEEndpointStore>(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_; +}; +} + +#endif // UB_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h new file mode 100644 index 0000000000..e96da7a740 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h @@ -0,0 +1,97 @@ +// 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/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h new file mode 100644 index 0000000000..956f8b4e12 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h @@ -0,0 +1,127 @@ +// 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_; +}; +} // mooncake + +#endif // UB_TRANSPORT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h new file mode 100644 index 0000000000..0bc7894df0 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h @@ -0,0 +1,204 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_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_; +}; +} +#endif // URMA_ENDPOINT_H 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..7ec2f7d154 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -25,8 +25,11 @@ 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; + config.num_jfc_per_ctx = val; + config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 + } 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..be8aab48ec 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/ub_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..6c825b1399 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,108 @@ 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"; + 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 : perfer 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 +470,18 @@ void Topology::clear() { int Topology::discover(const std::vector &filter) { matrix_.clear(); +#ifdef USE_UB + auto all_hca = listUBDevices(filter); + // 此处protocol带不进来,使用rdma时重新编译，注释本行，走上面分支 + 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..ac3e08d9d8 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,7 +216,30 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, } segmentJSON["buffers"] = buffersJSON; segmentJSON["priority_matrix"] = desc.topology.toJson(); - } else if (segmentJSON["protocol"] == "tcp") { + } 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) { Json::Value bufferJSON; @@ -381,7 +422,43 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "tcp") { + } 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) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + } + } else if (desc->protocol == "tcp") { for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 1b72a72159..abd3b23773 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(ub_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/ub_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt new file mode 100644 index 0000000000..3514560164 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_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/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp new file mode 100644 index 0000000000..e0a067da45 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp @@ -0,0 +1,536 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_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; + } + // LOG(INFO) << "Endpoint " << peer_nic_path << " not found in + // SIEVEEndpointStore"; + 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); + // remove endpoint but leaving it status unchanged + // in case it is setting up connection or submitting slice + 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); + // LOG(INFO) << victim << " evicted"; + 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();; +} +} \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp new file mode 100644 index 0000000000..b6fb57c3d6 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -0,0 +1,531 @@ +// 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/ub_transport/ub_context.h" +#include "transport/ub_transport/ub_transport.h" +#include "transport/ub_transport/ub_endpoint.h" +#include "transport/ub_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 + for (auto &entry : batch_desc_set_) delete entry.second; + 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) << "UrmaTransport: 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) << "UrmaTransport: cannot initialize URMA resources"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: initialize URMA 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) << "UrmaTransport: cannot start handshake daemon"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: start handshake daemon done"; + + ret = metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(ERROR) << "UrmaTransport: cannot publish segments"; + uninit(this); + return ret; + } + LOG(INFO) << "UrmaTransport: 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) << "UrmaTransport: 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) << "UrmaTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "UrmaTransport: 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) delete s; + LOG(ERROR) + << "UrmaTransport: Address not registered by any device(s) " + << source_addr; + return Status::AddressNotRegistered( + "UrmaTransport: not registered by any device(s), " + "address: " + + std::to_string(reinterpret_cast(source_addr))); + } + // start to submit batch reqeust 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( + "UrmaTransport::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 != URMA_SUCCESS && ret != URMA_EEXIST) { + 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) << "UrmaEndpoint init failed"; + return ERR_INVALID_ARGUMENT; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmEndpoint not support now."; + return ERR_INVALID_ARGUMENT; + } else { + LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; + return ERR_INVALID_ARGUMENT; + } + return 0; +} + +void UbTransport::uninit(UbTransport* transport) { + if (transport->endpoint_type_ == URMA_ENDPOINT) { + if (!UrmaContext::uninit()) { + LOG(ERROR) << "UrmaEndpoint uninit failed"; + } + } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { + LOG(ERROR) << "ObmmEndpoint 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) << "ObmmEndpoint not support now."; + return nullptr; + } else { + LOG(ERROR) << "invalid endpoint type : " << t->endpoint_type_; + return nullptr; + } +} +} + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp new file mode 100644 index 0000000000..67d2da421a --- /dev/null +++ b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp @@ -0,0 +1,982 @@ +// 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/ub_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:此处num_jfc_list和num_jfces使用相同数值，即一个jfc绑定一个jfce，max_jfc_e使用默认值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; + delete 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_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"; + + if (urma_delete_jetty(jetty_list_[i])) { + PLOG(ERROR) << "Failed to destroy jetty"; + return ERR_ENDPOINT; + } + // 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); +} +} 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..a6b6f9b295 --- /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 From 3da4eda51cde8674eaba66157431fdeecc3dacef Mon Sep 17 00:00:00 2001 From: zchuango Date: Mon, 23 Mar 2026 10:34:38 +0800 Subject: [PATCH 06/13] refine the code format --- .../example/transfer_engine_bench.cpp | 2 +- mooncake-transfer-engine/include/config.h | 2 +- .../include/transfer_metadata.h | 16 +- .../transport/ub_transport/ub_context.h | 61 ++--- .../transport/ub_transport/ub_endpoint.h | 19 +- .../transport/ub_transport/ub_transport.h | 32 +-- .../transport/ub_transport/urma_endpoint.h | 58 ++--- mooncake-transfer-engine/src/config.cpp | 5 +- mooncake-transfer-engine/src/topology.cpp | 19 +- .../src/transfer_metadata.cpp | 66 ++--- .../src/transport/ub_transport/ub_context.cpp | 79 +++--- .../transport/ub_transport/ub_transport.cpp | 114 +++----- .../transport/ub_transport/urma_endpoint.cpp | 245 ++++++++---------- .../tests/ub_transport_test.cpp | 6 +- 14 files changed, 325 insertions(+), 399 deletions(-) diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 820101c23f..76c5781aad 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -232,7 +232,7 @@ static void freeMemoryPool(void *addr, size_t size) { } #else if (FLAGS_protocol == "ub") { - munmap(addr, size); // for urma + munmap(addr, size); // for urma } else { numa_free(addr, size); } diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index d7c02d2483..8ceef0f271 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -70,7 +70,7 @@ struct GlobalConfig { size_t num_jfce_per_ctx = 2; int eid_index = 0; uint64_t max_seg_size = 0x10000000000; - size_t max_jfc_e = 4096;// urma暂时使用此默认值 + size_t max_jfc_e = 4096; // urma暂时使用此默认值 size_t num_jetty_per_ep = 1; }; diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index a152583243..c470bf14f7 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -46,19 +46,19 @@ class TransferMetadata { std::string name; uint16_t lid; std::string gid; - std::string eid; // for ub + 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 tseg; // for ub/urma - std::vector l_seg_index; // for ub/urma + 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 { @@ -116,7 +116,7 @@ class TransferMetadata { struct HandShakeDesc { std::string local_nic_path; std::string peer_nic_path; - std::vector jetty_num; // for ub/urma + std::vector jetty_num; // for ub/urma #ifdef USE_BAREX uint16_t barex_port; #endif diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h index b361f6f828..f42d93d640 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h @@ -33,7 +33,7 @@ class UbContext; // define worker pool class class UbWorkerPool { -public: + public: UbWorkerPool(UbContext& context, int numa_socket_id = 0); ~UbWorkerPool(); @@ -41,7 +41,7 @@ class UbWorkerPool { // Add slices to queue, called by Transport int submitPostSend(const std::vector& slice_list); -private: + private: void performPostSend(int thread_id); void performPoll(int thread_id); @@ -54,7 +54,7 @@ class UbWorkerPool { int doProcessContextEvents(); -private: + private: UbContext& context_; const int numa_socket_id_; std::vector worker_thread_; @@ -69,14 +69,14 @@ class UbWorkerPool { std::atomic slice_queue_count_[kShardCount]; TicketLock slice_queue_lock_[kShardCount]; std::vector> - collective_slice_queue_; + 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: + public: virtual ~UbEndpointStore() = default; virtual std::shared_ptr getEndpoint( const std::string& peer_nic_path) = 0; @@ -93,15 +93,14 @@ class UbEndpointStore { // NSDI 24, similar to clock with quick demotion class UbSIEVEEndpointStore : public UbEndpointStore { -public: + public: UbSIEVEEndpointStore(size_t max_size) - : waiting_list_len_(0), max_size_(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; + 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; @@ -110,12 +109,12 @@ class UbSIEVEEndpointStore : public UbEndpointStore { int destroy() override; int disconnect() override; -private: + private: RWSpinlock endpoint_map_lock_; // The bool represents visited - std::unordered_map< - std::string, std::pair, std::atomic_bool>> - endpoint_map_; + std::unordered_map, std::atomic_bool>> + endpoint_map_; std::unordered_map::iterator> fifo_map_; std::list fifo_list_; @@ -129,18 +128,16 @@ class UbSIEVEEndpointStore : public UbEndpointStore { // UbContext class class UbContext { -public: + 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) { - } + show_work_request_flushed_error_(false) {} - ~UbContext() { - } + ~UbContext() {} int doConstruct(GlobalConfig& config) { show_work_request_flushed_error_ = globalConfig().trace; @@ -149,8 +146,8 @@ class UbContext { return 1; } LOG(INFO) << "finish construct context " << toString(); - endpoint_store_ = std::make_shared< - UbSIEVEEndpointStore>(max_endpoints_); + endpoint_store_ = + std::make_shared(max_endpoints_); if (endpoint_store_ == nullptr) { LOG(INFO) << "failed create endpoint store."; return 1; @@ -166,12 +163,12 @@ class UbContext { virtual std::shared_ptr makeEndpoint() = 0; -private: + private: virtual int construct(GlobalConfig& config) = 0; virtual int deconstruct() = 0; -public: + public: virtual int registerMemoryRegion(uint64_t va, size_t length) = 0; virtual int unregisterMemoryRegion(uint64_t va) = 0; @@ -190,8 +187,8 @@ class UbContext { virtual int jfcCount() = 0; - virtual int submitPostSend(const std::vector& slice_list) - = 0; + virtual int submitPostSend( + const std::vector& slice_list) = 0; virtual int getAsyncFd() = 0; @@ -232,9 +229,7 @@ class UbContext { return endpoint_store_->deleteEndpoint(peer_nic_path); } - int disconnectAllEndpoints() { - return endpoint_store_->disconnect(); - } + int disconnectAllEndpoints() { return endpoint_store_->disconnect(); } // Device name, such as `mlx5_3` std::string deviceName() const { return device_name_; } @@ -253,8 +248,8 @@ class UbContext { int eventFd() const { return event_fd_; } int socketId() { - std::string path = "/sys/class/infiniband/" + - device_name_ + "/device/numa_node"; + std::string path = + "/sys/class/infiniband/" + device_name_ + "/device/numa_node"; std::ifstream file(path); if (file.is_open()) { int socket_id; @@ -306,7 +301,7 @@ class UbContext { } } -protected: + protected: static int joinNonblockingPollList(int& event_fd, int data_fd) { event_fd = epoll_create1(0); if (event_fd < 0) { @@ -335,7 +330,7 @@ class UbContext { return 0; } -protected: + protected: const std::string device_name_; UbTransport& engine_; int max_endpoints_; @@ -357,6 +352,6 @@ class UbContext { bool show_work_request_flushed_error_; }; -} +} // namespace mooncake #endif // UB_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h index e96da7a740..e95bd25825 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h @@ -22,7 +22,7 @@ namespace mooncake { // define the UbEndpoint class class UbEndPoint { -public: + public: enum Status { INITIALIZING, UNCONNECTED, @@ -31,10 +31,7 @@ class UbEndPoint { using HandShakeDesc = TransferMetadata::HandShakeDesc; - UbEndPoint() - : status_(INITIALIZING), - active_(true) { - } + UbEndPoint() : status_(INITIALIZING), active_(true) {} virtual int construct(GlobalConfig& config) = 0; @@ -49,9 +46,9 @@ class UbEndPoint { virtual bool hasOutstandingSlice() const = 0; - virtual int submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) = 0; + virtual int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) = 0; virtual const std::string toString() const = 0; @@ -82,16 +79,16 @@ class UbEndPoint { disconnectUnlocked(); } -private: + private: virtual void disconnectUnlocked() = 0; -protected: + protected: volatile uint64_t inactive_time_{}; volatile bool active_{}; std::atomic status_; RWSpinlock lock_; std::string peer_nic_path_; }; -} // namespace mooncake +} // namespace mooncake #endif // UB_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h index 956f8b4e12..223567c439 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h @@ -27,10 +27,7 @@ class UbEndPoint; class TransferMetadata; class UbWorkerpool; -enum UB_ENDPOINT_TYPE { - URMA_ENDPOINT = 0, - OBMM_ENDPOINT = 1 -}; +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. @@ -41,12 +38,12 @@ class UbTransport : public Transport { friend class UbEndPoint; friend class UbWorkerPool; -public: + public: using BufferDesc = TransferMetadata::BufferDesc; using SegmentDesc = TransferMetadata::SegmentDesc; using HandShakeDesc = TransferMetadata::HandShakeDesc; -public: + public: UbTransport(UB_ENDPOINT_TYPE endpoint_type = URMA_ENDPOINT); ~UbTransport(); @@ -82,46 +79,45 @@ class UbTransport : public Transport { SegmentID getSegmentID(const std::string& segment_name); -private: + private: int allocateLocalSegmentID(); -public: + 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); + return metadata_->sendHandshake(peer_server_name, local_desc, + peer_desc); } -private: + 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); + static std::shared_ptr buildContext( + UbTransport* transport, const std::string& device_name, + int max_endpoints); int startHandshakeDaemon(std::string& local_server_name); -public: + 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: + private: std::vector> context_list_; std::shared_ptr local_topology_; UB_ENDPOINT_TYPE endpoint_type_; }; -} // mooncake +} // namespace mooncake #endif // UB_TRANSPORT_H diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h index 0bc7894df0..a815c1d549 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h @@ -27,35 +27,30 @@ namespace mooncake { struct UrmaJFC { - UrmaJFC() : native(nullptr), outstanding(0) { - } + UrmaJFC() : native(nullptr), outstanding(0) {} urma_jfc_t* native; volatile int outstanding; }; struct UrmaJFR { - UrmaJFR() : native(nullptr), outstanding(0) { - } + 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 - } -}; + .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: + public: UrmaContext(UbTransport& engine, std::string device_name, int max_endpoints); ~UrmaContext(); @@ -84,7 +79,7 @@ class UrmaContext : public UbContext { static bool uninit(); static bool init(); -private: + private: int construct(GlobalConfig& config) override; int deconstruct() override; int openDevice(const std::string& device_name, uint8_t port, @@ -102,22 +97,23 @@ class UrmaContext : public UbContext { 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; + 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: + private: std::vector jfc_list_; urma_token_t urma_token = {.token = 0xACFE}; urma_context_t* urma_context_ = nullptr; - //ibv_pd *pd_ = nullptr; + // ibv_pd *pd_ = nullptr; uint64_t max_seg_size{}; urma_mtu active_mtu_; urma_eid_t eid_{}; @@ -152,11 +148,9 @@ class UrmaContext : public UbContext { // define the UrmaEndpoint class class UrmaEndpoint : public UbEndPoint { -public: + public: UrmaEndpoint(UrmaContext* context) - : context_(context), - jfc_outstanding_(nullptr) { - } + : context_(context), jfc_outstanding_(nullptr) {} int construct(GlobalConfig& config) override; @@ -171,16 +165,16 @@ class UrmaEndpoint : public UbEndPoint { bool hasOutstandingSlice() const override; - int submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) override; + int submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) override; const std::string toString() const override; -private: + private: void disconnectUnlocked() override; -private: + private: std::vector JettyNum() const; int doSetupConnection(const std::string& peer_eid, @@ -191,7 +185,7 @@ class UrmaEndpoint : public UbEndPoint { uint32_t peer_jetty_num, std::string* reply_msg = nullptr); -private: + private: UrmaContext* context_; urma_token_t urma_token = {.token = 0xACFE}; std::vector jetty_list_; @@ -200,5 +194,5 @@ class UrmaEndpoint : public UbEndPoint { volatile int* jfc_outstanding_; std::unordered_map imported_jetty_map_; }; -} -#endif // URMA_ENDPOINT_H +} // namespace mooncake +#endif // URMA_ENDPOINT_H diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index 7ec2f7d154..5c046ef823 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -28,9 +28,8 @@ void loadGlobalConfig(GlobalConfig &config) { if (val > 0 && val < 256) { config.num_cq_per_ctx = val; config.num_jfc_per_ctx = val; - config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 - } - else + config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 + } else LOG(WARNING) << "Ignore value from environment variable MC_NUM_CQ_PER_CTX"; } diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index 6c825b1399..f0f1f39dd0 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -235,25 +235,28 @@ static std::vector listUBDevices( 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; + 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"; + << " 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))); + snprintf(path, sizeof(path), "%s/numa", + dirname(dirname(resolved_path))); LOG(INFO) << "listUBDevices: numanodepath " << path; std::ifstream(path) >> numa_node; - LOG(INFO) << "UBDevices : perfer node ----" << device_list[i]->name <<" : " << 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}); + .pci_bus_id = std::move(pci_bus_id), + .numa_node = numa_node}); } urma_free_device_list(device_list); urma_uninit(); @@ -276,7 +279,7 @@ static std::vector discoverCpuTopology( 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; diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index ac3e08d9d8..d241c7ff76 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -70,9 +70,9 @@ struct TransferHandshakeUtil { 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(); + 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; } @@ -94,9 +94,9 @@ struct TransferHandshakeUtil { 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(); + 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; } @@ -216,30 +216,30 @@ 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") { + } 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) { Json::Value bufferJSON; @@ -422,7 +422,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "ub") { + } else if (desc->protocol == "ub") { for (const auto &deviceJSON : segmentJSON["devices"]) { DeviceDesc device; device.name = deviceJSON["name"].asString(); @@ -448,7 +448,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; return nullptr; - } + } desc->buffers.push_back(buffer); } @@ -458,7 +458,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol; } - } else if (desc->protocol == "tcp") { + } else if (desc->protocol == "tcp") { for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp index e0a067da45..885914553d 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp @@ -29,7 +29,7 @@ std::shared_ptr UbSIEVEEndpointStore::getEndpoint( 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 + true, std::memory_order_relaxed); // This is safe within read lock // because of idempotence return iter->second.first; } @@ -43,7 +43,7 @@ std::shared_ptr UbSIEVEEndpointStore::insertEndpoint( 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"; + << " already exists in SIEVEEndpointStore"; return endpoint_map_[peer_nic_path].first; } auto endpoint = context->makeEndpoint(); @@ -75,9 +75,8 @@ int UbSIEVEEndpointStore::deleteEndpoint(const std::string& peer_nic_path) { 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_iter == fifo_list_.begin() ? hand_ = std::nullopt + : hand_ = std::prev(fifo_iter); } fifo_list_.erase(fifo_iter); fifo_map_.erase(peer_nic_path); @@ -147,15 +146,12 @@ UbWorkerPool::UbWorkerPool(UbContext& context, int numa_socket_id) redispatch_counter_(0), submitted_slice_count_(0), processed_slice_count_(0) { - for (auto& i : slice_queue_count_) - i.store(0, std::memory_order_relaxed); + 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, i] { transferWorker(i); }); } - worker_thread_.emplace_back( - [this] { monitorWorker(); }); + worker_thread_.emplace_back([this] { monitorWorker(); }); } UbWorkerPool::~UbWorkerPool() { @@ -180,7 +176,7 @@ int UbWorkerPool::submitPostSend( tl_last_cache_ts = current_ts; } - for (auto &slice : slice_list) { + for (auto& slice : slice_list) { auto target_id = slice->target_id; if (!segment_desc_map.count(target_id)) { segment_desc_map[target_id] = @@ -194,8 +190,8 @@ int UbWorkerPool::submitPostSend( } } #else - std::unordered_map> + std::unordered_map> segment_desc_map; for (auto& slice : slice_list) { auto target_id = slice->target_id; @@ -222,22 +218,22 @@ int UbWorkerPool::submitPostSend( 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)) { + 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->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)) { + 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; @@ -252,8 +248,8 @@ int UbWorkerPool::submitPostSend( slice->markFailed(); continue; } - auto targetSegment = peer_segment_desc->buffers[buffer_id].tseg[ - device_id]; + 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, @@ -324,12 +320,12 @@ void UbWorkerPool::performPostSend(int thread_id) { if (entry.second.empty()) continue; #ifdef USE_FAKE_POST_SEND - for (auto &slice : entry.second) slice->markSuccess(); + 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]; + auto& endpoint = endpoint_map[entry.first]; if (endpoint == nullptr || !endpoint->active()) endpoint = context_.endpoint(entry.first); #else @@ -350,7 +346,7 @@ void UbWorkerPool::performPostSend(int thread_id) { } if (!endpoint->connected() && endpoint->setupConnectionsByActive()) { LOG(ERROR) << "Worker: Cannot make connection for endpoint: " - << entry.first << ", mark it inactive"; + << entry.first << ", mark it inactive"; for (auto& slice : entry.second) failed_slice_list.push_back(slice); endpoint->set_active(false); failed_nr_polls++; @@ -395,10 +391,10 @@ void UbWorkerPool::performPoll(int thread_id) { 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) { + 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_.nicPath() << ", mark it inactive"; context_.set_active(false); } context_.deleteEndpoint(slice->peer_nic_path); @@ -430,8 +426,8 @@ void UbWorkerPool::performPoll(int thread_id) { void UbWorkerPool::redispatch(std::vector& slice_list, int thread_id) { - std::unordered_map> + std::unordered_map> segment_desc_map; for (auto& slice : slice_list) { auto target_id = slice->target_id; @@ -448,16 +444,15 @@ void UbWorkerPool::redispatch(std::vector& slice_list, 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)) { + 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]; + 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, @@ -479,7 +474,7 @@ void UbWorkerPool::transferWorker(int thread_id) { std::array poll_jfc_buckets{}; } stats; bindToSocket(numa_socket_id_); - const static uint64_t kWaitPeriodInNano = 100000000; // 100ms + 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 = @@ -525,12 +520,12 @@ void UbWorkerPool::monitorWorker() { if (!(event.events & EPOLLIN)) continue; - if (event.data.fd == context_.getAsyncFd()) - doProcessContextEvents(); + if (event.data.fd == context_.getAsyncFd()) doProcessContextEvents(); } } int UbWorkerPool::doProcessContextEvents() { - return context_.doProcessContextEvents();; + return context_.doProcessContextEvents(); + ; } -} \ No newline at end of file +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp index b6fb57c3d6..53128deecb 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -22,13 +22,12 @@ #include "transport/ub_transport/urma_endpoint.h" namespace mooncake { -UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type): endpoint_type_( - endpoint_type) { -} +UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) + : endpoint_type_(endpoint_type) {} UbTransport::~UbTransport() { #ifdef CONFIG_USE_BATCH_DESC_SET - for (auto &entry : batch_desc_set_) delete entry.second; + for (auto& entry : batch_desc_set_) delete entry.second; batch_desc_set_.clear(); #endif metadata_->removeSegmentDesc(local_server_name_); @@ -57,7 +56,7 @@ int UbTransport::install(std::string& local_server_name, ret = allocateLocalSegmentID(); if (ret) { LOG(ERROR) << "Transfer engine cannot be initialized: cannot " - "allocate local segment"; + "allocate local segment"; uninit(this); return ret; } @@ -128,8 +127,7 @@ 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); + context->unregisterMemoryRegion((uint64_t)addr); return 0; } @@ -148,8 +146,8 @@ int UbTransport::registerLocalMemoryBatch( for (size_t i = 0; i < buffer_list.size(); ++i) { if (results[i].get()) { LOG(WARNING) << "UrmaTransport: Failed to register memory: addr " - << buffer_list[i].addr << " length " - << buffer_list[i].length; + << buffer_list[i].addr << " length " + << buffer_list[i].length; } } @@ -170,19 +168,18 @@ int UbTransport::unregisterLocalMemoryBatch( 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]; + << addr_list[i]; } return metadata_->updateLocalSegmentDesc(); } -Status UbTransport::submitTransfer(BatchID batch_id, - const std::vector& - entries) { +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) << "UrmaTransport: Exceed the limitation of current batch's " - "capacity"; + "capacity"; return Status::InvalidArgument( "UrmaTransport: Exceed the limitation of capacity, batch id: " + std::to_string(batch_id)); @@ -231,11 +228,11 @@ Status UbTransport::submitTransferTask( 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->length = + merge_final_slice ? request.length - offset : kBlockSize; slice->opcode = request.opcode; - // LOG(INFO) << "target_offset : " << request.target_offset << ", offset : " << offset; + // 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; @@ -259,15 +256,15 @@ Status UbTransport::submitTransferTask( buffer_id, device_id, retry_cnt++)) continue; assert(device_id >= 0 && - static_cast(device_id) < context_list_.size()); + 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()); + static_cast(buffer_id) < + local_segment_desc->buffers.size()); assert(local_segment_desc->buffers[buffer_id].tseg.size() == - context_list_.size()); + context_list_.size()); found_device = true; break; } @@ -283,16 +280,15 @@ Status UbTransport::submitTransferTask( "address: " + std::to_string(reinterpret_cast(source_addr))); } - // start to submit batch reqeust task + // 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"); + 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]; + 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; @@ -339,8 +335,8 @@ Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, return Status::OK(); } -Transport::SegmentID -UbTransport::getSegmentID(const std::string& segment_name) { +Transport::SegmentID UbTransport::getSegmentID( + const std::string& segment_name) { return metadata_->getSegmentID(segment_name); } @@ -387,8 +383,8 @@ int UbTransport::onSetupConnections(const HandShakeDesc& peer_desc, int UbTransport::startHandshakeDaemon(std::string& local_server_name) { return metadata_->startHandshakeDaemon( - std::bind(&UbTransport::onSetupConnections, this, - std::placeholders::_1, std::placeholders::_2), + std::bind(&UbTransport::onSetupConnections, this, std::placeholders::_1, + std::placeholders::_2), metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd); } @@ -407,8 +403,8 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, } const auto& buffers = desc->buffers; - for (buffer_id = 0; buffer_id < static_cast(buffers.size()); ++ - buffer_id) { + 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 || @@ -416,16 +412,15 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, continue; } - device_id = hint.empty() - ? desc->topology.selectDevice(buffer.name, retry_cnt) - : desc->topology.selectDevice( - buffer.name, hint, retry_cnt); + 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); + 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; @@ -487,11 +482,11 @@ void UbTransport::uninit(UbTransport* transport) { } } -std::shared_ptr UbTransport::buildContext(UbTransport* t, - const std::string& device_name, int max_endpoints) { +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); + auto context = + std::make_shared(*t, device_name, max_endpoints); if (!context) { LOG(ERROR) << "UrmaContext build failed"; return nullptr; @@ -505,27 +500,4 @@ std::shared_ptr UbTransport::buildContext(UbTransport* t, return nullptr; } } -} - - - - - - - - - - - - - - - - - - - - - - - +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp index 67d2da421a..2a79d13f10 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp @@ -32,8 +32,7 @@ UrmaContext::UrmaContext(UbTransport& engine, std::string device_name, next_jfce_index_(0), next_jfce_vector_index_(0), next_jfc_list_index_(0), - next_jfr_list_index_(0) { -} + next_jfr_list_index_(0) {} UrmaContext::~UrmaContext() { auto thisString = toString(); @@ -48,15 +47,12 @@ UrmaContext::~UrmaContext() { std::string UrmaContext::toString() { std::ostringstream ss; ss << "UrmaContext:{device_name : " << device_name_ - << " ,max_endpoints : " << max_endpoints_ - << " ,async_fd : " << getAsyncFd() - << " }"; + << " ,max_endpoints : " << max_endpoints_ + << " ,async_fd : " << getAsyncFd() << " }"; return ss.str(); } -int UrmaContext::getAsyncFd() { - return urma_context_->async_fd; -} +int UrmaContext::getAsyncFd() { return urma_context_->async_fd; } int UrmaContext::submitPostSend( const std::vector& slice_list) { @@ -70,8 +66,8 @@ int UrmaContext::construct(GlobalConfig& config) { size_t max_jfc_e = config.max_jfc_e; // urma:此处num_jfc_list和num_jfces使用相同数值，即一个jfc绑定一个jfce，max_jfc_e使用默认值DEFAULT_DEPTH if (openDevice(device_name_, port_, eid_index)) { - LOG(ERROR) << "Failed to open device : " << device_name_ << - " with EID index : " << eid_index; + LOG(ERROR) << "Failed to open device : " << device_name_ + << " with EID index : " << eid_index; return ERR_CONTEXT; } @@ -80,9 +76,9 @@ int UrmaContext::construct(GlobalConfig& config) { 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_; + PLOG(ERROR) << "Failed to create jetty for completion events queue " + "on device" + << device_name_; return ERR_CONTEXT; } } @@ -147,11 +143,10 @@ int UrmaContext::construct(GlobalConfig& config) { } LOG(INFO) << "create jfr done"; LOG(INFO) << "URMA device: " << urma_context_->dev->name - << ", EID: (EID_Index " << eid_index_ << ") " << eid(); + << ", EID: (EID_Index " << eid_index_ << ") " << eid(); - LOG(INFO) << "context_ == NULL ? " << (urma_context_ == nullptr - ? "TRUE" - : "FALSE"); + LOG(INFO) << "context_ == NULL ? " + << (urma_context_ == nullptr ? "TRUE" : "FALSE"); worker_pool_ = std::make_shared(*this, socketId()); LOG(INFO) << "create workerpool done"; return 0; @@ -197,8 +192,7 @@ int UrmaContext::deconstruct() { jfc_r_list_.clear(); if (event_fd_ >= 0) { - if (close(event_fd_)) - LOG(ERROR) << "Failed to close epoll fd"; + if (close(event_fd_)) LOG(ERROR) << "Failed to close epoll fd"; event_fd_ = -1; } @@ -206,8 +200,8 @@ int UrmaContext::deconstruct() { 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"; + LOG(ERROR) + << "Failed to destroy jetty for completion event queue"; delete[] jfce_; jfce_ = nullptr; } @@ -224,8 +218,8 @@ int UrmaContext::deconstruct() { 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) + 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; @@ -250,7 +244,7 @@ void* UrmaContext::localSegWithIndex(unsigned 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; + << "shrink it to " << globalConfig().max_seg_size; length = (size_t)globalConfig().max_seg_size; } LOG(INFO) << "Register memory region " << va << " length " << length; @@ -293,8 +287,8 @@ int UrmaContext::unregisterMemoryRegion(uint64_t addr) { 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; + LOG(ERROR) << "Failed to unregister memory " + << (*iter).first->seg.ubva.va; return ERR_CONTEXT; } seg_region_list_.erase(iter); @@ -306,32 +300,27 @@ int UrmaContext::unregisterMemoryRegion(uint64_t addr) { return 0; } -std::string UrmaContext::getEid() { - return eid(); -} +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_; + 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"; + 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"; + LOG(INFO) << "Worker: Context " << device_name_ << " is now active"; } urma_ack_async_event(&event); return 0; @@ -346,8 +335,8 @@ void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) { 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_); + 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; delete handle; @@ -386,8 +375,8 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, 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); + LOG(INFO) << "device_name : " << device_name + << " EID : " << eid(eid_list[j].eid); } if (eid_cnt > 0) { eid_index = eid_list[0].eid_index; @@ -397,8 +386,8 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, context = urma_create_context(devices[i], eid_index); if (!context) { - LOG(ERROR) << "Urma_create_device context(" << device_name << - ") failed failed"; + LOG(ERROR) << "Urma_create_device context(" << device_name + << ") failed failed"; urma_free_device_list(devices); return ERR_CONTEXT; } @@ -407,19 +396,19 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, if (ret) { PLOG(ERROR) << "Failed to query dev attr( " << device_name << " ) "; if (urma_delete_context(context)) { - PLOG(ERROR) << "urma_delete_context(" << device_name << - ") failed"; + 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 (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"; + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; } urma_free_device_list(devices); return ERR_CONTEXT; @@ -430,10 +419,10 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, #ifndef CONFIG_SKIP_NULL_GID_CHECK if (isNullEid(&eid_)) { LOG(WARNING) << "GID is NULL, please check your EID index by " - "specifying MC_EID_INDEX"; + "specifying MC_EID_INDEX"; if (urma_delete_context(context)) { - PLOG(ERROR) << "urma_delete_context(" << device_name << - ") failed"; + PLOG(ERROR) + << "urma_delete_context(" << device_name << ") failed"; } urma_free_device_list(devices); return ERR_CONTEXT; @@ -447,7 +436,7 @@ int UrmaContext::openDevice(const std::string& device_name, uint8_t port, 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_mtu_ = URMA_MTU_4096; // default mtu and speed active_speed_ = URMA_SP_100G; } @@ -502,14 +491,13 @@ bool UrmaContext::transEidFromString(const std::string& eid_str, 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_; + << device_name_; return ERR_CONTEXT; } Transport::Slice s[nr_poll]; @@ -522,22 +510,21 @@ int UrmaContext::poll(int num_entries, Transport::Slice** slices, } 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; + 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; } @@ -561,10 +548,7 @@ urma_jfce_t* UrmaContext::JFCE() { return jfce_[index]; } -int UrmaContext::jfcCount() { - return jfc_list_.size(); -} - +int UrmaContext::jfcCount() { return jfc_list_.size(); } bool UrmaContext::uninit() { urma_uninit(); @@ -597,19 +581,19 @@ int UrmaEndpoint::construct(GlobalConfig& config) { auto* jfc = context_->jfc(); jfc_outstanding_ = (volatile int*)jfc->jfc_cfg.user_ctx; - max_wr_depth_ = (int)max_wr_depth; // work request + 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) + .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 + .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 = {}; @@ -621,33 +605,31 @@ int UrmaEndpoint::construct(GlobalConfig& config) { 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.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; + 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_jetty = imported_jetty_map_[jetty_list_[i]]; ret = urma_unbind_jetty(jetty_list_[i]); - if (ret) - PLOG(ERROR) << "Failed to unbind jetty"; + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; ret = urma_unimport_jetty(imported_jetty); - if (ret) - PLOG(ERROR) << "Failed to unimport jetty"; + if (ret) PLOG(ERROR) << "Failed to unimport jetty"; if (urma_delete_jetty(jetty_list_[i])) { PLOG(ERROR) << "Failed to destroy jetty"; @@ -657,8 +639,8 @@ int UrmaEndpoint::deconstruct() { bool displayed = false; if (wr_depth_list_[i] != 0) { if (!displayed) { - LOG(WARNING) - << "Outstanding work requests found, CQ will not be generated"; + LOG(WARNING) << "Outstanding work requests found, CQ will not " + "be generated"; displayed = true; } __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); @@ -688,7 +670,6 @@ const std::string UrmaEndpoint::toString() const { return "EndPoint: local " + context_->nicPath() + " (unconnected)"; } - int UrmaEndpoint::setupConnectionsByActive() { RWSpinlock::WriteGuard guard(lock_); if (connected()) { @@ -706,8 +687,8 @@ int UrmaEndpoint::setupConnectionsByActive() { } } } - LOG(ERROR) << "Peer NIC " << context_->deviceName() << - " not found in localhost"; + LOG(ERROR) << "Peer NIC " << context_->deviceName() + << " not found in localhost"; return ERR_DEVICE_NOT_FOUND; } @@ -728,17 +709,17 @@ int UrmaEndpoint::setupConnectionsByActive() { if (rc) return rc; if (!peer_desc.reply_msg.empty()) { LOG(ERROR) << "Reject the handshake request by peer " - << local_desc.peer_nic_path; + << 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; + << ", 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; } @@ -752,7 +733,7 @@ int UrmaEndpoint::setupConnectionsByActive() { } } LOG(ERROR) << "Peer NIC " << peer_nic_name << " not found in " - << peer_server_name; + << peer_server_name; return ERR_DEVICE_NOT_FOUND; } @@ -763,21 +744,18 @@ void UrmaEndpoint::disconnectUnlocked() { 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"; + 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"; + if (ret) PLOG(ERROR) << "Failed to unbind jetty"; ret = urma_unimport_jetty(imported_jetty); - if (ret) - PLOG(ERROR) << "Failed to unimport 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"; + LOG(WARNING) << "Outstanding work requests found, JFC will not " + "be generated"; displayed = true; } __sync_fetch_and_sub(jfc_outstanding_, wr_depth_list_[i]); @@ -839,9 +817,9 @@ bool UrmaEndpoint::hasOutstandingSlice() const { return false; } -int UrmaEndpoint::submitPostSend(std::vector& slice_list, - std::vector& - failed_slice_list) { +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()); @@ -871,13 +849,11 @@ int UrmaEndpoint::submitPostSend(std::vector& slice_list, 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.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.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; @@ -910,8 +886,8 @@ int UrmaEndpoint::submitPostSend(std::vector& slice_list, std::vector UrmaEndpoint::JettyNum() const { std::vector ret; - for (int jetty_index = 0; jetty_index < (int)jetty_list_.size(); ++ - jetty_index) + 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; } @@ -928,11 +904,10 @@ int UrmaEndpoint::doSetupConnection(const std::string& peer_eid, 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); + 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; } @@ -959,8 +934,8 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, 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_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"; @@ -970,8 +945,8 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, 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; + LOG(INFO) << "Bind jetty success, local jetty id:" << jetty->jetty_id.id + << ", remote jetty id:" << peer_jetty_num; return 0; } @@ -979,4 +954,4 @@ int UrmaEndpoint::doSetupConnection(int jetty_index, std::shared_ptr UrmaContext::makeEndpoint() { return std::make_shared(this); } -} +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/ub_transport_test.cpp b/mooncake-transfer-engine/tests/ub_transport_test.cpp index a6b6f9b295..29e3474c0b 100644 --- a/mooncake-transfer-engine/tests/ub_transport_test.cpp +++ b/mooncake-transfer-engine/tests/ub_transport_test.cpp @@ -43,7 +43,7 @@ #include -static void checkCudaError(cudaError_t result, const char *message) { +static void checkCudaError(cudaError_t result, const char* message) { if (result != cudaSuccess) { LOG(ERROR) << message << " (Error code: " << result << " - " << cudaGetErrorString(result) << ")" << std::endl; @@ -83,7 +83,7 @@ static void* allocateMemoryPool(size_t size, int socket_id, #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) if (from_vram) { int gpu_id = FLAGS_gpu_id; - void *d_buf; + void* d_buf; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); checkCudaError(cudaMalloc(&d_buf, size), "Failed to allocate device memory"); @@ -184,7 +184,7 @@ int initiatorWorker(TransferEngine* engine, SegmentID segment_id, int thread_id, 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; From 147bc775f60dbdd873e6d66c6bddbe3c1bbded8c Mon Sep 17 00:00:00 2001 From: zchuango Date: Wed, 25 Mar 2026 17:46:59 +0800 Subject: [PATCH 07/13] change some log for ub transport --- .../transport/ub_transport/ub_transport.cpp | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp index 53128deecb..5feaad8fa9 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp @@ -39,7 +39,7 @@ int UbTransport::install(std::string& local_server_name, std::shared_ptr meta, std::shared_ptr topo) { if (topo == nullptr) { - LOG(ERROR) << "UrmaTransport: missing topology"; + LOG(ERROR) << "UbTransport: missing topology"; return ERR_INVALID_ARGUMENT; } metadata_ = meta; @@ -47,11 +47,11 @@ int UbTransport::install(std::string& local_server_name, local_topology_ = topo; auto ret = initializeUbResources(this); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot initialize URMA resources"; + LOG(ERROR) << "UbTransport: cannot initialize Ub resources"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: initialize URMA resources done"; + LOG(INFO) << "UbTransport: initialize Ub resources done"; ret = allocateLocalSegmentID(); if (ret) { @@ -64,19 +64,19 @@ int UbTransport::install(std::string& local_server_name, ret = startHandshakeDaemon(local_server_name); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot start handshake daemon"; + LOG(ERROR) << "UbTransport: cannot start handshake daemon"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: start handshake daemon done"; + LOG(INFO) << "UbTransport: start handshake daemon done"; ret = metadata_->updateLocalSegmentDesc(); if (ret) { - LOG(ERROR) << "UrmaTransport: cannot publish segments"; + LOG(ERROR) << "UbTransport: cannot publish segments"; uninit(this); return ret; } - LOG(INFO) << "UrmaTransport: publish segments done"; + LOG(INFO) << "UbTransport: publish segments done"; return 0; } @@ -145,7 +145,7 @@ int UbTransport::registerLocalMemoryBatch( for (size_t i = 0; i < buffer_list.size(); ++i) { if (results[i].get()) { - LOG(WARNING) << "UrmaTransport: Failed to register memory: addr " + LOG(WARNING) << "UbTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; } @@ -178,10 +178,10 @@ 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) << "UrmaTransport: Exceed the limitation of current batch's " + LOG(ERROR) << "UbTransport: Exceed the limitation of current batch's " "capacity"; return Status::InvalidArgument( - "UrmaTransport: Exceed the limitation of capacity, batch id: " + + "UbTransport: Exceed the limitation of capacity, batch id: " + std::to_string(batch_id)); } size_t task_id = batch_desc.task_list.size(); @@ -273,10 +273,10 @@ Status UbTransport::submitTransferTask( for (auto& entry : slices_to_post) for (auto s : entry.second) delete s; LOG(ERROR) - << "UrmaTransport: Address not registered by any device(s) " + << "UbTransport: Address not registered by any device(s) " << source_addr; return Status::AddressNotRegistered( - "UrmaTransport: not registered by any device(s), " + "UbTransport: not registered by any device(s), " "address: " + std::to_string(reinterpret_cast(source_addr))); } @@ -316,7 +316,7 @@ Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, const size_t task_count = batch_desc.task_list.size(); if (task_id >= task_count) { return Status::InvalidArgument( - "UrmaTransport::getTransportStatus invalid argument, batch id: " + + "UbTransport::getTransportStatus invalid argument, batch id: " + std::to_string(batch_id)); } auto& task = batch_desc.task_list[task_id]; @@ -428,7 +428,7 @@ int UbTransport::selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, int UbTransport::initializeUbResources(UbTransport* t) { auto ret = init(t); - if (ret != URMA_SUCCESS && ret != URMA_EEXIST) { + if (ret != 0) { LOG(ERROR) << "Failed to init, ret = " << ret; return -1; } @@ -457,15 +457,15 @@ int UbTransport::initializeUbResources(UbTransport* t) { int UbTransport::init(UbTransport* transport) { if (transport->endpoint_type_ == URMA_ENDPOINT) { if (!UrmaContext::init()) { - LOG(ERROR) << "UrmaEndpoint init failed"; - return ERR_INVALID_ARGUMENT; + LOG(ERROR) << "UrmaContext init failed"; + return -1; } } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; - return ERR_INVALID_ARGUMENT; + LOG(ERROR) << "ObmmContext not support now."; + return -1; } else { LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; - return ERR_INVALID_ARGUMENT; + return -1; } return 0; } @@ -473,10 +473,10 @@ int UbTransport::init(UbTransport* transport) { void UbTransport::uninit(UbTransport* transport) { if (transport->endpoint_type_ == URMA_ENDPOINT) { if (!UrmaContext::uninit()) { - LOG(ERROR) << "UrmaEndpoint uninit failed"; + LOG(ERROR) << "UrmaContext uninit failed"; } } else if (transport->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; + LOG(ERROR) << "ObmmContext not support now."; } else { LOG(ERROR) << "invalid endpoint type : " << transport->endpoint_type_; } @@ -493,7 +493,7 @@ std::shared_ptr UbTransport::buildContext( } return context; } else if (t->endpoint_type_ == OBMM_ENDPOINT) { - LOG(ERROR) << "ObmmEndpoint not support now."; + LOG(ERROR) << "ObmmContext not support now."; return nullptr; } else { LOG(ERROR) << "invalid endpoint type : " << t->endpoint_type_; From a1fea3c6ffca7849c8016c4ac045fe3f88af7b5e Mon Sep 17 00:00:00 2001 From: zchuango Date: Fri, 27 Mar 2026 19:51:59 +0800 Subject: [PATCH 08/13] change ub_transport to kunpeng_transport --- .../{ub_transport => kunpeng_transport}/ub_context.h | 0 .../{ub_transport => kunpeng_transport}/ub_endpoint.h | 0 .../{ub_transport => kunpeng_transport}/ub_transport.h | 0 .../{ub_transport => kunpeng_transport}/urma_endpoint.h | 4 ++-- mooncake-transfer-engine/src/multi_transport.cpp | 2 +- mooncake-transfer-engine/src/transport/CMakeLists.txt | 2 +- .../{ub_transport => kunpeng_transport}/CMakeLists.txt | 0 .../{ub_transport => kunpeng_transport}/ub_context.cpp | 4 ++-- .../{ub_transport => kunpeng_transport}/ub_transport.cpp | 8 ++++---- .../{ub_transport => kunpeng_transport}/urma_endpoint.cpp | 2 +- 10 files changed, 11 insertions(+), 11 deletions(-) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_context.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_endpoint.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/ub_transport.h (100%) rename mooncake-transfer-engine/include/transport/{ub_transport => kunpeng_transport}/urma_endpoint.h (98%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/CMakeLists.txt (100%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/ub_context.cpp (99%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/ub_transport.cpp (98%) rename mooncake-transfer-engine/src/transport/{ub_transport => kunpeng_transport}/urma_endpoint.cpp (99%) diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_context.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_endpoint.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h similarity index 100% rename from mooncake-transfer-engine/include/transport/ub_transport/ub_transport.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h diff --git a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h similarity index 98% rename from mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h rename to mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h index a815c1d549..e06f3f7927 100644 --- a/mooncake-transfer-engine/include/transport/ub_transport/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h @@ -22,8 +22,8 @@ #include "common.h" #include "config.h" #include "ub/umdk/urma/urma_api.h" -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_endpoint.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" namespace mooncake { struct UrmaJFC { diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index be8aab48ec..8ad5fb8bf8 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -56,7 +56,7 @@ #include "transport/efa_transport/efa_transport.h" #endif #ifdef USE_UB -#include "transport/ub_transport/ub_transport.h" +#include "transport/kunpeng_transport/ub_transport.h" #endif #include diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index abd3b23773..951c7049f5 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -6,7 +6,7 @@ add_library(transport OBJECT ${XPORT_SOURCES} $ $ target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread) if(USE_UB) - add_subdirectory(ub_transport) + add_subdirectory(kunpeng_transport) target_sources(transport PUBLIC $) endif() diff --git a/mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt similarity index 100% rename from mooncake-transfer-engine/src/transport/ub_transport/CMakeLists.txt rename to mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp similarity index 99% rename from mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index 885914553d..d3eb0cd8c9 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -19,8 +19,8 @@ #include #include #include "config.h" -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_endpoint.h" +#include "transport/kunpeng_transport/ub_context.h" +#include "transport/kunpeng_transport/ub_endpoint.h" namespace mooncake { std::shared_ptr UbSIEVEEndpointStore::getEndpoint( diff --git a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp similarity index 98% rename from mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 5feaad8fa9..e206c8688f 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -16,10 +16,10 @@ #include "config.h" #include "memory_location.h" #include -#include "transport/ub_transport/ub_context.h" -#include "transport/ub_transport/ub_transport.h" -#include "transport/ub_transport/ub_endpoint.h" -#include "transport/ub_transport/urma_endpoint.h" +#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) diff --git a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp similarity index 99% rename from mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp rename to mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index 2a79d13f10..b128587f8b 100644 --- a/mooncake-transfer-engine/src/transport/ub_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -16,7 +16,7 @@ #include #include #include "config.h" -#include "transport/ub_transport/urma_endpoint.h" +#include "transport/kunpeng_transport/urma_endpoint.h" namespace mooncake { static int isNullEid(urma_eid_t* eid) { From d4b2235f487861a93622fc7752d874eb898d6899 Mon Sep 17 00:00:00 2001 From: zchuango Date: Wed, 1 Apr 2026 16:42:19 +0800 Subject: [PATCH 09/13] clean some comments and unused code --- .../src/transport/kunpeng_transport/ub_context.cpp | 5 ----- .../src/transport/kunpeng_transport/urma_endpoint.cpp | 2 +- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index d3eb0cd8c9..37536c4e6e 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -33,8 +33,6 @@ std::shared_ptr UbSIEVEEndpointStore::getEndpoint( // because of idempotence return iter->second.first; } - // LOG(INFO) << "Endpoint " << peer_nic_path << " not found in - // SIEVEEndpointStore"; return nullptr; } @@ -67,8 +65,6 @@ std::shared_ptr UbSIEVEEndpointStore::insertEndpoint( int UbSIEVEEndpointStore::deleteEndpoint(const std::string& peer_nic_path) { RWSpinlock::WriteGuard guard(endpoint_map_lock_); auto iter = endpoint_map_.find(peer_nic_path); - // remove endpoint but leaving it status unchanged - // in case it is setting up connection or submitting slice if (iter != endpoint_map_.end()) { waiting_list_len_++; waiting_list_.insert(iter->second.first); @@ -103,7 +99,6 @@ void UbSIEVEEndpointStore::evictEndpoint() { hand_ = (o == fifo_list_.begin() ? --fifo_list_.end() : std::prev(o)); fifo_list_.erase(o); fifo_map_.erase(victim); - // LOG(INFO) << victim << " evicted"; auto victim_instance = endpoint_map_[victim].first; victim_instance->set_active(false); waiting_list_len_++; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index b128587f8b..088fb7f7c9 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -118,7 +118,7 @@ int UrmaContext::construct(GlobalConfig& config) { return ERR_CONTEXT; } jfc_r_list_[i] = jfc; - LOG(INFO) << "create jfc(send) done, jfc id : " << jfc->jfc_id.id; + LOG(INFO) << "create jfc(recv) done, jfc id : " << jfc->jfc_id.id; } jfr_list_.resize(num_jfc_list); urma_jfr_cfg_t jfr_cfg[num_jfc_list] = {}; From 48f9099adf93580074fe1a7a6d69b355bab1129c Mon Sep 17 00:00:00 2001 From: zchuango Date: Thu, 2 Apr 2026 15:20:30 +0800 Subject: [PATCH 10/13] fix the config.h merge error --- mooncake-transfer-engine/include/config.h | 7 ------- 1 file changed, 7 deletions(-) diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index f79823b22c..8ceef0f271 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -64,13 +64,6 @@ struct GlobalConfig { // ib_pci_relaxed_ordering_mode: 0: off, 1: on if supported, 2: auto int ib_pci_relaxed_ordering_mode = 0; bool ascend_use_fabric_mem = 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暂时使用此默认值 - size_t num_jetty_per_ep = 1; bool ascend_agent_mode = false; // ub config parameters size_t num_jfc_per_ctx = 2; From 6549c1192a9c7745c5ff5823cd9d293cb4e7fb36 Mon Sep 17 00:00:00 2001 From: zchuango Date: Thu, 2 Apr 2026 18:23:08 +0800 Subject: [PATCH 11/13] fix the code with review suggestion --- .../include/transport/kunpeng_transport/ub_context.h | 2 +- .../include/transport/kunpeng_transport/urma_endpoint.h | 6 +++--- mooncake-transfer-engine/src/topology.cpp | 1 - .../src/transport/kunpeng_transport/ub_transport.cpp | 3 +-- .../src/transport/kunpeng_transport/urma_endpoint.cpp | 8 ++++---- 5 files changed, 9 insertions(+), 11 deletions(-) diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h index f42d93d640..7240ff7bbd 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -137,7 +137,7 @@ class UbContext { active_(true), show_work_request_flushed_error_(false) {} - ~UbContext() {} + virtual ~UbContext() = default; int doConstruct(GlobalConfig& config) { show_work_request_flushed_error_ = globalConfig().trace; diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h index e06f3f7927..5c7f9b9006 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h @@ -87,13 +87,13 @@ class UrmaContext : public UbContext { urma_target_seg_t* seg(uint64_t addr); - std::vector remote_seg_list() { return remote_seg_list_; } + std::vector& remote_seg_list() { return remote_seg_list_; } - std::vector imported_seg_list() { + std::vector& imported_seg_list() { return imported_seg_list_; } - std::vector local_tseg_list() { + std::vector& local_tseg_list() { return local_tseg_list_; } diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index f0f1f39dd0..e8e23f4c37 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -475,7 +475,6 @@ int Topology::discover(const std::vector &filter) { matrix_.clear(); #ifdef USE_UB auto all_hca = listUBDevices(filter); - // 此处protocol带不进来,使用rdma时重新编译，注释本行，走上面分支 for (auto &ent : discoverCpuTopology(all_hca)) { matrix_[ent.name] = ent; } diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index e206c8688f..d3e9e8b9d9 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -27,7 +27,6 @@ UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) UbTransport::~UbTransport() { #ifdef CONFIG_USE_BATCH_DESC_SET - for (auto& entry : batch_desc_set_) delete entry.second; batch_desc_set_.clear(); #endif metadata_->removeSegmentDesc(local_server_name_); @@ -271,7 +270,7 @@ Status UbTransport::submitTransferTask( if (device_id < 0) { auto source_addr = slice->source_addr; for (auto& entry : slices_to_post) - for (auto s : entry.second) delete s; + for (auto s : entry.second) getSliceCache().deallocate(s); LOG(ERROR) << "UbTransport: Address not registered by any device(s) " << source_addr; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index b128587f8b..0a6c36b572 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -334,15 +334,15 @@ void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) { 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); + 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; - delete handle; + free(handle); return nullptr; } - imported_seg_list().push_back(import_tseg); + imported_seg_list_.push_back(import_tseg); import_tseg_map[remoteSegmentStr] = import_tseg; return import_tseg; } @@ -919,7 +919,7 @@ 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()) + if (jetty_index < 0 || jetty_index >= (int)jetty_list_.size()) return ERR_INVALID_ARGUMENT; auto& jetty = jetty_list_[jetty_index]; urma_eid_t eid; From 76c8e465cc1083ae10e0d7dbdb6da5b912978a69 Mon Sep 17 00:00:00 2001 From: zchuango Date: Fri, 3 Apr 2026 12:08:42 +0000 Subject: [PATCH 12/13] add the usage docs and refine some code for code review --- .../source/zh_archive/kunpeng_ub_transport.md | 90 +++++++++++++++++++ .../example/transfer_engine_bench.cpp | 3 +- mooncake-transfer-engine/src/topology.cpp | 1 + .../kunpeng_transport/ub_context.cpp | 1 - .../kunpeng_transport/urma_endpoint.cpp | 14 +-- mooncake-transfer-engine/tests/CMakeLists.txt | 2 +- 6 files changed, 100 insertions(+), 11 deletions(-) create mode 100644 docs/source/zh_archive/kunpeng_ub_transport.md 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-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 76c5781aad..ddcbb726d1 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -233,9 +233,8 @@ static void freeMemoryPool(void *addr, size_t size) { #else if (FLAGS_protocol == "ub") { munmap(addr, size); // for urma - } else { - numa_free(addr, size); } + numa_free(addr, size); #endif } diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index e8e23f4c37..abe026ed03 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -219,6 +219,7 @@ static std::vector listUBDevices( 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) { diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index 37536c4e6e..4814a183f8 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -521,6 +521,5 @@ void UbWorkerPool::monitorWorker() { int UbWorkerPool::doProcessContextEvents() { return context_.doProcessContextEvents(); - ; } } // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index 0a6c36b572..e0f5a3555b 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -625,15 +625,15 @@ int UrmaEndpoint::construct(GlobalConfig& config) { int UrmaEndpoint::deconstruct() { int ret = 0; for (size_t i = 0; i < jetty_list_.size(); ++i) { - auto imported_jetty = imported_jetty_map_[jetty_list_[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"; - ret = urma_unimport_jetty(imported_jetty); - if (ret) PLOG(ERROR) << "Failed to unimport jetty"; - - if (urma_delete_jetty(jetty_list_[i])) { - PLOG(ERROR) << "Failed to destroy jetty"; - return ERR_ENDPOINT; + 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; diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 47ecbafec9..7628618b38 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -68,7 +68,7 @@ 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) + # add_test(NAME ub_transport_test COMMAND ub_transport_test) endif() add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp) From 02ef725347d2a1f066c4bb2cf2599990f5ecaac3 Mon Sep 17 00:00:00 2001 From: zchuango Date: Sun, 5 Apr 2026 08:07:23 +0000 Subject: [PATCH 13/13] fix some comments with chinese and modify the CMakeLists.txt file --- mooncake-transfer-engine/include/config.h | 2 +- mooncake-transfer-engine/include/transfer_metadata.h | 2 ++ mooncake-transfer-engine/src/config.cpp | 3 ++- .../src/transport/kunpeng_transport/urma_endpoint.cpp | 4 +++- mooncake-transfer-engine/tests/CMakeLists.txt | 2 +- 5 files changed, 9 insertions(+), 4 deletions(-) diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index 8ceef0f271..df8812706a 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -70,7 +70,7 @@ struct GlobalConfig { size_t num_jfce_per_ctx = 2; int eid_index = 0; uint64_t max_seg_size = 0x10000000000; - size_t max_jfc_e = 4096; // urma暂时使用此默认值 + size_t max_jfc_e = 4096; // urma is temporarily using this default value. size_t num_jetty_per_ep = 1; }; diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index c470bf14f7..2cf0d6bf70 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -116,7 +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/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index 5c046ef823..5734ad2dcb 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -27,8 +27,9 @@ void loadGlobalConfig(GlobalConfig &config) { int val = atoi(num_cq_per_ctx_env); if (val > 0 && val < 256) { config.num_cq_per_ctx = val; + // In URMA, JFC and JFCE are bound one-to-one. config.num_jfc_per_ctx = val; - config.num_jfce_per_ctx = val; // urma中jfc和jfce是一对一绑定的 + 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/transport/kunpeng_transport/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp index e0f5a3555b..1d4af10763 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp @@ -64,7 +64,9 @@ int UrmaContext::construct(GlobalConfig& config) { 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:此处num_jfc_list和num_jfces使用相同数值，即一个jfc绑定一个jfce，max_jfc_e使用默认值DEFAULT_DEPTH + // 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; diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 7628618b38..47ecbafec9 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -68,7 +68,7 @@ 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) + add_test(NAME ub_transport_test COMMAND ub_transport_test) endif() add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp)