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

Code before:
        """

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

        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)

        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 (bug):
Is it possible that calling storage_to_mem with a list of keys will expose a potential deadlock?

https://github.com/LMCache/LMCache/blob/dev/lmcache/v1/storage_backend/nixl_storage_backend.py#L465

Performing the memory allocation while holding the key_lock seems risky - even if the allocation is not supposed to due to slab allocation.  I noticed that the weka backend takes a different approach to collect all the metadata values associated with key each in the list and then perform an allocation (outside the lock scope) for each metadata value.

Code after:
    async def storage_to_mem(
        self, keys: list[CacheEngineKey]
    ) -> list[Optional[MemoryObj]]:
        """
        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: