Skip to content

[Bugfix][Store] Store the full byte range of typed Python buffers in put/upsert #4327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
he-yufeng wants to merge 6 commits into from
Closed

[Bugfix][Store] Store the full byte range of typed Python buffers in put/upsert #4327

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d808a71
[Bugfix][Store] Store the full byte range of typed Python buffers in …
he-yufeng Sep 25, 2026
7351193
test(store): satisfy ruff on the touched store test file
he-yufeng Sep 26, 2026
0353baf
test(store): strip the whitespace-only line next to the touched print
he-yufeng Sep 26, 2026
af22755
[Bugfix][Store] Reject strided 1-D views in put_parts/upsert_parts
he-yufeng Sep 29, 2026
0838ebb
test(store): pin acceptance of extent-1 dims with non-canonical strides
he-yufeng Oct 1, 2026
177f756
Merge remote-tracking branch 'upstream/main' into fix/py-buffer-typed…
he-yufeng Oct 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 30 additions & 8 deletions mooncake-integration/store/store_py.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,20 @@ inline int to_py_ret(ErrorCode error_code) {
return static_cast<int>(error_code);
}

// pybind11's buffer `size` counts elements, not bytes, so a typed buffer
// (itemsize > 1) passed to put/upsert would be stored truncated while the
// call reports success (#4299). Store the full byte range of C-contiguous
// buffers and reject the rest loudly instead of writing the wrong bytes.
// PyBuffer_IsContiguous, not a hand-rolled stride walk: the buffer protocol
// leaves strides for extent-0/1 dimensions unspecified, so a manual walk
// rejects buffers that are contiguous in memory.
inline size_t contiguous_buffer_bytes(const py::buffer_info &info) {
if (!PyBuffer_IsContiguous(info.view(), 'C')) {
throw std::runtime_error("buffer must be C-contiguous");
}
return static_cast<size_t>(info.size) * static_cast<size_t>(info.itemsize);
}

#include "store_py_internal.h"

} // namespace
Expand Down Expand Up @@ -2870,7 +2884,7 @@ PYBIND11_MODULE(store, m) {
return self.store_->upsert(
key,
std::span<const char>(static_cast<char *>(info.ptr),
static_cast<size_t>(info.size)),
contiguous_buffer_bytes(info)),
config);
},
py::arg("key"), py::arg("value"),
Expand All @@ -2894,9 +2908,13 @@ PYBIND11_MODULE(store, m) {
py::buffer buf = py::reinterpret_borrow<py::buffer>(obj);
infos.emplace_back(buf.request(false));
const auto &info = infos.back();
if (info.ndim != 1 || info.itemsize != 1)
// A strided 1-D view (memoryview(b"..")[::2]) passes the
// ndim/itemsize check but is not contiguous in memory;
// without this it would be stored as the wrong bytes.
if (info.ndim != 1 || info.itemsize != 1 ||
!PyBuffer_IsContiguous(info.view(), 'C'))
throw std::runtime_error(
"parts must be 1-D bytes-like");
"parts must be 1-D C-contiguous bytes-like");

spans.emplace_back(static_cast<const char *>(info.ptr),
static_cast<size_t>(info.size));
Expand Down Expand Up @@ -2926,7 +2944,7 @@ PYBIND11_MODULE(store, m) {
infos.emplace_back(buf.request(/*writable=*/false));
const auto &info = infos.back();
spans.emplace_back(static_cast<const char *>(info.ptr),
static_cast<size_t>(info.size));
contiguous_buffer_bytes(info));

@Aionw Aionw Sep 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The same silent-wrong-bytes problem is still present in upsert_parts and put_parts (L2903–2909 and L3163–3169 in this revision; GitHub won't let me anchor there since those lines aren't in the diff). Their guard checks ndim == 1 && itemsize == 1 but not strides, so a strided 1-D view like memoryview(b"abcdefgh")[::2] passes and gets stored as 4 contiguous bytes starting at ptr.

Since the helper is already here, keeping the existing guard and swapping the length is enough:

spans.emplace_back(static_cast<const char *>(info.ptr),
                   contiguous_buffer_bytes(info));

With itemsize == 1 the byte count doesn't change; the only new behavior is the contiguity check. It would also be worth extending test_non_contiguous_buffer_rejected:

with self.assertRaises(RuntimeError):
    self.store.put_parts("test_non_contiguous_parts", view)
with self.assertRaises(RuntimeError):
    self.store.upsert_parts("test_non_contiguous_parts", view)

I think this belongs in the same PR rather than a follow-up: it's the same bug class, #4299 points at these guards as the reference behavior, and the PR description currently calls them "unchanged".

@he-yufeng he-yufeng Sep 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. Both put_parts and upsert_parts now also require PyBuffer_IsContiguous(info.view(), 'C'), so a strided 1-D view like memoryview(b"abcdefgh")[::2] is rejected instead of being stored as the wrong bytes. buf.request(false) passes PyBUF_STRIDES | PyBUF_FORMAT, so the strides in info.view() are real and the contiguity check actually sees them.

Added test_parts_strided_view_rejected next to test_non_contiguous_buffer_rejected covering both entry points.

}

py::gil_scoped_release release;
Expand Down Expand Up @@ -3134,7 +3152,7 @@ PYBIND11_MODULE(store, m) {
return self.store_->put(
key,
std::span<const char>(static_cast<char *>(info.ptr),
static_cast<size_t>(info.size)),
contiguous_buffer_bytes(info)),
config);
},
py::arg("key"), py::arg("value"),
Expand All @@ -3154,9 +3172,13 @@ PYBIND11_MODULE(store, m) {
py::buffer buf = py::reinterpret_borrow<py::buffer>(obj);
infos.emplace_back(buf.request(false));
const auto &info = infos.back();
if (info.ndim != 1 || info.itemsize != 1)
// A strided 1-D view (memoryview(b"..")[::2]) passes the
// ndim/itemsize check but is not contiguous in memory;
// without this it would be stored as the wrong bytes.
if (info.ndim != 1 || info.itemsize != 1 ||
!PyBuffer_IsContiguous(info.view(), 'C'))
throw std::runtime_error(
"parts must be 1-D bytes-like");
"parts must be 1-D C-contiguous bytes-like");

spans.emplace_back(static_cast<const char *>(info.ptr),
static_cast<size_t>(info.size));
Expand All @@ -3183,7 +3205,7 @@ PYBIND11_MODULE(store, m) {
infos.emplace_back(buf.request(/*writable=*/false));
const auto &info = infos.back();
spans.emplace_back(static_cast<const char *>(info.ptr),
static_cast<size_t>(info.size));
contiguous_buffer_bytes(info));
}

py::gil_scoped_release release;
Expand Down
68 changes: 65 additions & 3 deletions mooncake-wheel/tests/test_distributed_object_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import os
import time
import threading
import random
from mooncake.store import MooncakeDistributedStore, SoftPinAction

# The lease time of the kv object, should be set equal to
Expand Down Expand Up @@ -178,6 +177,69 @@ def test_basic_put_get_exist_operations(self):
time.sleep(default_kv_lease_ttl / 1000)
self.assertEqual(self.store.remove(key), 0)

def test_typed_buffer_roundtrip(self):
"""Non-byte buffers must store size * itemsize bytes, not just size."""
import array

key = "test_typed_buffer_key"
value = array.array("f", [float(i) for i in range(1024)])
expected = value.tobytes()

self.assertEqual(self.store.put(key, value), 0)
self.assertEqual(self.store.get_size(key), len(expected))
self.assertEqual(self.store.get(key), expected)

batch_keys = ["test_typed_batch_0", "test_typed_batch_1"]
batch_values = [
array.array("d", [1.5, 2.5, 3.5]),
array.array("i", [7, 8, 9, 10]),
]
self.assertEqual(self.store.put_batch(batch_keys, batch_values), 0)
for batch_key, batch_value in zip(batch_keys, batch_values):
self.assertEqual(self.store.get_size(batch_key), len(batch_value.tobytes()))
self.assertEqual(self.store.get(batch_key), batch_value.tobytes())

def test_non_contiguous_buffer_rejected(self):
"""Strided views must fail loudly instead of storing wrong bytes."""
view = memoryview(b"abcdefgh")[::2]
with self.assertRaises(RuntimeError):
self.store.put("test_non_contiguous_key", view)
with self.assertRaises(RuntimeError):
self.store.put_batch(["test_non_contiguous_batch"], [view])

def test_parts_strided_view_rejected(self):
"""put_parts/upsert_parts must reject strided 1-D views too (#4299)."""
view = memoryview(b"abcdefgh")[::2]
with self.assertRaises(RuntimeError):
self.store.put_parts("test_parts_strided_put", view)
with self.assertRaises(RuntimeError):
self.store.upsert_parts("test_parts_strided_upsert", view)

def test_extent1_dim_stride_accepted(self):
"""Strides of extent-0/1 dims are unspecified by the buffer protocol,
so buffers like a (1, N) broadcast view are C-contiguous and must be
accepted and stored in full."""
try:
import numpy as np
except ImportError:
self.skipTest("numpy not installed")

base = np.arange(4, dtype=np.uint8)
view = np.broadcast_to(base, (1, 4))
self.assertEqual(view.strides, (0, 1))

self.assertEqual(self.store.put("test_extent1_stride_key", view), 0)
self.assertEqual(self.store.get_size("test_extent1_stride_key"), 4)
self.assertEqual(self.store.get("test_extent1_stride_key"), base.tobytes())

typed = np.broadcast_to(np.array([1.5, 2.5], dtype=np.float32), (1, 2))
self.assertEqual(typed.strides, (0, 4))
self.assertEqual(self.store.put("test_extent1_typed_key", typed), 0)
self.assertEqual(
self.store.get("test_extent1_typed_key"),
np.array([1.5, 2.5], dtype=np.float32).tobytes(),
)

def test_soft_pin_config_forwarding(self):
"""Test soft-pin action and TTL forwarding through Store operations."""
from mooncake.store import ReplicateConfig
Expand Down Expand Up @@ -626,8 +688,8 @@ def worker(thread_id):
put_duration = system_stats['put_end'] - system_stats['put_start']
get_duration = system_stats['get_end'] - system_stats['get_start']
total_data_size_gb = (VALUE_SIZE * total_operations) / (1024**3)
print(f"\nConcurrent Stress Test Results:")

print("\nConcurrent Stress Test Results:")
print(f"Total threads: {NUM_THREADS}")
print(f"Operations per thread: {OPERATIONS_PER_THREAD}")
print(f"Total operations: {total_operations}")
Expand Down
Loading