From de56eea5787aa2a2e62c60bf688db96c7af96da0 Mon Sep 17 00:00:00 2001 From: "shenshuwei.3" Date: Sun, 15 Nov 2026 15:57:48 +0800 Subject: [PATCH] [Bugfix] Preserve transient Store GET errors during HA rebuild --- .../python-api-reference/mooncake-store.md | 16 ++- mooncake-integration/store/store_py.cpp | 108 ++++++++++++++---- mooncake-store/include/real_client.h | 13 ++- mooncake-store/src/real_client.cpp | 75 ++++++++---- mooncake-store/tests/pybind_client_test.cpp | 23 +++- 5 files changed, 186 insertions(+), 49 deletions(-) diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/python-api-reference/mooncake-store.md index c738f5ef1..865bd22a5 100644 --- a/docs/source/python-api-reference/mooncake-store.md +++ b/docs/source/python-api-reference/mooncake-store.md @@ -1172,9 +1172,23 @@ def get(self, key: str) -> bytes **Returns:** - `bytes`: Retrieved binary data +- Empty bytes if the key doesn't exist. **Raises:** -- Returns empty bytes if key doesn't exist +- `StoreUnavailableError` (a `RuntimeError` subclass) for transient Store + availability errors such as `RPC_FAIL`, `UNAVAILABLE_IN_CURRENT_STATUS`, or + `REPLICA_IS_NOT_READY`. The exception explicitly says that the key's + existence is unknown and that the result is not a confirmed cache miss; a + master HA rebuild or failover may be in progress. +- `StoreGetError` (a `RuntimeError` subclass) for other Store GET failures. +- GET does not automatically retry. A transient availability error is + reported immediately as `StoreUnavailableError`; the caller decides whether + to retry or fall back. + +Both exception types inherit from `StoreError`, which inherits from +`RuntimeError`, so existing `except RuntimeError` handlers continue to work. +Callers can catch `StoreUnavailableError` separately to distinguish an +unavailable/unknown key from a confirmed missing key. **Example:** diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index cf0353960..b3da89504 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -36,6 +37,30 @@ namespace py = pybind11; namespace mooncake { namespace { +// Keep the public Python exception hierarchy compatible with callers that +// already catch RuntimeError, while allowing callers to handle transient +// Store availability failures separately from other Store errors. +class StoreError : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +class StoreUnavailableError : public StoreError { + public: + using StoreError::StoreError; +}; + +class StoreGetError : public StoreError { + public: + using StoreError::StoreError; +}; + +bool is_transient_get_error(ErrorCode error) { + return error == ErrorCode::RPC_FAIL || + error == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS || + error == ErrorCode::REPLICA_IS_NOT_READY; +} + struct PyTensorInfo { uintptr_t data_ptr; size_t tensor_size; @@ -471,38 +496,70 @@ class MooncakeStorePyWrapper { pybind11::bytes get(const std::string &key) { if (!is_client_initialized()) { - LOG(ERROR) << "Client is not initialized"; - return pybind11::bytes("\\0", 0); + throw std::runtime_error( + "Mooncake Store client is not initialized"); } const auto kNullString = pybind11::bytes("\\0", 0); + std::shared_ptr buffer_handle; + ErrorCode get_error = ErrorCode::OK; { py::gil_scoped_release release_gil; - auto buffer_handle = store_->get_buffer(key); - if (!buffer_handle) { - py::gil_scoped_acquire acquire_gil; - return kNullString; + if (real_client_) { + auto result = real_client_->get_buffer_with_status(key); + if (result) { + buffer_handle = std::move(result.value()); + } else { + get_error = result.error(); + } + } else { + buffer_handle = store_->get_buffer(key); + if (!buffer_handle) { + const int exists = store_->isExist(key); + get_error = + exists == 0 + ? ErrorCode::OBJECT_NOT_FOUND + : (exists < 0 + ? ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS + : ErrorCode::INTERNAL_ERROR); + } } + } - py::gil_scoped_acquire acquire_gil; - auto runtime_accelerator = - mooncake::device::GetAcceleratorRegistry() - .RuntimeAccelerators(); - if (runtime_accelerator.FindDeviceForPointer( - buffer_handle->ptr())) { - std::string host_buf(buffer_handle->size(), '\0'); - if (!runtime_accelerator.CopyToHost(host_buf.data(), - buffer_handle->ptr(), - buffer_handle->size())) { - LOG(ERROR) << "Failed to copy buffer to host memory"; - return pybind11::none(); - } - return pybind11::bytes(host_buf); + if (!buffer_handle) { + if (get_error == ErrorCode::OBJECT_NOT_FOUND) { + return kNullString; + } + const auto error_details = + toString(get_error) + + " (code=" + std::to_string(static_cast(get_error)) + ")"; + if (is_transient_get_error(get_error)) { + throw StoreUnavailableError( + "Mooncake Store GET could not determine whether key '" + + key + "' exists: " + error_details + + ". This is not a confirmed OBJECT_NOT_FOUND; the key's " + "existence is unknown. A master HA rebuild or failover " + "may be in progress."); } - return pybind11::bytes((char *)buffer_handle->ptr(), - buffer_handle->size()); + throw StoreGetError("Mooncake Store GET for key '" + key + + "' failed: " + error_details + "."); + } + + auto runtime_accelerator = + mooncake::device::GetAcceleratorRegistry().RuntimeAccelerators(); + if (runtime_accelerator.FindDeviceForPointer(buffer_handle->ptr())) { + std::string host_buf(buffer_handle->size(), '\0'); + if (!runtime_accelerator.CopyToHost(host_buf.data(), + buffer_handle->ptr(), + buffer_handle->size())) { + LOG(ERROR) << "Failed to copy buffer to host memory"; + return pybind11::none(); + } + return pybind11::bytes(host_buf); } + return pybind11::bytes((char *)buffer_handle->ptr(), + buffer_handle->size()); } std::vector get_batch( @@ -1804,6 +1861,13 @@ class MooncakeDistributedNoFRegisterPyWrapper { }; PYBIND11_MODULE(store, m) { + auto store_error = + py::register_exception(m, "StoreError", PyExc_RuntimeError); + py::register_exception(m, "StoreUnavailableError", + store_error.ptr()); + py::register_exception(m, "StoreGetError", + store_error.ptr()); + m.def("_serialize_tensor", &serialize_tensor_metadata, "Inspect a torch tensor as Mooncake tensor metadata, data pointer, " "size, and owner."); diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 6b3382b3b..bb3ae056b 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -279,6 +279,16 @@ class RealClient : public PyClient { */ std::shared_ptr get_buffer(const std::string &key); + /** + * @brief Get a buffer and preserve the Store error on failure. + * + * Unlike get_buffer(), this distinguishes a missing object from transient + * master/rebuild errors. HA rebuild clients retry transient errors for the + * configured timeout before returning the last error. + */ + tl::expected, ErrorCode> + get_buffer_with_status(const std::string &key); + /** * @brief Get buffers containing the data for multiple keys (batch version) * @param keys Vector of keys to get data for @@ -664,7 +674,8 @@ class RealClient : public PyClient { std::shared_ptr get_buffer_internal( const std::string &key, const std::shared_ptr &client_buffer_allocator = - nullptr); + nullptr, + ErrorCode *failure_reason = nullptr); std::vector> batch_get_buffer_internal( const std::vector &keys, diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 314d7cf51..b42499c4e 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -46,7 +46,6 @@ DEFINE_bool(enable_http_server, false, DEFINE_int32(http_port, 9300, "Port for client HTTP server " "(only effective when --enable_http_server=true)."); - namespace mooncake { namespace { constexpr std::chrono::seconds kIpcRequestRecvTimeout{5}; @@ -2617,33 +2616,42 @@ tl::expected RealClient::unregister_shm_buffer_internal( // Implementation of get_buffer_internal method std::shared_ptr RealClient::get_buffer_internal( const std::string &key, - const std::shared_ptr &client_buffer_allocator) { + const std::shared_ptr &client_buffer_allocator, + ErrorCode *failure_reason) { + auto fail = [failure_reason](ErrorCode error) { + if (failure_reason) { + *failure_reason = error; + } + return std::shared_ptr{}; + }; + if (failure_reason) { + *failure_reason = ErrorCode::INTERNAL_ERROR; + } if (!client_) { LOG(ERROR) << "Client is not initialized"; - return nullptr; + return fail(ErrorCode::INVALID_PARAMS); } if (!client_buffer_allocator) { LOG(ERROR) << "Client buffer allocator is not provided"; - return nullptr; + return fail(ErrorCode::INVALID_PARAMS); } // Query the object info auto query_result = client_->Query(key); if (!query_result) { - if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND || - query_result.error() == ErrorCode::REPLICA_IS_NOT_READY) { - return nullptr; + const ErrorCode error = query_result.error(); + if (error != ErrorCode::OBJECT_NOT_FOUND) { + LOG(ERROR) << "Query failed for key: " << key + << " with error: " << toString(error); } - LOG(ERROR) << "Query failed for key: " << key - << " with error: " << toString(query_result.error()); - return nullptr; + return fail(error); } const std::vector &replica_list = query_result.value().replicas; if (replica_list.empty()) { LOG(ERROR) << "Empty replica list for key: " << key; - return nullptr; + return fail(ErrorCode::REPLICA_IS_NOT_READY); } // Select best replica: prefer local MEMORY, then any MEMORY, @@ -2654,21 +2662,21 @@ std::shared_ptr RealClient::get_buffer_internal( const auto *best_replica = SelectBestReplica(replica_list, local_endpoints); if (!best_replica) { LOG(ERROR) << "No usable replica for key: " << key; - return nullptr; + return fail(ErrorCode::REPLICA_IS_NOT_READY); } const auto &replica = *best_replica; uint64_t total_length = calculate_total_size(replica); if (total_length == 0) { - return nullptr; + return fail(ErrorCode::OBJECT_NOT_FOUND); } // Allocate buffer auto alloc_result = client_buffer_allocator->allocate(total_length); if (!alloc_result) { LOG(ERROR) << "Failed to allocate buffer for get_buffer, key: " << key; - return nullptr; + return fail(ErrorCode::INTERNAL_ERROR); } auto buffer_handle = @@ -2686,7 +2694,10 @@ std::shared_ptr RealClient::get_buffer_internal( if (!read_result) { LOG(ERROR) << "SSD read failed for key '" << key << "': " << toString(read_result.error()); - return nullptr; + return fail(read_result.error()); + } + if (failure_reason) { + *failure_reason = ErrorCode::OK; } return buffer_handle; } @@ -2711,22 +2722,38 @@ std::shared_ptr RealClient::get_buffer_internal( if (!get_result) { LOG(ERROR) << "Get failed for key: " << key << " with error: " << toString(get_result.error()); - return nullptr; + return fail(get_result.error()); } + if (failure_reason) { + *failure_reason = ErrorCode::OK; + } return buffer_handle; } // Implementation of get_buffer method std::shared_ptr RealClient::get_buffer(const std::string &key) { - return execute_timed_operation>( - [&]() { return get_buffer_internal(key, client_buffer_allocator_); }, - [](const auto &buffer) { return buffer != nullptr; }, - [&](uint64_t latency_us, const auto &buffer) { - client_->ObserveTransferOperation(TransferOperationKind::kRead, - "get_buffer", buffer->size(), - latency_us); - }); + auto result = get_buffer_with_status(key); + return result ? std::move(result.value()) : nullptr; +} + +tl::expected, ErrorCode> +RealClient::get_buffer_with_status(const std::string &key) { + if (!client_) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + const auto start_time = std::chrono::steady_clock::now(); + ErrorCode get_error = ErrorCode::INTERNAL_ERROR; + auto buffer = + get_buffer_internal(key, client_buffer_allocator_, &get_error); + if (buffer) { + client_->ObserveTransferOperation(TransferOperationKind::kRead, + "get_buffer", buffer->size(), + elapsed_us_since(start_time)); + return buffer; + } + return tl::make_unexpected(get_error); } tl::expected, ErrorCode> diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp index c08a48a8e..468832f76 100644 --- a/mooncake-store/tests/pybind_client_test.cpp +++ b/mooncake-store/tests/pybind_client_test.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -19,7 +20,6 @@ DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma"); - namespace mooncake { namespace testing { @@ -1140,6 +1140,27 @@ TEST_F(RealClientTest, ErrGetNonExistentKey) { } } +TEST_F(RealClientTest, GetBufferWithStatusDistinguishesUnavailableFromMissing) { + StartMasterAndSetupClient(); + auto missing = py_client_->get_buffer_with_status("missing_status_key"); + ASSERT_FALSE(missing); + EXPECT_EQ(missing.error(), ErrorCode::OBJECT_NOT_FOUND); + + const std::string key = "get_status_after_master_stop"; + const std::string value = "preserve the transient read error"; + ReplicateConfig config; + config.replica_num = 1; + ASSERT_EQ( + py_client_->put(key, std::span(value.data(), value.size()), + config), + 0); + + master_.Stop(); + auto unavailable = py_client_->get_buffer_with_status(key); + ASSERT_FALSE(unavailable); + EXPECT_NE(unavailable.error(), ErrorCode::OBJECT_NOT_FOUND); +} + // Get after remove TEST_F(RealClientTest, ErrGetAfterRemove) { const uint64_t kv_lease_ttl_ = 1;