handle_task catch all exception (#3441)

### What problem does this PR solve?

handle_task catch all exception
Report heartbeats

### Type of change

- [x] Refactoring

api/apps/system_app.py

@@ -13,7 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License
 #
-import json
+import logging
 from datetime import datetime
 
 from flask_login import login_required, current_user
@@ -154,26 +154,16 @@ def status():
             "error": str(e),
         }
 
+    task_executor_heartbeats = {}
     try:
-        v = REDIS_CONN.get("TASKEXE")
+        task_executors = REDIS_CONN.smembers("TASKEXE")
-        if not v:
+        now = datetime.now().timestamp()
-            raise Exception("No task executor running!")
+        for task_executor_id in task_executors:
-        obj = json.loads(v)
+            heartbeats = REDIS_CONN.zrangebyscore(task_executor_id, now - 60*30, now)
-        color = "green"
+            task_executor_heartbeats[task_executor_id] = heartbeats
-        for id in obj.keys():
+    except Exception:
-            arr = obj[id]
+        logging.exception("get task executor heartbeats failed!")
-            if len(arr) == 1:
+    res["task_executor_heartbeats"] = task_executor_heartbeats
-                obj[id] = [0]
-            else:
-                obj[id] = [arr[i + 1] - arr[i] for i in range(len(arr) - 1)]
-            elapsed = max(obj[id])
-            if elapsed > 50:
-                color = "yellow"
-            if elapsed > 120:
-                color = "red"
-        res["task_executor"] = {"status": color, "elapsed": obj}
-    except Exception as e:
-        res["task_executor"] = {"status": "red", "error": str(e)}
 
     return get_json_result(data=res)
 

api/db/services/task_service.py

@@ -36,7 +36,7 @@ class TaskService(CommonService):
 
     @classmethod
     @DB.connection_context()
-    def get_tasks(cls, task_id):
+    def get_task(cls, task_id):
         fields = [
             cls.model.id,
             cls.model.doc_id,
@@ -63,7 +63,7 @@ class TaskService(CommonService):
             .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id)) \
             .where(cls.model.id == task_id)
         docs = list(docs.dicts())
-        if not docs: return []
+        if not docs: return None
 
         msg = "\nTask has been received."
         prog = random.random() / 10.
@@ -77,9 +77,9 @@ class TaskService(CommonService):
                          ).where(
             cls.model.id == docs[0]["id"]).execute()
 
-        if docs[0]["retry_count"] >= 3: return []
+        if docs[0]["retry_count"] >= 3: return None
 
-        return docs
+        return docs[0]
 
     @classmethod
     @DB.connection_context()
@@ -108,7 +108,7 @@ class TaskService(CommonService):
             task = cls.model.get_by_id(id)
             _, doc = DocumentService.get_by_id(task.doc_id)
             return doc.run == TaskStatus.CANCEL.value or doc.progress < 0
-        except Exception as e:
+        except Exception:
             pass
         return False
 

rag/svr/task_executor.py

@@ -42,7 +42,6 @@ from multiprocessing.context import TimeoutError
 from timeit import default_timer as timer
 
 import numpy as np
-import pandas as pd
 
 from api.db import LLMType, ParserType
 from api.db.services.dialog_service import keyword_extraction, question_proposal
@@ -85,10 +84,9 @@ CONSUMER_NAME = "task_consumer_" + CONSUMER_NO
 PAYLOAD: Payload | None = None
 BOOT_AT = datetime.now().isoformat()
 DONE_TASKS = 0
-RETRY_TASKS = 0
+FAILED_TASKS = 0
 PENDING_TASKS = 0
-HEAD_CREATED_AT = ""
+LAG_TASKS = 0
-HEAD_DETAIL = ""
 
 
 def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing..."):
@@ -120,34 +118,35 @@ def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing...
 
 
 def collect():
-    global CONSUMER_NAME, PAYLOAD
+    global CONSUMER_NAME, PAYLOAD, DONE_TASKS, FAILED_TASKS
     try:
         PAYLOAD = REDIS_CONN.get_unacked_for(CONSUMER_NAME, SVR_QUEUE_NAME, "rag_flow_svr_task_broker")
         if not PAYLOAD:
             PAYLOAD = REDIS_CONN.queue_consumer(SVR_QUEUE_NAME, "rag_flow_svr_task_broker", CONSUMER_NAME)
         if not PAYLOAD:
             time.sleep(1)
-            return pd.DataFrame()
+            return None
     except Exception:
         logging.exception("Get task event from queue exception")
-        return pd.DataFrame()
+        return None
 
     msg = PAYLOAD.get_message()
     if not msg:
-        return pd.DataFrame()
+        return None
 
     if TaskService.do_cancel(msg["id"]):
+        DONE_TASKS += 1
         logging.info("Task {} has been canceled.".format(msg["id"]))
-        return pd.DataFrame()
+        return None
-    tasks = TaskService.get_tasks(msg["id"])
+    task = TaskService.get_task(msg["id"])
-    if not tasks:
+    if not task:
+        DONE_TASKS += 1
         logging.warning("{} empty task!".format(msg["id"]))
-        return []
+        return None
 
-    tasks = pd.DataFrame(tasks)
     if msg.get("type", "") == "raptor":
-        tasks["task_type"] = "raptor"
+        task["task_type"] = "raptor"
-    return tasks
+    return task
 
 
 def get_storage_binary(bucket, name):
@@ -176,14 +175,14 @@ def build(row):
         callback(-1, "Internal server error: Fetch file from minio timeout. Could you try it again.")
         logging.exception(
             "Minio {}/{} got timeout: Fetch file from minio timeout.".format(row["location"], row["name"]))
-        return
+        raise
     except Exception as e:
         if re.search("(No such file|not found)", str(e)):
             callback(-1, "Can not find file <%s> from minio. Could you try it again?" % row["name"])
         else:
             callback(-1, "Get file from minio: %s" % str(e).replace("'", ""))
         logging.exception("Chunking {}/{} got exception".format(row["location"], row["name"]))
-        return
+        raise
 
     try:
         cks = chunker.chunk(row["name"], binary=binary, from_page=row["from_page"],
@@ -194,7 +193,7 @@ def build(row):
         callback(-1, "Internal server error while chunking: %s" %
                  str(e).replace("'", ""))
         logging.exception("Chunking {}/{} got exception".format(row["location"], row["name"]))
-        return
+        raise
 
     docs = []
     doc = {
@@ -212,6 +211,7 @@ def build(row):
         d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
         d["create_timestamp_flt"] = datetime.now().timestamp()
         if not d.get("image"):
+            _ = d.pop("image", None)
             d["img_id"] = ""
             d["page_num_list"] = json.dumps([])
             d["position_list"] = json.dumps([])
@@ -232,6 +232,7 @@ def build(row):
         except Exception:
             logging.exception(
                 "Saving image of chunk {}/{}/{} got exception".format(row["location"], row["name"], d["_id"]))
+            raise
 
         d["img_id"] = "{}-{}".format(row["kb_id"], d["id"])
         del d["image"]
@@ -356,37 +357,29 @@ def run_raptor(row, chat_mdl, embd_mdl, callback=None):
     return res, tk_count, vector_size
 
 
-def main():
+def do_handle_task(r):
-    rows = collect()
-    if len(rows) == 0:
-        return
-
-    for _, r in rows.iterrows():
     callback = partial(set_progress, r["id"], r["from_page"], r["to_page"])
     try:
         embd_mdl = LLMBundle(r["tenant_id"], LLMType.EMBEDDING, llm_name=r["embd_id"], lang=r["language"])
     except Exception as e:
         callback(-1, msg=str(e))
-            logging.exception("LLMBundle got exception")
+        raise
-            continue
-
     if r.get("task_type", "") == "raptor":
         try:
             chat_mdl = LLMBundle(r["tenant_id"], LLMType.CHAT, llm_name=r["llm_id"], lang=r["language"])
             cks, tk_count, vector_size = run_raptor(r, chat_mdl, embd_mdl, callback)
         except Exception as e:
             callback(-1, msg=str(e))
-                logging.exception("run_raptor got exception")
+            raise
-                continue
     else:
         st = timer()
         cks = build(r)
         logging.info("Build chunks({}): {}".format(r["name"], timer() - st))
         if cks is None:
-                continue
+            return
         if not cks:
             callback(1., "No chunk! Done!")
-                continue
+            return
         # TODO: exception handler
         ## set_progress(r["did"], -1, "ERROR: ")
         callback(
@@ -400,9 +393,9 @@ def main():
             callback(-1, "Embedding error:{}".format(str(e)))
             logging.exception("run_rembedding got exception")
             tk_count = 0
+            raise
         logging.info("Embedding elapsed({}): {:.2f}".format(r["name"], timer() - st))
         callback(msg="Finished embedding (in {:.2f}s)! Start to build index!".format(timer() - st))
-
     # logging.info(f"task_executor init_kb index {search.index_name(r["tenant_id"])} embd_mdl {embd_mdl.llm_name} vector length {vector_size}")
     init_kb(r, vector_size)
     chunk_count = len(set([c["id"] for c in cks]))
@@ -413,16 +406,17 @@ def main():
         es_r = settings.docStoreConn.insert(cks[b:b + es_bulk_size], search.index_name(r["tenant_id"]), r["kb_id"])
         if b % 128 == 0:
             callback(prog=0.8 + 0.1 * (b + 1) / len(cks), msg="")
-
     logging.info("Indexing elapsed({}): {:.2f}".format(r["name"], timer() - st))
     if es_r:
-            callback(-1, "Insert chunk error, detail info please check log file. Please also check ES status!")
+        callback(-1, "Insert chunk error, detail info please check log file. Please also check Elasticsearch/Infinity status!")
         settings.docStoreConn.delete({"doc_id": r["doc_id"]}, search.index_name(r["tenant_id"]), r["kb_id"])
         logging.error('Insert chunk error: ' + str(es_r))
-        else:
+        raise Exception('Insert chunk error: ' + str(es_r))
+
     if TaskService.do_cancel(r["id"]):
         settings.docStoreConn.delete({"doc_id": r["doc_id"]}, search.index_name(r["tenant_id"]), r["kb_id"])
-                continue
+        return
+
     callback(msg="Indexing elapsed in {:.2f}s.".format(timer() - st))
     callback(1., "Done!")
     DocumentService.increment_chunk_num(
@@ -432,29 +426,42 @@ def main():
                 r["id"], tk_count, len(cks), timer() - st))
 
 
+def handle_task():
+    global PAYLOAD, DONE_TASKS, FAILED_TASKS
+    task = collect()
+    if task:
+        try:
+            logging.info(f"handle_task begin for task {json.dumps(task)}")
+            do_handle_task(task)
+            DONE_TASKS += 1
+            logging.exception(f"handle_task done for task {json.dumps(task)}")
+        except Exception:
+            FAILED_TASKS += 1
+            logging.exception(f"handle_task got exception for task {json.dumps(task)}")
+    if PAYLOAD:
+        PAYLOAD.ack()
+        PAYLOAD = None
+
+
 def report_status():
-    global CONSUMER_NAME, BOOT_AT, DONE_TASKS, RETRY_TASKS, PENDING_TASKS, HEAD_CREATED_AT, HEAD_DETAIL
+    global CONSUMER_NAME, BOOT_AT, DONE_TASKS, FAILED_TASKS, PENDING_TASKS, LAG_TASKS
     REDIS_CONN.sadd("TASKEXE", CONSUMER_NAME)
     while True:
         try:
             now = datetime.now()
-            PENDING_TASKS = REDIS_CONN.queue_length(SVR_QUEUE_NAME)
+            group_info = REDIS_CONN.queue_info(SVR_QUEUE_NAME, "rag_flow_svr_task_broker")
-            if PENDING_TASKS > 0:
+            if group_info is not None:
-                head_info = REDIS_CONN.queue_head(SVR_QUEUE_NAME)
+                PENDING_TASKS = int(group_info["pending"])
-                if head_info is not None:
+                LAG_TASKS = int(group_info["lag"])
-                    seconds = int(head_info[0].split("-")[0]) / 1000
-                    HEAD_CREATED_AT = datetime.fromtimestamp(seconds).isoformat()
-                    HEAD_DETAIL = head_info[1]
 
             heartbeat = json.dumps({
                 "name": CONSUMER_NAME,
                 "now": now.isoformat(),
                 "boot_at": BOOT_AT,
                 "done": DONE_TASKS,
-                "retry": RETRY_TASKS,
+                "failed": FAILED_TASKS,
                 "pending": PENDING_TASKS,
-                "head_created_at": HEAD_CREATED_AT,
+                "lag": LAG_TASKS,
-                "head_detail": HEAD_DETAIL,
             })
             REDIS_CONN.zadd(CONSUMER_NAME, heartbeat, now.timestamp())
             logging.info(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}")
@@ -466,14 +473,13 @@ def report_status():
             logging.exception("report_status got exception")
         time.sleep(30)
 
-
+def main():
-if __name__ == "__main__":
     background_thread = threading.Thread(target=report_status)
     background_thread.daemon = True
     background_thread.start()
 
     while True:
+        handle_task()
+
+if __name__ == "__main__":
     main()
-        if PAYLOAD:
-            PAYLOAD.ack()
-            PAYLOAD = None

rag/utils/redis_conn.py

@@ -225,14 +225,16 @@ class RedisDB:
             logging.exception("xpending_range: " + consumer_name + " got exception")
             self.__open__()
 
-    def queue_length(self, queue) -> int:
+    def queue_info(self, queue, group_name) -> dict:
         for _ in range(3):
             try:
-                num = self.REDIS.xlen(queue)
+                groups = self.REDIS.xinfo_groups(queue)
-                return num
+                for group in groups:
+                    if group["name"] == group_name:
+                        return group
             except Exception:
                 logging.exception("queue_length" + str(queue) + " got exception")
-        return 0
+        return None
 
     def queue_head(self, queue) -> int:
         for _ in range(3):