Merge branch 'master' into handle-grpc-exception

Author: shunping

.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions.yml

@@ -60,7 +60,7 @@ jobs:
       matrix:
         job_name: [beam_PostCommit_Java_ValidatesRunner_Dataflow_JavaVersions]
         job_phrase: [Run Dataflow ValidatesRunner Java]
-        java_version: ['8', '25']
+        java_version: ['25']
     if: |
       github.event_name == 'workflow_dispatch' ||
       github.event_name == 'pull_request_target' ||

.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Batch_DistilBert_Base_Uncased.txt

@@ -18,6 +18,7 @@
 --machine_type=n1-standard-2
 --num_workers=20
 --max_num_workers=250
+--timeout_ms=600000
 --disk_size_gb=50
 --autoscaling_algorithm=THROUGHPUT_BASED
 --staging_location=gs://temp-storage-for-perf-tests/loadtests
@@ -31,5 +32,7 @@
 --device=CPU
 --input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
 --runner=DataflowRunner
+--sdk_location=container
+--sdk_container_image=us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest
 --model_path=distilbert-base-uncased-finetuned-sst-2-english
 --model_state_dict_path=gs://apache-beam-ml/models/huggingface.sentiment.distilbert-base-uncased.pth

.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Streaming_DistilBert_Base_Uncased.txt

@@ -18,6 +18,7 @@
 --machine_type=n1-standard-2
 --num_workers=20
 --max_num_workers=250
+--timeout_ms=600000
 --disk_size_gb=50
 --autoscaling_algorithm=THROUGHPUT_BASED
 --staging_location=gs://temp-storage-for-perf-tests/loadtests
@@ -31,6 +32,8 @@
 --device=CPU
 --input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
 --runner=DataflowRunner
+--sdk_location=container
+--sdk_container_image=us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt:latest
 --dataflow_service_options=worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver
 --model_path=distilbert-base-uncased-finetuned-sst-2-english
 --model_state_dict_path=gs://apache-beam-ml/models/huggingface.sentiment.distilbert-base-uncased.pth

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -743,7 +743,7 @@ class BeamModulePlugin implements Plugin<Project> {
         google_api_common                           : "com.google.api:api-common", // google_cloud_platform_libraries_bom sets version
         google_api_services_bigquery                : "com.google.apis:google-api-services-bigquery:v2-rev20251012-2.0.0",  // [bomupgrader] sets version
         google_api_services_cloudresourcemanager    : "com.google.apis:google-api-services-cloudresourcemanager:v1-rev20250606-2.0.0",  // [bomupgrader] sets version
-        google_api_services_dataflow                : "com.google.apis:google-api-services-dataflow:v1b3-rev20260405-$google_clients_version",
+        google_api_services_dataflow                : "com.google.apis:google-api-services-dataflow:v1b3-rev20260503-$google_clients_version",
         google_api_services_healthcare              : "com.google.apis:google-api-services-healthcare:v1-rev20240130-$google_clients_version",
         google_api_services_pubsub                  : "com.google.apis:google-api-services-pubsub:v1-rev20220904-$google_clients_version",
         google_api_services_storage                 : "com.google.apis:google-api-services-storage:v1-rev20260204-2.0.0",  // [bomupgrader] sets version

runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java

@@ -37,6 +37,7 @@
 import org.apache.beam.runners.core.triggers.TriggerStateMachineRunner;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.options.ExperimentalOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.state.CombiningState;
 import org.apache.beam.sdk.state.TimeDomain;
@@ -94,6 +95,11 @@
 }) // TODO(https://github.com/apache/beam/issues/20497)
 public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
 
+  // Experiments guarding optimizations in development. No backward compatibility guarantees.
+  public static final String UNSTABLE_NOT_UPDATE_COMPATIBLE_NEW_WINDOW_OPTIMIZATION =
+      "unstable_not_update_compatible_new_window_optimization";
+  public static final String UNSTABLE_DISABLE_WATERMARK_KNOWN_EMPTY_OPTIMIZATION =
+      "unstable_disable_watermark_known_empty_optimization";
   /**
    * The {@link ReduceFnRunner} depends on most aspects of the {@link WindowingStrategy}.
    *
@@ -218,6 +224,9 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
    */
   private final NonEmptyPanes<K, W> nonEmptyPanes;
 
+  private final boolean useNewWindowOptimization;
+  private final boolean disableWatermarkKnownEmptyOptimization;
+
   public ReduceFnRunner(
       K key,
       WindowingStrategy<?, W> windowingStrategy,
@@ -244,6 +253,15 @@ public ReduceFnRunner(
 
     this.nonEmptyPanes = NonEmptyPanes.create(this.windowingStrategy, this.reduceFn);
 
+    this.useNewWindowOptimization =
+        windowingStrategy.getWindowFn().isNonMerging()
+            && ExperimentalOptions.hasExperiment(
+                options, UNSTABLE_NOT_UPDATE_COMPATIBLE_NEW_WINDOW_OPTIMIZATION);
+
+    this.disableWatermarkKnownEmptyOptimization =
+        ExperimentalOptions.hasExperiment(
+            options, UNSTABLE_DISABLE_WATERMARK_KNOWN_EMPTY_OPTIMIZATION);
+
     // Note this may incur I/O to load persisted window set data.
     this.activeWindows = createActiveWindowSet();
 
@@ -263,7 +281,8 @@ public ReduceFnRunner(
         new TriggerStateMachineRunner<>(
             triggerStateMachine,
             new TriggerStateMachineContextFactory<>(
-                windowingStrategy.getWindowFn(), stateInternals, activeWindows));
+                windowingStrategy.getWindowFn(), stateInternals, activeWindows),
+            this.useNewWindowOptimization);
   }
 
   private ActiveWindowSet<W> createActiveWindowSet() {
@@ -282,6 +301,16 @@ boolean hasNoActiveWindows() {
     return activeWindows.getActiveAndNewWindows().isEmpty();
   }
 
+  @VisibleForTesting
+  TriggerStateMachineRunner<W> getTriggerRunner() {
+    return triggerRunner;
+  }
+
+  @VisibleForTesting
+  ReduceFnContextFactory<K, InputT, OutputT, W> getContextFactory() {
+    return contextFactory;
+  }
+
   private Set<W> windowsThatAreOpen(Collection<W> windows) {
     Set<W> result = new HashSet<>();
     for (W window : windows) {
@@ -626,6 +655,20 @@ private void processElement(Map<W, W> windowToMergeResult, WindowedValue<InputT>
               StateStyle.RENAMED,
               value.causedByDrain());
 
+      // TODO: Make sure the NewWindowOptimization does not create unbounded trigger state
+      // in GlobalWindow
+      if (useNewWindowOptimization && triggerRunner.isNew(directContext.state())) {
+        // Blindly clear state to ensure Windmill doesn't do unnecessary reads.
+        // TODO: Instead of the clears here, we could mark these states as empty locally
+        //  in the state cache and/or explicitly tell that the entries are non-existent via api
+        reduceFn.clearState(renamedContext);
+        paneInfoTracker.clear(directContext.state());
+        if (!disableWatermarkKnownEmptyOptimization) {
+          watermarkHold.setKnownEmpty(renamedContext);
+        }
+        nonEmptyPanes.clearPane(renamedContext.state());
+      }
+
       nonEmptyPanes.recordContent(renamedContext.state());
       scheduleGarbageCollectionTimer(directContext);
 
@@ -761,7 +804,7 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
 
       // Perform prefetching of state to determine if the trigger should fire.
       if (windowActivation.isGarbageCollection) {
-        triggerRunner.prefetchIsClosed(directContext.state());
+        triggerRunner.prefetchFinishedSet(directContext.state());
       } else {
         triggerRunner.prefetchShouldFire(directContext.window(), directContext.state());
       }
@@ -941,7 +984,7 @@ private void prefetchEmit(
       ReduceFn<K, InputT, OutputT, W>.Context directContext,
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext) {
     triggerRunner.prefetchShouldFire(directContext.window(), directContext.state());
-    triggerRunner.prefetchIsClosed(directContext.state());
+    triggerRunner.prefetchFinishedSet(directContext.state());
     prefetchOnTrigger(directContext, renamedContext);
   }
 

runners/core-java/src/main/java/org/apache/beam/runners/core/WatermarkHold.java

@@ -466,6 +466,23 @@ public void clearHolds(ReduceFn<?, ?, ?, W>.Context context) {
     context.state().access(EXTRA_HOLD_TAG).clear();
   }
 
+  /**
+   * <b><i>For internal use only; no backwards-compatibility guarantees.</i></b>
+   *
+   * <p>Permit marking the watermark holds as empty locally, without necessarily clearing them in
+   * the backend.
+   */
+  public void setKnownEmpty(ReduceFn<?, ?, ?, W>.Context context) {
+    WindowTracing.debug(
+        "WatermarkHold.setKnownEmpty: For key:{}; window:{}; inputWatermark:{}; outputWatermark:{}",
+        context.key(),
+        context.window(),
+        timerInternals.currentInputWatermarkTime(),
+        timerInternals.currentOutputWatermarkTime());
+    context.state().access(elementHoldTag).setKnownEmpty();
+    context.state().access(EXTRA_HOLD_TAG).setKnownEmpty();
+  }
+
   /** Return the current data hold, or null if none. Does not clear. For debugging only. */
   public @Nullable Instant getDataCurrent(ReduceFn<?, ?, ?, W>.Context context) {
     return context.state().access(elementHoldTag).read();

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/FinishedTriggersBitSet.java

@@ -18,6 +18,7 @@
 package org.apache.beam.runners.core.triggers;
 
 import java.util.BitSet;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** A {@link FinishedTriggers} implementation based on an underlying {@link BitSet}. */
 public class FinishedTriggersBitSet implements FinishedTriggers {
@@ -60,4 +61,17 @@ public void clearRecursively(ExecutableTriggerStateMachine trigger) {
   public FinishedTriggersBitSet copy() {
     return new FinishedTriggersBitSet((BitSet) bitSet.clone());
   }
+
+  @Override
+  public boolean equals(@Nullable Object obj) {
+    if (!(obj instanceof FinishedTriggersBitSet)) {
+      return false;
+    }
+    return bitSet.equals(((FinishedTriggersBitSet) obj).bitSet);
+  }
+
+  @Override
+  public int hashCode() {
+    return bitSet.hashCode();
+  }
 }

runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachineRunner.java

@@ -63,13 +63,16 @@ public class TriggerStateMachineRunner<W extends BoundedWindow> {
 
   private final ExecutableTriggerStateMachine rootTrigger;
   private final TriggerStateMachineContextFactory<W> contextFactory;
+  private final boolean useNewWindowOptimization;
 
   public TriggerStateMachineRunner(
       ExecutableTriggerStateMachine rootTrigger,
-      TriggerStateMachineContextFactory<W> contextFactory) {
+      TriggerStateMachineContextFactory<W> contextFactory,
+      boolean useNewWindowOptimization) {
     checkState(rootTrigger.getTriggerIndex() == 0);
     this.rootTrigger = rootTrigger;
     this.contextFactory = contextFactory;
+    this.useNewWindowOptimization = useNewWindowOptimization;
   }
 
   private FinishedTriggersBitSet readFinishedBits(ValueState<BitSet> state) {
@@ -81,9 +84,11 @@ private FinishedTriggersBitSet readFinishedBits(ValueState<BitSet> state) {
     }
 
     @Nullable BitSet bitSet = state.read();
-    return bitSet == null
-        ? FinishedTriggersBitSet.emptyWithCapacity(rootTrigger.getFirstIndexAfterSubtree())
-        : FinishedTriggersBitSet.fromBitSet(bitSet);
+    if (bitSet == null) {
+      return FinishedTriggersBitSet.emptyWithCapacity(rootTrigger.getFirstIndexAfterSubtree());
+    }
+
+    return FinishedTriggersBitSet.fromBitSet(bitSet);
   }
 
   private void clearFinishedBits(ValueState<BitSet> state) {
@@ -99,19 +104,29 @@ public boolean isClosed(StateAccessor<?> state) {
     return readFinishedBits(state.access(FINISHED_BITS_TAG)).isFinished(rootTrigger);
   }
 
-  public void prefetchIsClosed(StateAccessor<?> state) {
+  /** Return true if the window is new (no trigger state has ever been persisted). */
+  public boolean isNew(StateAccessor<?> state) {
+    return isFinishedSetNeeded() && state.access(FINISHED_BITS_TAG).read() == null;
+  }
+
+  @VisibleForTesting
+  public BitSet getFinishedBits(StateAccessor<?> state) {
+    return readFinishedBits(state.access(FINISHED_BITS_TAG)).getBitSet();
+  }
+
+  public void prefetchFinishedSet(StateAccessor<?> state) {
     if (isFinishedSetNeeded()) {
       state.access(FINISHED_BITS_TAG).readLater();
     }
   }
 
   public void prefetchForValue(W window, StateAccessor<?> state) {
-    prefetchIsClosed(state);
+    prefetchFinishedSet(state);
     rootTrigger.invokePrefetchOnElement(contextFactory.createPrefetchContext(window, rootTrigger));
   }
 
   public void prefetchShouldFire(W window, StateAccessor<?> state) {
-    prefetchIsClosed(state);
+    prefetchFinishedSet(state);
     rootTrigger.invokePrefetchShouldFire(contextFactory.createPrefetchContext(window, rootTrigger));
   }
 
@@ -180,13 +195,23 @@ public void onFire(W window, Timers timers, StateAccessor<?> state) throws Excep
     persistFinishedSet(state, finishedSet);
   }
 
-  private void persistFinishedSet(
-      StateAccessor<?> state, FinishedTriggersBitSet modifiedFinishedSet) {
+  @VisibleForTesting
+  void persistFinishedSet(StateAccessor<?> state, FinishedTriggersBitSet modifiedFinishedSet) {
     if (!isFinishedSetNeeded()) {
       return;
     }
 
     ValueState<BitSet> finishedSetState = state.access(FINISHED_BITS_TAG);
+
+    if (useNewWindowOptimization) {
+      if (finishedSetState.read() == null
+          || !readFinishedBits(finishedSetState).equals(modifiedFinishedSet)) {
+        // Write a value even if the bitset was empty
+        finishedSetState.write(modifiedFinishedSet.getBitSet());
+      }
+      return;
+    }
+
     if (!readFinishedBits(finishedSetState).equals(modifiedFinishedSet)) {
       if (modifiedFinishedSet.getBitSet().isEmpty()) {
         finishedSetState.clear();

runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java

@@ -40,15 +40,19 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.BitSet;
+import java.util.Collections;
 import java.util.List;
 import java.util.Random;
 import java.util.concurrent.ThreadLocalRandom;
+import org.apache.beam.runners.core.ReduceFnContextFactory.StateStyle;
 import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
 import org.apache.beam.runners.core.triggers.DefaultTriggerStateMachine;
 import org.apache.beam.runners.core.triggers.TriggerStateMachine;
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.metrics.MetricName;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.apache.beam.sdk.options.ExperimentalOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.state.TimeDomain;
@@ -2343,4 +2347,58 @@ public interface TestOptions extends PipelineOptions {
 
     void setValue(int value);
   }
+
+  @Test
+  public void testNewWindowOptimization() throws Exception {
+    WindowingStrategy<?, IntervalWindow> strategy =
+        WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))
+            .withTrigger(AfterPane.elementCountAtLeast(2))
+            .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options
+        .as(ExperimentalOptions.class)
+        .setExperiments(
+            Collections.singletonList("unstable_not_update_compatible_new_window_optimization"));
+    ReduceFnTester<Integer, Iterable<Integer>, IntervalWindow> tester =
+        ReduceFnTester.nonCombining(strategy, options);
+
+    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
+
+    assertTrue(
+        "Window should be new",
+        tester
+            .createRunner()
+            .getTriggerRunner()
+            .isNew(
+                tester.createRunner().getContextFactory().base(window, StateStyle.DIRECT).state()));
+
+    // 1. First element for a new window.
+    tester.injectElements(TimestampedValue.of(1, new Instant(1)));
+
+    BitSet bitSet =
+        tester
+            .createRunner()
+            .getTriggerRunner()
+            .getFinishedBits(
+                tester.createRunner().getContextFactory().base(window, StateStyle.DIRECT).state());
+    assertTrue("Bitset should be empty", bitSet.isEmpty());
+    assertFalse("Trigger should not be finished", bitSet.get(0));
+
+    assertFalse(
+        "Window should no longer be new",
+        tester
+            .createRunner()
+            .getTriggerRunner()
+            .isNew(
+                tester.createRunner().getContextFactory().base(window, StateStyle.DIRECT).state()));
+
+    // 2. Second element for the same window.
+    tester.injectElements(TimestampedValue.of(2, new Instant(2)));
+
+    // Extract output.
+    List<WindowedValue<Iterable<Integer>>> output = tester.extractOutput();
+    // 2 elements fired at end of window.
+    assertThat(output, contains(isSingleWindowedValue(containsInAnyOrder(1, 2), 9, 0, 10)));
+  }
 }