From 90294fda8a24fe57d36c49c840d413c6acc6c514 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Wed, 15 Jul 2026 21:09:53 +0800 Subject: [PATCH 1/9] Store: pin host segments with quota --- .../python-api-reference/mooncake-store.md | 18 ++ mooncake-store/include/real_client.h | 4 + mooncake-store/src/CMakeLists.txt | 1 + mooncake-store/src/real_client.cpp | 34 ++- .../src/registered_pinned_memory.cpp | 256 ++++++++++++++++++ mooncake-store/src/registered_pinned_memory.h | 75 +++++ 6 files changed, 387 insertions(+), 1 deletion(-) create mode 100644 mooncake-store/src/registered_pinned_memory.cpp create mode 100644 mooncake-store/src/registered_pinned_memory.h diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/python-api-reference/mooncake-store.md index c738f5ef18..e3ae8467bb 100644 --- a/docs/source/python-api-reference/mooncake-store.md +++ b/docs/source/python-api-reference/mooncake-store.md @@ -1081,6 +1081,24 @@ def setup( - `ssd_offload_path` (str): SSD offload directory. When provided, overrides the storage path environment configuration. - `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`. +**Store segment pinned memory:** + +CUDA-enabled builds can register Store-managed host segments as pinned memory +when the process is given an explicit pinned-memory quota. This applies only to +Store segment memory for host transfer protocols allocated by `setup()` +(`global_segment_size`) and `allocateAndMountSegment()`. It does not pin +file-backed `mountSegment()` mappings, CXL/device segments, the +`local_buffer_size` client buffer, user buffers, dummy-client shared memory, or +temporary staging buffers. + +- `MC_STORE_PIN_MEMORY_MAX_BYTES`: positive process-wide quota, in bytes. Store + segment pinning is disabled when this is unset, empty, `0`, or invalid. +- `MC_STORE_PIN_MEMORY`: optional override. Set to `0`, `false`, `off`, or `no` + to disable Store segment pinning even when a quota is configured. + +If the quota is exhausted or CUDA registration fails, Mooncake continues with +pageable Store segment memory. + **Returns:** - `int`: Status code (0 = success, non-zero = error code) diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 6b3382b3b6..1492ff6779 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -29,6 +29,7 @@ namespace mooncake { class RealClient; +class RegisteredPinnedRegion; class UdsAcceptor; class UdsConnection; @@ -756,6 +757,7 @@ class RealClient : public PyClient { void *base = nullptr; size_t size = 0; std::string protocol; + std::shared_ptr pinned_region; }; std::unique_ptr port_binder_ = nullptr; @@ -816,6 +818,8 @@ class RealClient : public PyClient { std::vector> sunrise_segment_ptrs_; #endif + std::vector> + setup_segment_pinned_regions_; std::string protocol; std::string device_name; std::string local_hostname; diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 0b8104f965..c97cb5895a 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -27,6 +27,7 @@ set(MOONCAKE_STORE_SOURCES client_buffer.cpp aligned_client_buffer.cpp real_client.cpp + registered_pinned_memory.cpp dummy_client.cpp uds_transport.cpp shm_helper.cpp diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 314d7cf511..7c2ba64a0c 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -19,6 +19,7 @@ #include #include "real_client.h" +#include "registered_pinned_memory.h" #include "client_buffer.h" #include "replica_selection.h" #include "common.h" @@ -51,6 +52,21 @@ namespace mooncake { namespace { constexpr std::chrono::seconds kIpcRequestRecvTimeout{5}; +bool IsHostStoreSegmentProtocol(const std::string &protocol) { + return protocol.empty() || protocol == "tcp" || protocol == "rdma" || + protocol == "efa" || protocol == "cxi" || protocol == "rpc_only"; +} + +std::shared_ptr TryPinStoreSegment( + void *ptr, size_t size, const std::string &protocol, + const char *segment_owner) { + if (!IsHostStoreSegmentProtocol(protocol)) return nullptr; + return RegisteredPinnedMemoryManager::instance().try_pin( + ptr, size, + std::string("Store segment ") + segment_owner + + " protocol=" + protocol); +} + #ifdef USE_ASCEND_DIRECT bool checkAcl(aclError result, const char *message) { if (result != ACL_ERROR_NONE) { @@ -869,13 +885,20 @@ tl::expected RealClient::setup_internal( } } + auto pinned_region = + TryPinStoreSegment(ptr, mapped_size, this->protocol, "setup"); auto mount_result = client_->MountSegment(ptr, mapped_size, protocol, seg_location); if (!mount_result.has_value()) { + pinned_region.reset(); LOG(ERROR) << "Failed to mount segment: " << toString(mount_result.error()); return tl::unexpected(mount_result.error()); } + if (pinned_region) { + setup_segment_pinned_regions_.push_back( + std::move(pinned_region)); + } } if (total_glbseg_size == 0) { LOG(INFO) << "Global segment size is 0, skip mounting segment"; @@ -1187,6 +1210,7 @@ tl::expected RealClient::tearDownAll_internal() { ReleaseAllAllocatedSegmentRecords(); client_buffer_allocator_.reset(); port_binder_.reset(); + setup_segment_pinned_regions_.clear(); hugepage_segment_ptrs_.clear(); ub_segment_ptrs_.clear(); #if defined(USE_SUNRISE) @@ -1380,6 +1404,7 @@ void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { } } if (found && record.base) { + record.pinned_region.reset(); free_memory(record.protocol, record.base); } } @@ -1392,6 +1417,7 @@ void RealClient::ReleaseAllAllocatedSegmentRecords() { } for (auto &entry : records) { if (entry.second.base) { + entry.second.pinned_region.reset(); free_memory(entry.second.protocol, entry.second.base); } } @@ -1533,17 +1559,21 @@ int RealClient::allocateAndMountSegment( break; } + auto pinned_region = + TryPinStoreSegment(ptr, chunk_size, protocol, "allocated"); auto result = client_->MountSegmentAndGetId(ptr, chunk_size, protocol, location); if (!result.has_value()) { LOG(ERROR) << "MountSegmentAndGetId failed"; + pinned_region.reset(); free_memory(protocol, ptr); break; } std::string segment_id = UuidToString(result.value()); mounted_ids.push_back(segment_id); - allocated_records.push_back({ptr, chunk_size, protocol}); + allocated_records.push_back( + {ptr, chunk_size, protocol, std::move(pinned_region)}); remaining -= chunk_size; } @@ -1555,6 +1585,7 @@ int RealClient::allocateAndMountSegment( client_->UnmountSegmentById(id); } if (allocated_records[i].base) { + allocated_records[i].pinned_region.reset(); free_memory(allocated_records[i].protocol, allocated_records[i].base); } @@ -1643,6 +1674,7 @@ int RealClient::unmountAndFreeSegment( for (auto &p : to_cleanup) { if (p.second.base) { + p.second.pinned_region.reset(); free_memory(p.second.protocol, p.second.base); } } diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp new file mode 100644 index 0000000000..7df2dcd4f4 --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -0,0 +1,256 @@ +#include "registered_pinned_memory.h" + +#include +#include +#include +#include +#include +#include + +#include + +#if defined(USE_CUDA) +#include +#endif + +namespace mooncake { +namespace { + +bool ParsePinnedMemoryEnabled() { + const char* value = std::getenv("MC_STORE_PIN_MEMORY"); + if (!value) return true; + + std::string normalized(value); + std::transform( + normalized.begin(), normalized.end(), normalized.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + return !(normalized == "0" || normalized == "false" || + normalized == "off" || normalized == "no"); +} + +std::optional ParsePinnedMemoryLimit() { + const char* value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); + if (!value || value[0] == '\0') return 0; + + const char* number = value; + while (std::isspace(static_cast(*number))) { + ++number; + } + if (*number == '-') { + LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value + << "', disabling Store segment pinning"; + return std::nullopt; + } + + char* end = nullptr; + errno = 0; + unsigned long long parsed = std::strtoull(number, &end, 10); + while (end && std::isspace(static_cast(*end))) { + ++end; + } + if (end == number || (end && *end != '\0') || errno == ERANGE || + parsed > std::numeric_limits::max()) { + LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value + << "', disabling Store segment pinning"; + return std::nullopt; + } + return static_cast(parsed); +} + +std::pair ParsePinnedMemoryConfig() { + if (!ParsePinnedMemoryEnabled()) return {false, 0}; + + auto limit = ParsePinnedMemoryLimit(); + if (!limit.has_value() || *limit == 0) { + return {false, 0}; + } + return {true, *limit}; +} + +} // namespace + +RegisteredPinnedRegion::~RegisteredPinnedRegion() { + RegisteredPinnedMemoryManager::instance().release(this); +} + +RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { + static RegisteredPinnedMemoryManager manager; + return manager; +} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() + : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig()) {} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( + std::pair config) + : enabled_(config.first), limit_bytes_(config.second) { +#if defined(USE_CUDA) + LOG(INFO) << "Store segment pinned memory is " + << (enabled_ ? "enabled" : "disabled") + << ", max_bytes=" << limit_bytes_; +#else + if (enabled_) { + LOG(INFO) << "Store segment pinning requested but this build has no " + "CUDA runtime support"; + } +#endif +} + +std::shared_ptr RegisteredPinnedMemoryManager::try_pin( + void* addr, size_t size, const std::string& owner) { + if (!addr || size == 0 || !enabled_) return nullptr; + +#if !defined(USE_CUDA) + (void)owner; + return nullptr; +#else + RegionKey key{addr, size}; + const auto start = reinterpret_cast(addr); + const auto end = start + size; + if (end < start) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": address range overflow, size=" << size; + return nullptr; + } + + std::shared_ptr region; + try { + region.reset(new RegisteredPinnedRegion(addr, size, owner)); + } catch (...) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": failed to allocate pin tracking, size=" << size; + return nullptr; + } + + { + std::lock_guard lock(mutex_); + for (const auto& entry : regions_) { + const auto region_start = + reinterpret_cast(entry.first.addr); + const auto region_end = region_start + entry.first.size; + if (region_end < region_start) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": existing active range overflow, size=" + << entry.first.size; + return nullptr; + } + + const bool overlaps = start < region_end && end > region_start; + if (overlaps) { + LOG(WARNING) + << "Skip cudaHostRegister for " << owner + << ": overlaps an active pinned region, size=" << size; + return nullptr; + } + } + + if (size > limit_bytes_ || pinned_bytes_ > limit_bytes_ - size) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": quota exceeded, requested=" << size + << ", pinned=" << pinned_bytes_ + << ", limit=" << limit_bytes_; + return nullptr; + } + + try { + auto [_, inserted] = regions_.emplace(key, nullptr); + if (!inserted) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": active region already exists, size=" << size; + return nullptr; + } + } catch (...) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": failed to allocate pin tracking, size=" << size; + return nullptr; + } + pinned_bytes_ += size; + } + + cudaError_t err = cudaHostRegister(addr, size, cudaHostRegisterPortable); + if (err != cudaSuccess) { + { + std::lock_guard lock(mutex_); + drop_reservation_locked(key, size); + } + LOG(WARNING) << "cudaHostRegister failed for " << owner + << ", size=" << size + << ", error=" << cudaGetErrorString(err) + << ". Continue with pageable host memory."; + cudaGetLastError(); + return nullptr; + } + + bool tracking_ready = false; + uint64_t pinned_bytes = 0; + { + std::lock_guard lock(mutex_); + auto it = regions_.find(key); + if (it != regions_.end() && it->second == nullptr) { + it->second = region.get(); + pinned_bytes = pinned_bytes_; + tracking_ready = true; + } + } + if (!tracking_ready) { + err = cudaHostUnregister(addr); + if (err != cudaSuccess) { + LOG(FATAL) << "cudaHostUnregister failed after active range " + "tracking mismatch for " + << owner << ", size=" << size + << ", error=" << cudaGetErrorString(err); + } + std::lock_guard lock(mutex_); + drop_reservation_locked(key, size); + return nullptr; + } + + LOG(INFO) << "cudaHostRegister succeeded for " << owner << ", size=" << size + << ", pinned=" << pinned_bytes << ", limit=" << limit_bytes_; + return region; +#endif +} + +void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { + if (!region || !region->addr_ || region->size_ == 0) return; + + RegionKey key{region->addr_, region->size_}; + { + std::lock_guard lock(mutex_); + auto region_it = regions_.find(key); + if (region_it == regions_.end() || region_it->second != region) { + return; + } + region_it->second = nullptr; + } + +#if defined(USE_CUDA) + cudaError_t err = cudaHostUnregister(region->addr_); + if (err != cudaSuccess) { + LOG(FATAL) << "cudaHostUnregister failed for " << region->owner_ + << ", size=" << region->size_ + << ", error=" << cudaGetErrorString(err); + } +#endif + + std::lock_guard lock(mutex_); + auto it = regions_.find(key); + if (it == regions_.end() || it->second != nullptr) return; + regions_.erase(it); + if (pinned_bytes_ >= region->size_) { + pinned_bytes_ -= region->size_; + } else { + pinned_bytes_ = 0; + } +} + +void RegisteredPinnedMemoryManager::drop_reservation_locked( + const RegionKey& key, size_t size) { + auto it = regions_.find(key); + if (it != regions_.end() && it->second == nullptr) { + regions_.erase(it); + pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; + } +} + +} // namespace mooncake diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h new file mode 100644 index 0000000000..d499f1bf24 --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.h @@ -0,0 +1,75 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +class RegisteredPinnedMemoryManager; + +class RegisteredPinnedRegion { + public: + RegisteredPinnedRegion(const RegisteredPinnedRegion&) = delete; + RegisteredPinnedRegion& operator=(const RegisteredPinnedRegion&) = delete; + ~RegisteredPinnedRegion(); + + private: + friend class RegisteredPinnedMemoryManager; + + RegisteredPinnedRegion(void* addr, size_t size, std::string owner) + : addr_(addr), size_(size), owner_(std::move(owner)) {} + + void* addr_ = nullptr; + size_t size_ = 0; + std::string owner_; +}; + +class RegisteredPinnedMemoryManager { + public: + static RegisteredPinnedMemoryManager& instance(); + + std::shared_ptr try_pin(void* addr, size_t size, + const std::string& owner); + + private: + friend class RegisteredPinnedRegion; + + struct RegionKey { + void* addr = nullptr; + size_t size = 0; + + bool operator==(const RegionKey& other) const { + return addr == other.addr && size == other.size; + } + }; + + struct RegionKeyHash { + size_t operator()(const RegionKey& key) const { + const size_t addr_hash = std::hash{}(key.addr); + return addr_hash ^ + (std::hash{}(key.size) + 0x9e3779b97f4a7c15ULL + + (addr_hash << 6) + (addr_hash >> 2)); + } + }; + + RegisteredPinnedMemoryManager(); + explicit RegisteredPinnedMemoryManager(std::pair config); + + void release(RegisteredPinnedRegion* region); + void drop_reservation_locked(const RegionKey& key, size_t size); + + const bool enabled_; + const uint64_t limit_bytes_; + + mutable std::mutex mutex_; + uint64_t pinned_bytes_ = 0; + std::unordered_map + regions_; +}; + +} // namespace mooncake From 0c339467b28fb9e23dcdb5b494f01ab6ea250d78 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Thu, 16 Jul 2026 11:05:50 +0800 Subject: [PATCH 2/9] Store: address pinned memory review comments --- .../src/registered_pinned_memory.cpp | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 7df2dcd4f4..d9eeea6b9a 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -1,11 +1,11 @@ #include "registered_pinned_memory.h" #include -#include #include +#include #include -#include #include +#include #include @@ -42,19 +42,23 @@ std::optional ParsePinnedMemoryLimit() { return std::nullopt; } - char* end = nullptr; - errno = 0; - unsigned long long parsed = std::strtoull(number, &end, 10); - while (end && std::isspace(static_cast(*end))) { + const char* end = number; + while (*end != '\0') { ++end; } - if (end == number || (end && *end != '\0') || errno == ERANGE || - parsed > std::numeric_limits::max()) { + while (end > number && + std::isspace(static_cast(*(end - 1)))) { + --end; + } + + uint64_t parsed = 0; + auto [ptr, ec] = std::from_chars(number, end, parsed); + if (ec != std::errc{} || ptr != end) { LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value << "', disabling Store segment pinning"; return std::nullopt; } - return static_cast(parsed); + return parsed; } std::pair ParsePinnedMemoryConfig() { @@ -74,8 +78,9 @@ RegisteredPinnedRegion::~RegisteredPinnedRegion() { } RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { - static RegisteredPinnedMemoryManager manager; - return manager; + static RegisteredPinnedMemoryManager* manager = + new RegisteredPinnedMemoryManager(); + return *manager; } RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() From cfa0e9e1cc5202e19e62cb5688b1b0a9e3cb4863 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Thu, 16 Jul 2026 11:28:43 +0800 Subject: [PATCH 3/9] Store: harden pinned memory cleanup --- .../src/registered_pinned_memory.cpp | 48 ++++++++++++------- 1 file changed, 31 insertions(+), 17 deletions(-) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index d9eeea6b9a..f2ea092710 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -16,11 +17,30 @@ namespace mooncake { namespace { +std::pair TrimAsciiWhitespace(const char* value) { + const char* begin = value; + while (std::isspace(static_cast(*begin))) { + ++begin; + } + + const char* end = begin; + while (*end != '\0') { + ++end; + } + while (end > begin && + std::isspace(static_cast(*(end - 1)))) { + --end; + } + + return {begin, end}; +} + bool ParsePinnedMemoryEnabled() { const char* value = std::getenv("MC_STORE_PIN_MEMORY"); if (!value) return true; - std::string normalized(value); + auto [begin, end] = TrimAsciiWhitespace(value); + std::string normalized(begin, end); std::transform( normalized.begin(), normalized.end(), normalized.begin(), [](unsigned char ch) { return static_cast(std::tolower(ch)); }); @@ -32,25 +52,13 @@ std::optional ParsePinnedMemoryLimit() { const char* value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); if (!value || value[0] == '\0') return 0; - const char* number = value; - while (std::isspace(static_cast(*number))) { - ++number; - } + auto [number, end] = TrimAsciiWhitespace(value); if (*number == '-') { LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value << "', disabling Store segment pinning"; return std::nullopt; } - const char* end = number; - while (*end != '\0') { - ++end; - } - while (end > number && - std::isspace(static_cast(*(end - 1)))) { - --end; - } - uint64_t parsed = 0; auto [ptr, ec] = std::from_chars(number, end, parsed); if (ec != std::errc{} || ptr != end) { @@ -232,9 +240,15 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { #if defined(USE_CUDA) cudaError_t err = cudaHostUnregister(region->addr_); if (err != cudaSuccess) { - LOG(FATAL) << "cudaHostUnregister failed for " << region->owner_ - << ", size=" << region->size_ - << ", error=" << cudaGetErrorString(err); + if (err == cudaErrorCudartUnloading) { + LOG(WARNING) << "Skip cudaHostUnregister for " << region->owner_ + << " because CUDA runtime is unloading, size=" + << region->size_; + } else { + LOG(FATAL) << "cudaHostUnregister failed for " << region->owner_ + << ", size=" << region->size_ + << ", error=" << cudaGetErrorString(err); + } } #endif From 88d8bdaa1806caa06effac8af913e60b111a15ba Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 17 Jul 2026 10:43:17 +0800 Subject: [PATCH 4/9] Store: test pinned memory bookkeeping --- .../src/registered_pinned_memory.cpp | 83 ++++++++---- mooncake-store/src/registered_pinned_memory.h | 29 ++++- mooncake-store/tests/CMakeLists.txt | 1 + .../tests/registered_pinned_memory_test.cpp | 119 ++++++++++++++++++ 4 files changed, 207 insertions(+), 25 deletions(-) create mode 100644 mooncake-store/tests/registered_pinned_memory_test.cpp diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index f2ea092710..6b2f7371dd 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -79,10 +79,44 @@ std::pair ParsePinnedMemoryConfig() { return {true, *limit}; } +#if defined(USE_CUDA) +bool RegisterPinnedRegionWithCuda(void* addr, size_t size, + std::string* error_message) { + cudaError_t err = cudaHostRegister(addr, size, cudaHostRegisterPortable); + if (err == cudaSuccess) return true; + if (error_message) *error_message = cudaGetErrorString(err); + cudaGetLastError(); + return false; +} + +RegisteredPinnedMemoryManager::UnregisterResult UnregisterPinnedRegionWithCuda( + void* addr, std::string* error_message) { + cudaError_t err = cudaHostUnregister(addr); + if (err == cudaSuccess) { + return RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; + } + if (error_message) *error_message = cudaGetErrorString(err); + if (err == cudaErrorCudartUnloading) { + return RegisteredPinnedMemoryManager::UnregisterResult:: + kRuntimeUnloading; + } + return RegisteredPinnedMemoryManager::UnregisterResult::kError; +} + +#endif + +RegisteredPinnedMemoryManager::PinOps DefaultPinOps() { +#if defined(USE_CUDA) + return {RegisterPinnedRegionWithCuda, UnregisterPinnedRegionWithCuda}; +#else + return {}; +#endif +} + } // namespace RegisteredPinnedRegion::~RegisteredPinnedRegion() { - RegisteredPinnedMemoryManager::instance().release(this); + if (manager_) manager_->release(this); } RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { @@ -96,7 +130,11 @@ RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( std::pair config) - : enabled_(config.first), limit_bytes_(config.second) { + : RegisteredPinnedMemoryManager(config, DefaultPinOps()) {} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( + std::pair config, PinOps pin_ops) + : enabled_(config.first), limit_bytes_(config.second), pin_ops_(pin_ops) { #if defined(USE_CUDA) LOG(INFO) << "Store segment pinned memory is " << (enabled_ ? "enabled" : "disabled") @@ -112,11 +150,10 @@ RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( std::shared_ptr RegisteredPinnedMemoryManager::try_pin( void* addr, size_t size, const std::string& owner) { if (!addr || size == 0 || !enabled_) return nullptr; + if (!pin_ops_.register_region || !pin_ops_.unregister_region) { + return nullptr; + } -#if !defined(USE_CUDA) - (void)owner; - return nullptr; -#else RegionKey key{addr, size}; const auto start = reinterpret_cast(addr); const auto end = start + size; @@ -128,7 +165,7 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( std::shared_ptr region; try { - region.reset(new RegisteredPinnedRegion(addr, size, owner)); + region.reset(new RegisteredPinnedRegion(this, addr, size, owner)); } catch (...) { LOG(WARNING) << "Skip cudaHostRegister for " << owner << ": failed to allocate pin tracking, size=" << size; @@ -180,17 +217,17 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( pinned_bytes_ += size; } - cudaError_t err = cudaHostRegister(addr, size, cudaHostRegisterPortable); - if (err != cudaSuccess) { + std::string error_message; + const bool registered = + pin_ops_.register_region(addr, size, &error_message); + if (!registered) { { std::lock_guard lock(mutex_); drop_reservation_locked(key, size); } LOG(WARNING) << "cudaHostRegister failed for " << owner - << ", size=" << size - << ", error=" << cudaGetErrorString(err) + << ", size=" << size << ", error=" << error_message << ". Continue with pageable host memory."; - cudaGetLastError(); return nullptr; } @@ -206,12 +243,14 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( } } if (!tracking_ready) { - err = cudaHostUnregister(addr); - if (err != cudaSuccess) { + error_message.clear(); + auto unregister_result = + pin_ops_.unregister_region(addr, &error_message); + if (unregister_result != UnregisterResult::kSuccess) { LOG(FATAL) << "cudaHostUnregister failed after active range " "tracking mismatch for " << owner << ", size=" << size - << ", error=" << cudaGetErrorString(err); + << ", error=" << error_message; } std::lock_guard lock(mutex_); drop_reservation_locked(key, size); @@ -221,7 +260,6 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( LOG(INFO) << "cudaHostRegister succeeded for " << owner << ", size=" << size << ", pinned=" << pinned_bytes << ", limit=" << limit_bytes_; return region; -#endif } void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { @@ -237,20 +275,21 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { region_it->second = nullptr; } -#if defined(USE_CUDA) - cudaError_t err = cudaHostUnregister(region->addr_); - if (err != cudaSuccess) { - if (err == cudaErrorCudartUnloading) { + std::string error_message; + if (!pin_ops_.unregister_region) return; + auto unregister_result = + pin_ops_.unregister_region(region->addr_, &error_message); + if (unregister_result != UnregisterResult::kSuccess) { + if (unregister_result == UnregisterResult::kRuntimeUnloading) { LOG(WARNING) << "Skip cudaHostUnregister for " << region->owner_ << " because CUDA runtime is unloading, size=" << region->size_; } else { LOG(FATAL) << "cudaHostUnregister failed for " << region->owner_ << ", size=" << region->size_ - << ", error=" << cudaGetErrorString(err); + << ", error=" << error_message; } } -#endif std::lock_guard lock(mutex_); auto it = regions_.find(key); diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index d499f1bf24..ac82b79e48 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -21,9 +21,14 @@ class RegisteredPinnedRegion { private: friend class RegisteredPinnedMemoryManager; - RegisteredPinnedRegion(void* addr, size_t size, std::string owner) - : addr_(addr), size_(size), owner_(std::move(owner)) {} - + RegisteredPinnedRegion(RegisteredPinnedMemoryManager* manager, void* addr, + size_t size, std::string owner) + : manager_(manager), + addr_(addr), + size_(size), + owner_(std::move(owner)) {} + + RegisteredPinnedMemoryManager* manager_ = nullptr; void* addr_ = nullptr; size_t size_ = 0; std::string owner_; @@ -31,6 +36,15 @@ class RegisteredPinnedRegion { class RegisteredPinnedMemoryManager { public: + enum class UnregisterResult { kSuccess, kRuntimeUnloading, kError }; + + struct PinOps { + bool (*register_region)(void* addr, size_t size, + std::string* error_message) = nullptr; + UnregisterResult (*unregister_region)( + void* addr, std::string* error_message) = nullptr; + }; + static RegisteredPinnedMemoryManager& instance(); std::shared_ptr try_pin(void* addr, size_t size, @@ -59,12 +73,21 @@ class RegisteredPinnedMemoryManager { RegisteredPinnedMemoryManager(); explicit RegisteredPinnedMemoryManager(std::pair config); +#if defined(MOONCAKE_STORE_TEST) + public: +#endif + RegisteredPinnedMemoryManager(std::pair config, + PinOps pin_ops); +#if defined(MOONCAKE_STORE_TEST) + private: +#endif void release(RegisteredPinnedRegion* region); void drop_reservation_locked(const RegionKey& key, size_t size); const bool enabled_; const uint64_t limit_bytes_; + const PinOps pin_ops_; mutable std::mutex mutex_; uint64_t pinned_bytes_ = 0; diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index e7e66c9fbb..d32f8dfa8f 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -35,6 +35,7 @@ endfunction() add_store_test(buffer_allocator_test buffer_allocator_test.cpp) add_store_test(runtime_accelerator_test runtime_accelerator_test.cpp) +add_store_test(registered_pinned_memory_test registered_pinned_memory_test.cpp) add_store_test(allocation_strategy_test allocation_strategy_test.cpp) add_store_test(replica_selection_test replica_selection_test.cpp) add_test( diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp new file mode 100644 index 0000000000..be5cbea46f --- /dev/null +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -0,0 +1,119 @@ +#define MOONCAKE_STORE_TEST +#include "../src/registered_pinned_memory.h" + +#include +#include + +#include + +namespace mooncake { +namespace { + +struct FakePinState { + bool register_succeeds = true; + RegisteredPinnedMemoryManager::UnregisterResult unregister_result = + RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; + int register_calls = 0; + int unregister_calls = 0; +}; + +FakePinState& State() { + static FakePinState state; + return state; +} + +void ResetState() { State() = FakePinState(); } + +bool FakeRegister(void* addr, size_t size, std::string* error_message) { + (void)addr; + (void)size; + ++State().register_calls; + if (State().register_succeeds) return true; + if (error_message) *error_message = "fake register failure"; + return false; +} + +RegisteredPinnedMemoryManager::UnregisterResult FakeUnregister( + void* addr, std::string* error_message) { + (void)addr; + ++State().unregister_calls; + if (State().unregister_result == + RegisteredPinnedMemoryManager::UnregisterResult::kError && + error_message) { + *error_message = "fake unregister failure"; + } + return State().unregister_result; +} + +RegisteredPinnedMemoryManager::PinOps FakeOps() { + return {FakeRegister, FakeUnregister}; +} + +TEST(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { + ResetState(); + RegisteredPinnedMemoryManager manager({true, 64}, FakeOps()); + std::array buffer{}; + + auto first = manager.try_pin(buffer.data(), 64, "first"); + ASSERT_NE(first, nullptr); + + auto over_quota = manager.try_pin(buffer.data() + 64, 1, "over quota"); + EXPECT_EQ(over_quota, nullptr); + EXPECT_EQ(State().register_calls, 1); + + first.reset(); + EXPECT_EQ(State().unregister_calls, 1); + + auto second = manager.try_pin(buffer.data() + 64, 64, "second"); + ASSERT_NE(second, nullptr); + EXPECT_EQ(State().register_calls, 2); + + second.reset(); + EXPECT_EQ(State().unregister_calls, 2); +} + +TEST(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { + ResetState(); + RegisteredPinnedMemoryManager manager({true, 128}, FakeOps()); + std::array buffer{}; + + auto first = manager.try_pin(buffer.data() + 16, 32, "first"); + ASSERT_NE(first, nullptr); + + auto duplicate = manager.try_pin(buffer.data() + 16, 32, "duplicate"); + EXPECT_EQ(duplicate, nullptr); + + auto overlap = manager.try_pin(buffer.data() + 32, 16, "overlap"); + EXPECT_EQ(overlap, nullptr); + + auto adjacent = manager.try_pin(buffer.data() + 48, 16, "adjacent"); + ASSERT_NE(adjacent, nullptr); + + EXPECT_EQ(State().register_calls, 2); + adjacent.reset(); + first.reset(); + EXPECT_EQ(State().unregister_calls, 2); +} + +TEST(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { + ResetState(); + RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); + std::array buffer{}; + + State().register_succeeds = false; + auto failed = manager.try_pin(buffer.data(), 32, "failed"); + EXPECT_EQ(failed, nullptr); + EXPECT_EQ(State().register_calls, 1); + EXPECT_EQ(State().unregister_calls, 0); + + State().register_succeeds = true; + auto retried = manager.try_pin(buffer.data(), 32, "retried"); + ASSERT_NE(retried, nullptr); + EXPECT_EQ(State().register_calls, 2); + + retried.reset(); + EXPECT_EQ(State().unregister_calls, 1); +} + +} // namespace +} // namespace mooncake From 5133f212cdcae5066a00e397e7061c22de0704a8 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 17 Jul 2026 11:02:17 +0800 Subject: [PATCH 5/9] Store: make pinned unregister cleanup best-effort --- .../src/registered_pinned_memory.cpp | 12 ++++++---- .../tests/registered_pinned_memory_test.cpp | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 6b2f7371dd..0a63fecb08 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -247,10 +247,11 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( auto unregister_result = pin_ops_.unregister_region(addr, &error_message); if (unregister_result != UnregisterResult::kSuccess) { - LOG(FATAL) << "cudaHostUnregister failed after active range " + LOG(ERROR) << "cudaHostUnregister failed after active range " "tracking mismatch for " << owner << ", size=" << size - << ", error=" << error_message; + << ", error=" << error_message + << ". Continue with best-effort cleanup."; } std::lock_guard lock(mutex_); drop_reservation_locked(key, size); @@ -279,15 +280,18 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { if (!pin_ops_.unregister_region) return; auto unregister_result = pin_ops_.unregister_region(region->addr_, &error_message); + // Treat CUDA unregistration as best-effort cleanup: drop manager state so + // stale raw tracking pointers do not outlive the Store segment owner. if (unregister_result != UnregisterResult::kSuccess) { if (unregister_result == UnregisterResult::kRuntimeUnloading) { LOG(WARNING) << "Skip cudaHostUnregister for " << region->owner_ << " because CUDA runtime is unloading, size=" << region->size_; } else { - LOG(FATAL) << "cudaHostUnregister failed for " << region->owner_ + LOG(ERROR) << "cudaHostUnregister failed for " << region->owner_ << ", size=" << region->size_ - << ", error=" << error_message; + << ", error=" << error_message + << ". Continue with best-effort cleanup."; } } diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp index be5cbea46f..c2a7732917 100644 --- a/mooncake-store/tests/registered_pinned_memory_test.cpp +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -95,6 +95,30 @@ TEST(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { EXPECT_EQ(State().unregister_calls, 2); } +TEST(RegisteredPinnedMemoryManagerTest, + UnregisterFailureDropsTrackingAndRefunds) { + ResetState(); + RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); + std::array buffer{}; + + auto first = manager.try_pin(buffer.data(), 32, "first"); + ASSERT_NE(first, nullptr); + + State().unregister_result = + RegisteredPinnedMemoryManager::UnregisterResult::kError; + first.reset(); + EXPECT_EQ(State().unregister_calls, 1); + + State().unregister_result = + RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; + auto second = manager.try_pin(buffer.data(), 32, "second"); + ASSERT_NE(second, nullptr); + EXPECT_EQ(State().register_calls, 2); + + second.reset(); + EXPECT_EQ(State().unregister_calls, 2); +} + TEST(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { ResetState(); RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); From a3c0ec5168689511de8454ef7284453faccb9a40 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 17 Jul 2026 13:28:11 +0800 Subject: [PATCH 6/9] Store: simplify pinned region tracking --- .../src/registered_pinned_memory.cpp | 65 +++++++++---------- mooncake-store/src/registered_pinned_memory.h | 23 ++----- 2 files changed, 35 insertions(+), 53 deletions(-) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 0a63fecb08..1185d88513 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -154,7 +154,6 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( return nullptr; } - RegionKey key{addr, size}; const auto start = reinterpret_cast(addr); const auto end = start + size; if (end < start) { @@ -175,13 +174,12 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( { std::lock_guard lock(mutex_); for (const auto& entry : regions_) { - const auto region_start = - reinterpret_cast(entry.first.addr); - const auto region_end = region_start + entry.first.size; + const auto region_start = reinterpret_cast(entry.addr); + const auto region_end = region_start + entry.size; if (region_end < region_start) { - LOG(WARNING) << "Skip cudaHostRegister for " << owner - << ": existing active range overflow, size=" - << entry.first.size; + LOG(WARNING) + << "Skip cudaHostRegister for " << owner + << ": existing active range overflow, size=" << entry.size; return nullptr; } @@ -203,12 +201,7 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( } try { - auto [_, inserted] = regions_.emplace(key, nullptr); - if (!inserted) { - LOG(WARNING) << "Skip cudaHostRegister for " << owner - << ": active region already exists, size=" << size; - return nullptr; - } + regions_.push_back({addr, size, nullptr}); } catch (...) { LOG(WARNING) << "Skip cudaHostRegister for " << owner << ": failed to allocate pin tracking, size=" << size; @@ -223,7 +216,7 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( if (!registered) { { std::lock_guard lock(mutex_); - drop_reservation_locked(key, size); + remove_inactive_region_locked(addr, size); } LOG(WARNING) << "cudaHostRegister failed for " << owner << ", size=" << size << ", error=" << error_message @@ -235,11 +228,14 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( uint64_t pinned_bytes = 0; { std::lock_guard lock(mutex_); - auto it = regions_.find(key); - if (it != regions_.end() && it->second == nullptr) { - it->second = region.get(); + for (auto& entry : regions_) { + if (entry.addr != addr || entry.size != size || entry.region) { + continue; + } + entry.region = region.get(); pinned_bytes = pinned_bytes_; tracking_ready = true; + break; } } if (!tracking_ready) { @@ -254,7 +250,7 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( << ". Continue with best-effort cleanup."; } std::lock_guard lock(mutex_); - drop_reservation_locked(key, size); + remove_inactive_region_locked(addr, size); return nullptr; } @@ -266,15 +262,20 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { if (!region || !region->addr_ || region->size_ == 0) return; - RegionKey key{region->addr_, region->size_}; + bool should_unregister = false; { std::lock_guard lock(mutex_); - auto region_it = regions_.find(key); - if (region_it == regions_.end() || region_it->second != region) { - return; + for (auto& entry : regions_) { + if (entry.addr != region->addr_ || entry.size != region->size_ || + entry.region != region) { + continue; + } + entry.region = nullptr; + should_unregister = true; + break; } - region_it->second = nullptr; } + if (!should_unregister) return; std::string error_message; if (!pin_ops_.unregister_region) return; @@ -296,22 +297,16 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { } std::lock_guard lock(mutex_); - auto it = regions_.find(key); - if (it == regions_.end() || it->second != nullptr) return; - regions_.erase(it); - if (pinned_bytes_ >= region->size_) { - pinned_bytes_ -= region->size_; - } else { - pinned_bytes_ = 0; - } + remove_inactive_region_locked(region->addr_, region->size_); } -void RegisteredPinnedMemoryManager::drop_reservation_locked( - const RegionKey& key, size_t size) { - auto it = regions_.find(key); - if (it != regions_.end() && it->second == nullptr) { +void RegisteredPinnedMemoryManager::remove_inactive_region_locked(void* addr, + size_t size) { + for (auto it = regions_.begin(); it != regions_.end(); ++it) { + if (it->addr != addr || it->size != size || it->region) continue; regions_.erase(it); pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; + return; } } diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index ac82b79e48..7b31a39780 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -5,8 +5,8 @@ #include #include #include -#include #include +#include namespace mooncake { @@ -53,22 +53,10 @@ class RegisteredPinnedMemoryManager { private: friend class RegisteredPinnedRegion; - struct RegionKey { + struct ActiveRegion { void* addr = nullptr; size_t size = 0; - - bool operator==(const RegionKey& other) const { - return addr == other.addr && size == other.size; - } - }; - - struct RegionKeyHash { - size_t operator()(const RegionKey& key) const { - const size_t addr_hash = std::hash{}(key.addr); - return addr_hash ^ - (std::hash{}(key.size) + 0x9e3779b97f4a7c15ULL + - (addr_hash << 6) + (addr_hash >> 2)); - } + RegisteredPinnedRegion* region = nullptr; }; RegisteredPinnedMemoryManager(); @@ -83,7 +71,7 @@ class RegisteredPinnedMemoryManager { #endif void release(RegisteredPinnedRegion* region); - void drop_reservation_locked(const RegionKey& key, size_t size); + void remove_inactive_region_locked(void* addr, size_t size); const bool enabled_; const uint64_t limit_bytes_; @@ -91,8 +79,7 @@ class RegisteredPinnedMemoryManager { mutable std::mutex mutex_; uint64_t pinned_bytes_ = 0; - std::unordered_map - regions_; + std::vector regions_; }; } // namespace mooncake From 3aad53dfb49a281f6f47c68b5e6e04b99c7b836d Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 17 Jul 2026 13:46:29 +0800 Subject: [PATCH 7/9] Store: slim pinned memory manager changes --- .../python-api-reference/mooncake-store.md | 26 ++++------ .../src/registered_pinned_memory.cpp | 48 ++++++------------- mooncake-store/src/registered_pinned_memory.h | 6 +-- .../tests/registered_pinned_memory_test.cpp | 48 +++++++++---------- 4 files changed, 50 insertions(+), 78 deletions(-) diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/python-api-reference/mooncake-store.md index e3ae8467bb..f7f4782816 100644 --- a/docs/source/python-api-reference/mooncake-store.md +++ b/docs/source/python-api-reference/mooncake-store.md @@ -1081,23 +1081,15 @@ def setup( - `ssd_offload_path` (str): SSD offload directory. When provided, overrides the storage path environment configuration. - `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`. -**Store segment pinned memory:** - -CUDA-enabled builds can register Store-managed host segments as pinned memory -when the process is given an explicit pinned-memory quota. This applies only to -Store segment memory for host transfer protocols allocated by `setup()` -(`global_segment_size`) and `allocateAndMountSegment()`. It does not pin -file-backed `mountSegment()` mappings, CXL/device segments, the -`local_buffer_size` client buffer, user buffers, dummy-client shared memory, or -temporary staging buffers. - -- `MC_STORE_PIN_MEMORY_MAX_BYTES`: positive process-wide quota, in bytes. Store - segment pinning is disabled when this is unset, empty, `0`, or invalid. -- `MC_STORE_PIN_MEMORY`: optional override. Set to `0`, `false`, `off`, or `no` - to disable Store segment pinning even when a quota is configured. - -If the quota is exhausted or CUDA registration fails, Mooncake continues with -pageable Store segment memory. +**Store segment pinned memory:** CUDA-enabled builds can register Store-managed +host segments as pinned memory when `MC_STORE_PIN_MEMORY_MAX_BYTES` is set to a +positive process-wide quota; unset, empty, `0`, or invalid values disable it. +The scope is limited to host Store segments allocated by `setup()` +(`global_segment_size`) and `allocateAndMountSegment()`; it excludes file-backed +`mountSegment()` mappings, CXL/device segments, `local_buffer_size`, user +buffers, dummy-client shared memory, and temporary staging buffers. If the quota +is exhausted or CUDA registration fails, Mooncake continues with pageable Store +segment memory. **Returns:** - `int`: Status code (0 = success, non-zero = error code) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 1185d88513..9b6a9bd4b5 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -1,6 +1,5 @@ #include "registered_pinned_memory.h" -#include #include #include #include @@ -35,19 +34,6 @@ std::pair TrimAsciiWhitespace(const char* value) { return {begin, end}; } -bool ParsePinnedMemoryEnabled() { - const char* value = std::getenv("MC_STORE_PIN_MEMORY"); - if (!value) return true; - - auto [begin, end] = TrimAsciiWhitespace(value); - std::string normalized(begin, end); - std::transform( - normalized.begin(), normalized.end(), normalized.begin(), - [](unsigned char ch) { return static_cast(std::tolower(ch)); }); - return !(normalized == "0" || normalized == "false" || - normalized == "off" || normalized == "no"); -} - std::optional ParsePinnedMemoryLimit() { const char* value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); if (!value || value[0] == '\0') return 0; @@ -70,8 +56,6 @@ std::optional ParsePinnedMemoryLimit() { } std::pair ParsePinnedMemoryConfig() { - if (!ParsePinnedMemoryEnabled()) return {false, 0}; - auto limit = ParsePinnedMemoryLimit(); if (!limit.has_value() || *limit == 0) { return {false, 0}; @@ -229,13 +213,12 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( { std::lock_guard lock(mutex_); for (auto& entry : regions_) { - if (entry.addr != addr || entry.size != size || entry.region) { - continue; + if (entry.addr == addr && entry.size == size && !entry.region) { + entry.region = region.get(); + pinned_bytes = pinned_bytes_; + tracking_ready = true; + break; } - entry.region = region.get(); - pinned_bytes = pinned_bytes_; - tracking_ready = true; - break; } } if (!tracking_ready) { @@ -266,19 +249,17 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { { std::lock_guard lock(mutex_); for (auto& entry : regions_) { - if (entry.addr != region->addr_ || entry.size != region->size_ || - entry.region != region) { - continue; + if (entry.addr == region->addr_ && entry.size == region->size_ && + entry.region == region) { + entry.region = nullptr; + should_unregister = true; + break; } - entry.region = nullptr; - should_unregister = true; - break; } } if (!should_unregister) return; std::string error_message; - if (!pin_ops_.unregister_region) return; auto unregister_result = pin_ops_.unregister_region(region->addr_, &error_message); // Treat CUDA unregistration as best-effort cleanup: drop manager state so @@ -303,10 +284,11 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { void RegisteredPinnedMemoryManager::remove_inactive_region_locked(void* addr, size_t size) { for (auto it = regions_.begin(); it != regions_.end(); ++it) { - if (it->addr != addr || it->size != size || it->region) continue; - regions_.erase(it); - pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; - return; + if (it->addr == addr && it->size == size && !it->region) { + regions_.erase(it); + pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; + return; + } } } diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index 7b31a39780..d66da1e5a5 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -54,9 +54,9 @@ class RegisteredPinnedMemoryManager { friend class RegisteredPinnedRegion; struct ActiveRegion { - void* addr = nullptr; - size_t size = 0; - RegisteredPinnedRegion* region = nullptr; + void* addr; + size_t size; + RegisteredPinnedRegion* region; }; RegisteredPinnedMemoryManager(); diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp index c2a7732917..5a3180b83b 100644 --- a/mooncake-store/tests/registered_pinned_memory_test.cpp +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -2,7 +2,6 @@ #include "../src/registered_pinned_memory.h" #include -#include #include @@ -49,22 +48,27 @@ RegisteredPinnedMemoryManager::PinOps FakeOps() { return {FakeRegister, FakeUnregister}; } -TEST(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { - ResetState(); +class RegisteredPinnedMemoryManagerTest : public ::testing::Test { + protected: + void SetUp() override { ResetState(); } + + std::array buffer_{}; +}; + +TEST_F(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { RegisteredPinnedMemoryManager manager({true, 64}, FakeOps()); - std::array buffer{}; - auto first = manager.try_pin(buffer.data(), 64, "first"); + auto first = manager.try_pin(buffer_.data(), 64, "first"); ASSERT_NE(first, nullptr); - auto over_quota = manager.try_pin(buffer.data() + 64, 1, "over quota"); + auto over_quota = manager.try_pin(buffer_.data() + 64, 1, "over quota"); EXPECT_EQ(over_quota, nullptr); EXPECT_EQ(State().register_calls, 1); first.reset(); EXPECT_EQ(State().unregister_calls, 1); - auto second = manager.try_pin(buffer.data() + 64, 64, "second"); + auto second = manager.try_pin(buffer_.data() + 64, 64, "second"); ASSERT_NE(second, nullptr); EXPECT_EQ(State().register_calls, 2); @@ -72,21 +76,19 @@ TEST(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { EXPECT_EQ(State().unregister_calls, 2); } -TEST(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { - ResetState(); +TEST_F(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { RegisteredPinnedMemoryManager manager({true, 128}, FakeOps()); - std::array buffer{}; - auto first = manager.try_pin(buffer.data() + 16, 32, "first"); + auto first = manager.try_pin(buffer_.data() + 16, 32, "first"); ASSERT_NE(first, nullptr); - auto duplicate = manager.try_pin(buffer.data() + 16, 32, "duplicate"); + auto duplicate = manager.try_pin(buffer_.data() + 16, 32, "duplicate"); EXPECT_EQ(duplicate, nullptr); - auto overlap = manager.try_pin(buffer.data() + 32, 16, "overlap"); + auto overlap = manager.try_pin(buffer_.data() + 32, 16, "overlap"); EXPECT_EQ(overlap, nullptr); - auto adjacent = manager.try_pin(buffer.data() + 48, 16, "adjacent"); + auto adjacent = manager.try_pin(buffer_.data() + 48, 16, "adjacent"); ASSERT_NE(adjacent, nullptr); EXPECT_EQ(State().register_calls, 2); @@ -95,13 +97,11 @@ TEST(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { EXPECT_EQ(State().unregister_calls, 2); } -TEST(RegisteredPinnedMemoryManagerTest, - UnregisterFailureDropsTrackingAndRefunds) { - ResetState(); +TEST_F(RegisteredPinnedMemoryManagerTest, + UnregisterFailureDropsTrackingAndRefunds) { RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); - std::array buffer{}; - auto first = manager.try_pin(buffer.data(), 32, "first"); + auto first = manager.try_pin(buffer_.data(), 32, "first"); ASSERT_NE(first, nullptr); State().unregister_result = @@ -111,7 +111,7 @@ TEST(RegisteredPinnedMemoryManagerTest, State().unregister_result = RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; - auto second = manager.try_pin(buffer.data(), 32, "second"); + auto second = manager.try_pin(buffer_.data(), 32, "second"); ASSERT_NE(second, nullptr); EXPECT_EQ(State().register_calls, 2); @@ -119,19 +119,17 @@ TEST(RegisteredPinnedMemoryManagerTest, EXPECT_EQ(State().unregister_calls, 2); } -TEST(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { - ResetState(); +TEST_F(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); - std::array buffer{}; State().register_succeeds = false; - auto failed = manager.try_pin(buffer.data(), 32, "failed"); + auto failed = manager.try_pin(buffer_.data(), 32, "failed"); EXPECT_EQ(failed, nullptr); EXPECT_EQ(State().register_calls, 1); EXPECT_EQ(State().unregister_calls, 0); State().register_succeeds = true; - auto retried = manager.try_pin(buffer.data(), 32, "retried"); + auto retried = manager.try_pin(buffer_.data(), 32, "retried"); ASSERT_NE(retried, nullptr); EXPECT_EQ(State().register_calls, 2); From a1219df6bcf6afb67935b16c215bc6a118506c79 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Tue, 21 Jul 2026 10:37:02 +0800 Subject: [PATCH 8/9] Store: simplify pinned memory manager code --- .../src/registered_pinned_memory.cpp | 116 +++++------------- mooncake-store/src/registered_pinned_memory.h | 9 +- .../tests/registered_pinned_memory_test.cpp | 112 ++++++++--------- 3 files changed, 81 insertions(+), 156 deletions(-) diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index 9b6a9bd4b5..f7897de111 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -1,14 +1,13 @@ #include "registered_pinned_memory.h" -#include -#include #include -#include #include -#include +#include #include +#include "utils/type_util.h" + #if defined(USE_CUDA) #include #endif @@ -16,51 +15,31 @@ namespace mooncake { namespace { -std::pair TrimAsciiWhitespace(const char* value) { - const char* begin = value; - while (std::isspace(static_cast(*begin))) { - ++begin; - } - - const char* end = begin; - while (*end != '\0') { - ++end; - } - while (end > begin && - std::isspace(static_cast(*(end - 1)))) { - --end; - } - - return {begin, end}; +std::string_view TrimAsciiWhitespace(std::string_view value) { + constexpr std::string_view whitespace = " \t\r\n\f\v"; + size_t begin = value.find_first_not_of(whitespace); + if (begin == std::string_view::npos) return {}; + size_t end = value.find_last_not_of(whitespace); + return value.substr(begin, end - begin + 1); } -std::optional ParsePinnedMemoryLimit() { - const char* value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); - if (!value || value[0] == '\0') return 0; - - auto [number, end] = TrimAsciiWhitespace(value); - if (*number == '-') { - LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value - << "', disabling Store segment pinning"; - return std::nullopt; - } +std::pair ParsePinnedMemoryConfig() { + const char* raw_value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); + if (!raw_value || raw_value[0] == '\0') return {false, 0}; - uint64_t parsed = 0; - auto [ptr, ec] = std::from_chars(number, end, parsed); - if (ec != std::errc{} || ptr != end) { - LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << value + uint64_t limit = 0; + auto value = TrimAsciiWhitespace(raw_value); + if (value.empty() || !TypeUtil::ParseUint64(value, limit)) { + LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << raw_value << "', disabling Store segment pinning"; - return std::nullopt; + return {false, 0}; } - return parsed; + return {limit != 0, limit}; } -std::pair ParsePinnedMemoryConfig() { - auto limit = ParsePinnedMemoryLimit(); - if (!limit.has_value() || *limit == 0) { - return {false, 0}; - } - return {true, *limit}; +void LogPinSkip(const std::string& owner, const char* reason, size_t size) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner << ": " << reason + << ", size=" << size; } #if defined(USE_CUDA) @@ -110,11 +89,8 @@ RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { } RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() - : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig()) {} - -RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( - std::pair config) - : RegisteredPinnedMemoryManager(config, DefaultPinOps()) {} + : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig(), + DefaultPinOps()) {} RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( std::pair config, PinOps pin_ops) @@ -141,17 +117,15 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( const auto start = reinterpret_cast(addr); const auto end = start + size; if (end < start) { - LOG(WARNING) << "Skip cudaHostRegister for " << owner - << ": address range overflow, size=" << size; + LogPinSkip(owner, "address range overflow", size); return nullptr; } std::shared_ptr region; try { - region.reset(new RegisteredPinnedRegion(this, addr, size, owner)); + region.reset(new RegisteredPinnedRegion(this, addr, size)); } catch (...) { - LOG(WARNING) << "Skip cudaHostRegister for " << owner - << ": failed to allocate pin tracking, size=" << size; + LogPinSkip(owner, "failed to allocate pin tracking", size); return nullptr; } @@ -160,18 +134,9 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( for (const auto& entry : regions_) { const auto region_start = reinterpret_cast(entry.addr); const auto region_end = region_start + entry.size; - if (region_end < region_start) { - LOG(WARNING) - << "Skip cudaHostRegister for " << owner - << ": existing active range overflow, size=" << entry.size; - return nullptr; - } - const bool overlaps = start < region_end && end > region_start; if (overlaps) { - LOG(WARNING) - << "Skip cudaHostRegister for " << owner - << ": overlaps an active pinned region, size=" << size; + LogPinSkip(owner, "overlaps an active pinned region", size); return nullptr; } } @@ -187,8 +152,7 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( try { regions_.push_back({addr, size, nullptr}); } catch (...) { - LOG(WARNING) << "Skip cudaHostRegister for " << owner - << ": failed to allocate pin tracking, size=" << size; + LogPinSkip(owner, "failed to allocate pin tracking", size); return nullptr; } pinned_bytes_ += size; @@ -208,7 +172,6 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( return nullptr; } - bool tracking_ready = false; uint64_t pinned_bytes = 0; { std::lock_guard lock(mutex_); @@ -216,26 +179,10 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( if (entry.addr == addr && entry.size == size && !entry.region) { entry.region = region.get(); pinned_bytes = pinned_bytes_; - tracking_ready = true; break; } } } - if (!tracking_ready) { - error_message.clear(); - auto unregister_result = - pin_ops_.unregister_region(addr, &error_message); - if (unregister_result != UnregisterResult::kSuccess) { - LOG(ERROR) << "cudaHostUnregister failed after active range " - "tracking mismatch for " - << owner << ", size=" << size - << ", error=" << error_message - << ". Continue with best-effort cleanup."; - } - std::lock_guard lock(mutex_); - remove_inactive_region_locked(addr, size); - return nullptr; - } LOG(INFO) << "cudaHostRegister succeeded for " << owner << ", size=" << size << ", pinned=" << pinned_bytes << ", limit=" << limit_bytes_; @@ -266,12 +213,11 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { // stale raw tracking pointers do not outlive the Store segment owner. if (unregister_result != UnregisterResult::kSuccess) { if (unregister_result == UnregisterResult::kRuntimeUnloading) { - LOG(WARNING) << "Skip cudaHostUnregister for " << region->owner_ - << " because CUDA runtime is unloading, size=" + LOG(WARNING) << "Skip cudaHostUnregister because CUDA runtime " + "is unloading, size=" << region->size_; } else { - LOG(ERROR) << "cudaHostUnregister failed for " << region->owner_ - << ", size=" << region->size_ + LOG(ERROR) << "cudaHostUnregister failed, size=" << region->size_ << ", error=" << error_message << ". Continue with best-effort cleanup."; } diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index d66da1e5a5..d63ea524ee 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -22,16 +22,12 @@ class RegisteredPinnedRegion { friend class RegisteredPinnedMemoryManager; RegisteredPinnedRegion(RegisteredPinnedMemoryManager* manager, void* addr, - size_t size, std::string owner) - : manager_(manager), - addr_(addr), - size_(size), - owner_(std::move(owner)) {} + size_t size) + : manager_(manager), addr_(addr), size_(size) {} RegisteredPinnedMemoryManager* manager_ = nullptr; void* addr_ = nullptr; size_t size_ = 0; - std::string owner_; }; class RegisteredPinnedMemoryManager { @@ -60,7 +56,6 @@ class RegisteredPinnedMemoryManager { }; RegisteredPinnedMemoryManager(); - explicit RegisteredPinnedMemoryManager(std::pair config); #if defined(MOONCAKE_STORE_TEST) public: #endif diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp index 5a3180b83b..0313d011ce 100644 --- a/mooncake-store/tests/registered_pinned_memory_test.cpp +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -8,10 +8,12 @@ namespace mooncake { namespace { +using Manager = RegisteredPinnedMemoryManager; +using UnregisterResult = Manager::UnregisterResult; + struct FakePinState { bool register_succeeds = true; - RegisteredPinnedMemoryManager::UnregisterResult unregister_result = - RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; + UnregisterResult unregister_result = UnregisterResult::kSuccess; int register_calls = 0; int unregister_calls = 0; }; @@ -21,120 +23,102 @@ FakePinState& State() { return state; } -void ResetState() { State() = FakePinState(); } - -bool FakeRegister(void* addr, size_t size, std::string* error_message) { - (void)addr; - (void)size; +bool FakeRegister(void*, size_t, std::string* error_message) { ++State().register_calls; if (State().register_succeeds) return true; if (error_message) *error_message = "fake register failure"; return false; } -RegisteredPinnedMemoryManager::UnregisterResult FakeUnregister( - void* addr, std::string* error_message) { - (void)addr; +UnregisterResult FakeUnregister(void*, std::string* error_message) { ++State().unregister_calls; - if (State().unregister_result == - RegisteredPinnedMemoryManager::UnregisterResult::kError && + if (State().unregister_result == UnregisterResult::kError && error_message) { *error_message = "fake unregister failure"; } return State().unregister_result; } -RegisteredPinnedMemoryManager::PinOps FakeOps() { - return {FakeRegister, FakeUnregister}; -} - class RegisteredPinnedMemoryManagerTest : public ::testing::Test { protected: - void SetUp() override { ResetState(); } + void SetUp() override { State() = FakePinState(); } + + Manager MakeManager(size_t limit) { + return Manager({true, limit}, {FakeRegister, FakeUnregister}); + } + + std::shared_ptr Pin(Manager& manager, size_t offset, + size_t size) { + return manager.try_pin(buffer_.data() + offset, size, "segment"); + } + + void ExpectCalls(int register_calls, int unregister_calls) { + EXPECT_EQ(State().register_calls, register_calls); + EXPECT_EQ(State().unregister_calls, unregister_calls); + } std::array buffer_{}; }; TEST_F(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { - RegisteredPinnedMemoryManager manager({true, 64}, FakeOps()); + auto manager = MakeManager(64); - auto first = manager.try_pin(buffer_.data(), 64, "first"); + auto first = Pin(manager, 0, 64); ASSERT_NE(first, nullptr); - - auto over_quota = manager.try_pin(buffer_.data() + 64, 1, "over quota"); - EXPECT_EQ(over_quota, nullptr); - EXPECT_EQ(State().register_calls, 1); + EXPECT_EQ(Pin(manager, 64, 1), nullptr); + ExpectCalls(1, 0); first.reset(); - EXPECT_EQ(State().unregister_calls, 1); + ExpectCalls(1, 1); - auto second = manager.try_pin(buffer_.data() + 64, 64, "second"); + auto second = Pin(manager, 64, 64); ASSERT_NE(second, nullptr); - EXPECT_EQ(State().register_calls, 2); + ExpectCalls(2, 1); second.reset(); - EXPECT_EQ(State().unregister_calls, 2); + ExpectCalls(2, 2); } TEST_F(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { - RegisteredPinnedMemoryManager manager({true, 128}, FakeOps()); + auto manager = MakeManager(128); - auto first = manager.try_pin(buffer_.data() + 16, 32, "first"); + auto first = Pin(manager, 16, 32); ASSERT_NE(first, nullptr); - auto duplicate = manager.try_pin(buffer_.data() + 16, 32, "duplicate"); - EXPECT_EQ(duplicate, nullptr); - - auto overlap = manager.try_pin(buffer_.data() + 32, 16, "overlap"); - EXPECT_EQ(overlap, nullptr); + EXPECT_EQ(Pin(manager, 16, 32), nullptr); + EXPECT_EQ(Pin(manager, 32, 16), nullptr); - auto adjacent = manager.try_pin(buffer_.data() + 48, 16, "adjacent"); + auto adjacent = Pin(manager, 48, 16); ASSERT_NE(adjacent, nullptr); + ExpectCalls(2, 0); - EXPECT_EQ(State().register_calls, 2); adjacent.reset(); first.reset(); - EXPECT_EQ(State().unregister_calls, 2); + ExpectCalls(2, 2); } -TEST_F(RegisteredPinnedMemoryManagerTest, - UnregisterFailureDropsTrackingAndRefunds) { - RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); +TEST_F(RegisteredPinnedMemoryManagerTest, FailurePathsRefundReservations) { + auto manager = MakeManager(32); - auto first = manager.try_pin(buffer_.data(), 32, "first"); + auto first = Pin(manager, 0, 32); ASSERT_NE(first, nullptr); - State().unregister_result = - RegisteredPinnedMemoryManager::UnregisterResult::kError; + State().unregister_result = UnregisterResult::kError; first.reset(); - EXPECT_EQ(State().unregister_calls, 1); - - State().unregister_result = - RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; - auto second = manager.try_pin(buffer_.data(), 32, "second"); - ASSERT_NE(second, nullptr); - EXPECT_EQ(State().register_calls, 2); - - second.reset(); - EXPECT_EQ(State().unregister_calls, 2); -} - -TEST_F(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { - RegisteredPinnedMemoryManager manager({true, 32}, FakeOps()); + ExpectCalls(1, 1); + State().unregister_result = UnregisterResult::kSuccess; State().register_succeeds = false; - auto failed = manager.try_pin(buffer_.data(), 32, "failed"); - EXPECT_EQ(failed, nullptr); - EXPECT_EQ(State().register_calls, 1); - EXPECT_EQ(State().unregister_calls, 0); + EXPECT_EQ(Pin(manager, 0, 32), nullptr); + ExpectCalls(2, 1); State().register_succeeds = true; - auto retried = manager.try_pin(buffer_.data(), 32, "retried"); + auto retried = Pin(manager, 0, 32); ASSERT_NE(retried, nullptr); - EXPECT_EQ(State().register_calls, 2); + ExpectCalls(3, 1); retried.reset(); - EXPECT_EQ(State().unregister_calls, 1); + ExpectCalls(3, 2); } } // namespace From 23971d62d20802a4854418ba8fcc3408333ce6d9 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Tue, 21 Jul 2026 23:36:48 +0800 Subject: [PATCH 9/9] Store: retain pinned memory on unregister failure --- mooncake-store/include/real_client.h | 3 + mooncake-store/src/real_client.cpp | 71 ++++++++++++++----- .../src/registered_pinned_memory.cpp | 23 ++++-- mooncake-store/src/registered_pinned_memory.h | 5 +- .../tests/registered_pinned_memory_test.cpp | 33 +++++---- 5 files changed, 97 insertions(+), 38 deletions(-) diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 1492ff6779..c7502f3482 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -760,6 +760,8 @@ class RealClient : public PyClient { std::shared_ptr pinned_region; }; + void FreeAllocatedStoreSegment(AllocatedSegmentRecord &record); + std::unique_ptr port_binder_ = nullptr; struct SegmentDeleter { @@ -820,6 +822,7 @@ class RealClient : public PyClient { #endif std::vector> setup_segment_pinned_regions_; + bool setup_segment_memory_must_leak_ = false; std::string protocol; std::string device_name; std::string local_hostname; diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 7c2ba64a0c..be5cb138c9 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -67,6 +67,31 @@ std::shared_ptr TryPinStoreSegment( " protocol=" + protocol); } +bool ReleasePinnedRegionForFree( + std::shared_ptr &pinned_region, + const char *segment_owner) { + if (!pinned_region) return true; + const bool safe_to_free = pinned_region->release(); + pinned_region.reset(); + if (!safe_to_free) { + LOG(ERROR) << "Leaking " << segment_owner + << " backing memory because cudaHostUnregister failed"; + } + return safe_to_free; +} + +bool ReleasePinnedRegionsForFree( + std::vector> &pinned_regions, + const char *segment_owner) { + bool safe_to_free = true; + for (auto &pinned_region : pinned_regions) { + safe_to_free &= ReleasePinnedRegionForFree(pinned_region, + segment_owner); + } + pinned_regions.clear(); + return safe_to_free; +} + #ifdef USE_ASCEND_DIRECT bool checkAcl(aclError result, const char *message) { if (result != ACL_ERROR_NONE) { @@ -890,7 +915,10 @@ tl::expected RealClient::setup_internal( auto mount_result = client_->MountSegment(ptr, mapped_size, protocol, seg_location); if (!mount_result.has_value()) { - pinned_region.reset(); + if (!ReleasePinnedRegionForFree(pinned_region, + "Store setup segment")) { + setup_segment_memory_must_leak_ = true; + } LOG(ERROR) << "Failed to mount segment: " << toString(mount_result.error()); return tl::unexpected(mount_result.error()); @@ -1210,13 +1238,19 @@ tl::expected RealClient::tearDownAll_internal() { ReleaseAllAllocatedSegmentRecords(); client_buffer_allocator_.reset(); port_binder_.reset(); - setup_segment_pinned_regions_.clear(); + const bool setup_segments_safe_to_free = ReleasePinnedRegionsForFree( + setup_segment_pinned_regions_, "Store setup segment"); + if (!setup_segments_safe_to_free || setup_segment_memory_must_leak_) { + for (auto &ptr : hugepage_segment_ptrs_) ptr.release(); + for (auto &ptr : segment_ptrs_) ptr.release(); + setup_segment_memory_must_leak_ = false; + } hugepage_segment_ptrs_.clear(); + segment_ptrs_.clear(); ub_segment_ptrs_.clear(); #if defined(USE_SUNRISE) sunrise_segment_ptrs_.clear(); #endif - segment_ptrs_.clear(); local_hostname = ""; device_name = ""; protocol = ""; @@ -1391,6 +1425,14 @@ void RealClient::ReleaseAllMountedSegmentRecords() { } } +void RealClient::FreeAllocatedStoreSegment(AllocatedSegmentRecord &record) { + if (record.base && + ReleasePinnedRegionForFree(record.pinned_region, + "allocated Store segment")) { + free_memory(record.protocol, record.base); + } +} + void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { AllocatedSegmentRecord record; bool found = false; @@ -1404,8 +1446,7 @@ void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { } } if (found && record.base) { - record.pinned_region.reset(); - free_memory(record.protocol, record.base); + FreeAllocatedStoreSegment(record); } } @@ -1416,10 +1457,7 @@ void RealClient::ReleaseAllAllocatedSegmentRecords() { records.swap(allocated_segment_records_); } for (auto &entry : records) { - if (entry.second.base) { - entry.second.pinned_region.reset(); - free_memory(entry.second.protocol, entry.second.base); - } + FreeAllocatedStoreSegment(entry.second); } } @@ -1565,8 +1603,10 @@ int RealClient::allocateAndMountSegment( client_->MountSegmentAndGetId(ptr, chunk_size, protocol, location); if (!result.has_value()) { LOG(ERROR) << "MountSegmentAndGetId failed"; - pinned_region.reset(); - free_memory(protocol, ptr); + if (ReleasePinnedRegionForFree(pinned_region, + "allocated Store segment")) { + free_memory(protocol, ptr); + } break; } @@ -1585,9 +1625,7 @@ int RealClient::allocateAndMountSegment( client_->UnmountSegmentById(id); } if (allocated_records[i].base) { - allocated_records[i].pinned_region.reset(); - free_memory(allocated_records[i].protocol, - allocated_records[i].base); + FreeAllocatedStoreSegment(allocated_records[i]); } } out_segment_ids.clear(); @@ -1673,10 +1711,7 @@ int RealClient::unmountAndFreeSegment( } for (auto &p : to_cleanup) { - if (p.second.base) { - p.second.pinned_region.reset(); - free_memory(p.second.protocol, p.second.base); - } + FreeAllocatedStoreSegment(p.second); } return first_error; diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp index f7897de111..bacdac8beb 100644 --- a/mooncake-store/src/registered_pinned_memory.cpp +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -79,7 +79,15 @@ RegisteredPinnedMemoryManager::PinOps DefaultPinOps() { } // namespace RegisteredPinnedRegion::~RegisteredPinnedRegion() { - if (manager_) manager_->release(this); + release(); +} + +bool RegisteredPinnedRegion::release() { + if (!manager_) return release_succeeded_; + auto* manager = manager_; + manager_ = nullptr; + release_succeeded_ = manager->release(this); + return release_succeeded_; } RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { @@ -189,8 +197,8 @@ std::shared_ptr RegisteredPinnedMemoryManager::try_pin( return region; } -void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { - if (!region || !region->addr_ || region->size_ == 0) return; +bool RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { + if (!region || !region->addr_ || region->size_ == 0) return true; bool should_unregister = false; { @@ -204,13 +212,11 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { } } } - if (!should_unregister) return; + if (!should_unregister) return true; std::string error_message; auto unregister_result = pin_ops_.unregister_region(region->addr_, &error_message); - // Treat CUDA unregistration as best-effort cleanup: drop manager state so - // stale raw tracking pointers do not outlive the Store segment owner. if (unregister_result != UnregisterResult::kSuccess) { if (unregister_result == UnregisterResult::kRuntimeUnloading) { LOG(WARNING) << "Skip cudaHostUnregister because CUDA runtime " @@ -219,12 +225,15 @@ void RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { } else { LOG(ERROR) << "cudaHostUnregister failed, size=" << region->size_ << ", error=" << error_message - << ". Continue with best-effort cleanup."; + << ". Keep the range reserved; backing memory must not " + "be freed."; + return false; } } std::lock_guard lock(mutex_); remove_inactive_region_locked(region->addr_, region->size_); + return true; } void RegisteredPinnedMemoryManager::remove_inactive_region_locked(void* addr, diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h index d63ea524ee..cbc43c1388 100644 --- a/mooncake-store/src/registered_pinned_memory.h +++ b/mooncake-store/src/registered_pinned_memory.h @@ -18,6 +18,8 @@ class RegisteredPinnedRegion { RegisteredPinnedRegion& operator=(const RegisteredPinnedRegion&) = delete; ~RegisteredPinnedRegion(); + bool release(); + private: friend class RegisteredPinnedMemoryManager; @@ -28,6 +30,7 @@ class RegisteredPinnedRegion { RegisteredPinnedMemoryManager* manager_ = nullptr; void* addr_ = nullptr; size_t size_ = 0; + bool release_succeeded_ = true; }; class RegisteredPinnedMemoryManager { @@ -65,7 +68,7 @@ class RegisteredPinnedMemoryManager { private: #endif - void release(RegisteredPinnedRegion* region); + bool release(RegisteredPinnedRegion* region); void remove_inactive_region_locked(void* addr, size_t size); const bool enabled_; diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp index 0313d011ce..a1e4c32be4 100644 --- a/mooncake-store/tests/registered_pinned_memory_test.cpp +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -97,28 +97,37 @@ TEST_F(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { ExpectCalls(2, 2); } -TEST_F(RegisteredPinnedMemoryManagerTest, FailurePathsRefundReservations) { +TEST_F(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { auto manager = MakeManager(32); - auto first = Pin(manager, 0, 32); - ASSERT_NE(first, nullptr); - - State().unregister_result = UnregisterResult::kError; - first.reset(); - ExpectCalls(1, 1); - - State().unregister_result = UnregisterResult::kSuccess; State().register_succeeds = false; EXPECT_EQ(Pin(manager, 0, 32), nullptr); - ExpectCalls(2, 1); + ExpectCalls(1, 0); State().register_succeeds = true; auto retried = Pin(manager, 0, 32); ASSERT_NE(retried, nullptr); - ExpectCalls(3, 1); + ExpectCalls(2, 0); retried.reset(); - ExpectCalls(3, 2); + ExpectCalls(2, 1); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, + UnregisterFailureRetainsReservation) { + auto manager = MakeManager(32); + + auto first = Pin(manager, 0, 32); + ASSERT_NE(first, nullptr); + + State().unregister_result = UnregisterResult::kError; + EXPECT_FALSE(first->release()); + ExpectCalls(1, 1); + + State().unregister_result = UnregisterResult::kSuccess; + EXPECT_EQ(Pin(manager, 0, 32), nullptr); + EXPECT_EQ(Pin(manager, 32, 32), nullptr); + ExpectCalls(1, 1); } } // namespace