fix(streams): prevent unbounded queue memory spikes and task leaks in streams.concat - #167
fix(streams): prevent unbounded queue memory spikes and task leaks in streams.concat#167Rajeev91691 wants to merge 2 commits into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces backpressure support to the concat stream processor by adding a queue_maxsize parameter and improves task cancellation handling. The review feedback highlights a critical bug where concat could hang indefinitely if a background task raises an exception, suggesting the use of a task group manager (context.context()) to safely manage background tasks. Additionally, the reviewer points out a redundant try...except block in enqueue and an unnecessary sleep in the test that could lead to flakiness.
| async def concat( | ||
| *contents: AsyncIterable[_T], | ||
| queue_maxsize: int = 1, | ||
| ) -> AsyncIterable[_T]: | ||
| """Concatenate multiple streams into one. | ||
|
|
||
| The streams are looped over concurrently before being assembled into a single | ||
| output stream. | ||
|
|
||
| Args: | ||
| *contents: each stream to concat as a separate argument. | ||
| queue_maxsize: The maximum number of items to buffer in an internal queue | ||
| for each input stream. Set to 0 to use an unbounded queue. | ||
|
|
||
| Yields: | ||
| The concatenation of all streams. | ||
| """ | ||
| output_queues = [asyncio.Queue() for _ in contents] | ||
| output_queues = [asyncio.Queue(maxsize=queue_maxsize) for _ in contents] | ||
|
|
||
| async def _stream_outputs( | ||
| idx: int, | ||
| ): | ||
| async for c in contents[idx]: | ||
| output_queues[idx].put_nowait(c) | ||
| # Adds None to indicate end of output. | ||
| output_queues[idx].put_nowait(None) | ||
| async def _stream_outputs(idx: int): | ||
| try: | ||
| async for c in contents[idx]: | ||
| await output_queues[idx].put(c) | ||
| await output_queues[idx].put(None) | ||
| except asyncio.CancelledError: | ||
| raise | ||
|
|
||
| tasks = [] | ||
| for idx, _ in enumerate(contents): | ||
| tasks.append(context.create_task(_stream_outputs(idx))) | ||
| try: | ||
| for idx, _ in enumerate(contents): | ||
| tasks.append(context.create_task(_stream_outputs(idx))) | ||
|
|
||
| for q in output_queues: | ||
| while (part := await q.get()) is not None: | ||
| for q in output_queues: | ||
| while (part := await q.get()) is not None: | ||
| q.task_done() | ||
| yield part | ||
| q.task_done() | ||
| yield part | ||
| finally: | ||
| for t in tasks: | ||
| if not t.done(): | ||
| t.cancel() | ||
| if tasks: | ||
| await asyncio.gather(*tasks, return_exceptions=True) |
There was a problem hiding this comment.
Critical Bug: Potential Hang on Exception in Background Tasks
If one of the background _stream_outputs tasks raises an exception (for example, during iteration of contents[idx]), it will terminate early without putting None into its corresponding queue.
Because concat is currently manually managing tasks and awaiting q.get() sequentially, the main loop will block indefinitely on await q.get() for that queue. Since the background task is already dead and will never put None or any other item, the generator hangs forever, and the finally block is never reached.
Solution
Instead of manually managing tasks, cancelling them, and handling exceptions, we should leverage the custom context.context() task group manager. This ensures that if any background task fails, the task group automatically cancels all other tasks and the main task, propagating the exception correctly and preventing any hangs.
async def concat(
*contents: AsyncIterable[_T],
queue_maxsize: int = 1,
) -> AsyncIterable[_T]:
"""Concatenate multiple streams into one.
The streams are looped over concurrently before being assembled into a single
output stream.
Args:
*contents: each stream to concat as a separate argument.
queue_maxsize: The maximum number of items to buffer in an internal queue
for each input stream. Set to 0 to use an unbounded queue.
Yields:
The concatenation of all streams.
"""
output_queues = [asyncio.Queue(maxsize=queue_maxsize) for _ in contents]
async def _stream_outputs(idx: int):
async for c in contents[idx]:
await output_queues[idx].put(c)
await output_queues[idx].put(None)
async with context.context() as tg:
for idx in range(len(contents)):
tg.create_task(_stream_outputs(idx))
for q in output_queues:
while (part := await q.get()) is not None:
q.task_done()
yield part
q.task_done()| try: | ||
| async for part in content: | ||
| await queue.put(part) | ||
| finally: | ||
| await queue.put(None) | ||
| except asyncio.CancelledError: | ||
| raise |
There was a problem hiding this comment.
Redundant try...except Block
The try...except asyncio.CancelledError: raise block is completely redundant here. In Python, if a CancelledError (or any other exception) is raised during async for part in content or await queue.put(part), it will naturally propagate out of the function immediately.
Since there is no finally block or broad except Exception block that we need to bypass, wrapping the code in this try...except block has no functional effect and can be safely removed to simplify the code.
async for part in content:
await queue.put(part)
await queue.put(None)| await asyncio.sleep(0.05) | ||
| self.assertTrue(task_cancelled.is_set()) |
There was a problem hiding this comment.
Unnecessary Sleep and Potential Flakiness
Because concat is an async generator, when we break from the async for loop, Python automatically calls aclose() on the generator. The generator's cleanup (either via the finally block or the async with block) awaits the cancellation and completion of all background tasks before aclose() completes.
Therefore, by the time the async for loop exits, the background producer task is guaranteed to have been fully cancelled and completed. We can assert task_cancelled.is_set() immediately without any arbitrary sleep, which avoids potential flakiness in busy CI environments.
| await asyncio.sleep(0.05) | |
| self.assertTrue(task_cancelled.is_set()) | |
| self.assertTrue(task_cancelled.is_set()) |
…, simplify enqueue, deflake test
1 participant
This PR resolves #164.
Description
In
genai_processors/streams.py,streams.concat(*contents)is used to concatenate multiple asynchronous streams. Previously, it initialized internalasyncio.Queue()instances with unbounded size (maxsize=0) and spawned concurrent background tasks callingput_nowait(c). This caused all subsequent streams to buffer entirely in RAM while the consumer was still reading the first stream, leading to memory spikes and potential OOM issues.Furthermore:
breakor exception), the background tasks remained running, resulting in task/resource leaks.enqueuetask was cancelled while blocked on a queue put, thefinally: await queue.put(None)block blocked indefinitely on a full queue, causing the task to hang.This PR:
queue_maxsize: int = 1by default) instreams.concatto provide backpressure.try...finallyblock to cancel and clean up background tasks if iteration finishes early or errors.enqueueto prevent hanging on cancellation by avoiding blocking queue operations in theCancelledErrorpath.test_concat_backpressure_and_cancellationtogenai_processors/tests/streams_test.pyto verify this behavior.