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..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 @@ -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 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"}) @@ -82,6 +80,14 @@ public abstract static class Invoke builder(); @AutoValue.Builder @@ -94,6 +100,14 @@ abstract Builder setHandler( abstract Builder setBatchConfig(BatchElements.BatchConfig batchConfig); + abstract Builder setThrottleDelaySecs(Integer throttleDelaySecs); + + abstract Builder setSamplePeriodMs(Long samplePeriodMs); + + abstract Builder setSampleUpdateMs(Long sampleUpdateMs); + + abstract Builder setOverloadRatio(Double overloadRatio); + abstract Invoke build(); } @@ -113,6 +127,42 @@ public Invoke 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 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 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 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 withOverloadRatio(double overloadRatio) { + checkArgument(overloadRatio > 0, "overloadRatio must be positive"); + return builder().setOverloadRatio(overloadRatio).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -151,10 +201,19 @@ 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(); } @@ -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> 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> 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..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 @@ -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> 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 +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 inputs = new ArrayList<>(); + for (int i = 0; i < 30; i++) { + inputs.add(new TestInput("input" + i)); + } + + 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) + // 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> 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); + } }