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

Code before:

        self.manager_lock = threading.Lock()

        self.lmcache_worker = lmcache_worker
        self.instance_id = config.lmcache_instance_id
        self.worker_id = metadata.worker_id

        self.event_manager = event_manager

        self.async_lookup_server: Optional["LMCacheAsyncLookupServer"] = None

        # The cuda stream for internal copies during put
        if torch.cuda.is_available():
            self.internal_copy_stream = torch.cuda.Stream()
        else:
            self.internal_copy_stream = None

    def post_init(self, **kwargs) -> None:
        if "async_lookup_server" in kwargs:
            assert not self.config.save_unfull_chunk, (
                "save_unfull_chunk should be automatically set to False when using "
                "async loading."
            )
            self.async_lookup_server = kwargs.pop("async_lookup_server")
            self.async_serializer = AsyncSerializer(self.allocator_backend, self.loop)

    def _get_allocator_backend(
        self, config: LMCacheEngineConfig
    ) -> AllocatorBackendInterface:
        if self.enable_pd:
            allocator_backend = self.storage_backends["PDBackend"]
        else:
            allocator_backend = self.storage_backends["LocalCPUBackend"]
        assert isinstance(allocator_backend, AllocatorBackendInterface)
        return allocator_backend

    @_lmcache_nvtx_annotate
    def allocate(
        self,
        shape: torch.Size,
        dtype: torch.dtype,
        fmt: MemoryFormat = MemoryFormat.KV_2LTD,
        eviction=True,
        busy_loop=True,
    ) -> Optional[MemoryObj]:
        """
        Allocate memory object with memory allocator.
        Use LRU evictor if eviction is enabled.
        """
        # TODO (Jiayi): We might need to pre-allocate and management
        # disk in a similar way as CPU.

Review comment (suggestion):
<img width="2214" height="1258" alt="Clipboard_Screenshot_1761033530" src="https://github.com/user-attachments/assets/3e232de0-42d9-452e-ad8a-d86bdd0790fe" />

This will introduce amount of logs for each start_load_kv, especially in the decode node.

Code after:

        self.manager_lock = threading.Lock()

        self.lmcache_worker = lmcache_worker
        self.instance_id = config.lmcache_instance_id
        self.worker_id = metadata.worker_id

        self.event_manager = event_manager

        self.async_lookup_server: Optional["LMCacheAsyncLookupServer"] = None
        self.async_serializer: Optional[AsyncSerializer] = None

        # The cuda stream for internal copies during put
        if torch.cuda.is_available():
            self.internal_copy_stream = torch.cuda.Stream()
        else:
            self.internal_copy_stream = None

    def post_init(self, **kwargs) -> None:
        if "async_lookup_server" in kwargs:
            assert not self.config.save_unfull_chunk, (
                "save_unfull_chunk should be automatically set to False when using "
                "async loading."
            )
            self.async_lookup_server = kwargs.pop("async_lookup_server")
        self.async_serializer = AsyncSerializer(self.allocator_backend, self.loop)

    def _get_allocator_backend(
        self, config: LMCacheEngineConfig
    ) -> AllocatorBackendInterface:
        if self.enable_pd:
            allocator_backend = self.storage_backends["PDBackend"]
        else:
            allocator_backend = self.storage_backends["LocalCPUBackend"]
        assert isinstance(allocator_backend, AllocatorBackendInterface)
        return allocator_backend

    @_lmcache_nvtx_annotate
    def allocate(
        self,
        shape: torch.Size,
        dtype: torch.dtype,
        fmt: MemoryFormat = MemoryFormat.KV_2LTD,
        eviction=True,
        busy_loop=True,
    ) -> Optional[MemoryObj]:
        """
        Allocate memory object with memory allocator.
        Use LRU evictor if eviction is enabled.
        """
        # TODO (Jiayi): We might need to pre-allocate and management