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..14a8fc7785c4 --- /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) { + if (parameters == null) { + throw new NullPointerException("GeminiModelParameters must not be null"); + } + 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 { + 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 one is provided"); + } + 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..24f377afdbf3 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/main/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelParameters.java @@ -0,0 +1,56 @@ +/* + * 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 GeminiRequestFunction getRequestFn(); + + public static Builder builder() { + return new AutoValue_GeminiModelParameters.Builder(); + } + + @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 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/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; 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..42892580fd01 --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerIT.java @@ -0,0 +1,155 @@ +/* + * 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.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())); + + // 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())); + + // 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..c69b2433747e --- /dev/null +++ b/sdks/java/ml/inference/gemini/src/test/java/org/apache/beam/sdk/ml/inference/gemini/GeminiModelHandlerTest.java @@ -0,0 +1,119 @@ +/* + * 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 testNullParameters() { + GeminiModelHandler handler = new GeminiModelHandler<>(); + assertThrows(NullPointerException.class, () -> handler.createClient(null)); + } + + @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/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/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 c571911a484e..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; @@ -61,8 +62,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 +70,7 @@ Invoke invoke() { private RemoteInference() {} @AutoValue - public abstract static class Invoke + public abstract static class Invoke extends PTransform< PCollection, PCollection>>> { @@ -88,10 +88,12 @@ public abstract static class Invoke outputCoder(); + abstract Builder builder(); @AutoValue.Builder - abstract static class Builder { + abstract static class Builder { abstract Builder setHandler( Class> modelHandler); @@ -108,6 +110,8 @@ abstract Builder setHandler( abstract Builder setOverloadRatio(Double overloadRatio); + abstract Builder setOutputCoder(Coder outputCoder); + abstract Invoke build(); } @@ -163,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) { @@ -177,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; } /** @@ -194,7 +221,7 @@ public PCollection>> expand( * */ @SuppressWarnings("nullness") - static class RemoteInferenceFn + static class RemoteInferenceFn extends DoFn, Iterable>> { private final Class> handlerClass; 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..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 @@ -45,7 +45,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 +84,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) { 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")