This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 27d699eb721 [Gemini] Add client-side throttling to Java Remote
Inference (#39194)
27d699eb721 is described below
commit 27d699eb7215ca4ebda8d61984c17507e2937e87
Author: Jack McCluskey <[email protected]>
AuthorDate: Wed Jul 8 09:08:55 2026 -0400
[Gemini] Add client-side throttling to Java Remote Inference (#39194)
* [Gemini] Add client-side throttling to Java Remote Inference
* Expose other configuration fields
* Apply suggestions from code review
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
* argument validation
* unit test fix
* Docstrings
* disable on 0 throttlingDelaySecs
---------
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
sdks/java/ml/inference/remote/build.gradle | 1 +
.../sdk/ml/inference/remote/RemoteInference.java | 102 +++++++++++++++++++--
.../ml/inference/remote/RemoteInferenceTest.java | 100 ++++++++++++++++++++
3 files changed, 194 insertions(+), 9 deletions(-)
diff --git a/sdks/java/ml/inference/remote/build.gradle
b/sdks/java/ml/inference/remote/build.gradle
index 7e7bb61c959..6c2f4e0b617 100644
--- a/sdks/java/ml/inference/remote/build.gradle
+++ b/sdks/java/ml/inference/remote/build.gradle
@@ -29,6 +29,7 @@ ext.summary = "Base framework for remote ml inference"
dependencies {
// Core Beam SDK
implementation project(path: ":sdks:java:core", configuration: "shadow")
+ implementation project(":sdks:java:io:components")
compileOnly "com.google.auto.value:auto-value-annotations:1.11.0"
annotationProcessor "com.google.auto.value:auto-value:1.11.0"
diff --git
a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
index 193c83d7b3e..c571911a484 100644
---
a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
+++
b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
@@ -21,6 +21,7 @@ import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
import com.google.auto.value.AutoValue;
import java.util.List;
+import org.apache.beam.sdk.io.components.throttling.ReactiveThrottler;
import org.apache.beam.sdk.transforms.BatchElements;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
@@ -48,15 +49,12 @@ import org.checkerframework.checker.nullness.qual.Nullable;
* // Apply remote inference transform
* PCollection<OpenAIModelInput> inputs = pipeline.apply(Create.of(
* OpenAIModelInput.create("An excellent B2B SaaS solution that
streamlines business processes efficiently."),
- * OpenAIModelInput.create("Really impressed with the innovative
features!")
- * ));
+ * OpenAIModelInput.create("Really impressed with the innovative
features!")));
*
- * PCollection<Iterable<PredictionResult<OpenAIModelInput,
OpenAIModelResponse>>> results =
- * inputs.apply(
- * RemoteInference.<OpenAIModelInput, OpenAIModelResponse>invoke()
- * .handler(OpenAIModelHandler.class)
- * .withParameters(params)
- * );
+ * PCollection<Iterable<PredictionResult<OpenAIModelInput,
OpenAIModelResponse>>> results = inputs.apply(
+ * RemoteInference.<OpenAIModelInput, OpenAIModelResponse>invoke()
+ * .handler(OpenAIModelHandler.class)
+ * .withParameters(params));
* }</pre>
*/
@SuppressWarnings({"rawtypes", "unchecked"})
@@ -82,6 +80,14 @@ public class RemoteInference {
abstract BatchElements.@Nullable BatchConfig batchConfig();
+ abstract @Nullable Integer throttleDelaySecs();
+
+ abstract @Nullable Long samplePeriodMs();
+
+ abstract @Nullable Long sampleUpdateMs();
+
+ abstract @Nullable Double overloadRatio();
+
abstract Builder<InputT, OutputT> builder();
@AutoValue.Builder
@@ -94,6 +100,14 @@ public class RemoteInference {
abstract Builder<InputT, OutputT>
setBatchConfig(BatchElements.BatchConfig batchConfig);
+ abstract Builder<InputT, OutputT> setThrottleDelaySecs(Integer
throttleDelaySecs);
+
+ abstract Builder<InputT, OutputT> setSamplePeriodMs(Long samplePeriodMs);
+
+ abstract Builder<InputT, OutputT> setSampleUpdateMs(Long sampleUpdateMs);
+
+ abstract Builder<InputT, OutputT> setOverloadRatio(Double overloadRatio);
+
abstract Invoke<InputT, OutputT> build();
}
@@ -113,6 +127,42 @@ public class RemoteInference {
return builder().setBatchConfig(batchConfig).build();
}
+ /**
+ * Configures the throttling delay when the client is preemptively
throttled. Defaults to 5
+ * seconds. A value of 0 disables throttling. For more context, see {@link
ReactiveThrottler}
+ */
+ public Invoke<InputT, OutputT> withThrottleDelaySecs(int
throttleDelaySecs) {
+ checkArgument(throttleDelaySecs >= 0, "throttleDelaySecs must be
non-negative");
+ return builder().setThrottleDelaySecs(throttleDelaySecs).build();
+ }
+
+ /**
+ * Configures the length of history to consider when setting throttling
probability. Defaults to
+ * a sample period of 1000ms. For more context, see {@link
AdaptiveThrottler}
+ */
+ public Invoke<InputT, OutputT> withSamplePeriodMs(long samplePeriodMs) {
+ checkArgument(samplePeriodMs > 0, "samplePeriodMs must be positive");
+ return builder().setSamplePeriodMs(samplePeriodMs).build();
+ }
+
+ /**
+ * Configures the granularity of time buckets that we store data in for
throttling. Defaults to
+ * a sample period of 1000ms. For more context, see {@link
AdaptiveThrottler}
+ */
+ public Invoke<InputT, OutputT> withSampleUpdateMs(long sampleUpdateMs) {
+ checkArgument(sampleUpdateMs > 0, "sampleUpdateMs must be positive");
+ return builder().setSampleUpdateMs(sampleUpdateMs).build();
+ }
+
+ /**
+ * Configures the target ratio between requests sent and successful
requests. Defaults to an
+ * overload ratio of 2.0. For more context, see {@link AdaptiveThrottler}
+ */
+ public Invoke<InputT, OutputT> withOverloadRatio(double overloadRatio) {
+ checkArgument(overloadRatio > 0, "overloadRatio must be positive");
+ return builder().setOverloadRatio(overloadRatio).build();
+ }
+
@Override
public PCollection<Iterable<PredictionResult<InputT, OutputT>>> expand(
PCollection<InputT> input) {
@@ -151,10 +201,19 @@ public class RemoteInference {
private final BaseModelParameters parameters;
private transient @Nullable BaseModelHandler modelHandler;
private final RetryHandler retryHandler;
+ private final int throttleDelaySecs;
+ private final long samplePeriodMs;
+ private final long sampleUpdateMs;
+ private final double overloadRatio;
+ private transient @Nullable ReactiveThrottler throttler;
RemoteInferenceFn(Invoke<InputT, OutputT> spec) {
this.handlerClass = spec.handler();
this.parameters = spec.parameters();
+ this.throttleDelaySecs = spec.throttleDelaySecs() != null ?
spec.throttleDelaySecs() : 5;
+ this.samplePeriodMs = spec.samplePeriodMs() != null ?
spec.samplePeriodMs() : 1000L;
+ this.sampleUpdateMs = spec.sampleUpdateMs() != null ?
spec.sampleUpdateMs() : 1000L;
+ this.overloadRatio = spec.overloadRatio() != null ?
spec.overloadRatio() : 2.0;
retryHandler = RetryHandler.withDefaults();
}
@@ -164,15 +223,40 @@ public class RemoteInference {
try {
this.modelHandler =
handlerClass.getDeclaredConstructor().newInstance();
this.modelHandler.createClient(parameters);
+ if (throttleDelaySecs > 0) {
+ this.throttler =
+ new ReactiveThrottler(
+ samplePeriodMs,
+ sampleUpdateMs,
+ overloadRatio,
+ "RemoteInference",
+ throttleDelaySecs);
+ }
} catch (Exception e) {
throw new RuntimeException("Failed to instantiate handler: " +
handlerClass.getName(), e);
}
}
+
/** Perform Inference. */
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
Iterable<PredictionResult<InputT, OutputT>> response =
- retryHandler.execute(() -> modelHandler.request(c.element()));
+ retryHandler.execute(
+ () -> {
+ if (throttler != null) {
+ throttler.throttle();
+ }
+ long reqTime = System.currentTimeMillis();
+ if (modelHandler == null) {
+ throw new IllegalStateException("modelHandler is not
initialized");
+ }
+ Iterable<PredictionResult<InputT, OutputT>> result =
+ modelHandler.request(c.element());
+ if (throttler != null) {
+ throttler.successfulRequest(reqTime);
+ }
+ return result;
+ });
c.output(response);
}
}
diff --git
a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
index 261b87bfe3a..64691aa1fed 100644
---
a/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
+++
b/sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
@@ -23,6 +23,7 @@ import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
@@ -256,6 +257,30 @@ public class RemoteInferenceTest {
}
}
+ // Mock handler that fails repeatedly but eventually succeeds to trigger
throttling
+ public static class MockThrottlingHandler
+ implements BaseModelHandler<TestParameters, TestInput, TestOutput> {
+
+ 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());
+ }
+ }
+
private static boolean containsMessage(Throwable e, String message) {
Throwable current = e;
while (current != null) {
@@ -617,4 +642,79 @@ public class RemoteInferenceTest {
"Expected message to contain 'handler() is required', but got: " +
thrown.getMessage(),
thrown.getMessage().contains("handler() is required"));
}
+
+ @Test
+ public void testThrottlingBehavior() {
+ TestParameters params =
TestParameters.builder().setConfig("test-config").build();
+
+ // Create enough inputs to ensure throttling probabilistically triggers
+ List<TestInput> inputs = new ArrayList<>();
+ for (int i = 0; i < 30; i++) {
+ inputs.add(new TestInput("input" + i));
+ }
+
+ PCollection<TestInput> inputCollection =
+ pipeline.apply(
+ "CreateInputs",
Create.of(inputs).withCoder(SerializableCoder.of(TestInput.class)));
+
+ // Configure BatchElements to force a batch of exactly 1 so we get enough
requests
+ org.apache.beam.sdk.transforms.BatchElements.BatchConfig batchConfig =
+ org.apache.beam.sdk.transforms.BatchElements.BatchConfig.builder()
+ .withMinBatchSize(1)
+ .withMaxBatchSize(1)
+ .build();
+
+ PCollection<Iterable<PredictionResult<TestInput, TestOutput>>> results =
+ inputCollection.apply(
+ "RemoteInference",
+ RemoteInference.<TestInput, TestOutput>invoke()
+ .handler(MockThrottlingHandler.class)
+ .withBatchConfig(batchConfig)
+ // Use large sample periods so the 1s retry delay doesn't
flush the history
+ .withSamplePeriodMs(60000L)
+ .withSampleUpdateMs(60000L)
+ // Set to 1 second to minimize test wait time while still
verifying throttling
+ .withThrottleDelaySecs(1)
+ .withOverloadRatio(1.1)
+ .withParameters(params));
+
+ PAssert.that(results)
+ .satisfies(
+ batches -> {
+ int totalElements = 0;
+ for (Iterable<PredictionResult<TestInput, TestOutput>> batch :
batches) {
+ totalElements += (int)
StreamSupport.stream(batch.spliterator(), false).count();
+ }
+ assertEquals("Expected all 30 elements to succeed", 30,
totalElements);
+ return null;
+ });
+
+ org.apache.beam.sdk.PipelineResult result = pipeline.run();
+ result.waitUntilFinish();
+
+ // Verify that the throttling metrics were populated.
+ // The metric name is defined by Metrics.THROTTLE_TIME_COUNTER_NAME which
evaluates to
+ // "cumulativeThrottlingSeconds".
+ org.apache.beam.sdk.metrics.MetricQueryResults metrics =
+ result
+ .metrics()
+ .queryMetrics(
+ org.apache.beam.sdk.metrics.MetricsFilter.builder()
+ .addNameFilter(
+ org.apache.beam.sdk.metrics.MetricNameFilter.named(
+ "RemoteInference",
+
org.apache.beam.sdk.metrics.Metrics.THROTTLE_TIME_COUNTER_NAME))
+ .build());
+
+ // Throttling may not trigger if random numbers are very skewed, but with
30 elements * 2
+ // failures each = 60 failures,
+ // and overloadRatio=1.1, the chance of not throttling at least once is
very small.
+ // If this test becomes flaky, increase the number of inputs.
+ boolean hasThrottled =
+ StreamSupport.stream(metrics.getCounters().spliterator(), false)
+ .anyMatch(
+ metricResult -> metricResult.getAttempted() > 0 ||
metricResult.getCommitted() > 0);
+
+ assertTrue("Expected client-side throttling to trigger at least once",
hasThrottled);
+ }
}
