Add option to use asyncio for AsyncWrapper by AMOOOMA · Pull Request #38262 · apache/beam

AMOOOMA · 2026-04-21T18:17:15Z

Add option to use asyncio for AsyncWrapper to increase parallelism remote RPC calls.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

gemini-code-assist · 2026-04-21T18:17:24Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces an optional asyncio-based execution mode for the AsyncWrapper in Apache Beam's Python SDK. By allowing the use of coroutines instead of relying solely on ThreadPoolExecutor, this change aims to enhance the efficiency and parallelism of remote RPC calls within asynchronous DoFns. The implementation includes robust management of a background event loop and comprehensive updates to the test suite to validate the new functionality.

Highlights

Asyncio Support: Added an option to use asyncio for AsyncWrapper to improve parallelism for remote RPC calls.
Background Event Loop: Implemented a background thread to manage the asyncio event loop, ensuring it can be used alongside existing thread-based processing.
Test Coverage: Updated existing tests to support both ThreadPoolExecutor and asyncio modes, ensuring parity and stability across both execution strategies.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2026-04-21T20:23:34Z

Assigning reviewers:

R: @tvalentyn for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

AMOOOMA · 2026-04-28T19:34:12Z

R: @damccorm

github-actions · 2026-04-28T19:35:23Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

damccorm

Thanks! Generally the approach LGTM

damccorm · 2026-04-28T20:42:11Z

+    process_result = self._sync_fn.process(element, *args, **kwargs)
+    bundle_result = self._sync_fn.finish_bundle()
+
+    if not process_result:


Optional nit: thoughts on wrapping this in a small local function so that we just need to call something like:

process_result = gather_async_results(process_result) bundle_result = gather_async_results(bundle_result)

damccorm · 2026-04-28T20:47:13Z

      id_fn: A function that returns a hashable object from an element. This
        will be used to track items instead of the element's default hash.
+      use_asyncio: If true, use asyncio and coroutines to process items. If
+        false, use ThreadPoolExecutor.


In our tests, has asyncio generally performed better? Are there reasons to default it off?

Regardless, this doc should have info on why a user would choose one vs the other (or if its almost always one then we should just make this a kwarg)

I think mostly in cases of where we want to run CPU intensive work on threads the default will be better since we only have 1 main event loop in async mode. But yeah definitely would need to add some more info. Done.

damccorm · 2026-04-28T20:49:54Z



-class AsyncTest(unittest.TestCase):
+class _AsyncTestBase:


https://stackoverflow.com/questions/38729007/parametrize-class-tests-with-pytest is probably the more pythonic way to do this

Done. Turns out to use parameterized class instead of the pytest one since it seems like other test files in beam uses the parameterized one.

damccorm · 2026-04-28T20:55:22Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces an asyncio-based execution mode for AsyncWrapper, allowing elements to be processed via coroutines in addition to the existing ThreadPoolExecutor. The changes include the management of a background event loop thread and the implementation of async-specific processing logic. Feedback highlights critical bugs in the async_fn_process method, which currently fails to await coroutines and risks blocking the event loop. Additionally, a thread-safety issue was identified in reset_state, where shared class variables are modified without proper locking.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Co-authored-by: Danny McCormick <dannymccormick@google.com>

damccorm

Thanks!

Add option to use asyncio for AsyncWrapper

0772c50

github-actions Bot added the python label Apr 21, 2026

github-actions Bot added the Next Action: Reviewers label Apr 21, 2026

Merge branch 'apache:master' into remote

3f90f59

damccorm reviewed Apr 28, 2026

View reviewed changes

gemini-code-assist Bot reviewed Apr 28, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/transforms/async_dofn.py Outdated

Comment thread sdks/python/apache_beam/transforms/async_dofn.py Outdated

AMOOOMA and others added 5 commits April 28, 2026 14:05

Update sdks/python/apache_beam/transforms/async_dofn.py

2ca2ed4

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update sdks/python/apache_beam/transforms/async_dofn.py

015cb11

Co-authored-by: Danny McCormick <dannymccormick@google.com>

Address comments

d755306

Minor fixes

5d8c5ca

lint

5104609

damccorm approved these changes May 4, 2026

View reviewed changes

damccorm merged commit 637231c into apache:master May 4, 2026
144 of 156 checks passed

Conversation

AMOOOMA commented Apr 21, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Apr 21, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions Bot commented Apr 21, 2026

Uh oh!

AMOOOMA commented Apr 28, 2026

Uh oh!

github-actions Bot commented Apr 28, 2026

Uh oh!

damccorm left a comment

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

AMOOOMA Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

damccorm Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

AMOOOMA Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

AMOOOMA Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm commented Apr 28, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

damccorm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants