Creates a new Deallocate API in logical block manger

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<std::vector<int>> RaidenController::AllocateBlockIds(
@@ -451,7 +451,7 @@ absl::StatusOr<std::vector<int>> RaidenController::AllocateBlockIds(
 absl::Status RaidenController::DeallocateBlockIds(
     absl::Span<const int> block_ids) {
   absl::MutexLock lock(mutex_);
-  return block_manager_->Unlock(block_ids);
+  return block_manager_->Deallocate(block_ids);
 }
 
 absl::Status RaidenController::AllocateTargetBlockIds(

tpu_raiden/core/controller/raiden_controller.h

@@ -106,8 +106,8 @@ class RaidenController {
   // returns the corresponding Buffers. No gRPC call is made.
   absl::StatusOr<std::vector<Buffer>> 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<const Buffer> buffers);
 
   // Legacy API (Physical/BufferProto Mode):
@@ -121,8 +121,8 @@ class RaidenController {
   absl::StatusOr<std::vector<int>> 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<const int> block_ids);
 
   // New API (Logical/Block ID Mode):

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<std::string>& 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<std::string> first = {"hash_1", "hash_2"};
+  std::vector<RaidenBlockID> 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<std::string> second = {"hash_3"};
+  std::vector<RaidenBlockID> 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(

tpu_raiden/kv_cache/logical_block_manager.cc

@@ -159,6 +159,27 @@ absl::Status LogicalBlockManager::Unlock(absl::Span<const int> block_ids) {
   return absl::OkStatus();
 }
 
+absl::Status LogicalBlockManager::Deallocate(absl::Span<const int> 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(

tpu_raiden/kv_cache/logical_block_manager.h

@@ -55,9 +55,16 @@ class LogicalBlockManager {
   absl::Status AllocateTarget(absl::Span<const int> 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<const int> 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<const int> 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);

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());