Tighten checksum option flow

Author: charpty

kv_cache_manager/client/pybind/py_client_binding.cc

@@ -45,12 +45,14 @@ PYBIND11_MODULE(kvcm_py_client, module) {
         .value("ER_TRANSFERCLIENT_INIT_ERROR", kvcm::ClientErrorCode::ER_TRANSFERCLIENT_INIT_ERROR)
         .value("ER_MANAGERCLIENT_INIT_ERROR", kvcm::ClientErrorCode::ER_MANAGERCLIENT_INIT_ERROR)
         .value("ER_CLIENT_NOT_EXISTS", kvcm::ClientErrorCode::ER_CLIENT_NOT_EXISTS)
+        .value("ER_INIT_CHECK_BUFFER_ERROR", kvcm::ClientErrorCode::ER_INIT_CHECK_BUFFER_ERROR)
         .value("ER_SERVICE_NO_STATUS", kvcm::ClientErrorCode::ER_SERVICE_NO_STATUS)
         .value("ER_SERVICE_INTERNAL_ERROR", kvcm::ClientErrorCode::ER_SERVICE_INTERNAL_ERROR)
         .value("ER_SERVICE_UNSUPPORTED", kvcm::ClientErrorCode::ER_SERVICE_UNSUPPORTED)
         .value("ER_SERVICE_INVALID_ARGUMENT", kvcm::ClientErrorCode::ER_SERVICE_INVALID_ARGUMENT)
         .value("ER_SERVICE_DUPLICATE_ENTITY", kvcm::ClientErrorCode::ER_SERVICE_DUPLICATE_ENTITY)
         .value("ER_SERVICE_INSTANCE_NOT_EXIST", kvcm::ClientErrorCode::ER_SERVICE_INSTANCE_NOT_EXIST)
+        .value("ER_SERVICE_NOT_LEADER", kvcm::ClientErrorCode::ER_SERVICE_NOT_LEADER)
         .value("ER_SDK_TIMEOUT", kvcm::ClientErrorCode::ER_SDK_TIMEOUT)
         .value("ER_GETSDK_ERROR", kvcm::ClientErrorCode::ER_GETSDK_ERROR)
         .value("ER_CREATESDK_ERROR", kvcm::ClientErrorCode::ER_CREATESDK_ERROR)
@@ -69,6 +71,8 @@ PYBIND11_MODULE(kvcm_py_client, module) {
         .value("ER_CUDA_STREAM_SYNCHRONIZE_ERROR", kvcm::ClientErrorCode::ER_CUDA_STREAM_SYNCHRONIZE_ERROR)
         .value("ER_CUDA_STREAM_DESTROY_ERROR", kvcm::ClientErrorCode::ER_CUDA_STREAM_DESTROY_ERROR)
         .value("ER_CUDA_HOST_REGISTER_ERROR", kvcm::ClientErrorCode::ER_CUDA_HOST_REGISTER_ERROR)
+        .value("ER_CHECKSUM_MISMATCH", kvcm::ClientErrorCode::ER_CHECKSUM_MISMATCH)
+        .value("ER_INLINE_HEADER_INVALID", kvcm::ClientErrorCode::ER_INLINE_HEADER_INVALID)
         .finalize();
 
     py::native_enum<kvcm::MemoryType>(module, "MemoryType", "enum.Enum")

kv_cache_manager/client/src/internal/stub/grpc_stub.cc

@@ -409,12 +409,12 @@ ClientErrorCode GrpcStub::FinishWriteCache(const std::string &trace_id,
                                            const std::string write_session_id,
                                            const BlockMask &success_block,
                                            const Locations &locations,
-                                           const std::vector<int64_t> &checksums) {
+                                           const FinishWriteOptions &options) {
     auto stub = GET_AND_CHECK_STUB();
     proto::meta::FinishWriteCacheRequest request;
     SetCommonInfo(request, trace_id, instance_id);
     request.set_write_session_id(write_session_id);
-    for (auto checksum : checksums) {
+    for (auto checksum : options.checksums) {
         request.add_checksums(checksum);
     }
     ProtoConvert::BlockMaskToProto(success_block, request.mutable_success_blocks());

kv_cache_manager/client/src/internal/stub/grpc_stub.h

@@ -63,7 +63,7 @@ class GrpcStub : public Stub {
                                      const std::string write_session_id,
                                      const BlockMask &success_block,
                                      const Locations &locations,
-                                     const std::vector<int64_t> &checksums = {}) override;
+                                     const FinishWriteOptions &options = FinishWriteOptions{}) override;
 
     ClientErrorCode RemoveCache(const std::string &trace_id,
                                 const std::string &instance_id,

kv_cache_manager/client/src/internal/stub/stub.h

@@ -67,7 +67,7 @@ class Stub {
                                              const std::string write_session_id,
                                              const BlockMask &success_block,
                                              const Locations &locations,
-                                             const std::vector<int64_t> &checksums = {}) = 0;
+                                             const FinishWriteOptions &options = FinishWriteOptions{}) = 0;
 
     virtual ClientErrorCode RemoveCache(const std::string &trace_id,
                                         const std::string &instance_id,

kv_cache_manager/client/src/internal/stub/test/grpc_stub_test.cc

@@ -494,7 +494,12 @@ TEST_F(GrpcStubTest, TestFinishWriteCacheWithChecksums) {
         const std::vector<int64_t> checksums = {0x11, 0x22, 0x33, 0x44};
         ASSERT_EQ(ER_OK,
                   stub_->FinishWriteCache(
-                      "trace3", "instance1", write_session_id, success_block, unrelated_locations, checksums));
+                      "trace3",
+                      "instance1",
+                      write_session_id,
+                      success_block,
+                      unrelated_locations,
+                      FinishWriteOptions::WithChecksums(checksums)));
     }
     {
         auto [success, result] = stub_->GetCacheLocation("trace4",
@@ -525,7 +530,12 @@ TEST_F(GrpcStubTest, TestFinishWriteCacheRejectsChecksumSizeMismatch) {
     {
         BlockMask success_block = static_cast<size_t>(4);
         ASSERT_EQ(ER_SERVICE_INVALID_ARGUMENT,
-                  stub_->FinishWriteCache("trace3", "instance1", write_session_id, success_block, {}, {0x11}));
+                  stub_->FinishWriteCache("trace3",
+                                          "instance1",
+                                          write_session_id,
+                                          success_block,
+                                          {},
+                                          FinishWriteOptions::WithChecksums(std::vector<int64_t>{0x11})));
     }
 }
 

kv_cache_manager/client/src/meta_client_impl.cc

@@ -182,8 +182,7 @@ ClientErrorCode MetaClientImpl::FinishWrite(const std::string &trace_id,
                    DebugStringUtil::ToString(locations).c_str(),
                    options.checksums.size());
     const std::string &instance_id = CHECK_INSTANCE_STUB();
-    return stub_->FinishWriteCache(
-        trace_id, instance_id, write_session_id, success_block, locations, options.checksums);
+    return stub_->FinishWriteCache(trace_id, instance_id, write_session_id, success_block, locations, options);
 }
 
 ClientErrorCode MetaClientImpl::RemoveCache(const std::string &trace_id,

kv_cache_manager/client/src/transfer_client_impl.cc

@@ -292,14 +292,27 @@ ClientErrorCode TransferClientImpl::LoadKvCaches(const UriStrVec &uri_str_vec,
 #if defined(USING_CUDA) || defined(USING_MUSA)
     const auto &trace_info = options.trace_info;
     if (is_check_buffer_) {
-        bool need_print = (trace_info == nullptr) ? true : trace_info->need_print;
-        std::vector<int64_t> block_checksums;
+        const bool need_print = (trace_info == nullptr) ? true : trace_info->need_print;
         if (need_print) {
-            auto handle = sdk_buffer_check_pool_->GetCell();
-            block_checksums = SdkBufferCheckUtil::GetBlocksHash(
-                block_buffers, handle->d_iovs, handle->d_crcs, handle->h_iovs, max_check_iov_num_, handle->gpu_stream);
+            auto invalid_it = std::find_if(block_buffers.begin(), block_buffers.end(), [](const BlockBuffer &block) {
+                return !IsChecksumHashableBlock(block);
+            });
+            if (invalid_it != block_buffers.end()) {
+                const size_t idx = std::distance(block_buffers.begin(), invalid_it);
+                KVCM_LOG_WARN("block [%zu] has ignored, empty, null, or zero-size iovs; skip checksum print", idx);
+            } else if (sdk_buffer_check_pool_) {
+                auto handle = sdk_buffer_check_pool_->GetCell();
+                std::vector<int64_t> block_checksums;
+                if (HashBlocksByIovShape(block_buffers, handle, max_check_iov_num_, block_checksums)) {
+                    PrintBlockChecksumAndUri("get_", uri_str_vec, block_checksums, trace_info);
+                } else {
+                    KVCM_LOG_WARN("checksum print failed to hash blocks safely; skip checksum print");
+                }
+            } else {
+                KVCM_LOG_WARN("KVCM_SDK_CHECK is enabled but sdk_buffer_check_pool is not initialized; "
+                              "skip checksum print");
+            }
         }
-        PrintBlockChecksumAndUri("get_", uri_str_vec, block_checksums, trace_info);
     }
     // Read-side verification path: only kicks in when caller supplies expected checksums
     // (typically forwarded from CacheLocation.checksum returned by meta service).

kv_cache_manager/client/test/meta_client_test.cc

@@ -127,7 +127,7 @@ class MockStub : public Stub {
                  const std::string write_session_id,
                  const BlockMask &success_block,
                  const Locations &locations,
-                 const std::vector<int64_t> &checksums),
+                 const FinishWriteOptions &options),
                 (override));
 
     MOCK_METHOD(ClientErrorCode,
@@ -489,7 +489,7 @@ TEST_F(MetaClientTest, TestFinishWriteOptionsForwardToStub) {
                                  write_session_id,
                                  success_block,
                                  locations,
-                                 ::testing::IsEmpty()))
+                                 ::testing::Field(&FinishWriteOptions::checksums, ::testing::IsEmpty())))
         .Times(1)
         .WillOnce(::testing::Return(ER_OK));
 
@@ -502,7 +502,7 @@ TEST_F(MetaClientTest, TestFinishWriteOptionsForwardToStub) {
                                  write_session_id,
                                  success_block,
                                  locations,
-                                 ::testing::ElementsAre(0x44, 0)))
+                                 ::testing::Field(&FinishWriteOptions::checksums, ::testing::ElementsAre(0x44, 0))))
         .Times(1)
         .WillOnce(::testing::Return(ER_OK));
 

kv_cache_manager/manager/cache_manager.cc

@@ -701,7 +701,11 @@ CacheManager::StartWriteCache(RequestContext *request_context,
             RequestContext temp_request_context(trace_id + "_timeout_callback");
             BlockMaskOffset succeed_block = 0;
             auto ec = this->FinishWriteCache(
-                &temp_request_context, instance_id, write_session_id, succeed_block, std::move(write_location_info));
+                &temp_request_context,
+                instance_id,
+                write_session_id,
+                succeed_block,
+                CacheManager::FinishWriteCacheOptions::WithWriteLocationInfo(std::move(write_location_info)));
             static_cast<void>(ec);
         });
     KVCM_METRICS_COLLECTOR_CHRONO_MARK_END(service_metrics_collector, PutWriteLocationManager);
@@ -723,14 +727,13 @@ CacheManager::FinishWriteCache(RequestContext *request_context,
                                const std::string &instance_id,
                                const std::string &write_session_id,
                                const BlockMask &success_block_mask,
-                               std::unique_ptr<WriteLocationManager::WriteLocationInfo> write_location_info_internal,
-                               const std::vector<int64_t> &checksums) {
+                               CacheManager::FinishWriteCacheOptions options) {
     SPAN_TRACER(request_context);
     const std::string &trace_id = request_context->trace_id();
     auto *service_metrics_collector = dynamic_cast<ServiceMetricsCollector *>(request_context->metrics_collector());
     WriteLocationManager::WriteLocationInfo location_info;
-    if (write_location_info_internal != nullptr) {
-        location_info = std::move(*write_location_info_internal);
+    if (options.write_location_info_internal != nullptr) {
+        location_info = std::move(*options.write_location_info_internal);
     } else if (!write_location_manager_->GetAndDelete(write_session_id, location_info)) {
         request_context->error_tracer()->AddErrorMsg("write_session_id has been deleted");
         RETURN_IF_EC_NOT_OK_WITH_LOG(
@@ -746,12 +749,12 @@ CacheManager::FinishWriteCache(RequestContext *request_context,
     // checksums must be the same length as keys (full batch, not the success subset).
     // An empty vector means the client did not report any; existing CacheLocation
     // checksums are then preserved. Any length mismatch is treated as a client bug.
-    const bool has_checksums = !checksums.empty();
-    if (has_checksums && checksums.size() != location_info.keys.size()) {
+    const bool has_checksums = !options.checksums.empty();
+    if (has_checksums && options.checksums.size() != location_info.keys.size()) {
         RETURN_IF_EC_NOT_OK_WITH_LOG(WARN,
                                      EC_BADARGS,
                                      "checksums size (%zu) does not match keys size (%zu)",
-                                     checksums.size(),
+                                     options.checksums.size(),
                                      location_info.keys.size());
     }
 
@@ -771,7 +774,7 @@ CacheManager::FinishWriteCache(RequestContext *request_context,
             MetaSearcher::LocationUpdateTask task{
                 .location_id = location_info.location_ids[block_key_idx],
                 .new_status = CacheLocationStatus::CLS_SERVING,
-                .checksum = has_checksums ? checksums[block_key_idx] : int64_t{0},
+                .checksum = has_checksums ? options.checksums[block_key_idx] : int64_t{0},
             };
             success_batch_update_tasks.push_back({std::move(task)});
         } else {

kv_cache_manager/manager/cache_manager.h

@@ -5,6 +5,7 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <utility>
 #include <vector>
 
 #include "kv_cache_manager/common/error_code.h"
@@ -144,17 +145,34 @@ class CacheManager {
                                                               const std::vector<std::string> &location_spec_group_names,
                                                               int64_t write_timeout_seconds,
                                                               int32_t min_replica_count = 1);
-    // checksums must be the same length as the keys captured at StartWriteCache (full
-    // batch, not just the success subset within success_block_mask). Pass an empty
-    // vector when the client does not report a checksum; existing CacheLocation
-    // checksums are then preserved.
+    struct FinishWriteCacheOptions {
+        std::unique_ptr<WriteLocationManager::WriteLocationInfo> write_location_info_internal;
+        // checksums must be the same length as the keys captured at StartWriteCache
+        // (full batch, not just the success subset within success_block_mask). Pass
+        // an empty vector when the client does not report a checksum; existing
+        // CacheLocation checksums are then preserved.
+        std::vector<int64_t> checksums;
+
+        static FinishWriteCacheOptions WithWriteLocationInfo(
+            std::unique_ptr<WriteLocationManager::WriteLocationInfo> write_location_info_internal) {
+            FinishWriteCacheOptions options;
+            options.write_location_info_internal = std::move(write_location_info_internal);
+            return options;
+        }
+
+        static FinishWriteCacheOptions WithChecksums(std::vector<int64_t> checksums) {
+            FinishWriteCacheOptions options;
+            options.checksums = std::move(checksums);
+            return options;
+        }
+    };
+
     ErrorCode
     FinishWriteCache(RequestContext *request_context,
                      const std::string &instance_id,
                      const std::string &write_session_id,
                      const BlockMask &success_block_mask,
-                     std::unique_ptr<WriteLocationManager::WriteLocationInfo> write_location_info_internal = nullptr,
-                     const std::vector<int64_t> &checksums = {});
+                     FinishWriteCacheOptions options = FinishWriteCacheOptions{});
 
     ErrorCode RemoveCache(RequestContext *request_context,
                           const std::string &instance_id,