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

Code before:
                    "first_rank_max_local_cpu_size", cpu_size
                )

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

        # Calculate effective CPU memory size
        cpu_size = self._calculate_effective_cpu_size(cpu_size, config, metadata)

        if config.enable_p2p:
            # TODO(baoloongmao): Add lazy memory allocator support for P2P mode
            # For now, keep the original P2P implementation
            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()
            chunk_size_bytes = (
                new_shape.numel()
                * torch.tensor([], dtype=metadata.kv_dtype).element_size()
            )
            cpu_size_bytes = (
                int(cpu_size * 1024**3) // chunk_size_bytes * chunk_size_bytes
            )
            logger.info(f"Align CPU size bytes: {cpu_size_bytes}")
            paged_mem_allocator.init_cpu_memory_allocator(
                cpu_size_bytes,
                shape=new_shape,
                dtype=metadata.kv_dtype,
                fmt=MemoryFormat.KV_2LTD,  # TODO: remove this hardcode
                numa_mapping=numa_mapping,
            )
            return paged_mem_allocator
        else:
            # Check if lazy memory allocator should be enabled
            use_lazy = (
                config.enable_lazy_memory_allocator
                and cpu_size > config.lazy_memory_safe_size
            )

            if use_lazy:
                logger.info(
                    f"Using LazyMixedMemoryAllocator with "

Review comment (refactor):
You can abstract this to a method and place to `utils.py` to reduce redundant code, AFAIK, there is a same calculation in the `RemoteBackend`.

Code after:
                cpu_size = config.get_extra_config_value(
                    "first_rank_max_local_cpu_size", cpu_size
                )

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

        # Calculate effective CPU memory size
        cpu_size = self._calculate_effective_cpu_size(cpu_size, config, metadata)

        if config.enable_p2p:
            # TODO(baoloongmao): Add lazy memory allocator support for P2P mode
            # For now, keep the original P2P implementation
            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()
            chunk_size_bytes = get_size_bytes(new_shape, metadata.kv_dtype)
            origin_cpu_size_bytes = int(cpu_size * 1024**3)
            align_cpu_size_bytes = (
                origin_cpu_size_bytes // chunk_size_bytes * chunk_size_bytes
            )
            logger.info(
                f"Auto align cpu size bytes, origin: {origin_cpu_size_bytes}, "
                f"aligned: {align_cpu_size_bytes}, chunk size: {chunk_size_bytes}"
            )
            paged_mem_allocator.init_cpu_memory_allocator(
                align_cpu_size_bytes,
                shape=new_shape,
                dtype=metadata.kv_dtype,
                fmt=MemoryFormat.KV_2LTD,  # TODO: remove this hardcode
                numa_mapping=numa_mapping,
            )
            return paged_mem_allocator
        else:
            # Check if lazy memory allocator should be enabled
            use_lazy = (
                config.enable_lazy_memory_allocator
                and cpu_size > config.lazy_memory_safe_size
            )

            if use_lazy: