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

Code before:

        if future is None:
            return None

        obj_list = future.result()
        return obj_list[0]

    def batched_get_blocking(
        self,
        keys: List[CacheEngineKey],
    ) -> List[Optional[MemoryObj]]:
        """
        A blocking function to get the kv cache from the storage backend.

        :param List[CacheEngineKey] keys: The keys of the MemoryObjs.

        :return: a list of memory objects.
        """

        if not keys:
            return []

        future = asyncio.run_coroutine_threadsafe(self.storage_to_mem(keys), self.loop)

        if future is None:
            return []

        obj_list = future.result()
        return obj_list

    async def batched_get_non_blocking(
        self,
        lookup_id: str,
        keys: list[CacheEngineKey],
        transfer_spec: Any = None,
    ) -> list[MemoryObj]:
        obj_list = await self.storage_to_mem(keys)
        assert None not in obj_list
        return cast(list[MemoryObj], obj_list)

    def remove(self, key: CacheEngineKey, force: bool = True) -> bool:
        """
        Remove the key from the storage backend.

        :param key: The key to remove.
        """

        with self.key_lock:
            metadata = self.key_dict.pop(key, None)
            if metadata is None:
                return False

Review comment (performance):
I believe run_corutine_threadsafe never returns None, so this can be dropped. Might as well also remove the "future is None" case from the non-batched version too.

Code after:
        """
        Combined method: collect metadata with lock, then do NIXL transfer.
        """
        metadata_list = self._collect_metadata_with_lock(keys)
        return await self._nixl_transfer_async(metadata_list)

    def batched_submit_put_task(
        self,
        keys: Sequence[CacheEngineKey],
        memory_objs: List[MemoryObj],
        transfer_spec: Any = None,
    ) -> None:
        with self.key_lock:
            available_descs = self.pool.get_num_available_descs()
            num_evict = len(keys) - available_descs
            if num_evict > 0:
                evict_keys = self.cache_policy.get_evict_candidates(
                    self.key_dict, num_candidates=num_evict
                )

                if not evict_keys:
                    logger.warning(
                        "No eviction candidates found. Backend under pressure."
                    )
                    return None

                self.batched_remove(evict_keys, force=False)

        with self.progress_lock:
            for key in keys:
                self.progress_set.add(key)

        asyncio.run_coroutine_threadsafe(
            self.mem_to_storage(keys, memory_objs), self.loop
        )

    def get_blocking(self, key: CacheEngineKey) -> Optional[MemoryObj]:
        """
        A blocking function to get the kv cache from the storage backend.

        :param key: The key of the MemoryObj.

        :return: MemoryObj. None if the key does not exist.
        """

        future = asyncio.run_coroutine_threadsafe(self.storage_to_mem([key]), self.loop)

        if future is None:
            return None

        obj_list = future.result()