Copilot commented on code in PR #12370: URL: https://github.com/apache/gluten/pull/12370#discussion_r3552940188
########## shims/spark33/src/main/scala/org/apache/spark/storage/GlutenShuffleBlockFetcherIterator.scala: ########## @@ -0,0 +1,1503 @@ +/* + * 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.newKeySet[SuccessFetchResult]() + + /** + * Queue of fetch requests to issue; we'll pull requests off this gradually to make sure that the + * number of bytes in flight is limited to maxBytesInFlight. + */ + private[this] val fetchRequests = new Queue[FetchRequest] + + /** + * Queue of fetch requests which could not be issued the first time they were dequeued. These + * requests are tried again when the fetch constraints are satisfied. + */ + private[this] val deferredFetchRequests = new HashMap[BlockManagerId, Queue[FetchRequest]]() + + /** Current bytes in flight from our requests */ + private[this] var bytesInFlight = 0L + + /** Current number of requests in flight */ + private[this] var reqsInFlight = 0 + + /** Current number of blocks in flight per host:port */ + private[this] val numBlocksInFlightPerAddress = new HashMap[BlockManagerId, Int]() + + /** + * Count the retry times for the blocks due to Netty OOM. The block will stop retry if retry times + * has exceeded the [[maxAttemptsOnNettyOOM]]. + */ + private[this] val blockOOMRetryCounts = new HashMap[String, Int] + + /** + * The blocks that can't be decompressed successfully, it is used to guarantee that we retry at + * most once for those corrupted blocks. + */ + private[this] val corruptedBlocks = mutable.HashSet[BlockId]() + + /** + * Whether the iterator is still active. If isZombie is true, the callback interface will no + * longer place fetched blocks into [[results]]. + */ + @GuardedBy("this") + private[this] var isZombie = false + + /** + * A set to store the files used for shuffling remote huge blocks. Files in this set will be + * deleted when cleanup. This is a layer of defensiveness against disk file leaks. + */ + @GuardedBy("this") + private[this] val shuffleFilesSet = mutable.HashSet[DownloadFile]() + + private[this] val onCompleteCallback = new GlutenShuffleFetchCompletionListener(this) + + private[this] val pushBasedFetchHelper = new GlutenPushBasedFetchHelper( + this, + shuffleClient, + blockManager, + mapOutputTracker) + + initialize() + + override def createTempFile(transportConf: TransportConf): DownloadFile = { + // we never need to do any encryption or decryption here, regardless of configs, because that + // is handled at another layer in the code. When encryption is enabled, shuffle data is written + // to disk encrypted in the first place, and sent over the network still encrypted. + new SimpleDownloadFile( + blockManager.diskBlockManager.createTempLocalBlock()._2, + transportConf) + } + + override def registerTempFileToClean(file: DownloadFile): Boolean = synchronized { + if (isZombie) { + false + } else { + shuffleFilesSet += file + true + } + } + + /** + * Mark the iterator as zombie, and release all buffers that haven't been deserialized yet. + */ + private[storage] def cleanup(): Unit = { + synchronized { + isZombie = true + } + // Release all current result buffers from all threads + currentResults.forEach(result => result.buf.release()) + currentResults.clear() Review Comment: `cleanup()` releases buffers in `currentResults` without first removing them from the set. If a stream is closed concurrently (release callback removes + releases), this can double-release the same `ManagedBuffer`, which is unsafe and can crash. Make cleanup remove entries atomically before releasing (e.g., snapshot + remove) so each buffer is released exactly once regardless of close/cleanup interleavings. ########## cpp/velox/shuffle/ReaderThreadPool.cc: ########## @@ -0,0 +1,99 @@ +/* + * 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. + */ + +#include "shuffle/ReaderThreadPool.h" +#include <glog/logging.h> + +namespace gluten { + +ReaderThreadPool::ReaderThreadPool(size_t numThreads) : numThreads_(numThreads) { + workers_.reserve(numThreads); + for (size_t i = 0; i < numThreads; ++i) { + workers_.emplace_back([this]() { workerThread(); }); + } + LOG(WARNING) << "Created ReaderThreadPool with " << numThreads << " threads."; +} + +ReaderThreadPool::~ReaderThreadPool() { + shutdown(); +} + +void ReaderThreadPool::submitBatch(std::vector<Task> tasks, int32_t priority) { + std::lock_guard<std::mutex> lock(taskQueueMtx_); + if (stop_.load(std::memory_order_acquire)) { + return; + } + for (auto& task : tasks) { + tasks_.push({std::move(task), priority}); + } +} + +void ReaderThreadPool::start() { + // Wake up all worker threads to start processing. + wakeUpCV_.notify_all(); + LOG(WARNING) << "Started ReaderThreadPool execution."; +} Review Comment: `ReaderThreadPool::start()` logs at WARNING level even though starting the pool is expected behavior. With async GPU shuffle enabled, `start()` is invoked per shuffle reader creation, so this will generate a large volume of WARNING logs and can drown out real operational warnings. -- 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]
