ziting-openai commented on code in PR #5513: URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3876003185
########## spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala: ########## @@ -0,0 +1,337 @@ +/* + * 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.comet.shuffle + +import java.io.IOException +import java.lang.reflect.InvocationTargetException + +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.shuffle.ShuffleHandle +import org.apache.spark.storage.BlockManagerId + +import org.apache.comet.CometConf +import org.apache.comet.util.ClassLoaders + +/** Resolves an application's optional Celeborn client without adding a Celeborn dependency. */ +object CelebornShufflePusherFactory { + + private val CELEBORN_SHUFFLE_HANDLE = "org.apache.spark.shuffle.celeborn.CelebornShuffleHandle" + private val CELEBORN_SPARK_UTILS = "org.apache.spark.shuffle.celeborn.SparkUtils" + private val CELEBORN_SHUFFLE_CLIENT = "org.apache.celeborn.client.ShuffleClient" + private val CELEBORN_CONF = "org.apache.celeborn.common.CelebornConf" + private val CELEBORN_USER_IDENTIFIER = "org.apache.celeborn.common.identity.UserIdentifier" + private val MAX_STAGE_ATTEMPTS = 1 << 15 + private val MAX_TASK_ATTEMPTS = 1 << 16 + + private[shuffle] def encodeAttemptNumber(stageAttempt: Int, taskAttempt: Int): Int = { + require( + stageAttempt >= 0 && stageAttempt < MAX_STAGE_ATTEMPTS, + s"Celeborn stage attempt must be between 0 and ${MAX_STAGE_ATTEMPTS - 1}: $stageAttempt") + require( + taskAttempt >= 0 && taskAttempt < MAX_TASK_ATTEMPTS, + s"Celeborn task attempt must be between 0 and ${MAX_TASK_ATTEMPTS - 1}: $taskAttempt") + (stageAttempt << 16) | taskAttempt + } + + /** Bind one already-resolved Celeborn generation to one Spark map attempt. */ + def create( + conf: SparkConf, + client: AnyRef, + celebornShuffleId: Int, + numMappers: Int, + numPartitions: Int, + taskContext: TaskContext): CelebornShufflePartitionPusher = { + val frameEntry = CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES + val maxFrameBytes = conf.getSizeAsBytes(frameEntry.key, frameEntry.defaultValue.get.toString) + require( + maxFrameBytes >= 20 && maxFrameBytes <= Int.MaxValue - 16, + "Celeborn frame bytes must fit a complete Comet frame and a request header") + val limitEntry = CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES + val maxInFlightBytes = + conf.getSizeAsBytes(limitEntry.key, limitEntry.defaultValue.get.toString) + require( + maxInFlightBytes >= 76 && maxInFlightBytes <= Int.MaxValue, + "Celeborn executor in-flight bytes must fit three complete frames and a request header") + + new CelebornShufflePartitionPusher( + client, + celebornShuffleId, + taskContext.partitionId(), + encodeAttemptNumber(taskContext.stageAttemptNumber(), taskContext.attemptNumber()), + numMappers, + numPartitions, + maxFrameBytes.toInt, + maxInFlightBytes.toInt) + } + + /** + * Reuse the Celeborn client described by Spark's actual Celeborn handle. The acquisition hook + * runs before generation resolution so manager shutdown can release clients even if setup + * fails. + */ + def createFromHandle( + conf: SparkConf, + handle: ShuffleHandle, + taskContext: TaskContext, + onClientAcquired: AnyRef => Unit, + onShuffleGenerationResolved: (Int, Int) => Unit, + onShuffleGenerationInvalidated: (Int, Int) => Unit = (_, _) => (), + onShuffleGenerationInvalidationUnsafe: (Int, Int) => Boolean = (_, _) => false) + : ResolvedCelebornShufflePusher = { + try { + val handleClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_HANDLE) + require( + handleClass.isInstance(handle), + "Native Comet shuffle requires an actual Celeborn shuffle handle; " + + s"received ${handle.getClass.getName}") + + val sparkUtilsClass = ClassLoaders.loadClass(CELEBORN_SPARK_UTILS) + val shuffleClientClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_CLIENT) + val celebornConfClass = ClassLoaders.loadClass(CELEBORN_CONF) + val userIdentifierClass = ClassLoaders.loadClass(CELEBORN_USER_IDENTIFIER) + val celebornConf = + sparkUtilsClass.getMethod("fromSparkConf", classOf[SparkConf]).invoke(null, conf) + + def handleValue(name: String): AnyRef = handleClass.getMethod(name).invoke(handle) + + val client = acquireClient( + shuffleClientClass + .getMethod( Review Comment: [P1] Preserve Spark shuffle encryption when creating the Celeborn client This native path invokes the six-argument `ShuffleClient.get`, which in Celeborn 0.7 delegates to the seven-argument overload with `Optional.empty()` for its crypto handler. The normal Spark shuffle manager instead supplies `SparkCommonUtils.getCryptoHandler(conf)`. If a native Comet shuffle initializes an executor's shared Celeborn client first while `spark.io.encryption.enabled=true`, shuffle data is transmitted and stored as plaintext despite the application explicitly requiring encryption. Please prefer the crypto-aware overload with Spark's configured handler, retain a compatible fallback only for older clients, and add a native-first encrypted-shuffle regression. ########## native/shuffle/src/writers/rss/rss_partition_writer.rs: ########## @@ -125,36 +132,616 @@ impl RssPartitionWriter { where I: Iterator<Item = Result<RecordBatch>>, { - for batch in batches.by_ref() { - let batch = batch?; - self.frame.clear(); + let result = (|| { + for batch in batches.by_ref() { + self.push_batch_within_limit(partition_id, &batch?, metrics)?; + } + Ok(()) + })(); + if result.is_err() { + // Earlier frames may have been accepted remotely; a partial map cannot be resumed or + // committed after its input, encoding, reservation, or callback fails. + self.failed = true; + } + result + } - let encoded_size = self.block_writer.write_batch( - &batch, - &mut Cursor::new(&mut self.frame), - &mut self.compression_context, - &metrics.encode_time, - )?; + fn push_batch_within_limit( + &mut self, + partition_id: i32, + batch: &RecordBatch, + metrics: &ShufflePartitionerMetrics, + ) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + + // Estimate only live nested rows before allocation. Dictionary values remain charged + // because Arrow's dense garbage collection scans their complete value tables. + let original_size = Self::estimated_pre_compaction_ipc_data_size(batch)?; + let compaction_scratch = Self::estimated_compaction_scratch(batch)?; + let minimum_size = Self::estimated_minimum_compacted_ipc_data_size(batch)?; + if minimum_size > self.max_frame_size && batch.num_rows() > 1 { + return self.push_split_batch(partition_id, batch, metrics); + } - if encoded_size == 0 { - continue; + // Native IPC, its JNI byte array, and Celeborn's copied transport request overlap. + // Acquire all three copies before compaction/encoding to avoid allocation-before-admission + // and deadlocks caused by growing another task's reservation after partial acquisition. + let overlapping_copies = self.max_frame_size.checked_mul(3).ok_or_else(|| { Review Comment: [P2] Size encoding admission to the actual frame instead of its configured maximum Every nonempty batch reserves `3 * self.max_frame_size` before compaction and encoding, regardless of its actual size. With the default 64 MiB frame maximum and 256 MiB executor-wide in-flight limit, each tiny frame therefore acquires 192 MiB plus its header, so two map tasks cannot encode concurrently on the same executor. The reservation is reduced only after native IPC encoding/compression finishes and Java claims the actual frame. This serializes CPU-intensive shuffle encoding across all executor cores under default settings. Please reserve a conservative per-batch upper bound while preserving atomic three-copy accounting, and add a concurrent small-frame regression. -- 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]
