Repository: LMCache/LMCache
File: lmcache/v1/storage_backend/storage_manager.py (Python)

Code before:
from lmcache.utils import (
    CacheEngineKey,
    _lmcache_nvtx_annotate,
    start_loop_in_thread_with_exceptions,
)
from lmcache.v1.config import LMCacheEngineConfig
from lmcache.v1.lookup_server import LookupServerInterface
from lmcache.v1.memory_management import (
    MemoryAllocatorInterface,
    MemoryFormat,
    MemoryObj,
)
from lmcache.v1.storage_backend import CreateStorageBackends
from lmcache.v1.storage_backend.abstract_backend import StorageBackendInterface
from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
from lmcache.v1.storage_backend.storage_backend_listener import StorageBackendListener

if TYPE_CHECKING:
    # First Party
    from lmcache.v1.cache_controller.worker import LMCacheWorker

logger = init_logger(__name__)


# TODO: extend this class to implement caching policies and eviction policies
class StorageManager(StorageBackendListener):
    """
    The StorageManager is responsible for managing the storage backends.
    """

    def __init__(
        self,
        config: LMCacheEngineConfig,
        metadata: LMCacheEngineMetadata,
        allocator: MemoryAllocatorInterface,
        lmcache_worker: Optional["LMCacheWorker"] = None,
        lookup_server: Optional[LookupServerInterface] = None,
    ):
        self.loop = asyncio.new_event_loop()

        self.thread = threading.Thread(
            target=start_loop_in_thread_with_exceptions,
            args=(self.loop,),
            name="storage-manger-event-loop",
        )
        self.thread.start()

        dst_device = "cuda"
        self.storage_backends: OrderedDict[str, StorageBackendInterface] = (
            CreateStorageBackends(
                config,

Review comment (suggestion):
Let's try not to make `StorageManager` inherit from `StorageBackendListener`.

Alternatively, we can define a new sub class of `StorageBackendListener` (e.g., something like `EvictionListener`) that holds the reference to the storage manager object.
When the `on_evict` function of the `EvictionListener` is called, it will use that reference to access `StorageManager.contains`.

The listener can be initialized as a member in `StorageManager` at the `StorageManager.__init__` function.

Code after:
from lmcache.v1.lookup_server import LookupServerInterface
from lmcache.v1.memory_management import (
    MemoryAllocatorInterface,
    MemoryFormat,
    MemoryObj,
)
from lmcache.v1.storage_backend import CreateStorageBackends
from lmcache.v1.storage_backend.abstract_backend import (
    AllocatorBackendInterface,
    StorageBackendInterface,
)
from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
from lmcache.v1.storage_backend.storage_backend_listener import StorageBackendListener

if TYPE_CHECKING:
    # First Party
    from lmcache.v1.cache_controller.worker import LMCacheWorker

logger = init_logger(__name__)


# TODO: extend this class to implement caching policies and eviction policies
class StorageManager:
    """
    The StorageManager is responsible for managing the storage backends.
    """

    class _CPUDiskListener(StorageBackendListener):
        def __init__(self, storage_manager: "StorageManager"):
            self.storage_manager = storage_manager

        def on_evict(
            self, backend: StorageBackendInterface, keys: List[CacheEngineKey]
        ):
            """
            remove keys from lookup server only if they don't exist in local backends.

            :param StorageBackendInterface backend: The backend that evicted the keys.
            :param List[CacheEngineKey] keys: The keys that were evicted.
            """
            if self.storage_manager.lookup_server is None:
                return

            search_range = ["LocalCPUBackend", "LocalDiskBackend"]
            if str(backend) in search_range:
                search_range.remove(str(backend))

            keys_to_remove = []
            for key in keys:
                if not self.storage_manager.contains(key, search_range=search_range):
                    keys_to_remove.append(key)