fix semaphore

InternLM · Dec 25, 2024 · 3f07733 · 3f07733
1 parent 8fa85dc
commit 3f07733
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py
@@ -329,6 +329,7 @@ def __init__(self):
 
     async def acquire(self):
         if self.val:
+            self.val = 0
             return
         self.fut = self.loop.create_future()
         await self.fut
@@ -578,11 +579,11 @@ async def async_stream_infer(self,
         try:
             while True:
                 await sem.acquire()
-
                 state = shared_state.consume()
+
                 status, seq_len = state.status, state.seq_len
 
-                if status == 7:
+                if status in [7, 8]:  # finish / canceled
                     finish, status = True, 0
                 elif status:
                     yield self._get_error_output()