[Bugfix] Preserve transient Store GET errors during HA rebuild.

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:**
 

mooncake-integration/store/store_py.cpp

@@ -5,6 +5,7 @@
 #include <functional>
 #include <limits>
 #include <numeric>
+#include <stdexcept>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -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<BufferHandle> 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<int>(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<pybind11::bytes> get_batch(
@@ -1804,6 +1861,13 @@ class MooncakeDistributedNoFRegisterPyWrapper {
 };
 
 PYBIND11_MODULE(store, m) {
+    auto store_error =
+        py::register_exception<StoreError>(m, "StoreError", PyExc_RuntimeError);
+    py::register_exception<StoreUnavailableError>(m, "StoreUnavailableError",
+                                                  store_error.ptr());
+    py::register_exception<StoreGetError>(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.");

mooncake-store/include/real_client.h

@@ -279,6 +279,16 @@ class RealClient : public PyClient {
      */
     std::shared_ptr<BufferHandle> 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<std::shared_ptr<BufferHandle>, 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<BufferHandle> get_buffer_internal(
         const std::string &key,
         const std::shared_ptr<ClientBufferAllocator> &client_buffer_allocator =
-            nullptr);
+            nullptr,
+        ErrorCode *failure_reason = nullptr);
 
     std::vector<std::shared_ptr<BufferHandle>> batch_get_buffer_internal(
         const std::vector<std::string> &keys,

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<void, ErrorCode> RealClient::unregister_shm_buffer_internal(
 // Implementation of get_buffer_internal method
 std::shared_ptr<BufferHandle> RealClient::get_buffer_internal(
     const std::string &key,
-    const std::shared_ptr<ClientBufferAllocator> &client_buffer_allocator) {
+    const std::shared_ptr<ClientBufferAllocator> &client_buffer_allocator,
+    ErrorCode *failure_reason) {
+    auto fail = [failure_reason](ErrorCode error) {
+        if (failure_reason) {
+            *failure_reason = error;
+        }
+        return std::shared_ptr<BufferHandle>{};
+    };
+    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::Descriptor> &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<BufferHandle> 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<BufferHandle> 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<BufferHandle> 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<BufferHandle> RealClient::get_buffer(const std::string &key) {
-    return execute_timed_operation<std::shared_ptr<BufferHandle>>(
-        [&]() { 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<std::shared_ptr<BufferHandle>, 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<std::tuple<uint64_t, size_t>, ErrorCode>

mooncake-store/tests/pybind_client_test.cpp

@@ -5,6 +5,7 @@
 #include <barrier>
 #include <chrono>
 #include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <fcntl.h>
 #include <limits>
@@ -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<const char>(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;