From 6a771459d50b211cdcc55f85df145e8991edd08a Mon Sep 17 00:00:00 2001
From: irexyc <irexyc@gmail.com>
Date: Thu, 7 Mar 2024 11:59:36 +0000
Subject: [PATCH] fix close

---
 lmdeploy/serve/async_engine.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lmdeploy/serve/async_engine.py b/lmdeploy/serve/async_engine.py
index 783f6a4c2..a9258358a 100644
--- a/lmdeploy/serve/async_engine.py
+++ b/lmdeploy/serve/async_engine.py
@@ -123,10 +123,10 @@ def stream_response(self):
             self.history.append((self._prompt, resp.text))
 
     def close(self):
-        _ = self.response()
         if self._engine:
             inst = self._engine.create_instance()
-            inst.end(self._id)
+            inst.cancel(self._id)
+        _ = self.response()
 
     def __repr__(self) -> str:
         _ = self.response()
@@ -668,6 +668,7 @@ def _work_thread():
                         gen_config=gen_config,
                         stream_output=True):
                     _, res, tokens = outputs
+                    print(session._id, tokens)
                     response, state = self.tokenizer.detokenize_incrementally(
                         res,
                         state,