[Gemini] Add client-side throttling to Java Remote Inference (#39194)

jrmccluskey · gemini-code-assist[bot] · web-flow · commit 27d699eb7215 · 2026-07-08T09:08:55.000-04:00
* [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] &lt;176961590+gemini-code-assist[bot]@users.noreply.github.com&gt;

* argument validation

* unit test fix

* Docstrings

* disable on 0 throttlingDelaySecs

---------

Co-authored-by: gemini-code-assist[bot] &lt;176961590+gemini-code-assist[bot]@users.noreply.github.com&gt;
diff --git 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
@@ -21,6 +21,7 @@
 
 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 @@
  * // 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 abstract static class Invoke<InputT extends BaseInput, OutputT extends Ba
 
     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 @@ abstract Builder<InputT, OutputT> setHandler(
 
       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 Invoke<InputT, OutputT> withBatchConfig(BatchElements.BatchConfig batchCo
       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 @@ static class RemoteInferenceFn<InputT extends BaseInput, OutputT extends BaseRes
       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 void setupHandler() {
         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
@@ -23,6 +23,7 @@
 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 Iterable<PredictionResult<TestInput, TestOutput>> request(List<TestInput>
     }
   }
 
+  // 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 void testWithEmptyHandler() {
         "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);
+  }
 }
