From ff9d13508f28295381007c0b9c9f6d41d72a8a5f Mon Sep 17 00:00:00 2001 From: Benliang Wang Date: Wed, 5 Aug 2026 11:50:58 -0700 Subject: [PATCH] Creates a new Deallocate API in logical block manger PiperOrigin-RevId: 959785559 --- .../core/controller/raiden_controller.cc | 4 +- .../core/controller/raiden_controller.h | 8 +-- tpu_raiden/kv_cache/kv_cache_store_test.cc | 59 +++++++++++++++++++ tpu_raiden/kv_cache/logical_block_manager.cc | 21 +++++++ tpu_raiden/kv_cache/logical_block_manager.h | 9 ++- .../kv_cache/logical_block_manager_test.cc | 43 ++++++++++++++ 6 files changed, 137 insertions(+), 7 deletions(-) diff --git a/tpu_raiden/core/controller/raiden_controller.cc b/tpu_raiden/core/controller/raiden_controller.cc index 5b43e99c..54123117 100644 --- a/tpu_raiden/core/controller/raiden_controller.cc +++ b/tpu_raiden/core/controller/raiden_controller.cc @@ -439,7 +439,7 @@ absl::Status RaidenController::Deallocate( block_ids.push_back(sharded_buf.index()); } absl::MutexLock lock(mutex_); - return block_manager_->Unlock(block_ids); + return block_manager_->Deallocate(block_ids); } absl::StatusOr> RaidenController::AllocateBlockIds( @@ -451,7 +451,7 @@ absl::StatusOr> RaidenController::AllocateBlockIds( absl::Status RaidenController::DeallocateBlockIds( absl::Span block_ids) { absl::MutexLock lock(mutex_); - return block_manager_->Unlock(block_ids); + return block_manager_->Deallocate(block_ids); } absl::Status RaidenController::AllocateTargetBlockIds( diff --git a/tpu_raiden/core/controller/raiden_controller.h b/tpu_raiden/core/controller/raiden_controller.h index 1c8e707b..82ad71c8 100644 --- a/tpu_raiden/core/controller/raiden_controller.h +++ b/tpu_raiden/core/controller/raiden_controller.h @@ -106,8 +106,8 @@ class RaidenController { // returns the corresponding Buffers. No gRPC call is made. absl::StatusOr> AllocateBuffers(int num_blocks); - // Unlocks the specified Buffers in the local logical block manager so - // they can be reused. No gRPC call is made. + // Deallocates the specified Buffers in the local logical block manager, + // returning their blocks to the free pool. No gRPC call is made. absl::Status DeallocateBuffers(absl::Span buffers); // Legacy API (Physical/BufferProto Mode): @@ -121,8 +121,8 @@ class RaidenController { absl::StatusOr> AllocateBlockIds(int num_blocks); // New API (Logical/Block ID Mode): - // Unlocks/deallocates the given logical block IDs in the block manager. - // No RPC is performed. + // Deallocates the given logical block IDs in the block manager, returning + // them to the free pool. No RPC is performed. absl::Status DeallocateBlockIds(absl::Span block_ids); // New API (Logical/Block ID Mode): diff --git a/tpu_raiden/kv_cache/kv_cache_store_test.cc b/tpu_raiden/kv_cache/kv_cache_store_test.cc index ac1007c4..4785bd14 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_test.cc @@ -1165,6 +1165,65 @@ class KVCacheStoreEmbeddedControllerTest : public ::testing::Test { std::string orchestrator_address_; }; +TEST_F(KVCacheStoreEmbeddedControllerTest, SaveReusesFreedBlocksAfterEvict) { + // Evicted host blocks return to the free pool even when the directory has + // nothing left to evict: a full evict empties the directory and deallocates + // every block, and the save below must be served from those. + ::tpu_raiden::controller::MockTransferManager mock_mgr; + test_server_->service->SetTransferManager( + ::tpu_raiden::KVManagerHolder(&mock_mgr)); + + // A two-block pool, so the first save exhausts it. + auto controller = + std::make_unique<::tpu_raiden::controller::RaidenController>( + unit_, 2, 1, 512, orchestrator_address_, ""); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + auto* controller_ptr = controller.get(); + + RaidenId rid{"test_job", "0", "test_cache", 0}; + KVCacheStore store(2, std::move(controller), "", rid, std::nullopt, + /*store_server_ip=*/"127.0.0.1"); + + auto save_and_wait = + [&store](const std::vector& hashes) -> absl::Status { + absl::Status status = store.Save(hashes); + if (!status.ok()) return status; + while (true) { + auto [done, failed, pending] = store.PollSaveStatus(); + if (!failed.empty()) return absl::InternalError("async save failed"); + if (!done.empty()) return absl::OkStatus(); + absl::SleepFor(absl::Milliseconds(10)); + } + }; + + std::vector first = {"hash_1", "hash_2"}; + std::vector first_slices = { + RaidenBlockID(rid, -1, 0, BlockStatus::HBM), + RaidenBlockID(rid, -1, 1, BlockStatus::HBM)}; + ASSERT_TRUE(store.Insert(first, first_slices, false).first); + ASSERT_TRUE(store.Pin(first)); + ASSERT_TRUE(save_and_wait(first).ok()); + EXPECT_EQ(controller_ptr->block_manager()->num_locked_blocks(), 2); + + // Evict everything: the directory goes empty and both host blocks are + // deallocated back to the free pool. + store.Release(first); + ASSERT_EQ(KVCacheStoreTest::Evict(store, first), 2); + EXPECT_EQ(controller_ptr->block_manager()->num_free_blocks(), 2); + EXPECT_EQ(controller_ptr->block_manager()->num_locked_blocks(), 0); + + // The directory is empty and hash_3 gets pinned, so nothing is evictable: + // the host block for this save has to come from the freed pool. + std::vector second = {"hash_3"}; + std::vector second_slices = { + RaidenBlockID(rid, -1, 0, BlockStatus::HBM)}; + ASSERT_TRUE(store.Insert(second, second_slices, false).first); + ASSERT_TRUE(store.Pin(second)); + absl::Status status = save_and_wait(second); + EXPECT_TRUE(status.ok()) << status.message(); + EXPECT_EQ(controller_ptr->block_manager()->num_locked_blocks(), 1); +} + TEST_F(KVCacheStoreEmbeddedControllerTest, SaveSuccess) { ::tpu_raiden::controller::MockTransferManager mock_mgr; test_server_->service->SetTransferManager( diff --git a/tpu_raiden/kv_cache/logical_block_manager.cc b/tpu_raiden/kv_cache/logical_block_manager.cc index bf984933..d3e3ec1e 100644 --- a/tpu_raiden/kv_cache/logical_block_manager.cc +++ b/tpu_raiden/kv_cache/logical_block_manager.cc @@ -159,6 +159,27 @@ absl::Status LogicalBlockManager::Unlock(absl::Span block_ids) { return absl::OkStatus(); } +absl::Status LogicalBlockManager::Deallocate(absl::Span block_ids) { + // First validate all blocks. + for (int block_id : block_ids) { + if (block_id < 0 || block_id >= total_blocks_) { + return absl::InvalidArgumentError( + absl::StrCat("Invalid block ID: ", block_id)); + } + if (!blocks_[block_id].is_allocated) { + return absl::InvalidArgumentError( + absl::StrCat("Cannot deallocate unallocated block ID: ", block_id)); + } + } + + // Perform state update. + for (int block_id : block_ids) { + blocks_[block_id].is_allocated = false; + blocks_[block_id].is_locked = false; + } + return absl::OkStatus(); +} + absl::Status LogicalBlockManager::AccessBlock(int block_id) { if (block_id < 0 || block_id >= total_blocks_) { return absl::InvalidArgumentError( diff --git a/tpu_raiden/kv_cache/logical_block_manager.h b/tpu_raiden/kv_cache/logical_block_manager.h index 2b1acbe6..e12f0a47 100644 --- a/tpu_raiden/kv_cache/logical_block_manager.h +++ b/tpu_raiden/kv_cache/logical_block_manager.h @@ -55,9 +55,16 @@ class LogicalBlockManager { absl::Status AllocateTarget(absl::Span block_ids); // Unlocks the specified blocks, making them eligible for LRU eviction if - // future allocation requests require blocks. + // future allocation requests require blocks. The blocks stay allocated and + // their contents remain addressable until evicted. absl::Status Unlock(absl::Span block_ids); + // Deallocates the specified blocks, returning them to the free pool. Their + // contents must no longer be referenced. All IDs must be in range and + // currently allocated (locked or not); on any violation an error is + // returned and no state is modified. + absl::Status Deallocate(absl::Span block_ids); + // Updates the access time (logical counter) for the specified block, marking // it as the most recently used. absl::Status AccessBlock(int block_id); diff --git a/tpu_raiden/kv_cache/logical_block_manager_test.cc b/tpu_raiden/kv_cache/logical_block_manager_test.cc index b4686558..60ca08ba 100644 --- a/tpu_raiden/kv_cache/logical_block_manager_test.cc +++ b/tpu_raiden/kv_cache/logical_block_manager_test.cc @@ -195,6 +195,49 @@ TEST(LogicalBlockManagerTest, AllocateTargetValidatesAtomically) { EXPECT_EQ(manager.num_allocated_blocks(), 1); } +TEST(LogicalBlockManagerTest, DeallocateReturnsBlocksToFreePool) { + LogicalBlockManager manager(4); + ASSERT_TRUE(manager.Allocate(2, /*lock=*/true).ok()); + ASSERT_TRUE(manager.Allocate(1, /*lock=*/false).ok()); + EXPECT_EQ(manager.num_free_blocks(), 1); + + // Deallocation works on locked and unlocked blocks alike. + ASSERT_TRUE(manager.Deallocate({0, 2}).ok()); + EXPECT_EQ(manager.num_free_blocks(), 3); + EXPECT_EQ(manager.num_allocated_blocks(), 1); + EXPECT_EQ(manager.num_locked_blocks(), 1); + EXPECT_FALSE(manager.IsAllocated(0)); + EXPECT_FALSE(manager.IsLocked(0)); + EXPECT_FALSE(manager.IsAllocated(2)); + + // Deallocated blocks are free again for both allocation paths. + ASSERT_TRUE(manager.AllocateTarget({0}).ok()); + auto blocks_or = manager.Allocate(1); + ASSERT_TRUE(blocks_or.ok()); + EXPECT_THAT(*blocks_or, ElementsAre(2)); +} + +TEST(LogicalBlockManagerTest, DeallocateValidatesAtomically) { + LogicalBlockManager manager(3); + ASSERT_TRUE(manager.Allocate(1, /*lock=*/true).ok()); // Block 0. + + // Out of range. + EXPECT_EQ(manager.Deallocate({0, 3}).code(), + absl::StatusCode::kInvalidArgument); + // Not allocated. + EXPECT_EQ(manager.Deallocate({0, 1}).code(), + absl::StatusCode::kInvalidArgument); + + // Failed calls must not have modified any state. + EXPECT_TRUE(manager.IsAllocated(0)); + EXPECT_TRUE(manager.IsLocked(0)); + + ASSERT_TRUE(manager.Deallocate({0}).ok()); + // Double deallocation fails. + EXPECT_EQ(manager.Deallocate({0}).code(), + absl::StatusCode::kInvalidArgument); +} + TEST(LogicalBlockManagerTest, TargetAllocatedBlocksReusableAfterUnlock) { LogicalBlockManager manager(3); ASSERT_TRUE(manager.AllocateTarget({0, 1, 2}).ok());