From 16eae82736552cd04367c1b26141e8191b63d702 Mon Sep 17 00:00:00 2001 From: eric_song Date: Sun, 12 Jul 2026 23:48:28 +0800 Subject: [PATCH 1/3] [TransferEngine] Propagate batch memory operation errors --- .../src/transfer_engine_impl.cpp | 4 ++-- .../hccl_transport/hccl_transport.cpp | 12 +++++++++--- .../ubshmem_transport/ubshmem_transport.cpp | 4 ++-- .../transport/barex_transport/barex_transport.cpp | 9 +++++++-- .../src/transport/cxi_transport/cxi_transport.cpp | 13 +++++++++++-- .../src/transport/cxl_transport/cxl_transport.cpp | 12 +++++++++--- .../src/transport/efa_transport/efa_transport.cpp | 13 +++++++++++-- .../src/transport/hip_transport/hip_transport.cpp | 4 ++-- .../intranode_nvlink_transport.cpp | 12 +++++++++--- .../transport/kunpeng_transport/ub_transport.cpp | 13 +++++++++++-- .../transport/maca_transport/maca_transport.cpp | 12 +++++++++--- .../nvlink_transport/nvlink_transport.cpp | 12 +++++++++--- .../transport/rdma_transport/rdma_transport.cpp | 14 ++++++++++++-- .../sunrise_link/sunrise_link_transport.cpp | 9 +++++++-- .../src/transport/tcp_transport/tcp_transport.cpp | 12 +++++++++--- .../tests/transport_uint_test.cpp | 10 ++++++++++ 16 files changed, 129 insertions(+), 36 deletions(-) diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index 98b55596a6..e78cc0c122 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -788,7 +788,7 @@ int TransferEngineImpl::registerLocalMemoryBatch( for (auto transport : multi_transports_->listTransports()) { int ret = transport->registerLocalMemoryBatch(buffer_list, location); - if (ret < 0) { + if (ret) { releaseMemoryRegions(regions); return ret; } @@ -802,7 +802,7 @@ int TransferEngineImpl::unregisterLocalMemoryBatch( const std::vector& addr_list) { for (auto transport : multi_transports_->listTransports()) { int ret = transport->unregisterLocalMemoryBatch(addr_list); - if (ret < 0) return ret; + if (ret) return ret; } std::unique_lock lock(mutex_); diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp index a1c6958502..d37b0cac4f 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp @@ -607,14 +607,20 @@ int HcclTransport::allocateLocalSegmentID() { int HcclTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, -1); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int HcclTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, -1); + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp index 484c698335..9eccc5aa69 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp @@ -722,7 +722,7 @@ int UBShmemTransport::registerLocalMemoryBatch( for (auto &buffer : buffer_list) { int rc = registerLocalMemory(buffer.addr, buffer.length, location, true, false); - if (rc < 0) { + if (rc) { return rc; } } @@ -733,7 +733,7 @@ int UBShmemTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { for (auto &addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc < 0) { + if (rc) { return rc; } } diff --git a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp index e575dfc39d..9d94d56ed7 100644 --- a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp +++ b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp @@ -306,7 +306,7 @@ int BarexTransport::registerLocalMemoryBatch( if (ret) { LOG(ERROR) << "BarexTransport: Failed to register memory: addr " << buffer.addr << " length " << buffer.length; - return ERR_ADDRESS_NOT_REGISTERED; + return ret; } } @@ -323,11 +323,16 @@ int BarexTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "BarexTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp index c106beadaa..46a25373d3 100644 --- a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp +++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp @@ -552,13 +552,17 @@ int CxiTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "CxiTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } @@ -573,11 +577,16 @@ int CxiTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "CxiTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp index a12c649abb..fc2debb109 100644 --- a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp @@ -272,14 +272,20 @@ int CxlTransport::unregisterLocalMemory(void *addr, bool update_metadata) { int CxlTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int CxlTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp index 617a0ee005..ba57fb650c 100644 --- a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp @@ -636,13 +636,17 @@ int EfaTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "EfaTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } @@ -657,11 +661,16 @@ int EfaTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "EfaTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp index a15ffd450d..c3ec141edf 100644 --- a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp +++ b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp @@ -823,7 +823,7 @@ int HipTransport::registerLocalMemoryBatch( for (auto& buffer : buffer_list) { int rc = registerLocalMemory(buffer.addr, buffer.length, location, true, false); - if (rc < 0) return rc; + if (rc) return rc; } return metadata_->updateLocalSegmentDesc(); } @@ -832,7 +832,7 @@ int HipTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { for (auto& addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc < 0) return rc; + if (rc) return rc; } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp index 901fe24696..5e9e70731b 100644 --- a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp @@ -762,14 +762,20 @@ int IntraNodeNvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int IntraNodeNvlinkTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int IntraNodeNvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 2dd3dc42ca..b594059326 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -142,13 +142,17 @@ int UbTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "UbTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } @@ -164,11 +168,16 @@ int UbTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "UbTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp index b15fefea62..09818928bd 100644 --- a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp +++ b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp @@ -920,14 +920,20 @@ int MacaTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int MacaTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int MacaTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp index cdc6aea303..0b7cf8288f 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp @@ -855,14 +855,20 @@ int NvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int NvlinkTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int NvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index 9b76dacb03..766356486f 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -453,6 +453,7 @@ int RdmaTransport::registerLocalMemoryBatch( LOG(WARNING) << "RdmaTransport: Failed to register memory: addr " << buffer.addr << " length " << buffer.length; + return ret; } } } else { @@ -468,14 +469,18 @@ int RdmaTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "RdmaTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } + if (first_error) return first_error; #if defined(USE_CUDA) } // Environ::Get().GetWithNvidiaPeermem() #endif @@ -494,11 +499,16 @@ int RdmaTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "RdmaTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp index c91acb24b2..34227d1ad2 100644 --- a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp +++ b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp @@ -418,14 +418,19 @@ int SunriseLinkTransport::unregisterLocalMemory(void* addr, int SunriseLinkTransport::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { for (const auto& buffer : buffer_list) { - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; } return metadata_->updateLocalSegmentDesc(); } int SunriseLinkTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { - for (auto* addr : addr_list) unregisterLocalMemory(addr, false); + for (auto* addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index 253d7ed32c..ee711fad13 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -964,14 +964,20 @@ int TcpTransport::unregisterLocalMemory(void* addr, bool update_metadata) { int TcpTransport::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { - for (auto& buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto& buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int TcpTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { - for (auto& addr : addr_list) unregisterLocalMemory(addr, false); + for (auto& addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index 3c5b10319e..e5d4985409 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -366,6 +366,16 @@ TEST_F(TransportTest, FailedRegistrationReleasesReservedRegion) { 0); EXPECT_EQ(engine.unregisterLocalMemory(buffer.data()), 0); } + +TEST_F(TransportTest, UnregisterLocalMemoryBatchPropagatesTransportError) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + + std::array buffer{}; + EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), + ERR_ADDRESS_NOT_REGISTERED); +} } // namespace mooncake int main(int argc, char** argv) { From 86f72ec3b577e94ef4730f9eb6def35cabefc430 Mon Sep 17 00:00:00 2001 From: eric_song Date: Tue, 14 Jul 2026 21:02:33 +0800 Subject: [PATCH 2/3] [TransferEngine] Make batch unregister best-effort --- .../src/transfer_engine_impl.cpp | 4 +- .../ascend_direct_transport.cpp | 6 +- .../hccl_transport/hccl_transport.cpp | 6 +- .../heterogeneous_rdma_transport.cpp | 6 +- .../ubshmem_transport/ubshmem_transport.cpp | 8 +- .../barex_transport/barex_transport.cpp | 5 +- .../transport/cxi_transport/cxi_transport.cpp | 5 +- .../transport/cxl_transport/cxl_transport.cpp | 6 +- .../transport/efa_transport/efa_transport.cpp | 5 +- .../transport/hip_transport/hip_transport.cpp | 6 +- .../intranode_nvlink_transport.cpp | 6 +- .../kunpeng_transport/ub_transport.cpp | 5 +- .../maca_transport/maca_transport.cpp | 6 +- .../nvlink_transport/nvlink_transport.cpp | 6 +- .../rdma_transport/rdma_transport.cpp | 5 +- .../sunrise_link/sunrise_link_transport.cpp | 6 +- .../transport/tcp_transport/tcp_transport.cpp | 6 +- .../tests/transport_uint_test.cpp | 98 +++++++++++++++++++ 18 files changed, 155 insertions(+), 40 deletions(-) diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index e78cc0c122..a607ca94b5 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -800,10 +800,12 @@ int TransferEngineImpl::registerLocalMemoryBatch( int TransferEngineImpl::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto transport : multi_transports_->listTransports()) { int ret = transport->unregisterLocalMemoryBatch(addr_list); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } + if (first_error) return first_error; std::unique_lock lock(mutex_); for (auto& addr : addr_list) { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp index 51e97eeeb5..365505cbb0 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp @@ -420,16 +420,18 @@ int AscendDirectTransport::unregisterLocalMemoryBatch( "with addr count: " << addr_list.size(); + int first_error = 0; for (void *addr : addr_list) { int ret = unregisterLocalMemory(addr, false); if (ret != 0) { LOG(ERROR) << "Failed to unregister memory in batch, addr: " << addr; - return ret; + if (!first_error) first_error = ret; } } // Update metadata once for the entire batch - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp index d37b0cac4f..ef81cb6f30 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp @@ -617,11 +617,13 @@ int HcclTransport::registerLocalMemoryBatch( int HcclTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp index 943332a919..1f0962d245 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp @@ -208,16 +208,18 @@ int HeterogeneousRdmaTransport::registerLocalMemoryBatch( int HeterogeneousRdmaTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); if (ret) { LOG(ERROR) << "HeterogeneousRdmaTransport " "unregisterLocalMemoryBatch error, ret: " << ret; - return ret; + if (!first_error) first_error = ret; } } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int HeterogeneousRdmaTransport::checkAndCreateStreamCopy() { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp index 9eccc5aa69..1125e5d466 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp @@ -731,13 +731,13 @@ int UBShmemTransport::registerLocalMemoryBatch( int UBShmemTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc) { - return rc; - } + if (rc && !first_error) first_error = rc; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *UBShmemTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp index 9d94d56ed7..2301102149 100644 --- a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp +++ b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp @@ -332,9 +332,8 @@ int BarexTransport::unregisterLocalMemoryBatch( if (!first_error) first_error = ret; } } - if (first_error) return first_error; - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status BarexTransport::submitTransfer( diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp index 46a25373d3..3c42fe66da 100644 --- a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp +++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp @@ -586,9 +586,8 @@ int CxiTransport::unregisterLocalMemoryBatch( if (!first_error) first_error = ret; } } - if (first_error) return first_error; - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int CxiTransport::warmupSegment(const std::string& segment_name) { diff --git a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp index fc2debb109..23984f7d47 100644 --- a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp @@ -282,11 +282,13 @@ int CxlTransport::registerLocalMemoryBatch( int CxlTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status CxlTransport::getTransferStatus(BatchID batch_id, size_t task_id, diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp index ba57fb650c..7628165ca3 100644 --- a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp @@ -670,9 +670,8 @@ int EfaTransport::unregisterLocalMemoryBatch( if (!first_error) first_error = ret; } } - if (first_error) return first_error; - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int EfaTransport::warmupSegment(const std::string& segment_name) { diff --git a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp index c3ec141edf..71031c3b9f 100644 --- a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp +++ b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp @@ -830,11 +830,13 @@ int HipTransport::registerLocalMemoryBatch( int HipTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto& addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc) return rc; + if (rc && !first_error) first_error = rc; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void* HipTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp index 5e9e70731b..52ca987add 100644 --- a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp @@ -772,11 +772,13 @@ int IntraNodeNvlinkTransport::registerLocalMemoryBatch( int IntraNodeNvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *IntraNodeNvlinkTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index b594059326..61e8ea7ba4 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -177,9 +177,8 @@ int UbTransport::unregisterLocalMemoryBatch( if (!first_error) first_error = ret; } } - if (first_error) return first_error; - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status UbTransport::submitTransfer( diff --git a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp index 09818928bd..c769a9f53b 100644 --- a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp +++ b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp @@ -930,11 +930,13 @@ int MacaTransport::registerLocalMemoryBatch( int MacaTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *MacaTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp index 0b7cf8288f..bac4d625cd 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp @@ -865,11 +865,13 @@ int NvlinkTransport::registerLocalMemoryBatch( int NvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *NvlinkTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index 766356486f..50394094a9 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -508,9 +508,8 @@ int RdmaTransport::unregisterLocalMemoryBatch( if (!first_error) first_error = ret; } } - if (first_error) return first_error; - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status RdmaTransport::submitTransfer( diff --git a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp index 34227d1ad2..ebf3f654b7 100644 --- a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp +++ b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp @@ -427,11 +427,13 @@ int SunriseLinkTransport::registerLocalMemoryBatch( int SunriseLinkTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto* addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int SunriseLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr, diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index ee711fad13..9cf906a0cd 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -974,11 +974,13 @@ int TcpTransport::registerLocalMemoryBatch( int TcpTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto& addr : addr_list) { int ret = unregisterLocalMemory(addr, false); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status TcpTransport::getTransferStatus(BatchID batch_id, size_t task_id, diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index e5d4985409..59bf0be707 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -43,6 +44,56 @@ class TransferEngineImplTestPeer { engine.multi_transports_->transport_map_.emplace("blocking", std::move(transport)); } + + static void replaceTransports( + TransferEngineImpl& engine, + const std::vector>>& + transports) { + engine.multi_transports_->transport_map_.clear(); + for (const auto& [name, transport] : transports) { + engine.multi_transports_->transport_map_.emplace(name, transport); + } + } +}; + +class BatchResultTransport : public Transport { + public: + explicit BatchResultTransport(int unregister_result = 0) + : unregister_result_(unregister_result) {} + + int unregisterBatchCalls() const { return unregister_batch_calls_; } + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::OK(); + } + + Status getTransferStatus(BatchID, size_t, TransferStatus&) override { + return Status::OK(); + } + + private: + int registerLocalMemory(void*, size_t, const std::string&, bool, + bool) override { + return 0; + } + + int unregisterLocalMemory(void*, bool) override { return 0; } + + int registerLocalMemoryBatch(const std::vector&, + const std::string&) override { + return 0; + } + + int unregisterLocalMemoryBatch(const std::vector&) override { + ++unregister_batch_calls_; + return unregister_result_; + } + + const char* getName() const override { return "batch-result"; } + + int unregister_result_; + int unregister_batch_calls_ = 0; }; class BlockingRegistrationTransport : public Transport { @@ -376,6 +427,53 @@ TEST_F(TransportTest, UnregisterLocalMemoryBatchPropagatesTransportError) { EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), ERR_ADDRESS_NOT_REGISTERED); } + +TEST_F(TransportTest, UnregisterLocalMemoryBatchContinuesAcrossTransports) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto failing = std::make_shared(ERR_MEMORY); + auto succeeding = std::make_shared(); + TransferEngineImplTestPeer::replaceTransports( + engine, {{"a-failing", failing}, {"b-succeeding", succeeding}}); + + std::array buffer{}; + EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), ERR_MEMORY); + EXPECT_EQ(failing->unregisterBatchCalls(), 1); + EXPECT_EQ(succeeding->unregisterBatchCalls(), 1); +} + +TEST_F(TransportTest, UnregisterLocalMemoryBatchContinuesAfterAddressError) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + + std::array registered{}; + std::array missing{}; + std::vector entries = { + {registered.data(), 1}, + {registered.data() + 1, 1}, + }; + ASSERT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); + + auto metadata = engine.getMetadata(); + ASSERT_NE(metadata, nullptr); + auto contains_buffer = [&](void* addr) { + auto desc = metadata->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (!desc) return false; + auto value = reinterpret_cast(addr); + return std::any_of( + desc->buffers.begin(), desc->buffers.end(), + [value](const auto& buffer) { return buffer.addr == value; }); + }; + ASSERT_TRUE(contains_buffer(registered.data())); + ASSERT_TRUE(contains_buffer(registered.data() + 1)); + + EXPECT_EQ(engine.unregisterLocalMemoryBatch( + {missing.data(), registered.data(), registered.data() + 1}), + ERR_ADDRESS_NOT_REGISTERED); + EXPECT_FALSE(contains_buffer(registered.data())); + EXPECT_FALSE(contains_buffer(registered.data() + 1)); +} } // namespace mooncake int main(int argc, char** argv) { From 516b9495347f84d8af5d027bff45c65a9f35a51c Mon Sep 17 00:00:00 2001 From: eric_song Date: Tue, 14 Jul 2026 21:05:32 +0800 Subject: [PATCH 3/3] [TransferEngine] Roll back failed batch registrations --- .../src/transfer_engine_impl.cpp | 15 ++++++ .../tests/transport_uint_test.cpp | 53 ++++++++++++++++++- 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index a607ca94b5..f5aeb00b9b 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -776,9 +776,12 @@ int TransferEngineImpl::registerLocalMemoryBatch( } std::vector regions; + std::vector addr_list; regions.reserve(buffer_list.size()); + addr_list.reserve(buffer_list.size()); for (const auto& buffer : buffer_list) { regions.push_back({buffer.addr, buffer.length, location, true}); + addr_list.push_back(buffer.addr); } if (!tryReserveMemoryRegions(regions)) { LOG(ERROR) @@ -786,9 +789,21 @@ int TransferEngineImpl::registerLocalMemoryBatch( return ERR_ADDRESS_OVERLAPPED; } + std::vector attempted_transports; for (auto transport : multi_transports_->listTransports()) { + attempted_transports.push_back(transport); int ret = transport->registerLocalMemoryBatch(buffer_list, location); if (ret) { + for (auto it = attempted_transports.rbegin(); + it != attempted_transports.rend(); ++it) { + int rollback_ret = (*it)->unregisterLocalMemoryBatch(addr_list); + if (rollback_ret != 0 && + rollback_ret != ERR_ADDRESS_NOT_REGISTERED) { + LOG(WARNING) + << "Failed to roll back batch registration for " + << (*it)->getName() << ", ret=" << rollback_ret; + } + } releaseMemoryRegions(regions); return ret; } diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index 59bf0be707..cd7097f1d8 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -62,6 +62,8 @@ class BatchResultTransport : public Transport { : unregister_result_(unregister_result) {} int unregisterBatchCalls() const { return unregister_batch_calls_; } + size_t registeredBufferCount() const { return registered_buffers_.size(); } + void setRegisterResult(int result) { register_result_ = result; } Status submitTransfer(BatchID, const std::vector&) override { @@ -80,20 +82,38 @@ class BatchResultTransport : public Transport { int unregisterLocalMemory(void*, bool) override { return 0; } - int registerLocalMemoryBatch(const std::vector&, + int registerLocalMemoryBatch(const std::vector& buffer_list, const std::string&) override { + if (register_result_) { + if (!buffer_list.empty()) { + registered_buffers_.push_back(buffer_list.front().addr); + } + return register_result_; + } + for (const auto& buffer : buffer_list) { + registered_buffers_.push_back(buffer.addr); + } return 0; } - int unregisterLocalMemoryBatch(const std::vector&) override { + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override { ++unregister_batch_calls_; + for (void* addr : addr_list) { + registered_buffers_.erase( + std::remove(registered_buffers_.begin(), + registered_buffers_.end(), addr), + registered_buffers_.end()); + } return unregister_result_; } const char* getName() const override { return "batch-result"; } + int register_result_ = 0; int unregister_result_; int unregister_batch_calls_ = 0; + std::vector registered_buffers_; }; class BlockingRegistrationTransport : public Transport { @@ -474,6 +494,35 @@ TEST_F(TransportTest, UnregisterLocalMemoryBatchContinuesAfterAddressError) { EXPECT_FALSE(contains_buffer(registered.data())); EXPECT_FALSE(contains_buffer(registered.data() + 1)); } + +TEST_F(TransportTest, RegisterLocalMemoryBatchRollsBackAttemptedTransports) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto succeeding = std::make_shared(); + auto failing = std::make_shared(); + failing->setRegisterResult(ERR_MEMORY); + TransferEngineImplTestPeer::replaceTransports( + engine, {{"a-succeeding", succeeding}, {"b-failing", failing}}); + + std::array buffer{}; + std::vector entries = { + {buffer.data(), 1}, + {buffer.data() + 1, 1}, + }; + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), ERR_MEMORY); + EXPECT_EQ(succeeding->registeredBufferCount(), 0); + EXPECT_EQ(failing->registeredBufferCount(), 0); + EXPECT_EQ(succeeding->unregisterBatchCalls(), 1); + EXPECT_EQ(failing->unregisterBatchCalls(), 1); + + failing->setRegisterResult(0); + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); + EXPECT_EQ(succeeding->registeredBufferCount(), entries.size()); + EXPECT_EQ(failing->registeredBufferCount(), entries.size()); + EXPECT_EQ( + engine.unregisterLocalMemoryBatch({buffer.data(), buffer.data() + 1}), + 0); +} } // namespace mooncake int main(int argc, char** argv) {