gemini-code-assist[bot] commented on code in PR #39194:
URL: https://github.com/apache/beam/pull/39194#discussion_r3508659961
##########
sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java:
##########
@@ -113,6 +130,26 @@ public Invoke<InputT, OutputT>
withBatchConfig(BatchElements.BatchConfig batchCo
return builder().setBatchConfig(batchConfig).build();
}
+ /** Configures the throttling delay when the client is preemptively
throttled. */
+ public Invoke<InputT, OutputT> withThrottleDelaySecs(int
throttleDelaySecs) {
+ return builder().setThrottleDelaySecs(throttleDelaySecs).build();
+ }
+
+ /** Configures the length of history to consider when setting throttling
probability. */
+ public Invoke<InputT, OutputT> withSamplePeriodMs(long samplePeriodMs) {
+ return builder().setSamplePeriodMs(samplePeriodMs).build();
+ }
+
+ /** Configures the granularity of time buckets that we store data in for
throttling. */
+ public Invoke<InputT, OutputT> withSampleUpdateMs(long sampleUpdateMs) {
+ return builder().setSampleUpdateMs(sampleUpdateMs).build();
+ }
+
+ /** Configures the target ratio between requests sent and successful
requests. */
+ public Invoke<InputT, OutputT> withOverloadRatio(double overloadRatio) {
+ return builder().setOverloadRatio(overloadRatio).build();
+ }
Review Comment:
To ensure pipeline safety and catch configuration errors during pipeline
construction rather than at runtime on the workers, we should validate the
throttling parameters in the builder methods.
```java
public Invoke<InputT, OutputT> withThrottleDelaySecs(int
throttleDelaySecs) {
com.google.common.base.Preconditions.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<InputT, OutputT> withSamplePeriodMs(long samplePeriodMs) {
com.google.common.base.Preconditions.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<InputT, OutputT> withSampleUpdateMs(long sampleUpdateMs) {
com.google.common.base.Preconditions.checkArgument(
sampleUpdateMs > 0, "sampleUpdateMs must be positive");
return builder().setSampleUpdateMs(sampleUpdateMs).build();
}
/** Configures the target ratio between requests sent and successful
requests. */
public Invoke<InputT, OutputT> withOverloadRatio(double overloadRatio) {
com.google.common.base.Preconditions.checkArgument(
overloadRatio > 0.0, "overloadRatio must be positive");
return builder().setOverloadRatio(overloadRatio).build();
}
```
##########
sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java:
##########
@@ -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());
+ }
Review 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`.
```java
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());
}
```
##########
sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java:
##########
@@ -617,4 +642,75 @@ 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)
+ // 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.getCommitted() > 0);
Review Comment:
Some runners (such as the DirectRunner under certain configurations) might
not populate committed metrics, which can cause `getCommitted()` to return 0
and lead to flaky test failures. Checking both `getAttempted()` and
`getCommitted()` is much more robust.
```suggestion
boolean hasThrottled =
StreamSupport.stream(metrics.getCounters().spliterator(), false)
.anyMatch(metricResult -> metricResult.getAttempted() > 0 ||
metricResult.getCommitted() > 0);
```
##########
sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java:
##########
@@ -172,7 +225,19 @@ public void setupHandler() {
@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();
+ Iterable<PredictionResult<InputT, OutputT>> result =
+ modelHandler.request(c.element());
Review Comment:
The `modelHandler` field is marked as `@Nullable` and `transient`. To adhere
to defensive programming best practices and avoid a raw `NullPointerException`,
we should check if `modelHandler` is null before calling `request`.
```suggestion
if (modelHandler == null) {
throw new IllegalStateException("modelHandler is not
initialized");
}
Iterable<PredictionResult<InputT, OutputT>> result =
modelHandler.request(c.element());
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
