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

Code before:
            task = backend.get_non_blocking(key)
            if task:
                # TODO (Jiayi): add write-back logic here
                return task
        return None

    def batched_get(
        self,
        keys: List[CacheEngineKey],
        location: Optional[str] = None,
    ) -> Optional[List[Optional[MemoryObj]]]:
        """
        Blocking function to get the memory objects from the storages.
        """
        # TODO (ApostaC): remove the nested optional here
        for backend_name, storage_backend in self.get_active_storage_backends(location):
            memory_objs = storage_backend.batched_get_blocking(keys)
            if memory_objs:
                # Align with single-key `get()` logic:
                # auto-write remote data to local CPU cache
                if (
                    backend_name not in ["LocalCPUBackend", "PDBackend"]
                    and "LocalCPUBackend" in self.storage_backends
                ):
                    # Only write back if we found ALL keys (no Nones)
                    if None not in memory_objs:
                        logger.debug(
                            f"Storing {len(keys)} objects from"
                            f" {backend_name} to LocalCPUBackend"
                        )
                        local_cpu_backend = self.storage_backends["LocalCPUBackend"]
                        assert isinstance(local_cpu_backend, LocalCPUBackend)
                        # Type cast: Safe (we verified no Nones above)
                        # `batched_submit_put_task` expects list[MemoryObj]
                        # TODO (lisiG9): Refactor this write-back logic into caching
                        #  policy module
                        memory_objs_no_none = cast(List[MemoryObj], memory_objs)
                        local_cpu_backend.batched_submit_put_task(
                            keys, memory_objs_no_none
                        )
                return memory_objs
        return None

    def layerwise_batched_get(
        self,
        keys: List[List[CacheEngineKey]],
        location: Optional[str] = None,
    ) -> Generator[Future, None, None]:
        """
        Non-blocking function to get the memory objects into the storages
        in a layerwise manner.

Review comment (performance):
In practice, this could be pretty slow if `memory_objs` is a long list (it's an O(N)) operation.

One potential hacky solution is to check whether memory_objs[-1] is None. This is based on the assumption that if the returned memory objects should be something like `[obj1, obj2, ..., obj_i, None, None, None, ..., None]`. Not sure if it's true or not, though @YaoJiayi 

Code after:
            task = backend.get_non_blocking(key)
            if task:
                # TODO (Jiayi): add write-back logic here
                return task
        return None

    def batched_get(
        self,
        keys: List[CacheEngineKey],
        location: Optional[str] = None,
    ) -> Optional[List[Optional[MemoryObj]]]:
        """
        Blocking function to get the memory objects from the storages.
        """
        # TODO (ApostaC): remove the nested optional here
        for backend_name, storage_backend in self.get_active_storage_backends(location):
            memory_objs = storage_backend.batched_get_blocking(keys)
            if memory_objs:
                # Align with single-key `get()` logic:
                # auto-write remote data to local CPU cache
                if (
                    backend_name not in ["LocalCPUBackend", "PDBackend"]
                    and "LocalCPUBackend" in self.storage_backends
                    and None not in memory_objs
                ):
                    logger.debug(
                        "Storing %s objects from %s to LocalCPUBackend",
                        len(keys),
                        backend_name,
                    )
                    local_cpu_backend = self.storage_backends["LocalCPUBackend"]
                    assert isinstance(local_cpu_backend, LocalCPUBackend)
                    # Type cast: Safe (we verified no Nones above)
                    # `batched_submit_put_task` expects list[MemoryObj]
                    # TODO (lisiG9): Refactor this write-back logic into caching
                    #  policy module
                    memory_objs_no_none = cast(List[MemoryObj], memory_objs)
                    local_cpu_backend.batched_submit_put_task(keys, memory_objs_no_none)
                return memory_objs
        return None

    def layerwise_batched_get(
        self,
        keys: List[List[CacheEngineKey]],
        location: Optional[str] = None,
    ) -> Generator[Future, None, None]:
        """
        Non-blocking function to get the memory objects into the storages
        in a layerwise manner.
        Do not store if the same object is being stored (handled here by
        storage manager) or has been stored (handled by storage backend).