wehub-resource-sync
889 lines
31 KiB
Python
889 lines
31 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for MockL2Adapter.
|
|
|
|
Tests are written based on the L2AdapterInterface contract defined in base.py.
|
|
Tests only use public methods and do not access private fields.
|
|
"""
|
|
|
|
# Standard
|
|
import select
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.internal_api import L2AdapterListener
|
|
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import (
|
|
MockL2Adapter,
|
|
MockL2AdapterConfig,
|
|
)
|
|
from lmcache.v1.memory_management import (
|
|
MemoryFormat,
|
|
MemoryObjMetadata,
|
|
TensorMemoryObj,
|
|
)
|
|
from lmcache.v1.platform import consume_fd
|
|
|
|
_EMPTY_LAYOUT = MemoryLayoutDesc(shapes=[], dtypes=[])
|
|
|
|
|
|
class _RecordingListener(L2AdapterListener):
|
|
"""Listener that records all events for inspection in tests."""
|
|
|
|
def __init__(self):
|
|
self.stored: list[list[ObjectKey]] = []
|
|
self.accessed: list[list[ObjectKey]] = []
|
|
self.deleted: list[list[ObjectKey]] = []
|
|
|
|
def on_l2_keys_stored(self, keys: list[ObjectKey], sizes: list[int]):
|
|
self.stored.append(list(keys))
|
|
|
|
def on_l2_keys_accessed(self, keys: list[ObjectKey]):
|
|
self.accessed.append(list(keys))
|
|
|
|
def on_l2_keys_deleted(self, keys: list[ObjectKey]):
|
|
self.deleted.append(list(keys))
|
|
|
|
|
|
# =============================================================================
|
|
# Test Fixtures
|
|
# =============================================================================
|
|
|
|
|
|
def create_object_key(chunk_id: int, model_name: str = "test_model") -> ObjectKey:
|
|
"""Create a test ObjectKey with the given chunk ID."""
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name=model_name,
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def create_memory_obj(size: int = 1024, fill_value: float = 1.0) -> TensorMemoryObj:
|
|
"""Create a test TensorMemoryObj with the given size."""
|
|
raw_data = torch.empty(size, dtype=torch.float32)
|
|
raw_data.fill_(fill_value)
|
|
metadata = MemoryObjMetadata(
|
|
shape=torch.Size([size]),
|
|
dtype=torch.float32,
|
|
address=0,
|
|
phy_size=size * 4,
|
|
fmt=MemoryFormat.KV_2LTD,
|
|
ref_count=1,
|
|
)
|
|
return TensorMemoryObj(raw_data, metadata, parent_allocator=None)
|
|
|
|
|
|
def wait_for_event_fd(event_fd: int, timeout: float = 5.0) -> bool:
|
|
"""
|
|
Wait for an event fd to be signaled.
|
|
|
|
Returns True if signaled within timeout, False otherwise.
|
|
"""
|
|
poll = select.poll()
|
|
poll.register(event_fd, select.POLLIN)
|
|
events = poll.poll(timeout * 1000) # timeout in milliseconds
|
|
if events:
|
|
# Read and consume the event
|
|
try:
|
|
consume_fd(event_fd)
|
|
except BlockingIOError:
|
|
pass
|
|
return True
|
|
return False
|
|
|
|
|
|
@pytest.fixture
|
|
def adapter():
|
|
"""Create a MockL2Adapter with reasonable defaults for testing."""
|
|
config = MockL2AdapterConfig(
|
|
max_size_gb=0.001, # 1 MB
|
|
mock_bandwidth_gb=10.0, # 10 GB/s (fast for tests)
|
|
)
|
|
adapter = MockL2Adapter(config)
|
|
yield adapter
|
|
adapter.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def slow_adapter():
|
|
"""Create a MockL2Adapter with slow bandwidth for timing tests."""
|
|
config = MockL2AdapterConfig(
|
|
max_size_gb=0.001, # 1 MB
|
|
mock_bandwidth_gb=0.0001, # Very slow bandwidth
|
|
)
|
|
adapter = MockL2Adapter(config)
|
|
yield adapter
|
|
adapter.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Event Fd Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestEventFdInterface:
|
|
"""Test the event fd interface methods."""
|
|
|
|
def test_get_store_event_fd_returns_valid_fd(self, adapter):
|
|
"""get_store_event_fd should return a valid file descriptor."""
|
|
fd = adapter.get_store_event_fd()
|
|
assert isinstance(fd, int)
|
|
assert fd >= 0
|
|
|
|
def test_get_lookup_and_lock_event_fd_returns_valid_fd(self, adapter):
|
|
"""get_lookup_and_lock_event_fd should return a valid file descriptor."""
|
|
fd = adapter.get_lookup_and_lock_event_fd()
|
|
assert isinstance(fd, int)
|
|
assert fd >= 0
|
|
|
|
def test_get_load_event_fd_returns_valid_fd(self, adapter):
|
|
"""get_load_event_fd should return a valid file descriptor."""
|
|
fd = adapter.get_load_event_fd()
|
|
assert isinstance(fd, int)
|
|
assert fd >= 0
|
|
|
|
def test_event_fds_are_different(self, adapter):
|
|
"""Each operation should have a distinct event fd."""
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
assert store_fd != lookup_fd
|
|
assert store_fd != load_fd
|
|
assert lookup_fd != load_fd
|
|
|
|
|
|
# =============================================================================
|
|
# Store Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStoreInterface:
|
|
"""Test the store operation interface."""
|
|
|
|
def test_submit_store_task_returns_task_id(self, adapter):
|
|
"""submit_store_task should return a valid task ID."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
|
|
task_id = adapter.submit_store_task([key], [obj])
|
|
|
|
assert isinstance(task_id, int)
|
|
|
|
def test_submit_store_task_signals_event_fd(self, adapter):
|
|
"""Store completion should signal the store event fd."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
adapter.submit_store_task([key], [obj])
|
|
|
|
# Wait for the event fd to be signaled
|
|
assert wait_for_event_fd(store_fd, timeout=5.0), (
|
|
"Store event fd was not signaled within timeout"
|
|
)
|
|
|
|
def test_pop_completed_store_tasks_returns_completed(self, adapter):
|
|
"""pop_completed_store_tasks should return completed tasks
|
|
with success status.
|
|
"""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
task_id = adapter.submit_store_task([key], [obj])
|
|
|
|
# Wait for completion
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
|
|
completed = adapter.pop_completed_store_tasks()
|
|
|
|
assert task_id in completed
|
|
assert completed[task_id].is_successful() # Should be successful
|
|
|
|
def test_pop_completed_store_tasks_clears_completed(self, adapter):
|
|
"""pop_completed_store_tasks should clear the completed tasks."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
adapter.submit_store_task([key], [obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
|
|
# First pop should return the task
|
|
completed1 = adapter.pop_completed_store_tasks()
|
|
assert len(completed1) == 1
|
|
|
|
# Second pop should return empty
|
|
completed2 = adapter.pop_completed_store_tasks()
|
|
assert len(completed2) == 0
|
|
|
|
def test_submit_multiple_store_tasks(self, adapter):
|
|
"""Multiple store tasks should each get unique task IDs."""
|
|
keys = [create_object_key(i) for i in range(3)]
|
|
objs = [create_memory_obj(fill_value=float(i)) for i in range(3)]
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
task_ids = []
|
|
for key, obj in zip(keys, objs, strict=False):
|
|
task_id = adapter.submit_store_task([key], [obj])
|
|
task_ids.append(task_id)
|
|
|
|
# All task IDs should be unique
|
|
assert len(set(task_ids)) == 3
|
|
|
|
# Wait for all completions
|
|
completed = {}
|
|
while len(completed) < 3:
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
completed.update(adapter.pop_completed_store_tasks())
|
|
|
|
# All tasks should be completed
|
|
for task_id in task_ids:
|
|
assert task_id in completed
|
|
assert completed[task_id].is_successful()
|
|
|
|
def test_store_batch_of_objects(self, adapter):
|
|
"""A single store task can store multiple key-object pairs."""
|
|
keys = [create_object_key(i) for i in range(5)]
|
|
objs = [create_memory_obj(fill_value=float(i)) for i in range(5)]
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
task_id = adapter.submit_store_task(keys, objs)
|
|
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert task_id in completed
|
|
assert completed[task_id].is_successful()
|
|
|
|
|
|
# =============================================================================
|
|
# Lookup and Lock Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestLookupAndLockInterface:
|
|
"""Test the lookup and lock operation interface."""
|
|
|
|
def test_submit_lookup_returns_task_id(self, adapter):
|
|
"""submit_lookup_and_lock_task should return a valid task ID."""
|
|
key = create_object_key(1)
|
|
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
|
|
assert isinstance(task_id, int)
|
|
|
|
def test_lookup_signals_event_fd(self, adapter):
|
|
"""Lookup completion should signal the lookup event fd."""
|
|
key = create_object_key(1)
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
|
|
assert wait_for_event_fd(lookup_fd, timeout=5.0), (
|
|
"Lookup event fd was not signaled within timeout"
|
|
)
|
|
|
|
def test_lookup_nonexistent_key_returns_bitmap_with_zeros(self, adapter):
|
|
"""Looking up a non-existent key should return a bitmap with 0."""
|
|
key = create_object_key(999) # Never stored
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
|
|
assert bitmap is not None
|
|
assert bitmap.test(0) is False # Key not found
|
|
|
|
def test_lookup_existing_key_returns_bitmap_with_ones(self, adapter):
|
|
"""Looking up an existing key should return a bitmap with 1."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
# First store the object
|
|
adapter.submit_store_task([key], [obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Now lookup
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
|
|
assert bitmap is not None
|
|
assert bitmap.test(0) is True # Key found
|
|
|
|
def test_lookup_mixed_keys(self, adapter):
|
|
"""Lookup of mixed existing/non-existing keys returns correct bitmap."""
|
|
existing_key = create_object_key(1)
|
|
nonexistent_key = create_object_key(999)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
# Store only one key
|
|
adapter.submit_store_task([existing_key], [obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Lookup both keys
|
|
task_id = adapter.submit_lookup_and_lock_task(
|
|
[existing_key, nonexistent_key], _EMPTY_LAYOUT
|
|
)
|
|
wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
|
|
assert bitmap is not None
|
|
assert bitmap.test(0) is True # existing_key found
|
|
assert bitmap.test(1) is False # nonexistent_key not found
|
|
|
|
def test_query_lookup_result_returns_none_for_unknown_task(self, adapter):
|
|
"""Querying an unknown task ID should return None."""
|
|
result = adapter.query_lookup_and_lock_result(99999)
|
|
assert result is None
|
|
|
|
def test_query_lookup_result_clears_result(self, adapter):
|
|
"""Querying lookup result should remove it (can only query once)."""
|
|
key = create_object_key(1)
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
|
|
# First query returns result
|
|
result1 = adapter.query_lookup_and_lock_result(task_id)
|
|
assert result1 is not None
|
|
|
|
# Second query returns None
|
|
result2 = adapter.query_lookup_and_lock_result(task_id)
|
|
assert result2 is None
|
|
|
|
|
|
# =============================================================================
|
|
# Unlock Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestUnlockInterface:
|
|
"""Test the unlock operation interface."""
|
|
|
|
def test_submit_unlock_does_not_raise(self, adapter):
|
|
"""submit_unlock should not raise an exception."""
|
|
key = create_object_key(1)
|
|
|
|
# Should not raise even for non-existent key
|
|
adapter.submit_unlock([key])
|
|
|
|
def test_unlock_after_lock(self, adapter):
|
|
"""Unlocking after locking should work without error."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
# Store
|
|
adapter.submit_store_task([key], [obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Lookup and lock
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
adapter.query_lookup_and_lock_result(task_id)
|
|
|
|
# Unlock should not raise
|
|
adapter.submit_unlock([key])
|
|
|
|
|
|
# =============================================================================
|
|
# Load Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestLoadInterface:
|
|
"""Test the load operation interface."""
|
|
|
|
def test_submit_load_task_returns_task_id(self, adapter):
|
|
"""submit_load_task should return a valid task ID."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
|
|
task_id = adapter.submit_load_task([key], [obj])
|
|
|
|
assert isinstance(task_id, int)
|
|
|
|
def test_load_signals_event_fd(self, adapter):
|
|
"""Load completion should signal the load event fd."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
adapter.submit_load_task([key], [obj])
|
|
|
|
assert wait_for_event_fd(load_fd, timeout=5.0), (
|
|
"Load event fd was not signaled within timeout"
|
|
)
|
|
|
|
def test_load_nonexistent_key_returns_bitmap_with_zeros(self, adapter):
|
|
"""Loading a non-existent key should return a bitmap with 0."""
|
|
key = create_object_key(999) # Never stored
|
|
obj = create_memory_obj()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
task_id = adapter.submit_load_task([key], [obj])
|
|
wait_for_event_fd(load_fd, timeout=5.0)
|
|
|
|
bitmap = adapter.query_load_result(task_id)
|
|
|
|
assert bitmap is not None
|
|
assert bitmap.test(0) is False # Load failed
|
|
|
|
def test_load_existing_key_copies_data(self, adapter):
|
|
"""Loading an existing key should copy data to the provided buffer."""
|
|
key = create_object_key(1)
|
|
store_obj = create_memory_obj(size=100, fill_value=42.0)
|
|
load_obj = create_memory_obj(size=100, fill_value=0.0)
|
|
store_fd = adapter.get_store_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
# Store
|
|
adapter.submit_store_task([key], [store_obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Load
|
|
task_id = adapter.submit_load_task([key], [load_obj])
|
|
wait_for_event_fd(load_fd, timeout=5.0)
|
|
|
|
bitmap = adapter.query_load_result(task_id)
|
|
|
|
assert bitmap is not None
|
|
assert bitmap.test(0) is True # Load succeeded
|
|
|
|
# Data should be copied
|
|
assert torch.all(load_obj.tensor == 42.0)
|
|
|
|
def test_query_load_result_returns_none_for_unknown_task(self, adapter):
|
|
"""Querying an unknown task ID should return None."""
|
|
result = adapter.query_load_result(99999)
|
|
assert result is None
|
|
|
|
def test_query_load_result_clears_result(self, adapter):
|
|
"""Querying load result should remove it (can only query once)."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
task_id = adapter.submit_load_task([key], [obj])
|
|
wait_for_event_fd(load_fd, timeout=5.0)
|
|
|
|
# First query returns result
|
|
result1 = adapter.query_load_result(task_id)
|
|
assert result1 is not None
|
|
|
|
# Second query returns None
|
|
result2 = adapter.query_load_result(task_id)
|
|
assert result2 is None
|
|
|
|
|
|
# =============================================================================
|
|
# End-to-End Workflow Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestEndToEndWorkflow:
|
|
"""Test complete store-lookup-load workflows."""
|
|
|
|
def test_store_lookup_load_workflow(self, adapter):
|
|
"""Test the complete workflow: store -> lookup -> load."""
|
|
key = create_object_key(1)
|
|
store_obj = create_memory_obj(size=256, fill_value=123.0)
|
|
load_obj = create_memory_obj(size=256, fill_value=0.0)
|
|
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
# Step 1: Store
|
|
store_task_id = adapter.submit_store_task([key], [store_obj])
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert completed[store_task_id].is_successful()
|
|
|
|
# Step 2: Lookup and lock
|
|
lookup_task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
lookup_bitmap = adapter.query_lookup_and_lock_result(lookup_task_id)
|
|
assert lookup_bitmap.test(0) is True
|
|
|
|
# Step 3: Load
|
|
load_task_id = adapter.submit_load_task([key], [load_obj])
|
|
assert wait_for_event_fd(load_fd, timeout=5.0)
|
|
load_bitmap = adapter.query_load_result(load_task_id)
|
|
assert load_bitmap.test(0) is True
|
|
|
|
# Verify data
|
|
assert torch.all(load_obj.tensor == 123.0)
|
|
|
|
# Step 4: Unlock
|
|
adapter.submit_unlock([key])
|
|
|
|
def test_multiple_objects_workflow(self, adapter):
|
|
"""Test workflow with multiple objects."""
|
|
num_objects = 5
|
|
keys = [create_object_key(i) for i in range(num_objects)]
|
|
store_objs = [
|
|
create_memory_obj(size=64, fill_value=float(i * 10))
|
|
for i in range(num_objects)
|
|
]
|
|
load_objs = [
|
|
create_memory_obj(size=64, fill_value=0.0) for _ in range(num_objects)
|
|
]
|
|
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
# Store all
|
|
store_task_id = adapter.submit_store_task(keys, store_objs)
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert completed[store_task_id].is_successful()
|
|
|
|
# Lookup all
|
|
lookup_task_id = adapter.submit_lookup_and_lock_task(keys, _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
lookup_bitmap = adapter.query_lookup_and_lock_result(lookup_task_id)
|
|
for i in range(num_objects):
|
|
assert lookup_bitmap.test(i) is True
|
|
|
|
# Load all
|
|
load_task_id = adapter.submit_load_task(keys, load_objs)
|
|
assert wait_for_event_fd(load_fd, timeout=5.0)
|
|
load_bitmap = adapter.query_load_result(load_task_id)
|
|
for i in range(num_objects):
|
|
assert load_bitmap.test(i) is True
|
|
assert torch.all(load_objs[i].tensor == float(i * 10))
|
|
|
|
|
|
# =============================================================================
|
|
# Close Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestCloseInterface:
|
|
"""Test the close operation."""
|
|
|
|
def test_close_does_not_raise(self):
|
|
"""close() should not raise an exception."""
|
|
config = MockL2AdapterConfig(max_size_gb=0.001, mock_bandwidth_gb=10.0)
|
|
adapter = MockL2Adapter(config)
|
|
|
|
# Should not raise
|
|
adapter.close()
|
|
|
|
def test_close_after_operations(self):
|
|
"""close() should work after store/lookup/load operations."""
|
|
config = MockL2AdapterConfig(max_size_gb=0.001, mock_bandwidth_gb=10.0)
|
|
adapter = MockL2Adapter(config)
|
|
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
adapter.submit_store_task([key], [obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Should not raise
|
|
adapter.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Bandwidth Simulation Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestBandwidthSimulation:
|
|
"""Test that bandwidth simulation affects completion timing."""
|
|
|
|
def test_store_completion_is_delayed_by_bandwidth(self, slow_adapter):
|
|
"""Store completion should be delayed based on object size and bandwidth."""
|
|
key = create_object_key(1)
|
|
# Create a larger object to see noticeable delay
|
|
obj = create_memory_obj(size=10000) # ~40KB
|
|
store_fd = slow_adapter.get_store_event_fd()
|
|
|
|
start_time = time.time()
|
|
slow_adapter.submit_store_task([key], [obj])
|
|
|
|
# Should take some time due to bandwidth simulation
|
|
assert wait_for_event_fd(store_fd, timeout=60.0)
|
|
elapsed = time.time() - start_time
|
|
|
|
# With very slow bandwidth, should take noticeable time
|
|
# (exact time depends on bandwidth config)
|
|
assert elapsed > 0.1 # At least some delay
|
|
|
|
def test_fast_bandwidth_completes_quickly(self, adapter):
|
|
"""With fast bandwidth, completion should be quick."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj(size=1000)
|
|
store_fd = adapter.get_store_event_fd()
|
|
|
|
start_time = time.time()
|
|
adapter.submit_store_task([key], [obj])
|
|
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
elapsed = time.time() - start_time
|
|
|
|
# With fast bandwidth, should complete quickly
|
|
assert elapsed < 1.0
|
|
|
|
|
|
# =============================================================================
|
|
# Eviction Interface Tests
|
|
# =============================================================================
|
|
|
|
|
|
def _store_and_wait(adapter, key, obj):
|
|
"""Helper: store one key and wait for the store event fd to fire."""
|
|
store_fd = adapter.get_store_event_fd()
|
|
adapter.submit_store_task([key], [obj])
|
|
assert wait_for_event_fd(store_fd, timeout=5.0), "store timed out"
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
|
|
class TestEvictionInterface:
|
|
"""Tests for delete(), get_usage(), and listener notifications."""
|
|
|
|
def test_delete_removes_key(self, adapter):
|
|
"""delete() should make the key invisible to subsequent lookups."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
_store_and_wait(adapter, key, obj)
|
|
|
|
adapter.delete([key])
|
|
|
|
task_id = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
assert bitmap.test(0) is False
|
|
|
|
def test_delete_nonexistent_key_does_not_raise(self, adapter):
|
|
"""delete() on a key that was never stored should not raise."""
|
|
adapter.delete([create_object_key(999)])
|
|
|
|
def test_delete_multiple_keys(self, adapter):
|
|
"""delete() on a batch of keys removes all of them."""
|
|
keys = [create_object_key(i) for i in range(3)]
|
|
objs = [create_memory_obj() for _ in range(3)]
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
|
|
adapter.submit_store_task(keys, objs)
|
|
assert wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
adapter.delete(keys)
|
|
|
|
task_id = adapter.submit_lookup_and_lock_task(keys, _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(lookup_fd, timeout=5.0)
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
for i in range(len(keys)):
|
|
assert bitmap.test(i) is False
|
|
|
|
def test_get_usage_empty_adapter_is_zero(self, adapter):
|
|
"""get_usage() on a fresh adapter should report 0 bytes."""
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used == 0
|
|
assert usage.usage_fraction == 0.0
|
|
assert usage.bytes_by_cache_salt == {}
|
|
|
|
def test_get_usage_increases_after_store(self, adapter):
|
|
"""get_usage() should report positive bytes after a store."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj(size=1024)
|
|
_store_and_wait(adapter, key, obj)
|
|
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used > 0
|
|
assert 0.0 < usage.usage_fraction <= 1.0
|
|
|
|
def test_get_usage_decreases_after_delete(self, adapter):
|
|
"""get_usage() should drop back to 0 after deleting the only stored key."""
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj(size=1024)
|
|
_store_and_wait(adapter, key, obj)
|
|
|
|
assert adapter.get_usage().total_bytes_used > 0
|
|
|
|
adapter.delete([key])
|
|
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used == 0
|
|
assert usage.usage_fraction == 0.0
|
|
|
|
def test_bytes_by_cache_salt_populated_from_cache_salt(self, adapter):
|
|
"""End-to-end: storing keys with different ``cache_salt`` values
|
|
should drive the byte buckets in ``AdapterUsage`` — proves the
|
|
salt actually flows through the real adapter into the base-class
|
|
accounting (not just verified by the stub tests)."""
|
|
# Two keys per cache_salt so we know the totals are summed, not
|
|
# just per-key-overwritten.
|
|
alice_keys = [
|
|
ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(i),
|
|
model_name="m",
|
|
kv_rank=0,
|
|
cache_salt="alice",
|
|
)
|
|
for i in (1, 2)
|
|
]
|
|
bob_key = ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(3),
|
|
model_name="m",
|
|
kv_rank=0,
|
|
cache_salt="bob",
|
|
)
|
|
obj = create_memory_obj(size=128) # 128 floats * 4 bytes = 512 B
|
|
for k in alice_keys + [bob_key]:
|
|
_store_and_wait(adapter, k, obj)
|
|
|
|
usage = adapter.get_usage()
|
|
assert usage.bytes_by_cache_salt == {"alice": 1024, "bob": 512}
|
|
assert usage.total_bytes_used == 1536
|
|
|
|
# Deleting one of alice's keys should shrink alice's bucket but
|
|
# leave bob's untouched.
|
|
adapter.delete([alice_keys[0]])
|
|
usage = adapter.get_usage()
|
|
assert usage.bytes_by_cache_salt == {"alice": 512, "bob": 512}
|
|
assert usage.total_bytes_used == 1024
|
|
|
|
# Deleting alice's last key should drop the bucket entirely so
|
|
# the snapshot stays compact.
|
|
adapter.delete([alice_keys[1]])
|
|
usage = adapter.get_usage()
|
|
assert "alice" not in usage.bytes_by_cache_salt
|
|
assert usage.bytes_by_cache_salt == {"bob": 512}
|
|
|
|
def test_listener_notified_on_store(self, adapter):
|
|
"""Listener.on_l2_keys_stored should be called after a store completes."""
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
_store_and_wait(adapter, key, obj)
|
|
|
|
assert len(listener.stored) == 1
|
|
assert key in listener.stored[0]
|
|
assert listener.deleted == []
|
|
|
|
def test_listener_notified_on_delete(self, adapter):
|
|
"""Listener.on_l2_keys_deleted should be called after delete()."""
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
_store_and_wait(adapter, key, obj)
|
|
adapter.delete([key])
|
|
|
|
assert len(listener.deleted) == 1
|
|
assert key in listener.deleted[0]
|
|
|
|
def test_listener_delete_skips_missing_keys(self, adapter):
|
|
"""on_l2_keys_deleted should only include keys that were actually removed."""
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
real_key = create_object_key(1)
|
|
missing_key = create_object_key(999)
|
|
obj = create_memory_obj()
|
|
_store_and_wait(adapter, real_key, obj)
|
|
|
|
adapter.delete([real_key, missing_key])
|
|
|
|
assert len(listener.deleted) == 1
|
|
notified = listener.deleted[0]
|
|
assert real_key in notified
|
|
assert missing_key not in notified
|
|
|
|
def test_listener_notified_on_load(self, adapter):
|
|
"""Listener.on_l2_keys_accessed should be called after a load completes."""
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
key = create_object_key(1)
|
|
store_obj = create_memory_obj(size=100, fill_value=42.0)
|
|
load_obj = create_memory_obj(size=100, fill_value=0.0)
|
|
store_fd = adapter.get_store_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
# Store
|
|
adapter.submit_store_task([key], [store_obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Load
|
|
adapter.submit_load_task([key], [load_obj])
|
|
wait_for_event_fd(load_fd, timeout=5.0)
|
|
|
|
assert len(listener.accessed) == 1
|
|
assert key in listener.accessed[0]
|
|
|
|
def test_listener_load_skips_missing_keys(self, adapter):
|
|
"""on_l2_keys_accessed should only include keys that were actually loaded."""
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
real_key = create_object_key(1)
|
|
missing_key = create_object_key(999)
|
|
store_obj = create_memory_obj(size=100, fill_value=42.0)
|
|
load_obj1 = create_memory_obj(size=100, fill_value=0.0)
|
|
load_obj2 = create_memory_obj(size=100, fill_value=0.0)
|
|
store_fd = adapter.get_store_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
|
|
# Store only real_key
|
|
adapter.submit_store_task([real_key], [store_obj])
|
|
wait_for_event_fd(store_fd, timeout=5.0)
|
|
adapter.pop_completed_store_tasks()
|
|
|
|
# Load both keys
|
|
adapter.submit_load_task([real_key, missing_key], [load_obj1, load_obj2])
|
|
wait_for_event_fd(load_fd, timeout=5.0)
|
|
|
|
assert len(listener.accessed) == 1
|
|
assert real_key in listener.accessed[0]
|
|
assert missing_key not in listener.accessed[0]
|
|
|
|
def test_multiple_listeners_all_notified(self, adapter):
|
|
"""All registered listeners should receive the same store event."""
|
|
l1, l2 = _RecordingListener(), _RecordingListener()
|
|
adapter.register_listener(l1)
|
|
adapter.register_listener(l2)
|
|
|
|
key = create_object_key(1)
|
|
obj = create_memory_obj()
|
|
_store_and_wait(adapter, key, obj)
|
|
|
|
assert len(l1.stored) == 1
|
|
assert len(l2.stored) == 1
|