From c2542c53d3eb0f2dba7e488cf4ef998aed412209 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 09:51:07 -0400 Subject: [PATCH 1/6] [Gemini] Add Java GeminiModelHandler class --- sdks/java/ml/inference/gemini/build.gradle | 45 +++++ .../gemini/GeminiInferenceFunctions.java | 64 +++++++ .../inference/gemini/GeminiModelHandler.java | 84 +++++++++ .../gemini/GeminiModelParameters.java | 61 +++++++ .../gemini/GeminiRequestFunction.java | 28 +++ .../gemini/GeminiModelHandlerIT.java | 162 ++++++++++++++++++ .../gemini/GeminiModelHandlerTest.java | 124 ++++++++++++++ .../ml/inference/gemini/TestReflection.class | Bin 0 -> 1684 bytes .../ml/inference/gemini/TestReflection.java | 39 +++++ .../ml/inference/remote/BaseModelHandler.java | 3 +- .../ml/inference/remote/RemoteInference.java | 9 +- settings.gradle.kts | 1 + 12 files changed, 613 insertions(+), 7 deletions(-) create mode 100644 sdks/java/ml/inference/gemini/build.gradle create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class create mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java diff --git a/sdks/java/ml/inference/gemini/build.gradle b/sdks/java/ml/inference/gemini/build.gradle new file mode 100644 index 000000000000..c7248c7d91d8 --- /dev/null +++ b/sdks/java/ml/inference/gemini/build.gradle @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +plugins { + id 'org.apache.beam.module' +} + +applyJavaNature( + automaticModuleName: 'org.apache.beam.sdk.ml.inference.gemini', + requireJavaVersion: JavaVersion.VERSION_11 +) +provideIntegrationTestingDependencies() +enableJavaPerformanceTesting() + +description = "Apache Beam :: SDKs :: Java :: ML :: Inference :: Gemini" +ext.summary = "Gemini model handler for remote inference" + +dependencies { + implementation project(":sdks:java:ml:inference:remote") + implementation "com.google.genai:google-genai:1.59.0" + implementation library.java.jackson_databind + implementation library.java.jackson_annotations + implementation library.java.jackson_core + + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") + testImplementation project(path: ":sdks:java:core", configuration: "shadow") + testImplementation library.java.slf4j_api + testRuntimeOnly library.java.slf4j_simple + testImplementation library.java.junit + testImplementation library.java.mockito_core +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java new file mode 100644 index 000000000000..a76a7b9d8f90 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiInferenceFunctions.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; +import com.google.genai.types.GenerateImagesConfig; +import com.google.genai.types.GenerateImagesResponse; +import java.util.ArrayList; +import java.util.List; + +/** Common inference functions for Gemini. */ +public class GeminiInferenceFunctions { + + /** Generates content from string prompts using the standard generateContent API. */ + public static GeminiRequestFunction generateFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (String input : batch) { + GenerateContentResponse response = + client.models.generateContent( + modelName, input, GenerateContentConfig.builder().build()); + String text = response.text(); + results.add(text != null ? text : ""); + } + return results; + }; + } + + /** Generates images from string prompts using the generateImages API. */ + public static GeminiRequestFunction generateImageFromString() { + return (modelName, batch, client) -> { + List results = new ArrayList<>(); + for (String input : batch) { + GenerateImagesResponse response = + client.models.generateImages(modelName, input, GenerateImagesConfig.builder().build()); + // Retrieve the base64 string or bytes from the first generated image + List images = response.images(); + if (images != null && !images.isEmpty()) { + byte[] imageBytes = images.get(0).imageBytes().orElse(new byte[0]); + results.add(imageBytes); + } else { + results.add(new byte[0]); + } + } + return results; + }; + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java new file mode 100644 index 000000000000..3e1aafee3293 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.Client; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.BaseModelHandler; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; + +/** + * Model handler for Google Gemini API inference requests. + * + *

This handler manages communication with Google's Gemini API, including client initialization, + * request formatting, and response parsing. It allows executing a custom {@link + * GeminiRequestFunction} against a batch of inputs. + */ +@SuppressWarnings("nullness") +public class GeminiModelHandler + implements BaseModelHandler, InputT, OutputT> { + + private transient Client client; + private GeminiModelParameters modelParameters; + + @Override + public void createClient(GeminiModelParameters parameters) { + this.modelParameters = parameters; + + // Configure client based on vertex or API key + if (parameters.getApiKey() != null) { + if (parameters.getProject() != null || parameters.getLocation() != null) { + throw new IllegalArgumentException("Project and location must be null if API key is set"); + } + this.client = Client.builder().apiKey(parameters.getApiKey()).build(); + } else { + if (parameters.getProject() == null || parameters.getLocation() == null) { + throw new IllegalArgumentException( + "Project and location must both be provided if API key is not set"); + } + Client.Builder builder = + Client.builder() + .vertexAI(true) + .project(parameters.getProject()) + .location(parameters.getLocation()); + + this.client = builder.build(); + } + } + + @Override + public Iterable> request(List input) { + try { + GeminiRequestFunction requestFn = modelParameters.getRequestFn(); + List responses = requestFn.apply(modelParameters.getModelName(), input, client); + + if (responses.size() != input.size()) { + throw new IllegalStateException("Number of responses must match number of inputs"); + } + + List> results = new ArrayList<>(); + for (int i = 0; i < input.size(); i++) { + results.add(PredictionResult.create(input.get(i), responses.get(i))); + } + return results; + } catch (Exception e) { + throw new RuntimeException("Error during Gemini inference request", e); + } + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java new file mode 100644 index 000000000000..0f99a0720e22 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.auto.value.AutoValue; +import org.apache.beam.sdk.ml.inference.remote.BaseModelParameters; +import org.checkerframework.checker.nullness.qual.Nullable; + +@AutoValue +public abstract class GeminiModelParameters implements BaseModelParameters { + + public abstract @Nullable String getApiKey(); + + public abstract @Nullable String getProject(); + + public abstract @Nullable String getLocation(); + + public abstract String getModelName(); + + public abstract boolean getUseVertexFlexApi(); + + public abstract GeminiRequestFunction getRequestFn(); + + public static Builder builder() { + return new AutoValue_GeminiModelParameters.Builder() + .setUseVertexFlexApi(false); + } + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setApiKey(String apiKey); + + public abstract Builder setProject(String project); + + public abstract Builder setLocation(String location); + + public abstract Builder setModelName(String modelName); + + public abstract Builder setUseVertexFlexApi(boolean useVertexFlexApi); + + public abstract Builder setRequestFn( + GeminiRequestFunction requestFn); + + public abstract GeminiModelParameters build(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java new file mode 100644 index 000000000000..044e2047d07c --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiRequestFunction.java @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.Client; +import java.io.Serializable; +import java.util.List; + +/** Functional interface for custom request functions to the Gemini API. */ +@FunctionalInterface +public interface GeminiRequestFunction extends Serializable { + List apply(String modelName, List batch, Client client) throws Exception; +} diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java new file mode 100644 index 000000000000..688dc81b45a5 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java @@ -0,0 +1,162 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeNotNull; +import static org.junit.Assume.assumeTrue; + +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.apache.beam.sdk.ml.inference.remote.RemoteInference; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +/** + * Execute Gemini model handler integration test. + * + * 

+ * ./gradlew :sdks:java:ml:inference:gemini:integrationTest \
+ *   --tests org.apache.beam.sdk.ml.inference.gemini.GeminiModelHandlerIT \
+ *   --info
+ *
+ */ +public class GeminiModelHandlerIT { + + public static class GeminiStringContentHandler extends GeminiModelHandler {} + + public static class GeminiStringImageHandler extends GeminiModelHandler {} + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private String apiKey; + private static final String API_KEY_ENV = "GEMINI_API_KEY"; + private static final String DEFAULT_MODEL = "gemini-3.5-flash"; + + @Before + public void setUp() { + // Get API key + apiKey = System.getenv(API_KEY_ENV); + + // Skip tests if API key is not provided + assumeNotNull( + "Gemini API key not found. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + apiKey); + assumeTrue( + "Gemini API key is empty. Set " + + API_KEY_ENV + + " environment variable to run integration tests.", + !apiKey.trim().isEmpty()); + } + + @Test + public void testSentimentAnalysisWithSingleInput() { + String input = + "Classify the sentiment of the following text as either positive or negative: This product is absolutely amazing! I love it!"; + + PCollection inputs = pipeline.apply("CreateSingleInput", Create.of(input)); + + PCollection>> results = + inputs + .apply( + "SentimentInference", + RemoteInference.invoke() + .handler(GeminiStringContentHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName(DEFAULT_MODEL) + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build())) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + + String sentiment = result.getOutput().toLowerCase(); + assertTrue( + "Sentiment should be positive or negative, got: " + sentiment, + sentiment.contains("positive") || sentiment.contains("negative")); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } + + @Test + public void testGenerateImageFromString() { + String input = "A beautiful sunset over the mountains, digital art style."; + + PCollection inputs = pipeline.apply("CreateImageInput", Create.of(input)); + + PCollection>> results = + inputs + .apply( + "ImageInference", + RemoteInference.invoke() + .handler(GeminiStringImageHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName("imagen-4.0-generate-001") + .setRequestFn(GeminiInferenceFunctions.generateImageFromString()) + .build())) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + + // Verify results + PAssert.that(results) + .satisfies( + batches -> { + int count = 0; + for (Iterable> batch : batches) { + for (PredictionResult result : batch) { + count++; + assertNotNull("Input should not be null", result.getInput()); + assertNotNull("Output should not be null", result.getOutput()); + assertTrue("Output should have generated images", result.getOutput().length > 0); + } + } + assertEquals("Should have exactly 1 result", 1, count); + return null; + }); + + pipeline.run().waitUntilFinish(); + } +} diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java new file mode 100644 index 000000000000..340066db9f9b --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.ml.inference.remote.PredictionResult; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GeminiModelHandlerTest { + + @Test + public void testAllParamsSet() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setProject("test-project") + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexLocationParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setProject("test-project") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingVertexProjectParam() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setLocation("us-central1") + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testMissingAllParams() { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setModelName("gemini-model-123") + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + } + + @Test + public void testRequest() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn((modelName, batch, client) -> Arrays.asList("response1", "response2")) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = Arrays.asList("input1", "input2"); + Iterable> results = handler.request(input); + + int count = 0; + for (PredictionResult result : results) { + if (count == 0) { + assertEquals("input1", result.getInput()); + assertEquals("response1", result.getOutput()); + } else { + assertEquals("input2", result.getInput()); + assertEquals("response2", result.getOutput()); + } + count++; + } + assertEquals(2, count); + } + + @Test + public void testRequestMismatchedResponseSize() throws Exception { + GeminiModelParameters parameters = + GeminiModelParameters.builder() + .setApiKey("test-key") + .setModelName("gemini-model-123") + .setRequestFn((modelName, batch, client) -> Arrays.asList("response1")) + .build(); + GeminiModelHandler handler = new GeminiModelHandler<>(); + handler.createClient(parameters); + + List input = Arrays.asList("input1", "input2"); + assertThrows(RuntimeException.class, () -> handler.request(input)); + } +} diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class new file mode 100644 index 0000000000000000000000000000000000000000..dcbde61b6181132f3ab56d3f8886741ce7ee6141 GIT binary patch literal 1684 zcmaJ>O>-MX5Pc)bURl{#u@onE5+w>Ku@n<8HekSVOkyP_4zlA!j>E@+8m}huM*E@K z9i{M{Do$|bOcli$4o;O1`2pn0FX6<=(6hU-SA}2?vpq9CuV26Jnct5NUjev-FLOvC zts$dh1fvXRp7R&nu({hY9_~IDreqjhwp>fDFr=%soiSuFrXi=}6!Zc2)_x#`Ltb8= z4ox;gAj>nhe9M(B=?m`EvlwSM->@y=%DL6PWw(W2c7)vX+QCu|XE33mpreSh48Mj=A3gm1Z>dTmF;Ifg5y=NKK&>)68R2$x$%6tWVB%wY=C8ZPL#h)WFl zlbjni4+4f3D%$^!vule6{gJ< z)i_YGT1!BOiy>k~$2DB1bm`0{cLc-P0na#J6}+mWh8q$57Wc_n3cum?U8(|AY`s8r zen-bm%#rjC?}~=!np{4$tSA2$|} zZ%J-;H+V1f*RaZPG3h#3Z=#sH_l+rflA;D0?!D>guwRdN$gAb`eN(JiibHXb?VRc+ zx@D`LCj;s8Ufh&fJfPyTGRvaLaH?GPLs6H?EI*3P93Emz!y~2YV}{B}OIq#=uPclw zsi+UuxG6nLgYM4yqy#aZidCJu>xVm1Bqo zx9yfC1qrFIyDs&#x|@P}Rg15&h8>1$|H^I{X$_wgJodg-DJ=_W7HtEj#dbtk=x zG%nMuM1MyAksB|;ekFt97L96Kjbm8BCv-%){)sdl1b+ApTH*8o@`cF*oNvy)x|Loj zUHS<(XG;r9nd!^{N`=(Vm?Aedd0P!LhnNNE%(rRw%in(}1VbUr4RSlvWNi^qw!1y^3)>KptPyV{Bmp5(Rum^Y1Z<+jQar zMvqX`aEGz#Ul=(;A++boMuYY+9k88CtYfptVWNW1C_= K!Bc#J>3;ykg1B-3 literal 0 HcmV?d00001 diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java new file mode 100644 index 000000000000..b799ce743e6e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.gemini; + +import com.google.genai.Client; +import java.lang.reflect.Method; + +public class TestReflection { + public static void main(String[] args) throws Exception { + System.out.println("Client.Builder methods:"); + for (Method m : Client.Builder.class.getMethods()) { + if (m.getDeclaringClass() != Object.class) { + System.out.println(m.getName() + " " + m.getParameterCount()); + } + } + System.out.println("Client methods:"); + for (Method m : Client.class.getMethods()) { + if (m.getDeclaringClass() != Object.class) { + System.out.println(m.getName()); + System.out.println(" return: " + m.getReturnType().getName()); + } + } + } +} diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/BaseModelHandler.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/BaseModelHandler.java index 1a52703c745b..c7a0b5c2ca72 100644 --- a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/BaseModelHandler.java +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/BaseModelHandler.java @@ -56,8 +56,7 @@ *
  • Maintain input-output correspondence in {@link PredictionResult} * */ -public interface BaseModelHandler< - ParamT extends BaseModelParameters, InputT extends BaseInput, OutputT extends BaseResponse> { +public interface BaseModelHandler { /** Initializes the remote model client with the provided parameters. */ void createClient(ParamT parameters); 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 c571911a484e..2301b8df7ff8 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 @@ -61,8 +61,7 @@ public class RemoteInference { /** Invoke the model handler with model parameters. */ - public static - Invoke invoke() { + public static Invoke invoke() { return new AutoValue_RemoteInference_Invoke.Builder() // .setParameters(null) .build(); } @@ -70,7 +69,7 @@ Invoke invoke() { private RemoteInference() {} @AutoValue - public abstract static class Invoke + public abstract static class Invoke extends PTransform< PCollection, PCollection>>> { @@ -91,7 +90,7 @@ public abstract static class Invoke builder(); @AutoValue.Builder - abstract static class Builder { + abstract static class Builder { abstract Builder setHandler( Class> modelHandler); @@ -194,7 +193,7 @@ public PCollection>> expand( * */ @SuppressWarnings("nullness") - static class RemoteInferenceFn + static class RemoteInferenceFn extends DoFn, Iterable>> { private final Class> handlerClass; diff --git a/settings.gradle.kts b/settings.gradle.kts index d9dbbf9021ef..9941ad00726c 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -284,6 +284,7 @@ include(":sdks:java:maven-archetypes:gcp-bom-examples") include(":sdks:java:maven-archetypes:starter") include(":sdks:java:ml:inference:remote") include(":sdks:java:ml:inference:openai") +include(":sdks:java:ml:inference:gemini") include(":sdks:java:testing:nexmark") include(":sdks:java:testing:expansion-service") include(":sdks:java:testing:jpms-tests") From f1d81e53614cbb93de9c361f94f7790660828972 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 10:04:42 -0400 Subject: [PATCH 2/6] Remove model scratch files --- .../ml/inference/gemini/TestReflection.class | Bin 1684 -> 0 bytes .../ml/inference/gemini/TestReflection.java | 39 ------------------ 2 files changed, 39 deletions(-) delete mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class delete mode 100644 sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.class deleted file mode 100644 index dcbde61b6181132f3ab56d3f8886741ce7ee6141..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1684 zcmaJ>O>-MX5Pc)bURl{#u@onE5+w>Ku@n<8HekSVOkyP_4zlA!j>E@+8m}huM*E@K z9i{M{Do$|bOcli$4o;O1`2pn0FX6<=(6hU-SA}2?vpq9CuV26Jnct5NUjev-FLOvC zts$dh1fvXRp7R&nu({hY9_~IDreqjhwp>fDFr=%soiSuFrXi=}6!Zc2)_x#`Ltb8= z4ox;gAj>nhe9M(B=?m`EvlwSM->@y=%DL6PWw(W2c7)vX+QCu|XE33mpreSh48Mj=A3gm1Z>dTmF;Ifg5y=NKK&>)68R2$x$%6tWVB%wY=C8ZPL#h)WFl zlbjni4+4f3D%$^!vule6{gJ< z)i_YGT1!BOiy>k~$2DB1bm`0{cLc-P0na#J6}+mWh8q$57Wc_n3cum?U8(|AY`s8r zen-bm%#rjC?}~=!np{4$tSA2$|} zZ%J-;H+V1f*RaZPG3h#3Z=#sH_l+rflA;D0?!D>guwRdN$gAb`eN(JiibHXb?VRc+ zx@D`LCj;s8Ufh&fJfPyTGRvaLaH?GPLs6H?EI*3P93Emz!y~2YV}{B}OIq#=uPclw zsi+UuxG6nLgYM4yqy#aZidCJu>xVm1Bqo zx9yfC1qrFIyDs&#x|@P}Rg15&h8>1$|H^I{X$_wgJodg-DJ=_W7HtEj#dbtk=x zG%nMuM1MyAksB|;ekFt97L96Kjbm8BCv-%){)sdl1b+ApTH*8o@`cF*oNvy)x|Loj zUHS<(XG;r9nd!^{N`=(Vm?Aedd0P!LhnNNE%(rRw%in(}1VbUr4RSlvWNi^qw!1y^3)>KptPyV{Bmp5(Rum^Y1Z<+jQar zMvqX`aEGz#Ul=(;A++boMuYY+9k88CtYfptVWNW1C_= K!Bc#J>3;ykg1B-3 diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java deleted file mode 100644 index b799ce743e6e..000000000000 --- a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/TestReflection.java +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.sdk.ml.inference.gemini; - -import com.google.genai.Client; -import java.lang.reflect.Method; - -public class TestReflection { - public static void main(String[] args) throws Exception { - System.out.println("Client.Builder methods:"); - for (Method m : Client.Builder.class.getMethods()) { - if (m.getDeclaringClass() != Object.class) { - System.out.println(m.getName() + " " + m.getParameterCount()); - } - } - System.out.println("Client methods:"); - for (Method m : Client.class.getMethods()) { - if (m.getDeclaringClass() != Object.class) { - System.out.println(m.getName()); - System.out.println(" return: " + m.getReturnType().getName()); - } - } - } -} From 799de4d3cc19dfb02c8b3297bd7cff152b04626d Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 11:03:34 -0400 Subject: [PATCH 3/6] address review comments --- .../ml/inference/gemini/GeminiModelHandler.java | 16 ++++++++-------- .../inference/gemini/GeminiModelParameters.java | 7 +------ .../inference/gemini/GeminiModelHandlerTest.java | 9 ++------- 3 files changed, 11 insertions(+), 21 deletions(-) diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java index 3e1aafee3293..14a8fc7785c4 100644 --- a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandler.java @@ -39,6 +39,9 @@ public class GeminiModelHandler @Override public void createClient(GeminiModelParameters parameters) { + if (parameters == null) { + throw new NullPointerException("GeminiModelParameters must not be null"); + } this.modelParameters = parameters; // Configure client based on vertex or API key @@ -48,16 +51,13 @@ public void createClient(GeminiModelParameters parameters) { } this.client = Client.builder().apiKey(parameters.getApiKey()).build(); } else { - if (parameters.getProject() == null || parameters.getLocation() == null) { + Client.Builder builder = Client.builder(); + if (parameters.getProject() != null && parameters.getLocation() != null) { + builder.vertexAI(true).project(parameters.getProject()).location(parameters.getLocation()); + } else if (parameters.getProject() != null || parameters.getLocation() != null) { throw new IllegalArgumentException( - "Project and location must both be provided if API key is not set"); + "Project and location must both be provided if one is provided"); } - Client.Builder builder = - Client.builder() - .vertexAI(true) - .project(parameters.getProject()) - .location(parameters.getLocation()); - this.client = builder.build(); } } diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java index 0f99a0720e22..24f377afdbf3 100644 --- a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java @@ -32,13 +32,10 @@ public abstract class GeminiModelParameters implements BaseMode public abstract String getModelName(); - public abstract boolean getUseVertexFlexApi(); - public abstract GeminiRequestFunction getRequestFn(); public static Builder builder() { - return new AutoValue_GeminiModelParameters.Builder() - .setUseVertexFlexApi(false); + return new AutoValue_GeminiModelParameters.Builder(); } @AutoValue.Builder @@ -51,8 +48,6 @@ public abstract static class Builder { public abstract Builder setModelName(String modelName); - public abstract Builder setUseVertexFlexApi(boolean useVertexFlexApi); - public abstract Builder setRequestFn( GeminiRequestFunction requestFn); diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java index 340066db9f9b..c69b2433747e 100644 --- a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java @@ -69,14 +69,9 @@ public void testMissingVertexProjectParam() { } @Test - public void testMissingAllParams() { - GeminiModelParameters parameters = - GeminiModelParameters.builder() - .setModelName("gemini-model-123") - .setRequestFn(GeminiInferenceFunctions.generateFromString()) - .build(); + public void testNullParameters() { GeminiModelHandler handler = new GeminiModelHandler<>(); - assertThrows(IllegalArgumentException.class, () -> handler.createClient(parameters)); + assertThrows(NullPointerException.class, () -> handler.createClient(null)); } @Test From 703c158f8db74ad24e756be00bfc6f71f075e6f5 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 12:05:03 -0400 Subject: [PATCH 4/6] checkstyle issue --- .../sdk/ml/inference/gemini/package-info.java | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java diff --git a/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java new file mode 100644 index 000000000000..9dee9da9cc74 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/package-info.java @@ -0,0 +1,20 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** Gemini model handler for remote inference. */ +package org.apache.beam.sdk.ml.inference.gemini; From 8aee688ba01d837d3624c3f5b1a31677f395b7ed Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 13:29:04 -0400 Subject: [PATCH 5/6] unit test fix --- .../inference/remote/RemoteInferenceTest.java | 142 ++++++++++-------- 1 file changed, 82 insertions(+), 60 deletions(-) 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 64691aa1fedd..29f059b0595c 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 @@ -29,6 +29,8 @@ import java.util.List; import java.util.stream.Collectors; import java.util.stream.StreamSupport; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; @@ -45,7 +47,7 @@ public class RemoteInferenceTest { @Rule public final transient TestPipeline pipeline = TestPipeline.create(); // Test input class - public static class TestInput implements BaseInput { + public static class TestInput implements BaseInput, java.io.Serializable { private final String value; private TestInput(String value) { @@ -84,7 +86,7 @@ public String toString() { } // Test output class - public static class TestOutput implements BaseResponse { + public static class TestOutput implements BaseResponse, java.io.Serializable { private final String result; private TestOutput(String result) { @@ -300,11 +302,13 @@ public void testInvokeWithSingleElement() { PCollection inputCollection = pipeline.apply(Create.of(input)); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Verify the output contains expected predictions PAssert.thatSingleton(results) @@ -337,11 +341,13 @@ public void testInvokeWithMultipleElements() { "CreateInputs", Create.of(inputs).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Count total results across all batches PAssert.that(results) @@ -373,11 +379,13 @@ public void testInvokeWithEmptyCollection() { pipeline.apply("CreateEmptyInput", Create.empty(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // assertion for empty PCollection PAssert.that(results).empty(); @@ -395,11 +403,13 @@ public void testHandlerReturnsEmptyResults() { "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockEmptyResultHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockEmptyResultHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Verify we still get a result, but it's empty PAssert.thatSingleton(results) @@ -423,11 +433,13 @@ public void testHandlerSetupFailure() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockFailingSetupHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockFailingSetupHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Verify pipeline fails with expected error try { @@ -452,11 +464,13 @@ public void testHandlerRequestFailure() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockFailingRequestHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockFailingRequestHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Verify pipeline fails with expected error try { @@ -479,11 +493,13 @@ public void testHandlerWithoutDefaultConstructor() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockNoDefaultConstructorHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockNoDefaultConstructorHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); // Verify pipeline fails when handler cannot be instantiated try { @@ -519,11 +535,13 @@ public void testPredictionResultMapping() { "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); PAssert.thatSingleton(results) .satisfies( @@ -573,12 +591,14 @@ public void testBatchingProducesCombinedBatches() { .build(); PCollection>> results = - inputCollection.apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withBatchConfig(batchConfig) - .withParameters(params)); + inputCollection + .apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withBatchConfig(batchConfig) + .withParameters(params)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); PAssert.that(results) .satisfies( @@ -665,18 +685,20 @@ public void testThrottlingBehavior() { .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)); + 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)) + .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); PAssert.that(results) .satisfies( From 9309eca698996121c3c69fffbc4bced7e15263b0 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 15:11:20 -0400 Subject: [PATCH 6/6] Fix weird coder UX --- .../gemini/GeminiModelHandlerIT.java | 47 +++--- .../remote/PredictionResultCoder.java | 74 ++++++++++ .../ml/inference/remote/RemoteInference.java | 34 ++++- .../inference/remote/RemoteInferenceTest.java | 138 ++++++++---------- 4 files changed, 183 insertions(+), 110 deletions(-) create mode 100644 sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java diff --git a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java index 688dc81b45a5..42892580fd01 100644 --- a/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java @@ -23,9 +23,6 @@ import static org.junit.Assume.assumeNotNull; import static org.junit.Assume.assumeTrue; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.IterableCoder; -import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.ml.inference.remote.PredictionResult; import org.apache.beam.sdk.ml.inference.remote.RemoteInference; import org.apache.beam.sdk.testing.PAssert; @@ -83,18 +80,16 @@ public void testSentimentAnalysisWithSingleInput() { PCollection inputs = pipeline.apply("CreateSingleInput", Create.of(input)); PCollection>> results = - inputs - .apply( - "SentimentInference", - RemoteInference.invoke() - .handler(GeminiStringContentHandler.class) - .withParameters( - GeminiModelParameters.builder() - .setApiKey(apiKey) - .setModelName(DEFAULT_MODEL) - .setRequestFn(GeminiInferenceFunctions.generateFromString()) - .build())) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputs.apply( + "SentimentInference", + RemoteInference.invoke() + .handler(GeminiStringContentHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName(DEFAULT_MODEL) + .setRequestFn(GeminiInferenceFunctions.generateFromString()) + .build())); // Verify results PAssert.that(results) @@ -127,18 +122,16 @@ public void testGenerateImageFromString() { PCollection inputs = pipeline.apply("CreateImageInput", Create.of(input)); PCollection>> results = - inputs - .apply( - "ImageInference", - RemoteInference.invoke() - .handler(GeminiStringImageHandler.class) - .withParameters( - GeminiModelParameters.builder() - .setApiKey(apiKey) - .setModelName("imagen-4.0-generate-001") - .setRequestFn(GeminiInferenceFunctions.generateImageFromString()) - .build())) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputs.apply( + "ImageInference", + RemoteInference.invoke() + .handler(GeminiStringImageHandler.class) + .withParameters( + GeminiModelParameters.builder() + .setApiKey(apiKey) + .setModelName("imagen-4.0-generate-001") + .setRequestFn(GeminiInferenceFunctions.generateImageFromString()) + .build())); // Verify results PAssert.that(results) diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java new file mode 100644 index 000000000000..0d43870b880d --- /dev/null +++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/PredictionResultCoder.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.ml.inference.remote; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.StructuredCoder; + +/** A {@link Coder} for {@link PredictionResult}. */ +public class PredictionResultCoder + extends StructuredCoder> { + + private final Coder inputCoder; + private final Coder outputCoder; + + private PredictionResultCoder(Coder inputCoder, Coder outputCoder) { + this.inputCoder = inputCoder; + this.outputCoder = outputCoder; + } + + public static PredictionResultCoder of( + Coder inputCoder, Coder outputCoder) { + return new PredictionResultCoder<>(inputCoder, outputCoder); + } + + @Override + public void encode(PredictionResult value, OutputStream outStream) + throws CoderException, IOException { + if (value == null) { + throw new CoderException("cannot encode a null PredictionResult"); + } + inputCoder.encode(value.getInput(), outStream); + outputCoder.encode(value.getOutput(), outStream); + } + + @Override + public PredictionResult decode(InputStream inStream) + throws CoderException, IOException { + InputT input = inputCoder.decode(inStream); + OutputT output = outputCoder.decode(inStream); + return PredictionResult.create(input, output); + } + + @Override + public List> getCoderArguments() { + return Arrays.asList(inputCoder, outputCoder); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + verifyDeterministic(this, "Input coder must be deterministic", inputCoder); + verifyDeterministic(this, "Output coder must be deterministic", outputCoder); + } +} 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 2301b8df7ff8..78ec688241ef 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.coders.Coder; import org.apache.beam.sdk.io.components.throttling.ReactiveThrottler; import org.apache.beam.sdk.transforms.BatchElements; import org.apache.beam.sdk.transforms.DoFn; @@ -87,6 +88,8 @@ public abstract static class Invoke abstract @Nullable Double overloadRatio(); + abstract @Nullable Coder outputCoder(); + abstract Builder builder(); @AutoValue.Builder @@ -107,6 +110,8 @@ abstract Builder setHandler( abstract Builder setOverloadRatio(Double overloadRatio); + abstract Builder setOutputCoder(Coder outputCoder); + abstract Invoke build(); } @@ -162,6 +167,14 @@ public Invoke withOverloadRatio(double overloadRatio) { return builder().setOverloadRatio(overloadRatio).build(); } + /** + * Configures the coder for the output of the model. If not provided, it will fallback to using + * standard Java serialization for the output element. + */ + public Invoke withOutputCoder(Coder outputCoder) { + return builder().setOutputCoder(outputCoder).build(); + } + @Override public PCollection>> expand( PCollection input) { @@ -176,9 +189,24 @@ public PCollection>> expand( batchedInput = input.apply("BatchElements", BatchElements.withDefaults()); } - return batchedInput - // Pass the list to the inference function - .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + PCollection>> result = + batchedInput + // Pass the list to the inference function + .apply("RemoteInference", ParDo.of(new RemoteInferenceFn(this))); + + Coder outCoder = outputCoder(); + if (outCoder != null) { + result.setCoder( + org.apache.beam.sdk.coders.IterableCoder.of( + PredictionResultCoder.of(input.getCoder(), outCoder))); + } else { + result.setCoder( + (org.apache.beam.sdk.coders.Coder) + org.apache.beam.sdk.coders.IterableCoder.of( + org.apache.beam.sdk.coders.SerializableCoder.of(PredictionResult.class))); + } + + return result; } /** 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 29f059b0595c..5fe8ea1cce19 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 @@ -29,8 +29,6 @@ import java.util.List; import java.util.stream.Collectors; import java.util.stream.StreamSupport; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; @@ -302,13 +300,11 @@ public void testInvokeWithSingleElement() { PCollection inputCollection = pipeline.apply(Create.of(input)); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)); // Verify the output contains expected predictions PAssert.thatSingleton(results) @@ -341,13 +337,11 @@ public void testInvokeWithMultipleElements() { "CreateInputs", Create.of(inputs).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)); // Count total results across all batches PAssert.that(results) @@ -379,13 +373,11 @@ public void testInvokeWithEmptyCollection() { pipeline.apply("CreateEmptyInput", Create.empty(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)); // assertion for empty PCollection PAssert.that(results).empty(); @@ -403,13 +395,11 @@ public void testHandlerReturnsEmptyResults() { "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockEmptyResultHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockEmptyResultHandler.class) + .withParameters(params)); // Verify we still get a result, but it's empty PAssert.thatSingleton(results) @@ -433,13 +423,11 @@ public void testHandlerSetupFailure() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockFailingSetupHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockFailingSetupHandler.class) + .withParameters(params)); // Verify pipeline fails with expected error try { @@ -464,13 +452,11 @@ public void testHandlerRequestFailure() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockFailingRequestHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockFailingRequestHandler.class) + .withParameters(params)); // Verify pipeline fails with expected error try { @@ -493,13 +479,11 @@ public void testHandlerWithoutDefaultConstructor() { pipeline.apply( "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockNoDefaultConstructorHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockNoDefaultConstructorHandler.class) + .withParameters(params)); // Verify pipeline fails when handler cannot be instantiated try { @@ -535,13 +519,11 @@ public void testPredictionResultMapping() { "CreateInput", Create.of(input).withCoder(SerializableCoder.of(TestInput.class))); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withParameters(params)); PAssert.thatSingleton(results) .satisfies( @@ -591,14 +573,12 @@ public void testBatchingProducesCombinedBatches() { .build(); PCollection>> results = - inputCollection - .apply( - "RemoteInference", - RemoteInference.invoke() - .handler(MockSuccessHandler.class) - .withBatchConfig(batchConfig) - .withParameters(params)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + inputCollection.apply( + "RemoteInference", + RemoteInference.invoke() + .handler(MockSuccessHandler.class) + .withBatchConfig(batchConfig) + .withParameters(params)); PAssert.that(results) .satisfies( @@ -685,20 +665,18 @@ public void testThrottlingBehavior() { .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)) - .setCoder((Coder) IterableCoder.of(SerializableCoder.of(PredictionResult.class))); + 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(