wehub-resource-sync
393 lines
13 KiB
Python
393 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for runtime add/remove of L2 adapters.
|
|
|
|
Covers both the controller-level primitives (StoreController /
|
|
PrefetchController add_adapter / request_remove_adapter with graceful
|
|
drain) and the StorageManager orchestration (add_l2_adapter /
|
|
delete_l2_adapter), including stable-id semantics and the
|
|
active/draining observability surface.
|
|
"""
|
|
|
|
# Standard
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.config import (
|
|
EvictionConfig,
|
|
L1ManagerConfig,
|
|
L1MemoryManagerConfig,
|
|
StorageManagerConfig,
|
|
)
|
|
from lmcache.v1.distributed.l1_manager import L1Manager
|
|
from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig
|
|
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import (
|
|
MockL2Adapter,
|
|
MockL2AdapterConfig,
|
|
)
|
|
from lmcache.v1.distributed.storage_controllers.store_controller import StoreController
|
|
from lmcache.v1.distributed.storage_controllers.store_policy import (
|
|
AdapterDescriptor,
|
|
DefaultStorePolicy,
|
|
)
|
|
from lmcache.v1.distributed.storage_manager import StorageManager
|
|
|
|
# Skip all tests in this module if CUDA is not available
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="CUDA is not available"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Helpers
|
|
# =============================================================================
|
|
|
|
|
|
def make_object_key(chunk_id: int) -> ObjectKey:
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name="test_model",
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def make_layout() -> MemoryLayoutDesc:
|
|
return MemoryLayoutDesc(
|
|
shapes=[torch.Size([100, 2, 512])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
|
|
def should_use_lazy_alloc() -> bool:
|
|
return torch.cuda.is_available()
|
|
|
|
|
|
def wait_for_condition(predicate, timeout: float = 5.0, poll: float = 0.02) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(poll)
|
|
return False
|
|
|
|
|
|
def make_mock_config() -> MockL2AdapterConfig:
|
|
return MockL2AdapterConfig(max_size_gb=0.01, mock_bandwidth_gb=10.0)
|
|
|
|
|
|
def make_adapter() -> MockL2Adapter:
|
|
return MockL2Adapter(make_mock_config())
|
|
|
|
|
|
def make_descriptor(index: int) -> AdapterDescriptor:
|
|
return AdapterDescriptor(index=index, config=make_mock_config())
|
|
|
|
|
|
def adapter_by_id(sm: StorageManager, adapter_id: int) -> MockL2Adapter:
|
|
"""Fetch an active adapter by stable id via the public ``l2_adapters()``
|
|
API (no private-member access)."""
|
|
for desc, adapter in sm.l2_adapters():
|
|
if desc.index == adapter_id:
|
|
return adapter # type: ignore[return-value]
|
|
raise AssertionError(f"no active L2 adapter with id {adapter_id}")
|
|
|
|
|
|
def active_adapter_ids(sm: StorageManager) -> set[int]:
|
|
"""Set of active adapter ids exposed by the public ``l2_adapters()`` API."""
|
|
return {desc.index for desc, _ in sm.l2_adapters()}
|
|
|
|
|
|
def write_keys_to_l1(
|
|
l1_manager: L1Manager,
|
|
keys: list[ObjectKey],
|
|
layout: MemoryLayoutDesc,
|
|
) -> None:
|
|
results = l1_manager.reserve_write(
|
|
keys=keys,
|
|
is_temporary=[False] * len(keys),
|
|
layout_desc=layout,
|
|
mode="new",
|
|
)
|
|
written = [k for k, (e, m) in results.items() if m is not None]
|
|
if written:
|
|
l1_manager.finish_write(written)
|
|
|
|
|
|
@pytest.fixture
|
|
def l1_manager():
|
|
config = L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=128 * 1024 * 1024,
|
|
use_lazy=should_use_lazy_alloc(),
|
|
init_size_in_bytes=64 * 1024 * 1024,
|
|
align_bytes=0x1000,
|
|
),
|
|
write_ttl_seconds=600,
|
|
read_ttl_seconds=300,
|
|
)
|
|
mgr = L1Manager(config)
|
|
yield mgr
|
|
mgr.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def empty_storage_manager_config():
|
|
"""A StorageManagerConfig with no L2 adapters configured."""
|
|
return StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=128 * 1024 * 1024,
|
|
use_lazy=should_use_lazy_alloc(),
|
|
init_size_in_bytes=64 * 1024 * 1024,
|
|
align_bytes=0x1000,
|
|
),
|
|
write_ttl_seconds=600,
|
|
read_ttl_seconds=300,
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="LRU"),
|
|
l2_adapter_config=L2AdaptersConfig(adapters=[]),
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# StoreController-level tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStoreControllerRuntimeAdapters:
|
|
def test_add_adapter_routes_new_stores(self, l1_manager):
|
|
"""A key written after add_adapter is stored to the new adapter."""
|
|
ctrl = StoreController(
|
|
l1_manager=l1_manager,
|
|
l2_adapters=[],
|
|
adapter_descriptors=[],
|
|
policy=DefaultStorePolicy(),
|
|
)
|
|
ctrl.start()
|
|
try:
|
|
adapter = make_adapter()
|
|
ctrl.add_adapter(0, adapter, make_descriptor(0))
|
|
|
|
assert ctrl.report_status()["num_active_adapters"] == 1
|
|
assert ctrl.report_status()["num_draining_adapters"] == 0
|
|
|
|
layout = make_layout()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
write_keys_to_l1(l1_manager, keys, layout)
|
|
|
|
ok = wait_for_condition(
|
|
lambda: adapter.debug_get_stored_object_count() == 3
|
|
)
|
|
assert ok, "Keys should be stored to the runtime-added adapter"
|
|
finally:
|
|
ctrl.stop()
|
|
adapter.close()
|
|
|
|
def test_remove_adapter_drains_and_detaches(self, l1_manager):
|
|
"""request_remove_adapter drains in-flight work then detaches."""
|
|
adapter = make_adapter()
|
|
ctrl = StoreController(
|
|
l1_manager=l1_manager,
|
|
l2_adapters=[adapter],
|
|
adapter_descriptors=[make_descriptor(0)],
|
|
policy=DefaultStorePolicy(),
|
|
)
|
|
ctrl.start()
|
|
try:
|
|
layout = make_layout()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
write_keys_to_l1(l1_manager, keys, layout)
|
|
assert wait_for_condition(
|
|
lambda: adapter.debug_get_stored_object_count() == 3
|
|
)
|
|
|
|
done = ctrl.request_remove_adapter(0)
|
|
assert done.wait(timeout=5.0), "Adapter should drain and detach"
|
|
|
|
status = ctrl.report_status()
|
|
assert status["num_l2_adapters"] == 0
|
|
assert status["num_active_adapters"] == 0
|
|
assert status["num_draining_adapters"] == 0
|
|
finally:
|
|
ctrl.stop()
|
|
adapter.close()
|
|
|
|
def test_remove_stops_new_routing(self, l1_manager):
|
|
"""After removal, new writes are not routed to the removed adapter."""
|
|
keep = make_adapter()
|
|
drop = make_adapter()
|
|
ctrl = StoreController(
|
|
l1_manager=l1_manager,
|
|
l2_adapters=[keep, drop],
|
|
adapter_descriptors=[make_descriptor(0), make_descriptor(1)],
|
|
policy=DefaultStorePolicy(),
|
|
)
|
|
ctrl.start()
|
|
try:
|
|
done = ctrl.request_remove_adapter(1)
|
|
assert done.wait(timeout=5.0)
|
|
|
|
layout = make_layout()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
write_keys_to_l1(l1_manager, keys, layout)
|
|
|
|
assert wait_for_condition(lambda: keep.debug_get_stored_object_count() == 3)
|
|
# The dropped adapter must not receive the new keys.
|
|
time.sleep(0.2)
|
|
assert drop.debug_get_stored_object_count() == 0
|
|
finally:
|
|
ctrl.stop()
|
|
keep.close()
|
|
drop.close()
|
|
|
|
def test_double_remove_is_safe(self, l1_manager):
|
|
"""Removing an already-removed adapter signals immediately."""
|
|
adapter = make_adapter()
|
|
ctrl = StoreController(
|
|
l1_manager=l1_manager,
|
|
l2_adapters=[adapter],
|
|
adapter_descriptors=[make_descriptor(0)],
|
|
policy=DefaultStorePolicy(),
|
|
)
|
|
ctrl.start()
|
|
try:
|
|
assert ctrl.request_remove_adapter(0).wait(timeout=5.0)
|
|
# Second removal of the same (now-detached) id must not hang.
|
|
assert ctrl.request_remove_adapter(0).wait(timeout=5.0)
|
|
finally:
|
|
ctrl.stop()
|
|
adapter.close()
|
|
|
|
|
|
# =============================================================================
|
|
# StorageManager-level tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStorageManagerRuntimeAdapters:
|
|
def test_add_to_empty_manager(self, empty_storage_manager_config):
|
|
"""add_l2_adapter wires a new adapter into all controllers."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
assert sm.report_status()["num_l2_adapters"] == 0
|
|
|
|
adapter_id = sm.add_l2_adapter(make_mock_config())
|
|
assert adapter_id == 0
|
|
|
|
status = sm.report_status()
|
|
assert status["num_l2_adapters"] == 1
|
|
assert status["store_controller"]["num_active_adapters"] == 1
|
|
assert status["prefetch_controller"]["num_active_adapters"] == 1
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_add_then_store(self, empty_storage_manager_config):
|
|
"""Keys written after add_l2_adapter land in the new adapter."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
adapter_id = sm.add_l2_adapter(make_mock_config())
|
|
adapter = adapter_by_id(sm, adapter_id)
|
|
|
|
layout = make_layout()
|
|
keys = [make_object_key(i) for i in range(4)]
|
|
ret = sm.reserve_write(keys, layout, mode="new")
|
|
sm.finish_write(list(ret.keys()))
|
|
|
|
assert wait_for_condition(
|
|
lambda: all(adapter.debug_has_key(k) for k in keys), # type: ignore
|
|
timeout=10.0,
|
|
)
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_delete_detaches_everywhere(self, empty_storage_manager_config):
|
|
"""delete_l2_adapter drains and removes the adapter from all controllers."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
adapter_id = sm.add_l2_adapter(make_mock_config())
|
|
assert sm.report_status()["num_l2_adapters"] == 1
|
|
|
|
sm.delete_l2_adapter(adapter_id, timeout=10.0)
|
|
|
|
status = sm.report_status()
|
|
assert status["num_l2_adapters"] == 0
|
|
assert status["store_controller"]["num_l2_adapters"] == 0
|
|
assert status["prefetch_controller"]["num_l2_adapters"] == 0
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_delete_unknown_id_raises(self, empty_storage_manager_config):
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
with pytest.raises(ValueError):
|
|
sm.delete_l2_adapter(999)
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_stable_ids_across_add_delete(self, empty_storage_manager_config):
|
|
"""Ids are monotonic and never reused; delete does not shift others."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
id0 = sm.add_l2_adapter(make_mock_config())
|
|
id1 = sm.add_l2_adapter(make_mock_config())
|
|
assert (id0, id1) == (0, 1)
|
|
|
|
sm.delete_l2_adapter(id0, timeout=10.0)
|
|
|
|
id2 = sm.add_l2_adapter(make_mock_config())
|
|
assert id2 == 2, "Ids must not be reused after a delete"
|
|
|
|
# The surviving adapter keeps its original id.
|
|
assert active_adapter_ids(sm) == {1, 2}
|
|
assert sm.report_status()["num_l2_adapters"] == 2
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_add_after_delete_routes_to_new_adapter(self, empty_storage_manager_config):
|
|
"""A fresh adapter added after a delete receives new stores."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
first = sm.add_l2_adapter(make_mock_config())
|
|
sm.delete_l2_adapter(first, timeout=10.0)
|
|
second = sm.add_l2_adapter(make_mock_config())
|
|
adapter = adapter_by_id(sm, second)
|
|
|
|
layout = make_layout()
|
|
keys = [make_object_key(i) for i in range(3)]
|
|
ret = sm.reserve_write(keys, layout, mode="new")
|
|
sm.finish_write(list(ret.keys()))
|
|
|
|
assert wait_for_condition(
|
|
lambda: all(adapter.debug_has_key(k) for k in keys), # type: ignore
|
|
timeout=10.0,
|
|
)
|
|
finally:
|
|
sm.close()
|
|
|
|
def test_l2_adapters_reflects_runtime_changes(self, empty_storage_manager_config):
|
|
"""l2_adapters() tracks runtime add/delete in stable-id order."""
|
|
sm = StorageManager(empty_storage_manager_config)
|
|
try:
|
|
assert sm.l2_adapters() == []
|
|
|
|
id0 = sm.add_l2_adapter(make_mock_config())
|
|
id1 = sm.add_l2_adapter(make_mock_config())
|
|
pairs = sm.l2_adapters()
|
|
assert [d.index for d, _ in pairs] == [id0, id1]
|
|
a0, a1 = pairs[0][1], pairs[1][1]
|
|
assert a0 is not a1
|
|
|
|
# Deleting one leaves the survivor (same instance); order is by id.
|
|
sm.delete_l2_adapter(id0, timeout=10.0)
|
|
pairs = sm.l2_adapters()
|
|
assert len(pairs) == 1
|
|
assert pairs[0][0].index == id1
|
|
assert pairs[0][1] is a1
|
|
finally:
|
|
sm.close()
|