From 0bcbc7d0e2e425d87ae17e574dac185fe900c41f Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 1 Jul 2026 15:09:57 -0400 Subject: [PATCH 1/7] [Gemini] Add client-side throttling to Java Remote Inference --- sdks/java/ml/inference/remote/build.gradle | 1 + .../ml/inference/remote/RemoteInference.java | 29 ++++- .../inference/remote/RemoteInferenceTest.java | 102 ++++++++++++++++++ 3 files changed, 131 insertions(+), 1 deletion(-) diff --git a/sdks/java/ml/inference/remote/build.gradle b/sdks/java/ml/inference/remote/build.gradle index 7e7bb61c959c..6c2f4e0b6174 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 193c83d7b3ec..592f3283b58e 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 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; @@ -82,6 +83,8 @@ public abstract static class Invoke builder(); @AutoValue.Builder @@ -94,6 +97,8 @@ abstract Builder setHandler( abstract Builder setBatchConfig(BatchElements.BatchConfig batchConfig); + abstract Builder setThrottleDelaySecs(Integer throttleDelaySecs); + abstract Invoke build(); } @@ -113,6 +118,11 @@ public Invoke withBatchConfig(BatchElements.BatchConfig batchCo return builder().setBatchConfig(batchConfig).build(); } + /** Configures the throttling delay when the client is preemptively throttled. */ + public Invoke withThrottleDelaySecs(int throttleDelaySecs) { + return builder().setThrottleDelaySecs(throttleDelaySecs).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -151,10 +161,13 @@ static class RemoteInferenceFn spec) { this.handlerClass = spec.handler(); this.parameters = spec.parameters(); + this.throttleDelaySecs = spec.throttleDelaySecs() != null ? spec.throttleDelaySecs() : 5; retryHandler = RetryHandler.withDefaults(); } @@ -164,6 +177,8 @@ public void setupHandler() { try { this.modelHandler = handlerClass.getDeclaredConstructor().newInstance(); this.modelHandler.createClient(parameters); + this.throttler = + new ReactiveThrottler(1000, 1000, 2.0, "RemoteInference", throttleDelaySecs); } catch (Exception e) { throw new RuntimeException("Failed to instantiate handler: " + handlerClass.getName(), e); } @@ -172,7 +187,19 @@ public void setupHandler() { @ProcessElement public void processElement(ProcessContext c) throws Exception { Iterable> response = - retryHandler.execute(() -> modelHandler.request(c.element())); + retryHandler.execute( + () -> { + if (throttler != null) { + throttler.throttle(); + } + long reqTime = System.currentTimeMillis(); + Iterable> 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 261b87bfe3a2..65ceae230e17 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 @@ -256,6 +256,30 @@ public Iterable> request(List } } + // Mock handler that fails repeatedly but eventually succeeds to trigger throttling + public static class MockThrottlingHandler + implements BaseModelHandler { + + private int requestCount = 0; + + @Override + public void createClient(TestParameters parameters) {} + + @Override + public Iterable> request(List 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 +641,82 @@ 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 inputs = + Arrays.asList( + new TestInput("input1"), + new TestInput("input2"), + new TestInput("input3"), + new TestInput("input4"), + new TestInput("input5"), + new TestInput("input6"), + new TestInput("input7"), + new TestInput("input8"), + new TestInput("input9"), + new TestInput("input10")); + + PCollection 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>> results = + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockThrottlingHandler.class) + .withBatchConfig(batchConfig) + // Set to 1 second to minimize test wait time while still verifying throttling + .withThrottleDelaySecs(1) + .withParameters(params)); + + PAssert.that(results) + .satisfies( + batches -> { + int totalElements = 0; + for (Iterable> batch : batches) { + totalElements += (int) StreamSupport.stream(batch.spliterator(), false).count(); + } + assertEquals("Expected all 10 elements to succeed", 10, 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 10 elements * 2 + // failures each = 20 failures, + // and overloadRatio=2.0, 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.getCommitted() > 0); + + assertTrue("Expected client-side throttling to trigger at least once", hasThrottled); + } } From ca19c466158502af575569149411150f41d34419 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 1 Jul 2026 15:28:23 -0400 Subject: [PATCH 2/7] Expose other configuration fields --- .../ml/inference/remote/RemoteInference.java | 40 ++++++++++++++++++- .../inference/remote/RemoteInferenceTest.java | 26 +++++------- 2 files changed, 49 insertions(+), 17 deletions(-) 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 592f3283b58e..271da66f1925 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 @@ -85,6 +85,12 @@ public abstract static class Invoke builder(); @AutoValue.Builder @@ -99,6 +105,12 @@ abstract Builder setHandler( abstract Builder setThrottleDelaySecs(Integer throttleDelaySecs); + abstract Builder setSamplePeriodMs(Long samplePeriodMs); + + abstract Builder setSampleUpdateMs(Long sampleUpdateMs); + + abstract Builder setOverloadRatio(Double overloadRatio); + abstract Invoke build(); } @@ -123,6 +135,21 @@ public Invoke withThrottleDelaySecs(int throttleDelaySecs) { return builder().setThrottleDelaySecs(throttleDelaySecs).build(); } + /** Configures the length of history to consider when setting throttling probability. */ + public Invoke withSamplePeriodMs(long samplePeriodMs) { + return builder().setSamplePeriodMs(samplePeriodMs).build(); + } + + /** Configures the granularity of time buckets that we store data in for throttling. */ + public Invoke withSampleUpdateMs(long sampleUpdateMs) { + return builder().setSampleUpdateMs(sampleUpdateMs).build(); + } + + /** Configures the target ratio between requests sent and successful requests. */ + public Invoke withOverloadRatio(double overloadRatio) { + return builder().setOverloadRatio(overloadRatio).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -162,12 +189,18 @@ static class RemoteInferenceFn 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(); } @@ -178,7 +211,12 @@ public void setupHandler() { this.modelHandler = handlerClass.getDeclaredConstructor().newInstance(); this.modelHandler.createClient(parameters); this.throttler = - new ReactiveThrottler(1000, 1000, 2.0, "RemoteInference", throttleDelaySecs); + new ReactiveThrottler( + samplePeriodMs, + sampleUpdateMs, + overloadRatio, + "RemoteInference", + throttleDelaySecs); } catch (Exception e) { throw new RuntimeException("Failed to instantiate handler: " + handlerClass.getName(), e); } 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 65ceae230e17..c2ebca6e8e97 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.assertTrue; import static org.junit.Assert.fail; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -647,18 +648,10 @@ public void testThrottlingBehavior() { TestParameters params = TestParameters.builder().setConfig("test-config").build(); // Create enough inputs to ensure throttling probabilistically triggers - List inputs = - Arrays.asList( - new TestInput("input1"), - new TestInput("input2"), - new TestInput("input3"), - new TestInput("input4"), - new TestInput("input5"), - new TestInput("input6"), - new TestInput("input7"), - new TestInput("input8"), - new TestInput("input9"), - new TestInput("input10")); + List inputs = new ArrayList<>(); + for (int i = 0; i < 30; i++) { + inputs.add(new TestInput("input" + i)); + } PCollection inputCollection = pipeline.apply( @@ -679,6 +672,7 @@ public void testThrottlingBehavior() { .withBatchConfig(batchConfig) // Set to 1 second to minimize test wait time while still verifying throttling .withThrottleDelaySecs(1) + .withOverloadRatio(1.1) .withParameters(params)); PAssert.that(results) @@ -688,7 +682,7 @@ public void testThrottlingBehavior() { for (Iterable> batch : batches) { totalElements += (int) StreamSupport.stream(batch.spliterator(), false).count(); } - assertEquals("Expected all 10 elements to succeed", 10, totalElements); + assertEquals("Expected all 30 elements to succeed", 30, totalElements); return null; }); @@ -709,9 +703,9 @@ public void testThrottlingBehavior() { org.apache.beam.sdk.metrics.Metrics.THROTTLE_TIME_COUNTER_NAME)) .build()); - // Throttling may not trigger if random numbers are very skewed, but with 10 elements * 2 - // failures each = 20 failures, - // and overloadRatio=2.0, the chance of not throttling at least once is very small. + // 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) From 6618866de884cce15b34518131ea4a9a96bb55ef Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:51:17 -0400 Subject: [PATCH 3/7] Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../apache/beam/sdk/ml/inference/remote/RemoteInference.java | 3 +++ .../beam/sdk/ml/inference/remote/RemoteInferenceTest.java | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) 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 271da66f1925..2fe6f6ee1c2c 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 @@ -231,6 +231,9 @@ public void processElement(ProcessContext c) throws Exception { throttler.throttle(); } long reqTime = System.currentTimeMillis(); + if (modelHandler == null) { + throw new IllegalStateException("modelHandler is not initialized"); + } Iterable> result = modelHandler.request(c.element()); if (throttler != null) { 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 c2ebca6e8e97..811a1466aae6 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 @@ -709,7 +709,7 @@ public void testThrottlingBehavior() { // If this test becomes flaky, increase the number of inputs. boolean hasThrottled = StreamSupport.stream(metrics.getCounters().spliterator(), false) - .anyMatch(metricResult -> metricResult.getCommitted() > 0); + .anyMatch(metricResult -> metricResult.getAttempted() > 0 || metricResult.getCommitted() > 0); assertTrue("Expected client-side throttling to trigger at least once", hasThrottled); } From b275e30cddbe57bab162a6efafdc7ca194190291 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 1 Jul 2026 15:55:06 -0400 Subject: [PATCH 4/7] argument validation --- .../apache/beam/sdk/ml/inference/remote/RemoteInference.java | 4 ++++ .../beam/sdk/ml/inference/remote/RemoteInferenceTest.java | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) 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 2fe6f6ee1c2c..369f0bc2985f 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 @@ -132,21 +132,25 @@ public Invoke withBatchConfig(BatchElements.BatchConfig batchCo /** Configures the throttling delay when the client is preemptively throttled. */ public Invoke 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. */ public Invoke 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. */ public Invoke 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. */ public Invoke withOverloadRatio(double overloadRatio) { + checkArgument(overloadRatio > 0, "overloadRatio must be positive"); return builder().setOverloadRatio(overloadRatio).build(); } 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 811a1466aae6..8e7401a93318 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 @@ -709,7 +709,8 @@ public void testThrottlingBehavior() { // 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); + .anyMatch( + metricResult -> metricResult.getAttempted() > 0 || metricResult.getCommitted() > 0); assertTrue("Expected client-side throttling to trigger at least once", hasThrottled); } From e1c75f677ec3763ecb520a7384c8294cc63e3a6f Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 6 Jul 2026 09:38:08 -0400 Subject: [PATCH 5/7] unit test fix --- .../beam/sdk/ml/inference/remote/RemoteInferenceTest.java | 3 +++ 1 file changed, 3 insertions(+) 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 8e7401a93318..64691aa1fedd 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 @@ -670,6 +670,9 @@ public void testThrottlingBehavior() { RemoteInference.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) From 0669eb696af1939adafe0e5ea7d79e5efcfd0ca1 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 14:47:00 -0400 Subject: [PATCH 6/7] Docstrings --- .../ml/inference/remote/RemoteInference.java | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) 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 369f0bc2985f..45f8aa4a9387 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 @@ -130,25 +130,37 @@ public Invoke withBatchConfig(BatchElements.BatchConfig batchCo return builder().setBatchConfig(batchConfig).build(); } - /** Configures the throttling delay when the client is preemptively throttled. */ + /** + * Configures the throttling delay when the client is preemptively throttled. Defaults to 5 + * seconds. For more context, see {@link ReactiveThrottler} + */ public Invoke 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. */ + /** + * 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 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. */ + /** + * 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 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. */ + /** + * 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 withOverloadRatio(double overloadRatio) { checkArgument(overloadRatio > 0, "overloadRatio must be positive"); return builder().setOverloadRatio(overloadRatio).build(); From fee384a489dfaca597b832e55beaf8c9343ceca1 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 14:54:48 -0400 Subject: [PATCH 7/7] disable on 0 throttlingDelaySecs --- .../ml/inference/remote/RemoteInference.java | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) 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 45f8aa4a9387..c571911a484e 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 @@ -49,15 +49,12 @@ * // Apply remote inference transform * PCollection 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>> results = - * inputs.apply( - * RemoteInference.invoke() - * .handler(OpenAIModelHandler.class) - * .withParameters(params) - * ); + * PCollection>> results = inputs.apply( + * RemoteInference.invoke() + * .handler(OpenAIModelHandler.class) + * .withParameters(params)); * } */ @SuppressWarnings({"rawtypes", "unchecked"}) @@ -132,7 +129,7 @@ public Invoke withBatchConfig(BatchElements.BatchConfig batchCo /** * Configures the throttling delay when the client is preemptively throttled. Defaults to 5 - * seconds. For more context, see {@link ReactiveThrottler} + * seconds. A value of 0 disables throttling. For more context, see {@link ReactiveThrottler} */ public Invoke withThrottleDelaySecs(int throttleDelaySecs) { checkArgument(throttleDelaySecs >= 0, "throttleDelaySecs must be non-negative"); @@ -226,17 +223,20 @@ public void setupHandler() { try { this.modelHandler = handlerClass.getDeclaredConstructor().newInstance(); this.modelHandler.createClient(parameters); - this.throttler = - new ReactiveThrottler( - samplePeriodMs, - sampleUpdateMs, - overloadRatio, - "RemoteInference", - throttleDelaySecs); + 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 {