Fix: Prevent unbounded queue memory spikes and task leaks in streams.concat - #165
Fix: Prevent unbounded queue memory spikes and task leaks in streams.concat#165THRISHAL12345 wants to merge 2 commits into
streams.concat#165Conversation
|
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. |
There was a problem hiding this comment.
Code Review
This pull request introduces backpressure and cancellation safety to the concat stream utility by adding a queue_maxsize parameter and ensuring background tasks are cancelled when the generator exits. It also adds corresponding unit tests for bounded queue backpressure and early cancellation. The review feedback suggests leveraging the project's idiomatic context.context() task group in concat to handle background task exceptions robustly and simplify task cleanup. Additionally, it recommends removing a redundant try...except asyncio.CancelledError block in the enqueue function.
Summary of Changes
This pull request fixes a potential Out-Of-Memory (OOM) vulnerability and resource leak in
streams.concatwhen processing heavy multimodal streams (video, audio, PDF documents).Key Improvements:
streams.concat(genai_processors/streams.py):queue_maxsize: int = 1keyword argument tostreams.concat(*contents, queue_maxsize: int = 1).asyncio.Queue()(maxsize=0) with bounded queues (asyncio.Queue(maxsize=queue_maxsize)).put_nowait(c)calls withenqueue(c, output_queues[idx]), which usesawait queue.put(part)to pause upstream producers whenqueue_maxsizeitems are buffered ahead.queue_maxsize=0.streams.concatiteration in atry...finallyblock that cleanly cancels backgroundenqueuetasks (t.cancel()) when iteration completes or terminates early.enqueueso that when a task is cancelled (asyncio.CancelledError), it re-raises immediately without attempting to block onawait queue.put(None)into a full queue.genai_processors/tests/streams_test.py):test_concat_bounded_queue_backpressure: Verifies that bounded queueing correctly prevents downstream streams from running ahead and over-buffering.test_concat_early_cancellation: Verifies that early termination (aclose()) cleanly cancels background enqueue tasks.Verification
test_concat_bounded_queue_backpressureandtest_concat_early_cancellationpass cleanly.StreamsTestunit tests continue to pass without regression.Closes #164