"""Unit tests for EmbeddingCacheController — LRU eviction and RDMA ref counting.""" from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") import threading import time import unittest from unittest.mock import MagicMock import torch from sglang.srt.mem_cache.storage.mooncake_store.embedding_cache_controller import ( ContiguousMemoryAllocator, EmbeddingCacheController, EmbeddingInsertOperation, ) # --------------------------------------------------------------------------- # ContiguousMemoryAllocator tests # --------------------------------------------------------------------------- class TestContiguousMemoryAllocator(unittest.TestCase): def test_basic_alloc_free(self): alloc = ContiguousMemoryAllocator(1024) a = alloc.allocate(256) self.assertIsNotNone(a) self.assertEqual(a, 0) b = alloc.allocate(256) self.assertEqual(b, 256) alloc.free(a, 256) c = alloc.allocate(128) self.assertEqual(c, 0) # reused freed block def test_alloc_fails_when_full(self): alloc = ContiguousMemoryAllocator(256) a = alloc.allocate(256) self.assertIsNotNone(a) b = alloc.allocate(1) self.assertIsNone(b) def test_free_merges_adjacent(self): alloc = ContiguousMemoryAllocator(512) a = alloc.allocate(128) b = alloc.allocate(128) c = alloc.allocate(256) alloc.free(a, 128) alloc.free(b, 128) # The two 128-byte blocks should merge into one 256-byte free block d = alloc.allocate(256) self.assertIsNotNone(d) self.assertEqual(d, 0) def test_allocated_size_tracking(self): alloc = ContiguousMemoryAllocator(1024) self.assertEqual(alloc.get_allocated_size(), 0) a = alloc.allocate(300) self.assertEqual(alloc.get_allocated_size(), 300) b = alloc.allocate(200) self.assertEqual(alloc.get_allocated_size(), 500) alloc.free(a, 300) self.assertEqual(alloc.get_allocated_size(), 200) alloc.free(b, 200) self.assertEqual(alloc.get_allocated_size(), 0) def test_free_size_tracking(self): alloc = ContiguousMemoryAllocator(1024) self.assertEqual(alloc.get_free_size(), 1024) alloc.allocate(400) self.assertEqual(alloc.get_free_size(), 624) def test_double_free_is_safe(self): alloc = ContiguousMemoryAllocator(256) a = alloc.allocate(128) alloc.free(a, 128) # Second free of same offset — should not corrupt state alloc.free(a, 128) self.assertEqual(alloc.get_allocated_size(), 0) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_controller( pool_mb=1.0, enable_eviction=True, hidden_dims=None, max_eviction_batch=10 ): """Create an EmbeddingCacheController with a mocked MooncakeEmbeddingStore.""" ctrl = EmbeddingCacheController.__new__(EmbeddingCacheController) ctrl.tp_world_size = 1 ctrl.tp_group = None ctrl.tp_rank = 0 ctrl.all_rank_get = False ctrl.hidden_dims = hidden_dims or {"image": 1024} ctrl.element_size = torch.float32.itemsize ctrl.enable_eviction = enable_eviction ctrl.max_eviction_batch = max_eviction_batch # Small pool for testing (1 MB by default) ctrl.total_pool_size_bytes = int(pool_mb * 1024**2) ctrl.cpu_pool = torch.empty( ctrl.total_pool_size_bytes, dtype=torch.uint8, pin_memory=False ) # Mock the mooncake store — no real RDMA ctrl.mooncake_store = MagicMock() ctrl.mooncake_store.register_buffer = MagicMock() ctrl.allocator = ContiguousMemoryAllocator(ctrl.total_pool_size_bytes) ctrl.hash_to_metadata = {} ctrl.access_order = {} ctrl.access_lock = threading.Lock() ctrl.ref_counts = {} ctrl.stats = { "total_allocated": 0, "total_evicted": 0, "eviction_count": 0, "allocation_failures": 0, } ctrl.ongoing_prefetch = {} ctrl.prefetch_queue = MagicMock() ctrl.insert_queue = MagicMock() ctrl.lock = threading.Lock() ctrl.stop_event = threading.Event() # Do NOT start the IO thread — tests drive _io_loop logic manually ctrl.io_thread = MagicMock() ctrl.prefetch_tp_group = None return ctrl def _embedding_bytes(num_tokens, dim): return num_tokens * dim * torch.float32.itemsize # --------------------------------------------------------------------------- # LRU eviction tests # --------------------------------------------------------------------------- class TestLRUEviction(unittest.TestCase): def test_evict_oldest_first(self): dim = 64 size = _embedding_bytes(1, dim) # 256 bytes # Pool exactly fits 3 entries (768 bytes) pool_bytes = size * 3 ctrl = _make_controller(pool_mb=pool_bytes / (1024**2)) # Insert 3 entries — fills the pool for i in range(3): h = f"hash_{i}" with ctrl.lock: offset = ctrl.allocator.allocate(size) self.assertIsNotNone(offset) ctrl.hash_to_metadata[h] = (offset, 1, dim, size) ctrl._update_access_time(h) # Pool is full. Inserting a 4th should evict hash_0 (oldest). h_new = "hash_new" with ctrl.lock: offset = ctrl._allocate_with_eviction(size) self.assertIsNotNone(offset) with ctrl.lock: self.assertNotIn("hash_0", ctrl.hash_to_metadata) self.assertIn("hash_1", ctrl.hash_to_metadata) self.assertIn("hash_2", ctrl.hash_to_metadata) def test_eviction_disabled(self): dim = 64 size = _embedding_bytes(1, dim) # 256 bytes # Pool exactly fits 1 entry ctrl = _make_controller(pool_mb=size / (1024**2), enable_eviction=False) # Fill the pool with ctrl.lock: offset = ctrl.allocator.allocate(size) self.assertIsNotNone(offset) # Try to allocate more — should fail without eviction with ctrl.lock: offset2 = ctrl._allocate_with_eviction(size) self.assertIsNone(offset2) def test_access_time_updates_prevent_eviction(self): ctrl = _make_controller(pool_mb=0.01) dim = 64 size = _embedding_bytes(1, dim) # Insert 2 entries hashes = [] for i in range(2): h = f"hash_{i}" hashes.append(h) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata[h] = (offset, 1, dim, size) ctrl._update_access_time(h) # Touch hash_0 to make it recently used time.sleep(0.01) with ctrl.lock: ctrl._update_access_time("hash_0") # Trigger eviction — hash_1 should be evicted (older) with ctrl.lock: candidates = ctrl._select_eviction_candidates(size) self.assertIn("hash_1", candidates) self.assertNotIn("hash_0", candidates) def test_eviction_stats(self): ctrl = _make_controller(pool_mb=0.01) dim = 64 size = _embedding_bytes(1, dim) # Insert and then evict with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) ctrl._update_access_time("h") with ctrl.lock: freed = ctrl._evict_hashes(["h"]) self.assertGreater(freed, 0) self.assertEqual(ctrl.stats["eviction_count"], 1) self.assertGreater(ctrl.stats["total_evicted"], 0) def test_evict_nonexistent_hash(self): ctrl = _make_controller() with ctrl.lock: freed = ctrl._evict_hashes(["nonexistent"]) self.assertEqual(freed, 0) def test_max_eviction_batch(self): ctrl = _make_controller(pool_mb=1.0, max_eviction_batch=2) dim = 64 size = _embedding_bytes(1, dim) # Insert many small entries for i in range(10): h = f"hash_{i}" with ctrl.lock: offset = ctrl.allocator.allocate(size) if offset is None: break ctrl.hash_to_metadata[h] = (offset, 1, dim, size) ctrl._update_access_time(h) # _select_eviction_candidates should return at most 2 with ctrl.lock: candidates = ctrl._select_eviction_candidates(size * 100) self.assertLessEqual(len(candidates), 2) # --------------------------------------------------------------------------- # Ref counting tests # --------------------------------------------------------------------------- class TestRefCounting(unittest.TestCase): def test_protect_prevents_eviction(self): ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) ctrl._update_access_time("h") ctrl._protect_hash("h") # Should not be selected for eviction with ctrl.lock: candidates = ctrl._select_eviction_candidates(size) self.assertNotIn("h", candidates) # Release and retry with ctrl.lock: ctrl._release_hash("h") candidates = ctrl._select_eviction_candidates(size) self.assertIn("h", candidates) def test_evict_hashes_skips_protected(self): ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) ctrl._update_access_time("h") ctrl._protect_hash("h") # Attempt to evict — should be skipped with ctrl.lock: freed = ctrl._evict_hashes(["h"]) self.assertEqual(freed, 0) with ctrl.lock: self.assertIn("h", ctrl.hash_to_metadata) def test_ref_count_multiple_protects(self): ctrl = _make_controller() with ctrl.lock: ctrl._protect_hash("h") ctrl._protect_hash("h") self.assertEqual(ctrl.ref_counts["h"], 2) ctrl._release_hash("h") self.assertEqual(ctrl.ref_counts["h"], 1) # Still protected candidates = ctrl._select_eviction_candidates(1) self.assertNotIn("h", candidates) ctrl._release_hash("h") self.assertNotIn("h", ctrl.ref_counts) def test_release_nonexistent_is_safe(self): ctrl = _make_controller() with ctrl.lock: ctrl._release_hash("nonexistent") # should not raise def test_prefetch_sets_ref_count(self): """prefetch() should set ref_count=1 for each new entry.""" ctrl = _make_controller() dim = 64 ctrl.hidden_dims = {"image": dim} h = "img_hash_1" ctrl.prefetch("req1", [h], [1], modality="image") with ctrl.lock: self.assertEqual(ctrl.ref_counts.get(h), 1) # Simulate RDMA completion with ctrl.lock: ctrl._release_hash(h) self.assertNotIn(h, ctrl.ref_counts) def test_insert_batch_sets_ref_count(self): """insert_batch() should set ref_count=1 for each new entry.""" ctrl = _make_controller() dim = 64 h = "img_hash_1" tensor = torch.randn(1, dim) ctrl.insert_batch([h], [tensor]) with ctrl.lock: self.assertEqual(ctrl.ref_counts.get(h), 1) # Simulate RDMA completion with ctrl.lock: ctrl._release_hash(h) self.assertNotIn(h, ctrl.ref_counts) def test_get_embeddings_sets_ref_count(self): """get_embeddings() should set ref_count=1 per hash.""" ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) ctrl._update_access_time("h1") tensors = ctrl.get_embeddings(["h1"]) self.assertIsNotNone(tensors[0]) with ctrl.lock: self.assertEqual(ctrl.ref_counts.get("h1"), 1) # Protected — eviction should skip with ctrl.lock: candidates = ctrl._select_eviction_candidates(size) self.assertNotIn("h1", candidates) # Release ctrl.release_embeddings(["h1"]) with ctrl.lock: self.assertNotIn("h1", ctrl.ref_counts) def test_get_embeddings_missing_hash(self): """Missing hashes should return None and not set ref_count.""" ctrl = _make_controller() tensors = ctrl.get_embeddings(["nonexistent"]) self.assertIsNone(tensors[0]) with ctrl.lock: self.assertNotIn("nonexistent", ctrl.ref_counts) def test_release_embeddings_missing_hash_is_safe(self): """Releasing a hash that was never protected should be a no-op.""" ctrl = _make_controller() ctrl.release_embeddings(["nonexistent"]) # should not raise def test_io_loop_releases_prefetch_ref(self): """_io_loop should release ref_count after batch_get completes.""" ctrl = _make_controller() dim = 64 ctrl.hidden_dims = {"image": dim} h = "img_hash_1" ctrl.prefetch("req1", [h], [1], modality="image") with ctrl.lock: self.assertEqual(ctrl.ref_counts.get(h), 1) # Simulate _io_loop completing the RDMA GET op = ctrl.ongoing_prefetch.get("req1") self.assertIsNotNone(op) ctrl.mooncake_store.batch_get = MagicMock(return_value=[True]) # Manually execute what _io_loop does for prefetch results = ctrl.mooncake_store.batch_get(op.keys, op.ptrs, op.sizes) op.mark_done(all(results)) with ctrl.lock: for k in op.keys: ctrl._release_hash(k) with ctrl.lock: self.assertNotIn(h, ctrl.ref_counts) def test_io_loop_releases_insert_ref(self): """_io_loop should release ref_count after batch_put completes.""" ctrl = _make_controller() dim = 64 h = "img_hash_1" tensor = torch.randn(1, dim) ctrl.insert_batch([h], [tensor]) with ctrl.lock: self.assertEqual(ctrl.ref_counts.get(h), 1) # Get the enqueued insert operation ctrl.insert_queue.put.assert_called_once() insert_op = ctrl.insert_queue.put.call_args[0][0] self.assertIsInstance(insert_op, EmbeddingInsertOperation) # Simulate _io_loop completing the RDMA PUT ctrl.mooncake_store.batch_put = MagicMock() ctrl.mooncake_store.batch_put(insert_op.keys, insert_op.ptrs, insert_op.sizes) with ctrl.lock: for k in insert_op.keys: ctrl._release_hash(k) with ctrl.lock: self.assertNotIn(h, ctrl.ref_counts) # --------------------------------------------------------------------------- # Race condition prevention tests # --------------------------------------------------------------------------- class TestRDMAEvictionRacePrevention(unittest.TestCase): def test_eviction_during_prefetch_is_blocked(self): """An entry with in-flight RDMA GET cannot be evicted.""" ctrl = _make_controller(pool_mb=0.01) dim = 64 ctrl.hidden_dims = {"image": dim} size = _embedding_bytes(1, dim) # Prefetch sets ref_count=1 ctrl.prefetch("req1", ["h1"], [1], modality="image") # Now try to evict to make room — should skip h1 with ctrl.lock: candidates = ctrl._select_eviction_candidates(size * 10) self.assertNotIn("h1", candidates) # Direct eviction should also be skipped with ctrl.lock: freed = ctrl._evict_hashes(["h1"]) self.assertEqual(freed, 0) def test_eviction_during_insert_is_blocked(self): """An entry with in-flight RDMA PUT cannot be evicted.""" ctrl = _make_controller(pool_mb=0.01) dim = 64 tensor = torch.randn(1, dim) ctrl.insert_batch(["h1"], [tensor]) with ctrl.lock: candidates = ctrl._select_eviction_candidates(999999) self.assertNotIn("h1", candidates) with ctrl.lock: freed = ctrl._evict_hashes(["h1"]) self.assertEqual(freed, 0) def test_eviction_during_get_embeddings_is_blocked(self): """An entry returned by get_embeddings() cannot be evicted.""" ctrl = _make_controller(pool_mb=0.01) dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) ctrl._update_access_time("h1") tensors = ctrl.get_embeddings(["h1"]) self.assertIsNotNone(tensors[0]) # Try to evict — should be blocked with ctrl.lock: candidates = ctrl._select_eviction_candidates(999999) self.assertNotIn("h1", candidates) # Release and verify eviction is now possible ctrl.release_embeddings(["h1"]) with ctrl.lock: candidates = ctrl._select_eviction_candidates(999999) self.assertIn("h1", candidates) def test_concurrent_eviction_while_reading(self): """Simulate a concurrent eviction attempt while a read holds a ref.""" ctrl = _make_controller(pool_mb=0.05) dim = 64 size = _embedding_bytes(1, dim) num_entries = 10 # Insert entries for i in range(num_entries): h = f"hash_{i}" with ctrl.lock: offset = ctrl.allocator.allocate(size) if offset is None: break ctrl.hash_to_metadata[h] = (offset, 1, dim, size) ctrl._update_access_time(h) # Simulate get_embeddings holding refs on hash_0..hash_4 held_hashes = [f"hash_{i}" for i in range(5)] tensors = ctrl.get_embeddings(held_hashes) # Try to evict all — only unprotected entries should be candidates with ctrl.lock: candidates = ctrl._select_eviction_candidates(999999) for h in held_hashes: self.assertNotIn(h, candidates) # Unprotected entries should be candidates for i in range(5, num_entries): h = f"hash_{i}" if h in ctrl.hash_to_metadata: self.assertIn(h, candidates) # Release refs ctrl.release_embeddings(held_hashes) with ctrl.lock: candidates = ctrl._select_eviction_candidates(999999) for h in held_hashes: if h in ctrl.hash_to_metadata: self.assertIn(h, candidates) def test_evict_hashes_cleans_up_ref_counts(self): """After eviction, ref_counts for the evicted hash should be removed.""" ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) ctrl._update_access_time("h") # Stale ref_count (shouldn't happen normally, but test cleanup) ctrl.ref_counts["h"] = 0 with ctrl.lock: # ref_count is 0, so eviction should proceed freed = ctrl._evict_hashes(["h"]) self.assertGreater(freed, 0) with ctrl.lock: self.assertNotIn("h", ctrl.hash_to_metadata) self.assertNotIn("h", ctrl.ref_counts) # --------------------------------------------------------------------------- # get_embeddings view safety tests # --------------------------------------------------------------------------- class TestGetEmbeddingsViewSafety(unittest.TestCase): def test_get_embeddings_returns_view(self): """get_embeddings returns a view into cpu_pool, not a copy.""" ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) ctrl._update_access_time("h1") tensors = ctrl.get_embeddings(["h1"]) self.assertIsNotNone(tensors[0]) self.assertEqual(tensors[0].shape, (1, dim)) # Verify it's a view into cpu_pool (shares storage) self.assertTrue( tensors[0].storage().data_ptr() == ctrl.cpu_pool.storage().data_ptr() ) # Release ctrl.release_embeddings(["h1"]) def test_data_preserved_while_ref_held(self): """Data should remain intact as long as ref_count > 0.""" ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) # Write known data with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) ctrl._update_access_time("h1") view = ( ctrl.cpu_pool[offset : offset + size].view(torch.float32).view(1, dim) ) view.copy_(torch.ones(1, dim)) # Read via get_embeddings (holds ref) tensors = ctrl.get_embeddings(["h1"]) self.assertTrue(torch.all(tensors[0] == 1.0)) # Verify data is still valid self.assertTrue(torch.all(tensors[0] == 1.0)) # Release ctrl.release_embeddings(["h1"]) # --------------------------------------------------------------------------- # Stats tests # --------------------------------------------------------------------------- class TestGetStats(unittest.TestCase): def test_stats_include_num_protected(self): ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h1"] = (offset, 1, dim, size) ctrl._update_access_time("h1") ctrl._protect_hash("h1") ctrl._protect_hash("h2") # h2 not in metadata, but has ref stats = ctrl.get_stats() self.assertEqual(stats["num_protected"], 2) def test_stats_eviction_tracking(self): ctrl = _make_controller(pool_mb=0.01) dim = 64 size = _embedding_bytes(1, dim) with ctrl.lock: offset = ctrl.allocator.allocate(size) ctrl.hash_to_metadata["h"] = (offset, 1, dim, size) ctrl._update_access_time("h") with ctrl.lock: ctrl._evict_hashes(["h"]) stats = ctrl.get_stats() self.assertEqual(stats["eviction_count"], 1) self.assertGreater(stats["total_evicted"], 0) self.assertEqual(stats["num_cached"], 0) # --------------------------------------------------------------------------- # _select_eviction_candidates iterator safety test # --------------------------------------------------------------------------- class TestEvictionCandidateIteratorSafety(unittest.TestCase): def test_list_snapshot_prevents_concurrent_mutation(self): """sorted_hashes should be a list snapshot, not a live dict view.""" ctrl = _make_controller() dim = 64 size = _embedding_bytes(1, dim) # Insert entries for i in range(5): h = f"hash_{i}" with ctrl.lock: offset = ctrl.allocator.allocate(size) if offset is None: break ctrl.hash_to_metadata[h] = (offset, 1, dim, size) ctrl._update_access_time(h) # _select_eviction_candidates should work even if access_order # is modified during iteration (the snapshot via list() prevents this) with ctrl.lock: candidates = ctrl._select_eviction_candidates(size) # Should return candidates without RuntimeError self.assertIsInstance(candidates, list) if __name__ == "__main__": unittest.main()