marin-ma commented on code in PR #12370: URL: https://github.com/apache/gluten/pull/12370#discussion_r3535021092
########## shims/spark33/src/main/scala/org/apache/spark/storage/GlutenShuffleBlockFetcherIterator.scala: ########## @@ -0,0 +1,1506 @@ +/* + * 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.spark.storage + +import org.apache.spark.{MapOutputTracker, TaskContext} +import org.apache.spark.MapOutputTracker.SHUFFLE_PUSH_MAP_ID +import org.apache.spark.errors.SparkCoreErrors +import org.apache.spark.internal.Logging +import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer} +import org.apache.spark.network.shuffle._ +import org.apache.spark.network.shuffle.checksum.{Cause, ShuffleChecksumHelper} +import org.apache.spark.network.util.TransportConf +import org.apache.spark.shuffle.ShuffleReadMetricsReporter +import org.apache.spark.util.{TaskCompletionListener, Utils} + +import io.netty.util.internal.OutOfDirectMemoryError +import org.apache.commons.io.IOUtils + +import javax.annotation.concurrent.GuardedBy + +import java.io.{InputStream, IOException} +import java.nio.channels.ClosedByInterruptException +import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue, TimeUnit} +import java.util.zip.CheckedInputStream + +import scala.collection.mutable +import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Queue} +import scala.util.{Failure, Success} + +/** + * An iterator that fetches multiple blocks. For local blocks, it fetches from the local block + * manager. For remote blocks, it fetches them using the provided BlockTransferService. + * + * This creates an iterator of (BlockID, InputStream) tuples so the caller can handle blocks in a + * pipelined fashion as they are received. + * + * The implementation throttles the remote fetches so they don't exceed maxBytesInFlight to avoid + * using too much memory. + * + * @param context + * [[TaskContext]], used for metrics update + * @param shuffleClient + * [[BlockStoreClient]] for fetching remote blocks + * @param blockManager + * [[BlockManager]] for reading local blocks + * @param blocksByAddress + * list of blocks to fetch grouped by the [[BlockManagerId]]. For each block we also require two + * info: 1. the size (in bytes as a long field) in order to throttle the memory usage; 2. the + * mapIndex for this block, which indicate the index in the map stage. Note that zero-sized blocks + * are already excluded, which happened in + * [[org.apache.spark.MapOutputTracker.convertMapStatuses]]. + * @param mapOutputTracker + * [[MapOutputTracker]] for falling back to fetching the original blocks if we fail to fetch + * shuffle chunks when push based shuffle is enabled. + * @param streamWrapper + * A function to wrap the returned input stream. + * @param maxBytesInFlight + * max size (in bytes) of remote blocks to fetch at any given point. + * @param maxReqsInFlight + * max number of remote requests to fetch blocks at any given point. + * @param maxBlocksInFlightPerAddress + * max number of shuffle blocks being fetched at any given point for a given remote host:port. + * @param maxReqSizeShuffleToMem + * max size (in bytes) of a request that can be shuffled to memory. + * @param maxAttemptsOnNettyOOM + * The max number of a block could retry due to Netty OOM before throwing the fetch failure. + * @param detectCorrupt + * whether to detect any corruption in fetched blocks. + * @param checksumEnabled + * whether the shuffle checksum is enabled. When enabled, Spark will try to diagnose the cause of + * the block corruption. + * @param checksumAlgorithm + * the checksum algorithm that is used when calculating the checksum value for the block data. + * @param shuffleMetrics + * used to report shuffle metrics. + * @param doBatchFetch + * fetch continuous shuffle blocks from same executor in batch if the server side supports. + */ +final class GlutenShuffleBlockFetcherIterator( + context: TaskContext, + shuffleClient: BlockStoreClient, + blockManager: BlockManager, + mapOutputTracker: MapOutputTracker, + blocksByAddress: Iterator[(BlockManagerId, Seq[(BlockId, Long, Int)])], + streamWrapper: (BlockId, InputStream) => InputStream, + maxBytesInFlight: Long, + maxReqsInFlight: Int, + maxBlocksInFlightPerAddress: Int, + val maxReqSizeShuffleToMem: Long, + maxAttemptsOnNettyOOM: Int, + detectCorrupt: Boolean, + detectCorruptUseExtraMemory: Boolean, + checksumEnabled: Boolean, + checksumAlgorithm: String, + shuffleMetrics: ShuffleReadMetricsReporter, + doBatchFetch: Boolean) + extends GlutenShuffleBlockFetcherIteratorBase + with DownloadFileManager + with Logging { + + import ShuffleBlockFetcherIterator._ + + // Make remote requests at most maxBytesInFlight / 5 in length; the reason to keep them + // smaller than maxBytesInFlight is to allow multiple, parallel fetches from up to 5 + // nodes, rather than blocking on reading output from one node. + private val targetRemoteRequestSize = math.max(maxBytesInFlight / 5, 1L) + + /** + * Total number of blocks to fetch. + */ + private[this] var numBlocksToFetch = 0 + + /** + * The number of blocks processed by the caller. The iterator is exhausted when + * [[numBlocksProcessed]] == [[numBlocksToFetch]]. + */ + private[this] var numBlocksProcessed = 0 + + private[this] val startTimeNs = System.nanoTime() + + /** Host local blocks to fetch, excluding zero-sized blocks. */ + private[this] val hostLocalBlocks = scala.collection.mutable.LinkedHashSet[(BlockId, Int)]() + + /** + * A queue to hold our results. This turns the asynchronous model provided by + * [[org.apache.spark.network.BlockTransferService]] into a synchronous model (iterator). + */ + private[this] val results = new LinkedBlockingQueue[FetchResult] + + /** + * Current [[FetchResult]] being processed per thread. We track this so we can release the current + * buffer in case of a runtime exception when processing the current buffer. Using + * ConcurrentHashMap to support concurrent access from multiple threads while allowing cleanup + * from any thread. + */ + private[this] val currentResults: ConcurrentHashMap[Long, SuccessFetchResult] = + new ConcurrentHashMap[Long, SuccessFetchResult]() Review Comment: > it’s possible to have multiple shuffle block streams outstanding concurrently Yes, but a single thread can only handle one stream at a time. Therefore, using threadId to close the buffer is safe. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
