diff --git a/tpu_sync/api/jax/kv_cache_store.py b/tpu_sync/api/jax/kv_cache_store.py index dfe8a9f1..0a0934b1 100644 --- a/tpu_sync/api/jax/kv_cache_store.py +++ b/tpu_sync/api/jax/kv_cache_store.py @@ -215,9 +215,20 @@ def lookup( ) -> list[tuple[bytes, RaidenBlockID]]: """Checks the LRU directory for cached block hashes. + PINS every hash it returns, one pin each, so the answer cannot be evicted + between being given and being used. The operation the caller goes on to + perform consumes that pin: load() drops it on a successful local load, and + save() drops it on success. A caller that does neither -- one that only + wanted to know what is resident -- must release() what it was given. + + Locally cached hits are pinned. Hashes resolved only through the global + registry are not: they name a block on another node, and there is nothing + here to hold. + Args: block_hashes: Incoming block hashes to check. - enable_global: Whether to fallback to global registry on miss. + enable_global: Whether to fallback to global registry on miss. Defaults + to False; the fallback is a blocking RPC. Returns: A list of tuples containing the block hash and the matching @@ -368,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. @@ -424,35 +447,49 @@ def poll_load_status(self) -> tuple[list[bytes], list[bytes], list[bytes]]: def read_remote( self, block_hashes: list[bytes], - device_block_ids: list[int] | None = None, + slices: list[RaidenBlockID], + device_block_ids: list[int], ) -> bool: - """Launches an async receiver-initiated read of REMOTE blocks from peers. + """Reads REMOTE blocks from their owning peers straight into local HBM. Returns as soon as the reads are issued; poll with poll_remote_read_status(). + This store's cache is neither consulted nor modified. The hashes need not + be present locally and need not be pinned; nothing is inserted on success + and nothing is left behind on failure. The bytes land ONLY in the given + device blocks -- no local host copy is kept, so a later local load() of the + same hash is still a miss. + + Compare with load(): both bring a peer's block into local HBM. Use load() + when the hash may be resident locally and you want the store to decide; use + read_remote() when you already hold the source coordinates and want no + local record of the transfer. + Requires a global registry: it is what maps the owning peer to the controller address this store acquires its read lease from. A store built without a global_registry_address fails every read. Args: - block_hashes: Block hashes to read. Each must already be pinned, and must - stay pinned until poll_remote_read_status() reports it terminal. - Releasing early makes the entry deletable mid-read, in which case the - read is discarded and the WHOLE batch is reported failed. - device_block_ids: Optional. Omit (or pass None) to read into host DRAM, - leaving the entries HOST. Pass one device block id per hash to read - straight into HBM, leaving the entries HOST_AND_HBM (the host landing - blocks act as the staging hop, so a later load() can reuse them). - On FAILURE the contents of these device blocks are UNDEFINED: they are - written before the source's verdict is known. Treat them as scratch - until the read reports success -- nothing in the cache points at them - unless it does. + block_hashes: Block hashes to read. + slices: One REMOTE RaidenBlockID per hash, naming where to read from. + Only two fields are used -- raiden_id (the owning peer) and + host_block_id (the block on that peer) -- so a lookup() answer can be + passed straight through. + device_block_ids: One local device block id per hash. Mandatory. + On FAILURE their contents are UNDEFINED: they are written before the + source's verdict is known. Treat them as scratch until the read + reports success. Returns: True if successfully launched. """ - return self._impl.read_remote(block_hashes, device_block_ids or []) + raw_slices = [] + for s in slices: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + raw_slices.append(s._impl) # pylint: disable=protected-access + return self._impl.read_remote(block_hashes, raw_slices, device_block_ids) def poll_remote_read_status( self, 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 cd8569c7..698cc954 100644 --- a/tpu_sync/api/jax/kv_cache_store_e2e_test.py +++ b/tpu_sync/api/jax/kv_cache_store_e2e_test.py @@ -261,6 +261,7 @@ def _run_e2e_test(self, enable_multi_numa: bool, use_slices: bool = False): # Verify status in store is HBM lookup_res = store.lookup(hashes) + store.release(hashes) self.assertLen(lookup_res, 2) self.assertEqual(lookup_res[0][1].status, kv_cache_store.BlockStatus.HBM) self.assertEqual(lookup_res[0][1].device_block_id, 0) @@ -290,10 +291,10 @@ def get_slice_e2e(x): time.sleep(0.01) # Release them so we can test pinning before load - store.release(hashes) # Verify status in store is updated to HOST_AND_HBM lookup_res = store.lookup(hashes) + store.release(hashes) self.assertLen(lookup_res, 2) self.assertEqual( lookup_res[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -340,7 +341,6 @@ def get_slice_e2e(x): time.sleep(0.01) # Release at the very end - store.release(hashes) # 9. Verify device memory contains the original random data np.testing.assert_array_equal(np.asarray(tpu_cache), expected_ref) @@ -517,7 +517,6 @@ def get_slice(x): data_a = manager_a._impl.read_host_memory(0, 0, 16) print(f"DEBUG: Job A host memory (layer 0, shard 0) after Save: {data_a}") - store_a.release(hashes) # 5. Job B calls Lookup (enable_global=True) # Give some time for registry propagation @@ -540,6 +539,7 @@ def get_slice(x): # Verify correct source host block IDs lookup_res_a = store_a.lookup(hashes) + store_a.release(hashes) self.assertEqual( lookup_res_b[0][1].host_block_id, lookup_res_a[0][1].host_block_id ) @@ -577,12 +577,11 @@ def get_slice(x): if not done: time.sleep(0.01) else: - # 6. Job B controller calls insert_and_lock for the remote slices - self.assertTrue(store_b.insert_and_lock(hashes, slices_b, on_host=True)) - - # 7. Job B calls ReadRemote - self.assertTrue(store_b.read_remote(hashes)) - + # 6. Job B reads straight from Job A into its own device blocks. The + # source coordinates come from the lookup answer, so nothing needs to be + # inserted into Job B's cache first. + self.assertTrue(store_b.read_remote(hashes, slices_b, [0, 1])) + if not expect_read_success: # Strict node_id matching: the producer worker's node_id must equal the # consumer (destination) worker's node_id. A mismatch makes the source @@ -611,40 +610,15 @@ def get_slice(x): if not done: time.sleep(0.01) - data_b = manager_b._impl.read_host_memory(0, 0, 16) - print( - "DEBUG: Job B host memory (layer 0, shard 0) after ReadRemote:" - f" {data_b}" - ) - - # 8. Verify Job B's LRU block status becomes HOST - lookup_res_b_after = store_b.lookup(hashes) - self.assertLen(lookup_res_b_after, 2) - self.assertEqual( - lookup_res_b_after[0][1].status, kv_cache_store.BlockStatus.HOST - ) - self.assertEqual( - lookup_res_b_after[1][1].status, kv_cache_store.BlockStatus.HOST - ) - - # 9. Job B controller calls Load to transfer data to TPU blocks - self.assertTrue(store_b.load(hashes, [0, 1])) - - # Wait for Load completion - done = False - while not done: - load_done, load_failed, _ = store_b.poll_load_status() - if load_failed: - raise RuntimeError(f"Job B Load failed: {load_failed}") - if len(load_done) == 2: - done = True - if not done: - time.sleep(0.01) + # 8. The read is already in HBM -- there is no second Load step, and no + # local record of it either. Job B's cache is still a miss for these + # hashes: the bytes live only in the device blocks it named. + self.assertEmpty(store_b.lookup(hashes)) - store_b.release(hashes) - # 10. Verify byte-exact match on Job B TPU devices - np.testing.assert_array_equal(np.asarray(tpu_cache_b), host_data_a) + # 9. Verify byte-exact match on Job B TPU devices. The DMA landed behind + # JAX's back, so the buffer has to be re-read rather than np.asarray'd. + np.testing.assert_array_equal(self._reread_device(tpu_cache_b), host_data_a) # ========================================================================= # ReadRemote to HBM (receiver-initiated pull straight into device memory) @@ -777,7 +751,6 @@ def _run_remote_read_to_hbm_test(self, enable_multi_numa: bool, use_slices: bool self.assertTrue(store_a.pin(hashes)) store_a.save(hashes) self._await_terminal(store_a.poll_save_status, len(hashes), "Job A save") - store_a.release(hashes) # --- Job B: discover the blocks as REMOTE. ----------------------------- time.sleep(0.5) @@ -796,28 +769,27 @@ 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" ) + # 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: - # insert_and_lock pins the entries; the contract requires holding those - # pins until poll_remote_read_status reports the hashes terminal. + # --- The thing under test: pull straight into HBM. --------------------- + # No insert first: the lookup answer IS the source coordinate, and the + # read takes no pin because it records nothing. self.assertTrue( - store_b.insert_and_lock(hashes, [b for _, b in lookup_b], on_host=True) + store_b.read_remote( + hashes, [b for _, b in lookup_b], dst_device_blocks + ) ) - - # --- The thing under test: pull straight into HBM. --------------------- - self.assertTrue(store_b.read_remote(hashes, dst_device_blocks)) self._await_terminal( store_b.poll_remote_read_status, len(hashes), "Job B read_remote" ) - # The batch commits as a unit into HOST_AND_HBM: the bytes are in the - # caller's device blocks AND in the host landing blocks that were the - # staging hop. - after = store_b.lookup(hashes) - self.assertLen(after, 2) - for i, (_, blk) in enumerate(after): - expected_status = kv_cache_store.BlockStatus.HBM if use_slices else kv_cache_store.BlockStatus.HOST_AND_HBM - self.assertEqual(blk.status, expected_status) - self.assertEqual(blk.device_block_id, dst_device_blocks[i]) + # The bytes are in the caller's device blocks and nowhere else. The host + # blocks the transfer staged through went straight back to the pool, so + # there is no local entry and a later local lookup is still a miss. + self.assertEmpty(store_b.lookup(hashes)) # --- Byte-exact verification of device memory. ------------------------- actual_b = self._reread_device(tpu_cache_b) @@ -838,25 +810,26 @@ def _run_remote_read_to_hbm_test(self, enable_multi_numa: bool, use_slices: bool ) if use_slices: - store_b.release(hashes) return - # --- The host copy left behind by the staging hop must be usable. ------ - # load() the same hashes into the sentinel blocks; if the staging blocks - # did not really hold the data, this produces garbage. - self.assertTrue(store_b.load(hashes, sentinel_blocks)) - self._await_terminal(store_b.poll_load_status, len(hashes), "Job B load") - store_b.release(hashes) + # --- No host copy is left behind. -------------------------------------- + # The staging blocks were handed back, so the read is not a way to warm + # the local cache. A later local load() of the same hashes has nothing to + # read from and is refused; the caller that wants a host copy must save() + # what it pulled. This is the deliberate cost of read_remote leaving no + # local record. + self.assertFalse( + store_b.load(hashes, sentinel_blocks), + "read_remote must not leave a host copy behind", + ) - reloaded = self._reread_device(tpu_cache_b) - for src_blk, dst_blk in zip(src_device_blocks, sentinel_blocks): + # The sentinel blocks are therefore untouched, as they were before. + unchanged = self._reread_device(tpu_cache_b) + for blk in sentinel_blocks: np.testing.assert_array_equal( - reloaded[dst_blk], - host_data_a[src_blk], - err_msg=( - f"load() into device block {dst_blk} does not byte-match source" - f" block {src_blk}: the pull's host staging copy is not valid" - ), + unchanged[blk], + host_data_b[blk], + err_msg=f"sentinel device block {blk} must stay untouched", ) def _await_terminal(self, poll_fn, expected_done, what, timeout_s=120.0): @@ -1044,10 +1017,10 @@ def _run_remote_write_e2e_test( self.assertCountEqual(done, hashes) self.assertEmpty(failed) self.assertEmpty(existing) - store_a.release(hashes) # 3. Job B holds them locally, host-resident, as its own. lookup_b = store_b.lookup(hashes, enable_global=False) + store_b.release(hashes) self.assertLen(lookup_b, len(hashes)) for _, slice_b in lookup_b: self.assertEqual(slice_b.status, kv_cache_store.BlockStatus.HOST) @@ -1062,7 +1035,6 @@ def _run_remote_write_e2e_test( self.assertTrue(store_b.pin(hashes)) self.assertTrue(store_b.load(hashes, list(range(num_blocks)))) self._await_terminal(store_b.poll_load_status, len(hashes), "Job B load") - store_b.release(hashes) np.testing.assert_array_equal( self._reread_device(tpu_cache_b), @@ -1231,8 +1203,8 @@ def test_remote_read_e2e_source_missing_block_fails(self): ) time.sleep(1) - # Job B manually records a REMOTE block pointing at Job A for a hash Job A - # never saved, then tries to read it. + # Job B names a source coordinate on Job A for a hash Job A never saved, + # and tries to read it. ghost = [b"ghost_hash"] slices = [ kv_cache_store.RaidenBlockID( @@ -1242,9 +1214,8 @@ def test_remote_read_e2e_source_missing_block_fails(self): status=kv_cache_store.BlockStatus.REMOTE, ) ] - self.assertTrue(store_b.insert_and_lock(ghost, slices, on_host=True)) - self.assertTrue(store_b.read_remote(ghost)) + self.assertTrue(store_b.read_remote(ghost, slices, [0])) failed = False for _ in range(500): @@ -1343,22 +1314,16 @@ def test_remote_read_e2e_source_wrong_status_fails(self): ) time.sleep(1) - # Destination records a REMOTE reference to the source's HBM-only block. - self.assertTrue( - store_b.insert_and_lock( - hashes, - [ - kv_cache_store.RaidenBlockID( - rid_a, - host_block_id=0, - device_block_id=-1, - status=kv_cache_store.BlockStatus.REMOTE, - ) - ], - on_host=True, + # Destination names the source's HBM-only block as the read source. + slices = [ + kv_cache_store.RaidenBlockID( + rid_a, + host_block_id=0, + device_block_id=-1, + status=kv_cache_store.BlockStatus.REMOTE, ) - ) - self.assertTrue(store_b.read_remote(hashes)) + ] + self.assertTrue(store_b.read_remote(hashes, slices, [0])) failed = False for _ in range(500): @@ -1465,7 +1430,6 @@ def build_manager_after_a_delay(): break time.sleep(0.01) self.assertCountEqual(done, hashes) - store.release(hashes) del built diff --git a/tpu_sync/api/jax/kv_cache_store_recovery_e2e_test.py b/tpu_sync/api/jax/kv_cache_store_recovery_e2e_test.py index 2038d563..8e171f8f 100644 --- a/tpu_sync/api/jax/kv_cache_store_recovery_e2e_test.py +++ b/tpu_sync/api/jax/kv_cache_store_recovery_e2e_test.py @@ -153,14 +153,13 @@ def _phase_a(): ] inserted, evicted = store.insert(_HASHES, slices, on_host=False) assert inserted and not evicted - assert store.pin(_HASHES) store.save(_HASHES) _poll(store.poll_save_status, _NUM_BLOCKS, "save") - store.release(_HASHES) # The blocks are host-resident now; their bytes and the metadata table both # live in shared memory and must survive the crash below. lookup_res = store.lookup(_HASHES) + store.release(_HASHES) assert len(lookup_res) == _NUM_BLOCKS for i, (_, blk) in enumerate(lookup_res): assert blk.status == kv_cache_store.BlockStatus.HOST_AND_HBM, blk.status @@ -185,6 +184,7 @@ def _phase_b(expect_recovery: bool): del rid lookup_res = store.lookup(_HASHES) + store.release(_HASHES) if not expect_recovery: assert not lookup_res, f"expected a cold start, got hits: {lookup_res}" print(_PHASE_B_COLD_MARKER, flush=True) @@ -204,7 +204,6 @@ def _phase_b(expect_recovery: bool): assert store.pin(_HASHES) store.load(_HASHES, list(range(_NUM_BLOCKS))) _poll(store.poll_load_status, _NUM_BLOCKS, "load") - store.release(_HASHES) np.testing.assert_array_equal(np.asarray(tpu_cache), host_data) print(_PHASE_B_BYTES_MARKER, flush=True) diff --git a/tpu_sync/api/jax/kv_cache_store_test.py b/tpu_sync/api/jax/kv_cache_store_test.py index 2fa8edd9..412134eb 100644 --- a/tpu_sync/api/jax/kv_cache_store_test.py +++ b/tpu_sync/api/jax/kv_cache_store_test.py @@ -183,6 +183,7 @@ def test_basic_tests(self): # 2. Lookup with a partial miss at the end hashes_with_miss = [b"6001", b"6002", b"6003"] lookup_res = controller.lookup(hashes_with_miss) + controller.release(hashes_with_miss) self.assertLen(lookup_res, 2) self.assertEqual(lookup_res[0][0], b"6001") self.assertEqual(lookup_res[0][1].raiden_id.job_name, "inference_server") @@ -191,10 +192,12 @@ def test_basic_tests(self): # Lookup with an early miss hashes_early_miss = [b"6001", b"6003", b"6002"] lookup_res_early = controller.lookup(hashes_early_miss) + controller.release(hashes_early_miss) self.assertLen(lookup_res_early, 1) self.assertEqual(lookup_res_early[0][0], b"6001") # 3. Delete + controller.release(hashes) controller.delete(hashes, slices) self.assertTrue( controller.insert(hashes, slices, True)[0] @@ -230,8 +233,12 @@ def test_pin_and_release(self): slice_4 = [kv_cache_store.RaidenId("inference_server", "3", "kv_cache", 0)] controller.insert(hash_4, slice_4, True) - self.assertEmpty(controller.lookup([b"7001", b"7002"])) - self.assertLen(controller.lookup([b"7002"]), 1) + res_old = controller.lookup([b"7001", b"7002"]) + self.assertLen(res_old, 2) + controller.release([b"7001", b"7002"]) + res = controller.lookup([b"7002"]) + self.assertLen(res, 1) + controller.release([b"7002"]) def test_partial_pin_rollback(self): controller = kv_cache_store.KVCacheStore( @@ -261,8 +268,12 @@ def test_partial_pin_rollback(self): )[0] ) - self.assertEmpty(controller.lookup([b"8001", b"8002"])) - self.assertLen(controller.lookup([b"8004", b"8005"]), 2) + res_old = controller.lookup([b"8001", b"8002"]) + self.assertLen(res_old, 2) + controller.release([b"8001", b"8002"]) + res = controller.lookup([b"8004", b"8005"]) + self.assertLen(res, 2) + controller.release([b"8004", b"8005"]) def test_large_and_arbitrary_length_hashes(self): controller = kv_cache_store.KVCacheStore( @@ -281,6 +292,7 @@ def test_large_and_arbitrary_length_hashes(self): self.assertTrue(controller.insert(hashes, slices, True)[0]) lookup_res = controller.lookup(hashes) + controller.release(hashes) self.assertLen(lookup_res, 2) self.assertEqual(lookup_res[0][0], large_hash) self.assertEqual(lookup_res[1][0], long_hash) @@ -298,6 +310,7 @@ def test_global_lookup_case1_local_hit(self): self.assertTrue(controller.insert(hashes, slices, True)[0]) res = controller.lookup(hashes, enable_global=True) + controller.release(hashes) self.assertLen(res, 1) self.assertEqual(res[0][0], b"local_only") self.assertEqual(res[0][1].raiden_id.job_name, "local_job") @@ -321,6 +334,7 @@ def test_global_lookup_case2_and_3_mocked(self): mock_impl.lookup.return_value = [(b"shared_hash", local_id)] res = controller.lookup([b"shared_hash"], enable_global=True) + controller.release([b"shared_hash"]) self.assertLen(res, 1) self.assertEqual(res[0][0], b"shared_hash") self.assertEqual(res[0][1].raiden_id.job_name, "local_job") @@ -344,6 +358,7 @@ def test_global_lookup_case2_and_3_mocked(self): ] res = controller.lookup([b"global_1", b"global_2"], enable_global=True) + controller.release([b"global_1", b"global_2"]) self.assertLen(res, 2) self.assertEqual(res[0][0], b"global_1") self.assertEqual(res[0][1].raiden_id.job_name, "job1") @@ -386,6 +401,7 @@ def test_global_lookup_error_ignored(self): hashes = [b"9001"] # Should not fail, just return empty because the registry is now down. res = controller.lookup(hashes, enable_global=True) + controller.release(hashes) self.assertEmpty(res) def test_save_and_load_mocked(self): @@ -456,11 +472,15 @@ def test_insert_and_lock_release_and_delete(self): ] success = controller.insert_and_lock(remote_hashes, remote_slices, True) self.assertTrue(success) - self.assertEmpty(controller.lookup([b"local_1"])) + res_local = controller.lookup([b"local_1"]) + self.assertLen(res_local, 1) + controller.release([b"local_1"]) del_count = controller.release_and_delete(remote_hashes) self.assertEqual(del_count, 2) - self.assertLen(controller.lookup([b"local_1", b"local_2"]), 2) + res = controller.lookup([b"local_1", b"local_2"]) + self.assertLen(res, 2) + controller.release([b"local_1", b"local_2"]) def test_e2e_load(self): """Tests end-to-end load (H2D) on CPU.""" @@ -554,6 +574,7 @@ def test_e2e_load(self): # Verify LRU Status Upgrade in Store lookup_res1 = store.lookup([b"hash1"]) + store.release([b"hash1"]) self.assertLen(lookup_res1, 1) self.assertEqual( lookup_res1[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -562,6 +583,7 @@ def test_e2e_load(self): self.assertEqual(lookup_res1[0][1].device_block_id, 5) lookup_res2 = store.lookup([b"hash2"]) + store.release([b"hash2"]) self.assertLen(lookup_res2, 1) self.assertEqual( lookup_res2[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -680,6 +702,7 @@ def test_e2e_load_with_slices(self): (hashes[1], 4, 6), ): res = store.lookup([hash_val]) + store.release([hash_val]) self.assertLen(res, 1) self.assertEqual( res[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -687,7 +710,6 @@ def test_e2e_load_with_slices(self): self.assertEqual(res[0][1].host_block_id, expected_host) self.assertEqual(res[0][1].device_block_id, expected_device) - store.release(hashes) def test_e2e_save(self): """Tests end-to-end save (D2H) and load (H2D) back on TPU.""" @@ -751,7 +773,6 @@ def test_e2e_save(self): ) ] self.assertTrue(store.insert([b"hash1"], slices_1, False)[0]) - self.assertTrue(store.pin([b"hash1"])) slices_2 = [ kv_cache_store.RaidenBlockID( @@ -762,7 +783,6 @@ def test_e2e_save(self): ) ] self.assertTrue(store.insert([b"hash2"], slices_2, False)[0]) - self.assertTrue(store.pin([b"hash2"])) # 6. Trigger Save (D2H) self.assertTrue(store.save([b"hash1", b"hash2"])) @@ -781,6 +801,7 @@ def test_e2e_save(self): # Verify status in LRU is HOST_AND_HBM, and host_block_id is allocated lookup_res1 = store.lookup([b"hash1"]) + store.release([b"hash1"]) self.assertLen(lookup_res1, 1) self.assertEqual( lookup_res1[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -789,6 +810,7 @@ def test_e2e_save(self): self.assertGreaterEqual(host_block_id1, 0) lookup_res2 = store.lookup([b"hash2"]) + store.release([b"hash2"]) self.assertLen(lookup_res2, 1) self.assertEqual( lookup_res2[0][1].status, kv_cache_store.BlockStatus.HOST_AND_HBM @@ -809,6 +831,7 @@ def test_e2e_save(self): global_verified = False for _ in range(50): lookup_res = store2.lookup([b"hash1", b"hash2"], enable_global=True) + store2.release([b"hash1", b"hash2"]) if len(lookup_res) == 2: self.assertEqual(lookup_res[0][0], b"hash1") self.assertEqual( diff --git a/tpu_sync/api/torch/kv_cache_store.py b/tpu_sync/api/torch/kv_cache_store.py index 5c5938b1..facd49a8 100644 --- a/tpu_sync/api/torch/kv_cache_store.py +++ b/tpu_sync/api/torch/kv_cache_store.py @@ -233,9 +233,20 @@ def lookup( ) -> list[tuple[bytes, RaidenBlockID]]: """Checks the LRU directory for cached block hashes. + PINS every hash it returns, one pin each, so the answer cannot be evicted + between being given and being used. The operation the caller goes on to + perform consumes that pin: load() drops it on a successful local load, and + save() drops it on success. A caller that does neither -- one that only + wanted to know what is resident -- must release() what it was given. + + Locally cached hits are pinned. Hashes resolved only through the global + registry are not: they name a block on another node, and there is nothing + here to hold. + Args: block_hashes: Incoming block hashes to check. - enable_global: Whether to fallback to global registry on miss. + enable_global: Whether to fallback to global registry on miss. Defaults + to False; the fallback is a blocking RPC. Returns: A list of tuples containing the block hash and the matching @@ -349,35 +360,49 @@ def release(self, block_hashes: list[bytes]) -> None: def read_remote( self, block_hashes: list[bytes], - device_block_ids: list[int] | None = None, + slices: list[RaidenBlockID], + device_block_ids: list[int], ) -> bool: - """Launches an async receiver-initiated read of REMOTE blocks from peers. + """Reads REMOTE blocks from their owning peers straight into local HBM. Returns as soon as the reads are issued; poll with poll_remote_read_status(). + This store's cache is neither consulted nor modified. The hashes need not + be present locally and need not be pinned; nothing is inserted on success + and nothing is left behind on failure. The bytes land ONLY in the given + device blocks -- no local host copy is kept, so a later local load() of the + same hash is still a miss. + + Compare with load(): both bring a peer's block into local HBM. Use load() + when the hash may be resident locally and you want the store to decide; use + read_remote() when you already hold the source coordinates and want no + local record of the transfer. + Requires a global registry: it is what maps the owning peer to the controller address this store acquires its read lease from. A store built without a global_registry_address fails every read. Args: - block_hashes: Block hashes to read. Each must already be pinned, and must - stay pinned until poll_remote_read_status() reports it terminal. - Releasing early makes the entry deletable mid-read, in which case the - read is discarded and the WHOLE batch is reported failed. - device_block_ids: Optional. Omit (or pass None) to read into host DRAM, - leaving the entries HOST. Pass one device block id per hash to read - straight into HBM, leaving the entries HOST_AND_HBM (the host landing - blocks act as the staging hop, so a later load() can reuse them). - On FAILURE the contents of these device blocks are UNDEFINED: they are - written before the source's verdict is known. Treat them as scratch - until the read reports success -- nothing in the cache points at them - unless it does. + block_hashes: Block hashes to read. + slices: One REMOTE RaidenBlockID per hash, naming where to read from. + Only two fields are used -- raiden_id (the owning peer) and + host_block_id (the block on that peer) -- so a lookup() answer can be + passed straight through. + device_block_ids: One local device block id per hash. Mandatory. + On FAILURE their contents are UNDEFINED: they are written before the + source's verdict is known. Treat them as scratch until the read + reports success. Returns: True if successfully launched. """ - return self._impl.read_remote(block_hashes, device_block_ids or []) + raw_slices = [] + for s in slices: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + raw_slices.append(s._impl) # pylint: disable=protected-access + return self._impl.read_remote(block_hashes, raw_slices, device_block_ids) def poll_remote_read_status( self, @@ -413,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_e2e_test.py b/tpu_sync/api/torch/kv_cache_store_e2e_test.py index 1c5ccdd3..b955b71d 100644 --- a/tpu_sync/api/torch/kv_cache_store_e2e_test.py +++ b/tpu_sync/api/torch/kv_cache_store_e2e_test.py @@ -256,7 +256,6 @@ def get_slice_e2e(x): time.sleep(0.01) # Release them so we can test pinning before load - store.release(hashes) # Verify status in store is updated to HOST_AND_HBM lookup_res = store.lookup(hashes) @@ -298,7 +297,6 @@ def get_slice_e2e(x): time.sleep(0.01) # Release at the very end - store.release(hashes) try: torch.tpu.synchronize() @@ -443,7 +441,6 @@ def _run_remote_read_e2e_test( if not done: time.sleep(0.01) - store_a.release(hashes) # 5. Job B calls Lookup (enable_global=True) time.sleep(0.5) @@ -472,12 +469,11 @@ def _run_remote_read_e2e_test( lookup_res_b[1][1].host_block_id, lookup_res_a[1][1].host_block_id ) - # 6. Job B controller calls insert_and_lock for the remote slices + # 6. Job B reads straight from Job A into its own device blocks. The + # source coordinates come from the lookup answer, so nothing needs to be + # inserted into Job B's cache first. slices_b = [lookup_res_b[0][1], lookup_res_b[1][1]] - self.assertTrue(store_b.insert_and_lock(hashes, slices_b, on_host=True)) - - # 7. Job B calls ReadRemote - self.assertTrue(store_b.read_remote(hashes)) + self.assertTrue(store_b.read_remote(hashes, slices_b, [0, 1])) if not expect_read_success: failed = False @@ -505,33 +501,13 @@ def _run_remote_read_e2e_test( if not done: time.sleep(0.01) - # 8. Verify Job B's LRU block status becomes HOST - lookup_res_b_after = store_b.lookup(hashes) - self.assertLen(lookup_res_b_after, 2) - self.assertEqual( - lookup_res_b_after[0][1].status, kv_cache_store.BlockStatus.HOST - ) - self.assertEqual( - lookup_res_b_after[1][1].status, kv_cache_store.BlockStatus.HOST - ) - - # 9. Job B controller calls Load to transfer data to TPU blocks [0, 1] - self.assertTrue(store_b.load(hashes, [0, 1])) - - # Wait for Load completion - done = False - while not done: - load_done, load_failed, _ = store_b.poll_load_status() - if load_failed: - raise RuntimeError(f"Job B Load failed: {load_failed}") - if len(load_done) == 2: - done = True - if not done: - time.sleep(0.01) + # 8. The read is already in HBM -- there is no second Load step, and no + # local record of it either. Job B's cache is still a miss for these + # hashes: the bytes live only in the device blocks it named. + self.assertEmpty(store_b.lookup(hashes)) - store_b.release(hashes) - # 10. Verify byte-exact match on Job B TPU device + # 9. Verify byte-exact match on Job B TPU device try: torch.tpu.synchronize() except (AttributeError, RuntimeError): @@ -677,7 +653,6 @@ def _run_remote_write_e2e_test( self.assertCountEqual(done, hashes) self.assertEmpty(failed) self.assertEmpty(existing) - store_a.release(hashes) # 3. Job B holds them locally, host-resident, as its own. lookup_b = store_b.lookup(hashes, enable_global=False) @@ -700,7 +675,6 @@ def _run_remote_write_e2e_test( if time.time() > deadline: raise RuntimeError("Job B load did not complete in time") time.sleep(0.01) - store_b.release(hashes) try: torch.tpu.synchronize() @@ -819,8 +793,7 @@ def test_remote_read_e2e_source_missing_block_fails(self): status=kv_cache_store.BlockStatus.REMOTE, ) ] - self.assertTrue(store_b.insert_and_lock(ghost, slices, on_host=True)) - self.assertTrue(store_b.read_remote(ghost)) + self.assertTrue(store_b.read_remote(ghost, slices, [0])) failed = False for _ in range(500): @@ -933,8 +906,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self): status=kv_cache_store.BlockStatus.REMOTE, ), ] - self.assertTrue(store_b.insert_and_lock(hashes, slices_b, on_host=True)) - self.assertTrue(store_b.read_remote(hashes)) + self.assertTrue(store_b.read_remote(hashes, slices_b, [0, 1])) failed = False for _ in range(500): @@ -1022,7 +994,6 @@ def build_manager_after_a_delay(): break time.sleep(0.01) self.assertCountEqual(done, hashes) - store.release(hashes) del built diff --git a/tpu_sync/api/torch/kv_cache_store_mpmd_e2e_test.py b/tpu_sync/api/torch/kv_cache_store_mpmd_e2e_test.py index 61ec6e80..6daf06f5 100644 --- a/tpu_sync/api/torch/kv_cache_store_mpmd_e2e_test.py +++ b/tpu_sync/api/torch/kv_cache_store_mpmd_e2e_test.py @@ -264,7 +264,6 @@ def _worker_save_load_main(argv): done = True if not done: time.sleep(0.01) - store.release(hashes) if rank == 0: print("=== [Rank 0] Loading checkpoint from Host DRAM into TPU HBM blocks [2, 3] (store.load) ===") @@ -297,7 +296,6 @@ def _worker_save_load_main(argv): done = True if not done: time.sleep(0.01) - store.release(hashes) dist.barrier() try: @@ -452,7 +450,6 @@ def _worker_read_remote_main(argv): done = True if not done: time.sleep(0.01) - store_a.release(hashes) # Wait for global registry propagation deadline = time.time() + 10.0 @@ -480,15 +477,14 @@ def _worker_read_remote_main(argv): if not done: time.sleep(0.01) else: - assert store_b.insert_and_lock( - hashes, slices_b, on_host=True - ), "insert_and_lock failed on store_b" - + # The read goes straight into TPU blocks [0, 1]; the slices from the + # lookup are the source coordinate, so there is nothing to insert and + # no second Load step afterwards. print("=== [Rank 0] Launching ReadRemote from Job A to Job B ===") assert store_b.read_remote( - hashes + hashes, slices_b, [0, 1] ), "read_remote launch failed on store_b" - + done = False while not done: read_done, read_failed, _ = store_b.poll_remote_read_status() @@ -498,20 +494,12 @@ def _worker_read_remote_main(argv): done = True if not done: time.sleep(0.01) - - # Now load blocks [0, 1] from Job B's host pool into TPU HBM - assert store_b.load(hashes, [0, 1]), "load failed on store_b" - done = False - while not done: - load_done, load_failed, _ = store_b.poll_load_status() - if load_failed: - raise RuntimeError(f"Job B Load failed: {load_failed}") - if len(load_done) == 2: - done = True - if not done: - time.sleep(0.01) - store_b.release_and_delete(hashes) + # The read left no local entry, so there is nothing to release here. + assert not store_b.lookup(hashes), "read_remote must record nothing" + + if use_slices: + store_b.release_and_delete(hashes) dist.barrier() try: @@ -693,7 +681,6 @@ def _worker_write_remote_main(argv): break time.sleep(0.01) assert done, "Job A WriteRemote timed out" - store_a.release(hashes) # Destination holds blocks locally on host DRAM lookup_res_b = store_b.lookup(hashes, enable_global=False) diff --git a/tpu_sync/api/torch/kv_cache_store_recovery_e2e_test.py b/tpu_sync/api/torch/kv_cache_store_recovery_e2e_test.py index 42784ae9..7510deab 100644 --- a/tpu_sync/api/torch/kv_cache_store_recovery_e2e_test.py +++ b/tpu_sync/api/torch/kv_cache_store_recovery_e2e_test.py @@ -147,14 +147,13 @@ def _phase_a(): ] inserted, evicted = store.insert(_HASHES, slices, on_host=False) assert inserted and not evicted - assert store.pin(_HASHES) store.save(_HASHES) _poll(store.poll_save_status, _NUM_BLOCKS, "save") - store.release(_HASHES) # The blocks are host-resident now; their bytes and the metadata table both # live in shared memory and must survive the crash below. lookup_res = store.lookup(_HASHES) + store.release(_HASHES) assert len(lookup_res) == _NUM_BLOCKS for i, (_, blk) in enumerate(lookup_res): assert blk.status == kv_cache_store.BlockStatus.HOST_AND_HBM, blk.status @@ -181,6 +180,7 @@ def _phase_b(expect_recovery: bool): del rid lookup_res = store.lookup(_HASHES) + store.release(_HASHES) if not expect_recovery: assert not lookup_res, f"expected a cold start, got hits: {lookup_res}" print(_PHASE_B_COLD_MARKER, flush=True) @@ -200,7 +200,6 @@ def _phase_b(expect_recovery: bool): assert store.pin(_HASHES) store.load(_HASHES, list(range(_NUM_BLOCKS))) _poll(store.poll_load_status, _NUM_BLOCKS, "load") - store.release(_HASHES) np.testing.assert_array_equal(tpu_cache.cpu().numpy(), host_data) print(_PHASE_B_BYTES_MARKER, flush=True) diff --git a/tpu_sync/api/torch/kv_cache_store_test.py b/tpu_sync/api/torch/kv_cache_store_test.py index 99b28da5..79416b67 100644 --- a/tpu_sync/api/torch/kv_cache_store_test.py +++ b/tpu_sync/api/torch/kv_cache_store_test.py @@ -124,6 +124,7 @@ def test_basic_tests(self): # 2. Lookup with a partial miss at the end hashes_with_miss = [b"6001", b"6002", b"6003"] lookup_res = controller.lookup(hashes_with_miss) + controller.release(hashes_with_miss) self.assertLen(lookup_res, 2) self.assertEqual(lookup_res[0][0], b"6001") self.assertEqual(lookup_res[0][1].raiden_id.job_name, "inference_server") @@ -132,10 +133,12 @@ def test_basic_tests(self): # Lookup with an early miss hashes_early_miss = [b"6001", b"6003", b"6002"] lookup_res_early = controller.lookup(hashes_early_miss) + controller.release(hashes_early_miss) self.assertLen(lookup_res_early, 1) self.assertEqual(lookup_res_early[0][0], b"6001") # 3. Delete + controller.release(hashes) controller.delete(hashes, slices) self.assertTrue( controller.insert(hashes, slices, True)[0] @@ -171,8 +174,12 @@ def test_pin_and_release(self): slice_4 = [kv_cache_store.RaidenId("inference_server", "3", "kv_cache", 0)] controller.insert(hash_4, slice_4, True) - self.assertEmpty(controller.lookup([b"7001", b"7002"])) - self.assertLen(controller.lookup([b"7002"]), 1) + res_old = controller.lookup([b"7001", b"7002"]) + self.assertLen(res_old, 2) + controller.release([b"7001", b"7002"]) + res = controller.lookup([b"7002"]) + self.assertLen(res, 1) + controller.release([b"7002"]) def test_partial_pin_rollback(self): controller = kv_cache_store.KVCacheStore( @@ -202,8 +209,12 @@ def test_partial_pin_rollback(self): )[0] ) - self.assertEmpty(controller.lookup([b"8001", b"8002"])) - self.assertLen(controller.lookup([b"8004", b"8005"]), 2) + res_old = controller.lookup([b"8001", b"8002"]) + self.assertLen(res_old, 2) + controller.release([b"8001", b"8002"]) + res = controller.lookup([b"8004", b"8005"]) + self.assertLen(res, 2) + controller.release([b"8004", b"8005"]) def test_large_and_arbitrary_length_hashes(self): controller = kv_cache_store.KVCacheStore( @@ -222,6 +233,7 @@ def test_large_and_arbitrary_length_hashes(self): self.assertTrue(controller.insert(hashes, slices, True)[0]) lookup_res = controller.lookup(hashes) + controller.release(hashes) self.assertLen(lookup_res, 2) self.assertEqual(lookup_res[0][0], large_hash) self.assertEqual(lookup_res[1][0], long_hash) @@ -239,6 +251,7 @@ def test_global_lookup_case1_local_hit(self): self.assertTrue(controller.insert(hashes, slices, True)[0]) res = controller.lookup(hashes, enable_global=True) + controller.release(hashes) self.assertLen(res, 1) self.assertEqual(res[0][0], b"local_only") self.assertEqual(res[0][1].raiden_id.job_name, "local_job") @@ -262,6 +275,7 @@ def test_global_lookup_case2_and_3_mocked(self): mock_impl.lookup.return_value = [(b"shared_hash", local_id)] res = controller.lookup([b"shared_hash"], enable_global=True) + controller.release([b"shared_hash"]) self.assertLen(res, 1) self.assertEqual(res[0][0], b"shared_hash") self.assertEqual(res[0][1].raiden_id.job_name, "local_job") @@ -281,6 +295,7 @@ def test_global_lookup_case2_and_3_mocked(self): ] res = controller.lookup([b"global_1", b"global_2"], enable_global=True) + controller.release([b"global_1", b"global_2"]) self.assertLen(res, 2) self.assertEqual(res[0][0], b"global_1") self.assertEqual(res[0][1].raiden_id.job_name, "10.0.0.1:1234") @@ -321,6 +336,7 @@ def test_global_lookup_error_ignored(self): hashes = [b"9001"] # Should not fail, just return empty because the registry is now down. res = controller.lookup(hashes, enable_global=True) + controller.release(hashes) self.assertEmpty(res) def test_insert_and_lock_release_and_delete(self): @@ -358,11 +374,15 @@ def test_insert_and_lock_release_and_delete(self): ] success = controller.insert_and_lock(remote_hashes, remote_slices, True) self.assertTrue(success) - self.assertEmpty(controller.lookup([b"local_1"])) + res_local = controller.lookup([b"local_1"]) + self.assertLen(res_local, 1) + controller.release([b"local_1"]) del_count = controller.release_and_delete(remote_hashes) self.assertEqual(del_count, 2) - self.assertLen(controller.lookup([b"local_1", b"local_2"]), 2) + res = controller.lookup([b"local_1", b"local_2"]) + self.assertLen(res, 2) + controller.release([b"local_1", b"local_2"]) def test_save_and_load_mocked(self): controller = kv_cache_store.KVCacheStore( diff --git a/tpu_sync/frameworks/jax/kv_cache_store.pyi b/tpu_sync/frameworks/jax/kv_cache_store.pyi index 92130b2a..dbf95568 100644 --- a/tpu_sync/frameworks/jax/kv_cache_store.pyi +++ b/tpu_sync/frameworks/jax/kv_cache_store.pyi @@ -188,9 +188,10 @@ class KVCacheStore: def read_remote( self, block_hashes: list[bytes], - device_block_ids: list[int] = ..., + slices: list[RaidenBlockID], + device_block_ids: list[int], ) -> bool: - """Launches async H2H read from remote worker.""" + """Reads REMOTE blocks from their owning peers into local HBM.""" ... def poll_remote_read_status(self) -> tuple[list[bytes], list[bytes], list[bytes]]: """Polls status of active remote reads.""" diff --git a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc index e07408a0..7a251a00 100644 --- a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc @@ -701,12 +701,13 @@ NB_MODULE(_tpu_raiden_jax, m) { "read_remote", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, + const std::vector& slices, const std::vector& device_block_ids) -> bool { auto hashes = ToStdStringVector(block_hashes); - return self->ReadRemote(hashes, device_block_ids).ok(); + return self->ReadRemote(hashes, slices, device_block_ids).ok(); }, - nb::arg("block_hashes"), - nb::arg("device_block_ids") = std::vector()) + nb::arg("block_hashes"), nb::arg("slices"), + nb::arg("device_block_ids")) .def("poll_remote_read_status", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self) { // Released around the C++ call ONLY. The subsequent nb::bytes diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index c4e36b52..822b4ac8 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc @@ -746,12 +746,13 @@ NB_MODULE(_tpu_raiden_torch, m) { "read_remote", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, + const std::vector& slices, const std::vector& device_block_ids) -> bool { auto hashes = ToStdStringVector(block_hashes); - return self->ReadRemote(hashes, device_block_ids).ok(); + return self->ReadRemote(hashes, slices, device_block_ids).ok(); }, - nb::arg("block_hashes"), - nb::arg("device_block_ids") = std::vector()) + nb::arg("block_hashes"), nb::arg("slices"), + nb::arg("device_block_ids")) .def("poll_remote_read_status", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self) { // Poll drains whatever futures completed and can make blocking diff --git a/tpu_sync/kv_cache/host_offload_backend.cc b/tpu_sync/kv_cache/host_offload_backend.cc index 73943d0d..c22f8bce 100644 --- a/tpu_sync/kv_cache/host_offload_backend.cc +++ b/tpu_sync/kv_cache/host_offload_backend.cc @@ -183,10 +183,19 @@ absl::StatusOr HostOffloadBackend::Lookup( local_id = raiden_id_; for (size_t i = 0; i < block_hashes.size(); ++i) { const auto& hash = block_hashes[i]; - const RaidenBlockID* existing = options.pin_found - ? lru_cache_.GetAndPin(hash) - : lru_cache_.Peek(hash); + // Peek in BOTH cases, so pin_found decides only whether a hit is pinned, + // never whether it is a hit. GetAndPin would also resurrect an eviction + // candidate -- it splices the node out of evict_candidate_list_ -- which + // would make a lookup silently un-queue a block the store had already + // decided to reclaim, and make the space accounting behind it wrong. + // Candidates are invisible to lookup; pinning must not change that. + const RaidenBlockID* existing = lru_cache_.Peek(hash); if (existing != nullptr) { + if (options.pin_found) { + // Not a candidate (Peek just proved it), so this only moves the node + // from the active LRU list to the pinned list. + lru_cache_.Pin(hash); + } local_hit[i] = true; continue; } diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index 7af13013..6b247afd 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -755,6 +755,30 @@ KVCacheStore::~KVCacheStore() { poller_thread_->join(); } } + + // Abandon any remote write still outstanding, now that no poller can race + // us. This does NOT wait for them: a remote write goes terminal only when + // the destination answers or the HOLD (~30s) expires, so waiting would make + // destroying a store block for half a minute behind a slow or dead peer. + // + // Releasing the internal pin is the part that has to happen. `backends_` + // holds shared_ptrs, so a backend can outlive the store that pinned into it; + // a pin left behind there is a host block nothing can ever reclaim. + { + std::vector abandoned; + { + absl::MutexLock lock(mutex_); + abandoned.swap(active_remote_writes_); + polling_remote_writes_.clear(); + } + for (const auto& state : abandoned) { + LOG(WARNING) << "Store destroyed with remote write " << state.operation_id + << " still outstanding; releasing its source pin without " + "waiting for the destination's verdict."; + FinishRemoteWrite(state, /*succeeded=*/false, {}); + } + } + std::vector> futures_to_await; { absl::MutexLock lock(mutex_); @@ -772,7 +796,15 @@ KVCacheStore::~KVCacheStore() { absl::StatusOr KVCacheStore::Lookup( const std::vector& block_hashes, bool enable_global) { - return Lookup(block_hashes, LookupOptions{.enable_global = enable_global}); + // The application-facing overload PINS what it finds. A caller asks what is + // resident in order to use it, and between the answer and the use the entry + // would otherwise be evictable -- so the pin comes with the answer, and the + // operation the caller goes on to perform (load, save) consumes it. + // + // Only this overload. The LookupOptions overload leaves pin_found at its + // struct default of false, which is what every internal caller goes through. + return Lookup(block_hashes, + LookupOptions{.enable_global = enable_global, .pin_found = true}); } absl::StatusOr KVCacheStore::Lookup( @@ -1048,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); @@ -1061,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) { @@ -1115,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, }); } @@ -1137,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) { @@ -1163,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( @@ -1194,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, }); } @@ -1250,11 +1295,26 @@ void KVCacheStore::UnpinHostBlocks(absl::Span block_hashes) { absl::Status KVCacheStore::ReadRemote( const std::vector& block_hashes, + const std::vector& slices, const std::vector& device_block_ids) { if (block_hashes.empty()) { return absl::OkStatus(); } + // Validate before allocating anything: an early return past the allocation + // owes the cleanup below, and there is nothing to clean up yet here. + if (slices.size() != block_hashes.size()) { + return absl::InvalidArgumentError( + absl::StrCat("slices size ", slices.size(), + " must match block_hashes size ", block_hashes.size())); + } + if (device_block_ids.size() != block_hashes.size()) { + return absl::InvalidArgumentError(absl::StrCat( + "device_block_ids size ", device_block_ids.size(), + " must match block_hashes size ", block_hashes.size(), + ": read_remote always reads into local HBM")); + } + auto host_blocks_or = AllocateBlockIds(block_hashes.size()); if (!host_blocks_or.ok()) { return host_blocks_or.status(); @@ -1275,13 +1335,6 @@ absl::Status KVCacheStore::ReadRemote( } }); - const bool to_hbm = !device_block_ids.empty(); - if (to_hbm && device_block_ids.size() != block_hashes.size()) { - return absl::InvalidArgumentError(absl::StrCat( - "device_block_ids size ", device_block_ids.size(), - " must be empty (read to host) or match block_hashes size ", - block_hashes.size())); - } std::vector dst_host_block_ids = host_blocks_or.value(); struct RemoteReadGroup { @@ -1297,15 +1350,12 @@ absl::Status KVCacheStore::ReadRemote( std::vector groups; { + // The source coordinates come from the caller, not from this store's + // index: `slices[i].raiden_id` names the owning peer and + // `slices[i].host_block_id` the block on it. The lock still guards + // reading_hashes_, which is this store's own in-flight marker. absl::MutexLock lock(mutex_); - auto lookup_or = backend()->Lookup(block_hashes); - if (!lookup_or.ok()) return lookup_or.status(); - const auto& slices = lookup_or.value(); - if (slices.size() < block_hashes.size()) { - return absl::NotFoundError( - absl::StrCat("Block hash not found: ", block_hashes[slices.size()])); - } - for (size_t i = 0; i < slices.size(); ++i) { + for (size_t i = 0; i < block_hashes.size(); ++i) { const auto& hash = block_hashes[i]; if (!reading_hashes_.insert(hash).second) { return absl::FailedPreconditionError( @@ -1313,7 +1363,7 @@ absl::Status KVCacheStore::ReadRemote( } successfully_marked_as_reading.push_back(hash); - const auto& src_id = slices[i].second.raiden_id; + const auto& src_id = slices[i].raiden_id; auto it = std::find_if(groups.begin(), groups.end(), [&src_id](const RemoteReadGroup& g) { return g.src_raiden_id == src_id; @@ -1322,12 +1372,10 @@ absl::Status KVCacheStore::ReadRemote( groups.push_back(RemoteReadGroup{.src_raiden_id = src_id}); it = groups.end() - 1; } - it->src_host_block_ids.push_back(slices[i].second.host_block_id); + it->src_host_block_ids.push_back(slices[i].host_block_id); it->dst_host_block_ids.push_back(dst_host_block_ids[i]); it->block_hashes.push_back(hash); - if (to_hbm) { - it->device_block_ids.push_back(device_block_ids[i]); - } + it->device_block_ids.push_back(device_block_ids[i]); } } @@ -1381,10 +1429,6 @@ absl::Status KVCacheStore::ReadRemote( } } - // NOTE: the landing block ids are deliberately NOT stamped into the LRU - // entries here. The entry's host_block_id is the PEER's coordinate until the - // read commits; overwriting it up front (as this code used to) corrupts the - // entry on every failure path, because nothing restores it. // One lease per owning peer. The per-group futures are joined, so if ANY // group fails -- transfer error or a verdict other than HELD -- the whole // batch discards, including groups whose bytes landed perfectly. That is @@ -1417,11 +1461,10 @@ absl::Status KVCacheStore::ReadRemote( .block_hashes = block_hashes, .src_raiden_ids = std::move(peers), .host_block_ids = dst_host_block_ids, - .device_block_ids = device_block_ids, }); } - // Issued: the landing blocks now belong to the read, and the reading marks + // Issued: the staging blocks now belong to the read, and the reading marks // are cleared by the poller when it goes terminal. std::move(cleanup).Cancel(); return absl::OkStatus(); @@ -1753,15 +1796,34 @@ void KVCacheStore::FinishRemoteWrite(const RemoteWriteState& state, } void KVCacheStore::PollRemoteWritesInternal() { + // CLAIM the operations rather than copying the list. The poll below makes an + // RPC per operation and cannot hold mutex_ across it, so a plain copy lets + // two concurrent pollers observe the same operation as committed and both + // finish it -- releasing the internal pin twice (freeing blocks a later + // operation now owns) and reporting every hash twice. std::vector to_poll; { absl::MutexLock lock(mutex_); - to_poll = active_remote_writes_; + for (const auto& state : active_remote_writes_) { + if (polling_remote_writes_.insert(state.operation_id).second) { + to_poll.push_back(state); + } + } } if (to_poll.empty()) { return; } + // Every claim must come back, including on the paths that `continue` past a + // poll failure. A claim left behind is an operation no poller ever looks at + // again: it stays pending forever and its pin is never released. + auto release_claims = absl::MakeCleanup([this, &to_poll]() { + absl::MutexLock lock(mutex_); + for (const auto& state : to_poll) { + polling_remote_writes_.erase(state.operation_id); + } + }); + auto* backend = dynamic_cast(this->backend().get()); const absl::Time now = absl::Now(); for (auto& state : to_poll) { @@ -1918,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; @@ -1929,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); } @@ -1943,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 { @@ -1963,94 +2044,17 @@ void KVCacheStore::PollRemoteReadsInternal( absl::Status status = future.Await(); absl::MutexLock lock(mutex_); - // The batch commits as a UNIT: all hashes promoted, or none. Verify every - // entry is still present and still pinned BEFORE promoting anything. - // - // With pinned entries protected from erase, an entry can only vanish - // mid-read if the caller broke the contract by releasing its pin early -- - // so this is a bug detector, not a race handler. It replaces a loop that - // pushed a vanished hash onto done_remote_reads_ anyway, which told the - // caller "resident in HOST" about a landing block that had just been - // deallocated and reused. - std::vector contract_violations; if (status.ok()) { + // Nothing to record. The bytes are in the caller's device blocks and + // this store keeps no account of them: no LRU entry, no registry + // advertisement. Reporting the hashes done is the whole commit. for (const auto& hash : state.block_hashes) { - if (backend()->GetPinCount(hash) <= 0) { - contract_violations.push_back(hash); - } - } - if (!contract_violations.empty()) { - status = absl::FailedPreconditionError(absl::StrCat( - "read_remote caller released or deleted ", contract_violations.size(), - " of ", state.block_hashes.size(), - " entries before poll_remote_read_status reported them terminal; " - "the whole batch is discarded")); - LOG(ERROR) << status.message() << " First offending hash: " - << absl::BytesToHexString(contract_violations.front()); - } - } - - if (status.ok()) { - std::vector write_through_regs; - write_through_regs.reserve(state.block_hashes.size()); - auto lookup_or = backend()->Lookup(state.block_hashes); - if (lookup_or.ok()) { - const auto& slices = lookup_or.value(); - std::vector update_hashes; - std::vector update_slices; - for (size_t i = 0; i < state.block_hashes.size(); ++i) { - const auto& hash = state.block_hashes[i]; - if (i < slices.size()) { - RaidenBlockID block = slices[i].second; - block.host_block_id = state.host_block_ids[i]; - if (i < state.device_block_ids.size()) { - // Read-to-HBM: the bytes are in the caller's device blocks AND in - // the host landing blocks (which were the staging hop), so a - // later local load() can still reuse the host copy. - block.device_block_id = state.device_block_ids[i]; - block.status = BlockStatus::HOST_AND_HBM; - } else { - block.status = BlockStatus::HOST; - } - update_hashes.push_back(hash); - update_slices.push_back(block); - if (registry_client_) { - write_through_regs.push_back({ - .prefix_hash = hash, - .raiden_id = raiden_id_, - .block_id = state.host_block_ids[i], - }); - } - } else { - DeallocateBlockIds({state.host_block_ids[i]}); - } - done_remote_reads_.push_back(hash); - } - if (!update_hashes.empty()) { - backend()->Insert(update_hashes, update_slices, /*on_host=*/true); - } - } else { - DeallocateBlockIds(state.host_block_ids); - } - if (!write_through_regs.empty() && registry_client_ && - write_through_pool_) { - write_through_pool_->Schedule([client = registry_client_, - regs = std::move(write_through_regs)]() { - auto status = client->Register(regs); - if (!status.ok()) { - LOG(WARNING) << "Async write-through failed after ReadRemote: " - << status.message(); - } else { - LOG(INFO) << "Async write-through succeeded after ReadRemote for " - << regs.size() << " blocks"; - } - }); + done_remote_reads_.push_back(hash); } } else { - // Discard path. The entry keeps its ORIGINAL peer host_block_id (never - // stamped at issue time), so a retry is clean. In read-to-HBM mode the - // caller's device blocks may hold garbage -- by design: nothing in the - // LRU points at them, and the caller overwrites device blocks on reuse. + // The caller's device blocks may hold garbage -- by design: nothing + // points at them, and the caller treats them as scratch until this + // reports success. LOG(WARNING) << "Async ReadRemote failed: " << status.ToString(); // Drop these peers' cached controller addresses. Most failures are not // address failures, and dropping anyway is the point: re-resolving costs @@ -2059,11 +2063,15 @@ void KVCacheStore::PollRemoteReadsInternal( for (const auto& peer : state.src_raiden_ids) { resolved_peer_controllers_.erase(peer); } - DeallocateBlockIds(state.host_block_ids); for (const auto& hash : state.block_hashes) { failed_remote_reads_.push_back(hash); } } + // The staging blocks were a hop, not a destination, so they go back to the + // pool whichever way the read went. Success is not an exception: no LRU + // entry points at them, so leaking them here would burn a host block per + // read with nothing able to reclaim it. + DeallocateBlockIds(state.host_block_ids); for (const auto& hash : state.block_hashes) { reading_hashes_.erase(hash); } diff --git a/tpu_sync/kv_cache/kv_cache_store.h b/tpu_sync/kv_cache/kv_cache_store.h index 3f7e0c27..2ae1bbe4 100644 --- a/tpu_sync/kv_cache/kv_cache_store.h +++ b/tpu_sync/kv_cache/kv_cache_store.h @@ -189,11 +189,25 @@ class KVCacheStore { // matched replica pairs (block hash and RaidenBlockID) encountered // in sequence prior to the first miss. // If enable_global is true, it will query the global registry for any - // misses after the local lookup. + // misses after the local lookup. Defaults to false: the registry query is a + // blocking RPC, so a caller that only wants to know what is resident locally + // should not pay for one by omission. + // + // PINS every hash it returns, one pin each, so the answer cannot be evicted + // between being given and being used. The operation the caller goes on to + // perform consumes that pin -- Load() drops it on a successful local load, + // Save() on success. A caller that only wanted to observe residency, and + // performs neither, must Release() what it was given; the LookupOptions + // overload with pin_found = false is the way to observe without acquiring. + // + // Registry-only hits are not pinned: they name a block on another node, and + // there is nothing here to hold. absl::StatusOr Lookup( - const std::vector& block_hashes, bool enable_global = true); + const std::vector& block_hashes, bool enable_global = false); - // Overload accepting LookupOptions for granular control (e.g. pin_found). + // Overload accepting LookupOptions for granular control. Unlike the overload + // above, pin_found defaults to false here, so this is what internal callers + // and pure observers use. absl::StatusOr Lookup( const std::vector& block_hashes, const LookupOptions& options); @@ -260,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); @@ -276,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); @@ -339,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} @@ -352,34 +381,41 @@ class KVCacheStore { PollLoadStatus(); // Launches an async receiver-initiated read of REMOTE blocks from their - // owning peers. Returns as soon as the reads are issued; poll with - // PollRemoteReadStatus(). - // - // device_block_ids selects the destination: - // empty -> read to host. On success the entries become - // HOST. - // size == block_hashes -> read to HBM. The bytes land in the caller's - // device blocks, with the host landing blocks - // as the staging hop, so the entries become - // HOST_AND_HBM and a later load() can reuse - // the host copy. - // any other size -> InvalidArgument. - // - // CALLER CONTRACT: every requested hash must already be pinned, and must - // stay pinned until PollRemoteReadStatus() reports it terminal (done or - // failed). Releasing early makes the entry eligible for deletion mid-read; - // the read is then discarded and the WHOLE batch reported failed. - // - // In read-to-HBM mode the device blocks are written before the source's - // verdict is known, so on failure their contents are UNDEFINED -- treat - // supplied device blocks as scratch until the read reports success. Nothing - // in the cache ever points at them unless the read commits. + // owning peers straight into local HBM. Returns as soon as the reads are + // issued; poll with PollRemoteReadStatus(). + // + // The caller supplies the source coordinates directly: `slices[i]` is the + // REMOTE RaidenBlockID for `block_hashes[i]`, and only two of its fields are + // read -- `raiden_id` (which peer owns the block) and `host_block_id` (which + // block on that peer). A lookup() answer can be passed straight through. + // + // This store's LRU is not consulted and not modified. The hashes need not be + // present locally and need not be pinned, nothing is inserted on success, + // and nothing is left behind on failure. The bytes land ONLY in the caller's + // device blocks; the host blocks this call allocates are pure staging and + // are returned to the pool on both the success and the failure path. No + // local host copy is retained, so a later local load() of the same hash is + // still a miss. + // + // device_block_ids is mandatory and must match block_hashes in size, as must + // slices; any other size is InvalidArgument. + // + // The device blocks are written before the source's verdict is known, so on + // failure their contents are UNDEFINED -- treat them as scratch until the + // read reports success. + // + // Compare with Load(): both bring a peer's block into local HBM. Load() + // fetches through the store's own path and is the right call when the hash + // may be resident locally; ReadRemote() takes a lease on the source and is + // the right call when the caller already knows the source coordinates and + // wants no local record of the transfer. // // Requires a global registry: it is what maps the owning peer to the // controller address this store acquires its read lease from. A store built // without one fails every read with FailedPrecondition. absl::Status ReadRemote(const std::vector& block_hashes, - const std::vector& device_block_ids = {}); + const std::vector& slices, + const std::vector& device_block_ids); // Polls status of active remote reads. // Returns {done_hashes, failed_hashes, pending_hashes} @@ -486,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 { @@ -494,12 +534,12 @@ class KVCacheStore { // cached controller addresses -- the poller is where failure is observed, // and by then the grouping is gone. std::vector src_raiden_ids; - // The local landing blocks. These live HERE and nowhere else until the - // poller commits -- stamping them into the LRU entry at issue time would - // destroy the peer coordinate the entry needs for a retry. + // The local staging blocks the bytes hop through on their way to HBM. They + // live HERE and nowhere else: no LRU entry ever points at them, so the + // poller returns them to the pool on both the success and the failure + // path. The caller's device blocks are not tracked -- once the transfer is + // terminal this store has no further interest in them. std::vector host_block_ids; - // Empty for a read to host; otherwise the caller's device blocks. - std::vector device_block_ids; }; struct FutureHash { @@ -578,6 +618,12 @@ class KVCacheStore { active_remote_reads_ ABSL_GUARDED_BY(mutex_); std::vector active_remote_writes_ ABSL_GUARDED_BY(mutex_); + // Operations a poller has claimed and is currently asking the destination + // about. Polling drops mutex_ to make that RPC, so without a claim two + // concurrent pollers both see the same operation as committed, both call + // FinishRemoteWrite, and the internal pin is released twice while the hashes + // are reported twice. + absl::flat_hash_set polling_remote_writes_ ABSL_GUARDED_BY(mutex_); std::vector done_remote_writes_ ABSL_GUARDED_BY(mutex_); std::vector failed_remote_writes_ ABSL_GUARDED_BY(mutex_); std::vector existing_remote_writes_ ABSL_GUARDED_BY(mutex_); diff --git a/tpu_sync/kv_cache/kv_cache_store_test.cc b/tpu_sync/kv_cache/kv_cache_store_test.cc index 9aefbf06..0cb7dd7f 100644 --- a/tpu_sync/kv_cache/kv_cache_store_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_test.cc @@ -125,6 +125,16 @@ using TestHostOffloadBackend = HostOffloadBackendTest::Backend; namespace { +// Observes what is resident WITHOUT acquiring it. The application-level +// Lookup pins what it returns, which is right for a caller that goes on to +// load or save the answer but wrong for a case that is asserting on pin counts +// or eviction candidates -- there the pin is the thing under test, and a +// lookup taking one of its own would be measuring itself. +absl::StatusOr PeekLookup( + KVCacheStore& store, const std::vector& block_hashes) { + return store.Lookup(block_hashes, LookupOptions{.enable_global = false}); +} + // Publishes a peer so a remote read can resolve its controller. The store // server address is required by the registry but a read never dials it -- it // speaks to the controller. @@ -186,6 +196,12 @@ TEST(KVCacheStoreTest, BasicTests) { EXPECT_EQ(lookup_res_early->size(), 1); EXPECT_EQ((*lookup_res_early)[0].first, "4001"); + // Lookup pins what it returns, and this test only inspects the answers + // rather than going on to load or save them, so it gives the pins back + // itself. "4001" was returned by both lookups, so it holds two. + controller.Release({"4001", "4002"}); + controller.Release({"4001"}); + // 3. Delete controller.Delete(hashes, slices); EXPECT_TRUE( @@ -347,7 +363,8 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // Case 1: Full local hit, no global hit { - auto lookup_res = store.Lookup({"local_only_hash"}); + auto lookup_res = store.Lookup({"local_only_hash"}, + /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); @@ -358,7 +375,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // Case 2: Both local and global has the same hit, but we return local hit // results { - auto lookup_res = store.Lookup({"shared_hash"}); + auto lookup_res = store.Lookup({"shared_hash"}, /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "shared_hash"); @@ -369,7 +386,8 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // Case 3: No local hit, only global hits { - auto lookup_res = store.Lookup({"global_hash_1", "global_hash_2"}); + auto lookup_res = store.Lookup({"global_hash_1", "global_hash_2"}, + /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 2); @@ -400,7 +418,8 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // It should return both local and global { auto lookup_res = - store.Lookup({"local_only_hash", "global_hash_1", "global_hash_2"}); + store.Lookup({"local_only_hash", "global_hash_1", "global_hash_2"}, + /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 3); @@ -422,7 +441,8 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // It should stop at the first miss in registry { auto lookup_res = store.Lookup( - {"local_only_hash", "global_hash_1", "missing_hash", "global_hash_2"}); + {"local_only_hash", "global_hash_1", "missing_hash", "global_hash_2"}, + /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 2); // local_only_hash, global_hash_1 EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); @@ -728,7 +748,7 @@ TEST(KVCacheStoreTest, LookupCapLimitWithGlobal) { // Lookup 3 hashes, but capacity is 2. It should only return 2. std::vector lookup_hashes = {"global_hash_1", "global_hash_2", "global_hash_3"}; - auto lookup_res = store.Lookup(lookup_hashes); + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); EXPECT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); @@ -771,7 +791,7 @@ TEST(KVCacheStoreTest, LookupCapLimitMixed) { // global). std::vector lookup_hashes = {"local_hash_1", "global_hash_2", "global_hash_3"}; - auto lookup_res = store.Lookup(lookup_hashes); + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); ASSERT_TRUE(lookup_res.ok()); EXPECT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "local_hash_1"); @@ -887,11 +907,11 @@ TEST(KVCacheStoreTest, ReleaseAndDelete) { // remote_1 and remote_2 should be unpinned and deleted (since REMOTE) EXPECT_EQ(store.GetPinCount("remote_1"), 0); EXPECT_EQ(store.GetPinCount("remote_2"), 0); - EXPECT_EQ(store.Lookup({"remote_1"})->size(), 0); - EXPECT_EQ(store.Lookup({"remote_2"})->size(), 0); + EXPECT_EQ(PeekLookup(store, {"remote_1"})->size(), 0); + EXPECT_EQ(PeekLookup(store, {"remote_2"})->size(), 0); // local_1 and local_2 should be restored to the cache! - auto lookup_res = store.Lookup({"local_1", "local_2"}); + auto lookup_res = PeekLookup(store, {"local_1", "local_2"}); ASSERT_TRUE(lookup_res.ok()); EXPECT_EQ(lookup_res->size(), 2); @@ -901,7 +921,7 @@ TEST(KVCacheStoreTest, ReleaseAndDelete) { auto res_non_remote = store.ReleaseAndDelete({"local_1"}); EXPECT_EQ(res_non_remote, 0); EXPECT_EQ(store.GetPinCount("local_1"), 0); - EXPECT_EQ(store.Lookup({"local_1"})->size(), 1); + EXPECT_EQ(PeekLookup(store, {"local_1"})->size(), 1); // Test remote block pinned twice: after one ReleaseAndDelete, pin count is 1 // so it should NOT be deleted! @@ -911,7 +931,7 @@ TEST(KVCacheStoreTest, ReleaseAndDelete) { auto res_pinned = store.ReleaseAndDelete({"remote_1"}); EXPECT_EQ(res_pinned, 0); // 0 deleted because pin count was 2 -> 1 EXPECT_EQ(store.GetPinCount("remote_1"), 1); - EXPECT_EQ(store.Lookup({"remote_1"})->size(), 1); + EXPECT_EQ(PeekLookup(store, {"remote_1"})->size(), 1); store.Release({"remote_1"}); store.Delete({"remote_1"}, {remote_slices[0]}); @@ -1469,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)); @@ -1488,7 +1512,59 @@ 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 = + *::tpu_raiden::controller::RaidenController::Create(unit_, 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) { @@ -1607,13 +1683,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) { @@ -1690,14 +1767,96 @@ 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 = + *::tpu_raiden::controller::RaidenController::Create(unit_, 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) { @@ -1994,6 +2153,9 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictByHashesHostAndHbmToErased) { ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 2); } + // ...and give back the pins that check took, or the evict below has nothing + // it is allowed to reclaim. + store.Release({"hash_1", "hash_2"}); // 5. Check locked blocks on controller auto* controller_ptr = KVCacheStoreTest::GetController(store); @@ -2224,9 +2386,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { } store.Release(hashes); - // Verify both are HOST_AND_HBM + // Verify both are HOST_AND_HBM. Observed rather than looked up: this case is + // about which block evicts first, and a pin/unpin round trip moves a node + // through the pinned list and back, which reorders the very LRU under test. { - auto lookup_res = store.Lookup({"hash_B", "hash_A"}); + auto lookup_res = PeekLookup(store, {"hash_B", "hash_A"}); ASSERT_TRUE(lookup_res.ok()); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); @@ -2358,14 +2522,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { KVCacheStore store(10, std::move(dst_controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); - // Insert and pin remote block in local store + // The source coordinates come from the CALLER now: nothing is inserted into + // the local LRU, and the hash need not be known locally at all. std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); + const std::vector device_blocks = {7}; - // Trigger ReadRemote - absl::Status status = store.ReadRemote(hashes); + absl::Status status = store.ReadRemote(hashes, slices, device_blocks); ASSERT_TRUE(status.ok()) << status.message(); // Poll for completion @@ -2383,41 +2547,95 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { } ASSERT_TRUE(done); // The DESTINATION's worker executed the pull, against the source's - // authoritative block id (42, from the verify hook) and into the landing - // block the store allocated. - EXPECT_EQ(dst_transfer_mock_->vector_h2h_read_calls, 1); + // authoritative block id (42, from the verify hook), through the host + // staging block the store allocated, and into the CALLER's device block. + EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 1); EXPECT_THAT(dst_transfer_mock_->last_src_offsets, ::testing::ElementsAre(42)); - EXPECT_THAT(dst_transfer_mock_->last_dst_offsets, ::testing::ElementsAre(0)); + EXPECT_THAT(dst_transfer_mock_->last_staging_offsets, + ::testing::ElementsAre(0)); + EXPECT_THAT(dst_transfer_mock_->last_dst_offsets, ::testing::ElementsAre(7)); - // Verify status in LRU is HOST, host_block_id is 0 + // A successful read leaves NO local record: the bytes are in the caller's + // device block and nowhere else. A later local lookup is still a miss. 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::HOST); - EXPECT_EQ((*lookup_res)[0].second.host_block_id, 0); + EXPECT_TRUE(lookup_res->empty()); - // Verify registration in global registry (need to poll registry since - // registration is async) + // ...and nothing is advertised to the registry. There is no host-resident + // copy here to serve to a peer, so publishing one would advertise a block + // this node does not have. auto channel = grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); global_registry::GlobalRegistryClient registry_client(channel); + auto registry_lookup = registry_client.Lookup(hashes); + ASSERT_TRUE(registry_lookup.ok()); + EXPECT_TRUE(registry_lookup->empty()) + << "read_remote must not advertise the read block to the registry"; - bool registered = false; - std::vector metadata_results; - for (int attempt = 0; attempt < 100; ++attempt) { - auto lookup_res = registry_client.Lookup(hashes); - if (lookup_res.ok() && lookup_res->size() == 1) { - metadata_results = *std::move(lookup_res); - registered = true; - break; - } - absl::SleepFor(absl::Milliseconds(10)); + registry_server->Shutdown(); +} + +// The host blocks a read stages through are a hop, not a destination. Nothing +// in the LRU points at them, so if the poller did not return them to the pool +// each read would burn one host block permanently. Reading more blocks in +// total than the pool holds only works if every read gives its staging back. +TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { + 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); + + auto src_controller_server = core::controller::CreateTestControllerServer(); + kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; + ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, + src_controller_server->server_address)); + { + auto st = src_controller_server->client->RegisterWorker( + "worker_0", "src_worker_0_addr", {{"src_worker_0_transfer", {}}}); + ASSERT_TRUE(st.ok()) << st.message(); } - ASSERT_TRUE(registered) - << "Block hashes were not registered in global registry"; + src_controller_server->service->SetReadRemoteHooks( + [&](absl::Span h) + -> absl::StatusOr> { + return std::vector(h.size(), 42); + }, + [&](absl::Span /*h*/) {}); - EXPECT_EQ(metadata_results[0].raiden_id().job_name(), rid.job_name); - EXPECT_EQ(metadata_results[0].block_id(), 0); + // A deliberately small host pool: three reads of two blocks each cannot fit + // in four blocks unless each read's staging is reclaimed. + constexpr int kHostBlocks = 4; + auto dst_controller = MakeController(kHostBlocks); + RegisterAndInitWorker(*dst_controller, "worker_0", + test_server_->server_address); + RaidenId rid{"dst_job", "0", "dst_cache", 0}; + KVCacheStore store(kHostBlocks, std::move(dst_controller), registry_address, + rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); + + const std::vector slices = { + RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE), + RaidenBlockID(src_raiden_id, 43, BlockStatus::REMOTE)}; + + for (int round = 0; round < 3; ++round) { + std::vector hashes = {absl::StrCat("hash_", round, "_a"), + absl::StrCat("hash_", round, "_b")}; + absl::Status status = store.ReadRemote(hashes, slices, {7, 8}); + ASSERT_TRUE(status.ok()) + << "round " << round << " failed to launch: " << status.message() + << " -- staging blocks from an earlier round were not reclaimed"; + + bool done = false; + for (int attempt = 0; attempt < 200 && !done; ++attempt) { + auto [done_hashes, failed_hashes, pending] = store.PollRemoteReadStatus(); + ASSERT_TRUE(failed_hashes.empty()); + if (done_hashes.size() == hashes.size()) done = true; + if (!done) absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(done) << "round " << round << " never completed"; + } registry_server->Shutdown(); } @@ -2440,16 +2658,15 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteWithoutRegistryFails) { std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - absl::Status status = store.ReadRemote(hashes); + absl::Status status = store.ReadRemote(hashes, slices, {7}); EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); EXPECT_THAT(std::string(status.message()), ::testing::HasSubstr("global registry")); - EXPECT_EQ(dst_transfer_mock_->vector_h2h_read_calls, 0); - // Rejected cleanly: the same hashes are admissible again, and the landing + EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 0); + // Rejected cleanly: the same hashes are admissible again, and the staging // blocks went back to the pool. - EXPECT_EQ(store.ReadRemote(hashes).code(), + EXPECT_EQ(store.ReadRemote(hashes, slices, {7}).code(), absl::StatusCode::kFailedPrecondition); } @@ -2473,12 +2690,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - absl::Status status = store.ReadRemote(hashes); + absl::Status status = store.ReadRemote(hashes, slices, {7}); EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); EXPECT_THAT(std::string(status.message()), ::testing::HasSubstr("src_job")); - EXPECT_EQ(dst_transfer_mock_->vector_h2h_read_calls, 0); + EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 0); } // The peer's controller address is cached, so a repeat read costs no registry @@ -2528,8 +2744,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector hashes = {hash}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - EXPECT_TRUE(store.InsertAndLock(hashes, slices, true)); - EXPECT_TRUE(store.ReadRemote(hashes).ok()); + EXPECT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()); for (int attempt = 0; attempt < 300; ++attempt) { auto [done, failed, pending] = store.PollRemoteReadStatus(); if (!done.empty()) return true; @@ -2641,28 +2856,15 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { KVCacheStore store(2, std::move(dst_controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); - // Fill cache with two local blocks - std::vector local_hashes = {"local_1", "local_2"}; - std::vector local_slices = { - RaidenBlockID(rid, -1, BlockStatus::HOST), - RaidenBlockID(rid, -1, BlockStatus::HOST)}; - ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); - - // Unpin local_1 so it is evictable - store.Release({"local_1"}); - - // Insert and pin remote block (evicts local_1) std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); // The transfer now runs on the DESTINATION, so that is where the failure is // injected. dst_transfer_mock_->fail_transfers = true; - // Trigger ReadRemote - absl::Status status = store.ReadRemote(hashes); + absl::Status status = store.ReadRemote(hashes, slices, {7}); ASSERT_TRUE(status.ok()) << status.message(); // Poll for failure @@ -2680,29 +2882,22 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { } ASSERT_TRUE(failed); - // Verify hash_0 is still REMOTE + // A failed read leaves NOTHING to clean up. There is no entry to delete and + // no candidate to restore -- the read never touched the LRU. This is the + // whole point of taking the source coordinates as an argument: the old + // design stranded a REMOTE entry here that only release_and_delete could + // remove. { 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::REMOTE); + EXPECT_TRUE(lookup_res->empty()); } - // Caller calls ReleaseAndDelete to clean up failed remote read - size_t deleted = store.ReleaseAndDelete(hashes); - EXPECT_EQ(deleted, 1); - - // Verify hash_0 is deleted, local_1 is restored - { - auto lookup_res = store.Lookup(hashes); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 0); - } - { - auto lookup_res = store.Lookup({"local_1"}); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 1); - } + // The staging blocks went back to the pool, so the read can be retried. If + // the failure path leaked them this second launch would be the one to fail. + dst_transfer_mock_->fail_transfers = false; + EXPECT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()) + << "a failed read must return its staging blocks"; registry_server->Shutdown(); } @@ -2746,9 +2941,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - ASSERT_TRUE(store.ReadRemote(hashes).ok()); + ASSERT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()); bool failed = false; for (int attempt = 0; attempt < 100; ++attempt) { @@ -2766,77 +2960,15 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, // The block_hash flowed to the source and the transfer was never dispatched. EXPECT_THAT(validated, ::testing::ElementsAre("hash_0")); EXPECT_FALSE(transfer_ran); - // The block is still REMOTE on the destination (not promoted to HOST). + // Nothing was recorded locally, on this path as on every other. 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::REMOTE); + EXPECT_TRUE(lookup_res->empty()); } // The source verify hook accepts the hash -> the transfer runs and the // destination block is promoted to HOST. Confirms the block_hashes reach the // source verify path on the success flow. -TEST_F(KVCacheStoreEmbeddedControllerTest, - ReadRemoteSourceVerifySuccessTransfers) { - auto src_controller_server = core::controller::CreateTestControllerServer(); - ::tpu_sync::rpc::RaidenIdProto src_unit; - src_unit.set_job_name("src_job"); - src_unit.set_job_replica_id("0"); - src_unit.set_data_name("src_data"); - src_unit.set_data_replica_idx(0); - kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, - src_controller_server->server_address)); - - std::vector validated, unpinned; - src_controller_server->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - validated.assign(h.begin(), h.end()); - return std::vector{42}; // authoritative source id - }, - [&](absl::Span h) { - unpinned.assign(h.begin(), h.end()); - }); - // NOTE: the source no longer transfers anything. Under the pull design - // the DESTINATION's own worker (test_server_, backed by a mock transfer - // manager) executes the copy; the source only leases. - - auto dst_controller = MakeController(); - RegisterAndInitWorker(*dst_controller, "worker_0", - test_server_->server_address); - RaidenId rid{"dst_job", "0", "dst_cache", 0}; - KVCacheStore store(2, std::move(dst_controller), registry_address_, rid, - std::nullopt, - /*store_server_ip=*/"127.0.0.1"); - - std::vector hashes = {"hash_0"}; - std::vector slices = { - RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - - ASSERT_TRUE(store.ReadRemote(hashes).ok()); - - bool done = false; - for (int attempt = 0; attempt < 100; ++attempt) { - auto [done_hashes, failed_hashes, pending_hashes] = - store.PollRemoteReadStatus(); - ASSERT_TRUE(failed_hashes.empty()); - if (!done_hashes.empty()) { - EXPECT_THAT(done_hashes, ::testing::ElementsAre("hash_0")); - done = true; - break; - } - absl::SleepFor(absl::Milliseconds(10)); - } - ASSERT_TRUE(done); - EXPECT_THAT(validated, ::testing::ElementsAre("hash_0")); - EXPECT_THAT(unpinned, ::testing::ElementsAre("hash_0")); - 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::HOST); -} TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteDuplicateFails) { auto src_controller_server = core::controller::CreateTestControllerServer(); @@ -2893,14 +3025,13 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteDuplicateFails) { std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); // First call succeeds - absl::Status status1 = store.ReadRemote(hashes); + absl::Status status1 = store.ReadRemote(hashes, slices, {7}); ASSERT_TRUE(status1.ok()) << status1.message(); // Second call fails with FailedPreconditionError - absl::Status status2 = store.ReadRemote(hashes); + absl::Status status2 = store.ReadRemote(hashes, slices, {8}); EXPECT_FALSE(status2.ok()); EXPECT_EQ(status2.code(), absl::StatusCode::kFailedPrecondition); EXPECT_THAT(status2.message(), ::testing::HasSubstr("already reading")); @@ -2956,22 +3087,20 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteAllocationFailureAborts) { ASSERT_TRUE(save_done); // Now, 1 host block is allocated, free = 0. And local_1 remains pinned. - // Insert and pin remote block hash_0 kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_cache", 0}; std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - // ReadRemote should fail because allocation of host block fails (0 free, 0 - // evictable) - absl::Status status = store.ReadRemote(hashes); + // ReadRemote should fail because allocation of the staging block fails + // (0 free, 0 evictable) + absl::Status status = store.ReadRemote(hashes, slices, {7}); EXPECT_FALSE(status.ok()); EXPECT_EQ(status.code(), absl::StatusCode::kResourceExhausted); // Verify hash_0 is NOT in reading_hashes_ (so calling it again doesn't report // duplicate) - absl::Status status2 = store.ReadRemote(hashes); + absl::Status status2 = store.ReadRemote(hashes, slices, {7}); EXPECT_EQ(status2.code(), absl::StatusCode::kResourceExhausted); registry_server->Shutdown(); @@ -3065,15 +3194,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { KVCacheStore store(10, std::move(dst_controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); - // Insert and pin remote block hash_0 and hash_1 + // hash_0 lives on one peer and hash_1 on another; the caller names both. std::vector hashes = {"hash_0", "hash_1"}; std::vector slices = { RaidenBlockID(src_raiden_id_1, 10, BlockStatus::REMOTE), RaidenBlockID(src_raiden_id_2, 20, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); // Trigger ReadRemote for both - absl::Status status = store.ReadRemote(hashes); + absl::Status status = store.ReadRemote(hashes, slices, {7, 8}); ASSERT_TRUE(status.ok()) << status.message(); // A batch spanning two peers takes one lease per peer and joins the futures, @@ -3096,12 +3224,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { } ASSERT_TRUE(done); - // Verify both statuses are HOST + // Neither peer's block is recorded locally, however many peers were involved. auto lookup_res = store.Lookup(hashes); ASSERT_TRUE(lookup_res.ok()); - ASSERT_EQ(lookup_res->size(), 2); - EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST); - EXPECT_EQ((*lookup_res)[1].second.status, BlockStatus::HOST); + EXPECT_TRUE(lookup_res->empty()); registry_server->Shutdown(); } @@ -3162,95 +3288,6 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ::testing::FieldsAre(1, "hash_2", 1))); } -TEST_F(KVCacheStoreEmbeddedControllerTest, - ReadRemoteSetsMetadataEntriesOnCompletion) { - // Same setup as ReadRemoteSuccess: registry, src controller with a - // successful H2H callback, dest store — here with a metadata table attached. - 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); - - auto src_controller_server = core::controller::CreateTestControllerServer(); - - ::tpu_sync::rpc::RaidenIdProto src_unit; - src_unit.set_job_name("src_job"); - src_unit.set_job_replica_id("0"); - src_unit.set_data_name("src_data"); - src_unit.set_data_replica_idx(0); - - kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - src_controller_server->server_address)); - - auto worker_status = src_controller_server->client->RegisterWorker( - "worker_0", "src_worker_0_addr", {{"src_worker_0_transfer", {}}}); - ASSERT_TRUE(worker_status.ok()) << worker_status.message(); - - // Every read is now validated at the source by construction -- there is no - // longer any RPC that transfers without verifying and pinning first. Grant - // the lease and echo back authoritative ids. - src_controller_server->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); - // NOTE: the source no longer transfers anything. Under the pull design - // the DESTINATION's own worker (test_server_, backed by a mock transfer - // manager) executes the copy; the source only leases. - - auto dst_controller = MakeController(); - RegisterAndInitWorker(*dst_controller, "worker_0", - test_server_->server_address); - - MetadataRegion metadata_region(10); - auto metadata_or = KVCacheMetadata::Format(metadata_region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); - - RaidenId rid{"dst_job", "0", "dst_cache", 0}; - KVCacheStore store(10, std::move(dst_controller), registry_address, rid, - *metadata_or, /*store_server_ip=*/"127.0.0.1"); - - std::vector hashes = {"hash_0"}; - std::vector slices = { - RaidenBlockID(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.InsertAndLock(hashes, slices, true)); - - // InsertAndLock has already called SetMetadataEntry for the slice, but its - // REMOTE status fails the same data-lives-in-local-host-memory filter: a - // REMOTE entry owns no local data and must stay out of the table. - EXPECT_THAT(metadata_or->ValidEntries(), ::testing::IsEmpty()); - - absl::Status status = store.ReadRemote(hashes); - ASSERT_TRUE(status.ok()) << status.message(); - - // Poll for completion - bool done = false; - for (int attempt = 0; attempt < 100; ++attempt) { - auto [done_hashes, failed_hashes, pending_hashes] = - store.PollRemoteReadStatus(); - ASSERT_TRUE(failed_hashes.empty()); - if (!done_hashes.empty()) { - done = true; - break; - } - absl::SleepFor(absl::Milliseconds(10)); - } - ASSERT_TRUE(done); - - // Read completion lands the remote data on local host block 0, which is - // when the binding enters the table. - EXPECT_THAT(metadata_or->ValidEntries(), - ElementsAre(::testing::FieldsAre(0, "hash_0", 0))); - - registry_server->Shutdown(); -} TEST(KVCacheStoreTest, RecoverFromLocalManifestRebuildsLruCache) { RaidenId rid{"manifest_job", "0", "kv_cache", 0}; @@ -4452,6 +4489,125 @@ TEST_F(RemoteWriteSourceTest, TheOfferAsksForLessThanTheSourceWillHold) { absl::ToInt64Milliseconds(absl::Seconds(25))); } +// Polling a remote write means asking the destination, which cannot be done +// holding the store's lock. Every public poll entry point drives that ask, and +// so does the background loop, so several can be in it at once. If the poller +// works off a copy of the active list rather than claiming entries, two of +// them settle the same operation: the internal pin is released twice -- the +// second release landing on whatever now owns that block -- and every hash is +// reported to the caller twice. +TEST_F(RemoteWriteSourceTest, ConcurrentPollsSettleAnOfferExactlyOnce) { + RaidenId src{"rw_src_race", "0", "kv", 0}; + RaidenId dst{"rw_dst_race", "0", "kv", 0}; + auto src_store = MakeStore(src); + Populate(*src_store, src, {"a", "b"}); + StartFakeDestination(dst); + + proto::PollWriteRemoteResponse verdict; + verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); + verdict.add_committed_hashes("a"); + verdict.add_committed_hashes("b"); + fake_destination_.SetPollResponse(verdict); + + ASSERT_TRUE(src_store->WriteRemote({"a", "b"}, dst).ok()); + + // PollSaveStatus drives the write poller as a side effect (they share + // PollFuturesInternal), so these four threads plus the background loop are + // five callers racing into it. + std::atomic stop{false}; + std::vector drivers; + for (int t = 0; t < 4; ++t) { + drivers.emplace_back([&]() { + while (!stop.load(std::memory_order_relaxed)) { + (void)src_store->PollSaveStatus(); + } + }); + } + + std::vector settled; + const auto drain = [&]() { + auto [done, failed, pending, existing, unregistered] = + src_store->PollRemoteWriteStatus(); + settled.insert(settled.end(), done.begin(), done.end()); + settled.insert(settled.end(), failed.begin(), failed.end()); + }; + for (int i = 0; i < 300 && settled.size() < 2; ++i) { + drain(); + absl::SleepFor(absl::Milliseconds(10)); + } + // Keep draining after both arrive: a duplicate settle shows up late, and + // stopping at the first two hashes would never see it. + absl::SleepFor(absl::Milliseconds(200)); + drain(); + + stop.store(true, std::memory_order_relaxed); + for (auto& driver : drivers) driver.join(); + + EXPECT_THAT(settled, ::testing::UnorderedElementsAre("a", "b")) + << "the offer was settled more than once"; + // The caller's own pin survives; only the internal one was released, once. + EXPECT_EQ(src_store->GetPinCount("a"), 1); + EXPECT_EQ(src_store->GetPinCount("b"), 1); +} + +// A store destroyed with an offer still outstanding must not leave its +// internal pin behind. Backends are shared_ptrs and can outlive the store that +// pinned into them, and nothing else knows to release that pin -- the block +// would sit unreclaimable for the life of the backend. +TEST_F(RemoteWriteSourceTest, DestroyingAStoreMidOfferReleasesItsInternalPin) { + RaidenId src{"rw_src_dtor_pin", "0", "kv", 0}; + RaidenId dst{"rw_dst_dtor_pin", "0", "kv", 0}; + auto src_store = MakeStore(src); + Populate(*src_store, src, {"a"}); + StartFakeDestination(dst); + + // The destination never reaches a verdict, so the offer stays outstanding. + proto::PollWriteRemoteResponse verdict; + verdict.set_state(proto::PollWriteRemoteResponse::PENDING); + fake_destination_.SetPollResponse(verdict); + + ASSERT_TRUE(src_store->WriteRemote({"a"}, dst).ok()); + // Populate's pin plus the offer's internal one. + EXPECT_EQ(src_store->GetPinCount("a"), 2); + + // Outlive the store, exactly as a caller sharing a backend would. + std::shared_ptr backend = src_store->backend(); + src_store.reset(); + + EXPECT_EQ(backend->GetPinCount("a"), 1) + << "the offer's internal pin outlived the store that took it"; +} + +// The same drain must not WAIT for those offers. An offer goes terminal only +// when the destination answers or the ~30s HOLD expires, so waiting would turn +// destroying a store behind a dead peer into a half-minute stall. +TEST_F(RemoteWriteSourceTest, DestroyingAStoreMidOfferDoesNotWaitForTheHold) { + RaidenId src{"rw_src_dtor_fast", "0", "kv", 0}; + RaidenId dst{"rw_dst_dtor_fast", "0", "kv", 0}; + auto src_store = MakeStore(src); + Populate(*src_store, src, {"a"}); + StartFakeDestination(dst); + + proto::PollWriteRemoteResponse verdict; + verdict.set_state(proto::PollWriteRemoteResponse::PENDING); + fake_destination_.SetPollResponse(verdict); + + ASSERT_TRUE(src_store->WriteRemote({"a"}, dst).ok()); + { + auto [done, failed, pending, existing, unregistered] = + src_store->PollRemoteWriteStatus(); + ASSERT_THAT(pending, ::testing::ElementsAre("a")); + } + + const absl::Time before = absl::Now(); + src_store.reset(); + const absl::Duration took = absl::Now() - before; + + // The HOLD is 30s; anything near it means the drain waited. + EXPECT_LT(took, absl::Seconds(10)) + << "destruction blocked on the remote write's hold window"; +} + // A destination that vanished after registering. The source must get a prompt // error rather than waiting, and must drop the cached client so a restarted // peer is reachable. @@ -4787,7 +4943,8 @@ std::unique_ptr MakeStoreInterleaveFixture( TEST(KVCacheStoreTest, LookupInterleavesLocalAndRemoteThroughTheStore) { auto f = MakeStoreInterleaveFixture(); - auto res = f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}); + auto res = f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}, + /*enable_global=*/true); ASSERT_TRUE(res.ok()); ASSERT_EQ(res->size(), 4); @@ -4836,7 +4993,7 @@ TEST(KVCacheStoreTest, LookupInterleavedDisabledThroughTheStore) { ASSERT_EQ(legacy->size(), 1); EXPECT_EQ((*legacy)[0].first, "r1"); - auto interleaved = f->store->Lookup({"r1", "l1"}); + auto interleaved = f->store->Lookup({"r1", "l1"}, /*enable_global=*/true); ASSERT_TRUE(interleaved.ok()); ASSERT_EQ(interleaved->size(), 2); EXPECT_EQ((*interleaved)[1].first, "l1"); diff --git a/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc b/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc index d66277ea..f46920b0 100644 --- a/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc @@ -120,6 +120,7 @@ TEST_F(KVCacheStoreWrapperTest, ColdStartCreatesMetadataTable) { auto wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); EXPECT_TRUE(MetadataSegmentExists("_metadata")); auto lookup_or = (*wrapper)->Lookup({"host_1"}); + (*wrapper)->Release({"host_1"}); ASSERT_TRUE(lookup_or.ok()); EXPECT_THAT(*lookup_or, IsEmpty()); } @@ -142,6 +143,7 @@ TEST_F(KVCacheStoreWrapperTest, RecoversHostBlocksAfterRestart) { wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); auto lookup_or = (*wrapper)->Lookup({"host_1", "host_2"}); + (*wrapper)->Release({"host_1", "host_2"}); ASSERT_TRUE(lookup_or.ok()); ASSERT_EQ(lookup_or->size(), 2); EXPECT_EQ((*lookup_or)[0].first, "host_1"); @@ -162,6 +164,7 @@ TEST_F(KVCacheStoreWrapperTest, ModelUidMismatchColdStarts) { wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); auto lookup_or = (*wrapper)->Lookup({"host_1"}); + (*wrapper)->Release({"host_1"}); ASSERT_TRUE(lookup_or.ok()); EXPECT_THAT(*lookup_or, IsEmpty()); }