Introduce more abstract handling of throttling signaling #35984

jrmccluskey · 2025-08-27T18:27:06Z

Adds a class to abstract away the signaling of throttling to the SDK worker (for now this more clearly codifies the existing metric-based mechanism, but this also allows us to replace that mechanism in the future.) Also adds a new AdaptiveThrottler subclass that handles the throttling signaling itself, which streamlines code using it like the RemoteModelHandler, which now uses the new class.

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

Summary of Changes

Hello @jrmccluskey, 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 refactors the handling of throttling signaling within Apache Beam's Python SDK. It introduces a new ThrottlingSignaler class to abstract the mechanism for communicating throttling events to the SDK worker and a ReactiveThrottler class that wraps the existing AdaptiveThrottler to streamline logging and signaling. The RemoteModelHandler is updated to leverage this new ReactiveThrottler, simplifying its internal throttling logic.

Highlights

Abstracted Throttling Signaling: A new ThrottlingSignaler class is introduced to provide a more abstract way to signal throttling duration to the SDK worker, using a cumulative metric. This design allows for future changes to the signaling mechanism without impacting callers.
Encapsulated Reactive Throttling: A ReactiveThrottler class, inheriting from AdaptiveThrottler, is added. This new class centralizes the logic for logging throttling events and signaling them via the ThrottlingSignaler, making the throttling implementation more self-contained.
Simplified RemoteModelHandler Throttling: The RemoteModelHandler's run_inference method is simplified by delegating the entire throttling loop, including logging and metric updates, to the new ReactiveThrottler's throttle() method. This reduces boilerplate and improves readability in RemoteModelHandler.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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. ↩

jrmccluskey · 2025-08-27T20:23:14Z

@gemini-code-assist review

gemini-code-assist

Code Review

This pull request introduces a good abstraction for handling throttling signaling by adding ThrottlingSignaler and ReactiveThrottler classes. This simplifies the RemoteModelHandler and improves code organization and maintainability. The changes are well-structured. I've left a couple of minor suggestions for improvement regarding naming and a typo.

sdks/python/apache_beam/io/components/adaptive_throttler.py

codecov · 2025-08-27T21:51:21Z

Codecov Report

❌ Patch coverage is 82.60870% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 54.83%. Comparing base (f3d10e3) to head (669ab51).
⚠️ Report is 35 commits behind head on master.

Files with missing lines	Patch %	Lines
...on/apache_beam/io/components/adaptive_throttler.py	80.00%	4 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #35984      +/-   ##
============================================
+ Coverage     54.75%   54.83%   +0.07%     
  Complexity     1646     1646              
============================================
  Files          1060     1061       +1     
  Lines        163560   163946     +386     
  Branches       1189     1189              
============================================
+ Hits          89562    89901     +339     
- Misses        71845    71892      +47     
  Partials       2153     2153

Flag	Coverage Δ
python	`80.84% <82.60%> (+0.03%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

github-actions · 2025-08-27T22:06:36Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

jrmccluskey · 2025-08-28T18:26:55Z

test breakage on the ML tests is unrelated, workflow is red on master

jrmccluskey · 2025-08-28T18:27:00Z

assign set of reviewers

github-actions · 2025-08-28T18:28:17Z

Assigning reviewers:

R: @claudevdm 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).

jrmccluskey · 2025-08-28T18:31:28Z

CC: @damccorm

damccorm

Nice, this is very clean! It would be great to follow up with a small pr with an example in https://github.com/apache/beam/tree/master/sdks/python/apache_beam/examples that we can merge once the release is out. Also maybe a small update to https://beam.apache.org/documentation/io/developing-io-python/ to mention handling throttling.

While its a simple code feature, I think exposing this so cleanly will help a lot of users.

damccorm · 2025-08-29T13:17:34Z

sdks/python/apache_beam/io/components/adaptive_throttler.py

+
+class ReactiveThrottler(AdaptiveThrottler):
+  """ A wrapper around the AdaptiveThrottler that also handles logging and
+  signaling throttling to the SDK harness.


Could we include example usage here and be a little more descriptive on what this does?

I think we want people to use this instead of AdaptiveThrottler, so it would be better to be explicit about the UX

Also, could we update CHANGES.md?

updated CHANGE.md + added a pseudocode workflow for users in the docstring

sdks/python/apache_beam/ml/inference/base.py

sdks/python/apache_beam/io/components/adaptive_throttler.py

Co-authored-by: Danny McCormick <[email protected]>

Introduce more abstract handling of throttling signaling

6fcc80a

github-actions bot added python io labels Aug 27, 2025

gemini-code-assist bot reviewed Aug 27, 2025

View reviewed changes

linting

fd2ac43

gemini-code-assist bot reviewed Aug 27, 2025

View reviewed changes

sdks/python/apache_beam/io/components/adaptive_throttler.py Outdated Show resolved Hide resolved

sdks/python/apache_beam/io/components/adaptive_throttler.py Outdated Show resolved Hide resolved

gemini suggestions

669ab51

github-actions bot added the Next Action: Reviewers label Aug 28, 2025

damccorm approved these changes Aug 29, 2025

View reviewed changes

review suggestions

f3bd33e

damccorm reviewed Sep 2, 2025

View reviewed changes

sdks/python/apache_beam/io/components/adaptive_throttler.py Outdated Show resolved Hide resolved

Update sdks/python/apache_beam/io/components/adaptive_throttler.py

b64a9c6

Co-authored-by: Danny McCormick <[email protected]>

damccorm approved these changes Sep 2, 2025

View reviewed changes

jrmccluskey and others added 2 commits September 2, 2025 13:40

Merge branch 'master' into abstractThrottling

857592d

linting

58c572d

jrmccluskey merged commit 35969b3 into apache:master Sep 2, 2025
120 of 124 checks passed

jrmccluskey mentioned this pull request Sep 11, 2025

Add ThrottlingSignaler class to the Java SDK #36119

Merged

3 tasks

Introduce more abstract handling of throttling signaling #35984

Introduce more abstract handling of throttling signaling #35984

Uh oh!

Conversation

jrmccluskey commented Aug 27, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

jrmccluskey commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Aug 27, 2025

Uh oh!

jrmccluskey commented Aug 28, 2025

Uh oh!

jrmccluskey commented Aug 28, 2025

Uh oh!

github-actions bot commented Aug 28, 2025

Uh oh!

jrmccluskey commented Aug 28, 2025

Uh oh!

damccorm left a comment

Choose a reason for hiding this comment

Uh oh!

damccorm Aug 29, 2025

Choose a reason for hiding this comment

Uh oh!

damccorm Aug 29, 2025

Choose a reason for hiding this comment

Uh oh!

jrmccluskey Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jrmccluskey commented Aug 27, 2025 •

edited

Loading

codecov bot commented Aug 27, 2025 •

edited

Loading