[Store] Add put/get session APIs for ranged multi-buffer transfers

[ascend-direct-dev]

Description

WIP: add put/get session APIs so callers can reserve/query once, transfer object-byte ranges across multiple buffers, then finalize—without re-querying Master on every layer.

Related RFC: #2887 (Layerwise KV Cache Session-Based Ranged Transfer for Mooncake Store)

• Get: batch_get_session_start → batch_get_into_multi_buffer_ranges → batch_get_session_end
• Put: batch_put_session_start → batch_put_from_multi_buffer_ranges → batch_put_session_end / batch_put_session_revoke
• Get session caches a QueryResult with a single complete memory replica (FilterQueryResult); range paths only use the cache (no Master RPC)
• Object-range overflow checks share is_object_range_overflow with execute_ranged_read
• Put start reuses Client::StartBatchPut (same path as BatchPut); range transfers go through Client::BatchTransferReadRanges / BatchTransferWriteRanges → submitRangeRead / submitRangeWrite
• RealClient owns session state; transfer is delegated to Client (no TransferFuture leakage)

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test commands:

# Session unit tests (lease expiry, abnormal usage, ranged put/get)
./mooncake-store/tests/pybind_client_test --gtest_filter='*Session*'

# Real TCP transport e2e
bash mooncake-store/tests/e2e/run_session_ranges_tcp_e2e.sh

Test results:

• Unit tests pass (RealClientTest.TestPutGetSessionRanges, TestPutGetSessionAbnormal, TestGetSessionLeaseExpiredDropsSession)
• Integration tests pass (TCP e2e session_ranges_tcp_e2e.py PASSED)
• Manual testing done (layerwise multi-rank put/get example with per-layer timing + verification)

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue ([RFC]: Layerwise KV Cache Session-Based Ranged Transfer for Mooncake Store #2887)

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

Cursor agent assisted implementation and local validation. The human submitter has reviewed every changed line and can defend the change end-to-end.

[gemini-code-assist]

Code Review

This pull request introduces batch put and get session range APIs to support ranged reads and writes into multiple buffers using process-local sessions, optimizing data transfer by bypassing master RPCs. The review feedback identifies several critical issues that need to be addressed: an integer overflow vulnerability in submitRangeWrite during boundary checks, a concurrency bottleneck in RealClient::batch_put_start caused by executing network RPCs while holding a mutex, a lack of defensive checks against integer overflow in batch_get_into_multi_buffer_ranges, and potential out-of-bounds crashes in batch_put_end and batch_put_revoke due to unvalidated RPC response sizes.

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.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 71.96133% with 203 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-integration/store/store_py.cpp	0.00%	71 Missing ⚠️
mooncake-store/src/real_client.cpp	78.13%	61 Missing ⚠️
mooncake-store/src/client_service.cpp	64.95%	41 Missing ⚠️
mooncake-store/src/transfer_task.cpp	54.28%	16 Missing ⚠️
mooncake-store/include/pyclient.h	0.00%	14 Missing ⚠️

📢 Thoughts on this report? Let us know!

[zxpdemonio]

Thanks for the RFC and prototype. I think the session-based ranged transfer direction makes sense for layerwise KV cache load/save, especially for avoiding repeated Master queries after batch_get_start.

One concern is that the new batch_get_into_multi_buffer_ranges path seems to maintain a separate ranged-read implementation from the existing get_into_ranges / get_into_ranges_internal path. Since get_into_ranges_internal already has mature range validation, buffer-capacity checks, overflow handling, result mapping, and optional QueryResultCache support, could we consider factoring the common ranged-read machinery and reusing it for the session-based API? The session API can still keep its own lifecycle and cached replica/lease state, but the low-level range preparation and transfer execution should ideally be shared to avoid duplicated logic and behavior drift.

Relatedly, the PR currently introduces get_sessions_ as a session-level cache of replica and lease information, instead of reusing the existing request-scoped QueryResultCache. I agree that a long-lived session needs different lifecycle semantics from a per-request cache, but it would be helpful to clarify the relationship between these two mechanisms. For example, batch_get_start could populate session state, while the actual range transfer path reuses the same internal metadata/range execution helpers used by get_into_ranges. This would make the new API easier to review and reduce the chance of inconsistent error handling or lease/bounds behavior between the two ranged-read paths.

[ascend-direct-dev]

Thanks for the review — agreed on the direction.

What we reused / aligned:

1. Transfer path: session ranged get/put go through Client::BatchTransferReadRanges / BatchTransferWriteRanges, which call the existing SubmitRangeRead / SubmitRangeWrite → submitRangeRead / submitRangeWrite (same low-level RangeRead/Write as other paths).
2. QueryResult type: get_sessions_ now stores a QueryResult (same value type as QueryResultCache), built at batch_get_session_start via existing FilterQueryResult so the vector keeps only the selected complete memory replica + lease_timeout. Ranges use QueryResult::IsLeaseExpired.
3. Overflow check: object-byte range overflow is shared with execute_ranged_read via is_object_range_overflow.

What we intentionally do not call into:

• QueryResultCache as the session store: request-scoped cache vs long-lived start→end session; different miss policy (re-query vs fail). Same QueryResult type, different map lifecycle — session stays in get_sessions_.