From 4ca41761fc3af23d0f4f572ab6d4cdea5d38cb90 Mon Sep 17 00:00:00 2001 From: rongchenghao Date: Wed, 2 Sep 2026 15:22:23 +0800 Subject: [PATCH 1/2] [Conductor] conductor: tiered prefix indexer Third slice of the Mooncake Conductor upstreaming series, on top of the prefix index types and hash strategy. Adds PrefixCacheTable, the in-memory index that answers cache-hit queries across GPU, CPU and disk tiers Co-authored-by: Misak2333 <167268798+Misak2333@users.noreply.github.com> Co-authored-by: Chase-R --- mooncake-conductor/CMakeLists.txt | 6 +- .../conductor/prefixindex/prefix_indexer.h | 135 +++ .../src/prefixindex/prefix_indexer.cpp | 615 +++++++++++ mooncake-conductor/tests/CMakeLists.txt | 6 +- .../tests/prefix_indexer_test.cpp | 952 ++++++++++++++++++ .../tests/prefix_indexer_test_peer.h | 115 +++ 6 files changed, 1825 insertions(+), 4 deletions(-) create mode 100644 mooncake-conductor/include/conductor/prefixindex/prefix_indexer.h create mode 100644 mooncake-conductor/src/prefixindex/prefix_indexer.cpp create mode 100644 mooncake-conductor/tests/prefix_indexer_test.cpp create mode 100644 mooncake-conductor/tests/prefix_indexer_test_peer.h diff --git a/mooncake-conductor/CMakeLists.txt b/mooncake-conductor/CMakeLists.txt index 1603114435..9091bdae38 100644 --- a/mooncake-conductor/CMakeLists.txt +++ b/mooncake-conductor/CMakeLists.txt @@ -15,8 +15,10 @@ endif() find_package(Threads REQUIRED) find_package(OpenSSL REQUIRED COMPONENTS Crypto) -add_library(conductor_cpp_core STATIC src/common/utils.cpp - src/prefixindex/hash_strategy.cpp) +add_library( + conductor_cpp_core STATIC + src/common/utils.cpp src/prefixindex/hash_strategy.cpp + src/prefixindex/prefix_indexer.cpp) # mooncake-common headers are consumed directly instead of linking # mooncake_common: the helpers conductor needs (ascii_string.h, diff --git a/mooncake-conductor/include/conductor/prefixindex/prefix_indexer.h b/mooncake-conductor/include/conductor/prefixindex/prefix_indexer.h new file mode 100644 index 0000000000..0559d5c396 --- /dev/null +++ b/mooncake-conductor/include/conductor/prefixindex/prefix_indexer.h @@ -0,0 +1,135 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "conductor/prefixindex/types.h" + +namespace mooncake::conductor::prefixindex { + +struct RegistrationResult { + bool inserted = false; + std::string error; +}; + +struct BlockPresence { + std::set gpu_owners; + std::set cpu_owners; + std::set disk_owners; + + bool Empty() const { + return gpu_owners.empty() && cpu_owners.empty() && disk_owners.empty(); + } +}; + +// Maximum number of prefixes tracked for each context. +constexpr size_t kDefaultMaxBlocks = 200000; +// Target occupancy ratio for a batched capacity eviction. +constexpr double kEvictTargetRatio = 0.9; + +struct ContextState { + explicit ContextState(HashProfile registered_profile, + size_t block_limit = kDefaultMaxBlocks) + : profile(std::move(registered_profile)), max_blocks(block_limit) {} + + // Lock order is global context-map mutex, then this mutex. Code holding + // this mutex must never reacquire the global mutex. + mutable std::shared_mutex mutex; + const HashProfile profile; + std::map> instance_ranks; + std::unordered_map blocks; + + // Store order: newest prefix at the front, oldest at the back. Query does + // not update this order, so read-only queries can keep a shared lock. + std::list write_order; + std::unordered_map::iterator> + order_pos; + const size_t max_blocks; + // Cumulative number of entries removed by capacity eviction. + int64_t evicted_by_capacity = 0; +}; + +struct RankCacheHitResult { + int64_t gpu = 0; + int64_t cpu = 0; + int64_t disk = 0; + + bool operator==(const RankCacheHitResult&) const = default; +}; + +struct CacheHitResult { + int64_t longest_match_tokens = 0; + std::map dp; + std::map rank_matches; + int64_t gpu = 0; + int64_t cpu = 0; + int64_t disk = 0; +}; + +struct ContextView { + ContextKey context; + HashProfile profile; + std::map> instance_ranks; + size_t prefix_count = 0; +}; + +struct GlobalView { + int32_t context_count = 0; + std::vector contexts; +}; + +class PrefixCacheTable { + public: + PrefixCacheTable() = default; + // Per-context block limit; zero disables capacity eviction. + explicit PrefixCacheTable(size_t block_limit) : block_limit_(block_limit) {} + PrefixCacheTable(const PrefixCacheTable&) = delete; + PrefixCacheTable& operator=(const PrefixCacheTable&) = delete; + + static RegistrationResult ValidateRegistration( + const EngineRegistration& registration); + + RegistrationResult Register(const EngineRegistration& registration); + std::string ValidateProfileBinding(const ContextKey& context, + const HashProfile& profile) const; + std::string Unregister(const ContextKey& context, + const std::string& instance_id, int64_t dp_rank); + + std::string StoreGpu(const GpuMutation& mutation); + std::string RemoveGpu(const GpuMutation& mutation); + std::string ClearGpu(const GpuClear& clear); + + std::string StoreShared(const SharedMutation& mutation); + std::string RemoveShared(const SharedMutation& mutation); + std::string ClearShared(const SharedClear& clear); + + std::map Query( + const ContextKey& context, std::span token_ids, + std::optional cache_salt = std::nullopt, + std::optional instance_filter = std::nullopt) const; + + GlobalView GetGlobalView() const; + + private: + friend class PrefixCacheTableTestPeer; + + std::shared_ptr LoadContextState( + const ContextKey& context) const; + + mutable std::shared_mutex context_map_mutex_; + std::unordered_map> contexts_; + const size_t block_limit_ = kDefaultMaxBlocks; +}; + +} // namespace mooncake::conductor::prefixindex diff --git a/mooncake-conductor/src/prefixindex/prefix_indexer.cpp b/mooncake-conductor/src/prefixindex/prefix_indexer.cpp new file mode 100644 index 0000000000..2d636646b7 --- /dev/null +++ b/mooncake-conductor/src/prefixindex/prefix_indexer.cpp @@ -0,0 +1,615 @@ +#include "conductor/prefixindex/prefix_indexer.h" + +#include + +#include +#include +#include +#include + +#include "conductor/prefixindex/hash_strategy.h" + +namespace mooncake::conductor::prefixindex { + +namespace { + +std::string ValidateContext(const ContextKey& context) { + if (context.tenant_id.empty()) { + return "tenant_id is required"; + } + if (context.model_name.empty()) { + return "model_name is required"; + } + if (context.block_size <= 0) { + return "block_size must be positive"; + } + return ""; +} + +std::string ValidateLayout(const ContextKey& context, + int64_t effective_block_size, + std::optional cache_group) { + if (auto error = ValidateContext(context); !error.empty()) { + return error; + } + if (effective_block_size <= 0) { + return "effective_block_size must be positive"; + } + if (effective_block_size != context.block_size) { + return "effective_block_size must equal ContextKey block_size"; + } + if (cache_group.has_value() && *cache_group != 0) { + return "only cache group 0 is supported"; + } + return ""; +} + +std::string ValidateEngineOwner(const EngineOwner& owner) { + if (owner.source_stream.empty()) { + return "engine owner source_stream is required"; + } + if (owner.instance_id.empty()) { + return "engine owner instance_id is required"; + } + if (owner.dp_rank < 0) { + return "engine owner dp_rank must be non-negative"; + } + return ""; +} + +std::string ValidateSharedOwner(const SharedObjectOwner& owner) { + if (owner.source_stream.empty()) { + return "shared owner source_stream is required"; + } + if (owner.backend_id.empty()) { + return "shared owner backend_id is required"; + } + if (owner.object_id.empty()) { + return "shared owner object_id is required"; + } + return ""; +} + +std::string ValidateGpuMutation(const GpuMutation& mutation) { + if (auto error = + ValidateLayout(mutation.context, mutation.effective_block_size, + mutation.cache_group); + !error.empty()) { + return error; + } + return ValidateEngineOwner(mutation.owner); +} + +std::string ValidateGpuClear(const GpuClear& clear) { + if (auto error = ValidateLayout(clear.context, clear.effective_block_size, + clear.cache_group); + !error.empty()) { + return error; + } + return ValidateEngineOwner(clear.owner); +} + +bool IsSharedTier(StorageTier tier) { + return tier == StorageTier::kCpu || tier == StorageTier::kDisk; +} + +std::string ValidateSharedMutation(const SharedMutation& mutation) { + if (auto error = + ValidateLayout(mutation.context, mutation.effective_block_size, + mutation.cache_group); + !error.empty()) { + return error; + } + if (!IsSharedTier(mutation.tier)) { + return "shared mutation tier must be CPU or DISK"; + } + return ValidateSharedOwner(mutation.owner); +} + +std::string ValidateSharedClear(const SharedClear& clear) { + if (auto error = ValidateLayout(clear.context, clear.effective_block_size, + clear.cache_group); + !error.empty()) { + return error; + } + if (clear.tier.has_value() && !IsSharedTier(*clear.tier)) { + return "shared clear tier must be CPU, DISK, or omitted"; + } + return ValidateSharedOwner(clear.owner); +} + +std::set& SharedOwners(BlockPresence& presence, + StorageTier tier) { + return tier == StorageTier::kCpu ? presence.cpu_owners + : presence.disk_owners; +} + +// Remove empty blocks and their order metadata. The caller holds state.mutex. +void EraseEmptyBlocks(ContextState& state) { + std::erase_if(state.blocks, [&state](const auto& item) { + if (!item.second.Empty()) { + return false; + } + auto pos = state.order_pos.find(item.first); + if (pos != state.order_pos.end()) { + state.write_order.erase(pos->second); + state.order_pos.erase(pos); + } + return true; + }); +} + +// Remove a prefix from the insertion-order metadata. The caller holds +// state.mutex. +void ForgetOrder(ContextState& state, ProjectedPrefix prefix) { + auto pos = state.order_pos.find(prefix); + if (pos != state.order_pos.end()) { + state.write_order.erase(pos->second); + state.order_pos.erase(pos); + } +} + +// Move a prefix to the front of the insertion-order list, adding it if absent. +// The caller holds the state write lock. +void TouchOrder(ContextState& state, ProjectedPrefix prefix) { + auto pos = state.order_pos.find(prefix); + if (pos != state.order_pos.end()) { + state.write_order.splice(state.write_order.begin(), state.write_order, + pos->second); + return; + } + state.write_order.push_front(prefix); + state.order_pos.emplace(prefix, state.write_order.begin()); +} + +// Evict oldest entries in batches until the target occupancy is reached. The +// caller holds the state write lock. +void EvictIfOverCapacity(ContextState& state) { + if (state.max_blocks == 0 || state.blocks.size() <= state.max_blocks) { + return; + } + const size_t target = + static_cast(state.max_blocks * kEvictTargetRatio); + while (state.blocks.size() > target && !state.write_order.empty()) { + const ProjectedPrefix oldest = state.write_order.back(); + state.write_order.pop_back(); + state.order_pos.erase(oldest); + state.blocks.erase(oldest); + ++state.evicted_by_capacity; + } + LOG_EVERY_N(WARNING, 100) + << "Prefix index hit the capacity limit; oldest entries dropped." + << " limit=" << state.max_blocks << " now=" << state.blocks.size() + << " cumulative_evicted=" << state.evicted_by_capacity + << " (non-zero means stored/removed events are out of sync)"; +} + +int64_t TokensForBlocks(size_t block_count, int64_t block_size) { + const uint64_t max_blocks = + static_cast(std::numeric_limits::max() / block_size); + if (block_count > max_blocks) { + return std::numeric_limits::max(); + } + return static_cast(block_count) * block_size; +} + +} // namespace + +RegistrationResult PrefixCacheTable::ValidateRegistration( + const EngineRegistration& registration) { + if (auto error = ValidateLayout(registration.context, + registration.effective_block_size, + registration.cache_group); + !error.empty()) { + return {.error = std::move(error)}; + } + if (registration.instance_id.empty()) { + return {.error = "instance_id is required"}; + } + if (registration.dp_rank < 0) { + return {.error = "dp_rank must be non-negative"}; + } + if (auto error = ValidateHashProfile(registration.profile); + !error.empty()) { + return {.error = std::move(error)}; + } + return {}; +} + +RegistrationResult PrefixCacheTable::Register( + const EngineRegistration& registration) { + if (auto validation = ValidateRegistration(registration); + !validation.error.empty()) { + return validation; + } + + auto candidate = + std::make_shared(registration.profile, block_limit_); + candidate->instance_ranks[registration.instance_id].insert( + registration.dp_rank); + + std::shared_ptr state; + { + std::unique_lock map_lock(context_map_mutex_); + auto [it, inserted] = + contexts_.try_emplace(registration.context, std::move(candidate)); + if (inserted) { + return {.inserted = true, .error = ""}; + } + state = it->second; + } + + std::unique_lock state_lock(state->mutex); + if (state->profile != registration.profile) { + return {.error = + "registration conflicts with the ContextKey hash profile"}; + } + const bool inserted = state->instance_ranks[registration.instance_id] + .insert(registration.dp_rank) + .second; + return {.inserted = inserted, .error = ""}; +} + +std::shared_ptr PrefixCacheTable::LoadContextState( + const ContextKey& context) const { + std::shared_lock map_lock(context_map_mutex_); + auto it = contexts_.find(context); + return it == contexts_.end() ? nullptr : it->second; +} + +std::string PrefixCacheTable::ValidateProfileBinding( + const ContextKey& context, const HashProfile& profile) const { + if (auto error = ValidateContext(context); !error.empty()) { + return error; + } + if (auto error = ValidateHashProfile(profile); !error.empty()) { + return error; + } + auto state = LoadContextState(context); + if (!state) { + return "ContextKey is not registered"; + } + + std::shared_lock state_lock(state->mutex); + if (state->profile != profile) { + return "hash profile conflicts with the registered ContextKey profile"; + } + return ""; +} + +std::string PrefixCacheTable::Unregister(const ContextKey& context, + const std::string& instance_id, + int64_t dp_rank) { + if (auto error = ValidateContext(context); !error.empty()) { + return error; + } + if (instance_id.empty()) { + return "instance_id is required"; + } + if (dp_rank < 0) { + return "dp_rank must be non-negative"; + } + + auto state = LoadContextState(context); + if (!state) { + return ""; + } + + std::unique_lock state_lock(state->mutex); + auto instance = state->instance_ranks.find(instance_id); + if (instance != state->instance_ranks.end()) { + instance->second.erase(dp_rank); + if (instance->second.empty()) { + state->instance_ranks.erase(instance); + } + } + + for (auto& [unused_prefix, presence] : state->blocks) { + (void)unused_prefix; + std::erase_if(presence.gpu_owners, [&](const EngineOwner& owner) { + return owner.instance_id == instance_id && owner.dp_rank == dp_rank; + }); + } + EraseEmptyBlocks(*state); + return ""; +} + +std::string PrefixCacheTable::StoreGpu(const GpuMutation& mutation) { + if (auto error = ValidateGpuMutation(mutation); !error.empty()) { + return error; + } + auto state = LoadContextState(mutation.context); + if (!state) { + return "ContextKey is not registered"; + } + + std::unique_lock state_lock(state->mutex); + auto instance = state->instance_ranks.find(mutation.owner.instance_id); + if (instance == state->instance_ranks.end() || + !instance->second.contains(mutation.owner.dp_rank)) { + return "engine owner instance/rank is not registered"; + } + for (ProjectedPrefix prefix : mutation.prefixes) { + state->blocks[prefix].gpu_owners.insert(mutation.owner); + TouchOrder(*state, prefix); + } + EvictIfOverCapacity(*state); + return ""; +} + +std::string PrefixCacheTable::RemoveGpu(const GpuMutation& mutation) { + if (auto error = ValidateGpuMutation(mutation); !error.empty()) { + return error; + } + auto state = LoadContextState(mutation.context); + if (!state) { + return ""; + } + + std::unique_lock state_lock(state->mutex); + // Only prefixes in the mutation can become empty, so avoid a full scan. + for (ProjectedPrefix prefix : mutation.prefixes) { + auto block = state->blocks.find(prefix); + if (block != state->blocks.end()) { + block->second.gpu_owners.erase(mutation.owner); + if (block->second.Empty()) { + state->blocks.erase(block); + ForgetOrder(*state, prefix); + } + } + } + return ""; +} + +std::string PrefixCacheTable::ClearGpu(const GpuClear& clear) { + if (auto error = ValidateGpuClear(clear); !error.empty()) { + return error; + } + auto state = LoadContextState(clear.context); + if (!state) { + return ""; + } + + std::unique_lock state_lock(state->mutex); + for (auto& [unused_prefix, presence] : state->blocks) { + (void)unused_prefix; + presence.gpu_owners.erase(clear.owner); + } + EraseEmptyBlocks(*state); + return ""; +} + +std::string PrefixCacheTable::StoreShared(const SharedMutation& mutation) { + if (auto error = ValidateSharedMutation(mutation); !error.empty()) { + return error; + } + auto state = LoadContextState(mutation.context); + if (!state) { + return "ContextKey is not registered"; + } + + std::unique_lock state_lock(state->mutex); + for (ProjectedPrefix prefix : mutation.prefixes) { + SharedOwners(state->blocks[prefix], mutation.tier) + .insert(mutation.owner); + TouchOrder(*state, prefix); + } + EvictIfOverCapacity(*state); + return ""; +} + +std::string PrefixCacheTable::RemoveShared(const SharedMutation& mutation) { + if (auto error = ValidateSharedMutation(mutation); !error.empty()) { + return error; + } + auto state = LoadContextState(mutation.context); + if (!state) { + return ""; + } + + std::unique_lock state_lock(state->mutex); + // Only prefixes in the mutation can become empty, so avoid a full scan. + for (ProjectedPrefix prefix : mutation.prefixes) { + auto block = state->blocks.find(prefix); + if (block != state->blocks.end()) { + SharedOwners(block->second, mutation.tier).erase(mutation.owner); + if (block->second.Empty()) { + state->blocks.erase(block); + ForgetOrder(*state, prefix); + } + } + } + return ""; +} + +std::string PrefixCacheTable::ClearShared(const SharedClear& clear) { + if (auto error = ValidateSharedClear(clear); !error.empty()) { + return error; + } + auto state = LoadContextState(clear.context); + if (!state) { + return ""; + } + + std::unique_lock state_lock(state->mutex); + for (auto& [unused_prefix, presence] : state->blocks) { + (void)unused_prefix; + if (!clear.tier.has_value() || *clear.tier == StorageTier::kCpu) { + presence.cpu_owners.erase(clear.owner); + } + if (!clear.tier.has_value() || *clear.tier == StorageTier::kDisk) { + presence.disk_owners.erase(clear.owner); + } + } + EraseEmptyBlocks(*state); + return ""; +} + +std::map PrefixCacheTable::Query( + const ContextKey& context, std::span token_ids, + std::optional cache_salt, + std::optional instance_filter) const { + std::map results; + auto state = LoadContextState(context); + if (!state) { + return results; + } + + // The profile is immutable, and the shared_ptr keeps state alive while the + // hash strategy and chain are built without holding state.mutex. + std::string strategy_error; + auto strategy = CreateHashStrategy(state->profile, &strategy_error); + if (!strategy) { + LOG(ERROR) << "Registered hash profile became invalid: " + << strategy_error; + return results; + } + + std::string chain_error; + auto chain = strategy->CreateChain(context, token_ids, + std::move(cache_salt), &chain_error); + if (!chain) { + LOG(ERROR) << "Query hash chain setup failed: " << chain_error; + return results; + } + const size_t block_count = chain->BlockCount(); + + // Resolve the optional filter and copy rank sets before probing. The copies + // remain valid while the probe releases and reacquires state.mutex. + std::map> selected_instances; + { + std::shared_lock select_lock(state->mutex); + if (instance_filter.has_value()) { + auto instance = state->instance_ranks.find(*instance_filter); + if (instance == state->instance_ranks.end()) { + return results; + } + selected_instances.emplace(instance->first, instance->second); + } else { + selected_instances = state->instance_ranks; + } + } + if (selected_instances.empty()) { + return results; + } + + // Probe indexed block presence in chunks. Hashing runs outside the lock; + // each chunk holds a shared lock only for table lookups. + constexpr size_t kProbeChunkMin = 8; + constexpr size_t kProbeChunkMax = 512; + size_t chunk = kProbeChunkMin; + size_t probe_depth = 0; + bool probe_stalled = false; + while (!probe_stalled && probe_depth < block_count) { + const size_t chunk_end = std::min(probe_depth + chunk, block_count); + chunk = std::min(chunk * 2, kProbeChunkMax); + // Compute hashes without holding state.mutex; retain any error for the + // final check below. + for (size_t i = probe_depth; i < chunk_end; ++i) { + if (chain->At(i, &chain_error) == nullptr) { + probe_stalled = true; + break; + } + } + std::shared_lock probe_lock(state->mutex); + while (probe_depth < chunk_end) { + const HashBlock* hashed = chain->At(probe_depth, &chain_error); + if (hashed == nullptr || + !state->blocks.contains(hashed->projected)) { + probe_stalled = true; + break; + } + ++probe_depth; + } + } + + std::shared_lock state_lock(state->mutex); + + // Hashes needed by the probe are memoized, so the final read-locked walk + // performs only vector access and indexed lookups. + auto advance_cursor = [&](size_t& cursor, const auto& present) { + while (cursor < block_count) { + const HashBlock* hashed = chain->At(cursor, &chain_error); + if (hashed == nullptr) { + cursor = block_count; // stall every remaining walk + return; + } + auto block = state->blocks.find(hashed->projected); + if (block == state->blocks.end() || !present(block->second)) { + break; + } + ++cursor; + } + }; + + for (const auto& [instance_id, ranks] : selected_instances) { + CacheHitResult result; + + for (int64_t rank : ranks) { + auto gpu_present = [&](const BlockPresence& block) { + return std::any_of( + block.gpu_owners.begin(), block.gpu_owners.end(), + [&](const EngineOwner& owner) { + return owner.instance_id == instance_id && + owner.dp_rank == rank; + }); + }; + + size_t cursor = 0; + advance_cursor(cursor, gpu_present); + + RankCacheHitResult rank_match; + rank_match.gpu = TokensForBlocks(cursor, context.block_size); + + advance_cursor(cursor, [](const BlockPresence& block) { + return !block.cpu_owners.empty(); + }); + rank_match.cpu = TokensForBlocks(cursor, context.block_size); + + advance_cursor(cursor, [](const BlockPresence& block) { + return !block.disk_owners.empty(); + }); + rank_match.disk = TokensForBlocks(cursor, context.block_size); + + result.dp.emplace(rank, rank_match.gpu); + result.rank_matches.emplace(rank, rank_match); + result.gpu = std::max(result.gpu, rank_match.gpu); + result.cpu = std::max(result.cpu, rank_match.cpu); + result.disk = std::max(result.disk, rank_match.disk); + } + result.longest_match_tokens = result.disk; + results.emplace(instance_id, std::move(result)); + } + if (!chain_error.empty()) { + LOG(ERROR) << "Query hash computation failed: " << chain_error; + return {}; + } + return results; +} + +GlobalView PrefixCacheTable::GetGlobalView() const { + GlobalView view; + std::vector>> contexts; + { + std::shared_lock map_lock(context_map_mutex_); + contexts.reserve(contexts_.size()); + for (const auto& item : contexts_) { + contexts.push_back(item); + } + } + + view.context_count = static_cast(contexts.size()); + view.contexts.reserve(contexts.size()); + for (const auto& [context, state] : contexts) { + std::shared_lock state_lock(state->mutex); + ContextView context_view; + context_view.context = context; + context_view.profile = state->profile; + context_view.instance_ranks = state->instance_ranks; + context_view.prefix_count = state->blocks.size(); + view.contexts.push_back(std::move(context_view)); + } + return view; +} + +} // namespace mooncake::conductor::prefixindex diff --git a/mooncake-conductor/tests/CMakeLists.txt b/mooncake-conductor/tests/CMakeLists.txt index fa9b202c8d..2848da58d2 100644 --- a/mooncake-conductor/tests/CMakeLists.txt +++ b/mooncake-conductor/tests/CMakeLists.txt @@ -7,8 +7,10 @@ else() set(CONDUCTOR_GTEST_LIBS GTest::gtest GTest::gtest_main) endif() -add_executable(conductor_test common_utils_test.cpp json_uint64_test.cpp - model_context_test.cpp compute_hash_test.cpp) +add_executable( + conductor_test + common_utils_test.cpp json_uint64_test.cpp model_context_test.cpp + compute_hash_test.cpp prefix_indexer_test.cpp) target_link_libraries(conductor_test PRIVATE conductor_cpp_core ${CONDUCTOR_GTEST_LIBS}) diff --git a/mooncake-conductor/tests/prefix_indexer_test.cpp b/mooncake-conductor/tests/prefix_indexer_test.cpp new file mode 100644 index 0000000000..435c200ccb --- /dev/null +++ b/mooncake-conductor/tests/prefix_indexer_test.cpp @@ -0,0 +1,952 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "conductor/prefixindex/hash_strategy.h" +#include "conductor/prefixindex/prefix_indexer.h" +#include "prefix_indexer_test_peer.h" + +namespace { + +using mooncake::conductor::prefixindex::BlockPresenceSnapshot; +using mooncake::conductor::prefixindex::CacheHitResult; +using mooncake::conductor::prefixindex::ContextKey; +using mooncake::conductor::prefixindex::EngineOwner; +using mooncake::conductor::prefixindex::EngineRegistration; +using mooncake::conductor::prefixindex::GpuClear; +using mooncake::conductor::prefixindex::GpuMutation; +using mooncake::conductor::prefixindex::HashBlock; +using mooncake::conductor::prefixindex::HashProfile; +using mooncake::conductor::prefixindex::PrefixCacheTable; +using mooncake::conductor::prefixindex::PrefixCacheTableSnapshot; +using mooncake::conductor::prefixindex::PrefixCacheTableTestPeer; +using mooncake::conductor::prefixindex::ProjectedPrefix; +using mooncake::conductor::prefixindex::RankCacheHitResult; +using mooncake::conductor::prefixindex::SharedClear; +using mooncake::conductor::prefixindex::SharedMutation; +using mooncake::conductor::prefixindex::SharedObjectOwner; +using mooncake::conductor::prefixindex::StorageTier; + +constexpr char kRootDigest[] = + "4e1195df020de59e0d65a33a4279f1183e7ae4e5d980e309f8b55adff2e61c3e"; +constexpr char kPaddedSeedRootDigest[] = + "8d912e4e62b3cc377b1d1c7a14ef61dffbdaa0990237035c05401c29414c4172"; +constexpr char kPickleRootDigest[] = + "1973e23848344dc43a988a9b478663803cfffe1243480253f9a3cf004b14aa7c"; + +ContextKey TestContext(int64_t block_size = 16) { + return {.tenant_id = "tenant-a", + .model_name = "model-a", + .lora_name = "", + .block_size = block_size}; +} + +HashProfile TestProfile() { + return {.strategy = "vllm_v1", + .algorithm = "sha256_cbor", + .python_hash_seed = "0", + .root_digest = kRootDigest, + .index_projection = "low64_be"}; +} + +HashProfile PaddedSeedProfile() { + return {.strategy = "vllm_v1", + .algorithm = "sha256_cbor", + .python_hash_seed = "00", + .root_digest = kPaddedSeedRootDigest, + .index_projection = "low64_be"}; +} + +HashProfile PickleProfile() { + return {.strategy = "vllm_v1", + .algorithm = "sha256", + .python_hash_seed = "0", + .root_digest = kPickleRootDigest, + .index_projection = "low64_be"}; +} + +EngineRegistration Registration(const std::string& instance_id = "instance-a", + int64_t dp_rank = 0) { + const ContextKey context = TestContext(); + return {.context = context, + .profile = TestProfile(), + .instance_id = instance_id, + .dp_rank = dp_rank, + .effective_block_size = context.block_size, + .cache_group = 0}; +} + +EngineOwner GpuOwner(const std::string& instance_id = "instance-a", + int64_t dp_rank = 0, + const std::string& stream = "stream-a") { + return {.source_stream = stream, + .instance_id = instance_id, + .dp_rank = dp_rank}; +} + +SharedObjectOwner SharedOwner(const std::string& object_id = "object-a", + const std::string& stream = "pool-stream", + const std::string& backend = "backend-a") { + return { + .source_stream = stream, .backend_id = backend, .object_id = object_id}; +} + +ProjectedPrefix Prefix(uint64_t value) { return {.value = value}; } + +RankCacheHitResult RankMatch(int64_t gpu, int64_t cpu, int64_t disk) { + return {.gpu = gpu, .cpu = cpu, .disk = disk}; +} + +GpuMutation Gpu(const std::vector& prefixes, + EngineOwner owner = GpuOwner()) { + const ContextKey context = TestContext(); + return {.context = context, + .prefixes = prefixes, + .owner = std::move(owner), + .effective_block_size = context.block_size, + .cache_group = 0}; +} + +SharedMutation Shared(const std::vector& prefixes, + StorageTier tier, + SharedObjectOwner owner = SharedOwner()) { + const ContextKey context = TestContext(); + return {.context = context, + .prefixes = prefixes, + .tier = tier, + .owner = std::move(owner), + .effective_block_size = context.block_size, + .cache_group = 0}; +} + +GpuClear ClearFor(EngineOwner owner = GpuOwner()) { + const ContextKey context = TestContext(); + return {.context = context, + .owner = std::move(owner), + .effective_block_size = context.block_size, + .cache_group = 0}; +} + +SharedClear ClearFor(SharedObjectOwner owner, + std::optional tier = std::nullopt) { + const ContextKey context = TestContext(); + return {.context = context, + .owner = std::move(owner), + .tier = tier, + .effective_block_size = context.block_size, + .cache_group = 0}; +} + +std::vector Tokens(size_t count) { + std::vector tokens; + tokens.reserve(count); + for (size_t i = 0; i < count; ++i) { + tokens.push_back(static_cast(i + 1)); + } + return tokens; +} + +std::vector Hashes( + const std::vector& tokens, + std::optional cache_salt = std::nullopt) { + std::string error; + auto strategy = mooncake::conductor::prefixindex::CreateHashStrategy( + TestProfile(), &error); + EXPECT_TRUE(error.empty()) << error; + if (!strategy) { + return {}; + } + + std::vector blocks; + error = strategy->Compute(TestContext(), tokens, std::move(cache_salt), + &blocks); + EXPECT_TRUE(error.empty()) << error; + + std::vector prefixes; + prefixes.reserve(blocks.size()); + for (const HashBlock& block : blocks) { + prefixes.push_back(block.projected); + } + return prefixes; +} + +void RegisterOrFail(PrefixCacheTable& table, + const EngineRegistration& registration) { + const auto result = table.Register(registration); + ASSERT_TRUE(result.error.empty()) << result.error; +} + +BlockPresenceSnapshot Presence(const PrefixCacheTable& table, + ProjectedPrefix prefix) { + const PrefixCacheTableSnapshot table_snapshot = + PrefixCacheTableTestPeer::Snapshot(table); + return table_snapshot.contexts.at(TestContext()).blocks.at(prefix); +} + +TEST(Registration, InvalidInputsDoNotCreateContextState) { + std::vector invalid; + + auto non_positive = Registration(); + non_positive.context.block_size = 0; + non_positive.effective_block_size = 0; + invalid.push_back(non_positive); + + auto mismatch = Registration(); + mismatch.effective_block_size = 8; + invalid.push_back(mismatch); + + auto unsupported_group = Registration(); + unsupported_group.cache_group = 1; + invalid.push_back(unsupported_group); + + auto empty_instance = Registration(); + empty_instance.instance_id.clear(); + invalid.push_back(empty_instance); + + auto negative_rank = Registration(); + negative_rank.dp_rank = -1; + invalid.push_back(negative_rank); + + auto malformed_profile = Registration(); + malformed_profile.profile.root_digest = "not-a-digest"; + invalid.push_back(malformed_profile); + + PrefixCacheTable table; + for (const auto& registration : invalid) { + SCOPED_TRACE(registration.instance_id); + const auto validation = + PrefixCacheTable::ValidateRegistration(registration); + EXPECT_FALSE(validation.error.empty()); + const auto result = table.Register(registration); + EXPECT_FALSE(result.error.empty()); + EXPECT_FALSE(result.inserted); + } + EXPECT_EQ(table.GetGlobalView().context_count, 0); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table).contexts.empty()); +} + +TEST(Registration, ForgedSeedRootPairIsRejectedWithoutMutation) { + PrefixCacheTable table; + auto forged = Registration(); + forged.profile.root_digest = kPaddedSeedRootDigest; + + const auto validation = PrefixCacheTable::ValidateRegistration(forged); + EXPECT_NE(validation.error.find("does not match"), std::string::npos); + const auto rejected = table.Register(forged); + EXPECT_NE(rejected.error.find("does not match"), std::string::npos); + EXPECT_FALSE(rejected.inserted); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table).contexts.empty()); + + RegisterOrFail(table, Registration()); + const auto registered = PrefixCacheTableTestPeer::Snapshot(table); + EXPECT_NE(table.ValidateProfileBinding(TestContext(), forged.profile) + .find("does not match"), + std::string::npos); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), registered); + + forged.instance_id = "instance-b"; + const auto conflicting = table.Register(forged); + EXPECT_NE(conflicting.error.find("does not match"), std::string::npos); + EXPECT_FALSE(conflicting.inserted); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), registered); +} + +TEST(Registration, TracksEveryInstanceAndRankIdempotently) { + PrefixCacheTable table; + + auto first = table.Register(Registration("instance-a", 0)); + ASSERT_TRUE(first.error.empty()) << first.error; + EXPECT_TRUE(first.inserted); + + auto duplicate = table.Register(Registration("instance-a", 0)); + ASSERT_TRUE(duplicate.error.empty()) << duplicate.error; + EXPECT_FALSE(duplicate.inserted); + + auto omitted_group = Registration("instance-a", 0); + omitted_group.cache_group.reset(); + auto omitted_duplicate = table.Register(omitted_group); + ASSERT_TRUE(omitted_duplicate.error.empty()) << omitted_duplicate.error; + EXPECT_FALSE(omitted_duplicate.inserted); + + auto second_rank = table.Register(Registration("instance-a", 2)); + ASSERT_TRUE(second_rank.error.empty()) << second_rank.error; + EXPECT_TRUE(second_rank.inserted); + + auto second_instance = table.Register(Registration("instance-b", 1)); + ASSERT_TRUE(second_instance.error.empty()) << second_instance.error; + EXPECT_TRUE(second_instance.inserted); + + const auto snapshot = PrefixCacheTableTestPeer::Snapshot(table); + ASSERT_EQ(snapshot.contexts.size(), 1u); + const auto& state = snapshot.contexts.at(TestContext()); + EXPECT_EQ(state.profile, TestProfile()); + EXPECT_EQ(state.instance_ranks.at("instance-a"), (std::set{0, 2})); + EXPECT_EQ(state.instance_ranks.at("instance-b"), (std::set{1})); + EXPECT_TRUE(state.blocks.empty()); +} + +TEST(Registration, ConflictingProfilePreservesCompleteState) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(1)})), ""); + const auto before = PrefixCacheTableTestPeer::Snapshot(table); + + auto conflicting = Registration("instance-b", 1); + conflicting.profile = PaddedSeedProfile(); + const auto result = table.Register(conflicting); + + EXPECT_FALSE(result.error.empty()); + EXPECT_FALSE(result.inserted); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); +} + +TEST(Registration, ProfileBindingValidationIsExactAndLookupOnly) { + PrefixCacheTable table; + const auto empty_before = PrefixCacheTableTestPeer::Snapshot(table); + + EXPECT_FALSE( + table.ValidateProfileBinding(TestContext(), TestProfile()).empty()); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), empty_before); + + RegisterOrFail(table, Registration()); + const auto registered = PrefixCacheTableTestPeer::Snapshot(table); + EXPECT_EQ(table.ValidateProfileBinding(TestContext(), TestProfile()), ""); + + const HashProfile conflict = PaddedSeedProfile(); + EXPECT_FALSE(table.ValidateProfileBinding(TestContext(), conflict).empty()); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), registered); +} + +TEST(Registration, MixedAlgorithmsUnderOneContextAreRejected) { + // The resolved profile is immutable per ContextKey: the same seed under + // the other supported algorithm is still a conflict, in both orders. + { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto before = PrefixCacheTableTestPeer::Snapshot(table); + + auto conflicting = Registration("instance-b", 1); + conflicting.profile = PickleProfile(); + const auto result = table.Register(conflicting); + EXPECT_FALSE(result.error.empty()); + EXPECT_FALSE(result.inserted); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); + + EXPECT_FALSE( + table.ValidateProfileBinding(TestContext(), PickleProfile()) + .empty()); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); + } + { + PrefixCacheTable table; + auto pickle_registration = Registration(); + pickle_registration.profile = PickleProfile(); + RegisterOrFail(table, pickle_registration); + const auto before = PrefixCacheTableTestPeer::Snapshot(table); + EXPECT_EQ(table.ValidateProfileBinding(TestContext(), PickleProfile()), + ""); + + auto conflicting = Registration("instance-b", 1); + const auto result = table.Register(conflicting); + EXPECT_FALSE(result.error.empty()); + EXPECT_FALSE(result.inserted); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); + } +} + +TEST(Mutations, StoreRequiresKnownContextAndRegisteredGpuRank) { + PrefixCacheTable table; + const auto gpu = Gpu({Prefix(1)}); + const auto shared = Shared({Prefix(1)}, StorageTier::kCpu); + + EXPECT_FALSE(table.StoreGpu(gpu).empty()); + EXPECT_FALSE(table.StoreShared(shared).empty()); + EXPECT_EQ(table.GetGlobalView().context_count, 0); + + RegisterOrFail(table, Registration("instance-a", 1)); + EXPECT_FALSE(table.StoreGpu(gpu).empty()); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table) + .contexts.at(TestContext()) + .blocks.empty()); +} + +TEST(Mutations, InvalidGroupTierAndOwnersPreserveState) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(1)})), ""); + const auto before = PrefixCacheTableTestPeer::Snapshot(table); + + auto bad_group = Gpu({Prefix(2)}); + bad_group.cache_group = 3; + EXPECT_FALSE(table.StoreGpu(bad_group).empty()); + + auto bad_owner = Gpu({Prefix(2)}); + bad_owner.owner.source_stream.clear(); + EXPECT_FALSE(table.StoreGpu(bad_owner).empty()); + + auto bad_tier = Shared({Prefix(2)}, StorageTier::kGpu); + EXPECT_FALSE(table.StoreShared(bad_tier).empty()); + + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); +} + +TEST(Mutations, DuplicateGpuStoreAndRemoveAreIdempotent) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const ProjectedPrefix prefix = Prefix(7); + const auto mutation = Gpu({prefix}); + + ASSERT_EQ(table.StoreGpu(mutation), ""); + ASSERT_EQ(table.StoreGpu(mutation), ""); + EXPECT_EQ(Presence(table, prefix).gpu_owners, + (std::set{GpuOwner()})); + + auto absent_owner = Gpu({prefix}, GpuOwner("instance-b", 0, "stream-b")); + ASSERT_EQ(table.RemoveGpu(absent_owner), ""); + EXPECT_EQ(Presence(table, prefix).gpu_owners, + (std::set{GpuOwner()})); + + ASSERT_EQ(table.RemoveGpu(mutation), ""); + ASSERT_EQ(table.RemoveGpu(mutation), ""); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table) + .contexts.at(TestContext()) + .blocks.empty()); +} + +TEST(Mutations, CollidingSharedOwnersRemainIndependentlyRemovable) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const ProjectedPrefix collision = Prefix(0x123456789abcdef0ULL); + const SharedObjectOwner first = SharedOwner("object-a"); + const SharedObjectOwner second = SharedOwner("object-b"); + + ASSERT_EQ(table.StoreShared(Shared({collision}, StorageTier::kCpu, first)), + ""); + ASSERT_EQ(table.StoreShared(Shared({collision}, StorageTier::kCpu, second)), + ""); + ASSERT_EQ(table.StoreShared(Shared({collision}, StorageTier::kCpu, first)), + ""); + EXPECT_EQ(Presence(table, collision).cpu_owners, + (std::set{first, second})); + + ASSERT_EQ(table.RemoveShared(Shared({collision}, StorageTier::kCpu, first)), + ""); + EXPECT_EQ(Presence(table, collision).cpu_owners, + (std::set{second})); + + ASSERT_EQ( + table.RemoveShared(Shared({collision}, StorageTier::kCpu, second)), ""); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table) + .contexts.at(TestContext()) + .blocks.empty()); +} + +TEST(Mutations, BlockLivesUntilEveryTierOwnerSetIsEmpty) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const ProjectedPrefix prefix = Prefix(11); + const auto gpu = Gpu({prefix}); + const auto cpu = Shared({prefix}, StorageTier::kCpu, SharedOwner("cpu")); + const auto disk = Shared({prefix}, StorageTier::kDisk, SharedOwner("disk")); + + ASSERT_EQ(table.StoreGpu(gpu), ""); + ASSERT_EQ(table.StoreShared(cpu), ""); + ASSERT_EQ(table.StoreShared(disk), ""); + ASSERT_EQ(table.RemoveGpu(gpu), ""); + EXPECT_TRUE(Presence(table, prefix).gpu_owners.empty()); + EXPECT_FALSE(Presence(table, prefix).cpu_owners.empty()); + EXPECT_FALSE(Presence(table, prefix).disk_owners.empty()); + + ASSERT_EQ(table.RemoveShared(cpu), ""); + EXPECT_TRUE(Presence(table, prefix).cpu_owners.empty()); + EXPECT_FALSE(Presence(table, prefix).disk_owners.empty()); + + ASSERT_EQ(table.RemoveShared(disk), ""); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table) + .contexts.at(TestContext()) + .blocks.empty()); +} + +TEST(Mutations, GpuAndSharedClearAreExactlyOwnerScoped) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-b", 1)); + const ProjectedPrefix prefix = Prefix(21); + const EngineOwner engine_a = GpuOwner("instance-a", 0, "stream-a"); + const EngineOwner engine_a_other_stream = + GpuOwner("instance-a", 0, "stream-a-other"); + const EngineOwner engine_b = GpuOwner("instance-b", 1, "stream-b"); + const SharedObjectOwner shared_a = SharedOwner("object-a"); + const SharedObjectOwner shared_b = SharedOwner("object-b"); + + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, engine_a)), ""); + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, engine_a_other_stream)), ""); + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, engine_b)), ""); + ASSERT_EQ(table.StoreShared(Shared({prefix}, StorageTier::kCpu, shared_a)), + ""); + ASSERT_EQ(table.StoreShared(Shared({prefix}, StorageTier::kDisk, shared_a)), + ""); + ASSERT_EQ(table.StoreShared(Shared({prefix}, StorageTier::kCpu, shared_b)), + ""); + + ASSERT_EQ(table.ClearGpu(ClearFor(engine_a)), ""); + EXPECT_EQ(Presence(table, prefix).gpu_owners, + (std::set{engine_a_other_stream, engine_b})); + EXPECT_EQ(Presence(table, prefix).cpu_owners, + (std::set{shared_a, shared_b})); + + ASSERT_EQ(table.ClearShared(ClearFor(shared_a, StorageTier::kCpu)), ""); + EXPECT_EQ(Presence(table, prefix).cpu_owners, + (std::set{shared_b})); + EXPECT_EQ(Presence(table, prefix).disk_owners, + (std::set{shared_a})); + EXPECT_EQ(Presence(table, prefix).gpu_owners, + (std::set{engine_a_other_stream, engine_b})); + + ASSERT_EQ(table.ClearShared(ClearFor(shared_a)), ""); + EXPECT_TRUE(Presence(table, prefix).disk_owners.empty()); + EXPECT_EQ(Presence(table, prefix).gpu_owners, + (std::set{engine_a_other_stream, engine_b})); +} + +TEST(Mutations, UnknownRemoveClearAndUnregisterNeverCreateState) { + PrefixCacheTable table; + const ContextKey context = TestContext(); + + EXPECT_EQ(table.RemoveGpu(Gpu({Prefix(1)})), ""); + EXPECT_EQ(table.ClearGpu(ClearFor()), ""); + EXPECT_EQ(table.RemoveShared(Shared({Prefix(1)}, StorageTier::kCpu)), ""); + EXPECT_EQ(table.ClearShared(ClearFor(SharedOwner())), ""); + EXPECT_EQ(table.Unregister(context, "instance-a", 0), ""); + + EXPECT_FALSE(PrefixCacheTableTestPeer::ContextExists(table, context)); + EXPECT_TRUE(PrefixCacheTableTestPeer::Snapshot(table).contexts.empty()); +} + +TEST(Unregister, RemovesOnlySelectedRankGpuOwners) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-a", 1)); + RegisterOrFail(table, Registration("instance-b", 0)); + const ProjectedPrefix prefix = Prefix(31); + const EngineOwner a0 = GpuOwner("instance-a", 0, "stream-a0"); + const EngineOwner a0_second_stream = + GpuOwner("instance-a", 0, "stream-a0-second"); + const EngineOwner a1 = GpuOwner("instance-a", 1, "stream-a1"); + const EngineOwner b0 = GpuOwner("instance-b", 0, "stream-b0"); + const SharedObjectOwner shared = SharedOwner(); + + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, a0)), ""); + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, a0_second_stream)), ""); + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, a1)), ""); + ASSERT_EQ(table.StoreGpu(Gpu({prefix}, b0)), ""); + ASSERT_EQ(table.StoreShared(Shared({prefix}, StorageTier::kCpu, shared)), + ""); + + ASSERT_EQ(table.Unregister(TestContext(), "instance-a", 0), ""); + auto snapshot = PrefixCacheTableTestPeer::Snapshot(table); + const auto& state = snapshot.contexts.at(TestContext()); + EXPECT_EQ(state.instance_ranks.at("instance-a"), (std::set{1})); + EXPECT_EQ(state.instance_ranks.at("instance-b"), (std::set{0})); + EXPECT_EQ(state.blocks.at(prefix).gpu_owners, + (std::set{a1, b0})); + EXPECT_EQ(state.blocks.at(prefix).cpu_owners, + (std::set{shared})); + + ASSERT_EQ(table.Unregister(TestContext(), "instance-a", 1), ""); + snapshot = PrefixCacheTableTestPeer::Snapshot(table); + EXPECT_FALSE(snapshot.contexts.at(TestContext()) + .instance_ranks.contains("instance-a")); + EXPECT_EQ(snapshot.contexts.at(TestContext()).blocks.at(prefix).gpu_owners, + (std::set{b0})); + EXPECT_EQ(snapshot.contexts.at(TestContext()).blocks.at(prefix).cpu_owners, + (std::set{shared})); +} + +TEST(Query, ExactTwoInstanceSharedCacheExample) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-1", 0)); + RegisterOrFail(table, Registration("instance-2", 1)); + const auto tokens = Tokens(48); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 3u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[1]}, + GpuOwner("instance-1", 0, "engine-1"))), + ""); + ASSERT_EQ(table.StoreShared( + Shared(hashes, StorageTier::kCpu, SharedOwner("cpu-object"))), + ""); + ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kDisk, + SharedOwner("disk-object"))), + ""); + + const auto results = table.Query(TestContext(), tokens); + ASSERT_EQ(results.size(), 2u); + + const CacheHitResult& first = results.at("instance-1"); + EXPECT_EQ(first.longest_match_tokens, 48); + EXPECT_EQ(first.gpu, 32); + EXPECT_EQ(first.dp, (std::map{{0, 32}})); + EXPECT_EQ( + first.rank_matches, + (std::map{{0, RankMatch(32, 48, 48)}})); + EXPECT_EQ(first.cpu, 48); + EXPECT_EQ(first.disk, 48); + + const CacheHitResult& second = results.at("instance-2"); + EXPECT_EQ(second.longest_match_tokens, 48); + EXPECT_EQ(second.gpu, 0); + EXPECT_EQ(second.dp, (std::map{{1, 0}})); + EXPECT_EQ( + second.rank_matches, + (std::map{{1, RankMatch(0, 48, 48)}})); + EXPECT_EQ(second.cpu, 48); + EXPECT_EQ(second.disk, 48); +} + +TEST(Query, GpuCpuAndDiskExtendOneCumulativePrefix) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(64); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 4u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[1]})), ""); + ASSERT_EQ(table.StoreShared(Shared({hashes[2]}, StorageTier::kCpu)), ""); + ASSERT_EQ(table.StoreShared(Shared({hashes[3]}, StorageTier::kDisk)), ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 64); + EXPECT_EQ(result.gpu, 32); + EXPECT_EQ(result.dp, (std::map{{0, 32}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(32, 48, 64)}})); + EXPECT_EQ(result.cpu, 48); + EXPECT_EQ(result.disk, 64); +} + +TEST(Query, EmptyCpuPhaseFallsThroughToDiskAtSameBlock) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(48); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 3u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0]})), ""); + ASSERT_EQ( + table.StoreShared(Shared({hashes[1], hashes[2]}, StorageTier::kDisk)), + ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 48); + EXPECT_EQ(result.gpu, 16); + EXPECT_EQ(result.dp, (std::map{{0, 16}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(16, 16, 48)}})); + EXPECT_EQ(result.cpu, 16); + EXPECT_EQ(result.disk, 48); +} + +TEST(Query, CompleteGpuCoverageCarriesThroughLowerTierBoundaries) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(48); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 3u); + + ASSERT_EQ(table.StoreGpu(Gpu(hashes)), ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 48); + EXPECT_EQ(result.gpu, 48); + EXPECT_EQ(result.dp, (std::map{{0, 48}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(48, 48, 48)}})); + EXPECT_EQ(result.cpu, 48); + EXPECT_EQ(result.disk, 48); +} + +TEST(Query, DuplicateTierPresenceIsAttributedOnce) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(48); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 3u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0]})), ""); + ASSERT_EQ( + table.StoreShared(Shared({hashes[0], hashes[1]}, StorageTier::kCpu)), + ""); + ASSERT_EQ(table.StoreShared(Shared(hashes, StorageTier::kDisk)), ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 48); + EXPECT_EQ(result.gpu, 16); + EXPECT_EQ(result.dp, (std::map{{0, 16}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(16, 32, 48)}})); + EXPECT_EQ(result.cpu, 32); + EXPECT_EQ(result.disk, 48); +} + +TEST(Query, LowerTierPhaseNeverReturnsToHigherTier) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(64); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 4u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0]})), ""); + ASSERT_EQ(table.StoreShared(Shared({hashes[2]}, StorageTier::kCpu)), ""); + ASSERT_EQ( + table.StoreShared(Shared({hashes[1], hashes[3]}, StorageTier::kDisk)), + ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 32); + EXPECT_EQ(result.gpu, 16); + EXPECT_EQ(result.dp, (std::map{{0, 16}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(16, 16, 32)}})); + EXPECT_EQ(result.cpu, 16); + EXPECT_EQ(result.disk, 32); +} + +TEST(Query, DiskMissIgnoresAllLaterIsolatedBlocks) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(80); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 5u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[4]})), ""); + ASSERT_EQ( + table.StoreShared(Shared({hashes[1], hashes[4]}, StorageTier::kCpu)), + ""); + ASSERT_EQ( + table.StoreShared(Shared({hashes[3], hashes[4]}, StorageTier::kDisk)), + ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 32); + EXPECT_EQ(result.gpu, 16); + EXPECT_EQ(result.dp, (std::map{{0, 16}})); + EXPECT_EQ( + result.rank_matches, + (std::map{{0, RankMatch(16, 32, 32)}})); + EXPECT_EQ(result.cpu, 32); + EXPECT_EQ(result.disk, 32); +} + +TEST(Query, DifferentRanksNeverFabricateOneGpuPrefix) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-a", 1)); + const auto tokens = Tokens(32); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 2u); + + ASSERT_EQ( + table.StoreGpu(Gpu({hashes[0]}, GpuOwner("instance-a", 0, "rank-0"))), + ""); + ASSERT_EQ( + table.StoreGpu(Gpu({hashes[1]}, GpuOwner("instance-a", 1, "rank-1"))), + ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 16); + EXPECT_EQ(result.gpu, 16); + EXPECT_EQ(result.dp, (std::map{{0, 16}, {1, 0}})); + EXPECT_EQ(result.rank_matches, + (std::map{{0, RankMatch(16, 16, 16)}, + {1, RankMatch(0, 0, 0)}})); + EXPECT_EQ(result.dp.size(), result.rank_matches.size()); + EXPECT_EQ(result.cpu, 16); + EXPECT_EQ(result.disk, 16); +} + +TEST(Query, InstanceSummaryIsRealizedByMaximumGpuRank) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-a", 1)); + const auto tokens = Tokens(64); + const auto hashes = Hashes(tokens); + ASSERT_EQ(hashes.size(), 4u); + + ASSERT_EQ(table.StoreGpu(Gpu({hashes[0], hashes[1]}, + GpuOwner("instance-a", 0, "rank-0"))), + ""); + ASSERT_EQ( + table.StoreGpu(Gpu({hashes[0]}, GpuOwner("instance-a", 1, "rank-1"))), + ""); + ASSERT_EQ(table.StoreShared(Shared({hashes[2]}, StorageTier::kCpu)), ""); + ASSERT_EQ(table.StoreShared(Shared({hashes[3]}, StorageTier::kDisk)), ""); + + const auto result = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(result.dp, (std::map{{0, 32}, {1, 16}})); + EXPECT_EQ(result.rank_matches, + (std::map{ + {0, RankMatch(32, 48, 64)}, {1, RankMatch(16, 16, 16)}})); + EXPECT_EQ(result.dp.size(), result.rank_matches.size()); + for (const auto& [rank, gpu] : result.dp) { + ASSERT_TRUE(result.rank_matches.contains(rank)); + EXPECT_EQ(gpu, result.rank_matches.at(rank).gpu); + } + EXPECT_EQ(result.gpu, result.rank_matches.at(0).gpu); + EXPECT_EQ(result.cpu, result.rank_matches.at(0).cpu); + EXPECT_EQ(result.disk, result.rank_matches.at(0).disk); + EXPECT_EQ(result.longest_match_tokens, result.rank_matches.at(0).disk); +} + +TEST(Query, RegisteredZeroHitRanksAndIncompleteTailAreRetained) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-a", 2)); + const auto incomplete_tokens = Tokens(31); + + const auto results = table.Query(TestContext(), incomplete_tokens); + ASSERT_EQ(results.size(), 1u); + const auto& result = results.at("instance-a"); + EXPECT_EQ(result.longest_match_tokens, 0); + EXPECT_EQ(result.dp, (std::map{{0, 0}, {2, 0}})); + EXPECT_EQ(result.rank_matches, + (std::map{{0, RankMatch(0, 0, 0)}, + {2, RankMatch(0, 0, 0)}})); + EXPECT_EQ(result.gpu, 0); + EXPECT_EQ(result.cpu, 0); + EXPECT_EQ(result.disk, 0); +} + +TEST(Query, InstanceFilterAndUnknownContextAreLookupOnly) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-b", 1)); + const auto before = PrefixCacheTableTestPeer::Snapshot(table); + + const auto filtered = + table.Query(TestContext(), Tokens(16), std::nullopt, "instance-b"); + ASSERT_EQ(filtered.size(), 1u); + EXPECT_TRUE(filtered.contains("instance-b")); + + EXPECT_TRUE( + table.Query(TestContext(), Tokens(16), std::nullopt, "unknown-instance") + .empty()); + ContextKey unknown = TestContext(); + unknown.model_name = "missing"; + EXPECT_TRUE(table.Query(unknown, Tokens(16)).empty()); + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table), before); +} + +TEST(Query, CacheSaltChangesHashesWithoutChangingContextIdentity) { + PrefixCacheTable table; + RegisterOrFail(table, Registration()); + const auto tokens = Tokens(16); + const auto unsalted = Hashes(tokens); + ASSERT_EQ(table.StoreGpu(Gpu(unsalted)), ""); + + const auto hit = table.Query(TestContext(), tokens).at("instance-a"); + EXPECT_EQ(hit.longest_match_tokens, 16); + + const auto salted = + table.Query(TestContext(), tokens, std::string("request-salt")) + .at("instance-a"); + EXPECT_EQ(salted.longest_match_tokens, 0); + EXPECT_EQ(table.GetGlobalView().context_count, 1); +} + +TEST(GlobalView, ReportsProfileRegistrationAndOwnerMapSize) { + PrefixCacheTable table; + RegisterOrFail(table, Registration("instance-a", 0)); + RegisterOrFail(table, Registration("instance-b", 1)); + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(1), Prefix(2)})), ""); + + const auto view = table.GetGlobalView(); + ASSERT_EQ(view.context_count, 1); + ASSERT_EQ(view.contexts.size(), 1u); + EXPECT_EQ(view.contexts[0].context, TestContext()); + EXPECT_EQ(view.contexts[0].profile, TestProfile()); + EXPECT_EQ(view.contexts[0].instance_ranks.at("instance-a"), + (std::set{0})); + EXPECT_EQ(view.contexts[0].instance_ranks.at("instance-b"), + (std::set{1})); + EXPECT_EQ(view.contexts[0].prefix_count, 2u); +} + +TEST(Capacity, EvictsOldestWrittenPrefixesWhenOverLimit) { + PrefixCacheTable table(10); + RegisterOrFail(table, Registration()); + + for (uint64_t i = 1; i <= 14; ++i) { + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(i)})), ""); + } + + const auto snapshot = PrefixCacheTableTestPeer::Snapshot(table); + const auto& blocks = snapshot.contexts.at(TestContext()).blocks; + EXPECT_LE(blocks.size(), 10u); + EXPECT_TRUE(blocks.contains(Prefix(14))); + EXPECT_FALSE(blocks.contains(Prefix(1))); +} + +TEST(Capacity, UnlimitedWhenBlockLimitIsZero) { + PrefixCacheTable table(0); + RegisterOrFail(table, Registration()); + for (uint64_t i = 1; i <= 50; ++i) { + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(i)})), ""); + } + EXPECT_EQ(PrefixCacheTableTestPeer::Snapshot(table) + .contexts.at(TestContext()) + .blocks.size(), + 50u); +} + +TEST(Capacity, OrderTrackingStaysInSyncWithBlocks) { + PrefixCacheTable table(10); + RegisterOrFail(table, Registration()); + const ContextKey context = TestContext(); + + for (uint64_t i = 1; i <= 6; ++i) { + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(i)})), ""); + } + auto sizes = PrefixCacheTableTestPeer::Order(table, context); + EXPECT_EQ(sizes.blocks, 6u); + EXPECT_EQ(sizes.write_order, 6u); + EXPECT_EQ(sizes.order_pos, 6u); + + for (uint64_t i = 1; i <= 3; ++i) { + ASSERT_EQ(table.RemoveGpu(Gpu({Prefix(i)})), ""); + } + sizes = PrefixCacheTableTestPeer::Order(table, context); + EXPECT_EQ(sizes.blocks, 3u); + EXPECT_EQ(sizes.write_order, 3u); + EXPECT_EQ(sizes.order_pos, 3u); + + ASSERT_EQ(table.ClearGpu(ClearFor()), ""); + sizes = PrefixCacheTableTestPeer::Order(table, context); + EXPECT_EQ(sizes.blocks, 0u); + EXPECT_EQ(sizes.write_order, 0u); + EXPECT_EQ(sizes.order_pos, 0u); + + for (uint64_t i = 20; i <= 40; ++i) { + ASSERT_EQ(table.StoreGpu(Gpu({Prefix(i)})), ""); + } + sizes = PrefixCacheTableTestPeer::Order(table, context); + EXPECT_LE(sizes.blocks, 10u); + EXPECT_EQ(sizes.write_order, sizes.blocks); + EXPECT_EQ(sizes.order_pos, sizes.blocks); + EXPECT_GT(sizes.evicted_by_capacity, 0); +} + +} // namespace diff --git a/mooncake-conductor/tests/prefix_indexer_test_peer.h b/mooncake-conductor/tests/prefix_indexer_test_peer.h new file mode 100644 index 0000000000..5903fe7bee --- /dev/null +++ b/mooncake-conductor/tests/prefix_indexer_test_peer.h @@ -0,0 +1,115 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "conductor/prefixindex/prefix_indexer.h" + +namespace mooncake::conductor::prefixindex { + +struct BlockPresenceSnapshot { + std::set gpu_owners; + std::set cpu_owners; + std::set disk_owners; + + bool operator==(const BlockPresenceSnapshot&) const = default; +}; + +struct ContextStateSnapshot { + HashProfile profile; + std::map> instance_ranks; + std::unordered_map blocks; + + bool operator==(const ContextStateSnapshot&) const = default; +}; + +struct PrefixCacheTableSnapshot { + std::unordered_map contexts; + + bool operator==(const PrefixCacheTableSnapshot&) const = default; +}; + +class PrefixCacheTableTestPeer { + public: + static bool ContextExists(const PrefixCacheTable& table, + const ContextKey& context) { + return table.LoadContextState(context) != nullptr; + } + + static std::unique_lock LockContextState( + const PrefixCacheTable& table, const ContextKey& context) { + auto state = table.LoadContextState(context); + if (state == nullptr) return {}; + return std::unique_lock(state->mutex); + } + + static std::optional Presence( + const PrefixCacheTable& table, const ContextKey& context, + ProjectedPrefix prefix) { + auto state = table.LoadContextState(context); + if (state == nullptr) return std::nullopt; + + std::shared_lock state_lock(state->mutex); + const auto block = state->blocks.find(prefix); + if (block == state->blocks.end()) return std::nullopt; + return BlockPresenceSnapshot{block->second.gpu_owners, + block->second.cpu_owners, + block->second.disk_owners}; + } + + // Snapshot sizes for validating order metadata invariants. + struct OrderSizes { + size_t write_order = 0; + size_t order_pos = 0; + size_t blocks = 0; + int64_t evicted_by_capacity = 0; + }; + + static OrderSizes Order(const PrefixCacheTable& table, + const ContextKey& context) { + auto state = table.LoadContextState(context); + if (state == nullptr) return {}; + std::shared_lock state_lock(state->mutex); + return {.write_order = state->write_order.size(), + .order_pos = state->order_pos.size(), + .blocks = state->blocks.size(), + .evicted_by_capacity = state->evicted_by_capacity}; + } + + static PrefixCacheTableSnapshot Snapshot(const PrefixCacheTable& table) { + PrefixCacheTableSnapshot snapshot; + std::vector>> + contexts; + { + std::shared_lock map_lock(table.context_map_mutex_); + contexts.reserve(table.contexts_.size()); + for (const auto& item : table.contexts_) { + contexts.push_back(item); + } + } + + for (const auto& [context, state] : contexts) { + std::shared_lock state_lock(state->mutex); + ContextStateSnapshot state_snapshot; + state_snapshot.profile = state->profile; + state_snapshot.instance_ranks = state->instance_ranks; + for (const auto& [prefix, presence] : state->blocks) { + state_snapshot.blocks.emplace( + prefix, BlockPresenceSnapshot{presence.gpu_owners, + presence.cpu_owners, + presence.disk_owners}); + } + snapshot.contexts.emplace(context, std::move(state_snapshot)); + } + return snapshot; + } +}; + +} // namespace mooncake::conductor::prefixindex From 6de1beefdc47982a436e8700b3851d9a12bd5a8a Mon Sep 17 00:00:00 2001 From: Misak2333 <167268798+Misak2333@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:47:12 +0800 Subject: [PATCH 2/2] [bugfix] Fix formatting of SharedObjectOwner return statement --- mooncake-conductor/tests/prefix_indexer_test.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mooncake-conductor/tests/prefix_indexer_test.cpp b/mooncake-conductor/tests/prefix_indexer_test.cpp index 435c200ccb..57eccc325c 100644 --- a/mooncake-conductor/tests/prefix_indexer_test.cpp +++ b/mooncake-conductor/tests/prefix_indexer_test.cpp @@ -94,8 +94,7 @@ EngineOwner GpuOwner(const std::string& instance_id = "instance-a", SharedObjectOwner SharedOwner(const std::string& object_id = "object-a", const std::string& stream = "pool-stream", const std::string& backend = "backend-a") { - return { - .source_stream = stream, .backend_id = backend, .object_id = object_id}; + return {.source_stream = stream, .backend_id = backend, .object_id = object_id}; } ProjectedPrefix Prefix(uint64_t value) { return {.value = value}; }