From 9cb73d564b26225bd469c8a0c03c9074ec60cc77 Mon Sep 17 00:00:00 2001 From: Sarthak Deshpande Date: Mon, 25 Nov 2024 16:10:04 +0530 Subject: [PATCH] Optimized number of redis calls --- .../providers/utils/kvstore/redis/redis.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/llama_stack/providers/utils/kvstore/redis/redis.py b/llama_stack/providers/utils/kvstore/redis/redis.py index 3b6804d99..ef950a6d9 100644 --- a/llama_stack/providers/utils/kvstore/redis/redis.py +++ b/llama_stack/providers/utils/kvstore/redis/redis.py @@ -48,24 +48,27 @@ class RedisKVStoreImpl(KVStore): async def range(self, start_key: str, end_key: str) -> List[str]: start_key = self._namespaced_key(start_key) end_key = self._namespaced_key(end_key) - - result = [] cursor = 0 pattern = start_key + "*" # Match all keys starting with start_key prefix - + matching_keys = [] while True: - cursor, keys = await self.redis.scan(cursor, match=pattern) + cursor, keys = await self.redis.scan(cursor, match=pattern, count=1000) + for key in keys: key_str = key.decode("utf-8") if isinstance(key, bytes) else key if start_key <= key_str <= end_key: - value = await self.redis.get(key) - if value is not None: - value_str = ( - value.decode("utf-8") if isinstance(value, bytes) else value - ) - result.append(value_str) + matching_keys.append(key) if cursor == 0: break - return result + # Then fetch all values in a single MGET call + if matching_keys: + values = await self.redis.mget(matching_keys) + return [ + value.decode("utf-8") if isinstance(value, bytes) else value + for value in values + if value is not None + ] + + return []