Fix 'TOO MANY OPEN FILE' problem while using redis vector DB:

Enhance RedisKVStorage: Implement connection pooling and error handling. Refactor async methods to use context managers for Redis operations, improving resource management and error logging. Batch processing added for key operations to optimize performance.

lightrag/kg/redis_impl.py

@@ -3,12 +3,14 @@ from typing import Any, final
 from dataclasses import dataclass
 import pipmaster as pm
 import configparser
+from contextlib import asynccontextmanager
 
 if not pm.is_installed("redis"):
     pm.install("redis")
 
 # aioredis is a depricated library, replaced with redis
-from redis.asyncio import Redis
+from redis.asyncio import Redis, ConnectionPool
+from redis.exceptions import RedisError, ConnectionError
 from lightrag.utils import logger, compute_mdhash_id
 from lightrag.base import BaseKVStorage
 import json
@@ -17,6 +19,11 @@ import json
 config = configparser.ConfigParser()
 config.read("config.ini", "utf-8")
 
+# Constants for Redis connection pool
+MAX_CONNECTIONS = 50
+SOCKET_TIMEOUT = 5.0
+SOCKET_CONNECT_TIMEOUT = 3.0
+
 
 @final
 @dataclass
@@ -25,22 +32,71 @@ class RedisKVStorage(BaseKVStorage):
         redis_url = os.environ.get(
             "REDIS_URI", config.get("redis", "uri", fallback="redis://localhost:6379")
         )
-        self._redis = Redis.from_url(redis_url, decode_responses=True)
+        # Create a connection pool with limits
-        logger.info(f"Use Redis as KV {self.namespace}")
+        self._pool = ConnectionPool.from_url(
+            redis_url,
+            max_connections=MAX_CONNECTIONS,
+            decode_responses=True,
+            socket_timeout=SOCKET_TIMEOUT,
+            socket_connect_timeout=SOCKET_CONNECT_TIMEOUT
+        )
+        self._redis = Redis(connection_pool=self._pool)
+        logger.info(f"Initialized Redis connection pool for {self.namespace} with max {MAX_CONNECTIONS} connections")
+
+    @asynccontextmanager
+    async def _get_redis_connection(self):
+        """Safe context manager for Redis operations."""
+        try:
+            yield self._redis
+        except ConnectionError as e:
+            logger.error(f"Redis connection error in {self.namespace}: {e}")
+            raise
+        except RedisError as e:
+            logger.error(f"Redis operation error in {self.namespace}: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error in Redis operation for {self.namespace}: {e}")
+            raise
+
+    async def close(self):
+        """Close the Redis connection pool to prevent resource leaks."""
+        if hasattr(self, '_redis') and self._redis:
+            await self._redis.close()
+            await self._pool.disconnect()
+            logger.debug(f"Closed Redis connection pool for {self.namespace}")
+
+    async def __aenter__(self):
+        """Support for async context manager."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Ensure Redis resources are cleaned up when exiting context."""
+        await self.close()
 
     async def get_by_id(self, id: str) -> dict[str, Any] | None:
-        data = await self._redis.get(f"{self.namespace}:{id}")
+        async with self._get_redis_connection() as redis:
+            try:
+                data = await redis.get(f"{self.namespace}:{id}")
                 return json.loads(data) if data else None
+            except json.JSONDecodeError as e:
+                logger.error(f"JSON decode error for id {id}: {e}")
+                return None
 
     async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
-        pipe = self._redis.pipeline()
+        async with self._get_redis_connection() as redis:
+            try:
+                pipe = redis.pipeline()
                 for id in ids:
                     pipe.get(f"{self.namespace}:{id}")
                 results = await pipe.execute()
                 return [json.loads(result) if result else None for result in results]
+            except json.JSONDecodeError as e:
+                logger.error(f"JSON decode error in batch get: {e}")
+                return [None] * len(ids)
 
     async def filter_keys(self, keys: set[str]) -> set[str]:
-        pipe = self._redis.pipeline()
+        async with self._get_redis_connection() as redis:
+            pipe = redis.pipeline()
             for key in keys:
                 pipe.exists(f"{self.namespace}:{key}")
             results = await pipe.execute()
@@ -49,32 +105,30 @@ class RedisKVStorage(BaseKVStorage):
             return set(keys) - existing_ids
 
     async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
-        logger.info(f"Inserting {len(data)} to {self.namespace}")
         if not data:
             return
-        pipe = self._redis.pipeline()
             
+        logger.info(f"Inserting {len(data)} items to {self.namespace}")
+        async with self._get_redis_connection() as redis:
+            try:
+                pipe = redis.pipeline()
                 for k, v in data.items():
                     pipe.set(f"{self.namespace}:{k}", json.dumps(v))
                 await pipe.execute()
 
                 for k in data:
                     data[k]["_id"] = k
-
+            except json.JSONEncodeError as e:
-    async def index_done_callback(self) -> None:
+                logger.error(f"JSON encode error during upsert: {e}")
-        # Redis handles persistence automatically
+                raise
-        pass
 
     async def delete(self, ids: list[str]) -> None:
-        """Delete entries with specified IDs
+        """Delete entries with specified IDs"""
-
-        Args:
-            ids: List of entry IDs to be deleted
-        """
         if not ids:
             return
 
-        pipe = self._redis.pipeline()
+        async with self._get_redis_connection() as redis:
+            pipe = redis.pipeline()
             for id in ids:
                 pipe.delete(f"{self.namespace}:{id}")
 
@@ -85,20 +139,15 @@ class RedisKVStorage(BaseKVStorage):
             )
 
     async def delete_entity(self, entity_name: str) -> None:
-        """Delete an entity by name
+        """Delete an entity by name"""
-
-        Args:
-            entity_name: Name of the entity to delete
-        """
-
         try:
             entity_id = compute_mdhash_id(entity_name, prefix="ent-")
             logger.debug(
                 f"Attempting to delete entity {entity_name} with ID {entity_id}"
             )
 
-            # Delete the entity
+            async with self._get_redis_connection() as redis:
-            result = await self._redis.delete(f"{self.namespace}:{entity_id}")
+                result = await redis.delete(f"{self.namespace}:{entity_id}")
 
                 if result:
                     logger.debug(f"Successfully deleted entity {entity_name}")
@@ -108,42 +157,52 @@ class RedisKVStorage(BaseKVStorage):
             logger.error(f"Error deleting entity {entity_name}: {e}")
 
     async def delete_entity_relation(self, entity_name: str) -> None:
-        """Delete all relations associated with an entity
+        """Delete all relations associated with an entity"""
-
-        Args:
-            entity_name: Name of the entity whose relations should be deleted
-        """
         try:
-            # Get all keys in this namespace
+            async with self._get_redis_connection() as redis:
                 cursor = 0
                 relation_keys = []
                 pattern = f"{self.namespace}:*"
 
                 while True:
-                cursor, keys = await self._redis.scan(cursor, match=pattern)
+                    cursor, keys = await redis.scan(cursor, match=pattern)
                     
-                # For each key, get the value and check if it's related to entity_name
+                    # Process keys in batches
+                    pipe = redis.pipeline()
                     for key in keys:
-                    value = await self._redis.get(key)
+                        pipe.get(key)
+                    values = await pipe.execute()
+                    
+                    for key, value in zip(keys, values):
                         if value:
+                            try:
                                 data = json.loads(value)
-                        # Check if this is a relation involving the entity
                                 if (
                                     data.get("src_id") == entity_name
                                     or data.get("tgt_id") == entity_name
                                 ):
                                     relation_keys.append(key)
+                            except json.JSONDecodeError:
+                                logger.warning(f"Invalid JSON in key {key}")
+                                continue
 
-                # Exit loop when cursor returns to 0
                     if cursor == 0:
                         break
 
-            # Delete the relation keys
+                # Delete relations in batches
                 if relation_keys:
-                deleted = await self._redis.delete(*relation_keys)
+                    # Delete in chunks to avoid too many arguments
-                logger.debug(f"Deleted {deleted} relations for {entity_name}")
+                    chunk_size = 1000
+                    for i in range(0, len(relation_keys), chunk_size):
+                        chunk = relation_keys[i:i + chunk_size]
+                        deleted = await redis.delete(*chunk)
+                        logger.debug(f"Deleted {deleted} relations for {entity_name} (batch {i//chunk_size + 1})")
                 else:
                     logger.debug(f"No relations found for entity {entity_name}")
 
         except Exception as e:
             logger.error(f"Error deleting relations for {entity_name}: {e}")
+
+    async def index_done_callback(self) -> None:
+        # Redis handles persistence automatically
+        pass

lightrag/operate.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import traceback
 import json
 import re
 import os
@@ -994,6 +995,7 @@ async def mix_kg_vector_query(
 
         except Exception as e:
             logger.error(f"Error in get_kg_context: {str(e)}")
+            traceback.print_exc()
             return None
 
     async def get_vector_context():
@@ -1382,9 +1384,16 @@ async def _find_most_related_text_unit_from_entities(
                 all_text_units_lookup[c_id] = index
                 tasks.append((c_id, index, this_edges))
 
-    results = await asyncio.gather(
+    # Process in batches of 25 tasks at a time to avoid overwhelming resources
-        *[text_chunks_db.get_by_id(c_id) for c_id, _, _ in tasks]
+    batch_size = 25
+    results = []
+    
+    for i in range(0, len(tasks), batch_size):
+        batch_tasks = tasks[i:i + batch_size]
+        batch_results = await asyncio.gather(
+            *[text_chunks_db.get_by_id(c_id) for c_id, _, _ in batch_tasks]
         )
+        results.extend(batch_results)
 
     for (c_id, index, this_edges), data in zip(tasks, results):
         all_text_units_lookup[c_id] = {