From d808a7184a4bc54103489527fc272fa7099c6897 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 26 Sep 2026 04:00:43 +0800 Subject: [PATCH 1/5] [Bugfix][Store] Store the full byte range of typed Python buffers in put/upsert Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- mooncake-integration/store/store_py.cpp | 23 +++++++++++--- .../tests/test_distributed_object_store.py | 30 +++++++++++++++++++ 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 3fed10ade0..6fcf3f64c2 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -491,6 +491,21 @@ inline int to_py_ret(ErrorCode error_code) { return static_cast(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. +inline size_t contiguous_buffer_bytes(const py::buffer_info &info) { + py::ssize_t expected_stride = info.itemsize; + for (py::ssize_t i = info.ndim - 1; i >= 0; --i) { + if (info.strides[i] != expected_stride) { + throw std::runtime_error("buffer must be C-contiguous"); + } + expected_stride *= info.shape[i]; + } + return static_cast(info.size) * static_cast(info.itemsize); +} + #include "store_py_internal.h" } // namespace @@ -2862,7 +2877,7 @@ PYBIND11_MODULE(store, m) { return self.store_->upsert( key, std::span(static_cast(info.ptr), - static_cast(info.size)), + contiguous_buffer_bytes(info)), config); }, py::arg("key"), py::arg("value"), @@ -2918,7 +2933,7 @@ PYBIND11_MODULE(store, m) { infos.emplace_back(buf.request(/*writable=*/false)); const auto &info = infos.back(); spans.emplace_back(static_cast(info.ptr), - static_cast(info.size)); + contiguous_buffer_bytes(info)); } py::gil_scoped_release release; @@ -3126,7 +3141,7 @@ PYBIND11_MODULE(store, m) { return self.store_->put( key, std::span(static_cast(info.ptr), - static_cast(info.size)), + contiguous_buffer_bytes(info)), config); }, py::arg("key"), py::arg("value"), @@ -3175,7 +3190,7 @@ PYBIND11_MODULE(store, m) { infos.emplace_back(buf.request(/*writable=*/false)); const auto &info = infos.back(); spans.emplace_back(static_cast(info.ptr), - static_cast(info.size)); + contiguous_buffer_bytes(info)); } py::gil_scoped_release release; diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index bbb967ec6b..e633e9f8b8 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -178,6 +178,36 @@ 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_soft_pin_config_forwarding(self): """Test soft-pin action and TTL forwarding through Store operations.""" from mooncake.store import ReplicateConfig From 7351193c292a5dc2458977f38ae949ba523a8762 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 26 Sep 2026 08:23:48 +0800 Subject: [PATCH 2/5] test(store): satisfy ruff on the touched store test file Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- mooncake-wheel/tests/test_distributed_object_store.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index e633e9f8b8..f2394013c1 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -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 @@ -657,7 +656,7 @@ def worker(thread_id): 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}") From 0353baf2590d86018ccd8955ee8d319b2341f2c4 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 26 Sep 2026 11:56:18 +0800 Subject: [PATCH 3/5] test(store): strip the whitespace-only line next to the touched print Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- mooncake-wheel/tests/test_distributed_object_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index f2394013c1..f4360e7938 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -655,7 +655,7 @@ 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("\nConcurrent Stress Test Results:") print(f"Total threads: {NUM_THREADS}") print(f"Operations per thread: {OPERATIONS_PER_THREAD}") From af22755fcb6b139c806335ac7a169bfbad3e6ce4 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:39:47 +0800 Subject: [PATCH 4/5] [Bugfix][Store] Reject strided 1-D views in put_parts/upsert_parts Two review follow-ups: Use PyBuffer_IsContiguous instead of the hand-rolled stride walk. The buffer protocol leaves strides for extent-0/1 dimensions unspecified, so a manual walk can reject buffers that are contiguous in memory. put_parts/upsert_parts only checked ndim/itemsize, so a strided 1-D view (memoryview(b"abcdefgh")[::2]) passed and got stored as the wrong bytes. Add the contiguity check to both, plus a regression test. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- mooncake-integration/store/store_py.cpp | 27 ++++++++++++------- .../tests/test_distributed_object_store.py | 8 ++++++ 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 6fcf3f64c2..0799369093 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -495,13 +495,12 @@ inline int to_py_ret(ErrorCode error_code) { // (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) { - py::ssize_t expected_stride = info.itemsize; - for (py::ssize_t i = info.ndim - 1; i >= 0; --i) { - if (info.strides[i] != expected_stride) { - throw std::runtime_error("buffer must be C-contiguous"); - } - expected_stride *= info.shape[i]; + if (!PyBuffer_IsContiguous(info.view(), 'C')) { + throw std::runtime_error("buffer must be C-contiguous"); } return static_cast(info.size) * static_cast(info.itemsize); } @@ -2901,9 +2900,13 @@ PYBIND11_MODULE(store, m) { py::buffer buf = py::reinterpret_borrow(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(info.ptr), static_cast(info.size)); @@ -3161,9 +3164,13 @@ PYBIND11_MODULE(store, m) { py::buffer buf = py::reinterpret_borrow(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(info.ptr), static_cast(info.size)); diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index f4360e7938..88dcb99e20 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -207,6 +207,14 @@ def test_non_contiguous_buffer_rejected(self): 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_soft_pin_config_forwarding(self): """Test soft-pin action and TTL forwarding through Store operations.""" from mooncake.store import ReplicateConfig From 0838ebb710f675b7bcef004de3abb279f021c590 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 1 Oct 2026 20:28:51 +0800 Subject: [PATCH 5/5] test(store): pin acceptance of extent-1 dims with non-canonical strides The buffer protocol leaves strides for extent-0/1 dimensions unspecified, and PyBuffer_IsContiguous only checks strides where the extent is greater than 1. A (1, N) broadcast view with stride 0 on the size-1 dim is C-contiguous in memory, so put must accept it and store the full byte range; the hand-rolled stride walk this series replaced would have rejected it. Skipped when numpy is missing since run_tests.sh executes this file before numpy is installed. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../tests/test_distributed_object_store.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index 88dcb99e20..fb15ec0670 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -215,6 +215,31 @@ def test_parts_strided_view_rejected(self): 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