[BugFix] fix prefix tree updating timeout (#6615)

This commit is contained in:
Yonghua Li
2026-03-03 14:32:15 +08:00
committed by GitHub
parent 025697584a
commit 27ae02fd82
@@ -2087,32 +2087,33 @@ class PrefixCacheManager:
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 @@ class PrefixCacheManager:
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)