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

Code before:

        # In some scenarios, lookup is called first, and then the original tokens
        # is sliced based on the lookup result. In these scenarios, the tokens
        # passed in must exist in LMCache, and we can set skip_contains_check to True.
        # When skip_contains_check is True and there is only one backend, the `contains`
        # call can be skipped.
        skip_contains_check = (
            kwargs["skip_contains_check"] if "skip_contains_check" in kwargs else False
        )
        for start, end, key in self.token_database.process_tokens(
            tokens=tokens,
            mask=mask,
            request_configs=request_configs,
        ):
            assert isinstance(key, CacheEngineKey)

            location = None
            if key in self.lookup_cache:
                # TODO(Jiayi): we can reduce the number of `contains` calls
                # by checking the lookup cache first (should be updated in `lookup`)
                pass
            else:
                # NOTE: key should always be in the lookup cache once we support it.
                if (
                    skip_contains_check
                    and len(self.storage_manager.actual_storage_names) == 1
                ):
                    location = self.storage_manager.actual_storage_names[0]
                else:
                    location = self.storage_manager.contains(key)
                if location is None:
                    break

                # NOTE: Here we make the assumption that the underlying
                # storage backend support pin operation, and the memory
                # object is already pinned in the storage backend.
                ret_mask[start:end] = True

            assert location is not None

            block_mapping[location].append((key, start, end))

        last_failed_block_start = None
        for location, blocks in block_mapping.items():
            keys = [key for key, _, _ in blocks]
            memory_objs = self.storage_manager.batched_get(
                keys=keys,
                location=location,
            )
            assert memory_objs is not None, (
                "Failed to get memory objects from storage backend"

Review comment (suggestion):
in the future, to avoid this, we can store the location on the first contains in `get_num_new_matched_tokens()` (which should at least pin the local cpu and the disk) and then always skip this contains call. WDYT @maobaolong @chunxiaozheng

Code after:
        """

        tot_kv_size = 0
        # location -> [(CacheEngineKey, start, end)]
        block_mapping: dict[str, list[tuple[CacheEngineKey, int, int]]] = defaultdict(
            list
        )

        reordered_chunks: list[tuple[CacheEngineKey, MemoryObj, int, int]] = []

        request_configs = kwargs.get("request_configs")
        if request_configs is not None and len(request_configs) != 0:
            assert isinstance(request_configs, dict)

        # In some scenarios, lookup is called first, and then the original tokens
        # is sliced based on the lookup result. In these scenarios, the tokens
        # passed in must exist in LMCache, and we can set skip_contains_check to True.
        # When skip_contains_check is True and there is only one backend, the `contains`
        # call can be skipped.
        skip_contains_check = (
            kwargs["skip_contains_check"] if "skip_contains_check" in kwargs else False
        )
        for start, end, key in self.token_database.process_tokens(
            tokens=tokens,
            mask=mask,
            request_configs=request_configs,
        ):
            assert isinstance(key, CacheEngineKey)

            location = None
            if key in self.lookup_cache:
                # TODO(Jiayi): we can reduce the number of `contains` calls
                # by checking the lookup cache first (should be updated in `lookup`)
                pass
            else:
                # NOTE: key should always be in the lookup cache once we support it.
                # TODO: use lookup_cache to skip the contains
                if (
                    skip_contains_check
                    and len(self.storage_manager.non_allocator_backends) == 1
                ):
                    location = self.storage_manager.non_allocator_backends[0]
                else:
                    location = self.storage_manager.contains(key)
                if location is None:
                    break

                # NOTE: Here we make the assumption that the underlying
                # storage backend support pin operation, and the memory
                # object is already pinned in the storage backend.
                ret_mask[start:end] = True