Small async_llm_engine refactor (vllm-project#2618)

esmeetu · Jan 28, 2024 · 89be30f · 89be30f
1 parent f8ecb84
commit 89be30f
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/vllm/engine/async_llm_engine.py b/vllm/engine/async_llm_engine.py
@@ -53,7 +53,7 @@ def put(self, item: RequestOutput) -> None:
         self._queue.put_nowait(item)
 
     def finish(self) -> None:
-        self._queue.put_nowait(StopIteration)
+        self._queue.put_nowait(StopAsyncIteration())
         self._finished = True
 
     @property
@@ -65,9 +65,7 @@ def __aiter__(self):
 
     async def __anext__(self) -> RequestOutput:
         result = await self._queue.get()
-        if result is StopIteration:
-            raise StopAsyncIteration
-        elif isinstance(result, Exception):
+        if isinstance(result, Exception):
             raise result
         return result