[pull] master from apache:master

sdks/go.mod

@@ -60,9 +60,9 @@ require (
 	golang.org/x/sync v0.21.0
 	golang.org/x/sys v0.46.0
 	golang.org/x/text v0.38.0
-	google.golang.org/api v0.286.0
+	google.golang.org/api v0.287.0
 	google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68
-	google.golang.org/grpc v1.81.1
+	google.golang.org/grpc v1.82.0
 	google.golang.org/protobuf v1.36.11
 	gopkg.in/yaml.v2 v2.4.0
 	gopkg.in/yaml.v3 v3.0.1
@@ -174,7 +174,7 @@ require (
 	github.com/google/pprof v0.0.0-20260507013755-92041b743c96 // indirect
 	github.com/google/renameio/v2 v2.0.0 // indirect
 	github.com/google/s2a-go v0.1.9 // indirect
-	github.com/googleapis/enterprise-certificate-proxy v0.3.16 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.3.17 // indirect
 	github.com/googleapis/gax-go/v2 v2.22.0 // indirect
 	github.com/gorilla/handlers v1.5.2 // indirect
 	github.com/gorilla/mux v1.8.1 // indirect
@@ -204,5 +204,5 @@ require (
 	golang.org/x/tools v0.45.0 // indirect
 	golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20260622175928-b703f567277d // indirect
 )

sdks/go.sum

@@ -538,8 +538,8 @@ github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/wire v0.5.0/go.mod h1:ngWDr9Qvq3yZA10YrxfyGELY/AFWGVpy9c1LTRi1EoU=
-github.com/googleapis/enterprise-certificate-proxy v0.3.16 h1:F/VPrx0YPBdksZJQdCAp0WUsqnNmZpUZszzfYt0M5Dw=
-github.com/googleapis/enterprise-certificate-proxy v0.3.16/go.mod h1:9Yb0eAkH/Xqhvv3zbeKf/+wMJqCeocWc6KIhDvEAuYE=
+github.com/googleapis/enterprise-certificate-proxy v0.3.17 h1:73NfMHdiqo9JFU9+7a5ExpVa10/R29pXfZIaW559nrg=
+github.com/googleapis/enterprise-certificate-proxy v0.3.17/go.mod h1:rSEsBUemEBZEexP2y6jPp16LUmUbjmSbcPMQizR0o4k=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/googleapis/gax-go/v2 v2.1.0/go.mod h1:Q3nei7sK6ybPYH7twZdmQpAd1MKb7pfu6SK+H1/DsU0=
@@ -1326,8 +1326,8 @@ google.golang.org/api v0.69.0/go.mod h1:boanBiw+h5c3s+tBPgEzLDRHfFLWV0qXxRHz3ws7
 google.golang.org/api v0.70.0/go.mod h1:Bs4ZM2HGifEvXwd50TtW70ovgJffJYw2oRCOFU/SkfA=
 google.golang.org/api v0.71.0/go.mod h1:4PyU6e6JogV1f9eA4voyrTY2batOLdgZ5qZ5HOCc4j8=
 google.golang.org/api v0.74.0/go.mod h1:ZpfMZOVRMywNyvJFeqL9HRWBgAuRfSjJFpe9QtRRyDs=
-google.golang.org/api v0.286.0 h1:TdTXMvzYKnWV1/lPbCdbXRqBrkDqjPto22H2xeZZ8LI=
-google.golang.org/api v0.286.0/go.mod h1:NlOlUIr8MPoIhT9Bb/oUnRuHbJOLwxb6JSYJM8Yz+jQ=
+google.golang.org/api v0.287.0 h1:CQDMqUiqZZ0U/Yge3zyjAhNQ0OSYEH0PaA7l4xtEen4=
+google.golang.org/api v0.287.0/go.mod h1:pPW85yt3Iuc3unkpaMhFtMmOqnTdCwCqEOaUlnuxRlQ=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
 google.golang.org/appengine v1.5.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
@@ -1425,8 +1425,8 @@ google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68 h1:cTHF8xtqtBN5sQ4
 google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68/go.mod h1:RRHjglSYABVCWpQ7USCpdfhcd9t4PkajvVwyynZizTc=
 google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 h1:g0RAkxK/smSu/iRwC/KIX1mwUoVJtk2OjbgaeS4DmUM=
 google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324/go.mod h1:Z4WJ5pJOYWFWcHEQUelD5QaZDknIQkpIL/+fyJOT9+A=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 h1:9HZDLIdYBJXAnaFOr9WHrKVycfpY+75s9HGadC0305A=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260622175928-b703f567277d h1:mpAgMyM9vQHxycBlDq50y1VHpfSfVwzXvrQKtYbXuUY=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260622175928-b703f567277d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
 google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
@@ -1456,8 +1456,8 @@ google.golang.org/grpc v1.40.0/go.mod h1:ogyxbiOoUXAkP+4+xa6PZSE9DZgIHtSpzjDTB9K
 google.golang.org/grpc v1.40.1/go.mod h1:ogyxbiOoUXAkP+4+xa6PZSE9DZgIHtSpzjDTB9KAK34=
 google.golang.org/grpc v1.44.0/go.mod h1:k+4IHHFw41K8+bbowsex27ge2rCb65oeWqe4jJ590SU=
 google.golang.org/grpc v1.45.0/go.mod h1:lN7owxKUQEqMfSyQikvvk5tf/6zMPsrK+ONuO11+0rQ=
-google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ=
-google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I=
+google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU=
+google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA=
 google.golang.org/grpc/cmd/protoc-gen-go-grpc v1.1.0/go.mod h1:6Kw0yEErY5E/yWrBtf03jp27GLLJujG4z/JK95pnjjw=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=

sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java

@@ -20,13 +20,11 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 
 import com.google.auto.value.AutoValue;
-import java.util.Collections;
 import java.util.List;
+import org.apache.beam.sdk.transforms.BatchElements;
 import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SimpleFunction;
 import org.apache.beam.sdk.values.PCollection;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -82,6 +80,8 @@ public abstract static class Invoke<InputT extends BaseInput, OutputT extends Ba
 
     abstract @Nullable BaseModelParameters parameters();
 
+    abstract BatchElements.@Nullable BatchConfig batchConfig();
+
     abstract Builder<InputT, OutputT> builder();
 
     @AutoValue.Builder
@@ -92,6 +92,8 @@ abstract Builder<InputT, OutputT> setHandler(
 
       abstract Builder<InputT, OutputT> setParameters(BaseModelParameters modelParameters);
 
+      abstract Builder<InputT, OutputT> setBatchConfig(BatchElements.BatchConfig batchConfig);
+
       abstract Invoke<InputT, OutputT> build();
     }
 
@@ -106,21 +108,26 @@ public Invoke<InputT, OutputT> withParameters(BaseModelParameters modelParameter
       return builder().setParameters(modelParameters).build();
     }
 
+    /** Configures the batching behavior for the inputs. */
+    public Invoke<InputT, OutputT> withBatchConfig(BatchElements.BatchConfig batchConfig) {
+      return builder().setBatchConfig(batchConfig).build();
+    }
+
     @Override
     public PCollection<Iterable<PredictionResult<InputT, OutputT>>> expand(
         PCollection<InputT> input) {
       checkArgument(handler() != null, "handler() is required");
       checkArgument(parameters() != null, "withParameters() is required");
-      return input
-          .apply(
-              "WrapInputInList",
-              MapElements.via(
-                  new SimpleFunction<InputT, List<InputT>>() {
-                    @Override
-                    public List<InputT> apply(InputT element) {
-                      return Collections.singletonList(element);
-                    }
-                  }))
+
+      BatchElements.BatchConfig config = batchConfig();
+      PCollection<List<InputT>> batchedInput;
+      if (config != null) {
+        batchedInput = input.apply("BatchElements", BatchElements.withConfig(config));
+      } else {
+        batchedInput = input.apply("BatchElements", BatchElements.withDefaults());
+      }
+
+      return batchedInput
           // Pass the list to the inference function
           .apply("RemoteInference", ParDo.of(new RemoteInferenceFn<InputT, OutputT>(this)));
     }

sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RetryHandler.java

@@ -34,24 +34,42 @@ public class RetryHandler implements Serializable {
   private final Duration initialBackoff;
   private final Duration maxBackoff;
   private final Duration maxCumulativeBackoff;
+  private final RetryFilter retryFilter;
+
+  @FunctionalInterface
+  public interface RetryFilter extends Serializable {
+    boolean shouldRetry(Exception e);
+  }
 
   private RetryHandler(
-      int maxRetries, Duration initialBackoff, Duration maxBackoff, Duration maxCumulativeBackoff) {
+      int maxRetries,
+      Duration initialBackoff,
+      Duration maxBackoff,
+      Duration maxCumulativeBackoff,
+      RetryFilter retryFilter) {
     this.maxRetries = maxRetries;
     this.initialBackoff = initialBackoff;
     this.maxBackoff = maxBackoff;
     this.maxCumulativeBackoff = maxCumulativeBackoff;
+    this.retryFilter =
+        java.util.Objects.requireNonNull(retryFilter, "retryFilter must not be null");
   }
 
   public static RetryHandler withDefaults() {
     return new RetryHandler(
         3, // maxRetries
         Duration.standardSeconds(1), // initialBackoff
         Duration.standardSeconds(10), // maxBackoff per retry
-        Duration.standardMinutes(1) // maxCumulativeBackoff
+        Duration.standardMinutes(1), // maxCumulativeBackoff
+        e -> true // retryFilter default to retrying all exceptions
         );
   }
 
+  public RetryHandler withRetryFilter(RetryFilter retryFilter) {
+    return new RetryHandler(
+        maxRetries, initialBackoff, maxBackoff, maxCumulativeBackoff, retryFilter);
+  }
+
   public <T> T execute(RetryableRequest<T> request) throws Exception {
     BackOff backoff =
         FluentBackoff.DEFAULT
@@ -72,6 +90,11 @@ public <T> T execute(RetryableRequest<T> request) throws Exception {
       } catch (Exception e) {
         lastException = e;
 
+        if (retryFilter != null && !retryFilter.shouldRetry(e)) {
+          LOG.warn("Exception not eligible for retry. Failing immediately.", e);
+          throw e;
+        }
+
         long backoffMillis = backoff.nextBackOffMillis();
 
         if (backoffMillis == BackOff.STOP) {

sdks/java/ml/inference/remote/src/test/java/org/apache/beam/sdk/ml/inference/remote/RemoteInferenceTest.java

@@ -518,36 +518,54 @@ public void testPredictionResultMapping() {
     pipeline.run().waitUntilFinish();
   }
 
-  // Temporary behaviour until we introduce java BatchElements transform
-  // to batch elements in RemoteInference
-  @Test
-  public void testMultipleInputsProduceSeparateBatches() {
-    List<TestInput> inputs = Arrays.asList(new TestInput("input1"), new TestInput("input2"));
+  private static class GenerateInputsFn
+      extends org.apache.beam.sdk.transforms.DoFn<Integer, TestInput> {
+    @ProcessElement
+    public void processElement(ProcessContext c) {
+      c.output(new TestInput("input1"));
+      c.output(new TestInput("input2"));
+    }
+  }
 
+  @Test
+  public void testBatchingProducesCombinedBatches() {
     TestParameters params = TestParameters.builder().setConfig("test-config").build();
 
+    // Use a single element to trigger generation of inputs within the same bundle,
+    // ensuring DirectRunner doesn't split them before BatchElements processes them.
     PCollection<TestInput> inputCollection =
-        pipeline.apply(
-            "CreateInputs", Create.of(inputs).withCoder(SerializableCoder.of(TestInput.class)));
+        pipeline
+            .apply("CreateTrigger", Create.of(1))
+            .apply(
+                "GenerateInputs", org.apache.beam.sdk.transforms.ParDo.of(new GenerateInputsFn()))
+            .setCoder(SerializableCoder.of(TestInput.class));
+
+    // Configure BatchElements to force a batch of exactly 2
+    org.apache.beam.sdk.transforms.BatchElements.BatchConfig batchConfig =
+        org.apache.beam.sdk.transforms.BatchElements.BatchConfig.builder()
+            .withMinBatchSize(2)
+            .withMaxBatchSize(2)
+            .build();
 
     PCollection<Iterable<PredictionResult<TestInput, TestOutput>>> results =
         inputCollection.apply(
             "RemoteInference",
             RemoteInference.<TestInput, TestOutput>invoke()
                 .handler(MockSuccessHandler.class)
+                .withBatchConfig(batchConfig)
                 .withParameters(params));
 
     PAssert.that(results)
         .satisfies(
             batches -> {
               int batchCount = 0;
+              int totalElements = 0;
               for (Iterable<PredictionResult<TestInput, TestOutput>> batch : batches) {
                 batchCount++;
-                int elementCount = (int) StreamSupport.stream(batch.spliterator(), false).count();
-                // Each batch should contain exactly 1 element
-                assertEquals("Each batch should contain 1 element", 1, elementCount);
+                totalElements += (int) StreamSupport.stream(batch.spliterator(), false).count();
               }
-              assertEquals("Expected 2 batches", 2, batchCount);
+              assertEquals("Expected 1 batch", 1, batchCount);
+              assertEquals("Total output elements should be 2", 2, totalElements);
               return null;
             });