[bugfix]Add get_finished in patch

Author: qyh111

ucm/integration/vllm/patch/v0180/vllm_ascend/pc/distributed/kv_transfer/kv_pool/ucm_connector.py

@@ -74,6 +74,13 @@ def request_finished_all_groups(
     return engine.request_finished(request, [])
 
 
+def get_finished(
+    self: Any,
+    finished_req_ids: set[str],
+) -> tuple[set[str] | None, set[str] | None]:
+    return self._ucm_engine.get_finished(finished_req_ids)
+
+
 def build_connector_worker_meta(self: Any) -> Any:
     return self._call_ucm_reserved_hook("build_connector_worker_meta")
 
@@ -142,6 +149,7 @@ def patch_ucm_connector_v1(mod: Any) -> None:
         "request_finished_all_groups",
         request_finished_all_groups,
     )
+    patch_or_inject(target_cls, "get_finished", get_finished)
     patch_or_inject(
         target_cls,
         "build_connector_worker_meta",

ucm/integration/vllm/ucm_connector.py

@@ -711,15 +711,23 @@ def build_connector_meta(
         for request_id in scheduler_output.finished_req_ids:
             self.requests_meta.pop(request_id, None)
 
-        for request_id, dispatch_meta in requests_dispatch_meta.items():
-            if len(dispatch_meta.dump_block_ids[0]) > 0:
-                self._async_dump_req_ids.add(request_id)
+        self._track_async_dump_requests(requests_dispatch_meta)
 
         return UCMConnectorMetadata(
             requests_dispatch_meta,
             scheduler_output.preempted_req_ids or set(),
         )
 
+    def _track_async_dump_requests(
+        self,
+        requests_dispatch_meta: dict[str, RequestDispatchMeta],
+    ) -> None:
+        self._async_dump_req_ids.update(
+            request_id
+            for request_id, dispatch_meta in requests_dispatch_meta.items()
+            if len(dispatch_meta.dump_block_ids[0]) > 0
+        )
+
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         metadata = self._get_connector_metadata()
         assert isinstance(metadata, UCMConnectorMetadata)
@@ -2660,6 +2668,8 @@ def build_connector_meta(
         for request_id in scheduler_output.finished_req_ids:
             self.requests_meta.pop(request_id, None)
 
+        self._track_async_dump_requests(requests_dispatch_meta)
+
         return UCMConnectorMetadata(
             requests_dispatch_meta,
             scheduler_output.preempted_req_ids or set(),
@@ -2670,7 +2680,9 @@ def request_finished_all_groups(
         request: "Request",
         block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
-        return False, None
+        if block_ids:
+            return self.request_finished(request, block_ids[0])
+        return self.request_finished(request, [])
 
 
 def use_hybrid_linear_attention_layout(