Small async_llm_engine refactor (#2618)

89be30fa · Murali Andoorveedu · GitHub · f8ecb84c · 89be30fa
Unverified Commit 89be30fa authored Jan 27, 2024 by Murali Andoorveedu Committed by GitHub Jan 27, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 4 deletions

vllm/engine/async_llm_engine.py vllm/engine/async_llm_engine.py +2 -4

No files found.
--- a/vllm/engine/async_llm_engine.py
+++ b/vllm/engine/async_llm_engine.py
@@ -53,7 +53,7 @@ class AsyncStream:
        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 @@ class AsyncStream:

    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