434 lines
14 KiB
Python
434 lines
14 KiB
Python
"""Unit tests for EmbeddingCacheController paged host pool behavior."""
|
|
|
|
import threading
|
|
import unittest
|
|
from queue import Queue
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
from sglang.srt.managers.schedule_batch import Modality
|
|
from sglang.srt.mem_cache.storage.mooncake_store.embedding_cache_controller import (
|
|
EmbeddingCacheController,
|
|
EmbeddingCacheEntry,
|
|
EmbeddingPool,
|
|
EntryState,
|
|
EvictableLRU,
|
|
PageRun,
|
|
RangePageAllocator,
|
|
build_transfer_buffers,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
|
|
|
|
|
def _make_pool(num_pages=16, dim=4, page_size=2, modality="vision"):
|
|
total_tokens = num_pages * page_size
|
|
tensor = torch.empty((total_tokens, dim), dtype=torch.float32)
|
|
return EmbeddingPool(
|
|
modality=modality,
|
|
dim=dim,
|
|
dtype=torch.float32,
|
|
page_size=page_size,
|
|
tensor=tensor,
|
|
num_pages=num_pages,
|
|
allocator=RangePageAllocator(num_pages),
|
|
page_bytes=page_size * dim * torch.float32.itemsize,
|
|
pool_size_bytes=total_tokens * dim * torch.float32.itemsize,
|
|
pin_memory=False,
|
|
)
|
|
|
|
|
|
def _make_controller(num_pages=16, dim=4, page_size=2, enable_eviction=True):
|
|
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 = {
|
|
Modality.IMAGE: dim,
|
|
Modality.VIDEO: dim,
|
|
Modality.AUDIO: dim,
|
|
}
|
|
ctrl.dtype = torch.float32
|
|
ctrl.element_size = torch.float32.itemsize
|
|
ctrl.enable_eviction = enable_eviction
|
|
ctrl.max_eviction_batch = 10
|
|
ctrl.mooncake_store = MagicMock()
|
|
ctrl.total_pool_size_bytes = num_pages * page_size * dim * torch.float32.itemsize
|
|
ctrl.vision_pool = _make_pool(num_pages, dim, page_size)
|
|
ctrl.audio_pool = _make_pool(num_pages, dim, page_size, modality="audio")
|
|
ctrl.pools = {"vision": ctrl.vision_pool, "audio": ctrl.audio_pool}
|
|
ctrl.entries = {}
|
|
ctrl.vision_pool.evictable = EvictableLRU()
|
|
ctrl.audio_pool.evictable = EvictableLRU()
|
|
ctrl.stats = {
|
|
"total_allocated": 0,
|
|
"total_evicted": 0,
|
|
"eviction_count": 0,
|
|
"allocation_failures": 0,
|
|
}
|
|
ctrl.ongoing_prefetch = {}
|
|
ctrl.prefetch_queue = Queue()
|
|
ctrl.insert_queue = Queue()
|
|
ctrl.lock = threading.Lock()
|
|
ctrl.stop_event = threading.Event()
|
|
ctrl.io_thread = MagicMock()
|
|
ctrl.prefetch_tp_group = None
|
|
ctrl._copy_streams = {}
|
|
return ctrl
|
|
|
|
|
|
class TestRangePageAllocator(unittest.TestCase):
|
|
def test_prefers_single_contiguous_run(self):
|
|
allocator = RangePageAllocator(num_pages=8)
|
|
|
|
runs = allocator.allocate(num_tokens=6, page_size=2)
|
|
|
|
self.assertEqual(runs, [PageRun(start=0, length=3)])
|
|
self.assertEqual(allocator.free_ranges, [(3, 5)])
|
|
|
|
def test_free_merges_adjacent_ranges(self):
|
|
allocator = RangePageAllocator(num_pages=8)
|
|
first = allocator.allocate(num_tokens=4, page_size=2)
|
|
second = allocator.allocate(num_tokens=4, page_size=2)
|
|
|
|
allocator.free(first)
|
|
allocator.free(second)
|
|
|
|
self.assertEqual(allocator.free_ranges, [(0, 8)])
|
|
|
|
def test_scatter_fallback_returns_physical_order(self):
|
|
allocator = RangePageAllocator(num_pages=4)
|
|
a = allocator.allocate(num_tokens=2, page_size=2)
|
|
b = allocator.allocate(num_tokens=2, page_size=2)
|
|
c = allocator.allocate(num_tokens=2, page_size=2)
|
|
d = allocator.allocate(num_tokens=2, page_size=2)
|
|
allocator.free(a)
|
|
allocator.free(c)
|
|
|
|
runs = allocator.allocate(num_tokens=4, page_size=2)
|
|
|
|
self.assertEqual(runs, [PageRun(start=0, length=1), PageRun(start=2, length=1)])
|
|
self.assertEqual([run.start for run in runs], sorted(run.start for run in runs))
|
|
self.assertEqual(b, [PageRun(start=1, length=1)])
|
|
self.assertEqual(d, [PageRun(start=3, length=1)])
|
|
|
|
def test_allocate_fails_when_total_free_pages_are_insufficient(self):
|
|
allocator = RangePageAllocator(num_pages=2)
|
|
allocator.allocate(num_tokens=4, page_size=2)
|
|
|
|
self.assertIsNone(allocator.allocate(num_tokens=2, page_size=2))
|
|
self.assertEqual(allocator.free_pages, 0)
|
|
|
|
|
|
class TestEntryStateAndPins(unittest.TestCase):
|
|
def test_ready_entry_with_no_pins_is_evictable(self):
|
|
entry = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=2,
|
|
dim=4,
|
|
page_runs=[PageRun(0, 1)],
|
|
state=EntryState.READY,
|
|
)
|
|
|
|
self.assertTrue(entry.is_evictable())
|
|
|
|
def test_ready_entry_with_pin_is_not_evictable(self):
|
|
entry = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=2,
|
|
dim=4,
|
|
page_runs=[PageRun(0, 1)],
|
|
state=EntryState.READY,
|
|
)
|
|
|
|
entry.pin()
|
|
self.assertFalse(entry.is_evictable())
|
|
entry.unpin()
|
|
self.assertTrue(entry.is_evictable())
|
|
|
|
def test_multiple_pins_require_all_unpins(self):
|
|
entry = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=2,
|
|
dim=4,
|
|
page_runs=[PageRun(0, 1)],
|
|
state=EntryState.READY,
|
|
)
|
|
|
|
entry.pin()
|
|
entry.pin()
|
|
|
|
self.assertEqual(entry.ref_count, 2)
|
|
self.assertFalse(entry.is_evictable())
|
|
|
|
|
|
class TestEvictableLruInvariant(unittest.TestCase):
|
|
def _insert_entry(
|
|
self,
|
|
ctrl,
|
|
mm_hash,
|
|
modality=Modality.IMAGE,
|
|
state=EntryState.READY,
|
|
):
|
|
pool = ctrl._get_pool(modality)
|
|
page_runs = pool.allocator.allocate(2, pool.page_size)
|
|
entry = EmbeddingCacheEntry(
|
|
hash=mm_hash,
|
|
modality=modality,
|
|
num_tokens=2,
|
|
dim=pool.dim,
|
|
page_runs=page_runs,
|
|
state=state,
|
|
)
|
|
ctrl.entries[mm_hash] = entry
|
|
return entry
|
|
|
|
def test_filling_entry_is_not_in_evictable_lru_until_ready(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
entry = self._insert_entry(
|
|
ctrl,
|
|
"h",
|
|
state=EntryState.FILLING,
|
|
)
|
|
|
|
self.assertNotIn("h", ctrl.vision_pool.evictable)
|
|
|
|
with ctrl.lock:
|
|
ctrl._mark_ready(entry)
|
|
|
|
self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"])
|
|
|
|
def test_first_read_pin_removes_candidate_and_last_release_reinserts(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
entry = self._insert_entry(ctrl, "h")
|
|
with ctrl.lock:
|
|
ctrl._lru_touch("h")
|
|
|
|
with ctrl.lock:
|
|
ctrl._pin_read(entry)
|
|
ctrl._pin_read(entry)
|
|
|
|
self.assertNotIn("h", ctrl.vision_pool.evictable)
|
|
|
|
with ctrl.lock:
|
|
ctrl._unpin_read(entry)
|
|
self.assertNotIn("h", ctrl.vision_pool.evictable)
|
|
|
|
with ctrl.lock:
|
|
ctrl._unpin_read(entry)
|
|
self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"])
|
|
|
|
def test_evict_for_pool_pops_only_that_pool_candidates(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
self._insert_entry(ctrl, "vision_h", modality=Modality.IMAGE)
|
|
self._insert_entry(ctrl, "audio_h", modality=Modality.AUDIO)
|
|
with ctrl.lock:
|
|
ctrl._lru_touch("vision_h")
|
|
ctrl._lru_touch("audio_h")
|
|
|
|
required_pages = ctrl.vision_pool.allocator.free_pages + 1
|
|
with ctrl.lock:
|
|
ctrl._evict_for_pool(ctrl.vision_pool, required_pages)
|
|
|
|
self.assertNotIn("vision_h", ctrl.entries)
|
|
self.assertIn("audio_h", ctrl.entries)
|
|
self.assertEqual(list(ctrl.audio_pool.evictable.keys()), ["audio_h"])
|
|
|
|
|
|
class TestStoreToPool(unittest.TestCase):
|
|
def test_store_to_pool_async_raises_on_cpu_tensor(self):
|
|
ctrl = _make_controller(num_pages=8, dim=4, page_size=2)
|
|
tensor = torch.empty((2, 4), dtype=torch.float32)
|
|
|
|
with self.assertRaises(ValueError):
|
|
ctrl.store_to_pool_async(["h"], [tensor], Modality.IMAGE)
|
|
|
|
|
|
def _insert_ready_entry(ctrl, mm_hash, tensor, modality=Modality.IMAGE):
|
|
"""Manually write tensor into pool pages and create a READY entry."""
|
|
pool = ctrl._get_pool(modality)
|
|
if tensor.ndim != 2:
|
|
tensor = tensor.reshape(-1, tensor.shape[-1])
|
|
num_tokens = int(tensor.shape[0])
|
|
page_runs = pool.allocator.allocate(num_tokens, pool.page_size)
|
|
copied = 0
|
|
for run in page_runs:
|
|
valid = min(pool.page_size * run.length, num_tokens - copied)
|
|
start = run.start * pool.page_size
|
|
pool.tensor[start : start + valid].copy_(tensor[copied : copied + valid])
|
|
copied += valid
|
|
entry = EmbeddingCacheEntry(
|
|
hash=mm_hash,
|
|
modality=modality,
|
|
num_tokens=num_tokens,
|
|
dim=int(tensor.shape[1]),
|
|
page_runs=page_runs,
|
|
state=EntryState.READY,
|
|
)
|
|
ctrl.entries[mm_hash] = entry
|
|
pool.evictable.touch(mm_hash)
|
|
return entry
|
|
|
|
|
|
class TestMooncakeLifecycle(unittest.TestCase):
|
|
def test_prefetch_creates_filling_entry_and_get_success_marks_ready(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
|
|
ctrl.prefetch("req", ["h"], [2], Modality.IMAGE)
|
|
op = ctrl.ongoing_prefetch["req"]
|
|
ctrl._finish_get(op, [True])
|
|
|
|
entry = ctrl.entries["h"]
|
|
self.assertEqual(entry.state, EntryState.READY)
|
|
|
|
def test_prefetch_get_failure_frees_entry(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
|
|
ctrl.prefetch("req", ["h"], [2], Modality.IMAGE)
|
|
op = ctrl.ongoing_prefetch["req"]
|
|
ctrl._finish_get(op, [False])
|
|
|
|
self.assertNotIn("h", ctrl.entries)
|
|
self.assertEqual(ctrl.vision_pool.allocator.free_pages, 4)
|
|
|
|
def test_insert_batch_pins_and_releases_on_put(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4)
|
|
_insert_ready_entry(ctrl, "h", tensor)
|
|
|
|
ctrl.insert_batch(["h"], Modality.IMAGE)
|
|
op = ctrl.insert_queue.get_nowait()
|
|
|
|
entry = ctrl.entries["h"]
|
|
self.assertEqual(entry.ref_count, 1)
|
|
self.assertNotIn("h", ctrl.vision_pool.evictable)
|
|
|
|
ctrl._finish_put(op, [True])
|
|
|
|
self.assertEqual(entry.ref_count, 0)
|
|
self.assertEqual(entry.state, EntryState.READY)
|
|
self.assertIn("h", ctrl.vision_pool.evictable)
|
|
|
|
|
|
class TestGetPoolViews(unittest.TestCase):
|
|
def test_get_pool_views_returns_none_for_filling_entry(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
ctrl.entries["h"] = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=2,
|
|
dim=4,
|
|
page_runs=[PageRun(0, 1)],
|
|
state=EntryState.FILLING,
|
|
)
|
|
|
|
views = ctrl.get_pool_views(["h"])
|
|
self.assertIsNone(views[0])
|
|
|
|
def test_get_pool_views_returns_slices_and_release_unpins(self):
|
|
ctrl = _make_controller(num_pages=4, dim=4, page_size=2)
|
|
tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4)
|
|
_insert_ready_entry(ctrl, "h", tensor)
|
|
|
|
views = ctrl.get_pool_views(["h"])
|
|
self.assertIsNotNone(views[0])
|
|
entry = ctrl.entries["h"]
|
|
self.assertEqual(entry.ref_count, 1)
|
|
|
|
flat = torch.cat(views[0], dim=0)
|
|
self.assertTrue(torch.equal(flat, tensor))
|
|
|
|
ctrl.release_pool_views(["h"])
|
|
self.assertEqual(entry.ref_count, 0)
|
|
self.assertIn("h", ctrl.vision_pool.evictable)
|
|
|
|
|
|
class TestTransferBuffers(unittest.TestCase):
|
|
def test_build_transfer_buffers_for_single_run(self):
|
|
pool = _make_pool(num_pages=8, dim=4, page_size=2)
|
|
entry = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=5,
|
|
dim=4,
|
|
page_runs=[PageRun(2, 3)],
|
|
state=EntryState.READY,
|
|
)
|
|
|
|
ptrs, sizes = build_transfer_buffers(entry, pool)
|
|
|
|
self.assertEqual(ptrs, [pool.tensor[4].data_ptr()])
|
|
self.assertEqual(sizes, [5 * 4 * torch.float32.itemsize])
|
|
|
|
def test_build_transfer_buffers_for_multiple_runs(self):
|
|
pool = _make_pool(num_pages=8, dim=4, page_size=2)
|
|
entry = EmbeddingCacheEntry(
|
|
hash="h",
|
|
modality=Modality.IMAGE,
|
|
num_tokens=5,
|
|
dim=4,
|
|
page_runs=[PageRun(0, 1), PageRun(3, 2)],
|
|
state=EntryState.READY,
|
|
)
|
|
|
|
ptrs, sizes = build_transfer_buffers(entry, pool)
|
|
|
|
self.assertEqual(ptrs, [pool.tensor[0].data_ptr(), pool.tensor[6].data_ptr()])
|
|
self.assertEqual(
|
|
sizes,
|
|
[
|
|
2 * 4 * torch.float32.itemsize,
|
|
3 * 4 * torch.float32.itemsize,
|
|
],
|
|
)
|
|
|
|
|
|
class TestMooncakeEmbeddingStoreWrappers(unittest.TestCase):
|
|
def test_batch_put_multi_buffers_deduplicates_existing_keys(self):
|
|
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import (
|
|
MooncakeEmbeddingStore,
|
|
)
|
|
|
|
store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore)
|
|
store.store = MagicMock()
|
|
store.store.batch_is_exist.return_value = [1, 0]
|
|
store.store.batch_put_from_multi_buffers.return_value = [0]
|
|
|
|
results = store.batch_put_from_multi_buffers(
|
|
["a", "b"],
|
|
[[11], [22]],
|
|
[[4], [4]],
|
|
)
|
|
|
|
self.assertEqual(results, [True, True])
|
|
store.store.batch_put_from_multi_buffers.assert_called_once_with(
|
|
["emb_b"], [[22]], [[4]]
|
|
)
|
|
|
|
def test_batch_get_multi_buffers_maps_positive_result_to_true(self):
|
|
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import (
|
|
MooncakeEmbeddingStore,
|
|
)
|
|
|
|
store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore)
|
|
store.store = MagicMock()
|
|
store.store.batch_get_into_multi_buffers.return_value = [8, -1]
|
|
|
|
results = store.batch_get_into_multi_buffers(
|
|
["a", "b"],
|
|
[[11], [22]],
|
|
[[4], [4]],
|
|
)
|
|
|
|
self.assertEqual(results, [True, False])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|