[Gemini] Add client-side throttling to Java Remote Inference#39194
[Gemini] Add client-side throttling to Java Remote Inference#39194jrmccluskey wants to merge 4 commits into
Conversation
|
@gemini-code-assist review |
Summary of ChangesHello, I'm Gemini Code Assist1! 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 client-side throttling to the Java Remote Inference framework. By leveraging the existing ReactiveThrottler component, the changes ensure that remote inference requests are managed more robustly, preventing potential service overload by applying preemptive throttling based on request success and failure patterns. Highlights
New Features🧠 You can now enable Memory (public preview) 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 AssistThe 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
Customization To customize the 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. 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. Footnotes
|
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
/gemini review |
![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 integrates client-side reactive throttling into the RemoteInference transform. It adds a dependency on the :sdks:java:io:components project to utilize ReactiveThrottler, introduces a configurable throttling delay parameter (throttleDelaySecs), and wraps the remote model request execution with throttling logic. A new unit test and mock handler have also been added to verify that throttling is triggered and that the corresponding metrics are correctly recorded. There are no review comments, so I have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
There was a problem hiding this comment.
Code Review
This pull request integrates client-side throttling into the RemoteInference framework using ReactiveThrottler, introducing configuration options for throttling delay, sample period, sample update interval, and overload ratio, along with a new unit test to verify throttling behavior. The review feedback suggests several important improvements to ensure robustness: adding a null check for modelHandler to prevent potential NullPointerExceptions, validating throttling parameters in the builder methods to catch configuration errors early, using AtomicInteger in the mock handler to avoid race conditions during concurrent test execution, and checking both attempted and committed metrics to prevent test flakiness across different runners.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| private int requestCount = 0; | ||
|
|
||
| @Override | ||
| public void createClient(TestParameters parameters) {} | ||
|
|
||
| @Override | ||
| public Iterable<PredictionResult<TestInput, TestOutput>> request(List<TestInput> input) { | ||
| requestCount++; | ||
| // Fail 2 out of 3 requests. RetryHandler defaults to 3 max retries, | ||
| // so the 3rd attempt will succeed, avoiding pipeline failure while | ||
| // accumulating enough failures to trigger client-side throttling. | ||
| if (requestCount % 3 != 0) { | ||
| throw new RuntimeException("Intentional failure to trigger throttling"); | ||
| } | ||
| return input.stream() | ||
| .map(i -> PredictionResult.create(i, new TestOutput("processed-" + i.getModelInput()))) | ||
| .collect(Collectors.toList()); | ||
| } |
There was a problem hiding this comment.
To prevent potential race conditions and flaky test failures when multiple threads execute the pipeline concurrently, we should make MockThrottlingHandler thread-safe by using AtomicInteger instead of a plain int for requestCount.
private final java.util.concurrent.atomic.AtomicInteger requestCount =
new java.util.concurrent.atomic.AtomicInteger(0);
@Override
public void createClient(TestParameters parameters) {}
@Override
public Iterable<PredictionResult<TestInput, TestOutput>> request(List<TestInput> input) {
int count = requestCount.incrementAndGet();
// Fail 2 out of 3 requests. RetryHandler defaults to 3 max retries,
// so the 3rd attempt will succeed, avoiding pipeline failure while
// accumulating enough failures to trigger client-side throttling.
if (count % 3 != 0) {
throw new RuntimeException("Intentional failure to trigger throttling");
}
return input.stream()
.map(i -> PredictionResult.create(i, new TestOutput("processed-" + i.getModelInput())))
.collect(Collectors.toList());
}
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
1 participant
Takes the client-side throttling code introduced in #39021 and adds it into the remote inference base classes.
Part of #36253
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
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, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.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.