[BugFix] fix prefix tree updating timeout

Author: liyonghua0910

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -2087,32 +2087,33 @@ def recv_data_transfer_result(self):
                     logger.error(f"recv_data_transfer_result: {str(traceback.format_exc())}")
                     raise e
 
-    def reset(self):
+    def reset(self, wait_for_tasks_done=False):
         """
         Reset the RadixTree.
         """
-        logger.info(f"wait for cache_task_inflight_signal to reset {self.cache_task_inflight_signal.value}")
-        while np.sum(self.cache_task_inflight_signal.value) != 0:
-            time.sleep(0.1)
 
-        logger.info("wait for recv_data_transfer_result done")
-        while not self.cache_task_queue.result_queue_empty():
-            time.sleep(0.1)
+        if wait_for_tasks_done:
+            logger.info(f"wait for cache_task_inflight_signal to reset: {self.cache_task_inflight_signal.value}")
+            while np.sum(self.cache_task_inflight_signal.value) != 0:
+                time.sleep(0.1)
+
+            logger.info("wait for recv_data_transfer_result done")
+            while not self.cache_task_queue.result_queue_empty():
+                time.sleep(0.1)
+
+            logger.info("wait for cpu_free_future to finish")
+            if self.cpu_free_future is not None:
+                self.cpu_free_future.result()
+
+            logger.info("wait for gpu_free_task_future to finish")
+            if self.gpu_free_task_future is not None:
+                self.gpu_free_task_future.result()
 
         logger.info(f"Resetting the RadixTree! node_map len {len(self.node_map)}")
 
-        logger.info("waiting for cpu_free_future to finish")
-        if self.cpu_free_future is not None:
-            self.cpu_free_future.result()
+        # clear future & events
         self.cpu_free_future = None
-        logger.info("reset cpu_free_future")
-
-        logger.info("waiting for gpu_free_task_future to finish")
-        if self.gpu_free_task_future is not None:
-            self.gpu_free_task_future.result()
         self.gpu_free_task_future = None
-        logger.info("reset gpu_free_task_future")
-
         self.task_swapping_event.clear()
 
         # clear node map
@@ -2157,11 +2158,11 @@ def clear_prefix_cache(self):
         prefix_tree_status_signal = self.prefix_tree_status_signal
         while True:
             if prefix_tree_status_signal.value[0] == PrefixTreeStatus.CLEARING:
-                self.reset()
+                self.reset(wait_for_tasks_done=True)
                 prefix_tree_status_signal.value[0] = PrefixTreeStatus.CLEARED
                 logger.info("Prefix cache tree is cleared.")
             if prefix_tree_status_signal.value[0] == PrefixTreeStatus.UPDATING:
-                self.reset()
+                self.reset(wait_for_tasks_done=False)
                 prefix_tree_status_signal.value[0] = PrefixTreeStatus.NORMAL
                 logger.info("Prefix cache tree is updated.")
             time.sleep(0.01)