diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/python-api-reference/mooncake-store.md index c738f5ef18..f7f4782816 100644 --- a/docs/source/python-api-reference/mooncake-store.md +++ b/docs/source/python-api-reference/mooncake-store.md @@ -1081,6 +1081,16 @@ 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 `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/include/real_client.h b/mooncake-store/include/real_client.h index 6b3382b3b6..c7502f3482 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,8 +757,11 @@ class RealClient : public PyClient { void *base = nullptr; size_t size = 0; std::string protocol; + std::shared_ptr pinned_region; }; + void FreeAllocatedStoreSegment(AllocatedSegmentRecord &record); + std::unique_ptr port_binder_ = nullptr; struct SegmentDeleter { @@ -816,6 +820,9 @@ class RealClient : public PyClient { std::vector> sunrise_segment_ptrs_; #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/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..be5cb138c9 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,46 @@ 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); +} + +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) { @@ -869,13 +910,23 @@ 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()) { + 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()); } + 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,12 +1238,19 @@ tl::expected RealClient::tearDownAll_internal() { ReleaseAllAllocatedSegmentRecords(); client_buffer_allocator_.reset(); port_binder_.reset(); + 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 = ""; @@ -1367,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; @@ -1380,7 +1446,7 @@ void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { } } if (found && record.base) { - free_memory(record.protocol, record.base); + FreeAllocatedStoreSegment(record); } } @@ -1391,9 +1457,7 @@ void RealClient::ReleaseAllAllocatedSegmentRecords() { records.swap(allocated_segment_records_); } for (auto &entry : records) { - if (entry.second.base) { - free_memory(entry.second.protocol, entry.second.base); - } + FreeAllocatedStoreSegment(entry.second); } } @@ -1533,17 +1597,23 @@ 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"; - free_memory(protocol, ptr); + if (ReleasePinnedRegionForFree(pinned_region, + "allocated Store segment")) { + 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,8 +1625,7 @@ int RealClient::allocateAndMountSegment( client_->UnmountSegmentById(id); } if (allocated_records[i].base) { - free_memory(allocated_records[i].protocol, - allocated_records[i].base); + FreeAllocatedStoreSegment(allocated_records[i]); } } out_segment_ids.clear(); @@ -1642,9 +1711,7 @@ int RealClient::unmountAndFreeSegment( } for (auto &p : to_cleanup) { - if (p.second.base) { - 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 new file mode 100644 index 0000000000..bacdac8beb --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -0,0 +1,250 @@ +#include "registered_pinned_memory.h" + +#include +#include +#include + +#include + +#include "utils/type_util.h" + +#if defined(USE_CUDA) +#include +#endif + +namespace mooncake { +namespace { + +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::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 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 {false, 0}; + } + return {limit != 0, 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) +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() { + 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() { + static RegisteredPinnedMemoryManager* manager = + new RegisteredPinnedMemoryManager(); + return *manager; +} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() + : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig(), + 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") + << ", 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 (!pin_ops_.register_region || !pin_ops_.unregister_region) { + return nullptr; + } + + const auto start = reinterpret_cast(addr); + const auto end = start + size; + if (end < start) { + LogPinSkip(owner, "address range overflow", size); + return nullptr; + } + + std::shared_ptr region; + try { + region.reset(new RegisteredPinnedRegion(this, addr, size)); + } catch (...) { + LogPinSkip(owner, "failed to allocate pin tracking", size); + return nullptr; + } + + { + std::lock_guard lock(mutex_); + for (const auto& entry : regions_) { + const auto region_start = reinterpret_cast(entry.addr); + const auto region_end = region_start + entry.size; + const bool overlaps = start < region_end && end > region_start; + if (overlaps) { + LogPinSkip(owner, "overlaps an active pinned region", 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 { + regions_.push_back({addr, size, nullptr}); + } catch (...) { + LogPinSkip(owner, "failed to allocate pin tracking", size); + return nullptr; + } + pinned_bytes_ += size; + } + + std::string error_message; + const bool registered = + pin_ops_.register_region(addr, size, &error_message); + if (!registered) { + { + std::lock_guard lock(mutex_); + remove_inactive_region_locked(addr, size); + } + LOG(WARNING) << "cudaHostRegister failed for " << owner + << ", size=" << size << ", error=" << error_message + << ". Continue with pageable host memory."; + return nullptr; + } + + uint64_t pinned_bytes = 0; + { + std::lock_guard lock(mutex_); + for (auto& entry : regions_) { + if (entry.addr == addr && entry.size == size && !entry.region) { + entry.region = region.get(); + pinned_bytes = pinned_bytes_; + break; + } + } + } + + LOG(INFO) << "cudaHostRegister succeeded for " << owner << ", size=" << size + << ", pinned=" << pinned_bytes << ", limit=" << limit_bytes_; + return region; +} + +bool RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { + if (!region || !region->addr_ || region->size_ == 0) return true; + + bool should_unregister = false; + { + std::lock_guard lock(mutex_); + for (auto& entry : regions_) { + if (entry.addr == region->addr_ && entry.size == region->size_ && + entry.region == region) { + entry.region = nullptr; + should_unregister = true; + break; + } + } + } + if (!should_unregister) return true; + + std::string error_message; + 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 because CUDA runtime " + "is unloading, size=" + << region->size_; + } else { + LOG(ERROR) << "cudaHostUnregister failed, size=" << region->size_ + << ", error=" << error_message + << ". 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, + size_t size) { + for (auto it = regions_.begin(); it != regions_.end(); ++it) { + if (it->addr == addr && it->size == size && !it->region) { + regions_.erase(it); + pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; + return; + } + } +} + +} // 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..cbc43c1388 --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.h @@ -0,0 +1,83 @@ +#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(); + + bool release(); + + private: + friend class RegisteredPinnedMemoryManager; + + RegisteredPinnedRegion(RegisteredPinnedMemoryManager* manager, void* addr, + size_t size) + : manager_(manager), addr_(addr), size_(size) {} + + RegisteredPinnedMemoryManager* manager_ = nullptr; + void* addr_ = nullptr; + size_t size_ = 0; + bool release_succeeded_ = true; +}; + +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, + const std::string& owner); + + private: + friend class RegisteredPinnedRegion; + + struct ActiveRegion { + void* addr; + size_t size; + RegisteredPinnedRegion* region; + }; + + RegisteredPinnedMemoryManager(); +#if defined(MOONCAKE_STORE_TEST) + public: +#endif + RegisteredPinnedMemoryManager(std::pair config, + PinOps pin_ops); +#if defined(MOONCAKE_STORE_TEST) + private: +#endif + + bool release(RegisteredPinnedRegion* region); + void remove_inactive_region_locked(void* addr, size_t size); + + const bool enabled_; + const uint64_t limit_bytes_; + const PinOps pin_ops_; + + mutable std::mutex mutex_; + uint64_t pinned_bytes_ = 0; + std::vector regions_; +}; + +} // namespace mooncake 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..a1e4c32be4 --- /dev/null +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -0,0 +1,134 @@ +#define MOONCAKE_STORE_TEST +#include "../src/registered_pinned_memory.h" + +#include + +#include + +namespace mooncake { +namespace { + +using Manager = RegisteredPinnedMemoryManager; +using UnregisterResult = Manager::UnregisterResult; + +struct FakePinState { + bool register_succeeds = true; + UnregisterResult unregister_result = UnregisterResult::kSuccess; + int register_calls = 0; + int unregister_calls = 0; +}; + +FakePinState& State() { + static FakePinState state; + return state; +} + +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; +} + +UnregisterResult FakeUnregister(void*, std::string* error_message) { + ++State().unregister_calls; + if (State().unregister_result == UnregisterResult::kError && + error_message) { + *error_message = "fake unregister failure"; + } + return State().unregister_result; +} + +class RegisteredPinnedMemoryManagerTest : public ::testing::Test { + protected: + 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) { + auto manager = MakeManager(64); + + auto first = Pin(manager, 0, 64); + ASSERT_NE(first, nullptr); + EXPECT_EQ(Pin(manager, 64, 1), nullptr); + ExpectCalls(1, 0); + + first.reset(); + ExpectCalls(1, 1); + + auto second = Pin(manager, 64, 64); + ASSERT_NE(second, nullptr); + ExpectCalls(2, 1); + + second.reset(); + ExpectCalls(2, 2); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { + auto manager = MakeManager(128); + + auto first = Pin(manager, 16, 32); + ASSERT_NE(first, nullptr); + + EXPECT_EQ(Pin(manager, 16, 32), nullptr); + EXPECT_EQ(Pin(manager, 32, 16), nullptr); + + auto adjacent = Pin(manager, 48, 16); + ASSERT_NE(adjacent, nullptr); + ExpectCalls(2, 0); + + adjacent.reset(); + first.reset(); + ExpectCalls(2, 2); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { + auto manager = MakeManager(32); + + State().register_succeeds = false; + EXPECT_EQ(Pin(manager, 0, 32), nullptr); + ExpectCalls(1, 0); + + State().register_succeeds = true; + auto retried = Pin(manager, 0, 32); + ASSERT_NE(retried, nullptr); + ExpectCalls(2, 0); + + retried.reset(); + 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 +} // namespace mooncake