wehub-resource-sync
485 lines
17 KiB
Python
485 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Test cases for storage_plugin_launcher.
|
|
|
|
This module tests the storage plugins loading mechanism in CreateStorageBackends.
|
|
It creates a simple mock storage plugin extends StoragePluginInterface
|
|
and verifies that:
|
|
1. The backend is properly loaded via the configuration
|
|
2. The backend methods can be called through StorageManager
|
|
"""
|
|
|
|
# Standard
|
|
from typing import Any, Callable, List, Optional, Sequence
|
|
import asyncio
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.utils import CacheEngineKey
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.event_manager import EventManager
|
|
from lmcache.v1.memory_allocators.ad_hoc_memory_allocator import AdHocMemoryAllocator
|
|
from lmcache.v1.memory_management import MemoryObj
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
from lmcache.v1.storage_backend import CreateStorageBackends
|
|
from lmcache.v1.storage_backend.abstract_backend import (
|
|
AllocatorBackendInterface,
|
|
StoragePluginInterface,
|
|
)
|
|
from lmcache.v1.storage_backend.storage_manager import StorageManager
|
|
from tests.v1.utils import create_test_memory_obj
|
|
|
|
|
|
class MockStoragePlugin(StoragePluginInterface):
|
|
"""
|
|
A mock storage plugin for testing the storage plugin functionality.
|
|
This storage plugin extends StoragePluginInterface and provides simple
|
|
in-memory storage for testing purposes.
|
|
|
|
This plugin uses instance-level storage for call history and data storage
|
|
to enable proper testing of dynamically loaded storage plugins.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config=None,
|
|
metadata=None,
|
|
local_cpu_backend=None,
|
|
loop=None,
|
|
dst_device: str = "cpu",
|
|
):
|
|
super().__init__(
|
|
dst_device=dst_device,
|
|
config=config,
|
|
metadata=metadata,
|
|
local_cpu_backend=local_cpu_backend,
|
|
loop=loop,
|
|
)
|
|
self._allocator = AdHocMemoryAllocator(device="cpu")
|
|
# Use instance-level storage
|
|
self.call_history: List[str] = []
|
|
self.storage: dict = {}
|
|
|
|
def __str__(self):
|
|
return "MockStoragePlugin"
|
|
|
|
def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
|
|
"""Check if key exists in storage."""
|
|
self.call_history.append(f"contains:{key.chunk_hash}")
|
|
return key.chunk_hash in self.storage
|
|
|
|
def exists_in_put_tasks(self, key: CacheEngineKey) -> bool:
|
|
"""Check if key is in ongoing put tasks."""
|
|
self.call_history.append(f"exists_in_put_tasks:{key.chunk_hash}")
|
|
return False
|
|
|
|
def batched_submit_put_task(
|
|
self,
|
|
keys: Sequence[CacheEngineKey],
|
|
objs: List[MemoryObj],
|
|
transfer_spec: Any = None,
|
|
on_complete_callback: Optional[Callable[[CacheEngineKey], None]] = None,
|
|
) -> None:
|
|
"""Submit a batched put task."""
|
|
for key, obj in zip(keys, objs, strict=True):
|
|
self.call_history.append(f"batched_submit_put_task:{key.chunk_hash}")
|
|
self.storage[key.chunk_hash] = obj
|
|
|
|
def get_blocking(self, key: CacheEngineKey) -> Optional[MemoryObj]:
|
|
"""Blocking get for a single key."""
|
|
self.call_history.append(f"get_blocking:{key.chunk_hash}")
|
|
return self.storage.get(key.chunk_hash)
|
|
|
|
def pin(self, key: CacheEngineKey) -> bool:
|
|
"""Pin a key in storage."""
|
|
self.call_history.append(f"pin:{key.chunk_hash}")
|
|
return key.chunk_hash in self.storage
|
|
|
|
def unpin(self, key: CacheEngineKey) -> bool:
|
|
"""Unpin a key in storage."""
|
|
self.call_history.append(f"unpin:{key.chunk_hash}")
|
|
return key.chunk_hash in self.storage
|
|
|
|
def remove(self, key: CacheEngineKey, force: bool = True) -> bool:
|
|
"""Remove a key from storage."""
|
|
self.call_history.append(f"remove:{key.chunk_hash}")
|
|
if key.chunk_hash in self.storage:
|
|
del self.storage[key.chunk_hash]
|
|
return True
|
|
return False
|
|
|
|
def get_allocator_backend(self) -> "AllocatorBackendInterface":
|
|
"""Get the allocator backend."""
|
|
# Return the LocalCPUBackend if available, otherwise use self
|
|
if self.local_cpu_backend is not None:
|
|
return self.local_cpu_backend
|
|
raise RuntimeError("No allocator backend available")
|
|
|
|
def close(self) -> None:
|
|
"""Close the backend."""
|
|
self.call_history.append("close")
|
|
|
|
|
|
# Module-level constants for storage plugin configuration
|
|
MOCK_BACKEND_EXTRA_CONFIG = {
|
|
"storage_plugin.mock_backend.module_path": (
|
|
"tests.v1.storage_backend.test_storage_plugin"
|
|
),
|
|
"storage_plugin.mock_backend.class_name": "MockStoragePlugin",
|
|
}
|
|
MOCK_BACKEND_STORAGE_PLUGINS = ["mock_backend"]
|
|
|
|
|
|
def create_test_key(key_id: int = 0) -> CacheEngineKey:
|
|
"""Create a test CacheEngineKey."""
|
|
return CacheEngineKey(
|
|
model_name="test_model",
|
|
world_size=1,
|
|
worker_id=0,
|
|
chunk_hash=key_id,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
|
|
def create_test_config(
|
|
extra_config: Optional[dict] = None,
|
|
storage_plugins: Optional[List[str]] = None,
|
|
):
|
|
"""Create a test configuration for storage plugin testing."""
|
|
config = LMCacheEngineConfig.from_defaults(
|
|
chunk_size=256,
|
|
local_cpu=True,
|
|
max_local_cpu_size=0.1, # Small size for testing
|
|
lmcache_instance_id="test_storage_plugin_instance",
|
|
)
|
|
if extra_config:
|
|
config.extra_config = extra_config
|
|
if storage_plugins:
|
|
config.storage_plugins = storage_plugins
|
|
return config
|
|
|
|
|
|
def create_test_metadata():
|
|
"""Create test metadata for testing."""
|
|
return LMCacheMetadata(
|
|
model_name="test_model",
|
|
world_size=1,
|
|
local_world_size=1,
|
|
worker_id=0,
|
|
local_worker_id=0,
|
|
kv_dtype=torch.bfloat16,
|
|
kv_shape=(
|
|
4,
|
|
2,
|
|
256,
|
|
8,
|
|
128,
|
|
), # (num_layers, 2, chunk_size, num_heads, head_dim)
|
|
)
|
|
|
|
|
|
def get_mock_backend(storage_manager_or_backends):
|
|
"""
|
|
Get the mock storage plugin instance from storage manager or backends dict.
|
|
Handles both StorageManager and OrderedDict[str, StorageBackendInterface].
|
|
"""
|
|
if hasattr(storage_manager_or_backends, "storage_backends"):
|
|
backends = storage_manager_or_backends.storage_backends
|
|
else:
|
|
backends = storage_manager_or_backends
|
|
|
|
return backends.get("mock_backend")
|
|
|
|
|
|
class TestCreateDynamicBackends:
|
|
"""Test cases for storage_plugin_launcher functionality."""
|
|
|
|
@pytest.fixture
|
|
def event_manager(self):
|
|
"""Create an EventManager for testing."""
|
|
return EventManager()
|
|
|
|
@pytest.fixture
|
|
def async_loop(self):
|
|
"""Create an async event loop for testing."""
|
|
loop = asyncio.new_event_loop()
|
|
yield loop
|
|
loop.close()
|
|
|
|
@pytest.fixture
|
|
def storage_manager_with_mock_backend(self, event_manager):
|
|
"""Fixture to create a StorageManager with the mock storage plugin loaded."""
|
|
config = create_test_config(
|
|
extra_config=MOCK_BACKEND_EXTRA_CONFIG,
|
|
storage_plugins=MOCK_BACKEND_STORAGE_PLUGINS,
|
|
)
|
|
metadata = create_test_metadata()
|
|
|
|
storage_manager = StorageManager(
|
|
config=config,
|
|
metadata=metadata,
|
|
event_manager=event_manager,
|
|
)
|
|
|
|
yield storage_manager
|
|
|
|
storage_manager.close()
|
|
|
|
def test_storage_plugin_is_loaded(self, async_loop):
|
|
"""
|
|
Test that a storage plugin is properly loaded via CreateStorageBackends.
|
|
"""
|
|
config = create_test_config(
|
|
extra_config=MOCK_BACKEND_EXTRA_CONFIG,
|
|
storage_plugins=MOCK_BACKEND_STORAGE_PLUGINS,
|
|
)
|
|
metadata = create_test_metadata()
|
|
|
|
# Create storage backends
|
|
storage_backends = CreateStorageBackends(
|
|
config=config,
|
|
metadata=metadata,
|
|
loop=async_loop,
|
|
dst_device="cpu",
|
|
)
|
|
|
|
# Verify the dynamic backend is loaded
|
|
assert "mock_backend" in storage_backends, (
|
|
f"Expected 'mock_backend' in storage_backends, "
|
|
f"got: {list(storage_backends.keys())}"
|
|
)
|
|
|
|
# Verify it's an instance by checking class name
|
|
# (since the dynamically imported class is a different object)
|
|
mock_backend = storage_backends["mock_backend"]
|
|
assert mock_backend.__class__.__name__ == MockStoragePlugin.__name__, (
|
|
f"Expected MockStoragePlugin instance, got: {type(mock_backend)}"
|
|
)
|
|
|
|
# Verify it extends StoragePluginInterface
|
|
assert isinstance(mock_backend, StoragePluginInterface), (
|
|
f"Expected StoragePluginInterface instance, got: {type(mock_backend)}"
|
|
)
|
|
|
|
# Verify the backend has the expected attributes (call_history, storage)
|
|
assert hasattr(mock_backend, "call_history"), (
|
|
"Expected mock_backend to have 'call_history' attribute"
|
|
)
|
|
assert hasattr(mock_backend, "storage"), (
|
|
"Expected mock_backend to have 'storage' attribute"
|
|
)
|
|
|
|
# Close all backends
|
|
for backend in storage_backends.values():
|
|
backend.close()
|
|
|
|
def test_dynamic_backend_contains_called_via_storage_manager(
|
|
self, storage_manager_with_mock_backend
|
|
):
|
|
"""
|
|
Test that StorageManager calls the dynamic backend's contains method.
|
|
"""
|
|
storage_manager = storage_manager_with_mock_backend
|
|
|
|
# Verify MockStoragePlugin is in the storage backends
|
|
assert "mock_backend" in storage_manager.storage_backends, (
|
|
f"Expected 'mock_backend' in storage_backends, "
|
|
f"got: {list(storage_manager.storage_backends.keys())}"
|
|
)
|
|
|
|
# Get the mock backend instance to access its call history
|
|
mock_backend = get_mock_backend(storage_manager)
|
|
assert mock_backend is not None, "Expected mock_backend to be loaded"
|
|
|
|
# Clear the call history before the test
|
|
mock_backend.call_history.clear()
|
|
|
|
# Create a test key
|
|
test_key = create_test_key(key_id=12345)
|
|
|
|
# Call contains via storage_manager
|
|
# The contains method iterates over all backends
|
|
result = storage_manager.contains(test_key)
|
|
|
|
# Verify the dynamic backend's contains was called
|
|
assert any("contains:12345" in call for call in mock_backend.call_history), (
|
|
f"Expected 'contains:12345' in call history, "
|
|
f"got: {mock_backend.call_history}"
|
|
)
|
|
|
|
# The key doesn't exist, so result should be None
|
|
assert result is None, f"Expected None, got: {result}"
|
|
|
|
def test_dynamic_backend_get_blocking_called_via_storage_manager(
|
|
self, storage_manager_with_mock_backend
|
|
):
|
|
"""
|
|
Test that StorageManager can retrieve data through the dynamic backend.
|
|
"""
|
|
storage_manager = storage_manager_with_mock_backend
|
|
|
|
# Get the mock backend instance
|
|
mock_backend = get_mock_backend(storage_manager)
|
|
assert mock_backend is not None, "Expected mock_backend to be loaded"
|
|
|
|
# Pre-populate storage with a test memory object
|
|
test_key = create_test_key(key_id=67890)
|
|
test_memory_obj = create_test_memory_obj()
|
|
mock_backend.storage[test_key.chunk_hash] = test_memory_obj
|
|
|
|
# Clear the call history before the test
|
|
mock_backend.call_history.clear()
|
|
|
|
# Call get via storage_manager
|
|
# This should iterate through backends and call get_blocking
|
|
result = storage_manager.get(test_key)
|
|
|
|
# Verify get_blocking was called on our dynamic backend
|
|
assert any(
|
|
"get_blocking:67890" in call for call in mock_backend.call_history
|
|
), (
|
|
f"Expected 'get_blocking:67890' in call history, "
|
|
f"got: {mock_backend.call_history}"
|
|
)
|
|
|
|
# The result should be our test memory object
|
|
assert result is test_memory_obj, f"Expected test_memory_obj, got: {result}"
|
|
|
|
def test_dynamic_backend_remove_called_via_storage_manager(
|
|
self, storage_manager_with_mock_backend
|
|
):
|
|
"""
|
|
Test that StorageManager can remove data through the dynamic backend.
|
|
"""
|
|
storage_manager = storage_manager_with_mock_backend
|
|
|
|
# Get the mock backend instance
|
|
mock_backend = get_mock_backend(storage_manager)
|
|
assert mock_backend is not None, "Expected mock_backend to be loaded"
|
|
|
|
# Pre-populate storage with a test memory object
|
|
test_key = create_test_key(key_id=11111)
|
|
test_memory_obj = create_test_memory_obj()
|
|
mock_backend.storage[test_key.chunk_hash] = test_memory_obj
|
|
|
|
# Clear the call history before the test
|
|
mock_backend.call_history.clear()
|
|
|
|
# Call remove via storage_manager
|
|
num_removed = storage_manager.remove(test_key)
|
|
|
|
# Verify remove was called on our dynamic backend
|
|
assert any("remove:11111" in call for call in mock_backend.call_history), (
|
|
f"Expected 'remove:11111' in call history, got: {mock_backend.call_history}"
|
|
)
|
|
|
|
# Should have removed from at least one backend
|
|
assert num_removed >= 1, f"Expected at least 1 removal, got: {num_removed}"
|
|
|
|
# Verify the key is no longer in storage
|
|
assert test_key.chunk_hash not in mock_backend.storage, (
|
|
f"Expected key {test_key.chunk_hash} to be removed from storage"
|
|
)
|
|
|
|
def test_dynamic_backend_batched_contains_called_via_storage_manager(
|
|
self, storage_manager_with_mock_backend
|
|
):
|
|
"""
|
|
Test that StorageManager calls batched_contains on the dynamic backend.
|
|
"""
|
|
storage_manager = storage_manager_with_mock_backend
|
|
|
|
# Get the mock backend instance
|
|
mock_backend = get_mock_backend(storage_manager)
|
|
assert mock_backend is not None, "Expected mock_backend to be loaded"
|
|
|
|
# Pre-populate storage with test memory objects
|
|
test_keys = [create_test_key(key_id=i) for i in range(1000, 1003)]
|
|
for key in test_keys:
|
|
mock_backend.storage[key.chunk_hash] = create_test_memory_obj()
|
|
|
|
# Clear the call history before the test
|
|
mock_backend.call_history.clear()
|
|
|
|
# Call batched_contains via storage_manager
|
|
hit_chunks, block_mapping = storage_manager.batched_contains(test_keys)
|
|
|
|
# Verify contains was called for the keys
|
|
# (batched_contains internally calls contains for each key)
|
|
contains_calls = [
|
|
call for call in mock_backend.call_history if call.startswith("contains:")
|
|
]
|
|
assert len(contains_calls) >= 1, (
|
|
f"Expected at least 1 contains call, got: {mock_backend.call_history}"
|
|
)
|
|
|
|
# Should have found some hits in the mock backend
|
|
assert hit_chunks >= 1, f"Expected at least 1 hit, got: {hit_chunks}"
|
|
|
|
def test_dynamic_backend_without_storage_plugins_config(self, async_loop):
|
|
"""
|
|
Test that no dynamic backend is loaded when storage_plugins is not configured.
|
|
"""
|
|
# Configure extra_config but don't set storage_plugins
|
|
config = create_test_config(
|
|
extra_config=MOCK_BACKEND_EXTRA_CONFIG,
|
|
storage_plugins=None, # No storage_plugins
|
|
)
|
|
metadata = create_test_metadata()
|
|
|
|
# Create storage backends
|
|
storage_backends = CreateStorageBackends(
|
|
config=config,
|
|
metadata=metadata,
|
|
loop=async_loop,
|
|
dst_device="cpu",
|
|
)
|
|
|
|
# Verify no mock_backend is loaded
|
|
assert "mock_backend" not in storage_backends, (
|
|
f"Expected 'mock_backend' not to be in storage_backends, "
|
|
f"got: {list(storage_backends.keys())}"
|
|
)
|
|
|
|
# Close all backends
|
|
for backend in storage_backends.values():
|
|
backend.close()
|
|
|
|
def test_dynamic_backend_with_invalid_module_path(self, async_loop):
|
|
"""
|
|
Test that invalid module path is handled gracefully.
|
|
"""
|
|
# Configure with invalid module path
|
|
extra_config = {
|
|
"storage_plugin.invalid_backend.module_path": ("nonexistent.module.path"),
|
|
"storage_plugin.invalid_backend.class_name": "NonexistentClass",
|
|
}
|
|
storage_plugins = ["invalid_backend"]
|
|
|
|
config = create_test_config(
|
|
extra_config=extra_config,
|
|
storage_plugins=storage_plugins,
|
|
)
|
|
metadata = create_test_metadata()
|
|
|
|
# Create storage backends - should not raise exception
|
|
storage_backends = CreateStorageBackends(
|
|
config=config,
|
|
metadata=metadata,
|
|
loop=async_loop,
|
|
dst_device="cpu",
|
|
)
|
|
|
|
# Invalid backend should not be loaded
|
|
assert "invalid_backend" not in storage_backends, (
|
|
f"Expected 'invalid_backend' not to be in storage_backends, "
|
|
f"got: {list(storage_backends.keys())}"
|
|
)
|
|
|
|
# Close all backends
|
|
for backend in storage_backends.values():
|
|
backend.close()
|