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

Code before:
        """
        Allocate a memory object of shape and dtype
        evict if necessary. Storage manager should always call
        local_cpu_backend.allocate() to get memory objects
        regardless of whether local_cpu is True or False
        """
        logger.debug(
            f"Allocating memory in local cpu backend with busy loop: {busy_loop}"
        )
        if fmt is None:
            if self.layerwise:
                if self.enable_blending:
                    fmt = MemoryFormat.KV_2TD
                else:
                    fmt = MemoryFormat.KV_T2D
            else:
                fmt = MemoryFormat.KV_2LTD

        memory_obj = self.memory_allocator.allocate(shape, dtype, fmt)
        if memory_obj is not None or not eviction:
            return memory_obj

        assert (
            isinstance(self.memory_allocator, MixedMemoryAllocator)
            or isinstance(self.memory_allocator, NixlCPUMemoryAllocator)
            or isinstance(self.memory_allocator, CuFileMemoryAllocator)
        )

        evict_keys_count = 0
        num_attempts = 0
        while True:
            # whether or not this request needs to wait or other requests
            wait_other_requests = True
            if self.use_hot:
                # TODO(Jiayi): optimize `num_candidates` with estimation.
                # Accurate estimation is hard due to fragmentation
                num_candidates = 1
                evict_keys = None
                with self.cpu_lock:
                    evict_keys = self.cache_policy.get_evict_candidates(
                        self.hot_cache, num_candidates=num_candidates
                    )
                    if evict_keys:
                        # we can continue trying to evict from the hot_cache
                        # and don't need to wait for other requests yet
                        wait_other_requests = False
                        logger.debug(
                            f"Evicting {len(evict_keys)} chunks from cpu memory"
                        )
                        # remove
                        self.batched_remove(evict_keys, force=False)

Review comment (suggestion):
Here's a bit weird. Why does the local CPU backend need to handle `CuFileMemoryAllocator`? This should be for gds backend only.

IIUC, the storage manager should have the correct "allocator backend" to allocate and evict the memory objects. The allocator backend should be gds backend when using gds, which will handle CuFileMemoryAllocator by default.

cc @YaoJiayi @sammshen 

Code after:
            self.lmcache_worker.put_msg(
                KVEvictMsg(self.instance_id, key.worker_id, key.chunk_hash, str(self))
            )
        # NOTE (Jiayi): This `return True` might not accurately reflect
        # whether the key is removed from the actual memory because
        # other backends might still (temporarily) hold the memory object.
        return True

    def initialize_allocator(
        self,
        config: LMCacheEngineConfig,
        metadata: Optional[LMCacheEngineMetadata] = None,
    ) -> MemoryAllocatorInterface:
        cpu_size = config.max_local_cpu_size

        if metadata is not None:
            # save_only_first_rank only works when use mla
            save_only_first_rank = (
                config.get_extra_config_value("save_only_first_rank", metadata.use_mla)
                and metadata.use_mla
            )

            if save_only_first_rank and metadata.is_first_rank():
                # Only the first rank will save the cache,
                # so we need to set it lager than other ranks
                cpu_size = (
                    config.extra_config.get("first_rank_max_local_cpu_size", cpu_size)
                    if config.extra_config
                    else cpu_size
                )

        # Detect the numa mapping
        numa_mapping = NUMADetector.get_numa_mapping(config)
        logger.info(f"NUMA mapping {numa_mapping}")

        if config.enable_p2p:
            assert metadata is not None
            meta_shape = torch.Size(metadata.kv_shape)
            # TODO(Jiayi): remove this hardcode
            new_shape = torch.Size(
                [
                    meta_shape[1],
                    meta_shape[0],
                    meta_shape[2],
                    meta_shape[3] * meta_shape[4],
                ]
            )
            paged_mem_allocator = PagedCpuGpuMemoryAllocator()
            paged_mem_allocator.init_cpu_memory_allocator(
                int(cpu_size * 1024**3),
                shape=new_shape,