ziting-openai commented on code in PR #5513: URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3876194428
########## 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: Verified fixed at 26efda6a: native client initialization now selects the crypto-aware Celeborn 0.7 overload with the Spark handler, retains the 0.6 fallback, and covers native-first encrypted plus fail-closed behavior. ########## 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: Verified fixed at 26efda6a: admission now scales to a conservative per-batch bound, retries release their full reservation, and the new two-thread regression confirms concurrent small-frame encoding under production limits. -- 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]
