diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json index 455144f02a35..d6a91b7e2e86 100644 --- a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark3_Streaming.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 6 + "modification": 7 } diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json index 455144f02a35..d6a91b7e2e86 100644 --- a/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json +++ b/.github/trigger_files/beam_PostCommit_Java_PVR_Spark_Batch.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 6 + "modification": 7 } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json index 1efc8e9e4405..3f63c0c9975f 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 2 } diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json index f4ec72dc416b..6384446f50e4 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json @@ -3,5 +3,6 @@ "https://github.com/apache/beam/issues/35429": "testing", "trigger-2026-04-04": "portable_runner expand_sdf opt-in", "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test", - "modification": 1 + "modification": 1, + "https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization" } diff --git a/CHANGES.md b/CHANGES.md index a6fd20ad34dd..4c9b8700a8c1 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,8 @@ ## New Features / Improvements * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). +* Splittable DoFn self-checkpointing is now supported on the portable Spark runner, including unbounded reads in streaming mode ([#19468](https://github.com/apache/beam/issues/19468)). +* Bundle finalization is now supported on the portable Spark runner ([#19517](https://github.com/apache/beam/issues/19517)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/runners/spark/job-server/spark_job_server.gradle b/runners/spark/job-server/spark_job_server.gradle index 5240bb310d05..317947998338 100644 --- a/runners/spark/job-server/spark_job_server.gradle +++ b/runners/spark/job-server/spark_job_server.gradle @@ -199,11 +199,9 @@ def portableValidatesRunnerTask(String name, boolean streaming, boolean docker, excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo' excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration' excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream' - // TODO (https://github.com/apache/beam/issues/19468) SplittableDoFnTests - excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo' + // Unbounded SDF requires the streaming residual relay (#19468). excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo' excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering' - excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer' } testFilter = { // TODO (https://github.com/apache/beam/issues/20189) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java index 91a94896b89b..e20396e3af29 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java @@ -39,6 +39,7 @@ import org.apache.beam.runners.spark.translation.SparkStreamingPortablePipelineTranslator; import org.apache.beam.runners.spark.translation.SparkStreamingTranslationContext; import org.apache.beam.runners.spark.translation.SparkTranslationContext; +import org.apache.beam.runners.spark.translation.streaming.SdfResidualRelay; import org.apache.beam.runners.spark.util.GlobalWatermarkHolder; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.metrics.MetricsEnvironment; @@ -166,8 +167,10 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) jssc.awaitTerminationOrTimeout(timeout); } catch (InterruptedException e) { LOG.warn("Streaming context interrupted, shutting down.", e); + } finally { + jssc.stop(); + SdfResidualRelay.unregisterJob(jobInfo.jobId()); } - jssc.stop(); LOG.info("Job {} finished.", jobInfo.jobId()); }); result = new SparkPipelineResult.PortableStreamingMode(submissionFuture, jssc); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java index ba3aa0e4d24a..b7edfdaf9d87 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java @@ -262,7 +262,9 @@ private static void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); + windowCoder, + getWindowedValueCoder(inputPCollectionId, components), + false); staged = groupedByKey.flatMap(function.forPair()); } else { JavaRDD> inputRdd2 = ((BoundedDataset) inputDataset).getRDD(); @@ -275,7 +277,9 @@ private static void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); + windowCoder, + getWindowedValueCoder(inputPCollectionId, components), + false); staged = inputRdd2.mapPartitions(function2); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java index 757740e2df5a..2d7dd3debc68 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunction.java @@ -19,6 +19,7 @@ import java.io.IOException; import java.io.Serializable; +import java.util.ArrayList; import java.util.Collections; import java.util.EnumMap; import java.util.Iterator; @@ -27,6 +28,7 @@ import java.util.Map; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.stream.Collectors; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication; import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse; import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; @@ -36,6 +38,9 @@ import org.apache.beam.runners.core.TimerInternals; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.core.metrics.MetricsContainerImpl; +import org.apache.beam.runners.fnexecution.control.BundleCheckpointHandler; +import org.apache.beam.runners.fnexecution.control.BundleFinalizationHandlers; +import org.apache.beam.runners.fnexecution.control.BundleFinalizationHandlers.InMemoryFinalizer; import org.apache.beam.runners.fnexecution.control.BundleProgressHandler; import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; import org.apache.beam.runners.fnexecution.control.JobBundleFactory; @@ -59,6 +64,7 @@ import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.transforms.join.RawUnionValue; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.construction.Timer; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; import org.apache.beam.sdk.values.WindowedValue; @@ -85,6 +91,10 @@ class SparkExecutableStageFunction implements FlatMapFunction>, RawUnionValue> { + // Union tag carrying serialized SDF residuals to the streaming residual relay. Regular output + // tags start at 0. + static final int SDF_RESIDUAL_TAG = -1; + // Pipeline options for initializing the FileSystems private final SerializablePipelineOptions pipelineOptions; private final RunnerApi.ExecutableStagePayload stagePayload; @@ -95,6 +105,10 @@ class SparkExecutableStageFunction sideInputs; private final MetricsContainerStepMapAccumulator metricsAccumulator; private final Coder windowCoder; + // Coder for the stage input, used to re-feed SDF self-checkpoint residuals. + private final Coder> inputCoder; + // Streaming emits residuals to the cross-batch relay; batch drains them in place. + private final boolean emitSdfResiduals; private final JobInfo jobInfo; private transient InMemoryBagUserStateFactory bagUserStateHandlerFactory; @@ -108,7 +122,9 @@ class SparkExecutableStageFunction SparkExecutableStageContextFactory contextFactory, Map>, WindowedValueCoder>> sideInputs, MetricsContainerStepMapAccumulator metricsAccumulator, - Coder windowCoder) { + Coder windowCoder, + Coder> inputCoder, + boolean emitSdfResiduals) { this.pipelineOptions = pipelineOptions; this.stagePayload = stagePayload; this.jobInfo = jobInfo; @@ -117,6 +133,8 @@ class SparkExecutableStageFunction this.sideInputs = sideInputs; this.metricsAccumulator = metricsAccumulator; this.windowCoder = windowCoder; + this.inputCoder = inputCoder; + this.emitSdfResiduals = emitSdfResiduals; } /** Call the executable stage function on the values of a PairRDD, ignoring the key. */ @@ -146,7 +164,30 @@ public Iterator call(Iterator> inputs) thro executableStage, stageBundleFactory.getProcessBundleDescriptor()); if (executableStage.getTimers().size() == 0) { ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap); - processElements(stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs); + ResidualCollector residualCollector = new ResidualCollector(); + InMemoryFinalizer finalizer = + BundleFinalizationHandlers.inMemoryFinalizer( + stageBundleFactory.getInstructionRequestHandler()); + processElements( + stateRequestHandler, + receiverFactory, + null, + stageBundleFactory, + inputs, + residualCollector, + finalizer); + if (emitSdfResiduals) { + for (DelayedBundleApplication residual : residualCollector.drain()) { + collector.add(new RawUnionValue(SDF_RESIDUAL_TAG, residual.toByteArray())); + } + } else { + processResiduals( + stateRequestHandler, + receiverFactory, + stageBundleFactory, + residualCollector, + finalizer); + } return collector.iterator(); } // Used with Batch, we know that all the data is available for this key. We can't use the @@ -172,8 +213,18 @@ public Iterator call(Iterator> inputs) thro windowCoder); // Process inputs. + ResidualCollector residualCollector = new ResidualCollector(); + InMemoryFinalizer finalizer = + BundleFinalizationHandlers.inMemoryFinalizer( + stageBundleFactory.getInstructionRequestHandler()); processElements( - stateRequestHandler, receiverFactory, timerReceiverFactory, stageBundleFactory, inputs); + stateRequestHandler, + receiverFactory, + timerReceiverFactory, + stageBundleFactory, + inputs, + residualCollector, + finalizer); // Finish any pending windows by advancing the input watermark to infinity. timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE); @@ -189,12 +240,17 @@ public Iterator call(Iterator> inputs) thro receiverFactory, timerReceiverFactory, stateRequestHandler, - getBundleProgressHandler())) { + getBundleProgressHandler(), + finalizer, + residualCollector)) { PipelineTranslatorUtils.fireEligibleTimers( timerInternals, bundle.getTimerReceivers(), currentTimerKey); } + finalizer.finalizeAllOutstandingBundles(); } + processResiduals( + stateRequestHandler, receiverFactory, stageBundleFactory, residualCollector, finalizer); return collector.iterator(); } } @@ -207,14 +263,18 @@ private void processElements( ReceiverFactory receiverFactory, TimerReceiverFactory timerReceiverFactory, StageBundleFactory stageBundleFactory, - Iterator> inputs) + Iterator> inputs, + BundleCheckpointHandler checkpointHandler, + InMemoryFinalizer finalizer) throws Exception { try (RemoteBundle bundle = stageBundleFactory.getBundle( receiverFactory, timerReceiverFactory, stateRequestHandler, - getBundleProgressHandler())) { + getBundleProgressHandler(), + finalizer, + checkpointHandler)) { FnDataReceiver> mainReceiver = Iterables.getOnlyElement(bundle.getInputReceivers().values()); while (inputs.hasNext()) { @@ -222,6 +282,97 @@ private void processElements( mainReceiver.accept(input); } } + finalizer.finalizeAllOutstandingBundles(); + } + + // Re-feeds SDF self-checkpoint residuals in fresh bundles until the SDK returns none. Each + // residual resumes at its own requested time. Bounded restrictions always run out; unbounded ones + // never would, so they are rejected rather than drained forever. + private void processResiduals( + StateRequestHandler stateRequestHandler, + ReceiverFactory receiverFactory, + StageBundleFactory stageBundleFactory, + ResidualCollector residualCollector, + InMemoryFinalizer finalizer) + throws Exception { + List scheduled = schedule(residualCollector.drain()); + while (!scheduled.isEmpty()) { + List due = takeDue(scheduled); + try (RemoteBundle bundle = + stageBundleFactory.getBundle( + receiverFactory, + null, + stateRequestHandler, + getBundleProgressHandler(), + finalizer, + residualCollector)) { + FnDataReceiver> mainReceiver = + Iterables.getOnlyElement(bundle.getInputReceivers().values()); + for (ScheduledResidual residual : due) { + mainReceiver.accept(CoderUtils.decodeFromByteArray(inputCoder, residual.elementBytes)); + } + } + finalizer.finalizeAllOutstandingBundles(); + scheduled.addAll(schedule(residualCollector.drain())); + } + } + + private static List schedule(List residuals) { + long now = System.currentTimeMillis(); + List scheduled = new ArrayList<>(); + for (DelayedBundleApplication residual : residuals) { + if (residual.getApplication().getElement().isEmpty()) { + continue; + } + if (residual.getApplication().getIsBounded() == RunnerApi.IsBounded.Enum.UNBOUNDED) { + throw new UnsupportedOperationException( + "Unbounded splittable DoFn is not supported in batch mode on the Spark runner. See " + + "https://github.com/apache/beam/issues/19468."); + } + long delayMillis = + residual.hasRequestedTimeDelay() + ? residual.getRequestedTimeDelay().getSeconds() * 1000 + + residual.getRequestedTimeDelay().getNanos() / 1_000_000 + : 0; + scheduled.add( + new ScheduledResidual( + now + delayMillis, residual.getApplication().getElement().toByteArray())); + } + return scheduled; + } + + // Waits for the earliest resume time, then removes and returns everything due by then. + private static List takeDue(List scheduled) + throws InterruptedException { + long earliest = Long.MAX_VALUE; + for (ScheduledResidual residual : scheduled) { + earliest = Math.min(earliest, residual.dueMillis); + } + long waitMillis = earliest - System.currentTimeMillis(); + if (waitMillis > 0) { + Thread.sleep(waitMillis); + } + long now = System.currentTimeMillis(); + List due = new ArrayList<>(); + Iterator iterator = scheduled.iterator(); + while (iterator.hasNext()) { + ScheduledResidual residual = iterator.next(); + if (residual.dueMillis <= now) { + due.add(residual); + iterator.remove(); + } + } + return due; + } + + private static class ScheduledResidual { + private final long dueMillis; + private final byte[] elementBytes; + + ScheduledResidual(long dueMillis, byte[] elementBytes) { + this.dueMillis = dueMillis; + this.elementBytes = elementBytes; + } } private BundleProgressHandler getBundleProgressHandler() { @@ -295,6 +446,27 @@ interface JobBundleFactoryCreator extends Serializable { JobBundleFactory create(); } + /** Collects SDF self-checkpoint residuals returned by the SDK harness. */ + private static class ResidualCollector implements BundleCheckpointHandler { + + private final ConcurrentLinkedQueue residuals = + new ConcurrentLinkedQueue<>(); + + @Override + public void onCheckpoint(ProcessBundleResponse response) { + residuals.addAll(response.getResidualRootsList()); + } + + private List drain() { + List pending = new ArrayList<>(); + DelayedBundleApplication residual; + while ((residual = residuals.poll()) != null) { + pending.add(residual); + } + return pending; + } + } + /** * Receiver factory that wraps outgoing elements with the corresponding union tag for a * multiplexed PCollection. diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java index 9975c81b56a4..f23953b842e1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java @@ -37,6 +37,7 @@ import org.apache.beam.runners.spark.coders.CoderHelpers; import org.apache.beam.runners.spark.metrics.MetricsAccumulator; import org.apache.beam.runners.spark.stateful.SparkGroupAlsoByWindowViaWindowSet; +import org.apache.beam.runners.spark.translation.streaming.SdfResidualRelay; import org.apache.beam.runners.spark.translation.streaming.UnboundedDataset; import org.apache.beam.runners.spark.util.GlobalWatermarkHolder; import org.apache.beam.sdk.coders.ByteArrayCoder; @@ -47,6 +48,7 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; import org.apache.beam.sdk.util.construction.graph.PipelineNode; @@ -242,6 +244,37 @@ private static void translateExecutableStage( String, Tuple2>, WindowedValues.WindowedValueCoder>> broadcastVariables = ImmutableMap.copyOf(new HashMap<>()); + WindowedValues.WindowedValueCoder inputCoder = + getWindowedValueCoder(inputPCollectionId, components); + + boolean hasSdf = hasSdfProcess(stagePayload); + JavaDStream> stageInput = inputDStream; + List outputStreamSources = streamSources; + String relayId = null; + if (hasSdf) { + relayId = SdfResidualRelay.relayId(context.jobInfo.jobId(), transformNode.getId()); + SdfResidualRelay relay = + SdfResidualRelay.register( + relayId, + streamSources, + context + .getSerializableOptions() + .get() + .as(SparkPipelineOptions.class) + .getBatchIntervalMillis()); + SdfResidualRelay.ResidualInputDStream residualInput = + new SdfResidualRelay.ResidualInputDStream(context.getStreamingContext().ssc(), relayId); + relay.setSourceId(residualInput.id()); + JavaDStream residualBytes = + JavaDStream.fromDStream(residualInput, JavaSparkContext$.MODULE$.fakeClassTag()); + // Elements travel encoded through the relay; decode on the executors. + JavaDStream> residualStream = + residualBytes.map(bytes -> CoderUtils.decodeFromByteArray(inputCoder, bytes)); + stageInput = inputDStream.union(residualStream); + // The relay's watermark drives this stage's output from here on. + outputStreamSources = Collections.singletonList(residualInput.id()); + } + SparkExecutableStageFunction function = new SparkExecutableStageFunction<>( context.getSerializableOptions(), @@ -251,8 +284,25 @@ private static void translateExecutableStage( SparkExecutableStageContextFactory.getInstance(), broadcastVariables, MetricsAccumulator.getInstance(), - windowCoder); - JavaDStream staged = inputDStream.mapPartitions(function); + windowCoder, + inputCoder, + hasSdf); + JavaDStream staged = stageInput.mapPartitions(function); + if (hasSdf) { + // The relay and the output extraction both consume the stage; never execute it twice. + staged.persist(StorageLevel.MEMORY_ONLY()); + final String residualRelayId = relayId; + staged.foreachRDD( + (rdd, time) -> { + List residuals = + rdd.filter( + value -> + value.getUnionTag() == SparkExecutableStageFunction.SDF_RESIDUAL_TAG) + .map(value -> (byte[]) value.getValue()) + .collect(); + SdfResidualRelay.onBatchResiduals(residualRelayId, residuals, time.milliseconds()); + }); + } String intermediateId = getExecutableStageIntermediateId(transformNode); context.pushDataset( @@ -281,7 +331,7 @@ public void setName(String name) { for (String outputId : outputs.values()) { JavaDStream> outStream = staged.flatMap(new SparkExecutableStageExtractionFunction<>(outputMap.get(outputId))); - context.pushDataset(outputId, new UnboundedDataset<>(outStream, streamSources)); + context.pushDataset(outputId, new UnboundedDataset<>(outStream, outputStreamSources)); } // Add sink to ensure stage is executed @@ -290,8 +340,20 @@ public void setName(String name) { staged.flatMap((rawUnionValue) -> Collections.emptyIterator()); context.pushDataset( String.format("EmptyOutputSink_%d", context.nextSinkId()), - new UnboundedDataset<>(outStream, streamSources)); + new UnboundedDataset<>(outStream, outputStreamSources)); + } + } + + private static boolean hasSdfProcess(RunnerApi.ExecutableStagePayload stagePayload) { + for (String transformId : stagePayload.getTransformsList()) { + RunnerApi.PTransform transform = + stagePayload.getComponents().getTransformsOrThrow(transformId); + if (PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals( + transform.getSpec().getUrn())) { + return true; + } } + return false; } @SuppressWarnings("unchecked") diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java new file mode 100644 index 000000000000..d2762d1724b4 --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java @@ -0,0 +1,268 @@ +/* + * 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.runners.spark.translation.streaming; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication; +import org.apache.beam.runners.spark.util.GlobalWatermarkHolder; +import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.api.java.JavaSparkContext$; +import org.apache.spark.rdd.RDD; +import org.apache.spark.streaming.StreamingContext; +import org.apache.spark.streaming.Time; +import org.apache.spark.streaming.dstream.ConstantInputDStream; +import org.joda.time.Instant; +import scala.Option; + +/** + * Relays SDF self-checkpoint residuals across micro-batches on the driver. + * + *

Residuals returned by the SDK harness in one micro-batch are collected via {@link + * #onBatchResiduals}, held until their requested resume time, and re-emitted as encoded stage input + * by {@link ResidualInputDStream} in a later micro-batch. Elements stay in their coder-encoded byte + * form on this path; decoding happens on the executors. The relay also advances the {@link + * GlobalWatermarkHolder} watermark for its stream from the residuals' output watermarks, so + * downstream event-time windows fire correctly. + * + *

State lives in a static driver-side registry (like {@link GlobalWatermarkHolder}) so DStream + * checkpointing never needs to serialize it. Residuals are lost on driver failure; the portable + * streaming path has no driver recovery. + * + *

This relay reports a watermark per micro-batch, while an impulse reports once. A {@link + * org.apache.beam.sdk.transforms.GroupByKey} whose inputs are flattened from both therefore sees + * two sources with different synchronized processing times, which {@code + * SparkTimerInternals#forStreamFromSources} rejects. Grouping an SDF's output on its own is + * unaffected, since this relay replaces the stage's stream sources. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class SdfResidualRelay { + + private static final Map RELAYS = new ConcurrentHashMap<>(); + + // Residuals waiting for their requested resume time; guarded by this. + private final List pending = new ArrayList<>(); + // Residuals taken by a generated batch, keyed by batch time, until that batch reports back. + private final Map> inFlight = new HashMap<>(); + // The stage's original inputs; this stage can never be ahead of them. + private final List upstreamSourceIds; + private final long batchDurationMillis; + private Instant highWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + private boolean seenResidual = false; + // Set during translation, read from the streaming scheduler threads. + private volatile int sourceId = -1; + + private SdfResidualRelay(List upstreamSourceIds, long batchDurationMillis) { + this.upstreamSourceIds = new ArrayList<>(upstreamSourceIds); + this.batchDurationMillis = batchDurationMillis; + } + + /** Builds a registry key that is unique across jobs sharing a job server JVM. */ + public static String relayId(String jobId, String transformId) { + return jobId + "/" + transformId; + } + + public static SdfResidualRelay register( + String relayId, List upstreamSourceIds, long batchDurationMillis) { + SdfResidualRelay relay = new SdfResidualRelay(upstreamSourceIds, batchDurationMillis); + if (RELAYS.putIfAbsent(relayId, relay) != null) { + throw new IllegalStateException("Duplicate SDF residual relay registration: " + relayId); + } + return relay; + } + + /** Drops every relay belonging to a finished job. */ + public static void unregisterJob(String jobId) { + RELAYS.keySet().removeIf(relayId -> relayId.startsWith(jobId + "/")); + } + + /** Feeds the residuals of a completed micro-batch into the relay. */ + public static void onBatchResiduals( + String relayId, List serializedResiduals, long batchTimeMillis) { + SdfResidualRelay relay = RELAYS.get(relayId); + if (relay != null) { + relay.onBatch(serializedResiduals, batchTimeMillis); + } + } + + public void setSourceId(int sourceId) { + this.sourceId = sourceId; + } + + private synchronized void onBatch(List serializedResiduals, long batchTimeMillis) { + inFlight.remove(batchTimeMillis); + long now = System.currentTimeMillis(); + for (byte[] serializedResidual : serializedResiduals) { + try { + DelayedBundleApplication residual = DelayedBundleApplication.parseFrom(serializedResidual); + if (residual.getApplication().getElement().isEmpty()) { + continue; + } + long delayMillis = + residual.hasRequestedTimeDelay() + ? residual.getRequestedTimeDelay().getSeconds() * 1000 + + residual.getRequestedTimeDelay().getNanos() / 1_000_000 + : 0; + pending.add( + new ScheduledResidual( + now + delayMillis, + residual.getApplication().getElement().toByteArray(), + residualWatermark(residual))); + seenResidual = true; + } catch (Exception e) { + throw new RuntimeException("Failed to parse SDF residual", e); + } + } + advanceWatermark(batchTimeMillis); + } + + // The source watermark a residual promises for its future output; absent means unknown (hold). + private static Instant residualWatermark(DelayedBundleApplication residual) { + if (residual.getApplication().getOutputWatermarksMap().isEmpty()) { + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + long watermark = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + for (org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Timestamp outputWatermark : + residual.getApplication().getOutputWatermarksMap().values()) { + watermark = Math.min(watermark, outputWatermark.getSeconds() * 1000); + } + return Instant.ofEpochMilli(watermark); + } + + private void advanceWatermark(long batchTimeMillis) { + Instant newHigh; + if (pending.isEmpty() && inFlight.isEmpty()) { + // The SDF completed if it ever ran; otherwise hold until the first residual arrives. + newHigh = seenResidual ? BoundedWindow.TIMESTAMP_MAX_VALUE : highWatermark; + } else { + newHigh = BoundedWindow.TIMESTAMP_MAX_VALUE; + for (ScheduledResidual scheduled : pending) { + newHigh = earlier(newHigh, scheduled.watermark); + } + for (List taken : inFlight.values()) { + for (ScheduledResidual scheduled : taken) { + newHigh = earlier(newHigh, scheduled.watermark); + } + } + } + // This stage's output can never be ahead of the input still to come. + newHigh = earlier(newHigh, upstreamHighWatermark()); + if (newHigh.isBefore(highWatermark)) { + newHigh = highWatermark; + } + GlobalWatermarkHolder.add( + sourceId, new SparkWatermarks(highWatermark, newHigh, new Instant(batchTimeMillis))); + highWatermark = newHigh; + } + + // Slowest watermark among the upstreams still reporting. A source drops out of the holder once it + // stops reporting, which is how an impulse behaves after its single emission, so an absent + // upstream constrains nothing and the residual holds stay in charge. + private Instant upstreamHighWatermark() { + Map committed = + upstreamSourceIds.isEmpty() ? null : GlobalWatermarkHolder.get(batchDurationMillis); + Instant high = BoundedWindow.TIMESTAMP_MAX_VALUE; + if (committed == null) { + return high; + } + for (Integer upstreamSourceId : upstreamSourceIds) { + SparkWatermarks upstream = committed.get(upstreamSourceId); + if (upstream != null) { + high = earlier(high, upstream.getHighWatermark()); + } + } + return high; + } + + private static Instant earlier(Instant a, Instant b) { + return a.isBefore(b) ? a : b; + } + + private synchronized List takeDue(long validTimeMillis) { + List due = new ArrayList<>(); + List taken = new ArrayList<>(); + Iterator iterator = pending.iterator(); + while (iterator.hasNext()) { + ScheduledResidual scheduled = iterator.next(); + if (scheduled.dueMillis <= validTimeMillis) { + due.add(scheduled.elementBytes); + taken.add(scheduled); + iterator.remove(); + } + } + if (!taken.isEmpty()) { + inFlight.merge( + validTimeMillis, + taken, + (existing, added) -> { + existing.addAll(added); + return existing; + }); + } + return due; + } + + private static class ScheduledResidual { + private final long dueMillis; + private final byte[] elementBytes; + private final Instant watermark; + + ScheduledResidual(long dueMillis, byte[] elementBytes, Instant watermark) { + this.dueMillis = dueMillis; + this.elementBytes = elementBytes; + this.watermark = watermark; + } + } + + /** Input stream emitting the encoded residuals due for resumption at each micro-batch. */ + public static class ResidualInputDStream extends ConstantInputDStream { + + private final String relayId; + + public ResidualInputDStream(StreamingContext ssc, String relayId) { + super(ssc, emptyRdd(ssc), JavaSparkContext$.MODULE$.fakeClassTag()); + this.relayId = relayId; + } + + private static RDD emptyRdd(StreamingContext ssc) { + return ssc.sparkContext().emptyRDD(JavaSparkContext$.MODULE$.fakeClassTag()); + } + + @Override + public Option> compute(Time validTime) { + SdfResidualRelay relay = RELAYS.get(relayId); + if (relay == null) { + return Option.apply(emptyRdd(context())); + } + List due = relay.takeDue(validTime.milliseconds()); + if (due.isEmpty()) { + return Option.apply(emptyRdd(context())); + } + JavaSparkContext jsc = JavaSparkContext.fromSparkContext(context().sparkContext()); + return Option.apply(jsc.parallelize(due).rdd()); + } + } +} diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java index 98601389f5c9..59ba8a8d406c 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkExecutableStageFunctionTest.java @@ -20,6 +20,10 @@ import static org.apache.beam.sdk.util.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.verify; @@ -33,6 +37,9 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication; +import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.model.pipeline.v1.RunnerApi.Components; import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload; @@ -59,6 +66,7 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.junit.Before; import org.junit.Test; @@ -102,7 +110,8 @@ public void setUpMocks() throws Exception { MockitoAnnotations.initMocks(this); when(contextFactory.get(any())).thenReturn(stageContext); when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory); - when(stageBundleFactory.getBundle(any(), any(), any(), any(BundleProgressHandler.class))) + when(stageBundleFactory.getBundle( + any(), any(), any(), any(BundleProgressHandler.class), any(), any())) .thenReturn(remoteBundle); @SuppressWarnings("unchecked") ImmutableMap inputReceiver = @@ -126,7 +135,8 @@ public void expectedInputsAreSent() throws Exception { SparkExecutableStageFunction function = getFunction(Collections.emptyMap()); RemoteBundle bundle = Mockito.mock(RemoteBundle.class); - when(stageBundleFactory.getBundle(any(), any(), any(), any(BundleProgressHandler.class))) + when(stageBundleFactory.getBundle( + any(), any(), any(), any(BundleProgressHandler.class), any(), any())) .thenReturn(bundle); @SuppressWarnings("unchecked") @@ -247,7 +257,9 @@ public void testStageBundleClosed() throws Exception { List> inputs = new ArrayList<>(); inputs.add(WindowedValues.valueInGlobalWindow(0)); function.call(inputs.iterator()); - verify(stageBundleFactory).getBundle(any(), any(), any(), any(BundleProgressHandler.class)); + verify(stageBundleFactory) + .getBundle(any(), any(), any(), any(BundleProgressHandler.class), any(), any()); + verify(stageBundleFactory).getInstructionRequestHandler(); verify(stageBundleFactory).getProcessBundleDescriptor(); verify(stageBundleFactory).close(); verifyNoMoreInteractions(stageBundleFactory); @@ -260,8 +272,72 @@ public void testNoCallOnEmptyInputIterator() throws Exception { verifyNoInteractions(stageBundleFactory); } + @Test + public void sdfResidualsAreEmittedInStreamingMode() throws Exception { + DelayedBundleApplication residual = + DelayedBundleApplication.newBuilder() + .setApplication( + BundleApplication.newBuilder() + .setElement(ByteString.copyFromUtf8("residual-element"))) + .build(); + when(stageBundleFactory.getBundle( + any(), any(), any(), any(BundleProgressHandler.class), any(), any())) + .thenAnswer( + invocation -> { + BundleCheckpointHandler handler = invocation.getArgument(5); + handler.onCheckpoint( + ProcessBundleResponse.newBuilder().addResidualRoots(residual).build()); + return remoteBundle; + }); + + SparkExecutableStageFunction function = getFunction(Collections.emptyMap(), true); + List> inputs = new ArrayList<>(); + inputs.add(WindowedValues.valueInGlobalWindow(0)); + Iterator outputs = function.call(inputs.iterator()); + + RawUnionValue only = outputs.next(); + assertEquals(SparkExecutableStageFunction.SDF_RESIDUAL_TAG, only.getUnionTag()); + assertArrayEquals(residual.toByteArray(), (byte[]) only.getValue()); + assertFalse(outputs.hasNext()); + // Streaming mode must not drain residuals in place with extra bundles. + verify(stageBundleFactory, Mockito.times(1)) + .getBundle(any(), any(), any(), any(BundleProgressHandler.class), any(), any()); + } + + @Test + public void unboundedResidualIsRejectedInBatchMode() throws Exception { + DelayedBundleApplication residual = + DelayedBundleApplication.newBuilder() + .setApplication( + BundleApplication.newBuilder() + .setElement(ByteString.copyFromUtf8("residual-element")) + .setIsBounded(RunnerApi.IsBounded.Enum.UNBOUNDED)) + .build(); + when(stageBundleFactory.getBundle( + any(), any(), any(), any(BundleProgressHandler.class), any(), any())) + .thenAnswer( + invocation -> { + BundleCheckpointHandler handler = invocation.getArgument(5); + handler.onCheckpoint( + ProcessBundleResponse.newBuilder().addResidualRoots(residual).build()); + return remoteBundle; + }); + + SparkExecutableStageFunction function = getFunction(Collections.emptyMap()); + List> inputs = new ArrayList<>(); + inputs.add(WindowedValues.valueInGlobalWindow(0)); + + // Draining an unbounded residual would never terminate, so batch mode must fail fast. + assertThrows(UnsupportedOperationException.class, () -> function.call(inputs.iterator())); + } + private SparkExecutableStageFunction getFunction( Map outputMap) { + return getFunction(outputMap, false); + } + + private SparkExecutableStageFunction getFunction( + Map outputMap, boolean emitSdfResiduals) { return new SparkExecutableStageFunction<>( pipelineOptions, stagePayload, @@ -270,6 +346,8 @@ private SparkExecutableStageFunction ge contextFactory, Collections.emptyMap(), metricsAccumulator, - null); + null, + null, + emitSdfResiduals); } } diff --git a/sdks/python/apache_beam/runners/portability/spark_runner_test.py b/sdks/python/apache_beam/runners/portability/spark_runner_test.py index 4152b8d09f4f..6b07ec76c9ab 100644 --- a/sdks/python/apache_beam/runners/portability/spark_runner_test.py +++ b/sdks/python/apache_beam/runners/portability/spark_runner_test.py @@ -144,37 +144,11 @@ def test_metrics(self): # Skip until Spark runner supports metrics. raise unittest.SkipTest("https://github.com/apache/beam/issues/19496") - def test_sdf(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - def test_unbounded_source_read(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_sdf_with_watermark_tracking(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_sdf_with_sdf_initiated_checkpointing(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_sdf_synthetic_source(self): - # Skip until Spark runner supports SDF. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") - - def test_callbacks_with_exception(self): - # Skip until Spark runner supports bundle finalization. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") - - def test_register_finalizations(self): - # Skip until Spark runner supports bundle finalization. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") - - def test_sdf_with_dofn_as_watermark_estimator(self): - # Skip until Spark runner supports SDF and self-checkpoint. - raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") + # The source self-terminates, but a streaming pipeline on this runner runs + # until its streaming timeout elapses, so it never reports completion. + raise unittest.SkipTest( + "Spark portable streaming pipelines do not self-terminate.") def test_pardo_dynamic_timer(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20179")