Perf: make do_cancel quicker. (#8846)

### What problem does this PR solve?

### Type of change

- [x] Performance Improvement

api/utils/api_utils.py

@@ -606,6 +606,7 @@ TimeoutException = Union[Type[BaseException], BaseException]
 OnTimeoutCallback = Union[Callable[..., Any], Coroutine[Any, Any, Any]]
 def timeout(
     seconds: float |int = None,
+    attempts: int = 2,
     *,
     exception: Optional[TimeoutException] = None,
     on_timeout: Optional[OnTimeoutCallback] = None
@@ -625,23 +626,28 @@ def timeout(
             thread.daemon = True
             thread.start()
 
+            for a in range(attempts):
                 try:
                     result = result_queue.get(timeout=seconds)
                     if isinstance(result, Exception):
                         raise result
                     return result
                 except queue.Empty:
-                raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds")
+                    pass
+            raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds and {attempts} attempts.")
 
         @wraps(func)
         async def async_wrapper(*args, **kwargs) -> Any:
             if seconds is None:
                 return await func(*args, **kwargs)
 
+            for a in range(attempts):
                 try:
                     with trio.fail_after(seconds):
                         return await func(*args, **kwargs)
                 except trio.TooSlowError:
+                    if a < attempts -1:
+                        continue
                     if on_timeout is not None:
                         if callable(on_timeout):
                             result = on_timeout()
@@ -651,13 +657,13 @@ def timeout(
                         return on_timeout
 
                     if exception is None:
-                    raise TimeoutError(f"Operation timed out after {seconds} seconds")
+                        raise TimeoutError(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
 
                     if isinstance(exception, BaseException):
                         raise exception
 
                     if isinstance(exception, type) and issubclass(exception, BaseException):
-                    raise exception(f"Operation timed out after {seconds} seconds")
+                        raise exception(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
 
                     raise RuntimeError("Invalid exception type provided")
 

rag/raptor.py

@@ -42,6 +42,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
         self._prompt = prompt
         self._max_token = max_token
 
+    @timeout(60)
     async def _chat(self, system, history, gen_conf):
         response = get_llm_cache(self._llm_model.llm_name, system, history, gen_conf)
         if response:

rag/svr/task_executor.py

@@ -214,7 +214,7 @@ async def collect():
     canceled = False
     task = TaskService.get_task(msg["id"])
     if task:
-        canceled = TaskService.do_cancel(task["id"])
+        canceled = DocumentService.do_cancel(task["doc_id"])
     if not task or canceled:
         state = "is unknown" if not task else "has been cancelled"
         FAILED_TASKS += 1
@@ -382,7 +382,7 @@ async def build_chunks(task, progress_callback):
 
         docs_to_tag = []
         for d in docs:
-            task_canceled = TaskService.do_cancel(task["id"])
+            task_canceled = DocumentService.do_cancel(task["doc_id"])
             if task_canceled:
                 progress_callback(-1, msg="Task has been canceled.")
                 return
@@ -531,7 +531,7 @@ async def do_handle_task(task):
         progress_callback(-1, msg=error_message)
         raise Exception(error_message)
 
-    task_canceled = TaskService.do_cancel(task_id)
+    task_canceled = DocumentService.do_cancel(task_doc_id)
     if task_canceled:
         progress_callback(-1, msg="Task has been canceled.")
         return
@@ -609,7 +609,7 @@ async def do_handle_task(task):
 
     for b in range(0, len(chunks), DOC_BULK_SIZE):
         doc_store_result = await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert(chunks[b:b + DOC_BULK_SIZE], search.index_name(task_tenant_id), task_dataset_id))
-        task_canceled = TaskService.do_cancel(task_id)
+        task_canceled = DocumentService.do_cancel(task_doc_id)
         if task_canceled:
             progress_callback(-1, msg="Task has been canceled.")
             return