From ca197a34d6da12bf32274b4a0e59627741cc18e6 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 7 Aug 2024 13:18:57 -0700 Subject: [PATCH] [FrontEnd] Make merge_async_iterators is_cancelled arg optional https://github.com/vllm-project/vllm/pull/7111 made a change to the merge_async_iterators utils function to add an is_cancelled arg. It would be good for this new arg to be optional to retain backwards compatibility for other server front-ends that might already be using this utility function. --- vllm/utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm/utils.py b/vllm/utils.py index 61e3bb0bfc33..70106dc9c665 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -404,7 +404,7 @@ async def iterate_with_cancellation( async def merge_async_iterators( *iterators: AsyncGenerator[T, None], - is_cancelled: Callable[[], Awaitable[bool]], + is_cancelled: Optional[Callable[[], Awaitable[bool]]] = None, ) -> AsyncGenerator[Tuple[int, T], None]: """Merge multiple asynchronous iterators into a single iterator. @@ -412,8 +412,8 @@ async def merge_async_iterators( When it yields, it yields a tuple (i, item) where i is the index of the iterator that yields the item. - It also polls the provided function at least once per second to check - for client cancellation. + It also optionally polls a provided function at least once per second + to check for client cancellation. """ # Can use anext() in python >= 3.10 @@ -421,12 +421,13 @@ async def merge_async_iterators( ensure_future(pair[1].__anext__()): pair for pair in enumerate(iterators) } + timeout = None if is_cancelled is None else 1 try: while awaits: done, pending = await asyncio.wait(awaits.keys(), return_when=FIRST_COMPLETED, - timeout=1) - if await is_cancelled(): + timeout=timeout) + if is_cancelled is not None and await is_cancelled(): raise asyncio.CancelledError("client cancelled") for d in done: pair = awaits.pop(d)