Fix entrypoints unit test

AlibabaPAI · Feb 7, 2025 · 322862b · 322862b
1 parent 75af824
commit 322862b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/llumnix/entrypoints/vllm/client.py b/llumnix/entrypoints/vllm/client.py
@@ -48,7 +48,7 @@ async def generate(self,
             raise ValueError("Unsupported feature: multiple sequence decoding")
         logger.info("entrypoints receive request {}".format(request_id))
         # pylint: disable=unexpected-keyword-arg
-        results_generator = AsyncStream(request_id, cancel=partial(self.abort, verbose=False))
+        results_generator = AsyncStream(request_id, cancel=self.abort_request)
         self.request_streams[request_id] = results_generator
         server_info_copy = copy.deepcopy(self.server_info)
 
@@ -117,6 +117,10 @@ async def abort(self, request_id: str) -> None:
         except ray.exceptions.RayActorError:
             logger.warning("Manager is unavailable.")
 
+    def abort_request(self, request_id: str) -> None:
+        logger.info("Abort request: {}.".format(request_id))
+        self.manager.abort.remote(request_id)
+
     async def is_ready(self) -> bool:
         ready_status = await self.manager.is_ready.remote()
         return ready_status