[Store] Python store API supports batch operation#511
[Store] Python store API supports batch operation#511xinranwang17 wants to merge 2 commits intokvcache-ai:mainfrom
Conversation
1. expose batch operation via python api 2. add BatchIsExist support
|
Since the PR in SGLang requires further refactoring, we can proceed with reviewing and merging this PR in Mooncake as a preparatory step. |
Thx! |
| * @param keys Key to check | ||
| * @return Map of keys to booleans | ||
| */ | ||
| std::unordered_map<std::string, bool> BatchIsExist( |
There was a problem hiding this comment.
I’d prefer to return `vector that matches up one-to-one with the keys in the input parameters.
There was a problem hiding this comment.
According to SGLang integration, is it better to return <key, is_exist> pair so that the existed key can be easily appended instead of addressing them by index again.
There was a problem hiding this comment.
I was looking through the sglang code and noticed this part:
mooncake_exist_keys = self.mooncake_l3_kv_pool.is_batch_exist(
fragment_keys
)
non_exist_keys = []
non_exist_value = []
for i in range(len(fragment_keys)):
if not mooncake_exist_keys[fragment_keys[i]]:I think we could simplify it to:
mooncake_exist_keys = self.mooncake_l3_kv_pool.is_batch_exist(
fragment_keys
)
non_exist_keys = []
non_exist_value = []
for i in range(len(fragment_keys)):
if mooncake_exist_keys[i] == 1:This way we're using direct indexing instead of querying the hash map, which should be slightly more efficient.
There was a problem hiding this comment.
LMCache needs same interface, so I’ve submitted a separate PR for batch exist: #542.
Appreciate it if you could take a look. Thanks!
There was a problem hiding this comment.
great work! Since array is more efficient, let's use array as the return value of batch exist interface.
| * @param keys Keys to check | ||
| * @return existence map, true if exists, false if not | ||
| */ | ||
| std::unordered_map<std::string, ErrorCode> BatchIsExist( |
| test_client_->BatchIsExist(keys); | ||
| end = std::chrono::high_resolution_clock::now(); | ||
| LOG(INFO) << "Time taken for BatchIsExist: " | ||
| << std::chrono::duration_cast<std::chrono::microseconds>(end - |
There was a problem hiding this comment.
Could you share a quick performance result here?
|
Since we have had an agreement on using array as the return value of batch exist, this PR will be accept. I'd like to continue contribute other implementation, such as batch put/get interface, in a separate PR. |
|
hi @xinranwang17 .When can this pr be merge. thanks |
|
hi @xiaguan When does lmcache have such a batch interface. thanks |
Part of this PR has been merged here: #542 |
After this LMCache/LMCache#924 got merged, I will submit mooncake's impl |
could you creat a PR about this so that i can preview. thanks. |
|
|
since #556 is merged, close this pr |
This PR introduces batch APIs (put_batch, get_batch, is_batch_exist) in the Python module to accelerate data transfer. These APIs are concurrently utilized in the Mooncake store implementation as the L3 cache in SGLang.
Related PR in sglang