lmcache--lmcache
wehub-resource-sync
152 行
4.3 KiB
Python
152 行
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Integration tests for the Aerospike L2 adapter (native connector).
|
|
|
|
Requires Aerospike CE and BUILD_AEROSPIKE=1 extension. Skipped otherwise.
|
|
"""
|
|
|
|
# Standard
|
|
import os
|
|
import select
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.l2_adapters.factory import create_l2_adapter_from_registry
|
|
from lmcache.v1.memory_management import (
|
|
MemoryFormat,
|
|
MemoryObjMetadata,
|
|
TensorMemoryObj,
|
|
)
|
|
from lmcache.v1.platform import consume_fd
|
|
|
|
_EMPTY_LAYOUT = MemoryLayoutDesc(shapes=[], dtypes=[])
|
|
|
|
AEROSPIKE_HOST = os.environ.get("AEROSPIKE_TEST_HOST", "127.0.0.1")
|
|
AEROSPIKE_PORT = int(os.environ.get("AEROSPIKE_TEST_PORT", "3000"))
|
|
AEROSPIKE_NAMESPACE = os.environ.get("AEROSPIKE_TEST_NAMESPACE", "lmcache")
|
|
RUN_AEROSPIKE_IT = os.environ.get("RUN_AEROSPIKE_INTEGRATION") == "1"
|
|
|
|
|
|
def _aerospike_available() -> bool:
|
|
if not RUN_AEROSPIKE_IT:
|
|
return False
|
|
try:
|
|
# Third Party
|
|
import aerospike
|
|
|
|
client = aerospike.client(
|
|
{"hosts": [(AEROSPIKE_HOST, AEROSPIKE_PORT)]}
|
|
).connect()
|
|
info = client.info_random_node(f"namespace/{AEROSPIKE_NAMESPACE}")
|
|
client.close()
|
|
return "nsup-period" in info
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _native_extension_available() -> bool:
|
|
try:
|
|
# First Party
|
|
from lmcache.lmcache_aerospike import LMCacheAerospikeClient # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
requires_aerospike = pytest.mark.skipif(
|
|
not _aerospike_available(),
|
|
reason=(
|
|
f"Aerospike not available at {AEROSPIKE_HOST}:{AEROSPIKE_PORT} "
|
|
"(set RUN_AEROSPIKE_INTEGRATION=1)"
|
|
),
|
|
)
|
|
requires_native = pytest.mark.skipif(
|
|
not _native_extension_available(),
|
|
reason="lmcache.lmcache_aerospike extension not built",
|
|
)
|
|
|
|
|
|
def _wait_fd(fd: int, timeout: float = 30.0) -> None:
|
|
poller = select.poll()
|
|
poller.register(fd, select.POLLIN)
|
|
events = poller.poll(timeout * 1000)
|
|
assert events, "timed out waiting for eventfd"
|
|
try:
|
|
consume_fd(fd)
|
|
except BlockingIOError:
|
|
pass
|
|
|
|
|
|
def _make_tensor_obj(size: int, fill: float) -> TensorMemoryObj:
|
|
raw_data = torch.empty(size, dtype=torch.float32)
|
|
raw_data.fill_(fill)
|
|
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 _object_key(suffix: int) -> ObjectKey:
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(suffix),
|
|
model_name="aerospike-it",
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def _adapter_config():
|
|
# First Party
|
|
from lmcache.v1.distributed.l2_adapters.aerospike_l2_adapter import (
|
|
AerospikeL2AdapterConfig,
|
|
)
|
|
|
|
return AerospikeL2AdapterConfig(
|
|
hosts=f"{AEROSPIKE_HOST}:{AEROSPIKE_PORT}",
|
|
namespace=AEROSPIKE_NAMESPACE,
|
|
set_name="kv_chunks_aerospike_it",
|
|
num_workers=2,
|
|
)
|
|
|
|
|
|
@requires_aerospike
|
|
@requires_native
|
|
class TestAerospikeL2Integration:
|
|
def test_store_lookup_load_roundtrip(self):
|
|
adapter = create_l2_adapter_from_registry(_adapter_config())
|
|
try:
|
|
key = _object_key(9001)
|
|
store_obj = _make_tensor_obj(64, 42.0)
|
|
load_obj = _make_tensor_obj(64, 0.0)
|
|
|
|
tid = adapter.submit_store_task([key], [store_obj])
|
|
_wait_fd(adapter.get_store_event_fd())
|
|
done = adapter.pop_completed_store_tasks()
|
|
assert done[tid].is_successful()
|
|
|
|
lookup_tid = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
_wait_fd(adapter.get_lookup_and_lock_event_fd())
|
|
lookup_bm = adapter.query_lookup_and_lock_result(lookup_tid)
|
|
assert lookup_bm is not None
|
|
assert lookup_bm.test(0)
|
|
|
|
load_tid = adapter.submit_load_task([key], [load_obj])
|
|
_wait_fd(adapter.get_load_event_fd())
|
|
load_bm = adapter.query_load_result(load_tid)
|
|
assert load_bm is not None
|
|
assert load_bm.test(0)
|
|
assert torch.all(load_obj.tensor == 42.0)
|
|
|
|
adapter.submit_unlock([key])
|
|
finally:
|
|
adapter.close()
|