Repository navigation
[Store] Add session IDs to KV events for Dynamo session prefix indexer - #4683
Conversation
magicYang1573
changed the title
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
github-actions
Bot
added
documentation
stmatengss
left a comment
stmatengss
left a comment
There was a problem hiding this comment.
Reviewed the session-attribution changes at 180099d. I found three correctness issues, detailed inline: structured-object expansion resets unrelated write policies; filtered client batch retries retain the original session-ID vector; and invalid UTF-8 session IDs can prevent strict subscribers from decoding a whole event batch.
Validation: compiled an isolated native binding using this revision's ReplicateConfig definition and pybind declarations, exercised the production Python config helpers, and exercised the production dangling-disk batch-healing function with RPC test doubles and the actual session validator. Also reproduced the MessagePack strict-decoding failure. These are focused component reproductions, not a full Store build or live RPC/ZMQ/failover test. All 30 reported GitHub checks currently pass.
AI-assisted review: CodeRabbit raised one issue (UTF-8 validation), which I checked independently; the other two came from direct code review and reproduction.
| if key_count != 1 or isinstance(config.session_ids, str): | ||
| config = _copy_write_config(config) | ||
| config.session_ids = session_ids * key_count |
There was a problem hiding this comment.
[P2] Preserve all write options when expanding session IDs
This newly copies an otherwise ungrouped config whenever a logical object expands to multiple physical keys. The native pybind ReplicateConfig is not copyable with copy.copy, so _copy_write_config falls back to _WRITE_CONFIG_COPY_FIELDS, which omits dfs_replica_num, soft_pin_action, and soft_pin_ttl_ms (it still lists the obsolete with_soft_pin). I reproduced that a config with group_ids=None, dfs_replica_num=1, soft_pin_action=ENABLE, and soft_pin_ttl_ms=60000 keeps these settings without a session tag, but adding session_ids=['s'] and expanding to two keys changes them to 0, PRESERVE, and None. Merely adding attribution therefore silently removes requested DFS replication and eviction protection. Preserve these fields in the fallback copy, or expose a complete native copy operation, and cover this with the native config rather than only SimpleNamespace.
There was a problem hiding this comment.
Thanks for the detailed review and reproductions! Addressed all three findings and the documentation suggestion in 218299d:
- Preserve DFS replica and soft-pin settings when expanding session IDs.
- Keep session IDs aligned during batch filtering and disk-repair retries, and reject invalid original tag vectors before filtering.
- Reject invalid UTF-8 at the shared write-validation boundary.
- Clarify the documentation’s UTF-8 requirement and 256-byte limit.
Added regression coverage and confirmed the failures before the fixes. All 122 C++ and 34 Python test cases run for this update passed locally.
Please take another look when you have time. Thanks!
| } else if (!config.ValidSessionIds(keys.size())) { | ||
| results.assign(keys.size(), | ||
| tl::make_unexpected(ErrorCode::INVALID_PARAMS)); |
There was a problem hiding this comment.
[P2] Slice session IDs for the client’s filtered batch retries
The validation is needed, but not every caller keeps the new per-key vector aligned. Client::healDanglingLocalDiskBatchStarts (client_service.cpp) evicts stale LOCAL_DISK entries, builds a smaller retry_keys vector, and calls BatchPutStart(retry_keys, retry_slices, config) with the original config. For a tagged two-key batch containing one fresh/healthy key and one wiped disk-only key, the retry has one key but two session IDs, so this check returns INVALID_PARAMS after the stale replica has already been evicted. The untagged batch heals successfully. I reproduced this using the production healing function with RPC doubles and this validator. Index the session IDs by the original positions when constructing the retry config; also handle StartBatchPut/StartBatchUpsert filtering of checksum-failed operations. Extending BatchPutAfterPhysicalWipeHealsDanglingLocalDiskReplica with per-key tags would cover the retry regression.
| std::all_of(session_ids->begin(), session_ids->end(), | ||
| [](const auto& id) { return id.size() <= 256; })); |
There was a problem hiding this comment.
[P2] Reject invalid UTF-8 before publishing session IDs as text
Length validation accepts arbitrary bytes: the native Python setter accepts config.session_ids = [b'\xff'], and ValidSessionIds(1) returns true (C++ callers can supply the same bytes directly). PublishBatch then emits the ID as a MessagePack STR. With the strict decoding used by the new integration test (msgpack.unpackb(payload, raw=False)), that value raises UnicodeDecodeError, preventing the entire event batch, including unrelated valid events, from being decoded. I reproduced both native acceptance and the decoder failure. Validate UTF-8 alongside the byte limit at the shared write-validation boundary, and add an invalid-byte test; if arbitrary binary IDs are intended instead, define a matching wire/consumer encoding rather than emitting them as text.
| } | ||
| if (session_ids) { | ||
| key_config.session_ids = | ||
| std::vector<std::string>{session_ids->at(key_index)}; |
There was a problem hiding this comment.
[P2] Avoid copying the entire session-ID vector for every key
ForSingleKey() first executes ReplicateConfig key_config = *this, which deep-copies all N session IDs, and only then replaces that vector with this one entry. Both WrappedMasterService::BatchPutStart and MasterService::BatchUpsertStart call this once per key, so adding session attribution to an otherwise ungrouped batch introduces O(N²) string copies/allocations in the Master request path.
I reproduced this by extracting the production ReplicateConfig declaration/methods and using struct_pack::compatible from the project's pinned YLT revision. With GCC 13.3, -O2, no group IDs and 64-byte session IDs, three runs gave:
| Keys | Cumulative allocated bytes | Time for all ForSingleKey() calls |
|---|---|---|
| 1,024 | 101,877,760 | 16.6–18.8 ms |
| 4,096 | 1,628,053,504 (16,793,600 allocations) | 303–314 ms |
The same untagged configuration performed zero heap allocations in this loop. These are isolated config-splitting measurements, not end-to-end throughput results or peak-memory figures.
Please prepare a common configuration without the per-key batch vectors once outside the loop, then populate only the selected group/session entry for each key, or otherwise avoid copying the full vectors before slicing. Merely clearing the copied vector after key_config = *this still incurs the quadratic cost. A focused allocation/scaling check with tags present and group_ids absent would cover this regression.
AI-assisted review: OpenAI Codex; the reproduction does not constitute a full Store build or RPC/HA test.
There was a problem hiding this comment.
Thanks for catching this! Fixed in 9536f6b.
The Master batch handlers now copy the configuration once outside the loop and update only the selected key’s group/session IDs.
With 4,096 keys and 64-byte session IDs, cumulative allocation for config splitting dropped from 1.63 GB to 1.06 MB. Increasing the batch size 4× now increases allocation exactly 4×.
Added an allocation regression test. All 132 C++ and 5 native Python/RPC/ZMQ test cases passed locally.
Please take another look when you have time. Thanks!
…fix index Mooncake can tag stored objects with the writing session (ReplicateConfig.session_ids, kvcache-ai/Mooncake#4683) and publishes it as session_id on stored events. The HiCache shared cache already subscribes to those events; feed the tagged ones into the session prefix index so it knows how much of each session's prefix lives in the shared store, independent of GPU eviction. - The object key's page hash maps to SGLang's event block hash (first 64 bits, as hash_str_to_int64), so Mooncake stores land on the nodes the device events created. - Residency is recorded under one shared owner (worker id u64::MAX) in tier External, or Disk for disk replicas, reusing the parent the device events recorded. A block not yet known becomes a root that a later device event grafts. - Removals recede that tier only; a sequence gap or clear drops the owner. - The router hands its index to the shared cache through a default SharedKvCache hook. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
…4895) session_tag_test replaces operator new(size_t) and operator delete to count allocations. Under ASan the forms it does not replace come from ASan's operator new; gtest's std::stable_sort allocates its temporary buffer with nothrow new, which is then freed through the replaced sized delete (free), and ASan aborts with alloc-dealloc-mismatch. This has failed every nightly-coverage run since #4683. Replace the rest of the replaceable set (nothrow, array and nothrow delete forms) on malloc/free as well. The nothrow form counts bytes like operator new(size_t), matching what the default nothrow new does in a non-ASan build. Co-authored-by: apex-mochen <265981728+apex-mochen@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Description
Allow Dynamo's session prefix indexer to associate Mooncake stored KV events with the session that wrote the object. The C++ and Python write APIs accept optional per-key
ReplicateConfig.session_ids; Mooncake retains the write attribution and emits a singlesession_idstring (or nil) on stored events, matching SGLang's attribution model. Dynamo can use this event to update its session prefix index; tracking other sessions that share a block remains on the Dynamo side.The Master commits the tag when the write becomes readable. Duplicate Put/read operations retain the existing tag; Upsert with a nonempty tag replaces it, while absent/empty tags preserve it. Tags survive the existing metadata persistence and replication paths. Removed/cleared event schemas are unchanged, and RPC/metadata extensions retain serialization compatibility. Session IDs must be valid UTF-8, at most 256 bytes.
Batch filtering and disk-repair retries preserve key/tag alignment. Tensor and structured-object expansion propagate the writing session to physical keys; copying a structured write config preserves its DFS replica and soft-pin options. Invalid original tag vectors are rejected before filtering. Master batch handlers copy the configuration once and reuse it while selecting each key's IDs, avoiding quadratic copying of session/group vectors.
Related to #2819. This PR supplies write attribution and events for SPI integration; it does not introduce multi-session membership or implement the Dynamo consumer. Based on main
1d97a1df97d5277836e0e297b03289d7a39e708d.Module
mooncake-store)mooncake-integration)mooncake-wheel)Type of Change
How Has This Been Tested?
GCC 13.3, Python 3.12, local CPU/TCP,
ENABLE_KV_EVENTS=ON. Rebuilt the Master, native Python extension and affected C++ tests. The added regressions reproduce the review findings before the fixes.Allocation-fix validation (2026-10-19):
Growing the batch 4x now grows allocation exactly 4x. These are config-splitting allocations, not peak memory or end-to-end throughput measurements. The test enforces a linear allocation bound without timing thresholds.
All applicable pre-commit hooks for this update's five files pass without skip overrides. No CMake changes were needed; the earlier PR-wide CMake exception below remains documented.
Earlier correctness-fix validation (2026-10-18):
cmake-formatexception below.Earlier PR validation also covered Master groups, OpLog/standby/snapshots, repository HA tests and the disabled publisher stub. The earlier full structured-object suite had 172 passed, 2 skipped and 2 confirmed upstream failures (
_RawDestinationBuffer() takes no arguments); this review-fix round runs the affected subset above. This round does not repeat the full SGLang/Dynamo inference E2E or live multi-node/RDMA/DFS failure injection.Checklist
./scripts/code_format.shsession_idandpriority#2819)The full hook run showed that
cmake-formatrewrites unrelated pre-existing regions ofmooncake-store/tests/CMakeLists.txt. Those changes were discarded to keep the PR focused; the changed CMake line itself is unchanged by the formatter. All other applicable hooks pass withSKIP=cmake-format, andgit diff --checkpasses.AI Assistance Disclosure
OpenAI Codex assisted with implementation, review, documentation and test execution. The human submitter remains responsible for reviewing, understanding and defending the changes.
"$BUILD/mooncake-store/tests/file_storage_test"
--gtest_filter='FileStorageTest.AfterPhysicalWipe'
MOONCAKE_TEST_MASTER_BINARY="$BUILD/mooncake-store/src/mooncake_master"
python -m pytest -q mooncake-wheel/tests/test_session_kv_events.py
python -m pytest -q python/tests/store/test_structured_object_store.py
-k 'session or config or group'
(cd docs && make html)