From d5ac7cafd7356ac66846fb084d8ebc7975126909 Mon Sep 17 00:00:00 2001 From: jcgu Date: Mon, 17 Aug 2026 14:58:57 -0700 Subject: [PATCH] Remote loads record nothing locally & local loads consume pins PiperOrigin-RevId: 966199810 --- tpu_sync/api/jax/kv_cache_store.py | 22 ++- tpu_sync/api/jax/kv_cache_store_e2e_test.py | 12 +- tpu_sync/api/jax/kv_cache_store_test.py | 9 +- tpu_sync/api/torch/kv_cache_store.py | 22 ++- tpu_sync/api/torch/kv_cache_store_test.py | 9 +- tpu_sync/kv_cache/kv_cache_store.cc | 50 ++++-- tpu_sync/kv_cache/kv_cache_store.h | 37 +++-- tpu_sync/kv_cache/kv_cache_store_test.cc | 161 ++++++++++++++++++-- 8 files changed, 261 insertions(+), 61 deletions(-) diff --git a/tpu_sync/api/jax/kv_cache_store.py b/tpu_sync/api/jax/kv_cache_store.py index 0b1229c5..0a0934b1 100644 --- a/tpu_sync/api/jax/kv_cache_store.py +++ b/tpu_sync/api/jax/kv_cache_store.py @@ -379,11 +379,23 @@ def load( `device_block_ids` is the destination and must name one device block per hash. - NOTE: The block_hashes must be pinned in the LRU cache before calling load - when loading from local host. Once the operation is complete (as reported by - poll_load_status), the caller must manually release/unpin them. - Blocks provided in `slices` must be already pinned externally, and remote - loads will re-resolve hashes at the peer, ignoring `slices`. + PIN CONTRACT: + local source -- every hash must be pinned on entry (lookup() is what + normally grants that pin), and a SUCCESSFUL load + consumes exactly one pin per hash. Do not release + afterwards. A FAILED load does not consume it: the entry + stays pinned so you can retry, or release it + deliberately. Giving up is your decision, not the + store's. + remote source -- no pin is required and none is consumed. A hash resolved + only through the registry never entered the local cache, + so there is nothing here to have pinned. + + A load from a peer records NOTHING locally: no host copy is kept, so a + later lookup() of that hash is still a miss. Your own block manager is what + remembers you already own the device block. + + Remote loads re-resolve hashes at the peer, ignoring the rest of `slices`. Args: block_hashes: List of block hashes to load. diff --git a/tpu_sync/api/jax/kv_cache_store_e2e_test.py b/tpu_sync/api/jax/kv_cache_store_e2e_test.py index a636a171..5df0b383 100644 --- a/tpu_sync/api/jax/kv_cache_store_e2e_test.py +++ b/tpu_sync/api/jax/kv_cache_store_e2e_test.py @@ -777,14 +777,10 @@ def _run_remote_read_to_hbm_test(self, enable_multi_numa: bool, use_slices: bool self._await_terminal( store_b.poll_load_status, len(hashes), "Job B peer-fetch load" ) - # The batch commits as a unit; the entry records where in HBM it landed. - after = store_b.lookup(hashes) - - store_b.release(hashes) - self.assertLen(after, 2) - for i, (_, blk) in enumerate(after): - self.assertEqual(blk.status, kv_cache_store.BlockStatus.HBM) - self.assertEqual(blk.device_block_id, dst_device_blocks[i]) + # A load from a peer records nothing locally: no host copy was kept, so + # there is no residency to describe. The bytes are in the device blocks + # the caller named and the cache is a miss for these hashes. + self.assertEmpty(store_b.lookup(hashes)) else: # --- The thing under test: pull straight into HBM. --------------------- # No insert first: the lookup answer IS the source coordinate, and the diff --git a/tpu_sync/api/jax/kv_cache_store_test.py b/tpu_sync/api/jax/kv_cache_store_test.py index c59df959..c75b4af9 100644 --- a/tpu_sync/api/jax/kv_cache_store_test.py +++ b/tpu_sync/api/jax/kv_cache_store_test.py @@ -234,8 +234,7 @@ def test_pin_and_release(self): controller.insert(hash_4, slice_4, True) res_old = controller.lookup([b"7001", b"7002"]) - self.assertLen(res_old, 2) - controller.release([b"7001", b"7002"]) + self.assertEmpty(res_old) res = controller.lookup([b"7002"]) self.assertLen(res, 1) controller.release([b"7002"]) @@ -269,8 +268,7 @@ def test_partial_pin_rollback(self): ) res_old = controller.lookup([b"8001", b"8002"]) - self.assertLen(res_old, 2) - controller.release([b"8001", b"8002"]) + self.assertEmpty(res_old) res = controller.lookup([b"8004", b"8005"]) self.assertLen(res, 2) controller.release([b"8004", b"8005"]) @@ -473,8 +471,7 @@ def test_insert_and_lock_release_and_delete(self): success = controller.insert_and_lock(remote_hashes, remote_slices, True) self.assertTrue(success) res_local = controller.lookup([b"local_1"]) - self.assertLen(res_local, 1) - controller.release([b"local_1"]) + self.assertEmpty(res_local) del_count = controller.release_and_delete(remote_hashes) self.assertEqual(del_count, 2) diff --git a/tpu_sync/api/torch/kv_cache_store.py b/tpu_sync/api/torch/kv_cache_store.py index d271aee6..facd49a8 100644 --- a/tpu_sync/api/torch/kv_cache_store.py +++ b/tpu_sync/api/torch/kv_cache_store.py @@ -438,11 +438,23 @@ def load( `device_block_ids` is the destination and must name one device block per hash. - NOTE: The block_hashes must be pinned in the LRU cache before calling load - when loading from local host. Once the operation is complete (as reported by - poll_load_status), the caller must manually release/unpin them. - Blocks provided in `slices` must be already pinned externally, and remote - loads will re-resolve hashes at the peer, ignoring `slices`. + PIN CONTRACT: + local source -- every hash must be pinned on entry (lookup() is what + normally grants that pin), and a SUCCESSFUL load + consumes exactly one pin per hash. Do not release + afterwards. A FAILED load does not consume it: the entry + stays pinned so you can retry, or release it + deliberately. Giving up is your decision, not the + store's. + remote source -- no pin is required and none is consumed. A hash resolved + only through the registry never entered the local cache, + so there is nothing here to have pinned. + + A load from a peer records NOTHING locally: no host copy is kept, so a + later lookup() of that hash is still a miss. Your own block manager is what + remembers you already own the device block. + + Remote loads re-resolve hashes at the peer, ignoring the rest of `slices`. Args: block_hashes: List of block hashes to load. diff --git a/tpu_sync/api/torch/kv_cache_store_test.py b/tpu_sync/api/torch/kv_cache_store_test.py index 79416b67..c0e38764 100644 --- a/tpu_sync/api/torch/kv_cache_store_test.py +++ b/tpu_sync/api/torch/kv_cache_store_test.py @@ -175,8 +175,7 @@ def test_pin_and_release(self): controller.insert(hash_4, slice_4, True) res_old = controller.lookup([b"7001", b"7002"]) - self.assertLen(res_old, 2) - controller.release([b"7001", b"7002"]) + self.assertEmpty(res_old) res = controller.lookup([b"7002"]) self.assertLen(res, 1) controller.release([b"7002"]) @@ -210,8 +209,7 @@ def test_partial_pin_rollback(self): ) res_old = controller.lookup([b"8001", b"8002"]) - self.assertLen(res_old, 2) - controller.release([b"8001", b"8002"]) + self.assertEmpty(res_old) res = controller.lookup([b"8004", b"8005"]) self.assertLen(res, 2) controller.release([b"8004", b"8005"]) @@ -375,8 +373,7 @@ def test_insert_and_lock_release_and_delete(self): success = controller.insert_and_lock(remote_hashes, remote_slices, True) self.assertTrue(success) res_local = controller.lookup([b"local_1"]) - self.assertLen(res_local, 1) - controller.release([b"local_1"]) + self.assertEmpty(res_local) del_count = controller.release_and_delete(remote_hashes) self.assertEqual(del_count, 2) diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index 52c72126..6b247afd 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -1080,6 +1080,7 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, } RaidenId remote_id; + bool from_remote = false; { absl::MutexLock lock(mutex_); auto lookup_or = backend()->Lookup(block_hashes); @@ -1093,6 +1094,7 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, BlockStatus first_status = slices[0].second.status; if (first_status == BlockStatus::REMOTE) { remote_id = slices[0].second.raiden_id; + from_remote = true; } for (size_t i = 0; i < slices.size(); ++i) { @@ -1147,6 +1149,7 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, std::vector(block_hashes.begin(), block_hashes.end()), .device_block_ids = std::vector(device_block_ids.begin(), device_block_ids.end()), + .from_remote = from_remote, }); } @@ -1169,12 +1172,14 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, } RaidenId remote_id; + bool from_remote = false; { absl::MutexLock lock(mutex_); BlockStatus first_status = slices[0].status; if (first_status == BlockStatus::REMOTE) { remote_id = slices[0].raiden_id; + from_remote = true; } for (size_t i = 0; i < slices.size(); ++i) { @@ -1195,6 +1200,13 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, "Mixed remote node IDs in a single Load call"); } } else { + // The caller's pin is what a successful local load consumes, so it has + // to exist. The no-slices form has always required it; this form did + // not, which left one signature hiding two different pin contracts. + if (backend()->GetPinCount(hash) <= 0) { + return absl::FailedPreconditionError( + absl::StrCat("Block is not pinned: ", hash)); + } if (existing.status != BlockStatus::HOST && existing.status != BlockStatus::HOST_AND_HBM) { return absl::FailedPreconditionError( @@ -1226,6 +1238,7 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, std::vector(block_hashes.begin(), block_hashes.end()), .device_block_ids = std::vector(device_block_ids.begin(), device_block_ids.end()), + .from_remote = from_remote, }); } @@ -1967,8 +1980,25 @@ void KVCacheStore::PollLoadsInternal(std::vector ready_loads) { for (auto& state : ready_loads) { absl::Status status = state.future.Await(); absl::MutexLock lock(mutex_); - if (status.ok()) { - auto lookup_or = backend()->Lookup(state.block_hashes); + if (status.ok() && state.from_remote) { + // A load from a peer records NOTHING locally. The bytes went to the + // caller's device blocks and no local host copy was kept, so there is no + // residency to describe: an entry here would claim HBM with + // host_block_id -1, which eviction cannot reclaim (it only takes HOST and + // HOST_AND_HBM) and which nothing left in the API can delete. + // + // The consequence is deliberate: a later lookup() of the same hash is a + // miss, and a repeat request re-fetches unless the caller's own block + // manager remembers it already owns the device block. + for (const auto& hash : state.block_hashes) { + done_loads_.push_back(hash); + } + } else if (status.ok()) { + // Local source: the entry exists here by construction, so this lookup is + // purely local -- no registry fallback, which would otherwise put a + // blocking RPC inside the poller while it holds mutex_. + auto lookup_or = backend()->Lookup(state.block_hashes, + LookupOptions{.enable_global = false}); if (lookup_or.ok()) { const auto& slices = lookup_or.value(); std::vector update_hashes; @@ -1978,13 +2008,7 @@ 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]; - 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; - } + block.status = BlockStatus::HOST_AND_HBM; update_hashes.push_back(hash); update_slices.push_back(block); } @@ -1992,6 +2016,14 @@ void KVCacheStore::PollLoadsInternal(std::vector ready_loads) { } if (!update_hashes.empty()) { backend()->Insert(update_hashes, update_slices, /*on_host=*/true); + // The load is done with the block, so the pin the caller acquired to + // keep it alive across the transfer is consumed here. Released AFTER + // the index update, so the entry cannot be evicted between the two. + // + // Only on success, and only for a local source: a failed load stays + // pinned so the caller can retry or release deliberately, and a + // remote load never had a caller pin to consume. + backend()->Release(update_hashes); } } } else { diff --git a/tpu_sync/kv_cache/kv_cache_store.h b/tpu_sync/kv_cache/kv_cache_store.h index 00c041f6..2ae1bbe4 100644 --- a/tpu_sync/kv_cache/kv_cache_store.h +++ b/tpu_sync/kv_cache/kv_cache_store.h @@ -274,9 +274,13 @@ class KVCacheStore { // `device_block_ids` is the destination and must name one device block per // hash. // - // NOTE: The block_hashes must be pinned in the LRU cache before calling Load. - // Once the operation is complete (as reported by PollLoadStatus), the caller - // must manually release/unpin them via Release. + // PIN CONTRACT: every hash must be pinned on entry -- Lookup() is what + // normally grants that pin -- and a SUCCESSFUL load consumes exactly one pin + // per hash. The caller does not release afterwards. + // + // A FAILED load does not: the entry stays pinned so the caller can retry, or + // release it deliberately. Deciding to give up is the caller's, not this + // store's. absl::Status Load(absl::Span block_hashes, absl::Span device_block_ids); @@ -290,9 +294,16 @@ class KVCacheStore { // hash. // // If `slices` is non-empty, the caller's pre-looked up RaidenBlockIDs are - // used directly. Note that blocks in `slices` must be already pinned - // externally (when Load from local host), and remote loads will re-resolve - // hashes at the peer, ignoring `slices`. + // used directly. Remote loads re-resolve hashes at the peer, ignoring the + // rest of `slices`. + // + // PIN CONTRACT, same as the overload above and now enforced the same way: + // local source -- every hash must be pinned on entry, and a successful + // load consumes one pin per hash. + // remote source -- no pin is required and none is consumed. A hash + // resolved only through the registry never entered the + // local index, so there is nothing here to have pinned, + // and a load from a peer records nothing either. absl::Status Load(absl::Span block_hashes, absl::Span slices, absl::Span device_block_ids); @@ -353,10 +364,14 @@ class KVCacheStore { // Polls the status of all active/inflight Load operations. // Updates cache metadata upon successful H2D transfers: // - Loaded from local host DRAM -> HOST_AND_HBM - // - Loaded from a peer -> HBM, with host_block_id -1. + // - Loaded from a peer -> nothing is recorded at all. // - // Note: HBM-only entries hold a slot in the LRU but own no host block, and - // Evict only reclaims HOST and HOST_AND_HBM entries. They must be explicitly deleted. + // A peer load leaves no entry because there is nothing here to describe: no + // local host copy is kept, so the entry could only say HBM with + // host_block_id -1 -- which Evict cannot reclaim (it takes HOST and + // HOST_AND_HBM only) and which nothing would ever remove. A later lookup() + // of such a hash is therefore a miss, and the caller's own block manager is + // what remembers it already owns the device block. // // Returns: // A tuple of {done_block_hashes, failed_block_hashes, pending_block_hashes} @@ -507,6 +522,10 @@ class KVCacheStore { tsl::Future<> future; std::vector block_hashes; std::vector device_block_ids; + // Whether the source was a peer. Decided at submit time and carried here + // because the poller cannot re-derive it: a remote load records nothing + // locally, so by completion there is no entry to read a status off. + bool from_remote = false; }; struct RemoteReadState { diff --git a/tpu_sync/kv_cache/kv_cache_store_test.cc b/tpu_sync/kv_cache/kv_cache_store_test.cc index 9e8c140f..ebf4aab3 100644 --- a/tpu_sync/kv_cache/kv_cache_store_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_test.cc @@ -1489,7 +1489,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesSizeMismatch) { EXPECT_THAT(std::string(status.message()), ::testing::HasSubstr("mismatch")); } -TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesUnpinnedSucceeds) { +// The slices form used to accept an unpinned local block, while the no-slices +// form required a pin -- one signature, two pin contracts. It requires the pin +// now, for the same reason the other form always did: a successful local load +// CONSUMES one, so there has to be one to consume. +TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesUnpinnedFails) { ::tpu_raiden::controller::MockTransferManager mock_mgr; test_server_->service->SetTransferManager( ::tpu_raiden::KVManagerHolder(&mock_mgr)); @@ -1508,7 +1512,57 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesUnpinnedSucceeds) { ASSERT_TRUE(store.Insert(hashes, slices, true).first); absl::Status status = store.Load(hashes, slices, {2}); - EXPECT_TRUE(status.ok()) << status.message(); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("not pinned")); + + // With the pin the caller was supposed to hold, it goes through. + ASSERT_TRUE(store.Pin(hashes)); + EXPECT_TRUE(store.Load(hashes, slices, {2}).ok()); +} + +// A successful local load consumes the caller's pin, so the caller never +// releases after a load. A failed one does not: giving up is the caller's +// decision, and an entry silently unpinned under a retry would be evictable +// while the caller still believed it held it. +TEST_F(KVCacheStoreEmbeddedControllerTest, LocalLoadConsumesTheCallerPin) { + ::tpu_raiden::controller::MockTransferManager mock_mgr; + test_server_->service->SetTransferManager( + ::tpu_raiden::KVManagerHolder(&mock_mgr)); + + auto controller = MakeController(10, 1, 512, ""); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + RaidenId rid{"test_job", "0", "test_cache", 0}; + KVCacheStore store(10, std::move(controller), "", rid, std::nullopt, + /*store_server_ip=*/"127.0.0.1"); + + std::vector hashes = {"hash_1"}; + std::vector slices = { + RaidenBlockID(rid, 0, -1, BlockStatus::HOST)}; + ASSERT_TRUE(store.Insert(hashes, slices, true).first); + ASSERT_TRUE(store.Pin(hashes)); + ASSERT_EQ(store.GetPinCount("hash_1"), 1); + + ASSERT_OK(store.Load(hashes, {2})); + + bool done = false; + for (int attempt = 0; attempt < 100 && !done; ++attempt) { + auto [load_done, load_failed, load_pending] = store.PollLoadStatus(); + ASSERT_TRUE(load_failed.empty()); + if (!load_done.empty()) done = true; + if (!done) absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(done); + + EXPECT_EQ(store.GetPinCount("hash_1"), 0) + << "a successful local load must consume the caller's pin"; + // The entry survives the unpin and carries the load's result. + auto after = PeekLookup(store, hashes); + ASSERT_TRUE(after.ok()); + ASSERT_EQ(after->size(), 1); + EXPECT_EQ((*after)[0].second.status, BlockStatus::HOST_AND_HBM); + EXPECT_EQ((*after)[0].second.device_block_id, 2); } TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesAlreadyLoadingFails) { @@ -1627,13 +1681,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesRemoteSuccess) { } ASSERT_TRUE(done); - auto lookup_res = store.Lookup(hashes); + // Nothing is recorded for a peer source, so the entry this test inserted up + // front is left exactly as it was: still REMOTE, still naming the peer. + auto lookup_res = PeekLookup(store, 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); + EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::REMOTE); + EXPECT_EQ((*lookup_res)[0].second.host_block_id, 42); + EXPECT_EQ((*lookup_res)[0].second.raiden_id, remote_rid); } TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteSuccess) { @@ -1710,14 +1765,94 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteSuccess) { } ASSERT_TRUE(done); - // 7. Verify status in store is updated to HBM and device_block_id is 5 - auto lookup_res = store.Lookup(hashes); + // 7. A load from a peer records NOTHING. This entry was put here by the + // caller before the load, and the load leaves it exactly as it found it -- + // still REMOTE, still naming the peer's block 42. Promoting it to HBM with + // host_block_id -1, as this used to, produced an entry describing no local + // residency that Evict could not reclaim and nothing could delete. + auto lookup_res = PeekLookup(store, 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); + EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::REMOTE); + EXPECT_EQ((*lookup_res)[0].second.host_block_id, 42); + EXPECT_EQ((*lookup_res)[0].second.raiden_id, remote_rid); +} + +// The flow the API actually serves for a peer source: lookup() resolves the +// hash through the registry and hands back a REMOTE slice, load() takes that +// slice directly. Nothing is inserted before, and -- the point of this case -- +// nothing is left after. The local cache is untouched from start to finish. +TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteWithSlicesRecordsNothing) { + auto registry_server = global_registry::CreateTestGlobalRegistryServer(); + std::string registry_address = registry_server->server_address; + + RaidenId local_rid{"local_job", "0", "local_cache", 0}; + RaidenId remote_rid{"remote_job", "0", "remote_cache", 0}; + + auto controller = MakeController(10, 1, 512, ""); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + 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); + remote_backend->Insert({"slice_load_hash"}, + {RaidenBlockID(remote_rid, 42, BlockStatus::HOST)}, + /*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(), + controller->controller_address())); + + KVCacheStore store(10, std::move(controller), registry_address, local_rid, + std::nullopt, /*store_server_ip=*/"127.0.0.1"); + + std::vector hashes = {"slice_load_hash"}; + + // The registry answers, but a registry-only hit never enters the local index + // and is never pinned -- so this load has no caller pin to consume either. + auto resolved = store.Lookup(hashes, /*enable_global=*/true); + ASSERT_TRUE(resolved.ok()); + ASSERT_EQ(resolved->size(), 1); + EXPECT_EQ((*resolved)[0].second.status, BlockStatus::REMOTE); + EXPECT_TRUE(PeekLookup(store, hashes)->empty()) + << "a registry-only hit must not have entered the local index"; + + std::vector slices = {(*resolved)[0].second}; + ASSERT_OK(store.Load(hashes, slices, {5})); + + bool done = false; + for (int attempt = 0; attempt < 100 && !done; ++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("slice_load_hash")); + done = true; + break; + } + absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(done); + + // The whole point: the bytes are in device block 5 and the cache is as empty + // as it was before the load. + EXPECT_TRUE(PeekLookup(store, hashes)->empty()) + << "a load from a peer must leave no local entry"; } TEST_F(KVCacheStoreEmbeddedControllerTest, LoadUnpinnedRemoteBlockFails) {