From d99c843287461d6b18174a9fd3488c6046cd62c7 Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Wed, 15 Jul 2026 12:04:01 +1000 Subject: [PATCH] [Spark] Support SDF self-checkpointing and bundle finalization on the portable runner The portable Spark runner never registered a bundle checkpoint handler, so a splittable DoFn that self-checkpoints failed on its first bundle and unbounded reads could not run at all. In batch, the executable stage function now collects residual roots and re-feeds them in fresh bundles until the SDK returns none, resuming each residual at its requested time. Unbounded residuals are rejected, since draining them would never terminate. The stage also registers an in-memory bundle finalizer. In streaming, residuals leave the stage on a reserved union tag and are relayed on the driver: they are held until their resume time, fed back as stage input in a later micro-batch, and their output watermarks advance the stage's watermark so downstream event-time windows fire. The relay bounds its watermark by the stage's upstream sources, and is keyed per job so concurrent jobs stay isolated. Unskips the splittable DoFn and bundle finalization tests for the Spark runner. --- ..._PostCommit_Java_PVR_Spark3_Streaming.json | 2 +- .../beam_PostCommit_Java_PVR_Spark_Batch.json | 2 +- ...PostCommit_Java_ValidatesRunner_Spark.json | 2 +- ...stCommit_Python_ValidatesRunner_Spark.json | 3 +- CHANGES.md | 2 + .../spark/job-server/spark_job_server.gradle | 4 +- .../runners/spark/SparkPipelineRunner.java | 5 +- .../SparkBatchPortablePipelineTranslator.java | 8 +- .../SparkExecutableStageFunction.java | 184 +++++++++++- ...rkStreamingPortablePipelineTranslator.java | 70 ++++- .../streaming/SdfResidualRelay.java | 268 ++++++++++++++++++ .../SparkExecutableStageFunctionTest.java | 86 +++++- .../runners/portability/spark_runner_test.py | 34 +-- 13 files changed, 616 insertions(+), 54 deletions(-) create mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SdfResidualRelay.java 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")