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

Code before:

        :param dst_device: the device where the blocking retrieved KV is stored,
            could be either "cpu", "cuda", or "cuda:0", "cuda:1", etc.

        :raise: RuntimeError if the device is not valid
        """
        try:
            torch.device(dst_device)
        except RuntimeError:
            raise

        self.dst_device = dst_device
        self._listener: Optional[StorageBackendListener] = None

    def set_listener(self, listener: StorageBackendListener):
        """
        Set the listener to receive events.
        """
        self._listener = listener

    def _on_evict(self, keys: List[CacheEngineKey]) -> None:
        """
        Evict keys from the storage backend.
        """
        if hasattr(self, "_listener") and self._listener is not None:
            self._listener.on_evict(self, keys)

    @abc.abstractmethod
    def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
        """
        Check whether key is in the storage backend.

        :param CacheEngineKey key: The key of the MemoryObj.

        :param bool pin: Whether to pin the key.
            If True, the corresponding KV cache will be
            pinned in the storage backend.

        :return: True if the key exists, False otherwise.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def exists_in_put_tasks(self, key: CacheEngineKey) -> bool:
        """
        Check whether key is in the ongoing put tasks.
        """
        raise NotImplementedError

    # NOTE (Jiayi): Using batched interface allows the underlying implementation
    # have more flexibility to do optimizations.

Review comment (suggestion):
How about we introduce a new method `get_listener(self) -> Optional[StorageBackendListener]` so that we don't need to use `hasattr` here.

Can do something like:
```python
if listener := self.get_listener():
    listener.on_evict(self, keys)
```

Code after:

        :param dst_device: the device where the blocking retrieved KV is stored,
            could be either "cpu", "cuda", or "cuda:0", "cuda:1", etc.

        :raise: RuntimeError if the device is not valid
        """
        try:
            torch.device(dst_device)
        except RuntimeError:
            raise

        self.dst_device = dst_device
        self._listener: Optional[StorageBackendListener] = None

    def set_listener(self, listener: StorageBackendListener):
        """
        Set the listener to receive events.
        """
        self._listener = listener

    def _on_evict(self, keys: List[CacheEngineKey]) -> None:
        """
        Evict keys from the storage backend.
        """
        if self._listener is not None:
            self._listener.on_evict(self, keys)

    @abc.abstractmethod
    def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
        """
        Check whether key is in the storage backend.

        :param CacheEngineKey key: The key of the MemoryObj.

        :param bool pin: Whether to pin the key.
            If True, the corresponding KV cache will be
            pinned in the storage backend.

        :return: True if the key exists, False otherwise.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def exists_in_put_tasks(self, key: CacheEngineKey) -> bool:
        """
        Check whether key is in the ongoing put tasks.
        """
        raise NotImplementedError

    # NOTE (Jiayi): Using batched interface allows the underlying implementation
    # have more flexibility to do optimizations.