Perf: Enhance timeout handling. (#8826)

### What problem does this PR solve?


### Type of change

- [x] Performance Improvement

rag/raptor.py

@@ -20,6 +20,7 @@ import numpy as np
 from sklearn.mixture import GaussianMixture
 import trio
 
+from api.utils.api_utils import timeout
 from graphrag.utils import (
     get_llm_cache,
     get_embed_cache,
@@ -54,6 +55,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
         set_llm_cache(self._llm_model.llm_name, system, response, history, gen_conf)
         return response
 
+    @timeout(2)
     async def _embedding_encode(self, txt):
         response = get_embed_cache(self._embd_model.llm_name, txt)
         if response is not None:
@@ -83,6 +85,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
         layers = [(0, len(chunks))]
         start, end = 0, len(chunks)
 
+        @timeout(60)
         async def summarize(ck_idx: list[int]):
             nonlocal chunks
             texts = [chunks[i][0] for i in ck_idx]

rag/svr/task_executor.py

@@ -21,6 +21,7 @@ import sys
 import threading
 import time
 
+from api.utils.api_utils import timeout
 from api.utils.log_utils import init_root_logger, get_project_base_directory
 from graphrag.general.index import run_graphrag
 from graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache
@@ -275,6 +276,7 @@ async def build_chunks(task, progress_callback):
         doc[PAGERANK_FLD] = int(task["pagerank"])
     st = timer()
 
+    @timeout(60)
     async def upload_to_minio(document, chunk):
         try:
             d = copy.deepcopy(document)
@@ -415,6 +417,7 @@ def init_kb(row, vector_size: int):
     return settings.docStoreConn.createIdx(idxnm, row.get("kb_id", ""), vector_size)
 
 
+@timeout(60*20)
 async def embedding(docs, mdl, parser_config=None, callback=None):
     if parser_config is None:
         parser_config = {}
@@ -461,6 +464,7 @@ async def embedding(docs, mdl, parser_config=None, callback=None):
     return tk_count, vector_size
 
 
+@timeout(3600)
 async def run_raptor(row, chat_mdl, embd_mdl, vector_size, callback=None):
     chunks = []
     vctr_nm = "q_%d_vec"%vector_size
@@ -502,6 +506,7 @@ async def run_raptor(row, chat_mdl, embd_mdl, vector_size, callback=None):
     return res, tk_count
 
 
+@timeout(60*60*1.5)
 async def do_handle_task(task):
     task_id = task["id"]
     task_from_page = task["from_page"]

rag/utils/redis_conn.py

@@ -220,40 +220,43 @@ class RedisDB:
                 logging.exception(
                     "RedisDB.queue_product " + str(queue) + " got exception: " + str(e)
                 )
+                self.__open__()
         return False
 
     def queue_consumer(self, queue_name, group_name, consumer_name, msg_id=b">") -> RedisMsg:
         """https://redis.io/docs/latest/commands/xreadgroup/"""
-        try:
-            group_info = self.REDIS.xinfo_groups(queue_name)
-            if not any(gi["name"] == group_name for gi in group_info):
-                self.REDIS.xgroup_create(queue_name, group_name, id="0", mkstream=True)
-            args = {
-                "groupname": group_name,
-                "consumername": consumer_name,
-                "count": 1,
-                "block": 5,
-                "streams": {queue_name: msg_id},
-            }
-            messages = self.REDIS.xreadgroup(**args)
-            if not messages:
-                return None
-            stream, element_list = messages[0]
-            if not element_list:
-                return None
-            msg_id, payload = element_list[0]
-            res = RedisMsg(self.REDIS, queue_name, group_name, msg_id, payload)
-            return res
-        except Exception as e:
-            if str(e) == 'no such key':
-                pass
-            else:
-                logging.exception(
-                    "RedisDB.queue_consumer "
-                    + str(queue_name)
-                    + " got exception: "
-                    + str(e)
-                )
+        for _ in range(3):
+            try:
+                group_info = self.REDIS.xinfo_groups(queue_name)
+                if not any(gi["name"] == group_name for gi in group_info):
+                    self.REDIS.xgroup_create(queue_name, group_name, id="0", mkstream=True)
+                args = {
+                    "groupname": group_name,
+                    "consumername": consumer_name,
+                    "count": 1,
+                    "block": 5,
+                    "streams": {queue_name: msg_id},
+                }
+                messages = self.REDIS.xreadgroup(**args)
+                if not messages:
+                    return None
+                stream, element_list = messages[0]
+                if not element_list:
+                    return None
+                msg_id, payload = element_list[0]
+                res = RedisMsg(self.REDIS, queue_name, group_name, msg_id, payload)
+                return res
+            except Exception as e:
+                if str(e) == 'no such key':
+                    pass
+                else:
+                    logging.exception(
+                        "RedisDB.queue_consumer "
+                        + str(queue_name)
+                        + " got exception: "
+                        + str(e)
+                    )
+                    self.__open__()
         return None
 
     def get_unacked_iterator(self, queue_names: list[str], group_name, consumer_name):
@@ -294,26 +297,30 @@ class RedisDB:
         return []
 
     def requeue_msg(self, queue: str, group_name: str, msg_id: str):
-        try:
-            messages = self.REDIS.xrange(queue, msg_id, msg_id)
-            if messages:
-                self.REDIS.xadd(queue, messages[0][1])
-                self.REDIS.xack(queue, group_name, msg_id)
-        except Exception as e:
-            logging.warning(
-                "RedisDB.get_pending_msg " + str(queue) + " got exception: " + str(e)
-            )
+        for _ in range(3):
+            try:
+                messages = self.REDIS.xrange(queue, msg_id, msg_id)
+                if messages:
+                    self.REDIS.xadd(queue, messages[0][1])
+                    self.REDIS.xack(queue, group_name, msg_id)
+            except Exception as e:
+                logging.warning(
+                    "RedisDB.get_pending_msg " + str(queue) + " got exception: " + str(e)
+                )
+                self.__open__()
 
     def queue_info(self, queue, group_name) -> dict | None:
-        try:
-            groups = self.REDIS.xinfo_groups(queue)
-            for group in groups:
-                if group["name"] == group_name:
-                    return group
-        except Exception as e:
-            logging.warning(
-                "RedisDB.queue_info " + str(queue) + " got exception: " + str(e)
-            )
+        for _ in range(3):
+            try:
+                groups = self.REDIS.xinfo_groups(queue)
+                for group in groups:
+                    if group["name"] == group_name:
+                        return group
+            except Exception as e:
+                logging.warning(
+                    "RedisDB.queue_info " + str(queue) + " got exception: " + str(e)
+                )
+                self.__open__()
         return None
 
     def delete_if_equal(self, key: str, expected_value: str) -> bool: