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 2f7d95d1cf9 [Gemini] Add batching support to Java Remote Inference
(#39067)
2f7d95d1cf9 is described below
commit 2f7d95d1cf9ff136df394f4f56b0a4b84ba467c1
Author: Jack McCluskey <[email protected]>
AuthorDate: Wed Jul 1 11:52:02 2026 -0400
[Gemini] Add batching support to Java Remote Inference (#39067)
* [Gemini] Add batching support to Java Remote Inference
* Update
sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
* Fix batching test issue
---------
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
.../sdk/ml/inference/remote/RemoteInference.java | 33 +++++++++++-------
.../ml/inference/remote/RemoteInferenceTest.java | 40 ++++++++++++++++------
2 files changed, 49 insertions(+), 24 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 9092fc9910d..193c83d7b3e 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
@@ -20,13 +20,11 @@ package org.apache.beam.sdk.ml.inference.remote;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import com.google.auto.value.AutoValue;
-import java.util.Collections;
import java.util.List;
+import org.apache.beam.sdk.transforms.BatchElements;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.PCollection;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -82,6 +80,8 @@ public class RemoteInference {
abstract @Nullable BaseModelParameters parameters();
+ abstract BatchElements.@Nullable BatchConfig batchConfig();
+
abstract Builder<InputT, OutputT> builder();
@AutoValue.Builder
@@ -92,6 +92,8 @@ public class RemoteInference {
abstract Builder<InputT, OutputT> setParameters(BaseModelParameters
modelParameters);
+ abstract Builder<InputT, OutputT>
setBatchConfig(BatchElements.BatchConfig batchConfig);
+
abstract Invoke<InputT, OutputT> build();
}
@@ -106,21 +108,26 @@ public class RemoteInference {
return builder().setParameters(modelParameters).build();
}
+ /** Configures the batching behavior for the inputs. */
+ public Invoke<InputT, OutputT> withBatchConfig(BatchElements.BatchConfig
batchConfig) {
+ return builder().setBatchConfig(batchConfig).build();
+ }
+
@Override
public PCollection<Iterable<PredictionResult<InputT, OutputT>>> expand(
PCollection<InputT> input) {
checkArgument(handler() != null, "handler() is required");
checkArgument(parameters() != null, "withParameters() is required");
- return input
- .apply(
- "WrapInputInList",
- MapElements.via(
- new SimpleFunction<InputT, List<InputT>>() {
- @Override
- public List<InputT> apply(InputT element) {
- return Collections.singletonList(element);
- }
- }))
+
+ BatchElements.BatchConfig config = batchConfig();
+ PCollection<List<InputT>> batchedInput;
+ if (config != null) {
+ batchedInput = input.apply("BatchElements",
BatchElements.withConfig(config));
+ } else {
+ batchedInput = input.apply("BatchElements",
BatchElements.withDefaults());
+ }
+
+ return batchedInput
// Pass the list to the inference function
.apply("RemoteInference", ParDo.of(new RemoteInferenceFn<InputT,
OutputT>(this)));
}
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 eef47dddfcd..261b87bfe3a 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
@@ -518,36 +518,54 @@ public class RemoteInferenceTest {
pipeline.run().waitUntilFinish();
}
- // Temporary behaviour until we introduce java BatchElements transform
- // to batch elements in RemoteInference
- @Test
- public void testMultipleInputsProduceSeparateBatches() {
- List<TestInput> inputs = Arrays.asList(new TestInput("input1"), new
TestInput("input2"));
+ private static class GenerateInputsFn
+ extends org.apache.beam.sdk.transforms.DoFn<Integer, TestInput> {
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ c.output(new TestInput("input1"));
+ c.output(new TestInput("input2"));
+ }
+ }
+ @Test
+ public void testBatchingProducesCombinedBatches() {
TestParameters params =
TestParameters.builder().setConfig("test-config").build();
+ // Use a single element to trigger generation of inputs within the same
bundle,
+ // ensuring DirectRunner doesn't split them before BatchElements processes
them.
PCollection<TestInput> inputCollection =
- pipeline.apply(
- "CreateInputs",
Create.of(inputs).withCoder(SerializableCoder.of(TestInput.class)));
+ pipeline
+ .apply("CreateTrigger", Create.of(1))
+ .apply(
+ "GenerateInputs", org.apache.beam.sdk.transforms.ParDo.of(new
GenerateInputsFn()))
+ .setCoder(SerializableCoder.of(TestInput.class));
+
+ // Configure BatchElements to force a batch of exactly 2
+ org.apache.beam.sdk.transforms.BatchElements.BatchConfig batchConfig =
+ org.apache.beam.sdk.transforms.BatchElements.BatchConfig.builder()
+ .withMinBatchSize(2)
+ .withMaxBatchSize(2)
+ .build();
PCollection<Iterable<PredictionResult<TestInput, TestOutput>>> results =
inputCollection.apply(
"RemoteInference",
RemoteInference.<TestInput, TestOutput>invoke()
.handler(MockSuccessHandler.class)
+ .withBatchConfig(batchConfig)
.withParameters(params));
PAssert.that(results)
.satisfies(
batches -> {
int batchCount = 0;
+ int totalElements = 0;
for (Iterable<PredictionResult<TestInput, TestOutput>> batch :
batches) {
batchCount++;
- int elementCount = (int)
StreamSupport.stream(batch.spliterator(), false).count();
- // Each batch should contain exactly 1 element
- assertEquals("Each batch should contain 1 element", 1,
elementCount);
+ totalElements += (int)
StreamSupport.stream(batch.spliterator(), false).count();
}
- assertEquals("Expected 2 batches", 2, batchCount);
+ assertEquals("Expected 1 batch", 1, batchCount);
+ assertEquals("Total output elements should be 2", 2,
totalElements);
return null;
});
