Files
lmcache--lmcache/tests/v1/distributed/test_native_connector_l2_adapter.py
2026-07-13 12:24:33 +08:00

1361 lines
44 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""
Unit tests for NativeConnectorL2Adapter.
Uses a mock native connector (pure Python) that simulates the pybind-wrapped
C++ IStorageConnector interface, so no Redis or C++ build is needed.
"""
# Standard
import ctypes
import select
import threading
# Third Party
import pytest
import torch
# First Party
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
from lmcache.v1.distributed.l2_adapters.native_connector_l2_adapter import (
NativeConnectorL2Adapter,
_object_key_to_string,
)
from lmcache.v1.memory_management import (
MemoryFormat,
MemoryObjMetadata,
TensorMemoryObj,
)
from lmcache.v1.platform import consume_fd, create_event_notifier
_EMPTY_LAYOUT = MemoryLayoutDesc(shapes=[], dtypes=[])
# =============================================================================
# Mock Native Connector (simulates the pybind C++ IStorageConnector interface)
# =============================================================================
class MockNativeConnector:
"""
Pure-Python mock that implements the same interface as a pybind-wrapped
C++ IStorageConnector. Stores data in-memory dicts.
Methods:
- event_fd() -> int
- submit_batch_get(keys, memoryviews) -> int
- submit_batch_set(keys, memoryviews) -> int
- submit_batch_exists(keys) -> int
- drain_completions() -> list[tuple[int, bool, str, list[bool] | None]]
- close()
"""
def __init__(self):
self._efd = create_event_notifier()
self._store: dict[str, bytes] = {}
self._next_id = 1
self._completions: list[tuple[int, bool, str, list[bool] | None]] = []
self._lock = threading.Lock()
self._closed = False
def event_fd(self) -> int:
return self._efd.fileno()
def submit_batch_set(self, keys: list[str], memoryviews: list) -> int:
with self._lock:
fid = self._next_id
self._next_id += 1
try:
for key, mv in zip(keys, memoryviews, strict=False):
self._store[key] = bytes(mv)
self._push_completion(fid, True, "", None)
except Exception as e:
self._push_completion(fid, False, str(e), None)
return fid
def submit_batch_get(self, keys: list[str], memoryviews: list) -> int:
with self._lock:
fid = self._next_id
self._next_id += 1
try:
all_ok = True
for key, mv in zip(keys, memoryviews, strict=False):
data = self._store.get(key)
if data is None:
all_ok = False
break
if len(data) != mv.nbytes:
all_ok = False
break
# Copy data into the buffer using ctypes (same as C++ void* write)
dest_ptr = ctypes.c_char_p(
ctypes.addressof(ctypes.c_char.from_buffer(mv))
)
ctypes.memmove(dest_ptr, data, len(data))
self._push_completion(fid, all_ok, "", None)
except Exception as e:
self._push_completion(fid, False, str(e), None)
return fid
def submit_batch_exists(self, keys: list[str]) -> int:
with self._lock:
fid = self._next_id
self._next_id += 1
results = [key in self._store for key in keys]
self._push_completion(fid, True, "", results)
return fid
def submit_batch_delete(self, keys: list[str]) -> int:
with self._lock:
fid = self._next_id
self._next_id += 1
results = []
for key in keys:
if key in self._store:
del self._store[key]
results.append(True)
else:
results.append(False)
self._push_completion(fid, True, "", results)
return fid
def drain_completions(self) -> list[tuple[int, bool, str, list[bool] | None]]:
# Drain the eventfd
try:
self._efd.consume()
except BlockingIOError:
pass
with self._lock:
completions = list(self._completions)
self._completions.clear()
return completions
def close(self):
if not self._closed:
self._closed = True
self._efd.close()
def _push_completion(
self, fid: int, ok: bool, error: str, result_bools: list[bool] | None
):
with self._lock:
self._completions.append((fid, ok, error, result_bools))
# Signal the eventfd
try:
self._efd.notify()
except OSError:
pass
# =============================================================================
# Test Fixtures
# =============================================================================
def create_object_key(chunk_id: int, model_name: str = "test_model") -> ObjectKey:
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:
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:
poll = select.poll()
poll.register(event_fd, select.POLLIN)
events = poll.poll(timeout * 1000)
if events:
try:
consume_fd(event_fd)
except BlockingIOError:
pass
return True
return False
@pytest.fixture
def adapter():
mock_client = MockNativeConnector()
adp = NativeConnectorL2Adapter(mock_client)
yield adp
adp.close()
# =============================================================================
# ObjectKey Serialization Tests
# =============================================================================
class TestObjectKeySerialization:
def test_serialization_is_deterministic(self):
key = create_object_key(42, "my_model")
s1 = _object_key_to_string(key)
s2 = _object_key_to_string(key)
assert s1 == s2
def test_different_keys_produce_different_strings(self):
k1 = create_object_key(1)
k2 = create_object_key(2)
assert _object_key_to_string(k1) != _object_key_to_string(k2)
def test_serialization_format(self):
"""Unsalted keys use the 4-field shape with object_group_id
embedded right after kv_rank."""
key = ObjectKey(
chunk_hash=b"\x00\x01\x02\x03",
model_name="llama",
kv_rank=255,
)
s = _object_key_to_string(key)
assert s == "llama@000000ff@0@00010203"
def test_object_group_id_embedded(self):
key = ObjectKey(
chunk_hash=b"\x00\x01\x02\x03",
model_name="llama",
kv_rank=255,
object_group_id=5,
)
s = _object_key_to_string(key)
assert s == "llama@000000ff@5@00010203"
def test_salted_serialization_format(self):
"""Salted keys append ``@<cache_salt>`` as a 5th field."""
key = ObjectKey(
chunk_hash=b"\x00\x01\x02\x03",
model_name="llama",
kv_rank=255,
cache_salt="alice",
)
s = _object_key_to_string(key)
assert s == "llama@000000ff@0@00010203@alice"
def test_different_salts_produce_different_strings(self):
base = {
"chunk_hash": b"\x00\x01\x02\x03",
"model_name": "llama",
"kv_rank": 0,
}
k_empty = ObjectKey(**base)
k_alice = ObjectKey(**base, cache_salt="alice")
k_bob = ObjectKey(**base, cache_salt="bob")
s_empty = _object_key_to_string(k_empty)
s_alice = _object_key_to_string(k_alice)
s_bob = _object_key_to_string(k_bob)
assert s_empty != s_alice
assert s_alice != s_bob
# Empty salt has no trailing "@salt", salted keys do.
assert s_empty.count("@") == 3 # 4 fields (model, kv_rank, group, hash)
assert s_alice.endswith("@alice")
assert s_bob.endswith("@bob")
class TestObjectKeyModelNameValidation:
"""model_name must not contain ``@`` — the L2 adapters split keys
and filenames on ``@`` and rely on this invariant."""
def test_reject_at_in_model_name(self):
with pytest.raises(ValueError, match="model_name"):
ObjectKey(
chunk_hash=b"\x00",
model_name="ns@model",
kv_rank=0,
)
def test_slash_in_model_name_is_accepted(self):
# '/' is sanitized to '-SEP-' by the FS adapter; the invariant
# is only about '@'.
key = ObjectKey(
chunk_hash=b"\x00",
model_name="meta-llama/Llama-3",
kv_rank=0,
)
assert key.model_name == "meta-llama/Llama-3"
class TestObjectKeyCacheSaltValidation:
"""cache_salt must not contain ``@``, ``/``, ``\\``, or NUL, and must
be <= 128 chars. The invariant is enforced at construction time so
all downstream serializers (Python + C++) can rely on it."""
def test_reject_at_in_salt(self):
with pytest.raises(ValueError, match="cache_salt"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="alice@bob",
)
def test_reject_leading_at_in_salt(self):
with pytest.raises(ValueError, match="cache_salt"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="@user",
)
def test_reject_slash_in_salt(self):
with pytest.raises(ValueError, match="cache_salt"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="tenant/alice",
)
def test_reject_backslash_in_salt(self):
with pytest.raises(ValueError, match="cache_salt"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="tenant\\alice",
)
def test_reject_nul_in_salt(self):
with pytest.raises(ValueError, match="cache_salt"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="bad\x00salt",
)
def test_reject_too_long_salt(self):
with pytest.raises(ValueError, match="max length"):
ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="x" * 129,
)
def test_max_length_salt_accepted(self):
key = ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="x" * 128,
)
assert len(key.cache_salt) == 128
def test_empty_salt_is_accepted(self):
# Default (unsalted) path.
key = ObjectKey(chunk_hash=b"\x00", model_name="m", kv_rank=0)
assert key.cache_salt == ""
def test_non_salt_chars_are_accepted(self):
# Common identifier chars are fine.
key = ObjectKey(
chunk_hash=b"\x00",
model_name="m",
kv_rank=0,
cache_salt="user-abc_123.xyz:42",
)
assert key.cache_salt == "user-abc_123.xyz:42"
class TestObjectKeyIsolation:
"""cache_salt must participate in eq/hash so the L1/L2 caches treat
same-content/different-user entries as distinct."""
def test_different_salts_are_unequal(self):
base = {"chunk_hash": b"x", "model_name": "m", "kv_rank": 0}
a = ObjectKey(**base, cache_salt="alice")
b = ObjectKey(**base, cache_salt="bob")
assert a != b
assert hash(a) != hash(b)
def test_empty_salt_is_unequal_to_any_salted(self):
base = {"chunk_hash": b"x", "model_name": "m", "kv_rank": 0}
unsalted = ObjectKey(**base)
salted = ObjectKey(**base, cache_salt="alice")
assert unsalted != salted
def test_same_salt_are_equal(self):
base = {"chunk_hash": b"x", "model_name": "m", "kv_rank": 0}
a = ObjectKey(**base, cache_salt="alice")
b = ObjectKey(**base, cache_salt="alice")
assert a == b
assert hash(a) == hash(b)
# =============================================================================
# Event Fd Interface Tests
# =============================================================================
class TestEventFdInterface:
def test_three_distinct_event_fds(self, adapter):
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 >= 0
assert lookup_fd >= 0
assert load_fd >= 0
assert len({store_fd, lookup_fd, load_fd}) == 3
# =============================================================================
# Store Interface Tests
# =============================================================================
class TestStoreInterface:
def test_submit_store_returns_task_id(self, adapter):
key = create_object_key(1)
obj = create_memory_obj()
task_id = adapter.submit_store_task([key], [obj])
assert isinstance(task_id, int)
def test_store_signals_event_fd_and_completes(self, adapter):
key = create_object_key(1)
obj = create_memory_obj()
store_fd = adapter.get_store_event_fd()
task_id = adapter.submit_store_task([key], [obj])
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()
def test_pop_clears_completed_tasks(self, adapter):
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)
completed1 = adapter.pop_completed_store_tasks()
assert len(completed1) == 1
completed2 = adapter.pop_completed_store_tasks()
assert len(completed2) == 0
def test_multiple_store_tasks_get_unique_ids(self, adapter):
store_fd = adapter.get_store_event_fd()
task_ids = set()
for i in range(5):
key = create_object_key(i)
obj = create_memory_obj(fill_value=float(i))
task_ids.add(adapter.submit_store_task([key], [obj]))
assert len(task_ids) == 5
# Wait for all completions
completed = {}
while len(completed) < 5:
wait_for_event_fd(store_fd, timeout=5.0)
completed.update(adapter.pop_completed_store_tasks())
for tid in task_ids:
assert completed[tid].is_successful()
def test_batch_store(self, adapter):
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_id = adapter.submit_store_task(keys, objs)
assert wait_for_event_fd(store_fd, timeout=5.0)
completed = adapter.pop_completed_store_tasks()
assert completed[task_id].is_successful()
# =============================================================================
# Lookup and Lock Interface Tests
# =============================================================================
class TestLookupAndLockInterface:
def test_lookup_nonexistent_key(self, adapter):
key = create_object_key(999)
lookup_fd = adapter.get_lookup_and_lock_event_fd()
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 is not None
assert bitmap.test(0) is False
def test_lookup_existing_key(self, adapter):
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 first
adapter.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adapter.pop_completed_store_tasks()
# Lookup
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 is not None
assert bitmap.test(0) is True
def test_lookup_mixed_keys(self, adapter):
existing = create_object_key(1)
missing = create_object_key(999)
obj = create_memory_obj()
store_fd = adapter.get_store_event_fd()
lookup_fd = adapter.get_lookup_and_lock_event_fd()
adapter.submit_store_task([existing], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adapter.pop_completed_store_tasks()
task_id = adapter.submit_lookup_and_lock_task(
[existing, missing], _EMPTY_LAYOUT
)
assert 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
assert bitmap.test(1) is False
def test_query_is_one_shot(self, adapter):
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)
result1 = adapter.query_lookup_and_lock_result(task_id)
assert result1 is not None
result2 = adapter.query_lookup_and_lock_result(task_id)
assert result2 is None
def test_query_unknown_task_returns_none(self, adapter):
assert adapter.query_lookup_and_lock_result(99999) is None
# =============================================================================
# Unlock Interface Tests
# =============================================================================
class TestUnlockInterface:
def test_unlock_does_not_raise(self, adapter):
key = create_object_key(1)
adapter.submit_unlock([key]) # should not raise
def test_unlock_after_lock(self, adapter):
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()
adapter.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adapter.pop_completed_store_tasks()
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)
adapter.submit_unlock([key]) # should not raise
# =============================================================================
# Load Interface Tests
# =============================================================================
class TestLoadInterface:
def test_submit_load_returns_task_id(self, adapter):
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):
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)
def test_load_existing_key_copies_data(self, adapter):
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])
assert 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
# Verify data was copied into the load buffer
assert torch.all(load_obj.tensor == 42.0)
def test_load_nonexistent_key_fails(self, adapter):
key = create_object_key(999)
obj = create_memory_obj()
load_fd = adapter.get_load_event_fd()
task_id = adapter.submit_load_task([key], [obj])
assert wait_for_event_fd(load_fd, timeout=5.0)
bitmap = adapter.query_load_result(task_id)
assert bitmap is not None
# Batch GET failed → no bits set
assert bitmap.test(0) is False
def test_query_load_is_one_shot(self, adapter):
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)
result1 = adapter.query_load_result(task_id)
assert result1 is not None
result2 = adapter.query_load_result(task_id)
assert result2 is None
def test_query_unknown_task_returns_none(self, adapter):
assert adapter.query_load_result(99999) is None
# =============================================================================
# End-to-End Workflow Tests
# =============================================================================
class TestEndToEndWorkflow:
def test_store_lookup_load_workflow(self, adapter):
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()
# Store
store_tid = adapter.submit_store_task([key], [store_obj])
assert wait_for_event_fd(store_fd, timeout=5.0)
assert adapter.pop_completed_store_tasks()[store_tid].is_successful()
# Lookup
lookup_tid = 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(lookup_tid)
assert bitmap.test(0) is True
# Load
load_tid = adapter.submit_load_task([key], [load_obj])
assert wait_for_event_fd(load_fd, timeout=5.0)
bitmap = adapter.query_load_result(load_tid)
assert bitmap.test(0) is True
assert torch.all(load_obj.tensor == 123.0)
# Unlock
adapter.submit_unlock([key])
def test_multiple_objects_workflow(self, adapter):
n = 5
keys = [create_object_key(i) for i in range(n)]
store_objs = [
create_memory_obj(size=64, fill_value=float(i * 10)) for i in range(n)
]
load_objs = [create_memory_obj(size=64, fill_value=0.0) for _ in range(n)]
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_tid = adapter.submit_store_task(keys, store_objs)
assert wait_for_event_fd(store_fd, timeout=5.0)
assert adapter.pop_completed_store_tasks()[store_tid].is_successful()
# Lookup all
lookup_tid = 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(lookup_tid)
for i in range(n):
assert bitmap.test(i) is True
# Load all
load_tid = adapter.submit_load_task(keys, load_objs)
assert wait_for_event_fd(load_fd, timeout=5.0)
bitmap = adapter.query_load_result(load_tid)
for i in range(n):
assert bitmap.test(i) is True
assert torch.all(load_objs[i].tensor == float(i * 10))
# =============================================================================
# Close Tests
# =============================================================================
class TestClose:
def test_close_does_not_raise(self):
mock_client = MockNativeConnector()
adp = NativeConnectorL2Adapter(mock_client)
adp.close()
def test_close_after_operations(self):
mock_client = MockNativeConnector()
adp = NativeConnectorL2Adapter(mock_client)
key = create_object_key(1)
obj = create_memory_obj()
store_fd = adp.get_store_event_fd()
adp.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
adp.close()
# =============================================================================
# Config Tests
# =============================================================================
class TestRESPL2AdapterConfig:
def test_from_dict_minimal(self):
# First Party
from lmcache.v1.distributed.l2_adapters.resp_l2_adapter import (
RESPL2AdapterConfig,
)
config = RESPL2AdapterConfig.from_dict(
{
"type": "resp",
"host": "localhost",
"port": 6379,
}
)
assert config.host == "localhost"
assert config.port == 6379
assert config.num_workers == 8
assert config.username == ""
assert config.password == ""
def test_from_dict_full(self):
# First Party
from lmcache.v1.distributed.l2_adapters.resp_l2_adapter import (
RESPL2AdapterConfig,
)
config = RESPL2AdapterConfig.from_dict(
{
"type": "resp",
"host": "10.0.0.1",
"port": 6380,
"num_workers": 16,
"username": "user",
"password": "pass",
}
)
assert config.host == "10.0.0.1"
assert config.port == 6380
assert config.num_workers == 16
assert config.username == "user"
assert config.password == "pass"
def test_from_dict_missing_host_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.resp_l2_adapter import (
RESPL2AdapterConfig,
)
with pytest.raises(ValueError, match="host"):
RESPL2AdapterConfig.from_dict({"type": "resp", "port": 6379})
def test_from_dict_missing_port_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.resp_l2_adapter import (
RESPL2AdapterConfig,
)
with pytest.raises(ValueError, match="port"):
RESPL2AdapterConfig.from_dict({"type": "resp", "host": "localhost"})
def test_registered_as_resp(self):
# First Party
from lmcache.v1.distributed.l2_adapters.config import (
get_registered_l2_adapter_types,
)
assert "resp" in get_registered_l2_adapter_types()
# =============================================================================
# NativePluginL2AdapterConfig Tests
# =============================================================================
class TestNativePluginL2AdapterConfig:
def test_from_dict_minimal(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
config = NativePluginL2AdapterConfig.from_dict(
{
"type": "native_plugin",
"module_path": "my_ext.connector",
"class_name": "MyClient",
}
)
assert config.module_path == "my_ext.connector"
assert config.class_name == "MyClient"
assert config.adapter_params == {}
def test_from_dict_full(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
config = NativePluginL2AdapterConfig.from_dict(
{
"type": "native_plugin",
"module_path": "my_ext.connector",
"class_name": "MyClient",
"adapter_params": {
"host": "localhost",
"port": 1234,
},
}
)
assert config.module_path == "my_ext.connector"
assert config.class_name == "MyClient"
assert config.adapter_params == {
"host": "localhost",
"port": 1234,
}
def test_from_dict_missing_module_path_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
with pytest.raises(ValueError, match="module_path"):
NativePluginL2AdapterConfig.from_dict(
{
"type": "native_plugin",
"class_name": "X",
}
)
def test_from_dict_missing_class_name_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
with pytest.raises(ValueError, match="class_name"):
NativePluginL2AdapterConfig.from_dict(
{
"type": "native_plugin",
"module_path": "my_ext",
}
)
def test_from_dict_invalid_adapter_params_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
with pytest.raises(ValueError, match="adapter_params"):
NativePluginL2AdapterConfig.from_dict(
{
"type": "native_plugin",
"module_path": "my_ext",
"class_name": "X",
"adapter_params": "not_a_dict",
}
)
def test_registered_as_native_plugin(self):
# First Party
from lmcache.v1.distributed.l2_adapters.config import (
get_registered_l2_adapter_types,
)
assert "native_plugin" in get_registered_l2_adapter_types()
def test_help_returns_string(self):
# First Party
from lmcache.v1.distributed.l2_adapters.native_plugin_l2_adapter import (
NativePluginL2AdapterConfig,
)
h = NativePluginL2AdapterConfig.help()
assert isinstance(h, str)
assert "module_path" in h
assert "class_name" in h
assert "adapter_params" in h
# =============================================================================
# FSNativeL2AdapterConfig Tests
# =============================================================================
class TestFSNativeL2AdapterConfig:
def test_from_dict_minimal(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
config = FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/lmcache_test",
}
)
assert config.base_path == "/tmp/lmcache_test"
assert config.num_workers == 4
assert config.relative_tmp_dir == ""
assert config.use_odirect is False
assert config.read_ahead_size is None
def test_from_dict_full(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
config = FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/data/kv_cache",
"num_workers": 16,
"relative_tmp_dir": ".tmp",
"use_odirect": True,
"read_ahead_size": 4096,
}
)
assert config.base_path == "/data/kv_cache"
assert config.num_workers == 16
assert config.relative_tmp_dir == ".tmp"
assert config.use_odirect is True
assert config.read_ahead_size == 4096
def test_from_dict_missing_base_path_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="base_path"):
FSNativeL2AdapterConfig.from_dict({"type": "fs_native"})
def test_from_dict_empty_base_path_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="base_path"):
FSNativeL2AdapterConfig.from_dict({"type": "fs_native", "base_path": ""})
def test_from_dict_invalid_num_workers_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="num_workers"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"num_workers": 0,
}
)
def test_from_dict_zero_num_workers_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="num_workers"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"num_workers": -1,
}
)
def test_from_dict_invalid_relative_tmp_dir_raises(
self,
):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="relative_tmp_dir"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"relative_tmp_dir": 123,
}
)
def test_from_dict_invalid_use_odirect_raises(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="use_odirect"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"use_odirect": "yes",
}
)
def test_from_dict_invalid_read_ahead_size_raises(
self,
):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="read_ahead_size"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"read_ahead_size": -1,
}
)
def test_from_dict_zero_read_ahead_size_raises(
self,
):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
with pytest.raises(ValueError, match="read_ahead_size"):
FSNativeL2AdapterConfig.from_dict(
{
"type": "fs_native",
"base_path": "/tmp/x",
"read_ahead_size": 0,
}
)
def test_registered_as_fs_native(self):
# First Party
from lmcache.v1.distributed.l2_adapters.config import (
get_registered_l2_adapter_types,
)
assert "fs_native" in get_registered_l2_adapter_types()
def test_help_returns_string(self):
# First Party
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
h = FSNativeL2AdapterConfig.help()
assert isinstance(h, str)
assert "base_path" in h
assert "num_workers" in h
assert "use_odirect" in h
assert "read_ahead_size" in h
def test_type_name_lookup(self):
# First Party
from lmcache.v1.distributed.l2_adapters.config import (
get_type_name_for_config,
)
from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
FSNativeL2AdapterConfig,
)
cfg = FSNativeL2AdapterConfig(
base_path="/tmp/test",
)
assert get_type_name_for_config(cfg) == "fs_native"
# =============================================================================
# Delete Interface Tests
# =============================================================================
class TestDeleteInterface:
def test_delete_existing_key(self, adapter):
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()
# Verify exists
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.test(0) is True
adapter.submit_unlock([key])
# Delete (synchronous)
adapter.delete([key])
# Verify gone
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.test(0) is False
def test_delete_nonexistent_key(self, adapter):
key = create_object_key(999)
adapter.delete([key]) # should not raise
def test_delete_empty_keys(self, adapter):
adapter.delete([]) # should not raise
def test_delete_batch(self, adapter):
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()
lookup_fd = adapter.get_lookup_and_lock_event_fd()
# Store all
adapter.submit_store_task(keys, objs)
wait_for_event_fd(store_fd, timeout=5.0)
adapter.pop_completed_store_tasks()
# Delete first 3
adapter.delete(keys[:3])
# Verify: first 3 gone, last 2 remain
task_id = adapter.submit_lookup_and_lock_task(keys, _EMPTY_LAYOUT)
wait_for_event_fd(lookup_fd, timeout=5.0)
bitmap = adapter.query_lookup_and_lock_result(task_id)
for i in range(3):
assert bitmap.test(i) is False
for i in range(3, 5):
assert bitmap.test(i) is True
adapter.submit_unlock(keys[3:])
# =============================================================================
# Delete Backward Compatibility Tests
# =============================================================================
class TestDeleteBackwardCompatibility:
def test_delete_noop_without_submit_batch_delete(self):
"""Connector without submit_batch_delete => delete is no-op."""
class NoDeleteConnector:
"""Mock connector that only has the 6 original methods."""
def __init__(self):
self._efd = create_event_notifier()
self._closed = False
def event_fd(self) -> int:
return self._efd.fileno()
def submit_batch_get(self, keys, memoryviews):
return 0
def submit_batch_set(self, keys, memoryviews):
return 0
def submit_batch_exists(self, keys):
return 0
def drain_completions(self):
return []
def close(self):
if not self._closed:
self._closed = True
self._efd.close()
client = NoDeleteConnector()
adp = NativeConnectorL2Adapter(client)
try:
key = create_object_key(1)
adp.delete([key]) # should not raise, just no-op
finally:
adp.close()
# =============================================================================
# Usage Tracking Tests
# =============================================================================
@pytest.fixture
def adapter_with_capacity():
"""Adapter with max_capacity_gb set for usage tracking tests."""
mock_client = MockNativeConnector()
# 100 floats * 4 bytes = 400 bytes per obj; capacity = 2000 bytes = 2000/1024^3 GB
adp = NativeConnectorL2Adapter(mock_client, max_capacity_gb=2000 / (1024**3))
yield adp
adp.close()
class TestUsageTracking:
def test_get_usage_without_capacity(self, adapter):
"""Without max_capacity_bytes, usage_fraction == -1 (sentinel)."""
usage = adapter.get_usage()
assert usage.usage_fraction == -1.0
assert usage.total_capacity_bytes == 0
def test_get_usage_starts_at_zero(self, adapter_with_capacity):
usage = adapter_with_capacity.get_usage()
assert usage.usage_fraction == 0.0
assert usage.total_bytes_used == 0
def test_get_usage_after_store(self, adapter_with_capacity):
adp = adapter_with_capacity
store_fd = adp.get_store_event_fd()
key = create_object_key(1)
obj = create_memory_obj(size=100, fill_value=1.0) # 100 floats = 400 bytes
adp.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
usage = adp.get_usage()
# 400 bytes / 2000 bytes = 0.2
assert usage.usage_fraction == pytest.approx(0.2)
assert usage.total_bytes_used == 400
def test_get_usage_after_delete(self, adapter_with_capacity):
adp = adapter_with_capacity
store_fd = adp.get_store_event_fd()
key = create_object_key(1)
obj = create_memory_obj(size=100, fill_value=1.0)
# Store
adp.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
assert adp.get_usage().usage_fraction == pytest.approx(0.2)
# Delete
adp.delete([key])
assert adp.get_usage().usage_fraction == pytest.approx(0.0)
assert adp.get_usage().total_bytes_used == 0
def test_get_usage_store_delete_cycle(self, adapter_with_capacity):
adp = adapter_with_capacity
store_fd = adp.get_store_event_fd()
# Store 3 objects (3 * 400 = 1200 bytes)
keys = [create_object_key(i) for i in range(3)]
objs = [create_memory_obj(size=100, fill_value=float(i)) for i in range(3)]
adp.submit_store_task(keys, objs)
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
usage = adp.get_usage()
assert usage.usage_fraction == pytest.approx(1200 / 2000)
assert usage.total_bytes_used == 1200
# Delete 2
adp.delete(keys[:2])
usage = adp.get_usage()
assert usage.usage_fraction == pytest.approx(400 / 2000)
assert usage.total_bytes_used == 400
def test_idempotent_store_no_double_count(self, adapter_with_capacity):
adp = adapter_with_capacity
store_fd = adp.get_store_event_fd()
key = create_object_key(1)
obj = create_memory_obj(size=100, fill_value=1.0)
# Store same key twice
adp.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
adp.submit_store_task([key], [obj])
wait_for_event_fd(store_fd, timeout=5.0)
adp.pop_completed_store_tasks()
# Should only count once
usage = adp.get_usage()
assert usage.usage_fraction == pytest.approx(0.2)
assert usage.total_bytes_used == 400