[Bugfix][Core] Close underlying iterator in merge_async_iterators single-iterator fast path (#44726)
Signed-off-by: Ting Sun <suntcrick@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Ting SUN committed
Jul 7, 2026 at 19:13 UTC
8b91cd5b203598bb06595f3a8cb991fecde94f52
2 files changed
+31
-2
tests/utils_/test_async_utils.py
+22
@@ -40,3 +40,25 @@ async def test_merge_async_iterators():
40
print("Iterator was cancelled normally")
41
except (Exception, asyncio.CancelledError) as e:
42
raise AssertionError() from e
43
+
44
+
45
+@pytest.mark.asyncio
46
+async def test_merge_async_iterators_single_closes_underlying():
47
+ # The single-iterator fast path must close the underlying generator when
48
+ # the merged generator is closed, matching the multi-iterator path. On the
49
+ # buggy fast path the underlying generator is left running.
50
+ closed = False
51
+
52
+ async def gen():
53
+ nonlocal closed
54
+ try:
55
+ while True:
56
+ yield "x"
57
+ await asyncio.sleep(0.01)
58
+ finally:
59
+ closed = True
60
+
61
+ merged = merge_async_iterators(gen())
62
+ assert await anext(merged) == (0, "x")
63
+ await merged.aclose()
64
+ assert closed
vllm/utils/async_utils.py
+9
-2
@@ -103,8 +103,15 @@ async def merge_async_iterators(
103
"""
104
if len(iterators) == 1:
105
# Fast-path single iterator case.
106
- async for item in iterators[0]:
107
- yield 0, item
106
+ iterator: AsyncGenerator[T, None] | None = iterators[0]
107
+ try:
108
+ async for item in iterator: # type: ignore[union-attr]
109
+ yield 0, item
110
+ iterator = None
111
+ finally:
112
+ if iterator is not None:
113
+ with contextlib.suppress(BaseException):
114
+ await iterator.aclose()
115
return
116
117
loop = asyncio.get_running_loop()