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

Code before:
        if self.enable_p2p:
            self.distributed_server.close()

        if self.lmcache_worker is not None:
            self.lmcache_worker.close()

        self.storage_manager.close()

        self.memory_allocator.close()

        logger.info("LMCacheEngine closed.")

    def _process_tokens_internal(
        self,
        tokens,
        mask,
        reordered_keys,
        reordered_memory_objs,
        reordered_starts,
        reordered_ends,
        ret_mask,
    ):
        key_mapping = {}
        start_mapping = {}
        end_mapping = {}

        for start, end, key in self.token_database.process_tokens(
            tokens=tokens, mask=mask
        ):
            assert isinstance(key, CacheEngineKey)

            if key in self.lookup_cache:
                # TODO(Jiayi): we can reduce the number of `contains` calls
                # by checking the lookup cache first (should be updated in `lookup`)
                pass
            else:
                # NOTE: key should always be in the lookup cache once
                # we support it.
                location = self.storage_manager.contains(key)
                if location is None:
                    # TODO(Jiayi): Need to refactor P2P as a storage backend to
                    # clean up the following code.
                    if self.enable_p2p:
                        future_memory_obj = asyncio.run_coroutine_threadsafe(
                            self.distributed_server.issue_get(key),
                            self.distributed_loop,
                        )
                        memory_obj = future_memory_obj.result()
                        if memory_obj:
                            reordered_keys.append(key)
                            reordered_memory_objs.append(memory_obj)

Review comment (suggestion):
Can we add a docstring for this function to briefly describe what this function does, and what side-effect this function will have on the input parameters (`reordered_keys`, `reordered_memory_objs`, ...)

Code after:
    @_lmcache_nvtx_annotate
    def clear(
        self,
        tokens: Optional[Union[torch.Tensor, List[int]]] = None,
        locations: Optional[List[str]] = None,
        tags: OrderedDict = None,  # TODO: need to clean by tags
    ) -> int:
        if self.save_only_first_rank:
            if self.metadata.is_first_rank():
                num_removed = self._clear(tokens, locations)
                self.broadcast_object_fn(num_removed, self.metadata.first_rank)
                return num_removed
            else:
                num_removed = self.broadcast_object_fn(None, self.metadata.first_rank)
                return int(num_removed)
        return self._clear(tokens, locations)

    def _clear(
        self,
        tokens: Optional[Union[torch.Tensor, List[int]]] = None,
        locations: Optional[List[str]] = None,
        tags: OrderedDict = None,  # TODO: need to clean by tags
    ) -> int:
        assert isinstance(self.storage_manager, StorageManager)
        # Clear all caches if tokens is None
        if tokens is None or len(tokens) == 0:
            num_cleared = self.storage_manager.clear(locations)
            return num_cleared

        num_removed = 0
        # Only remove the caches for the given tokens
        for start, end, key in self.token_database.process_tokens(
            tokens=tokens, tags=tags
        ):
            assert isinstance(key, CacheEngineKey)
            removed = self.storage_manager.remove(key, locations)
            num_removed += removed
        return num_removed

    @_lmcache_nvtx_annotate
    def health(
        self,
    ) -> int:
        """
        Check the health of the cache engine.
        return: 0 if healthy, otherwise the error code
        """
        return 0 if self.memory_allocator.memcheck() else -1

    def close(self) -> None:
        """Close the cache engine and free all the resources"""