refine4

Author: niushengxiao

lightllm/server/multi_level_kv_cache/disk_cache_worker.py

@@ -85,10 +85,14 @@ def run(self) -> None:
             payload_groups = self._gather_offload_payloads()
             if not payload_groups:
                 continue
+            tasks_and_pages = []
             for payloads in payload_groups:
                 if not payloads:
                     continue
-                self._persist_pages_to_disk(payloads)
+                result = self._persist_pages_to_disk(payloads)
+                if result is not None:
+                    tasks_and_pages.append(result)
+            self._wait_and_update_pages(tasks_and_pages)
 
     def _gather_offload_payloads(self) -> List[List[_PagePayload]]:
         self.cpu_cache_client.lock.acquire_sleep1ms()
@@ -111,13 +115,13 @@ def _gather_offload_payloads(self) -> List[List[_PagePayload]]:
         return payload_groups
 
     # 数据写入磁盘
-    def _persist_pages_to_disk(self, payloads: List[_PagePayload]) -> None:
+    def _persist_pages_to_disk(self, payloads: List[_PagePayload]):
         if not payloads:
-            return
+            return None
         page_indexes = [payload.index for payload in payloads]
         tokens = [payload.hash_key for payload in payloads]
         if not page_indexes:
-            return
+            return None
 
         kv_indexer = torch.tensor(page_indexes, dtype=torch.int32, device="cpu")
         query_result = self.service.query(tokens)
@@ -129,13 +133,16 @@ def _persist_pages_to_disk(self, payloads: List[_PagePayload]) -> None:
                 time.sleep(0.001)
 
             task = self.service.create(tokens=tokens, kv_page_indexer=kv_indexer, mode="w")
-            # 数据安全即可结束等待，无需写入完成
+            return (task, page_indexes)
+        return None
+
+    def _wait_and_update_pages(self, tasks_and_pages: List) -> None:
+        for task, page_indexes in tasks_and_pages:
             while not task.data_safe():
                 time.sleep(0.001)
-
-        self.cpu_cache_client.lock.acquire_sleep1ms()
-        self.cpu_cache_client.update_pages_status_to_ready_recycle(page_list=page_indexes, deref=True)
-        self.cpu_cache_client.lock.release()
+            self.cpu_cache_client.lock.acquire_sleep1ms()
+            self.cpu_cache_client.update_pages_status_to_ready_recycle(page_list=page_indexes, deref=True)
+            self.cpu_cache_client.lock.release()
 
     def query_loadable_pages(self, tokens: List[int], start_pos: int) -> int:
         """