From 74d6643dc892ff81fbe535ccf5874529ec0ca159 Mon Sep 17 00:00:00 2001 From: justinlu Date: Wed, 5 Aug 2026 14:43:18 -0700 Subject: [PATCH] Expose and enhance Load in KVCacheStoreBackend and HostOffloadBackend. PiperOrigin-RevId: 959879231 --- tpu_raiden/kv_cache/BUILD | 2 + tpu_raiden/kv_cache/host_offload_backend.cc | 224 +++++++++++------- tpu_raiden/kv_cache/host_offload_backend.h | 2 +- .../kv_cache/host_offload_backend_test.cc | 98 ++++++++ tpu_raiden/kv_cache/kv_cache_store.cc | 74 +++--- tpu_raiden/kv_cache/kv_cache_store_backend.h | 7 + .../kv_cache_store_backend_factory_test.cc | 7 + tpu_raiden/kv_cache/kv_cache_store_test.cc | 113 +++++++++ 8 files changed, 404 insertions(+), 123 deletions(-) diff --git a/tpu_raiden/kv_cache/BUILD b/tpu_raiden/kv_cache/BUILD index 2eede1c6..5240cb36 100644 --- a/tpu_raiden/kv_cache/BUILD +++ b/tpu_raiden/kv_cache/BUILD @@ -280,6 +280,7 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", + "@xla//xla/tsl/concurrency:future", ], ) @@ -354,6 +355,7 @@ cc_test( "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/concurrency:future", ], ) diff --git a/tpu_raiden/kv_cache/host_offload_backend.cc b/tpu_raiden/kv_cache/host_offload_backend.cc index 63c5552e..3a3e69ab 100644 --- a/tpu_raiden/kv_cache/host_offload_backend.cc +++ b/tpu_raiden/kv_cache/host_offload_backend.cc @@ -946,9 +946,11 @@ tsl::Future<> HostOffloadBackend::Load( } controller::RaidenController* ctrl = nullptr; + bool is_remote = false; { absl::MutexLock lock(mutex_); ctrl = raiden_controller_; + is_remote = !remote_id.empty() && remote_id != raiden_id_; } if (ctrl == nullptr) { @@ -956,109 +958,151 @@ tsl::Future<> HostOffloadBackend::Load( absl::FailedPreconditionError("RaidenController is null")); } - auto client_or = GetKVCacheStoreClient(remote_id); - if (!client_or.ok()) { - return tsl::Future<>(client_or.status()); - } - std::shared_ptr client = std::move(client_or.value()); + if (is_remote) { + auto client_or = GetKVCacheStoreClient(remote_id); + if (!client_or.ok()) { + return tsl::Future<>(client_or.status()); + } + std::shared_ptr client = std::move(client_or.value()); - auto host_blocks_or = ctrl->AllocateBlockIds(block_hashes.size()); - if (!host_blocks_or.ok()) { - return tsl::Future<>(host_blocks_or.status()); - } - std::vector dst_host_block_ids(host_blocks_or.value().begin(), - host_blocks_or.value().end()); + auto host_blocks_or = ctrl->AllocateBlockIds(block_hashes.size()); + if (!host_blocks_or.ok()) { + return tsl::Future<>(host_blocks_or.status()); + } + std::vector dst_host_block_ids(host_blocks_or.value().begin(), + host_blocks_or.value().end()); + + auto [load_promise, load_future] = tsl::MakePromise<>(); + + tpu_raiden::rpc::RaidenIdProto client_raiden_id = ctrl->unit(); + std::vector<::tpu_raiden::proto::RaidenWorkerEndpointsProto> + client_worker_endpoints = BuildLocalWorkerEndpoints(ctrl); + tsl::Future fetch_future = + client->Fetch(block_hashes, device_block_ids, dst_host_block_ids, + client_raiden_id, client_worker_endpoints); + + fetch_future.OnReady( + [this, remote_id, dst_host_block_ids, + dev_ids_vec = std::vector(device_block_ids.begin(), + device_block_ids.end()), + load_promise = std::move(load_promise)]( + const absl::StatusOr& response_or) mutable { + controller::RaidenController* ctrl_cb = nullptr; + { + absl::MutexLock lock(mutex_); + ctrl_cb = raiden_controller_; + } + if (!response_or.ok()) { + if (ctrl_cb) { + (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + } + // The peer may have restarted on a new port; drop the cached client + // so the next attempt re-resolves instead of redialling a dead one. + InvalidateStoreClient(remote_id); + load_promise.Set(response_or.status()); + return; + } - auto [load_promise, load_future] = tsl::MakePromise<>(); + const auto& response = response_or.value(); + if (!response.failed_block_hashes().empty()) { + if (ctrl_cb) { + (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + } + std::string err_msg = response.error_message().empty() + ? "Fetch RPC returned failed blocks" + : response.error_message(); + load_promise.Set(absl::InternalError(err_msg)); + return; + } - tpu_raiden::rpc::RaidenIdProto client_raiden_id = ctrl->unit(); - std::vector<::tpu_raiden::proto::RaidenWorkerEndpointsProto> - client_worker_endpoints = BuildLocalWorkerEndpoints(ctrl); - tsl::Future fetch_future = client->Fetch( - block_hashes, device_block_ids, dst_host_block_ids, client_raiden_id, - client_worker_endpoints); + if (dev_ids_vec.empty()) { + if (ctrl_cb) { + (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + } + load_promise.Set(absl::OkStatus()); + return; + } - fetch_future.OnReady( - [this, remote_id, dst_host_block_ids, - dev_ids_vec = std::vector(device_block_ids.begin(), - device_block_ids.end()), - load_promise = std::move(load_promise)]( - const absl::StatusOr& response_or) mutable { - controller::RaidenController* ctrl_cb = nullptr; - { - absl::MutexLock lock(mutex_); - ctrl_cb = raiden_controller_; - } - if (!response_or.ok()) { - if (ctrl_cb) { - (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + std::vector src_buffers; + src_buffers.reserve(dst_host_block_ids.size()); + for (int id : dst_host_block_ids) { + src_buffers.emplace_back(id, std::vector{}, + std::nullopt, rpc::MEMORY_TYPE_DRAM); } - // The peer may have restarted on a new port; drop the cached client - // so the next attempt re-resolves instead of redialling a dead one. - InvalidateStoreClient(remote_id); - load_promise.Set(response_or.status()); - return; - } - const auto& response = response_or.value(); - if (!response.failed_block_hashes().empty()) { - if (ctrl_cb) { - (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + std::vector dst_buffers; + dst_buffers.reserve(dev_ids_vec.size()); + for (int id : dev_ids_vec) { + dst_buffers.emplace_back(id, std::vector{}, + std::nullopt, rpc::MEMORY_TYPE_HBM); } - std::string err_msg = response.error_message().empty() - ? "Fetch RPC returned failed blocks" - : response.error_message(); - load_promise.Set(absl::InternalError(err_msg)); - return; - } - if (dev_ids_vec.empty()) { - if (ctrl_cb) { - (void)ctrl_cb->DeallocateBlockIds(dst_host_block_ids); + if (!ctrl_cb) { + load_promise.Set( + absl::FailedPreconditionError("RaidenController is null")); + return; } - load_promise.Set(absl::OkStatus()); - return; - } - std::vector src_buffers; - src_buffers.reserve(dst_host_block_ids.size()); - for (int id : dst_host_block_ids) { - src_buffers.emplace_back(id, std::vector{}, std::nullopt, - rpc::MEMORY_TYPE_DRAM); - } + tsl::Future<> h2d_future = + ctrl_cb->TransferBuffers(src_buffers, dst_buffers); + + h2d_future.OnReady([this, dst_host_block_ids, + load_promise = std::move(load_promise)]( + absl::Status status) mutable { + controller::RaidenController* ctrl_h2d = nullptr; + { + absl::MutexLock lock(mutex_); + ctrl_h2d = raiden_controller_; + } + if (ctrl_h2d) { + (void)ctrl_h2d->DeallocateBlockIds(dst_host_block_ids); + } + load_promise.Set(status); + }); + }); - std::vector dst_buffers; - dst_buffers.reserve(dev_ids_vec.size()); - for (int id : dev_ids_vec) { - dst_buffers.emplace_back(id, std::vector{}, std::nullopt, - rpc::MEMORY_TYPE_HBM); - } + return load_future; + } - if (!ctrl_cb) { - load_promise.Set( - absl::FailedPreconditionError("RaidenController is null")); - return; - } + // --- Local Host DRAM Branch --- + std::vector src_host_block_ids; + src_host_block_ids.reserve(block_hashes.size()); + { + absl::MutexLock lock(mutex_); + for (const auto& hash : block_hashes) { + const RaidenBlockID* entry = lru_cache_.Peek(hash); + if (entry == nullptr) { + return tsl::Future<>(absl::NotFoundError( + absl::StrCat("Block hash not found in host backend: ", hash))); + } + if (entry->status != BlockStatus::HOST && + entry->status != BlockStatus::HOST_AND_HBM) { + return tsl::Future<>(absl::FailedPreconditionError( + absl::StrCat("Block is not on host: ", hash))); + } + if (entry->host_block_id == -1) { + return tsl::Future<>(absl::FailedPreconditionError( + absl::StrCat("Block host_block_id is -1: ", hash))); + } + src_host_block_ids.push_back(entry->host_block_id); + } + } + + std::vector src_buffers; + src_buffers.reserve(src_host_block_ids.size()); + for (int64_t id : src_host_block_ids) { + src_buffers.emplace_back(id, std::vector{}, std::nullopt, + rpc::MEMORY_TYPE_DRAM); + } + + std::vector dst_buffers; + dst_buffers.reserve(device_block_ids.size()); + for (int id : device_block_ids) { + dst_buffers.emplace_back(id, std::vector{}, std::nullopt, + rpc::MEMORY_TYPE_HBM); + } - tsl::Future<> h2d_future = - ctrl_cb->TransferBuffers(src_buffers, dst_buffers); - - h2d_future.OnReady( - [this, dst_host_block_ids, load_promise = std::move(load_promise)]( - absl::Status status) mutable { - controller::RaidenController* ctrl_h2d = nullptr; - { - absl::MutexLock lock(mutex_); - ctrl_h2d = raiden_controller_; - } - if (ctrl_h2d) { - (void)ctrl_h2d->DeallocateBlockIds(dst_host_block_ids); - } - load_promise.Set(status); - }); - }); - - return load_future; + return ctrl->TransferBuffers(src_buffers, dst_buffers); } void HostOffloadBackend::SetMetadataEntry(absl::string_view hash, diff --git a/tpu_raiden/kv_cache/host_offload_backend.h b/tpu_raiden/kv_cache/host_offload_backend.h index 49da0fc0..3a0546f0 100644 --- a/tpu_raiden/kv_cache/host_offload_backend.h +++ b/tpu_raiden/kv_cache/host_offload_backend.h @@ -131,7 +131,7 @@ class HostOffloadBackend : public KVCacheStoreBackend { tsl::Future<> Load(const RaidenId& remote_id, absl::Span block_hashes, - absl::Span device_block_ids = {}); + absl::Span device_block_ids = {}) override; // --- Remote write (WriteRemote); see KVCacheStoreBackend for why each of // these exists rather than reusing Lookup/Insert/Delete. diff --git a/tpu_raiden/kv_cache/host_offload_backend_test.cc b/tpu_raiden/kv_cache/host_offload_backend_test.cc index fded8696..e2a3ad63 100644 --- a/tpu_raiden/kv_cache/host_offload_backend_test.cc +++ b/tpu_raiden/kv_cache/host_offload_backend_test.cc @@ -562,6 +562,104 @@ TEST(HostOffloadBackendTest, LoadSuccess) { server->Shutdown(); } +TEST(HostOffloadBackendTest, LoadLocalSuccess) { + RaidenId node_id{"node_job", "0", "data", 0}; + rpc::RaidenIdProto unit_proto; + unit_proto.set_job_name(node_id.job_name); + unit_proto.set_job_replica_id(node_id.job_replica_id); + unit_proto.set_data_name(node_id.data_name); + unit_proto.set_data_replica_idx(node_id.data_replica_idx); + + controller::RaidenController controller(unit_proto, /*num_blocks=*/100, + /*num_shards=*/1, + /*shard_size_bytes=*/1024); + auto test_worker_server = controller::CreateTestWorkerServer(); + auto transfer_mock = + std::make_unique(); + test_worker_server->service->SetTransferManager( + KVManagerHolder(transfer_mock.get())); + + core::controller::RaidenControllerClient controller_client( + controller.controller_address()); + ASSERT_OK(controller_client.RegisterWorker( + "worker_0", test_worker_server->server_address, + {{test_worker_server->server_address, {}}})); + + BackendConfig config; + config.type = "HostOffloadBackend"; + config.capacity = 100; + config.raiden_id = node_id; + + auto backend_or = HostOffloadBackend::Create(config, &controller); + ASSERT_OK(backend_or.status()); + auto backend = std::dynamic_pointer_cast(*backend_or); + ASSERT_NE(backend, nullptr); + + backend->Insert({"local_hash_1"}, + {RaidenBlockID(node_id, 10, BlockStatus::HOST)}, + /*on_host=*/true); + + auto load_future = backend->Load(RaidenId{}, {"local_hash_1"}, {5}); + EXPECT_OK(load_future.Await()); +} + +TEST(HostOffloadBackendTest, LoadLocalMissingBlockError) { + RaidenId node_id{"node_job", "0", "data", 0}; + rpc::RaidenIdProto unit_proto; + unit_proto.set_job_name(node_id.job_name); + unit_proto.set_job_replica_id(node_id.job_replica_id); + unit_proto.set_data_name(node_id.data_name); + unit_proto.set_data_replica_idx(node_id.data_replica_idx); + + controller::RaidenController controller(unit_proto, /*num_blocks=*/100, + /*num_shards=*/1, + /*shard_size_bytes=*/1024); + BackendConfig config; + config.type = "HostOffloadBackend"; + config.capacity = 100; + config.raiden_id = node_id; + + auto backend_or = HostOffloadBackend::Create(config, &controller); + ASSERT_OK(backend_or.status()); + auto backend = std::dynamic_pointer_cast(*backend_or); + ASSERT_NE(backend, nullptr); + + auto load_future = backend->Load(RaidenId{}, {"missing_hash"}, {5}); + EXPECT_THAT(load_future.Await(), + absl_testing::StatusIs(absl::StatusCode::kNotFound)); +} + +TEST(HostOffloadBackendTest, LoadLocalNonHostBlockError) { + RaidenId node_id{"node_job", "0", "data", 0}; + rpc::RaidenIdProto unit_proto; + unit_proto.set_job_name(node_id.job_name); + unit_proto.set_job_replica_id(node_id.job_replica_id); + unit_proto.set_data_name(node_id.data_name); + unit_proto.set_data_replica_idx(node_id.data_replica_idx); + + controller::RaidenController controller(unit_proto, /*num_blocks=*/100, + /*num_shards=*/1, + /*shard_size_bytes=*/1024); + BackendConfig config; + config.type = "HostOffloadBackend"; + config.capacity = 100; + config.raiden_id = node_id; + + auto backend_or = HostOffloadBackend::Create(config, &controller); + ASSERT_OK(backend_or.status()); + auto backend = std::dynamic_pointer_cast(*backend_or); + ASSERT_NE(backend, nullptr); + + RaidenId remote_id{"remote_job", "0", "data", 0}; + backend->Insert({"remote_hash"}, + {RaidenBlockID(remote_id, 10, BlockStatus::REMOTE)}, + /*on_host=*/true); + + auto load_future = backend->Load(RaidenId{}, {"remote_hash"}, {5}); + EXPECT_THAT(load_future.Await(), + absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition)); +} + TEST(HostOffloadBackendTest, StoreServerOverride) { RaidenId local_node_id{"override_job", "0", "cache", 0}; rpc::RaidenIdProto local_unit; diff --git a/tpu_raiden/kv_cache/kv_cache_store.cc b/tpu_raiden/kv_cache/kv_cache_store.cc index 0cde11d7..59be5851 100644 --- a/tpu_raiden/kv_cache/kv_cache_store.cc +++ b/tpu_raiden/kv_cache/kv_cache_store.cc @@ -874,10 +874,11 @@ absl::Status KVCacheStore::Load(const std::vector& block_hashes, if (!raiden_controller_) { return absl::FailedPreconditionError("RaidenController is not initialized"); } + if (block_hashes.empty()) { + return absl::OkStatus(); + } - std::vector src_host_block_ids; - src_host_block_ids.reserve(block_hashes.size()); - + RaidenId remote_id; { absl::MutexLock lock(mutex_); auto lookup_or = backend()->Lookup(block_hashes); @@ -887,18 +888,15 @@ absl::Status KVCacheStore::Load(const std::vector& block_hashes, return absl::NotFoundError( absl::StrCat("Block hash not found: ", block_hashes[slices.size()])); } + + BlockStatus first_status = slices[0].second.status; + if (first_status == BlockStatus::REMOTE) { + remote_id = slices[0].second.raiden_id; + } + for (size_t i = 0; i < slices.size(); ++i) { const auto& hash = block_hashes[i]; const auto& existing = slices[i].second; - if (existing.status != BlockStatus::HOST && - existing.status != BlockStatus::HOST_AND_HBM) { - return absl::FailedPreconditionError( - absl::StrCat("Block is not on host: ", hash)); - } - if (existing.host_block_id == -1) { - return absl::FailedPreconditionError( - absl::StrCat("Block host_block_id is -1: ", hash)); - } if (backend()->GetPinCount(hash) <= 0) { return absl::FailedPreconditionError( absl::StrCat("Block is not pinned: ", hash)); @@ -907,29 +905,35 @@ absl::Status KVCacheStore::Load(const std::vector& block_hashes, return absl::FailedPreconditionError( absl::StrCat("Block is already loading: ", hash)); } - src_host_block_ids.push_back(existing.host_block_id); + + if (first_status == BlockStatus::REMOTE) { + if (existing.status != BlockStatus::REMOTE) { + return absl::InvalidArgumentError( + "Mixed block statuses in a single Load call"); + } + if (existing.raiden_id != remote_id) { + return absl::InvalidArgumentError( + "Mixed remote node IDs in a single Load call"); + } + } else { + if (existing.status != BlockStatus::HOST && + existing.status != BlockStatus::HOST_AND_HBM) { + return absl::FailedPreconditionError( + absl::StrCat("Block is not on host: ", hash)); + } + if (existing.host_block_id == -1) { + return absl::FailedPreconditionError( + absl::StrCat("Block host_block_id is -1: ", hash)); + } + } } for (const auto& hash : block_hashes) { loading_hashes_.insert(hash); } } - std::vector src_buffers; - src_buffers.reserve(src_host_block_ids.size()); - for (int64_t id : src_host_block_ids) { - src_buffers.emplace_back(id, std::vector{}, std::nullopt, - rpc::MEMORY_TYPE_DRAM); - } - std::vector dst_buffers; - dst_buffers.reserve(device_block_ids.size()); - for (int id : device_block_ids) { - dst_buffers.emplace_back(id, std::vector{}, std::nullopt, - rpc::MEMORY_TYPE_HBM); - } - - tsl::Future<> future = raiden_controller_->TransferBuffers( - src_buffers, dst_buffers, /*staging_host_buffers=*/{}, - /*copy_sizes=*/{}); + tsl::Future<> future = backend()->Load(remote_id, block_hashes, + absl::MakeConstSpan(device_block_ids)); { absl::MutexLock lock(mutex_); @@ -1183,7 +1187,7 @@ size_t KVCacheStore::Evict(const std::vector& block_hashes) { } std::vector host_ids_to_deallocate; { - absl::MutexLock lock(&mutex_); + absl::MutexLock lock(mutex_); host_ids_to_deallocate = backend()->Evict(block_hashes); } @@ -1207,7 +1211,7 @@ size_t KVCacheStore::Evict(const std::vector& block_hashes) { absl::StatusOr> KVCacheStore::AllocateBlockIds(int needed) { std::vector hashes_to_deallocate; { - absl::MutexLock lock(&mutex_); + absl::MutexLock lock(mutex_); int free_count = raiden_controller_->block_manager()->num_free_blocks(); int to_free = needed - free_count; if (to_free > 0) { @@ -1620,7 +1624,13 @@ void KVCacheStore::PollLoadsInternal(std::vector ready_loads) { if (i < slices.size()) { RaidenBlockID block = slices[i].second; block.device_block_id = state.device_block_ids[i]; - block.status = BlockStatus::HOST_AND_HBM; + if (block.status == BlockStatus::REMOTE) { + block.raiden_id = raiden_id_; + block.host_block_id = -1; + block.status = BlockStatus::HBM; + } else { + block.status = BlockStatus::HOST_AND_HBM; + } update_hashes.push_back(hash); update_slices.push_back(block); } diff --git a/tpu_raiden/kv_cache/kv_cache_store_backend.h b/tpu_raiden/kv_cache/kv_cache_store_backend.h index fe066c38..f763b874 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_backend.h +++ b/tpu_raiden/kv_cache/kv_cache_store_backend.h @@ -25,6 +25,7 @@ #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/tsl/concurrency/future.h" #include "tpu_raiden/kv_cache/raiden_id.h" namespace tpu_raiden { @@ -109,6 +110,12 @@ class KVCacheStoreBackend { absl::Span block_hashes, const LookupOptions& options = {}) = 0; + // Asynchronously loads KV cache blocks from a peer node into local + // storage / device. + virtual tsl::Future<> Load( + const RaidenId& remote_id, absl::Span block_hashes, + absl::Span device_block_ids = {}) = 0; + // Inserts key-block mappings into the backend. // Returns: // - bool: true if all hashes were newly inserted (none already existed) diff --git a/tpu_raiden/kv_cache/kv_cache_store_backend_factory_test.cc b/tpu_raiden/kv_cache/kv_cache_store_backend_factory_test.cc index e63e5204..bc5771cb 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_backend_factory_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_backend_factory_test.cc @@ -14,6 +14,7 @@ #include "tpu_raiden/kv_cache/kv_cache_store_backend_factory.h" +#include #include #include #include @@ -21,6 +22,7 @@ #include #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "xla/tsl/concurrency/future.h" #include "tpu_raiden/kv_cache/host_offload_backend.h" #include "tpu_raiden/kv_cache/kv_cache_store.h" #include "tpu_raiden/kv_cache/kv_cache_store_backend.h" @@ -37,6 +39,11 @@ class CustomTestBackend : public KVCacheStoreBackend { const LookupOptions&) override { return BlockSliceList{}; } + tsl::Future<> Load(const RaidenId& remote_id, + absl::Span block_hashes, + absl::Span device_block_ids = {}) override { + return tsl::Future<>(absl::UnimplementedError("Load is not supported.")); + } std::pair Insert(absl::Span, absl::Span, bool) override { diff --git a/tpu_raiden/kv_cache/kv_cache_store_test.cc b/tpu_raiden/kv_cache/kv_cache_store_test.cc index 4785bd14..99642305 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_test.cc @@ -1353,6 +1353,119 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadSuccess) { EXPECT_EQ((*lookup_res)[1].second.device_block_id, 3); } +TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteSuccess) { + // 1. Setup GlobalRegistry server + auto service = std::make_unique(); + grpc::ServerBuilder registry_builder; + int registry_port = 0; + registry_builder.AddListeningPort( + "localhost:0", grpc::InsecureServerCredentials(), ®istry_port); + registry_builder.RegisterService(service.get()); + auto registry_server = registry_builder.BuildAndStart(); + std::string registry_address = "localhost:" + std::to_string(registry_port); + + RaidenId local_rid{"local_job", "0", "local_cache", 0}; + RaidenId remote_rid{"remote_job", "0", "remote_cache", 0}; + + // 2. Setup local RaidenController & KVCacheStore + auto controller = + std::make_unique<::tpu_raiden::controller::RaidenController>( + unit_, 10, 1, 512, orchestrator_address_, ""); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + // 3. Setup remote node's backend & server + BackendConfig remote_config; + remote_config.type = "HostOffloadBackend"; + remote_config.capacity = 100; + remote_config.global_registry_address = registry_address; + remote_config.raiden_id = remote_rid; + + auto remote_backend_or = + HostOffloadBackend::Create(remote_config, controller.get()); + ASSERT_OK(remote_backend_or.status()); + auto remote_backend = + std::dynamic_pointer_cast(*remote_backend_or); + ASSERT_NE(remote_backend, nullptr); + + std::vector remote_slices = { + RaidenBlockID(remote_rid, 42, BlockStatus::HOST), + }; + remote_backend->Insert({"load_remote_hash_1"}, remote_slices, + /*on_host=*/true); + + auto remote_server = KVCacheStoreServer::Create(); + ASSERT_OK(remote_server->StartServer(remote_backend.get(), controller.get(), + "127.0.0.1")); + + auto channel = + grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); + auto registry_client = + std::make_shared(channel); + ASSERT_OK(registry_client->RegisterStore( + remote_rid, remote_server->GetServerAddress(), orchestrator_address_)); + + // 4. Create store and insert remote block entry + KVCacheStore store(10, std::move(controller), registry_address, local_rid, + std::nullopt, /*store_server_ip=*/"127.0.0.1"); + + std::vector hashes = {"load_remote_hash_1"}; + std::vector slices = { + RaidenBlockID(remote_rid, 42, BlockStatus::REMOTE)}; + + ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/false).first); + ASSERT_TRUE(store.Pin(hashes)); + + // 5. Load remote block into local device block 5 + absl::Status status = store.Load(hashes, {5}); + ASSERT_TRUE(status.ok()) << status.message(); + + // 6. Poll for completion + bool done = false; + for (int attempt = 0; attempt < 100; ++attempt) { + auto [load_done, load_failed, load_pending] = store.PollLoadStatus(); + ASSERT_TRUE(load_failed.empty()); + if (!load_done.empty()) { + EXPECT_THAT(load_done, + ::testing::UnorderedElementsAre("load_remote_hash_1")); + done = true; + break; + } + absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(done); + + // 7. Verify status in store is updated to HBM and device_block_id is 5 + auto lookup_res = store.Lookup(hashes); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HBM); + EXPECT_EQ((*lookup_res)[0].second.host_block_id, -1); + EXPECT_EQ((*lookup_res)[0].second.device_block_id, 5); + EXPECT_EQ((*lookup_res)[0].second.raiden_id, local_rid); +} + +TEST_F(KVCacheStoreEmbeddedControllerTest, LoadUnpinnedRemoteBlockFails) { + auto controller = + std::make_unique<::tpu_raiden::controller::RaidenController>( + unit_, 10, 1, 512, orchestrator_address_, ""); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + RaidenId local_rid{"local_job", "0", "local_cache", 0}; + RaidenId remote_rid{"remote_job", "0", "remote_cache", 0}; + KVCacheStore store(10, std::move(controller), "", local_rid, std::nullopt, + /*store_server_ip=*/"127.0.0.1"); + + std::vector hashes = {"unpinned_remote_hash"}; + std::vector slices = { + RaidenBlockID(remote_rid, /*host_block_id=*/-1, /*device_block_id=*/-1, + BlockStatus::REMOTE)}; + ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/false).first); + + absl::Status status = store.Load(hashes, {0}); + EXPECT_TRUE(absl::IsFailedPrecondition(status)) << status; + EXPECT_THAT(status.message(), ::testing::HasSubstr("is not pinned")); +} + TEST_F(KVCacheStoreEmbeddedControllerTest, SaveMultiWorkerSuccess) { auto test_server_0 = ::tpu_raiden::controller::CreateTestWorkerServer(); auto test_server_1 = ::tpu_raiden::controller::CreateTestWorkerServer();