raise error for wait

Author: qyh111

ucm/integration/vllm/ucm_connector.py

@@ -482,11 +482,13 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         for request_id, task in request_to_task.items():
             # TODO error handling
             if self.global_rank == 0 or not self.load_only_first_rank:
-                if self.store.wait(task) != 0:
+                try:
+                    self.store.wait(task)
+                except RuntimeError as e:
+                    logger.error("request {request_id} load kv cache failed.:", e)
                     self._invalid_block_ids.update(
                         metadata.request_meta[request_id].load_block_ids[1]
                     )
-                    logger.error(f"request {request_id} load kv cache failed.")
             if self.load_only_first_rank:
                 self._broadcast(req_broadcast_addr[request_id])
         load_end_time = time.perf_counter() * 1000
@@ -557,7 +559,10 @@ def wait_for_save(self) -> None:
             )
 
         for request_id, task in request_to_task.items():
-            self.store.wait(task)
+            try:
+                self.store.wait(task)
+            except RuntimeError as e:
+                logger.error("request {request_id} dump kv cache failed.:", e)
         save_end_time = time.perf_counter() * 1000
         save_speed = (
             num_saved_block

ucm/store/pcstore/pcstore_connector.py

@@ -107,7 +107,9 @@ def dump_data(
         return task_id
 
     def wait(self, task: Task) -> None:
-        self.store.Wait(task.task_id)
+        ret = self.store.Wait(task.task_id)
+        if ret != 0:
+            raise RuntimeError(f"Wait failed for task {task.task_id}, return={ret}")
 
     def check(self, task: Task) -> bool:
         return self.store.Check(task.task_id)