Repository navigation
[Store]Add async batch get interface - #1861
zhangzuo21 wants to merge 2 commits into
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 introduces a two-phase asynchronous batch get mechanism and an AsyncGetContext class for concurrent operations, including Python bindings. Feedback focuses on improving the performance of result re-indexing from O(N^2) to O(N) and utilizing C++20 features for cleaner code in the completion polling logic.
| for (size_t i = 0; i < num_keys; ++i) { | ||
| auto it = std::find(batch_to_orig.begin(), batch_to_orig.end(), i); | ||
| if (it != batch_to_orig.end()) { | ||
| size_t batch_idx = std::distance(batch_to_orig.begin(), it); | ||
| state.results[i] = batch_results[batch_idx]; | ||
| state.lease_timeouts[i] = batch_leases[batch_idx]; | ||
| } else { | ||
| state.results[i] = tl::unexpected(pre_results[i].error()); | ||
| } | ||
| } |
There was a problem hiding this comment.
This loop for re-indexing results has a potential performance issue. It iterates num_keys times and calls std::find within the loop, leading to a complexity that can approach O(N^2) for large batches, where N is num_keys. This can be inefficient.
A more performant O(N) approach would be to first initialize all results with their pre-check outcomes, and then iterate through the batch_to_orig vector once to update the results for the operations that were actually submitted.
// Initialize all results with their pre-determined outcomes from pre-checks.
for (size_t i = 0; i < num_keys; ++i) {
if (!pre_results[i]) {
state.results[i] = tl::unexpected(pre_results[i].error());
}
}
// Overwrite with results from the actual batch operation for keys that were submitted.
for (size_t batch_idx = 0; batch_idx < batch_to_orig.size(); ++batch_idx) {
size_t orig_idx = batch_to_orig[batch_idx];
state.results[orig_idx] = batch_results[batch_idx];
state.lease_timeouts[orig_idx] = batch_leases[batch_idx];
}| for (auto it = active_slots_.begin(); it != active_slots_.end(); ) { | ||
| bool all_ready = true; | ||
| for (auto& pt : it->state.pending_transfers) { | ||
| if (!pt.future.isReady()) { | ||
| all_ready = false; | ||
| break; | ||
| } | ||
| } | ||
|
|
||
| if (all_ready) { | ||
| ready_slots.push_back(Slot{it->token, std::move(it->state)}); | ||
| it = active_slots_.erase(it); | ||
| } else { | ||
| ++it; | ||
| } | ||
| } |
There was a problem hiding this comment.
For improved code clarity and maintainability, consider using the C++20 std::erase_if algorithm here. It simplifies the code for filtering and moving elements from active_slots_ to ready_slots_, making the intent clearer and reducing the risk of manual iterator errors.
std::erase_if(active_slots_, [&](Slot& slot) {
bool all_ready = true;
for (const auto& pt : slot.state.pending_transfers) {
if (!pt.future.isReady()) {
all_ready = false;
break;
}
}
if (all_ready) {
ready_slots.push_back({slot.token, std::move(slot.state)});
}
return all_ready;
});
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
This PR has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Please update the PR or comment to keep it open. |
|
This PR was closed because it has been inactive for 104 days. Feel free to reopen if you'd like to continue the work. |
Description
Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-pg)mooncake-rl)Type of Change
How Has This Been Tested?
Checklist
./scripts/code_format.shbefore submitting.Motivation
When sglang's HiCache prefetch IO thread fetches KV cache pages from Mooncake, it processes operations serially — one
batch_getmust fully complete before the next begins. This underutilizes RDMA/network bandwidth when multiple prefetch operations are queued.This PR adds a two-phase async
batch_getinterface that allows callers to submit multiple concurrent get operations and wait for completions in any order, enabling parallel RDMA transfers across different prefetch requests.Changes
Two-phase split of
BatchGet(Client layer)Client::BatchGetSubmit(): Performs metadata query, replica selection, hot cache redirect, and submits RDMA transfers viaTransferSubmitter::submit(). Returns aBatchGetStateholding the in-flightTransferFutures.Client::BatchGetComplete(): Waits for allTransferFutures, handles hot cache admission, lease expiry checks, and metrics.BatchGet(). The new methods are additive.AsyncGetContext (RealClient layer)
submit(keys, buffers, sizes) → token: Callsbatch_get_submit(synchronous but fast — metadata RPC + RDMA submit), stores state in an active slot. ReturnsINVALID_TOKENonly when max concurrency is reached. All key-level errors are deferred towait_any.wait_any() → (token, results): Blocks on a condition variable until any operation completes. Results contain per-key success (0) or failure (-1).TransferFuture::isReady()on active slots, moves ready slots out of the lock before callingBatchGetComplete, then enqueues results.submit()to prevent exceedingmax_concurrency.Python bindings (pybind11)
AsyncGetContextclass exposed withsubmit(),wait_any(),in_flightproperty.MooncakeDistributedStore.create_async_context(max_concurrency)factory method.Limitations
batch_get_intocontinues to handle SSD.Tests
Files changed
mooncake-store/include/client_service.hBatchGetStatestruct,BatchGetSubmit/BatchGetCompletedeclarationsmooncake-store/src/client_service.cppmooncake-store/include/real_client.hAsyncGetContextclass,RealClientnew methodsmooncake-store/src/real_client.cppAsyncGetContextimplementation,batch_get_submit/batch_get_completewrappersmooncake-integration/store/store_py.cppAsyncGetContextmooncake-store/tests/CMakeLists.txtmooncake-store/tests/async_batch_get_test.cppmooncake-store/tests/test_async_python.py