Abacn commented on code in PR #39971: URL: https://github.com/apache/beam/pull/39971#discussion_r3927356127
########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java: ########## @@ -0,0 +1,184 @@ +/* + * 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.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.Path; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager.CancellableFSDataOutputStream; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Durable state of one Beam unbounded source under the per source checkpoint location Spark hands + * to {@code toMicroBatchStream}. + * + * <p>{@code <location>/splits} pins the split list, written once by the driver. {@code + * <location>/marks/<splitId>/<epoch>} holds the coded checkpoint mark of a split at the end of the + * batch ending at that epoch. All IO goes through Spark's {@link CheckpointFileManager}, writes are + * atomic renames. + */ +public final class BeamSourceCheckpoint { + + private static final Logger LOG = LoggerFactory.getLogger(BeamSourceCheckpoint.class); + + private static final String SPLITS_FILE = "splits"; + private static final String MARKS_DIR = "marks"; + + /** A purge further than this above the last one lists the directory instead of probing epochs. */ + private static final long MAX_BLIND_PURGE_RANGE = 1_000L; + + private final String location; + private final CheckpointFileManager fm; + private final Path root; + private final Path splitsPath; + private final Path marksRoot; + + /** Splits whose mark directory this instance already created. */ + private final Set<Integer> preparedSplits = ConcurrentHashMap.newKeySet(); + + /** Per split, every mark epoch strictly below the value is known to be deleted. */ + private final ConcurrentMap<Integer, Long> purgeFloors = new ConcurrentHashMap<>(); + + public BeamSourceCheckpoint(String checkpointLocation, Configuration hadoopConf) { + this.location = checkpointLocation; + this.root = new Path(checkpointLocation); + this.fm = CheckpointFileManager.create(root, hadoopConf); + this.splitsPath = new Path(root, SPLITS_FILE); + this.marksRoot = new Path(root, MARKS_DIR); + } + + public String location() { + return location; + } + + /** The pinned split list, or null if none was pinned yet. */ + public @Nullable List<UnboundedSource<?, ?>> readSplits() throws IOException { + if (!fm.exists(splitsPath)) { + return null; + } + @SuppressWarnings("unchecked") // written by writeSplits as an ArrayList of sources + List<UnboundedSource<?, ?>> splits = + (List<UnboundedSource<?, ?>>) + SerializableUtils.deserializeFromByteArray(read(splitsPath), "splits at " + splitsPath); + return splits; + } + + /** Pins the split list, fails if one is pinned already. */ + public void writeSplits(List<? extends UnboundedSource<?, ?>> splits) throws IOException { + fm.mkdirs(root); + if (fm.exists(splitsPath)) { + throw new IOException("Split list already pinned at " + splitsPath); + } + write(splitsPath, SerializableUtils.serializeToByteArray(new ArrayList<>(splits)), false); + LOG.info("Pinned {} split(s) at {}.", splits.size(), splitsPath); + } + + public void writeMark(int splitId, long epoch, byte[] codedMark) throws IOException { + if (preparedSplits.add(splitId)) { + fm.mkdirs(marksDir(splitId)); + } + write(markPath(splitId, epoch), codedMark, true); + } + + /** The coded mark of a split at an epoch, or null if absent. */ + public byte @Nullable [] readMark(int splitId, long epoch) throws IOException { + Path path = markPath(splitId, epoch); + if (!fm.exists(path)) { + return null; + } + return read(path); + } + + /** + * Deletes every mark of a split with an epoch strictly below {@code epoch}. Lists the directory + * once per split, later calls delete the range above the previous floor only. Idempotent. + */ + public void purgeMarksBelow(int splitId, long epoch) throws IOException { + Long floor = purgeFloors.get(splitId); + if (floor != null && epoch - floor > MAX_BLIND_PURGE_RANGE) { + floor = null; + } + if (floor == null) { + Path dir = marksDir(splitId); + if (!fm.exists(dir)) { + return; + } + for (FileStatus status : fm.list(dir)) { + long existing = parseEpoch(status.getPath().getName()); + if (existing >= 0 && existing < epoch) { + fm.delete(status.getPath()); + } + } + purgeFloors.put(splitId, epoch); + return; + } + for (long e = floor; e < epoch; e++) { Review Comment: It issues individual synchronous delete() RPCs for every epoch. On cloud object stores like GCS/S3, this can result in hundreds of sequential HTTP calls on every commit. Use directory listing or batch deletions instead. ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java: ########## @@ -0,0 +1,277 @@ +/* + * 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.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.spark.SparkEnv; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.connector.read.streaming.Offset; +import org.apache.spark.storage.BlockManager; +import org.apache.spark.storage.BlockManagerId; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import scala.collection.Iterator; + +/** + * Driver side {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * + * <p>Offsets are opaque epochs, {@link #latestOffset()} advances by one every trigger. The source + * is split once and the splits are pinned under the checkpoint location, every batch of every run + * plans the same partitions. {@link #commit} purges marks below the committed epoch on a background + * thread, Spark never asks for those again. + * + * <p>Each split prefers the executor rendezvous hashing assigns it, so an executor joining or + * leaving moves only the splits of that executor. Locality is a hint, the reader cache restores a + * split from its durable mark wherever it lands. + */ +public class BeamMicroBatchStream<T> implements MicroBatchStream { + + private static final Logger LOG = LoggerFactory.getLogger(BeamMicroBatchStream.class); + + private final BeamSourceSpec<T> spec; + private final String checkpointLocation; + private final BeamSourceCheckpoint checkpoint; + + private final ExecutorService purger = + Executors.newSingleThreadExecutor( + runnable -> { + Thread thread = new Thread(runnable, "beam-source-mark-purge"); + thread.setDaemon(true); + return thread; + }); + private final AtomicBoolean purgeInFlight = new AtomicBoolean(); + private final AtomicLong purgeRequested = new AtomicLong(); + + private long epoch; + private @Nullable List<UnboundedSource<T, ?>> splits; + + BeamMicroBatchStream(BeamSourceSpec<T> spec, String checkpointLocation) { + this.spec = spec; + this.checkpointLocation = checkpointLocation; + this.checkpoint = + new BeamSourceCheckpoint(checkpointLocation, spec.hadoopConf().value().value()); + } + + @Override + public Offset initialOffset() { + return BeamOffset.ZERO; + } + + @Override + public synchronized Offset latestOffset() { + return new BeamOffset(++epoch); + } + + @Override + public Offset deserializeOffset(String json) { + BeamOffset offset = BeamOffset.fromJson(json); + fastForwardEpoch(offset.epoch()); + return offset; + } + + @Override + public InputPartition[] planInputPartitions(Offset start, Offset end) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); + List<UnboundedSource<T, ?>> pinned = splits(); + long[] quotas = splitQuotas(spec.maxRecordsPerBatch(), pinned.size()); + List<String> executors = sortedExecutors(); Review Comment: (from AI reivew) sortedExecutors uses `.master().getPeers(...)` and assumes naming format of Spark internals, both are unsupported and fragile in managed environments. We already support restoring from durable marks when a partition runs on another executor, so we should let Spark handle partition locality. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
