[Store][POC] Add TTL-bounded group retention leases#2835
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a group retention feature (RetainGroups) across the Mooncake store components, including the Python bindings, client, master client, and master service. This feature allows users to retain objects within specific groups for a requested TTL to prevent eviction, supported by new configuration limits and background pruning of expired retentions. The review feedback highlights three key areas for improvement: optimizing the RetainGroups implementation by replacing an inefficient O(N) full scan of retention deadlines with lazy O(1) pruning, ensuring the Python wrapper correctly returns negative error codes on failure, and validating empty group IDs early to avoid redundant shard routing and locking.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| for (auto it = tenant_state.group_retention_deadlines.begin(); | ||
| it != tenant_state.group_retention_deadlines.end();) { | ||
| if (it->second <= now) { | ||
| it = tenant_state.group_retention_deadlines.erase(it); | ||
| retained_group_count_.fetch_sub(1, std::memory_order_relaxed); | ||
| } else { | ||
| ++it; | ||
| } | ||
| } | ||
| for (const size_t i : group_indices) { | ||
| if (group_ids[i].empty()) { | ||
| results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS); | ||
| continue; | ||
| } | ||
| const auto deadline = now + std::chrono::milliseconds(ttl_ms); | ||
| auto retention_it = | ||
| tenant_state.group_retention_deadlines.find(group_ids[i]); | ||
| if (retention_it == tenant_state.group_retention_deadlines.end()) { |
There was a problem hiding this comment.
Scanning the entire group_retention_deadlines map of the tenant on every RetainGroups call is highly inefficient. With max_retained_groups defaulting to 65,536, this
Since the background thread TaskCleanupThreadFunc already periodically calls PruneExpiredGroupRetentions() to clean up expired entries globally, we do not need to perform a full scan here. Instead, we can lazily prune only the specific group being requested if it has expired, which is an
for (const size_t i : group_indices) {
if (group_ids[i].empty()) {
results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
continue;
}
const auto deadline = now + std::chrono::milliseconds(ttl_ms);
auto retention_it =
tenant_state.group_retention_deadlines.find(group_ids[i]);
if (retention_it != tenant_state.group_retention_deadlines.end() &&
retention_it->second <= now) {
tenant_state.group_retention_deadlines.erase(retention_it);
retained_group_count_.fetch_sub(1, std::memory_order_relaxed);
retention_it = tenant_state.group_retention_deadlines.end();
}
if (retention_it == tenant_state.group_retention_deadlines.end()) {| if (!self.real_client_) { | ||
| return std::vector<int>( | ||
| group_ids.size(), | ||
| static_cast<int>(ErrorCode::INVALID_PARAMS)); | ||
| } |
There was a problem hiding this comment.
Directly casting ErrorCode::INVALID_PARAMS to int via static_cast<int> might bypass the standard error-to-integer mapping (such as the toInt() helper used in RealClient), potentially returning a positive integer instead of a negative error code as documented.
To ensure consistency with the rest of the client and adhere to the 'negative error code on failure' contract, we should use a negative value or a proper conversion helper. If ErrorCode values are positive, we should negate it or use the appropriate mapping function.
| if (!self.real_client_) { | |
| return std::vector<int>( | |
| group_ids.size(), | |
| static_cast<int>(ErrorCode::INVALID_PARAMS)); | |
| } | |
| if (!self.real_client_) { | |
| return std::vector<int>( | |
| group_ids.size(), | |
| -static_cast<int>(ErrorCode::INVALID_PARAMS)); | |
| } |
| std::vector<std::vector<size_t>> indices_by_shard(kNumShards); | ||
| for (size_t i = 0; i < group_ids.size(); ++i) { | ||
| indices_by_shard[getShardIndex(group_ids[i])].push_back(i); | ||
| } |
There was a problem hiding this comment.
We can validate and filter out empty group_ids before distributing them into indices_by_shard. This avoids routing empty strings to shards, acquiring locks for them, and performing redundant lookups.
std::vector<std::vector<size_t>> indices_by_shard(kNumShards);
for (size_t i = 0; i < group_ids.size(); ++i) {
if (group_ids[i].empty()) {
results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
} else {
indices_by_shard[getShardIndex(group_ids[i])].push_back(i);
}
}![qoderai[bot]](https://avatars.githubusercontent.com/in/1396158?s=60&v=4){kind=small}
There was a problem hiding this comment.
👋 Review Summary
This POC for TTL-bounded group retention is thoughtfully integrated: it plumbs retain_groups through master, client, RPC, and the Python binding, and the new tests cover key behaviors like future-member protection, TTL expiry, admission bounds, and concurrent admission.
🛡️ Key Risks & Issues
- Global admission cap:
retained_group_count_is a single global counter enforced againstmax_retained_groups_across all tenants and shards (mooncake-store/src/master_service.cpp:2206). This matches the current design but means one noisy tenant can monopolize the retention budget and effectively block others from using group retention. For multi-tenant deployments, this is a fairness and operability concern rather than a correctness bug. - Retention accounting and lifecycle: Retentions are decremented in both
RegisterGroupMemberandPruneExpiredGroupRetentions, and the counter is reset inMetadataSerializer::Reset. The logic looks consistent, but given the mix of batch calls, background cleanup, and tenant eviction, it would be good to have more explicit tests stressing multi-tenant scenarios and ensuring the counter never goes negative or leaks.
🧪 Verification Advice
- Multi-tenant simulation: Before rolling this out broadly, consider running a workload where multiple tenants issue
retain_groupscalls with different TTLs and group counts, and verify that the observed behavior (who gets accepted vs rejected) matches the intended “global pool” semantics. - Capacity and expiry behavior: In a staging environment, exercise the configured
max_retained_groupsandmax_group_retention_ttl_msby:- Filling the retention pool, verifying that new groups start returning 0 while extensions on existing groups remain accepted.
- Allowing groups to expire and observing that eviction resumes as expected and that subsequent
retain_groupscalls are accepted again.
💡 Thoughts & Suggestions
- If multi-tenant fairness is a concern for your deployment, it may be worth documenting that
max_retained_groupsis currently a global cap and considering per-tenant limits or at least exposing telemetry that breaks down retained-group usage by tenant. - The Python binding and client-side APIs look consistent and ergonomic; once the POC stabilizes, adding end-to-end tests for
retain_groupsat those layers (including error-code mappings) will help keep the behavior stable as you iterate on the feature.
🤖 Generated by Qoder • View workflow run
| return results; | ||
| } | ||
|
|
||
| std::vector<tl::expected<bool, ErrorCode>> MasterService::RetainGroups( |
There was a problem hiding this comment.
The admission and accounting logic around retained_group_count_ and group_retention_deadlines looks correct for a single-tenant workload, but it may have surprising behavior in multi-tenant setups. RetainGroups uses a single global retained_group_count_ across all shards and tenants, and PruneExpiredGroupRetentions plus RegisterGroupMember both decrement that counter. This means a heavy tenant can fill max_retained_groups_ and effectively block other tenants from reserving groups, even if those tenants are otherwise light users, and there is no explicit per-tenant cap or fairness mechanism.
If the intended semantics are “global pool shared across tenants”, it could help to make that explicit in the configuration and docs, and consider whether we want to expose per-tenant limits or instrumentation so operators can see when one tenant is monopolizing group retention. At minimum, it might be worth adding tests that exercise RetainGroups across multiple tenants and validate that retained_group_count_ and group_retention_deadlines remain consistent when retentions are created and expire under concurrent load.
🤖 Generated by Qoder • Fix in Qoder
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
2 participants
This draft POC adds master-owned, TTL-bounded retention leases for logical Mooncake object groups and supports intent recorded before group members are published. It is the L3 ownership layer for the companion Dynamo and SGLang prompt-KV retention path, without pinning GPU/host cache or importing provider/session semantics into Mooncake.
Tracking: DYN-3309
Companion PRs: Dynamo #11534, SGLang #30796, and frontend-crates #107. Parent design: SGLang RFC #27574.
How This Was Implemented
RetainGroups(group_ids, ttl_ms)through the master RPC, client interfaces, Store client, RealClient, and Python binding.Walkthrough
Mental model
Mooncake receives only exact logical group IDs and TTLs. Dynamo owns provider normalization, SGLang owns prompt-to-page resolution, and Mooncake owns the resulting L3 lease.
flowchart LR A["OpenAI prompt_cache_options or Anthropic cache_control"] --> B["Dynamo KvHintEnvelope"] B --> C["SGLang resolves committed HiCache page hashes"] C --> D["Mooncake retain_groups(group_ids, ttl_ms)"] D --> E["Master extends current leases"] D --> F["Master records future-member intent"] E --> G["Grouped eviction skips retained objects"] F --> GState and ordering
The master stores one retention deadline per logical group. Repeated calls take the maximum of the existing deadline and
now + ttl, and expired entries are lazily removed as normal master operations encounter them.SGLang can finish radix insertion before its asynchronous write-through publication reaches Mooncake.
retain_groupstherefore accepts a new group with no members, and a later grouped put applies the remaining deadline to every newly published member. This closes the completion/publication race without SGLang lease-renewal polling.Request lifecycle
retain_groupsrequest from HiCache's storage worker.New groups return
falsewhen the admission cap is full; invalid TTLs returnINVALID_PARAMS; partial acceptance is observable to SGLang while inference remains fail-open.Boundaries and limitations
Validation
PATH=/tmp/mooncake-clang-format:$PATH pre-commit run --from-ref origin/main --to-ref HEAD— passed all scoped hooks.cmake --build build-retention --target master_service_test -j2— passed.master_service_test --gtest_filter='MasterServiceTest.*Group*'— 24/24 passed, including future-member inheritance, expiry, bounds, and 32-thread concurrent admission.DYN_REQUEST_TRACE, and Tachometer — passed functional retention attribution.Benchmark Results
The current A/B seeded and probed on different SGLang workers, applied 70 × 16,384-token pressure requests, waited for drain, flushed worker-local caches, and observed
router_cached=0on every final probe.Retention survived measured eviction, but Mooncake materialization was slower than recompute on this single-node topology. The POC establishes ownership and correctness; it does not justify blanket retention policy enablement.
This draft POC adds a master-owned, TTL-bounded
retain_groupsoperation for logical Mooncake object groups. It lets agentic endpoints preserve explicitly retained prompt KV through ordinary store eviction without pinning GPU/host cache or polling Mooncake from SGLang.Related design discussion: SGLang programmatic KV cache RFC #27574. Companion Dynamo and SGLang PRs will be released soon.
How This Was Implemented
RetainGroups(group_ids, ttl_ms)through the Mooncake master, RPC client, Store client, RealClient, and Python binding.retain_until = max(existing_deadline, now + ttl), automatic expiry, and bounded admission with a 65,536-group default cap and 24-hour maximum TTL.falsewhen admission is full, invalid TTLs returnINVALID_PARAMS, and existing groups can still extend while full.Walkthrough
Mental model
Provider APIs express retention intent, Dynamo normalizes that intent, SGLang resolves the token prefix into exact HiCache page groups, and Mooncake owns the resulting L3 lease. Mooncake never receives provider, agent, session, or routing semantics.
flowchart LR A["OpenAI prompt_cache_retention or Anthropic cache_control"] --> B["Dynamo normalizes KvHintEnvelope"] B --> C["SGLang receives prefix_tokens and ttl_seconds"] C --> D["HiCache page-aligns the prefix and resolves group_ids"] D --> E["Mooncake retain_groups(group_ids, ttl_ms)"] E --> F["Master extends current leases and records future-member intent"] F --> G["Existing Mooncake eviction skips leased groups"]prompt_cache_keyalone is not treated as retention intent. The normalized envelope is carried after Dynamo routing/admission, so it does not change worker selection or queue policy.State and ordering
The master stores one deadline per logical group. Repeated requests use max semantics: if a group already has 60 minutes remaining, a later 5-minute request cannot shorten it.
SGLang applies the hint after finished-request radix insertion, but HiCache publication to L3 is asynchronous. Therefore
RetainGroupsaccepts a group even when it has no current members; objects published later inherit the remaining group deadline. This closes the request-completion/publication race without SGLang renewal polling.Request lifecycle
prefix_tokens + ttl_seconds.retain_groupscall records or extends the master-owned deadline and grants ordinary leases to current members.Admission is globally count-bounded and concurrency-safe. Extensions remain accepted at capacity; new groups return
false, allowing SGLang to log partial acceptance while inference continues.Boundaries and limitations
Validation
PATH=/tmp/mooncake-clang-format:$PATH pre-commit run --from-ref origin/main --to-ref HEAD— passed all hooks.cmake --build build-retention --target master_service_test -j2— passed on current upstreammain.master_service_test --gtest_filter='MasterServiceTest.*Group*'— 24/24 passed, including future-member inheritance, TTL expiry, bounded admission, and a 32-thread concurrent cap test.Benchmark Results
Fresh MiniMax M2.7 H100 validation used two TP4 workers, an 8 GiB-per-worker embedded Mooncake tier, a retained 10k-token prefix, and 60 seconds of c=8 unique-prompt pressure.
PR Classification
Checklist
AI Assistance Disclosure