From e0ddd91629c67cd7e37708eba9bbc474c2029cc9 Mon Sep 17 00:00:00 2001 From: EkiRui <374389385@qq.com> Date: Tue, 30 Dec 2025 16:12:40 +0800 Subject: [PATCH] [Store] add HugePage support (#1300) * [Store] add HugePage support Co-authored-by: Teng Ma Signed-off-by: Xingrui Yi * doc: add huge page env introduction Signed-off-by: Xingrui Yi --------- Signed-off-by: Xingrui Yi Co-authored-by: Teng Ma --- docs/source/design/mooncake-store.md | 2 + docs/source/zh_archive/mooncake-store.md | 2 + mooncake-store/include/client_buffer.hpp | 7 +- mooncake-store/include/dummy_client.h | 3 + mooncake-store/include/real_client.h | 12 +++ mooncake-store/include/utils.h | 108 ++++++++++++++++++++--- mooncake-store/src/client_buffer.cpp | 21 +++-- mooncake-store/src/dummy_client.cpp | 36 ++++++-- mooncake-store/src/real_client.cpp | 29 ++++-- mooncake-store/src/utils.cpp | 37 ++++++++ 10 files changed, 223 insertions(+), 34 deletions(-) diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md index 08677b60..58720f93 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -637,6 +637,8 @@ The HTTP metadata server can be configured using the following parameters: - MC_STORE_MEMCPY: Enables or disables local memcpy optimization, set to 1/true to enable, 0/false to disable. - MC_STORE_CLIENT_METRIC: Enables client metric reporting, enabled by default; set to 0/false to disable. - MC_STORE_CLIENT_METRIC_INTERVAL: Reporting interval in seconds, default 0 (collects but does not report). +- MC_STORE_USE_HUGEPAGE: Enables huge page support, disabled by default. +- MC_STORE_HUGEPAGE_SIZE: Specifies the page size of the huge page to use, default 2M. #### Usage Example To start the master service with the HTTP metadata server enabled: ```bash diff --git a/docs/source/zh_archive/mooncake-store.md b/docs/source/zh_archive/mooncake-store.md index 78cb296f..babf8cca 100644 --- a/docs/source/zh_archive/mooncake-store.md +++ b/docs/source/zh_archive/mooncake-store.md @@ -647,6 +647,8 @@ HTTP 元数据服务器可通过以下参数进行配置： - **MC_STORE_MEMCPY**: 控制是否启用本地 memcpy 优化, 1/true 启用, 0/false 禁用 - **MC_STORE_CLIENT_METRIC**: 启用客户端指标上报, 默认启用；设为 0/false 禁用 - **MC_STORE_CLIENT_METRIC_INTERVAL**: 指标上报间隔(秒), 默认 0(仅收集不上报) +- **MC_STORE_USE_HUGEPAGE**: 启用 hugepage 优化, 默认禁用, 设置为 1/true 启用 +- **MC_STORE_HUGEPAGE_SIZE**: hugepage 页大小, 默认 2M #### 使用示例 diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.hpp index 86d56814..beb78eff 100644 --- a/mooncake-store/include/client_buffer.hpp +++ b/mooncake-store/include/client_buffer.hpp @@ -28,7 +28,8 @@ class ClientBufferAllocator public: // Create for heap-allocated memory static std::shared_ptr create( - size_t size, const std::string& protocol = ""); + size_t size, const std::string& protocol = "", + bool use_hugepage = false); // Create for shared memory static std::shared_ptr create( @@ -52,7 +53,8 @@ class ClientBufferAllocator private: // Private constructors for different memory types - ClientBufferAllocator(size_t size, const std::string& protocol); + ClientBufferAllocator(size_t size, const std::string& protocol, + bool use_hugepage); ClientBufferAllocator(void* addr, size_t size, const std::string& protocol); std::shared_ptr allocator_; @@ -61,6 +63,7 @@ class ClientBufferAllocator void* buffer_; size_t buffer_size_; bool is_external_memory_ = false; + bool use_hugepage_ = false; }; /** diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index fde0089e..b2895bd3 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -34,6 +34,8 @@ class ShmHelper { return shms_; } + bool is_hugepage() const { return use_hugepage_; } + ShmHelper(const ShmHelper &) = delete; ShmHelper &operator=(const ShmHelper &) = delete; @@ -43,6 +45,7 @@ class ShmHelper { std::vector> shms_; static std::mutex shm_mutex_; + bool use_hugepage_ = false; }; class DummyClient : public PyClient { diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 02585d83..12c76415 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -428,6 +428,15 @@ class RealClient : public PyClient { } }; + struct HugepageSegmentDeleter { + size_t size = 0; + void operator()(void *ptr) const { + if (ptr && size > 0) { + free_buffer_mmap_memory(ptr, size); + } + } + }; + struct AscendSegmentDeleter { void operator()(void *ptr) { if (ptr) { @@ -436,12 +445,15 @@ class RealClient : public PyClient { } }; + std::vector> + hugepage_segment_ptrs_; std::vector> segment_ptrs_; std::vector> ascend_segment_ptrs_; std::string protocol; std::string device_name; std::string local_hostname; + bool use_hugepage_ = false; struct MappedShm { std::string shm_name; diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 6b8064cf..794ed632 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -2,10 +2,13 @@ #include #include +#include +#include #include #include #include +#include "rpc_types.h" #include "types.h" namespace mooncake { @@ -119,17 +122,13 @@ std::string expected_to_str(const tl::expected& expected) { return oss.str(); } -/* - @brief Allocates memory for the `BufferAllocator` class. - @param total_size The total size of the memory to allocate. - @return A pointer to the allocated memory. -*/ -void* allocate_buffer_allocator_memory( - size_t total_size, const std::string& protocol = "", - size_t alignment = facebook::cachelib::Slab::kSize); - -void free_memory(const std::string& protocol, void* ptr); +// String utility functions +/** + * @brief Convert a byte size to a human-readable string + * @param bytes Number of bytes + * @return std::string Human-readable string representation of size + */ [[nodiscard]] inline std::string byte_size_to_string(uint64_t bytes) { const double KB = 1024.0; const double MB = KB * 1024.0; @@ -230,8 +229,6 @@ void free_memory(const std::string& protocol, void* ptr); } } -// String utility functions - /** * @brief Split a string by delimiter into a vector of strings * @param str The string to split @@ -245,6 +242,93 @@ std::vector splitString(const std::string& str, bool trim_spaces = true, bool keep_empty = false); +// Buffer allocator functions + +constexpr size_t SZ_2MB = 2 * 1024 * 1024; +constexpr size_t SZ_1GB = 1024 * 1024 * 1024; + +/** + * @brief Allocates memory for the `BufferAllocator` class. + * @param total_size The total size of the memory to allocate. + * @return A pointer to the allocated memory. + */ +void* allocate_buffer_allocator_memory( + size_t total_size, const std::string& protocol = "", + size_t alignment = facebook::cachelib::Slab::kSize); + +inline size_t align_up(size_t size, size_t alignment) { + if (alignment == 0) { + return size; + } + return ((size + alignment - 1) / alignment) * alignment; +} + +/** + * @brief Get hugepage size from env and optionally set the corresponding memfd + * flags. + * * @param out_flags Optional pointer to an int. If provided, + * MAP_HUGETLB and MAP_HUGE_2MB/1GB will be OR-ed into it. + * @return size_t Hugepage size in bytes, or 0 if disabled. + */ +[[nodiscard]] inline size_t get_hugepage_size_from_env( + unsigned int* out_flags = nullptr, bool use_memfd = false) { + const char* use_hp_env = std::getenv("MC_STORE_USE_HUGEPAGE"); + if (use_hp_env == nullptr) { + return 0; + } + + constexpr size_t SZ_2MB = 2 * 1024 * 1024; + constexpr size_t SZ_1GB = 1024 * 1024 * 1024; + + size_t size = SZ_2MB; // Default to 2MB + + const char* size_env = std::getenv("MC_STORE_HUGEPAGE_SIZE"); + if (size_env != nullptr) { + size_t parsed_size = string_to_byte_size(size_env); + + if (parsed_size == SZ_2MB || parsed_size == SZ_1GB) { + size = parsed_size; + } else { + LOG(WARNING) << "Invalid MC_STORE_HUGEPAGE_SIZE='" << size_env + << "'. Supported: 2MB, 1GB. Fallback to 2MB."; + size = SZ_2MB; + } + } + + if (out_flags != nullptr) { + if (use_memfd) { + *out_flags |= MFD_HUGETLB; + } else { + *out_flags |= MAP_HUGETLB; + } + + // Add size specific flag + if (size == SZ_2MB) { + if (use_memfd) { + *out_flags |= MFD_HUGE_2MB; + } else { + *out_flags |= MAP_HUGE_2MB; + } + } else if (size == SZ_1GB) { + if (use_memfd) { + *out_flags |= MFD_HUGE_1GB; + } else { + *out_flags |= MAP_HUGE_1GB; + } + } + LOG(INFO) << "Using hugepage size: " + << (size == SZ_2MB ? "2MB" : "1GB"); + } + + return size; +} + +// Hugepage-backed allocation helpers (MAP_HUGETLB + MADV_HUGEPAGE) +void* allocate_buffer_mmap_memory(size_t total_size, size_t alignment); +void free_buffer_mmap_memory(void* ptr, size_t total_size); + +void free_memory(const std::string& protocol, void* ptr); + // Network utility functions /** diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp index f51d5ef3..a78b9d0c 100644 --- a/mooncake-store/src/client_buffer.cpp +++ b/mooncake-store/src/client_buffer.cpp @@ -13,9 +13,9 @@ namespace mooncake { std::shared_ptr ClientBufferAllocator::create( - size_t size, const std::string& protocol) { + size_t size, const std::string& protocol, bool use_hugepage) { return std::shared_ptr( - new ClientBufferAllocator(size, protocol)); + new ClientBufferAllocator(size, protocol, use_hugepage)); } std::shared_ptr ClientBufferAllocator::create( @@ -25,8 +25,9 @@ std::shared_ptr ClientBufferAllocator::create( } ClientBufferAllocator::ClientBufferAllocator(size_t size, - const std::string& protocol) - : protocol(protocol), buffer_size_(size) { + const std::string& protocol, + bool use_hugepage) + : protocol(protocol), buffer_size_(size), use_hugepage_(use_hugepage) { if (size == 0) { buffer_ = nullptr; allocator_ = nullptr; @@ -34,7 +35,11 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size, } // Align to 64 bytes(cache line size) for better cache performance constexpr size_t alignment = 64; - buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment); + if (use_hugepage_) { + buffer_ = allocate_buffer_mmap_memory(size, alignment); + } else { + buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment); + } if (!buffer_) { throw std::bad_alloc(); } @@ -55,7 +60,11 @@ ClientBufferAllocator::ClientBufferAllocator(void* addr, size_t size, ClientBufferAllocator::~ClientBufferAllocator() { // Free the aligned allocated memory or unmap shared memory if (!is_external_memory_ && buffer_) { - free_memory(protocol, buffer_); + if (use_hugepage_) { + free_buffer_mmap_memory(buffer_, buffer_size_); + } else { + free_memory(protocol, buffer_); + } } } diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index 41c113a6..f9f444e7 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -6,9 +6,11 @@ #include // For S_IRUSR, S_IWUSR #include // For O_CREAT, O_RDWR #include // For ftruncate, close, shm_unlink +#include #include "real_client.h" #include "dummy_client.h" +#include "utils.h" #include "utils/scoped_vlog_timer.h" #include "rpc_types.h" #include "types.h" @@ -31,7 +33,10 @@ ShmHelper* ShmHelper::getInstance() { return &instance; } -ShmHelper::ShmHelper() {} +ShmHelper::ShmHelper() { + const char* hp = std::getenv("MC_STORE_USE_HUGEPAGE"); + use_hugepage_ = (hp != nullptr); +} ShmHelper::~ShmHelper() { cleanup(); } @@ -59,10 +64,20 @@ bool ShmHelper::cleanup() { void* ShmHelper::allocate(size_t size) { std::lock_guard lock(shm_mutex_); + unsigned int flags = MFD_CLOEXEC; + if (use_hugepage_) { + bool use_memfd = true; + size = align_up(size, get_hugepage_size_from_env(&flags, use_memfd)); + LOG(INFO) << "Using huge pages for shared memory, size: " << size; + } + // Create memfd - int fd = memfd_create_wrapper(MOONCAKE_SHM_NAME, MFD_CLOEXEC); + int fd = memfd_create_wrapper(MOONCAKE_SHM_NAME, flags); if (fd == -1) { - throw std::runtime_error("Failed to create anonymous shared memory: " + + std::string extra_msg = + use_hugepage_ ? " (Check /proc/sys/vm/nr_hugepages?)" : ""; + throw std::runtime_error("Failed to create anonymous shared memory" + + extra_msg + ": " + std::string(strerror(errno))); } @@ -74,8 +89,8 @@ void* ShmHelper::allocate(size_t size) { } // Map memory - void* base_addr = - mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + void* base_addr = mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_POPULATE, fd, 0); if (base_addr == MAP_FAILED) { close(fd); throw std::runtime_error("Failed to map shared memory: " + @@ -413,14 +428,17 @@ int DummyClient::register_buffer(void* buffer, size_t size) { LOG(ERROR) << "Buffer is not in any registered shared memory"; return -1; } + if (shm_helper_->is_hugepage()) { + size = align_up(size, get_hugepage_size_from_env()); + } // Check bounds if (reinterpret_cast(buffer) != reinterpret_cast(shm->base_addr) || size != shm->size) { - LOG(ERROR) - << "Invalid buffer address or size for registration: Buffer addr: " - << buffer << ", need addr: " << shm->base_addr - << ", buffer size: " << size << ", need size: " << shm->size; + LOG(ERROR) << "Invalid buffer address or size for registration: " + "Buffer addr: " + << buffer << ", need addr: " << shm->base_addr + << ", buffer size: " << size << ", need size: " << shm->size; return -1; } diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 8d1065c8..92d53e53 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -9,6 +9,8 @@ #include #include // for atexit +#include +#include #include #include @@ -143,6 +145,8 @@ void ResourceTracker::startSignalThread() { RealClient::RealClient() { // Initialize logging severity (leave as before) mooncake::init_ylt_log_level(); + const char *hp = std::getenv("MC_STORE_USE_HUGEPAGE"); + use_hugepage_ = (hp != nullptr); } RealClient::~RealClient() { @@ -165,6 +169,8 @@ tl::expected RealClient::setup_internal( const std::string &ipc_socket_path, bool enable_offload) { this->protocol = protocol; this->ipc_socket_path_ = ipc_socket_path; + const bool should_use_hugepage = + use_hugepage_ && this->protocol != "ascend"; // Remove port if hostname already contains one std::string hostname = local_hostname; @@ -200,8 +206,8 @@ tl::expected RealClient::setup_internal( // fail in some rdma implementations. // Dummy Client can create shm and share it with Real Client, so Real Client // can create client buffer allocator on the shared memory later. - client_buffer_allocator_ = - ClientBufferAllocator::create(local_buffer_size, this->protocol); + client_buffer_allocator_ = ClientBufferAllocator::create( + local_buffer_size, this->protocol, should_use_hugepage); if (local_buffer_size > 0) { LOG(INFO) << "Registering local memory: " << local_buffer_size << " bytes"; @@ -229,18 +235,30 @@ tl::expected RealClient::setup_internal( current_glbseg_size += segment_size; LOG(INFO) << "Mounting segment: " << segment_size << " bytes, " << current_glbseg_size << " of " << total_glbseg_size; - void *ptr = - allocate_buffer_allocator_memory(segment_size, this->protocol); + size_t mapped_size = segment_size; + void *ptr = nullptr; + if (should_use_hugepage) { + mapped_size = align_up(segment_size, get_hugepage_size_from_env()); + ptr = allocate_buffer_mmap_memory(mapped_size, + get_hugepage_size_from_env()); + } else { + ptr = + allocate_buffer_allocator_memory(segment_size, this->protocol); + } + if (!ptr) { LOG(ERROR) << "Failed to allocate segment memory"; return tl::unexpected(ErrorCode::INVALID_PARAMS); } if (this->protocol == "ascend") { ascend_segment_ptrs_.emplace_back(ptr); + } else if (should_use_hugepage) { + hugepage_segment_ptrs_.emplace_back( + ptr, HugepageSegmentDeleter{mapped_size}); } else { segment_ptrs_.emplace_back(ptr); } - auto mount_result = client_->MountSegment(ptr, segment_size); + auto mount_result = client_->MountSegment(ptr, mapped_size); if (!mount_result.has_value()) { LOG(ERROR) << "Failed to mount segment: " << toString(mount_result.error()); @@ -324,6 +342,7 @@ tl::expected RealClient::tearDownAll_internal() { client_.reset(); client_buffer_allocator_.reset(); port_binder_.reset(); + hugepage_segment_ptrs_.clear(); segment_ptrs_.clear(); local_hostname = ""; device_name = ""; diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 9d255070..f24b4064 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -7,7 +7,11 @@ #include #include +#include #include +#include +#include +#include #ifdef USE_ASCEND_DIRECT #include "acl/acl.h" #include "config.h" @@ -136,6 +140,39 @@ void *allocate_buffer_allocator_memory(size_t total_size, return aligned_alloc(alignment, total_size); } +void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) { + if (total_size == 0) { + LOG(ERROR) << "Total size must be greater than 0 for hugepage mmap"; + return nullptr; + } + + unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE; + const size_t effective_alignment = + std::max(alignment, get_hugepage_size_from_env(&flags)); + const size_t map_size = align_up(total_size, effective_alignment); + + void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, flags, -1, 0); + if (ptr == MAP_FAILED) { + LOG(ERROR) << "Hugepage mmap failed, size=" << map_size + << ", errno=" << errno << " (" << strerror(errno) << ")"; + return nullptr; + } + + return ptr; +} + +void free_buffer_mmap_memory(void *ptr, size_t total_size) { + if (!ptr || total_size == 0) { + return; + } + + const size_t map_size = align_up(total_size, get_hugepage_size_from_env()); + if (munmap(ptr, map_size) != 0) { + LOG(ERROR) << "munmap hugepage failed, size=" << map_size + << ", errno=" << errno << " (" << strerror(errno) << ")"; + } +} + void free_memory(const std::string &protocol, void *ptr) { #ifdef USE_ASCEND_DIRECT if (protocol == "ascend") {