diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0d62838cd27..645717e281f 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -374,6 +374,10 @@ jobs: org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleFallbackSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornConcurrentMaterializationSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleStatisticsSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornLocalFetchFailureSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index f784dc0e510..673fac33a8d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -147,6 +147,10 @@ jobs: org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleFallbackSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornConcurrentMaterializationSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleStatisticsSuite + org.apache.spark.sql.comet.execution.shuffle.CometCelebornLocalFetchFailureSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite diff --git a/common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java b/common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java new file mode 100644 index 00000000000..ff257fb0fb7 --- /dev/null +++ b/common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java @@ -0,0 +1,27 @@ +/* + * 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; + +/** A remote shuffle frame or its encoding workspace cannot fit the configured byte limits. */ +public final class CometShuffleSizeLimitException extends CometNativeException { + public CometShuffleSizeLimitException(String message) { + super(message); + } +} diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md index a268691a3c5..49ee6dd48df 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -45,6 +45,25 @@ These settings can be used to determine which parts of the plan are accelerated ## Shuffle Configuration Settings +For native remote shuffle, `spark.comet.shuffle.rss.maxFrameBytes` limits one complete +encoded frame, while `spark.comet.shuffle.rss.maxInFlightBytes` limits the memory reserved +by map attempts sharing an executor's remote shuffle client. The reservation includes +encoding workspace and overlapping frame copies. An ordinary uncompressed frame needs +approximately seven times its size plus schema and transport overhead. The default 512 MiB +reservation budget accommodates ordinary frames up to the default 64 MiB frame limit. +Compressed frames still need workspace for their uncompressed data. Increase the reservation +budget when larger rows or schemas need more workspace. + +If a row cannot fit the remote limits, Comet materializes a replacement shuffle using its +local writer before publishing the exchange to downstream tasks. The replacement has a separate +shuffle and scheduling identity, so late remote results or failures cannot affect local output. +Independent exchanges can materialize concurrently, and runtime output statistics describe only +the selected destination. Subsequent fetch failures retain Spark's normal recovery behavior. +With dynamic allocation enabled, native Celeborn shuffle also requires either +`spark.shuffle.service.enabled=true` or `spark.dynamicAllocation.shuffleTracking.enabled=true` +(the Spark default) to preserve fallback files. Otherwise exchanges retain ordinary Spark/Celeborn +shuffle, including applications that rely only on remote reliable storage or decommissioning. + diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index decb47a8c2c..0ad8ec992c7 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -257,7 +257,7 @@ even when both its parent and child are non-Comet operators. Applications using Apache Celeborn can use Comet's composite shuffle manager to retain ordinary Spark/Celeborn shuffle while accelerating other operators with Comet. -Native shuffle requires reliable completion tracking for in-flight payloads. Released Celeborn +Native shuffle also requires reliable completion tracking for in-flight payloads. Released Celeborn 0.6.0 and 0.7.0 clients do not provide the required guarantee, so these versions retain ordinary Spark/Celeborn shuffle even when `spark.comet.shuffle.mode=native`. Native shuffle support for these clients requires a safe Celeborn push-completion API. The following settings request @@ -302,8 +302,29 @@ Celeborn to prohibit local fallback for ordinary Spark shuffles. Native frames retain Comet's configured compression; the raw Celeborn client path bypasses Celeborn's additional row compression and decompression. Use `spark.comet.shuffle.rss.maxFrameBytes` and `spark.comet.shuffle.rss.maxInFlightBytes` to bound -encoded frame size and executor-side push admission. These limits include framing and overlapping -native/JNI/client copies; a frame that cannot fit is rejected rather than split across requests. +encoded frame size and executor-side push admission. The defaults are 64 MiB and 512 MiB, +respectively. Admission includes Arrow encoding workspace as well as overlapping native, JNI, +and client frame copies. An uncompressed frame needs roughly seven times its size plus schema +and codec overhead. Compression reduces the transmitted bytes but still needs uncompressed +encoding workspace. + +Comet splits large batches between rows. If a single row, its schema, or its encoding workspace +cannot fit the remote limits, Comet abandons the remote shuffle and materializes a replacement +using its local shuffle writer before downstream tasks can consume the exchange. The replacement +has a separate shuffle and scheduling identity, so late remote results cannot overwrite or skip +local map output, and remote stage failures cannot abort the replacement. Independent exchanges +can materialize concurrently; readers wait for their storage decisions before execution. Runtime +output statistics count only the selected destination. All reads and retries for the replacement +use local files and Spark's block transfer +service, including normal recovery after later fetch failures. Native operators and Comet's +Arrow shuffle format are preserved, and remote admission limits remain enforced. Once remote +output has been published, subsequent failures use the existing Spark/Celeborn recovery path; +Comet does not change that shuffle's destination. Local fallback uses executor disk. When `spark.dynamicAllocation.enabled=true`, native Celeborn shuffle requires +`spark.shuffle.service.enabled=true` or `spark.dynamicAllocation.shuffleTracking.enabled=true` +(the Spark default) so those files remain available. Applications using dynamic allocation with +both settings disabled retain ordinary Spark/Celeborn shuffle, even if remote reliable storage or +decommissioning enables dynamic allocation. Executor shutdown preserves fallback files for the +external shuffle service; explicit shuffle unregister retains the normal local cleanup behavior. AQE reducer coalescing and mapper-range reads are supported, but Celeborn physical-skew chunk reads are not. diff --git a/native/jni-bridge/src/errors.rs b/native/jni-bridge/src/errors.rs index 4b8cb439a1e..a072be95ed1 100644 --- a/native/jni-bridge/src/errors.rs +++ b/native/jni-bridge/src/errors.rs @@ -90,6 +90,11 @@ pub enum CometError { #[error("Comet Internal Error: {0}")] Internal(String), + /// A remote shuffle frame or its encoding workspace cannot fit the configured limits. + /// Preserve this classification so Spark can restart the shuffle with a local writer. + #[error("{0}")] + ShuffleSizeLimit(String), + #[error(transparent)] Arrow { #[from] @@ -214,7 +219,9 @@ impl From for DataFusionError { // own codegen inside the JVM UDF kernel) as an `External` error so it survives the trip // back through DataFusion and can be re-thrown with its exact type at the JNI boundary. // Flattening it to a string here would surface it as a generic CometNativeException. - value @ CometError::JavaException { .. } => DataFusionError::External(Box::new(value)), + value @ (CometError::JavaException { .. } | CometError::ShuffleSizeLimit(_)) => { + DataFusionError::External(Box::new(value)) + } _ => DataFusionError::Execution(value.to_string()), } } @@ -338,6 +345,10 @@ impl jni::errors::ToException for CometError { class: spark_err.exception_class().to_string(), msg: spark_err.to_string(), }, + CometError::ShuffleSizeLimit(message) => Exception { + class: "org/apache/comet/CometShuffleSizeLimitException".to_string(), + msg: message.clone(), + }, _other => Exception { class: "org/apache/comet/CometNativeException".to_string(), msg: self.to_string(), @@ -473,6 +484,15 @@ pub fn unwrap_or_throw_default( fn throw_exception(env: &mut Env, error: &CometError, backtrace: Option) { // If there isn't already an exception? if !env.exception_check() { + // DataFusion operators can wrap the original failure in Context, Shared, or External + // errors. Keep capacity failures typed across those wrappers and the JNI boundary. + if let Some(message) = shuffle_size_limit_message(error) { + let _ = env.throw_new( + jni::jni_str!("org/apache/comet/CometShuffleSizeLimitException"), + JNIString::new(message), + ); + return; + } // ... then throw new exception // Note: in jni 0.22.x, throw/throw_new return Err(JavaException) on success // (to signal the pending exception to Rust callers via `?`). We discard the @@ -554,6 +574,17 @@ fn throw_exception(env: &mut Env, error: &CometError, backtrace: Option) } } +fn shuffle_size_limit_message<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a str> { + let mut cause = Some(error); + while let Some(error) = cause { + if let Some(CometError::ShuffleSizeLimit(message)) = error.downcast_ref::() { + return Some(message); + } + cause = error.source(); + } + None +} + /// Generic fallback throw for an error that isn't a structured `SparkError`. Recognises a /// file-not-found arriving through non-typed wrapping paths and duplicate-field errors; otherwise /// throws the error's natural JVM exception (with the captured backtrace when available). @@ -912,6 +943,36 @@ mod tests { } } + #[test] + #[cfg_attr(miri, ignore)] // miri cannot create a JVM. + fn shuffle_size_limit_survives_datafusion_wrappers_and_jni() { + let message = "Remote shuffle exceeds spark.comet.shuffle.rss.maxInFlightBytes"; + let error = DataFusionError::from(CometError::ShuffleSizeLimit(message.to_string())); + let error = DataFusionError::Shared(Arc::new(DataFusionError::Context( + "executing shuffle writer".to_string(), + Box::new(error), + ))); + jvm() + .attach_current_thread(|env| -> jni::errors::Result<()> { + unwrap_or_throw_default::<()>(env, Err(CometError::from(error))); + assert_pending_java_exception_detailed( + env, + Some("org/apache/comet/CometShuffleSizeLimitException"), + Some(message), + ); + Ok(()) + }) + .unwrap(); + } + + #[test] + fn shuffle_size_limit_is_not_inferred_from_error_text() { + let error = CometError::from(DataFusionError::Execution( + "Remote shuffle exceeds spark.comet.shuffle.rss.maxInFlightBytes".to_string(), + )); + assert!(shuffle_size_limit_message(&error).is_none()); + } + #[test] #[cfg_attr(miri, ignore)] // miri can't call foreign function `dlopen` pub fn error_from_panic() { diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs index a151b6831fc..bf7cbaa4648 100644 --- a/native/shuffle/src/writers/rss/mod.rs +++ b/native/shuffle/src/writers/rss/mod.rs @@ -33,6 +33,7 @@ mod tests { use arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, Result}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time}; + use datafusion_comet_jni_bridge::errors::CometError; use datafusion_comet_jni_bridge::ShufflePartitionPusher; use std::collections::HashMap; use std::io::{self, Cursor}; @@ -411,6 +412,28 @@ mod tests { .unwrap() } + fn single_string_row(bytes: usize) -> RecordBatch { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])), + vec![Arc::new(StringArray::from(vec!["x".repeat(bytes)]))], + ) + .unwrap() + } + + fn size_limit_message(error: &DataFusionError) -> &str { + let DataFusionError::External(cause) = error else { + panic!("expected a typed shuffle size limit, got {error}"); + }; + let Some(CometError::ShuffleSizeLimit(message)) = cause.downcast_ref::() else { + panic!("expected a typed shuffle size limit, got {error}"); + }; + message + } + #[test] #[cfg_attr(miri, ignore)] fn round_trips_all_supported_compression_codecs() { @@ -633,7 +656,15 @@ mod tests { frame_size - 1, ); - assert!(write_batches(&mut writer, 0, vec![batch], &metrics()).is_err()); + let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err(); + let message = size_limit_message(&error); + assert!(message.contains(&format!( + "effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \ + and spark.comet.shuffle.rss.maxInFlightBytes is {} bytes", + frame_size - 1 + ))); + assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes")); + assert!(message.contains(&format!("frame needs at least {frame_size} bytes"))); assert!( pusher.frames().is_empty(), "a single oversized Arrow IPC row must never be fragmented" @@ -786,7 +817,7 @@ mod tests { #[test] fn small_batches_encode_concurrently_with_default_frame_and_admission_limits() { let frame_limit = 64 * 1024 * 1024; - let admission_limit = 256 * 1024 * 1024; + let admission_limit = 512 * 1024 * 1024; let pusher = Arc::new(ConcurrentAdmissionPusher::new(admission_limit)); thread::scope(|scope| { @@ -815,6 +846,115 @@ mod tests { assert_eq!(pusher.state.lock().unwrap().total, 0); } + #[test] + #[cfg_attr(miri, ignore)] + fn large_single_rows_fit_default_frame_and_admission_limits() { + let frame_limit = 64 * 1024 * 1024; + // Celeborn reserves its 16-byte request header separately from the native reservation. + let reservation_limit = 512 * 1024 * 1024 - 16; + for row_size_mib in [38, 50, 63] { + let batch = single_string_row(row_size_mib * 1024 * 1024); + let pusher = Arc::new(ReservationRecordingPusher { + max_reservation: Some(reservation_limit), + capacity: Some(reservation_limit), + ..ReservationRecordingPusher::default() + }); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + frame_limit, + ); + let metrics = metrics(); + + finish_partition(&mut writer, 0, vec![batch.clone()], &metrics).unwrap(); + writer.finish_all(&metrics).unwrap(); + + let frames = pusher.frames.lock().unwrap(); + assert_eq!(frames.len(), 1); + assert!(frames[0].1.len() <= frame_limit); + assert_eq!(decode_frame(&frames[0].1), batch); + assert!(pusher.outstanding.lock().unwrap().is_none()); + assert!(pusher + .reservations + .lock() + .unwrap() + .iter() + .all(|bytes| *bytes <= reservation_limit)); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn reports_admission_limit_when_a_single_row_fits_the_configured_frame_limit() { + let batch = single_string_row(38 * 1024 * 1024); + let frame_limit = 64 * 1024 * 1024; + let reservation_limit = 256 * 1024 * 1024 - 16; + let pusher = Arc::new(ReservationRecordingPusher { + max_reservation: Some(reservation_limit), + capacity: Some(reservation_limit), + ..ReservationRecordingPusher::default() + }); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + frame_limit, + ); + assert!(encoded_frame_size(&batch) < frame_limit); + + let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err(); + + let message = size_limit_message(&error); + assert!( + message.contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes") + ); + assert!(message.contains(&format!(": {frame_limit} bytes)"))); + assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes")); + assert!(message.contains(&format!("{reservation_limit}-byte reservation"))); + assert!(message.contains("encoded bytes fit alongside")); + assert!(message.contains("required reservation is at least")); + assert!(pusher.frames.lock().unwrap().is_empty()); + assert!(pusher.outstanding.lock().unwrap().is_none()); + } + + #[test] + #[cfg_attr(miri, ignore)] + fn reports_workspace_limit_before_encoding_a_single_row() { + let batch = single_string_row(20 * 1024 * 1024); + let frame_limit = 64 * 1024 * 1024; + let reservation_limit = 64 * 1024 * 1024 - 16; + let pusher = Arc::new(ReservationRecordingPusher { + max_reservation: Some(reservation_limit), + capacity: Some(reservation_limit), + ..ReservationRecordingPusher::default() + }); + let mut writer = writer( + &batch, + CompressionCodec::None, + pusher.clone(), + 1, + frame_limit, + ); + + let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err(); + + let message = size_limit_message(&error); + assert!(message.contains("encoding workspace for a single row requires")); + assert!(message.contains("bytes of reservation are needed")); + assert!(message.contains(&format!("only {reservation_limit} bytes per frame"))); + assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes")); + assert!( + message.contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes") + ); + assert!(message.contains(&format!("is {frame_limit} bytes"))); + assert_eq!(pusher.reservations.lock().unwrap().len(), 1); + assert!(pusher.frames.lock().unwrap().is_empty()); + assert!(pusher.outstanding.lock().unwrap().is_none()); + } + #[test] fn codec_workspace_is_rejected_before_encoding_under_a_tiny_budget() { let batch = sample_batch(0, 1); @@ -830,7 +970,12 @@ mod tests { }); let mut writer = writer(&batch, codec, pusher.clone(), 1, 5_456); let error = write_batches(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap_err(); - assert!(error.to_string().contains("encoding workspace")); + let message = size_limit_message(&error); + assert!(message.contains("encoding workspace")); + assert!(message.contains("spark.comet.shuffle.rss.maxInFlightBytes")); + assert!(message + .contains("effective frame limit from spark.comet.shuffle.rss.maxFrameBytes")); + assert!(message.contains("is 5456 bytes")); assert!(pusher.reservations.lock().unwrap().is_empty()); assert!(pusher.frames.lock().unwrap().is_empty()); } diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index 76503011dbf..a7e374f7021 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -29,6 +29,7 @@ use arrow::datatypes::{DataType, Field, Int16Type, Int32Type, Int64Type}; use arrow::ipc::writer::CompressionContext; use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion::common::{DataFusionError, Result}; +use datafusion_comet_jni_bridge::errors::CometError; use datafusion_comet_jni_bridge::ShufflePartitionPusher; use std::io::{self, Cursor, Seek, SeekFrom, Write}; use std::sync::Arc; @@ -160,14 +161,19 @@ impl RssPartitionWriter { let (metadata_scratch, planning_scratch) = Self::estimated_ipc_metadata_scratch(batch); let codec_workspace = self.block_writer.rss_codec_workspace()?; let reservation_limit = self.pusher.max_reservation_size(); - if metadata_scratch + let minimum_reservation = metadata_scratch .saturating_add(codec_workspace) - .saturating_add(60) - > reservation_limit - { + .saturating_add(60); + if minimum_reservation > reservation_limit { // Neither schema descriptors nor codec workspace gets smaller when rows are split. - return Err(DataFusionError::Execution(format!( - "Remote shuffle frame schema and encoding workspace exceed the byte admission budget of {reservation_limit} bytes" + return Err(Self::size_limit_error(format!( + "Remote shuffle schema, encoding workspace, and minimum frame copies require at least \ + {minimum_reservation} bytes of reservation, but only {reservation_limit} \ + bytes per frame are available under spark.comet.shuffle.rss.maxInFlightBytes. \ + Increase that setting to allow the required reservation plus transport overhead. \ + The effective frame limit from spark.comet.shuffle.rss.maxFrameBytes and \ + spark.comet.shuffle.rss.maxInFlightBytes is {} bytes", + self.max_frame_size ))); } @@ -199,7 +205,7 @@ impl RssPartitionWriter { .and_then(|bytes| bytes.checked_add(compaction_scratch)) .and_then(|bytes| bytes.checked_add(codec_workspace)) .ok_or_else(|| { - DataFusionError::Execution( + Self::size_limit_error( "Remote shuffle encoding workspace exceeds the native integer limit" .to_string(), ) @@ -210,8 +216,16 @@ impl RssPartitionWriter { if batch.num_rows() > 1 { return self.push_split_batch(partition_id, batch, metrics); } - return Err(DataFusionError::Execution(format!( - "Remote shuffle frame for a single row and encoding workspace exceed the byte admission budget of {reservation_limit} bytes" + return Err(Self::size_limit_error(format!( + "Remote shuffle encoding workspace for a single row requires {ipc_scratch} bytes; \ + at least {} bytes of reservation are needed including a minimum frame, but only \ + {reservation_limit} bytes per frame are available under \ + spark.comet.shuffle.rss.maxInFlightBytes. Increase that setting to allow the \ + required reservation plus transport overhead. \ + The effective frame limit from spark.comet.shuffle.rss.maxFrameBytes and \ + spark.comet.shuffle.rss.maxInFlightBytes is {} bytes", + ipc_scratch.saturating_add(60), + self.max_frame_size ))); } @@ -231,7 +245,7 @@ impl RssPartitionWriter { .checked_mul(3) .and_then(|bytes| bytes.checked_add(ipc_scratch)) .ok_or_else(|| { - DataFusionError::Execution( + Self::size_limit_error( "Remote shuffle frame-copy reservation exceeds the native integer limit" .to_string(), ) @@ -277,6 +291,7 @@ impl RssPartitionWriter { &metrics.encode_time, ) { let exceeded = output.exceeded; + let minimum_frame_size = output.minimum_size; drop(output); drop(compacted); self.pusher.release_partition_data_reservation()?; @@ -288,9 +303,32 @@ impl RssPartitionWriter { continue; } if batch.num_rows() <= 1 { - return Err(DataFusionError::Execution(format!( - "Remote shuffle frame exceeds its configured maximum: a single row exceeds {} bytes", - admitted_frame_limit + let minimum_reservation = + ipc_scratch.saturating_add(minimum_frame_size.saturating_mul(3)); + let reason = if admitted_frame_limit < self.max_frame_size { + format!( + "only {admitted_frame_limit} encoded bytes fit alongside the \ + {ipc_scratch}-byte encoding workspace in the {reservation_limit}-byte \ + reservation available per frame under \ + spark.comet.shuffle.rss.maxInFlightBytes \ + (effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \ + and spark.comet.shuffle.rss.maxInFlightBytes: {} bytes)", + self.max_frame_size + ) + } else { + format!( + "the effective frame limit from spark.comet.shuffle.rss.maxFrameBytes \ + and spark.comet.shuffle.rss.maxInFlightBytes is {} bytes; encoding \ + workspace requires {ipc_scratch} bytes and the reservation available \ + per frame is {reservation_limit} bytes", + self.max_frame_size + ) + }; + return Err(Self::size_limit_error(format!( + "Remote shuffle cannot encode a single row: the frame needs at least \ + {minimum_frame_size} bytes, but {reason}. Increase the limiting setting; \ + the required reservation is at least {minimum_reservation} bytes plus \ + transport overhead" ))); } return self.push_split_batch(partition_id, batch, metrics); @@ -316,6 +354,10 @@ impl RssPartitionWriter { } } + fn size_limit_error(message: String) -> DataFusionError { + DataFusionError::External(Box::new(CometError::ShuffleSizeLimit(message))) + } + fn push_split_batch( &mut self, partition_id: i32, @@ -1092,6 +1134,8 @@ struct BoundedBuffer { inner: Cursor>, limit: usize, exceeded: bool, + // The first rejected write/seek establishes a lower bound without allocating more output. + minimum_size: usize, } impl BoundedBuffer { @@ -1113,6 +1157,7 @@ impl BoundedBuffer { inner: Cursor::new(bytes), limit, exceeded: false, + minimum_size: 0, }) } @@ -1128,6 +1173,11 @@ impl Write for BoundedBuffer { fn write(&mut self, bytes: &[u8]) -> io::Result { let end = self.inner.position().checked_add(bytes.len() as u64); if end.is_none_or(|end| end > self.limit as u64) { + if !self.exceeded { + self.minimum_size = end + .and_then(|end| usize::try_from(end).ok()) + .unwrap_or(usize::MAX); + } self.exceeded = true; return Err(Self::limit_error()); } @@ -1145,6 +1195,9 @@ impl Seek for BoundedBuffer { let next = self.inner.seek(position)?; if next > self.limit as u64 { self.inner.set_position(previous); + if !self.exceeded { + self.minimum_size = usize::try_from(next).unwrap_or(usize::MAX); + } self.exceeded = true; return Err(Self::limit_error()); } @@ -1186,6 +1239,7 @@ mod buffer_tests { assert_eq!(output.inner.get_ref().capacity(), 67_996); assert!(output.write(&[1]).is_err()); assert!(output.exceeded); + assert_eq!(output.minimum_size, 67_997); assert_eq!(output.inner.get_ref().capacity(), 67_996); } } diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java index 4240755412e..c634b28e856 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -45,7 +45,7 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES; private static final int MINIMUM_COMET_FRAME_BYTES = 2 * Long.BYTES; private static final int MAX_JVM_ARRAY_BYTES = Integer.MAX_VALUE - 8; - private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 256 * 1024 * 1024; + private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 512 * 1024 * 1024; private static final long RECONCILIATION_INTERVAL_MILLIS = 10; // This daemon never owns client state: reconciliations are cancelled and removed once the diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f4aeacc8478..309ac7d47a8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -659,14 +659,17 @@ object CometConf extends ShimCometConf { "an executor-side remote shuffle client. Admission includes native encoding " + "scratch and overlapping native, JNI, and remote shuffle frame copies. " + "A frame must fit its codec and Arrow workspace as well as its encoded bytes; " + - "too-small limits fail before encoding. Encrypted native RSS is not supported; " + + "ordinary uncompressed frames need approximately seven times their size plus " + + "schema and transport overhead. Compressed frames also reserve workspace for " + + "their uncompressed data. Admission is acquired before encoding. " + + "Encrypted native RSS is not supported; " + "use ordinary Spark shuffle when spark.io.encryption.enabled is true.") .bytesConf(ByteUnit.BYTE) .checkValue( value => value >= 76 && value <= Int.MaxValue, "Remote shuffle in-flight byte limit must fit three complete frame copies and a " + "Celeborn request header") - .createWithDefault(256L * 1024 * 1024) + .createWithDefault(512L * 1024 * 1024) val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.debug.enabled") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index 5b0516d0763..71fec591ee8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -24,17 +24,19 @@ import java.util.Locale import java.util.concurrent.ConcurrentHashMap import scala.collection.mutable +import scala.concurrent.ExecutionContext import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.internal.config.{DYN_ALLOCATION_ENABLED, DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, SHUFFLE_SERVICE_ENABLED} import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint} import org.apache.spark.scheduler.OutputCommitCoordinator import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} import org.apache.spark.util.RpcUtils import org.apache.comet.CometConf -import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory} +import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, ResolvedCelebornShufflePusher} import org.apache.comet.util.ClassLoaders /** @@ -51,7 +53,8 @@ class CometCelebornShuffleManager private[shuffle] ( backendFactory: (SparkConf, Boolean) => ShuffleManager, readerApi: CelebornRawPartitionReader.Api = CelebornRawPartitionReader.reflectedApi, planningSupportFactory: SparkConf => CelebornNativeShufflePlanningSupport = conf => - CometCelebornShuffleManager.nativeShufflePlanningSupport(conf)) + CometCelebornShuffleManager.nativeShufflePlanningSupport(conf), + localManagerFactory: SparkConf => ShuffleManager = conf => new CometShuffleManager(conf)) extends ShuffleManager { /** Constructor used by Spark on both the driver and executors. */ @@ -72,16 +75,35 @@ class CometCelebornShuffleManager private[shuffle] ( def nativeShuffleFallbackReason(numPartitions: Int): Option[String] = nativePlanningSupport.fallbackReason(numPartitions) + protected[shuffle] def materializationExecutionContext: ExecutionContext = + ExecutionContext.global + private val nativeShuffleClients = new ConcurrentHashMap[Int, ConcurrentHashMap[Int, AnyRef]]() + private val sizeLimitFallbacks = new ConcurrentHashMap[Int, () => Boolean]() + private val localShuffleIds = ConcurrentHashMap.newKeySet[Integer]() private val ownedNativeClients = new ConcurrentHashMap[AnyRef, java.lang.Boolean]() @volatile private var nativeGenerationCoordinator: CelebornShuffleGenerationCoordinator = _ @volatile private var nativeGenerationEndpoint: RpcEndpointRef = _ + @volatile private var localShuffleManagerInstance: ShuffleManager = _ + + private def localShuffleManager: ShuffleManager = synchronized { + if (localShuffleManagerInstance == null) { + localShuffleManagerInstance = Option(localManagerFactory(conf)).getOrElse { + throw new IllegalStateException("Local Comet shuffle manager factory returned null") + } + } + localShuffleManagerInstance + } override def registerShuffle[K, V, C]( shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { dependency match { + case native: CometShuffleDependency[_, _, _] + if native.shuffleType == CometNativeShuffle && native.useLocalShuffle => + localShuffleIds.add(shuffleId) + localShuffleManager.registerShuffle(shuffleId, dependency) case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle => val handle = backend.registerShuffle(shuffleId, dependency) if (!CometCelebornShuffleManager.isCelebornHandle(handle)) { @@ -107,6 +129,10 @@ class CometCelebornShuffleManager private[shuffle] ( mapId: Long, context: TaskContext, metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + if (isLocalNativeHandle(handle)) { + localShuffleIds.add(handle.shuffleId) + return localShuffleManager.getWriter(handle, mapId, context, metrics) + } nativeDependency(handle) match { case Some(dependency) => val earlyClaim = claimNativeShuffleAttempt(handle.shuffleId, context) @@ -115,11 +141,9 @@ class CometCelebornShuffleManager private[shuffle] ( throw CelebornShufflePusherFactory.commitDenied(context) } var preparedClaim = earlyClaim - val resolved = CelebornShufflePusherFactory.createFromHandle( - conf, + val resolved = createRemotePusher( handle, context, - client => ownedNativeClients.put(client, java.lang.Boolean.TRUE), (celebornShuffleId, numMappers) => preparedClaim = prepareNativeShuffleGeneration( handle.shuffleId, @@ -171,7 +195,15 @@ class CometCelebornShuffleManager private[shuffle] ( handle.shuffleId, resolved.celebornShuffleId, context, - preparedClaim)))) + preparedClaim), + onSizeLimitExceeded = failure => + requestLocalShuffle( + handle.shuffleId, + resolved.celebornShuffleId, + resolved.client, + context, + preparedClaim, + failure)))) case None => rejectCometHandle(handle) @@ -188,6 +220,17 @@ class CometCelebornShuffleManager private[shuffle] ( endPartition: Int, context: TaskContext, metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + if (isLocalNativeHandle(handle)) { + localShuffleIds.add(handle.shuffleId) + return localShuffleManager.getReader( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + } nativeDependency(handle) match { case Some(dependency) => if (startMapIndex > endMapIndex) { @@ -238,32 +281,60 @@ class CometCelebornShuffleManager private[shuffle] ( override def shuffleBlockResolver: ShuffleBlockResolver = backend.shuffleBlockResolver override def unregisterShuffle(shuffleId: Int): Boolean = { - Option(nativeShuffleClients.remove(shuffleId)).foreach { generations => - generations.forEach { (celebornShuffleId, client) => - CelebornShufflePusherFactory.cleanupShuffle(client, celebornShuffleId) - } - } - if (isDriver) { - Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId)) + sizeLimitFallbacks.remove(shuffleId) + if (localShuffleIds.remove(shuffleId)) { + return localShuffleManager.unregisterShuffle(shuffleId) } - backend.unregisterShuffle(shuffleId) + val generations = Option(nativeShuffleClients.remove(shuffleId)).toSeq + .flatMap(_.asScala.toSeq) + var removed = false + val cleanup = Seq[() => Unit](() => + Option(localShuffleManagerInstance).foreach(_.unregisterShuffle(shuffleId))) ++ + generations.map { case (celebornShuffleId, client) => + () => CelebornShufflePusherFactory.cleanupShuffle(client, celebornShuffleId) + } ++ Seq[() => Unit]( + () => { + if (isDriver) { + Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId)) + } + }, + () => removed = backend.unregisterShuffle(shuffleId)) + cleanupAll(cleanup) + removed } override def stop(): Unit = { - try backend.stop() - finally { - try { - if (isDriver) { - Option(nativeGenerationEndpoint).foreach { endpoint => - SparkEnv.get.rpcEnv.stop(endpoint) + // Spark's disk block manager owns executor shutdown cleanup. In particular, an external + // shuffle service must retain local fallback output after this executor has stopped. + val localCleanup = Option(localShuffleManagerInstance).toSeq.map { local => () => + local.stop() + } + sizeLimitFallbacks.clear() + cleanupAll( + localCleanup ++ Seq[() => Unit]( + () => backend.stop(), + () => { + if (isDriver) { + Option(nativeGenerationEndpoint).foreach { endpoint => + SparkEnv.get.rpcEnv.stop(endpoint) + } } - } - } finally { - ownedNativeClients.keySet().asScala.foreach(CelebornShufflePusherFactory.releaseClient) - ownedNativeClients.clear() - nativeShuffleClients.clear() + }) ++ ownedNativeClients.keySet().asScala.toSeq.map { client => () => + CelebornShufflePusherFactory.releaseClient(client) + } ++ Seq[() => Unit](() => ownedNativeClients.clear(), () => nativeShuffleClients.clear())) + } + + private def cleanupAll(actions: Seq[() => Unit]): Unit = { + var failure: Throwable = null + actions.foreach { action => + try action() + catch { + case cleanupFailure: Throwable => + if (failure == null) failure = cleanupFailure + else if (failure ne cleanupFailure) failure.addSuppressed(cleanupFailure) } } + if (failure != null) throw failure } private def initializeNativeGenerationCoordinator(): Unit = synchronized { @@ -273,7 +344,8 @@ class CometCelebornShuffleManager private[shuffle] ( } val coordinator = new CelebornShuffleGenerationCoordinator( env.outputCommitCoordinator, - CelebornShufflePusherFactory.shouldReportShuffleFetchFailure) + shouldReportShuffleFetchFailure, + shuffleId => Option(sizeLimitFallbacks.get(shuffleId)).exists(callback => callback())) val endpoint = env.rpcEnv.setupEndpoint( CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT, new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator)) @@ -310,6 +382,91 @@ class CometCelebornShuffleManager private[shuffle] ( taskContext.partitionId(), taskContext.attemptNumber())) + protected[shuffle] def createRemotePusher( + handle: ShuffleHandle, + context: TaskContext, + onGenerationResolved: (Int, Int) => Unit, + onGenerationInvalidated: (Int, Int) => Unit, + onInvalidationUnsafe: (Int, Int) => Boolean): ResolvedCelebornShufflePusher = + CelebornShufflePusherFactory.createFromHandle( + conf, + handle, + context, + client => ownedNativeClients.put(client, java.lang.Boolean.TRUE), + onGenerationResolved, + onGenerationInvalidated, + onInvalidationUnsafe) + + protected[shuffle] def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean = + CelebornShufflePusherFactory.shouldReportShuffleFetchFailure(taskAttemptId) + + /** Registers the driver-owned materialization that can replace an unpublished shuffle. */ + private[shuffle] def registerSizeLimitFallback( + shuffleId: Int, + callback: () => Boolean): Unit = { + require(isDriver, "Shuffle fallback belongs to the driver") + require( + sizeLimitFallbacks.putIfAbsent(shuffleId, callback) == null, + s"Shuffle $shuffleId already has a materialization") + } + + private[shuffle] def removeSizeLimitFallback(shuffleId: Int): Unit = { + sizeLimitFallbacks.remove(shuffleId) + } + + private def isLocalNativeHandle(handle: ShuffleHandle): Boolean = handle match { + case native: CometNativeShuffleHandle[_, _] => + native.dependency match { + case dependency: CometShuffleDependency[_, _, _] => dependency.useLocalShuffle + case _ => false + } + case _ => false + } + + private def requestLocalShuffle( + shuffleId: Int, + celebornShuffleId: Int, + client: AnyRef, + context: TaskContext, + claim: CelebornMapAttemptClaim, + failure: Throwable): Unit = { + val accepted = generationEndpoint.askSync[Boolean]( + RequestLocalCometShuffle( + ValidateCelebornMapAttempt( + shuffleId, + celebornShuffleId, + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber(), + claim.epoch), + context.taskAttemptId())) + if (!accepted) { + // The remote dependency may already have been published to readers. Preserve the typed + // failure instead of silently changing the identity of a shuffle they are consuming. + throw failure + } + + // The driver is materializing a fresh local dependency. Cleanup is best effort: neither a + // lost executor nor a delayed remote success can affect that dependency's map outputs. + try { + client.getClass + .getMethod( + "reportShuffleFetchFailure", + java.lang.Integer.TYPE, + java.lang.Integer.TYPE, + java.lang.Long.TYPE) + .invoke( + client, + Int.box(shuffleId), + Int.box(celebornShuffleId), + Long.box(context.taskAttemptId())) + } catch { + case NonFatal(cleanupFailure) => failure.addSuppressed(cleanupFailure) + } + throw failure + } + private def prepareNativeShuffleGeneration( shuffleId: Int, celebornShuffleId: Int, @@ -452,6 +609,16 @@ private[shuffle] object CometCelebornShuffleManager { return CelebornNativeShufflePlanningSupport( Some("Native Celeborn shuffle does not support spark.io.encryption.enabled=true")) } + if (conf.get(DYN_ALLOCATION_ENABLED) && !conf.get(SHUFFLE_SERVICE_ENABLED) && + !conf.get(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED)) { + // A reliable remote ShuffleDataIO plugin can enable dynamic allocation without retaining + // executor-local output. Size-limit recovery must also preserve its local fallback files. + return CelebornNativeShufflePlanningSupport( + Some( + "Native Celeborn shuffle with dynamic allocation requires " + + "spark.shuffle.service.enabled=true or " + + "spark.dynamicAllocation.shuffleTracking.enabled=true for local fallback")) + } try { val completionUnavailable = pushCompletionUnavailableReason() if (completionUnavailable.nonEmpty) { @@ -577,12 +744,18 @@ private[shuffle] final case class AbandonCelebornMapAttempt( taskAttemptId: Long) extends Serializable +private[shuffle] final case class RequestLocalCometShuffle( + validation: ValidateCelebornMapAttempt, + taskAttemptId: Long) + extends Serializable + /** * Keeps Comet map admission aligned with Spark task failures and Celeborn shuffle generations. */ private[shuffle] final class CelebornShuffleGenerationCoordinator( outputCommitCoordinator: OutputCommitCoordinator, - shouldReportShuffleFetchFailure: Long => Boolean = _ => true) { + shouldReportShuffleFetchFailure: Long => Boolean = _ => true, + requestFallback: Int => Boolean = _ => false) { private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration] private val invalidatedGenerations = mutable.HashSet.empty[Int] @@ -590,6 +763,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( private val claimOwners = mutable.HashMap.empty[(Int, Int, Int, Int), (Int, Long)] private val deniedAttempts = mutable.HashMap.empty[(Int, Int, Int, Int), mutable.HashSet[Int]] + // An abandoned remote shuffle never becomes local. Its replacement owns a new shuffle ID, + // so even a success already queued in Spark cannot overwrite or skip a replacement map. + private val abandonedShuffles = mutable.HashSet.empty[Int] private def currentEpoch(shuffleId: Int): Long = generationEpochs.getOrElse(shuffleId, 0L) @@ -664,6 +840,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( } def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized { + if (abandonedShuffles.contains(claim.shuffleId)) { + return CelebornMapAttemptClaim(false, currentEpoch(claim.shuffleId)) + } val previousGeneration = generations.get(claim.shuffleId) val stale = previousGeneration.exists { generation => generation.stageId == claim.stageId && @@ -724,6 +903,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized { require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive") + if (abandonedShuffles.contains(generation.shuffleId)) { + return false + } generations.get(generation.shuffleId) match { case Some(previous) @@ -867,7 +1049,23 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( abandon } + def requestLocalShuffle(request: RequestLocalCometShuffle): Boolean = synchronized { + val validation = request.validation + if (abandonedShuffles.contains(validation.shuffleId) || + !validateMapAttempt(validation) || + !shouldReportShuffleFetchFailure(request.taskAttemptId) || + !requestFallback(validation.shuffleId)) { + return false + } + + abandonedShuffles.add(validation.shuffleId) + invalidatedGenerations.add(validation.shuffleId) + invalidateOwners(validation.shuffleId) + true + } + def unregisterShuffle(shuffleId: Int): Unit = synchronized { + abandonedShuffles.remove(shuffleId) generations.remove(shuffleId) invalidatedGenerations.remove(shuffleId) generationEpochs.remove(shuffleId) @@ -894,5 +1092,7 @@ private[shuffle] final class CelebornShuffleGenerationEndpoint( context.reply(coordinator.invalidateGeneration(invalidation)) case abandoned: AbandonCelebornMapAttempt => context.reply(coordinator.abandonMapAttempt(abandoned)) + case request: RequestLocalCometShuffle => + context.reply(coordinator.requestLocalShuffle(request)) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala new file mode 100644 index 00000000000..f9f6770b407 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala @@ -0,0 +1,359 @@ +/* + * 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.sql.comet.execution.shuffle + +import java.util.Properties + +import scala.collection.mutable +import scala.collection.mutable.ArrayBuffer +import scala.concurrent.{CanAwait, ExecutionContext, Future, Promise} +import scala.concurrent.duration.Duration +import scala.util.{Failure, Success, Try} +import scala.util.control.NonFatal + +import org.apache.spark.{FutureAction, MapOutputStatistics, ShuffleDependency, SparkException} +import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart} +import org.apache.spark.sql.comet.shims.ShimCometShuffleMaterialization +import org.apache.spark.util.ThreadUtils + +/** + * Driver-owned materialization of one native Celeborn shuffle. Storage is chosen before any + * downstream RDD can depend on its output. A size-limit failure cancels the remote map-stage job + * and materializes a fresh local dependency with an independent scheduling RDD. Their separate + * RDD, shuffle and stage identities isolate remote failures and late task completions from the + * replacement. Once output is published, its destination is fixed. + */ +private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( + remoteDependency: CometShuffleDependency[K, V, C], + manager: CometCelebornShuffleManager) + extends FutureAction[MapOutputStatistics] + with ShimCometShuffleMaterialization { + + import CometCelebornShuffleMaterialization._ + + private val sparkContext = remoteDependency.rdd.context + private val executionContext = manager.materializationExecutionContext + private val localProperties = sparkContext.getLocalProperties.clone().asInstanceOf[Properties] + private val contextClassLoader = Thread.currentThread().getContextClassLoader + private val completion = Promise[MapOutputStatistics]() + private val actions = ArrayBuffer.empty[FutureAction[MapOutputStatistics]] + private var state: State = RunningRemote + private var activeAction: Option[FutureAction[MapOutputStatistics]] = None + private var remoteAction: Option[FutureAction[MapOutputStatistics]] = None + private var remoteCancellationIssued = false + private var selected: Option[CometShuffleDependency[K, V, C]] = None + + private val jobListener = new SparkListener { + override def onJobStart(event: SparkListenerJobStart): Unit = { + localJobStarted(event.jobId) + } + } + sparkContext.addSparkListener(jobListener) + try { + manager.registerSizeLimitFallback(remoteDependency.shuffleId, () => requestLocalShuffle()) + submit(remoteDependency, RunningRemote) + } catch { + case NonFatal(failure) => + removeCallbacks() + throw failure + } + + private def withCapturedProperties[T](body: => T): T = withCapturedSession { + val thread = Thread.currentThread() + val previousProperties = sparkContext.getLocalProperties + val previousClassLoader = thread.getContextClassLoader + sparkContext.setLocalProperties(localProperties.clone().asInstanceOf[Properties]) + thread.setContextClassLoader(contextClassLoader) + try body + finally { + thread.setContextClassLoader(previousClassLoader) + sparkContext.setLocalProperties(previousProperties) + } + } + + private def submit(dependency: CometShuffleDependency[K, V, C], expected: State): Unit = { + implicit val executor: ExecutionContext = executionContext + val ready = Future { + withCapturedProperties(inputMaterializations(dependency.rdd)) + }.flatMap(upstream => Future.sequence(upstream)) + ready.onComplete { + case Failure(failure) => + fail(expected, failure) + cancelActions(failure) + case Success(_) => + try + withCapturedProperties { + // Upstream storage decisions have completed without occupying preparation workers. + // Cache the graph before taking our lock: submitMapStage also traverses it eagerly. + prepareInput(dependency.rdd) + val action = synchronized { + if (state != expected) { + None + } else { + if (expected == RunningLocal) remoteFailure.foreach(failure => throw failure) + // Keep job admission and action assignment atomic with cancellation and size + // reports. + val submitted = sparkContext.submitMapStage(dependency) + activeAction = Some(submitted) + if (expected == RunningRemote) remoteAction = Some(submitted) + actions += submitted + Some(submitted) + } + } + action.foreach(_.onComplete(result => finish(dependency, expected, result))) + } + catch { + case NonFatal(failure) => + fail(expected, failure) + cancelActions(failure) + } + } + } + + private def inputMaterializations(input: RDD[_]): Seq[Future[MapOutputStatistics]] = { + val visited = mutable.Set.empty[Int] + val upstream = mutable.Set.empty[Future[MapOutputStatistics]] + val pending = mutable.Stack[RDD[_]](input) + while (pending.nonEmpty && synchronized { state != Finished && state != Cancelled }) { + val next = pending.pop() + if (visited.add(next.id)) { + next match { + // Calling this RDD's dependencies would wait under Spark's RDD lock. Compose its + // completion instead; its materialization already owns preparation of its ancestors. + case shuffled: CometShuffledBatchRDD if !shuffled.isCheckpointed => + forDependency(shuffled.dependency) match { + case Some(materialization) => upstream += materialization + case None => pending.push(shuffled.dependency.rdd) + } + case other => other.dependencies.foreach(dependency => pending.push(dependency.rdd)) + } + } + } + upstream.toSeq + } + + private def prepareInput(input: RDD[_]): Unit = { + val visited = mutable.Set.empty[Int] + val pending = mutable.Stack[RDD[_]](input) + while (pending.nonEmpty && synchronized { state != Finished && state != Cancelled }) { + val next = pending.pop() + if (visited.add(next.id)) { + next.partitions + next.dependencies.foreach(dependency => pending.push(dependency.rdd)) + } + } + } + + private def finish( + dependency: CometShuffleDependency[K, V, C], + expected: State, + result: Try[MapOutputStatistics]): Unit = { + val finished = synchronized { + if (state == RunningLocal && expected == RunningRemote && + !remoteCancellationIssued && result.failed.toOption.exists(isNonSizeLimitFailure)) { + // A group cancellation processed before the replacement job starts must not disappear + // just because fallback has already won the storage decision. + state = Finished + completion.tryComplete(result) + cancelActions(result.failed.get) + true + } else if (state != expected) { + false + } else { + // JobStart listeners are asynchronous. If local output finishes before that notification, + // it still proves the replacement was active and lets us retire the old job here. + val outcome = if (expected == RunningLocal && result.isSuccess) { + cancelRemoteAfterLocalStarted().map(Failure(_)).getOrElse(result) + } else { + result + } + val published = outcome.map { statistics => + dependency.outputMetrics.foreach(_.publish(sparkContext)) + statistics + } + published match { + case Success(_) => selected = Some(dependency) + case _ => + } + state = Finished + completion.tryComplete(published) + published.failed.toOption.foreach(cancelActions) + true + } + } + if (finished) removeCallbacks() + } + + private def fail(expected: State, failure: Throwable): Unit = + finish(remoteDependency, expected, Failure(failure)) + + private def cancelActions(failure: Throwable): Unit = synchronized { + actions.foreach { action => + try action.cancel() + catch { + case NonFatal(cancellationFailure) => + if (failure ne cancellationFailure) failure.addSuppressed(cancellationFailure) + } + } + } + + private def removeCallbacks(): Unit = { + try manager.removeSizeLimitFallback(remoteDependency.shuffleId) + finally sparkContext.removeSparkListener(jobListener) + } + + private def isNonSizeLimitFailure(failure: Throwable): Boolean = + !CometNativeShuffleWriter.isSizeLimitFailure(failure) + + private def remoteFailure: Option[Throwable] = + remoteAction.flatMap(_.value).flatMap(_.failed.toOption).filter(isNonSizeLimitFailure) + + // Called with this materialization's lock held, after Spark has made the local job active. + private def cancelRemoteAfterLocalStarted(): Option[Throwable] = { + if (remoteCancellationIssued) return None + val previousFailure = remoteFailure + if (previousFailure.nonEmpty) return previousFailure + remoteCancellationIssued = true + try { + remoteAction.foreach(_.cancel()) + None + } catch { + case NonFatal(failure) => Some(failure) + } + } + + private def localJobStarted(jobId: Int): Unit = synchronized { + if (state == RunningLocal && activeAction.exists(_.jobIds.contains(jobId))) { + cancelRemoteAfterLocalStarted().foreach { failure => + fail(RunningLocal, failure) + cancelActions(failure) + } + } + } + + private def requestLocalShuffle(): Boolean = synchronized { + if (state != RunningRemote || activeAction.exists(_.isCompleted)) { + false + } else { + state = RunningLocal + try { + // Register the replacement before cancelling the old job. Its upstream inputs were + // already materialized for the remote submission. + // The JobStart listener retires the remote job only after Spark has made the replacement + // active. Until then, any external cancellation of the remote job also cancels fallback. + val localDependency = + withCapturedProperties(remoteDependency.createLocalShuffleDependency()) + remoteFailure match { + case Some(failure) => + fail(RunningLocal, failure) + cancelActions(failure) + case None => submit(localDependency, RunningLocal) + } + true + } catch { + case NonFatal(failure) => + fail(RunningLocal, failure) + cancelActions(failure) + // The storage transition won even if creating the replacement failed. The coordinator + // must still fence the abandoned remote ID; the materialization retains this failure. + true + } + } + } + + def completedDependency: Option[CometShuffleDependency[K, V, C]] = synchronized { selected } + + def selectedDependency: CometShuffleDependency[K, V, C] = { + try ThreadUtils.awaitResult(this, Duration.Inf) + catch { + case interrupted: InterruptedException => + cancel() + Thread.currentThread().interrupt() + throw interrupted + } + synchronized { + selected.getOrElse(throw new IllegalStateException("Shuffle materialization has no output")) + } + } + + override def cancel(): Unit = cancel(None) + + // Spark 4 calls this overload; Spark 3's FutureAction only declares cancel(). + def cancel(reason: Option[String]): Unit = { + val cancelled = synchronized { + if (state == Finished || state == Cancelled) { + false + } else { + state = Cancelled + val failure = + new SparkException(reason.getOrElse("Comet shuffle materialization cancelled")) + cancelActions(failure) + completion.tryFailure(failure) + true + } + } + if (cancelled) removeCallbacks() + } + + override def isCancelled: Boolean = synchronized { state == Cancelled } + + override def jobIds: Seq[Int] = synchronized { actions.iterator.flatMap(_.jobIds).toVector } + + override def isCompleted: Boolean = completion.isCompleted + + override def value: Option[Try[MapOutputStatistics]] = completion.future.value + + override def ready(atMost: Duration)(implicit permit: CanAwait): this.type = { + completion.future.ready(atMost) + this + } + + override def result(atMost: Duration)(implicit permit: CanAwait): MapOutputStatistics = + completion.future.result(atMost) + + override def onComplete[U](f: Try[MapOutputStatistics] => U)(implicit + executor: ExecutionContext): Unit = completion.future.onComplete(f) + + override def transform[S](f: Try[MapOutputStatistics] => Try[S])(implicit + executor: ExecutionContext): Future[S] = completion.future.transform(f) + + override def transformWith[S](f: Try[MapOutputStatistics] => Future[S])(implicit + executor: ExecutionContext): Future[S] = completion.future.transformWith(f) +} + +private[shuffle] object CometCelebornShuffleMaterialization { + private sealed trait State + private case object RunningRemote extends State + private case object RunningLocal extends State + private case object Finished extends State + private case object Cancelled extends State + + def forDependency(dependency: ShuffleDependency[_, _, _]) + : Option[CometCelebornShuffleMaterialization[_, _, _]] = dependency match { + case comet: CometShuffleDependency[_, _, _] => comet.materialization + case _ => None + } + + def selectForRead(dependency: ShuffleDependency[Int, _, _]): ShuffleDependency[Int, _, _] = + forDependency(dependency) + .map(_.selectedDependency.asInstanceOf[ShuffleDependency[Int, _, _]]) + .getOrElse(dependency) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index 41d1d6cc6a4..bced55b34df 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -45,6 +45,20 @@ private[shuffle] class CometNativeShuffleInputRDD( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { + /** + * Give local fallback its own scheduling RDD over the original upstream inputs. Spark aborts + * jobs whose RDD ancestry contains the failed stage's RDD, so reusing this instance or wrapping + * it in a narrow dependency lets a late remote failure abort the local replacement as well. + */ + private[shuffle] def copyForLocalShuffle(): CometNativeShuffleInputRDD = + new CometNativeShuffleInputRDD( + context, + inputRDDs, + numPartitionsParam, + shuffleScanIndices, + spillMetricNode, + perPartitionByKey) + override protected def getPartitions: Array[Partition] = (0 until numPartitionsParam).map { i => // Resolve leaf-RDD partitions on the driver here (where their @transient fields are still diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index 89be100dbd3..21dd7686302 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructField import org.apache.spark.util.{ThreadUtils, Utils} -import org.apache.comet.{CometConf, CometExecIterator} +import org.apache.comet.{CometConf, CometExecIterator, CometShuffleSizeLimitException} import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde} import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator} import org.apache.comet.serde.operator.schema2Proto @@ -103,6 +103,9 @@ class CometNativeShuffleWriter[K, V]( catch { case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure) } + if (CometNativeShuffleWriter.isSizeLimitFailure(failure)) { + destination.onSizeLimitExceeded(failure) + } } throw failure } @@ -484,6 +487,16 @@ class CometNativeShuffleWriter[K, V]( } private[shuffle] object CometNativeShuffleWriter { + private[shuffle] def isSizeLimitFailure(failure: Throwable): Boolean = { + var cause = failure + val visited = new java.util.IdentityHashMap[Throwable, java.lang.Boolean]() + while (cause != null && visited.put(cause, java.lang.Boolean.TRUE) == null) { + if (cause.isInstanceOf[CometShuffleSizeLimitException]) return true + cause = cause.getCause + } + false + } + def drainAndClose(iterator: Iterator[_], close: () => Unit): Unit = { Utils.tryWithSafeFinally { while (iterator.hasNext) { @@ -503,7 +516,8 @@ private[shuffle] final case class CelebornNativeShuffleDestination( maxFrameBytes: Int, numPartitions: Int, commitAuthorized: Boolean = false, - commitValidator: () => Boolean = () => true) { + commitValidator: () => Boolean = () => true, + onSizeLimitExceeded: Throwable => Unit = _ => ()) { require(pusher != null, "The Celeborn shuffle partition pusher must not be null") require(maxFrameBytes >= 20, "The Celeborn shuffle frame limit must fit a complete frame") require( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 2a058430074..34f1407fc9b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.StructType @@ -68,7 +69,9 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( val shuffleWriteMetrics: Map[String, SQLMetric] = Map.empty, val numParts: Int = 0, val rangePartitionBounds: Option[Seq[InternalRow]] = None, - val nativeShuffleSpec: Option[NativeShuffleSpec] = None) + val nativeShuffleSpec: Option[NativeShuffleSpec] = None, + val useLocalShuffle: Boolean = false, + private[shuffle] val outputMetrics: Option[CometShuffleOutputMetrics] = None) extends ShuffleDependency[K, V, C]( _rdd, partitioner, @@ -76,7 +79,66 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( keyOrdering, aggregator, mapSideCombine, - shuffleWriterProcessor) {} + shuffleWriterProcessor) { + + @transient @volatile private var materializationInstance + : CometCelebornShuffleMaterialization[K, V, C] = _ + + // One driver-owned materialization selects storage before any reader depends on this shuffle. + // Local Comet shuffle and ordinary Spark shuffle retain their existing lazy scheduling. + @transient private[shuffle] lazy val materialization + : Option[CometCelebornShuffleMaterialization[K, V, C]] = { + if (shuffleType == CometNativeShuffle && !useLocalShuffle && rdd.getNumPartitions > 0) { + rdd.context.env.shuffleManager match { + case manager: CometCelebornShuffleManager => + val instance = new CometCelebornShuffleMaterialization(this, manager) + materializationInstance = instance + Some(instance) + case _ => None + } + } else { + None + } + } + + private[shuffle] def currentShuffleDependency: CometShuffleDependency[K, V, C] = + Option(materializationInstance).flatMap(_.completedDependency).getOrElse(this) + + private[shuffle] def createLocalShuffleDependency(): CometShuffleDependency[K, V, C] = { + // Spark aborts dependent jobs by input RDD identity, not just shuffle or stage identity. + // Use a sibling scheduling RDD so a failure in the abandoned remote stage cannot abort + // the replacement. The upstream inputs remain shared. + val localRDD = rdd + .asInstanceOf[CometNativeShuffleInputRDD] + .copyForLocalShuffle() + .asInstanceOf[RDD[_ <: Product2[K, V]]] + val localOutputMetrics = outputMetrics.map(_.newDestination(rdd.context)) + val localWriteMetrics = shuffleWriteMetrics ++ localOutputMetrics.toSeq.flatMap(_.metrics) + new CometShuffleDependency[K, V, C]( + localRDD, + partitioner, + serializer, + keyOrdering, + aggregator, + mapSideCombine, + if (localOutputMetrics.nonEmpty) { + ShuffleExchangeExec.createShuffleWriteProcessor(localWriteMetrics) + } else { + shuffleWriterProcessor + }, + shuffleType, + schema, + decodeTime, + outputPartitioning, + outputAttributes, + localWriteMetrics, + numParts, + rangePartitionBounds, + nativeShuffleSpec, + useLocalShuffle = true, + outputMetrics = localOutputMetrics) + } +} /** Indicates shuffle type */ sealed trait ShuffleType diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 46fe621e481..39cd5186a73 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -148,7 +148,10 @@ case class CometShuffleExchangeExec( if (inputRDD.getNumPartitions == 0) { Future.successful(null) } else { - sparkContext.submitMapStage(shuffleDependency) + CometCelebornShuffleMaterialization.forDependency(shuffleDependency) match { + case Some(materialization) => materialization + case None => sparkContext.submitMapStage(shuffleDependency) + } } } @@ -167,7 +170,13 @@ case class CometShuffleExchangeExec( } // TODO: add `override` keyword after dropping Spark-3.x supports - def shuffleId: Int = getShuffleId(shuffleDependency) + def shuffleId: Int = { + val current = shuffleDependency match { + case comet: CometShuffleDependency[Int @unchecked, _, _] => comet.currentShuffleDependency + case other => other + } + getShuffleId(current) + } /** * A [[ShuffleDependency]] that will partition rows of its child based on the partitioning @@ -882,19 +891,31 @@ object CometShuffleExchangeExec None) } + // The remote stage can finish some maps before an oversized row requires a complete local + // replacement. Keep output statistics separate from the public counters so neither those + // completed maps nor late remote updates inflate the selected shuffle's AQE statistics. + val outputMetrics = thinRDD.context.env.shuffleManager match { + case _: CometCelebornShuffleManager if numParts > 0 => + Some(CometShuffleOutputMetrics(thinRDD.context, metrics)) + case _ => None + } + val destinationMetrics = metrics ++ outputMetrics.toSeq.flatMap(_.metrics) + new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( thinRDD, serializer = serializer, - shuffleWriterProcessor = ShuffleExchangeExec.createShuffleWriteProcessor(metrics), + shuffleWriterProcessor = + ShuffleExchangeExec.createShuffleWriteProcessor(destinationMetrics), shuffleType = CometNativeShuffle, partitioner = partitioner, decodeTime = metrics("decode_time"), outputPartitioning = Some(outputPartitioning), outputAttributes = outputAttributes, - shuffleWriteMetrics = metrics, + shuffleWriteMetrics = destinationMetrics, numParts = numParts, rangePartitionBounds = rangePartitionBounds, - nativeShuffleSpec = Some(augmentedSpec)) + nativeShuffleSpec = Some(augmentedSpec), + outputMetrics = outputMetrics) } /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala new file mode 100644 index 00000000000..029220812d4 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala @@ -0,0 +1,69 @@ +/* + * 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.sql.comet.execution.shuffle + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleWriteMetricsReporter} + +/** + * Output statistics for one shuffle destination. Remote and replacement tasks update distinct + * accumulators; only the successful destination publishes output statistics to the exchange. + * Timing and operator metrics still account for the work done by both destinations. + */ +private[shuffle] final class CometShuffleOutputMetrics private ( + val metrics: Map[String, SQLMetric], + @transient private val publishedMetrics: Map[String, SQLMetric], + private val executionId: String) + extends Serializable { + + def newDestination(sc: SparkContext): CometShuffleOutputMetrics = + CometShuffleOutputMetrics.create(sc, publishedMetrics, executionId) + + def publish(sc: SparkContext): Unit = { + metrics.foreach { case (name, metric) => publishedMetrics(name).set(metric.value) } + if (executionId != null) { + SQLMetrics.postDriverMetricUpdates(sc, executionId, publishedMetrics.values.toSeq) + } + } +} + +private[shuffle] object CometShuffleOutputMetrics { + import SQLShuffleWriteMetricsReporter.{SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN} + + def apply(sc: SparkContext, metrics: Map[String, SQLMetric]): CometShuffleOutputMetrics = { + val outputNames = Set("dataSize", SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN) + create( + sc, + metrics.filter { case (name, _) => outputNames.contains(name) }, + sc.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)) + } + + private def create( + sc: SparkContext, + publishedMetrics: Map[String, SQLMetric], + executionId: String): CometShuffleOutputMetrics = { + val metrics = Map( + "dataSize" -> SQLMetrics.createSizeMetric(sc, "data size"), + SHUFFLE_BYTES_WRITTEN -> SQLMetrics.createSizeMetric(sc, "shuffle bytes written"), + SHUFFLE_RECORDS_WRITTEN -> SQLMetrics.createMetric(sc, "shuffle records written")) + new CometShuffleOutputMetrics(metrics, publishedMetrics, executionId) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index 2a598a7872c..17f84ed6626 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -51,7 +51,18 @@ class CometShuffledBatchRDD( SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY, SQLConf.get.fetchShuffleBlocksInBatch.toString) - override def getDependencies: Seq[Dependency[_]] = List(dependency) + // Start without waiting so constructing a join or union can start every independent input. + // Spark resolves dependencies on the submitting thread before queueing a downstream job; + // freeze the selected destination there, before Spark caches this RDD's dependency graph. + // Cache fixed partition metadata first: parents such as UnionRDD read it while discovering + // dependencies, and must not wait for the RDD lock held by a concurrent getDependencies call. + partitions + CometCelebornShuffleMaterialization.forDependency(dependency) + + override def getDependencies: Seq[Dependency[_]] = { + dependency = CometCelebornShuffleMaterialization.selectForRead(dependency) + List(dependency) + } override val partitioner: Option[Partitioner] = if (partitionSpecs.forall(_.isInstanceOf[CoalescedPartitionSpec])) { diff --git a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala new file mode 100644 index 00000000000..26ba626a665 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala @@ -0,0 +1,44 @@ +/* + * 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.sql.comet.shims + +import org.apache.spark.sql.SparkSession + +/** + * Restores the session context captured before shuffle materialization leaves the caller thread. + */ +trait ShimCometShuffleMaterialization { + private val activeSession = SparkSession.getActiveSession + + protected def withCapturedSession[T](body: => T): T = { + val previousSession = SparkSession.getActiveSession + activeSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + try body + finally { + previousSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } +} diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala new file mode 100644 index 00000000000..440212dcdfb --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala @@ -0,0 +1,47 @@ +/* + * 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.sql.comet.shims + +import org.apache.spark.JobArtifactSet +import org.apache.spark.sql.SparkSession + +/** + * Restores the session context captured before shuffle materialization leaves the caller thread. + */ +trait ShimCometShuffleMaterialization { + private val activeSession = SparkSession.getActiveSession + private val artifactState = JobArtifactSet.getCurrentJobArtifactState.orNull + + protected def withCapturedSession[T](body: => T): T = + JobArtifactSet.withActiveJobArtifactState(artifactState) { + val previousSession = SparkSession.getActiveSession + activeSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + try body + finally { + previousSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } +} diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala new file mode 100644 index 00000000000..81fa37f9f66 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala @@ -0,0 +1,50 @@ +/* + * 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.sql.comet.shims + +import org.apache.spark.JobArtifactSet +import org.apache.spark.sql.classic.SparkSession + +/** + * Restores the session context captured before shuffle materialization leaves the caller thread. + */ +trait ShimCometShuffleMaterialization { + private val activeSession = SparkSession.getActiveSession + // Spark 4 also associates artifacts with ordinary, non-Connect sessions. + private val artifactState = JobArtifactSet.getCurrentJobArtifactState + .orElse(activeSession.map(_.artifactManager.state)) + .orNull + + protected def withCapturedSession[T](body: => T): T = + JobArtifactSet.withActiveJobArtifactState(artifactState) { + val previousSession = SparkSession.getActiveSession + activeSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + try body + finally { + previousSession match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala index d6db22b2825..0265f304817 100644 --- a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala @@ -161,7 +161,7 @@ class CometConfSuite extends AnyFunSuite { val conf = new SQLConf assert(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get(conf) == 64L * 1024 * 1024) - assert(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.get(conf) == 256L * 1024 * 1024) + assert(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.get(conf) == 512L * 1024 * 1024) } test("remote shuffle frame limit rejects incomplete frames and oversized JVM requests") { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 10cdbc7932e..7012ae1ba53 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -43,7 +43,7 @@ import org.apache.spark.sql.functions.{col, count, sum} import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native} +import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass} import org.apache.comet.shuffle.ShufflePartitionPusher @@ -249,8 +249,11 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper iterator.hasNext Iterator.single((false, callbacks)) } catch { - case failure: Exception => - Iterator.single((failure.getMessage.contains("admission budget"), callbacks)) + case failure: CometShuffleSizeLimitException => + Iterator.single( + ( + failure.getMessage.contains("spark.comet.shuffle.rss.maxInFlightBytes"), + callbacks)) } } finally { iterator.close() diff --git a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala index 6104c320b93..d2f825764c5 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleHandle.scala @@ -27,4 +27,6 @@ class CelebornShuffleHandle[K, V, C]( shuffleId: Int, dependency: ShuffleDependency[K, V, C], val stageRerunEnabled: Boolean = true) - extends BaseShuffleHandle[K, V, C](shuffleId, dependency) + extends BaseShuffleHandle[K, V, C](shuffleId, dependency) { + val numMappers: Int = dependency.rdd.getNumPartitions +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala new file mode 100644 index 00000000000..60132c984fd --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala @@ -0,0 +1,338 @@ +/* + * 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.sql.comet.execution.shuffle + +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import scala.collection.mutable +import scala.concurrent.{Await, ExecutionContext, Future} +import scala.concurrent.duration.DurationInt +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{SparkConf, SparkEnv, TaskContext} +import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle +import org.apache.spark.sql.{CometTestBase, DataFrame, SparkSession} +import org.apache.spark.sql.comet.shims.ShimCometShuffleMaterialization +import org.apache.spark.util.ThreadUtils + +import org.apache.comet.CometConf + +class CometCelebornConcurrentMaterializationSuite extends CometTestBase { + + import testImplicits._ + + override protected val shuffleManager: String = + classOf[CometCelebornConcurrentMaterializationTestManager].getName + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SHUFFLE_MODE.key, "native") + .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false") + .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k") + .set("spark.sql.adaptive.enabled", "false") + .set("spark.shuffle.compress", "false") + .set("spark.io.encryption.enabled", "false") + .set("spark.celeborn.client.spark.stageRerun.enabled", "true") + + private def manager: CometCelebornConcurrentMaterializationTestManager = + SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornConcurrentMaterializationTestManager] + + private def branch(start: Long, oversized: Boolean): DataFrame = { + // Keep the large value dependent on the input, so Spark cannot replace the inner shuffle's + // payload with a small key and reconstruct a folded string above that exchange. + val payload = if (oversized) "concat(CAST(id AS STRING), repeat('x', 1048576))" else "'small'" + spark + .range(start, start + 1, 1, 1) + .selectExpr("id", s"$payload AS payload") + .repartition(2, $"id") + } + + test("materialization restores the captured session without changing the worker's session") { + class CapturedSession extends ShimCometShuffleMaterialization { + def current: Option[SparkSession] = withCapturedSession(SparkSession.getActiveSession) + } + val captured = spark.withActive { new CapturedSession } + val worker = Future { + val previous = SparkSession.getActiveSession + SparkSession.clearActiveSession() + try { + assert(captured.current.contains(spark)) + assert(SparkSession.getActiveSession.isEmpty) + } finally { + previous.foreach(SparkSession.setActiveSession) + } + }(ExecutionContext.global) + Await.result(worker, 20.seconds) + } + + test("non-AQE sibling shuffles start before either remote materialization finishes") { + val gate = new ConcurrentRemoteMapStarts + manager.remoteMapStarts = Some(gate) + val query = branch(0, oversized = false).union(branch(100, oversized = false)) + val plan = query.queryExecution.executedPlan + val exchanges = collect(plan) { case exchange: CometShuffleExchangeExec => exchange } + assert(exchanges.size == 2) + val localBefore = manager.localAttempts.size() + val construction = Future(exchanges.map(_.executeColumnar()))(ExecutionContext.global) + try { + // Both maps remain blocked. Reaching this point proves independent branches overlap, + // without relying on elapsed task timings or allowing one branch to finish first. + assert(gate.bothStarted.await(20, TimeUnit.SECONDS)) + Await.result(construction, 20.seconds) + assert(gate.shuffleCount == 2) + gate.release.countDown() + exchanges.foreach { exchange => + val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val materialization = original.materialization.get + val statistics = Await.result(materialization, 20.seconds) + assert(statistics.shuffleId == original.shuffleId) + assert(materialization.completedDependency.exists(_ eq original)) + } + assert(manager.localAttempts.size() == localBefore) + } finally { + gate.release.countDown() + Await.ready(construction, 20.seconds) + manager.remoteMapStarts = None + exchanges.foreach { exchange => + val dependency = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + dependency.materialization.foreach(_.cancel()) + manager.unregisterShuffle(dependency.shuffleId) + } + } + } + + test("cancelling while an upstream shuffle materializes prevents downstream job admission") { + val gate = new ConcurrentRemoteMapStarts + manager.remoteMapStarts = Some(gate) + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val query = branch(0, oversized = true) + .selectExpr("id + 10 AS id", "payload") + .repartition(3, $"id") + val exchanges = collect(query.queryExecution.executedPlan) { + case exchange: CometShuffleExchangeExec => exchange + } + assert(exchanges.size == 2) + val outer = exchanges.find(_.outputPartitioning.numPartitions == 3).get + val construction = Future(outer.executeColumnar())(ExecutionContext.global) + try { + assert(gate.firstStarted.await(20, TimeUnit.SECONDS)) + Await.result(construction, 20.seconds) + val dependency = outer.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val materialization = dependency.materialization.get + assert(!materialization.isCompleted) + assert(materialization.jobIds.isEmpty) + + // The outer worker is waiting for the blocked inner shuffle. It must not hold the state + // lock during that wait, and cancellation must fence its not-yet-submitted map-stage job. + val cancellation = Future(materialization.cancel(Some("Cancelled before job admission")))( + ExecutionContext.global) + Await.result(cancellation, 5.seconds) + assert(materialization.isCancelled) + assert( + materialization.value.exists(_.failed.get.getMessage == + "Cancelled before job admission")) + gate.release.countDown() + exchanges.filterNot(_ eq outer).foreach { inner => + val original = inner.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + Await.result(original.materialization.get, 20.seconds) + } + assert(materialization.jobIds.isEmpty) + } finally { + gate.release.countDown() + Await.ready(construction, 20.seconds) + manager.remoteMapStarts = None + exchanges.foreach { exchange => + exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[_, _, _]] + .materialization + .foreach(_.cancel()) + } + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + manager.localAttempts.asScala.drop(localBefore).foreach { attempt => + manager.unregisterShuffle(attempt.shuffleId) + } + } + } + + test("shared upstream materialization does not exhaust a bounded execution context") { + val workers = ThreadUtils.newDaemonFixedThreadPool(2, "comet-materialization-test") + val executionContext = ExecutionContext.fromExecutorService(workers) + val gate = new ConcurrentRemoteMapStarts + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val branchCount = 6 + val shared = branch(0, oversized = true) + val query = (1 to branchCount) + .map { index => + shared.selectExpr(s"id + $index AS id", "payload").repartition(3, $"id") + } + .reduce(_.union(_)) + val exchanges = collect(query.queryExecution.executedPlan) { + case exchange: CometShuffleExchangeExec => exchange + } + val materializations = + mutable.ArrayBuffer.empty[CometCelebornShuffleMaterialization[_, _, _]] + val completedJobs = new CompletedMaterializationJobs + spark.sparkContext.addSparkListener(completedJobs) + manager.testMaterializationExecutionContext = Some(executionContext) + manager.remoteMapStarts = Some(gate) + try { + // Exchange reuse gives every sibling the same unfinished upstream destination. There are + // more dependents than workers, so blocking even a small bounded pool would starve the + // upstream completion callback as well as unrelated work scheduled on that pool. + assert(exchanges.size == branchCount + 1) + exchanges.foreach { exchange => + exchange.executeColumnar() + materializations += exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[_, _, _]] + .materialization + .get + } + assert(gate.firstStarted.await(20, TimeUnit.SECONDS)) + assert(gate.shuffleCount == 1) + assert(materializations.forall(!_.isCompleted)) + Await.result(Future(())(executionContext), 5.seconds) + + gate.release.countDown() + materializations.foreach { materialization => + Await.result(materialization, 30.seconds) + assert(materialization.completedDependency.exists(_.useLocalShuffle)) + assert(materialization.jobIds.size == 2) + } + val actual = query.collect().map(row => (row.getLong(0), row.getString(1))).toSeq + val expected = (1 to branchCount).map(index => (index.toLong, "0" + "x" * 1048576)) + assert(actual.sortBy(_._1) == expected) + } finally { + gate.release.countDown() + materializations.foreach(_.cancel()) + completedJobs.awaitCompletion(materializations.flatMap(_.jobIds).toSeq) + Await.result(Future(())(executionContext), 20.seconds) + spark.sparkContext.removeSparkListener(completedJobs) + manager.remoteMapStarts = None + manager.testMaterializationExecutionContext = None + executionContext.shutdown() + assert(executionContext.awaitTermination(20, TimeUnit.SECONDS)) + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + manager.localAttempts.asScala.drop(localBefore).foreach { attempt => + manager.unregisterShuffle(attempt.shuffleId) + } + } + } + + for (nested <- Seq(false, true)) { + test(s"non-AQE fallback preserves concurrent branches and downstream reads: nested=$nested") { + val gate = new ConcurrentRemoteMapStarts + manager.remoteMapStarts = Some(gate) + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val left = if (nested) { + // The changed partition key keeps a second shuffle above the first one. Submitting its + // materialization must wait for the first destination without blocking the right branch. + branch(0, oversized = true) + .selectExpr("id + 10 AS id", "payload") + .repartition(3, $"id") + } else { + branch(0, oversized = true) + } + val query = left.union(branch(100, oversized = true)) + val exchanges = collect(query.queryExecution.executedPlan) { + case exchange: CometShuffleExchangeExec => exchange + } + assert(exchanges.size == (if (nested) 3 else 2)) + assert(exchanges.forall(_.output.exists(_.name == "payload"))) + val execution = Future(query.collect())(ExecutionContext.global) + try { + assert(gate.bothStarted.await(20, TimeUnit.SECONDS)) + assert(gate.shuffleCount == 2) + gate.release.countDown() + val actual = + Await.result(execution, 30.seconds).map(row => (row.getLong(0), row.getString(1))).toSeq + val leftId = if (nested) 10L else 0L + assert( + actual + .sortBy(_._1) == Seq((leftId, "0" + "x" * 1048576), (100L, "100" + "x" * 1048576))) + assert(manager.localAttempts.size() > localBefore) + exchanges.foreach { exchange => + val dependency = + exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + assert(dependency.materialization.get.completedDependency.exists(_.useLocalShuffle)) + } + } finally { + gate.release.countDown() + Await.ready(execution, 30.seconds) + manager.remoteMapStarts = None + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + manager.localAttempts.asScala.drop(localBefore).foreach { attempt => + manager.unregisterShuffle(attempt.shuffleId) + } + } + } + } +} + +/** Holds two distinct remote map stages so the tests can prove their execution overlaps. */ +private[shuffle] class ConcurrentRemoteMapStarts { + val firstStarted = new CountDownLatch(1) + val bothStarted = new CountDownLatch(2) + val release = new CountDownLatch(1) + private val shuffleIds = mutable.Set.empty[Int] + + def shuffleCount: Int = synchronized { shuffleIds.size } + + def mapStarted(shuffleId: Int): Unit = { + synchronized { + if (shuffleIds.add(shuffleId)) { + firstStarted.countDown() + bothStarted.countDown() + } + } + require(release.await(20, TimeUnit.SECONDS), "The test must release concurrent remote maps") + } +} + +class CometCelebornConcurrentMaterializationTestManager(conf: SparkConf, isDriver: Boolean) + extends CometCelebornFallbackTestShuffleManager(conf, isDriver) { + + @volatile private[shuffle] var remoteMapStarts: Option[ConcurrentRemoteMapStarts] = None + @volatile private[shuffle] var testMaterializationExecutionContext: Option[ExecutionContext] = + None + + override protected[shuffle] def materializationExecutionContext: ExecutionContext = + testMaterializationExecutionContext.getOrElse(super.materializationExecutionContext) + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]) { + remoteMapStarts.foreach(_.mapStarted(handle.shuffleId)) + } + super.getWriter(handle, mapId, context, metrics) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala new file mode 100644 index 00000000000..85829401d4d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala @@ -0,0 +1,139 @@ +/* + * 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.sql.comet.execution.shuffle + +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{SparkConf, SparkEnv, TaskContext} +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} +import org.apache.spark.shuffle.{FetchFailedException, ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter} +import org.apache.spark.sql.CometTestBase + +import org.apache.comet.CometConf + +class CometCelebornLocalFetchFailureSuite extends CometTestBase { + + import testImplicits._ + + override protected val shuffleManager: String = + classOf[CometCelebornLocalFetchFailureTestManager].getName + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SHUFFLE_MODE.key, "native") + .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false") + .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k") + .set("spark.shuffle.compress", "false") + .set("spark.io.encryption.enabled", "false") + .set("spark.celeborn.client.spark.stageRerun.enabled", "true") + .set("spark.stage.maxConsecutiveAttempts", "4") + + test("local fetch failure after fallback recovers a partially completed result stage") { + val manager = SparkEnv.get.shuffleManager + .asInstanceOf[CometCelebornLocalFetchFailureTestManager] + val listener = new SparkListener { + override def onTaskEnd(event: SparkListenerTaskEnd): Unit = { + if (event.taskType == "ResultTask" && event.taskInfo.successful) { + manager.completedResultTasks.incrementAndGet() + } + } + } + spark.sparkContext.addSparkListener(listener) + try { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + // Both deterministic input partitions exceed the remote frame limit, forcing fresh + // local materialization. Keys 1 and 2 reach different reducers, so the result partition + // that completes before the fetch failure contains an actual row. + val query = spark + .range(1, 3, 1, 2) + .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload") + .repartition(2, $"key") + + val actual = query.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + + assert(actual.sortBy(_._1) == Seq((1, "x" * (1024 * 1024)), (2, "x" * (1024 * 1024)))) + assert(manager.injectedFetchFailures.get() == 1) + assert(manager.completedResultsBeforeFailure.get() > 0) + assert(!manager.remoteAttempts.isEmpty) + val failedShuffleId = manager.failedLocalShuffleId.get() + assert(manager.remoteAttempts.asScala.forall(_._1.shuffleId != failedShuffleId)) + assert(manager.localAttempts.asScala.exists { attempt => + attempt.shuffleId == failedShuffleId && attempt.stageAttempt > 0 + }) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } +} + +/** Injects one local block loss only after Spark accepts another result partition. */ +class CometCelebornLocalFetchFailureTestManager(conf: SparkConf, isDriver: Boolean) + extends CometCelebornFallbackTestShuffleManager(conf, isDriver) { + + // This local-mode fixture keeps state in its SparkContext-owned shuffle manager. + val completedResultTasks = new AtomicInteger() + val completedResultsBeforeFailure = new AtomicInteger() + val injectedFetchFailures = new AtomicInteger() + val failedLocalShuffleId = new AtomicInteger(-1) + private val failNextLocalFetch = new AtomicBoolean(true) + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]] && + context.partitionId() == 1 && + failNextLocalFetch.compareAndSet(true, false)) { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (completedResultTasks.get() == 0 && System.nanoTime() < deadline) { + Thread.sleep(10) + } + val completed = completedResultTasks.get() + require(completed > 0, "Another result partition must finish before the local fetch fails") + completedResultsBeforeFailure.set(completed) + failedLocalShuffleId.set(handle.shuffleId) + injectedFetchFailures.incrementAndGet() + throw new FetchFailedException( + SparkEnv.get.blockManager.shuffleServerId, + handle.shuffleId, + -1L, + 0, + startPartition, + "A local shuffle block was lost after another result partition completed", + null) + } + super.getReader( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala index c3f291af857..a300bf8de9f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, IndexS import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.CometConf +import org.apache.comet.{CometConf, CometShuffleSizeLimitException} import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, RecordingCelebornPushClient} /** Exercises real Spark map tasks, native RSS planning, and the Celeborn map lifecycle. */ @@ -105,6 +105,22 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase { } } + test("default pushers and the configured factory can share executor byte admission") { + val client = new RecordingCelebornPushClient + val defaultPusher = new CelebornShufflePartitionPusher(client, 19, 3, 0, 12, 9) + val configuredPusher = CelebornShufflePusherFactory.create( + new SparkConf(false), + client, + 19, + 12, + 9, + TaskContext.empty()) + + assert(defaultPusher.maxReservationBytes() == configuredPusher.maxReservationBytes()) + defaultPusher.abort() + configuredPusher.abort() + } + test("configured frame limits and executor admission both constrain native RSS callbacks") { val context = TaskContext.empty() val configured = new SparkConf(false) @@ -272,6 +288,282 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase { } } + test("native frame limit failures abort before requesting a shuffle fallback") { + withNativeShuffleDependency() { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornPushClient + val pusher = CelebornShufflePusherFactory.create( + SparkEnv.get.conf.clone().set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "20"), + client, + 96, + numMappers, + dependency.partitioner.numPartitions, + context) + val restart = + new FetchFailedException( + null, + dependency.shuffleId, + -1L, + -1, + 0, + "restart locally", + null) + var fallbackCalls = 0 + var reportedFailure: Throwable = null + var cleanupBeforeFallback = false + val writer = CometCelebornNativeShuffleWriterSuite.newWriter( + dependency, + context, + pusher, + onSizeLimitExceeded = failure => { + fallbackCalls += 1 + reportedFailure = failure + cleanupBeforeFallback = client.cleanupCalls.get() == 1 + throw restart + }) + val actual = + try { + writer.write(inputs) + null + } catch { + case failure: FetchFailedException => failure + } + + ( + actual eq restart, + fallbackCalls, + reportedFailure.isInstanceOf[CometShuffleSizeLimitException], + cleanupBeforeFallback, + client.pushCount, + client.mapperEndCalls.get(), + writer.mapStatus == null, + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 1)) + assert(results.forall(_._3)) + assert(results.forall(_._4)) + assert(results.forall(_._5 == 0)) + assert(results.forall(_._6 == 0)) + assert(results.forall(_._7)) + assert(results.forall(_._8)) + } + } + + test("a size failure after a successful push discards the partial remote map") { + withNativeShuffleDependency() { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val expected = new CometShuffleSizeLimitException("single row exceeds RSS frame limit") + val client = new RecordingCelebornPushClient { + override def pushOrMergeData( + shuffleId: Int, + mapId: Int, + attemptId: Int, + partitionId: Int, + bytes: Array[Byte], + offset: Int, + length: Int, + numMappers: Int, + numPartitions: Int, + doPush: Boolean, + skipCompress: Boolean): Int = { + if (pushCount == 1) failure = expected + super.pushOrMergeData( + shuffleId, + mapId, + attemptId, + partitionId, + bytes, + offset, + length, + numMappers, + numPartitions, + doPush, + skipCompress) + } + } + val pusher = CelebornShufflePusherFactory.create( + SparkEnv.get.conf, + client, + 97, + numMappers, + dependency.partitioner.numPartitions, + context) + var fallbackCalls = 0 + var cleanupBeforeFallback = false + val writer = CometCelebornNativeShuffleWriterSuite.newWriter( + dependency, + context, + pusher, + onSizeLimitExceeded = failure => { + assert(failure eq expected) + fallbackCalls += 1 + cleanupBeforeFallback = client.cleanupCalls.get() == 1 + }) + val actual = + try { + writer.write(inputs) + null + } catch { + case failure: CometShuffleSizeLimitException => failure + } + + ( + actual eq expected, + fallbackCalls, + cleanupBeforeFallback, + client.pushCount, + client.mapperEndCalls.get(), + writer.mapStatus == null, + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 1)) + assert(results.forall(_._3)) + assert(results.forall(_._4 == 2)) + assert(results.forall(_._5 == 0)) + assert(results.forall(_._6)) + assert(results.forall(_._7)) + } + } + + test("transport failures do not request fallback based on their error message") { + withNativeShuffleDependency() { dependency => + val numMappers = dependency.rdd.getNumPartitions + val results = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val expected = new IOException("single row exceeds RSS frame limit") + val client = new RecordingCelebornPushClient + client.failure = expected + val pusher = CelebornShufflePusherFactory.create( + SparkEnv.get.conf, + client, + 98, + numMappers, + dependency.partitioner.numPartitions, + context) + var fallbackCalls = 0 + val writer = CometCelebornNativeShuffleWriterSuite.newWriter( + dependency, + context, + pusher, + onSizeLimitExceeded = _ => fallbackCalls += 1) + val actual = + try { + writer.write(inputs) + null + } catch { + case failure: IOException => failure + } + + ( + actual eq expected, + fallbackCalls, + client.cleanupCalls.get(), + writer.mapStatus == null, + writer.stop(success = false).isEmpty) + }) + + assert(results.nonEmpty) + assert(results.forall(_._1)) + assert(results.forall(_._2 == 0)) + assert(results.forall(_._3 == 1)) + assert(results.forall(_._4)) + assert(results.forall(_._5)) + } + } + + for (adaptive <- Seq(false, true)) { + test( + s"a row exceeding the RSS frame limit roundtrips through local Comet shuffle: AQE=$adaptive") { + // Exercise the reported 70 MiB row once; the AQE variant uses the same limit crossing at a + // smaller scale. Disabling compression makes the encoded size deterministic. + val payloadBytes = if (adaptive) 1024 * 1024 else 70 * 1024 * 1024 + val payload = "x" * payloadBytes + val rows = Seq((0, "small"), (1, payload)) + val runtimeConf = SparkEnv.get.conf + val previousCompression = runtimeConf.getOption("spark.shuffle.compress") + runtimeConf.set("spark.shuffle.compress", "false") + try { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + "spark.sql.adaptive.enabled" -> adaptive.toString) { + withTempPath { path => + rows + .toDF("key", "payload") + .coalesce(1) + .write + .option("parquet.enable.dictionary", "false") + .parquet(path.getCanonicalPath) + val shuffled = spark.read.parquet(path.getCanonicalPath).repartition(3, $"key") + val dependency = collect(shuffled.queryExecution.executedPlan) { + case exchange: CometShuffleExchangeExec => + exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]] + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val numMappers = dependency.rdd.getNumPartitions + val frameBytes = if (adaptive) "32k" else "64m" + val remoteAttempts = spark.sparkContext.runJob( + dependency.rdd, + (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { + val client = new RecordingCelebornPushClient + val pusher = CelebornShufflePusherFactory.create( + SparkEnv.get.conf + .clone() + .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, frameBytes), + client, + 99, + numMappers, + dependency.partitioner.numPartitions, + context) + var fallbackRequested = false + val writer = CometCelebornNativeShuffleWriterSuite.newWriter( + dependency, + context, + pusher, + onSizeLimitExceeded = _ => fallbackRequested = true) + try { + writer.write(inputs) + } catch { + case _: CometShuffleSizeLimitException => + } + val failedWithoutCommit = + fallbackRequested && client.cleanupCalls.get() == 1 && + client.mapperEndCalls.get() == 0 && writer.mapStatus == null + writer.stop(success = false) + failedWithoutCommit + }) + assert(remoteAttempts.exists(identity)) + + // The dependency is registered with the suite's real local Comet manager. Running + // it again exercises native file output, Spark map-status publication, and the Comet + // reader, including a frame larger than the remote limit and normal AQE stage reads. + val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + assert(actual.sortBy(_._1) == rows) + } + } + } finally { + previousCompression match { + case Some(value) => runtimeConf.set("spark.shuffle.compress", value) + case None => runtimeConf.remove("spark.shuffle.compress") + } + } + } + } + test("low-cardinality range shuffle plans use the actual reducer partition count") { withNativeShuffleDependency(rangePartitioning = true) { dependency => val actualPartitions = dependency.partitioner.numPartitions @@ -424,7 +716,9 @@ private[shuffle] object CometCelebornNativeShuffleWriterSuite { context: TaskContext, pusher: CelebornShufflePartitionPusher, commitAuthorized: Boolean = false, - commitValidator: () => Boolean = () => true): CometNativeShuffleWriter[Int, ColumnarBatch] = + commitValidator: () => Boolean = () => true, + onSizeLimitExceeded: Throwable => Unit = _ => ()) + : CometNativeShuffleWriter[Int, ColumnarBatch] = new CometNativeShuffleWriter[Int, ColumnarBatch]( dependency.nativeShuffleSpec.get, dependency.outputPartitioning.get, @@ -442,7 +736,8 @@ private[shuffle] object CometCelebornNativeShuffleWriterSuite { pusher.maxFrameBytes(), dependency.partitioner.numPartitions, commitAuthorized, - commitValidator))) + commitValidator, + onSizeLimitExceeded))) final class LocalFallbackShuffleManager extends ShuffleManager { var unregisteredShuffle: Option[Int] = None diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala new file mode 100644 index 00000000000..826ac7e86df --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala @@ -0,0 +1,808 @@ +/* + * 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.sql.comet.execution.shuffle + +import java.util.Properties +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.collection.mutable +import scala.concurrent.Await +import scala.concurrent.duration.DurationInt +import scala.jdk.CollectionConverters._ +import scala.language.existentials + +import org.apache.spark.{FutureAction, MapOutputStatistics, MapOutputTrackerMaster, ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.scheduler.{JobFailed, MapStatus, SparkListener, SparkListenerJobEnd, SparkListenerJobStart} +import org.apache.spark.shuffle.{ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle +import org.apache.spark.sql.CometTestBase + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.shuffle.{CelebornShufflePusherFactory, RecordingCelebornPushClient, ResolvedCelebornShufflePusher} + +/** + * Real Spark stage recovery and local Comet I/O; only the optional Celeborn client is replaced. + */ +class CometCelebornShuffleFallbackSuite extends CometTestBase { + + import testImplicits._ + + override protected val shuffleManager: String = + classOf[CometCelebornFallbackTestShuffleManager].getName + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SHUFFLE_MODE.key, "native") + .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false") + .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k") + .set("spark.shuffle.compress", "false") + .set("spark.io.encryption.enabled", "false") + .set("spark.celeborn.client.spark.stageRerun.enabled", "true") + .set("spark.stage.maxConsecutiveAttempts", "3") + + private def manager: CometCelebornFallbackTestShuffleManager = + SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornFallbackTestShuffleManager] + + for (adaptive <- Seq(false, true)) { + test(s"an oversized row materializes a fresh local shuffle and succeeds: AQE=$adaptive") { + withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { + val rows = Seq((0, "small"), (1, "x" * (1024 * 1024))) + withTempPath { path => + rows + .toDF("key", "payload") + .coalesce(1) + .write + .option("parquet.enable.dictionary", "false") + .parquet(path.getCanonicalPath) + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val readsBefore = manager.nativeReads.get() + val shuffled = spark.read.parquet(path.getCanonicalPath).repartition(3, $"key") + + val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + + assert(actual.sortBy(_._1) == rows) + val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq + assert(remote.size == 1, s"expected one failed remote map attempt, got $remote") + val local = manager.localAttempts.asScala.drop(localBefore).toSeq + assert(local.size == 1, s"expected a complete local map stage: $local") + assert(local.head.shuffleId != remote.head._1.shuffleId) + assert(local.head.stageId != remote.head._1.stageId) + assert(local.head.stageAttempt == 0) + assert(manager.nativeReads.get() > readsBefore) + val client = remote.head._2 + assert(client.cleanupCalls.get() == 1) + assert(client.mapperEndCalls.get() == 0) + assert(client.fetchFailureReports.get() == 1) + + // Each destination has its own shuffle lifecycle after the successful replacement. + assert(manager.unregisterShuffle(remote.head._1.shuffleId)) + assert(manager.unregisterShuffle(local.head.shuffleId)) + assert(client.shuffleCleanupCalls.get() == 1) + } + } + } + + test(s"fallback recomputes a previously completed remote map: AQE=$adaptive") { + withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val readsBefore = manager.nativeReads.get() + manager.remoteMapsBeforeFailure.set(0) + manager.waitForFirstRemoteMap = true + try { + // Range partition zero produces the small row; partition one produces the oversized + // row. Delay the second writer until Spark has registered the first remote map output. + val shuffled = spark + .range(0, 2, 1, 2) + .selectExpr( + "CAST(id AS INT) AS key", + "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload") + .repartition(3, $"key") + val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + + assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024)))) + assert(manager.remoteMapsBeforeFailure.get() == 1) + val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq + assert(remote.size == 2) + assert(remote.map(_._1.shuffleId).distinct.size == 1) + assert(remote.forall(_._1.stageAttempt == 0)) + assert(remote.count(_._2.mapperEndCalls.get() == 1) == 1) + assert(remote.map(_._2.fetchFailureReports.get()).sum == 1) + val local = manager.localAttempts.asScala.drop(localBefore).toSeq + assert(local.map(_.partitionId).sorted == Seq(0, 1)) + assert(local.map(_.shuffleId).distinct.size == 1) + assert(local.map(_.stageId).distinct.size == 1) + assert(local.forall(_.shuffleId != remote.head._1.shuffleId)) + assert(local.forall(_.stageId != remote.head._1.stageId)) + assert(local.forall(_.stageAttempt == 0)) + assert(manager.nativeReads.get() > readsBefore) + assert(manager.unregisterShuffle(remote.head._1.shuffleId)) + assert(manager.unregisterShuffle(local.head.shuffleId)) + } finally { + manager.waitForFirstRemoteMap = false + } + } + } + + test(s"delayed remote map completion cannot replace fresh local output: AQE=$adaptive") { + withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val delayed = new DelayedRemoteMapCompletion + manager.delayedRemoteCompletion = Some(delayed) + try { + val shuffled = spark + .range(0, 2, 1, 2) + .selectExpr( + "CAST(id AS INT) AS key", + "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload") + .repartition(3, $"key") + + val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + + assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024)))) + assert(delayed.remoteReturned.await(20, TimeUnit.SECONDS)) + val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq + val local = manager.localAttempts.asScala.drop(localBefore).toSeq + assert(remote.size == 2) + assert(remote.count(_._2.mapperEndCalls.get() == 1) == 1) + assert(local.map(_.partitionId).sorted == Seq(0, 1)) + assert(local.map(_.shuffleId).distinct.size == 1) + assert(local.map(_.stageId).distinct.size == 1) + assert(local.forall(_.shuffleId != remote.head._1.shuffleId)) + assert(local.forall(_.stageId != remote.head._1.stageId)) + assert(local.forall(_.stageAttempt == 0)) + assert(manager.unregisterShuffle(remote.head._1.shuffleId)) + assert(manager.unregisterShuffle(local.head.shuffleId)) + } finally { + // Unblock an old writer even if the query fails before the local stage starts. + delayed.localStageStarted.countDown() + manager.delayedRemoteCompletion = None + } + } + } + } + + test("successful remote materialization publishes the original shuffle identity") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val readsBefore = manager.nativeReads.get() + val query = spark.range(0, 2, 1, 2).repartition(2, $"id") + val exchange = collect(query.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val action = exchange.mapOutputStatisticsFuture + .asInstanceOf[FutureAction[MapOutputStatistics]] + try { + val statistics = Await.result(action, 20.seconds) + + assert(statistics.shuffleId == original.shuffleId) + assert(statistics.bytesByPartitionId.sum > 0L) + assert(original.currentShuffleDependency eq original) + // Spark 4 exposes this ID to AQE; the Spark 3 shim supplies an unused placeholder. + if (isSpark40Plus) assert(exchange.shuffleId == original.shuffleId) + assert(action.jobIds.size == 1) + val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq + assert(remote.map(_._1.partitionId).sorted == Seq(0, 1)) + assert(remote.forall(_._1.shuffleId == original.shuffleId)) + assert(remote.forall(_._2.mapperEndCalls.get() == 1)) + assert(manager.localAttempts.size() == localBefore) + assert(manager.nativeReads.get() == readsBefore) + } finally { + action.cancel() + manager.unregisterShuffle(original.shuffleId) + } + } + } + + test("a remote stage failure cannot abort an active local replacement") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val sc = spark.sparkContext + val previousProperties = sc.getLocalProperties.clone().asInstanceOf[Properties] + val groupId = "comet-remote-failure-isolation" + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val paused = new PausedShuffleMapCompletion(local = true) + val jobStart = new PausedReplacementJobStart(groupId) + val completedJobs = new CompletedMaterializationJobs + manager.pausedMapCompletion = Some(paused) + manager.remoteFailureAfterLocalMap = Some(paused) + // Delay the materialization's JobStart listener, so the remote size failure reaches + // Spark before Comet can cancel that job. Observe job failures on an independent queue. + sc.addSparkListener(jobStart) + sc.listenerBus.addToQueue(completedJobs, groupId) + var active: Option[FutureAction[MapOutputStatistics]] = None + try { + sc.setJobGroup(groupId, "Isolate the local replacement from remote failure", true) + val query = spark + .range(0, 1, 1, 1) + .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload") + .repartition(2, $"key") + val exchange = collect(query.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.head + val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val action = exchange.mapOutputStatisticsFuture + .asInstanceOf[FutureAction[MapOutputStatistics]] + active = Some(action) + assert(jobStart.entered.await(20, TimeUnit.SECONDS)) + assert(paused.stopped.await(20, TimeUnit.SECONDS)) + assert(action.jobIds.size == 2) + val remoteJobId = action.jobIds.head + completedJobs.awaitCompletion(Seq(remoteJobId)) + assert(completedJobs.hasFailed(remoteJobId)) + + jobStart.release.countDown() + paused.release.countDown() + val statistics = Await.result(action, 20.seconds) + + val selected = original.currentShuffleDependency + assert(selected.useLocalShuffle) + assert(selected.rdd ne original.rdd) + assert(statistics.shuffleId == selected.shuffleId) + completedJobs.awaitCompletion(action.jobIds) + assert(!completedJobs.hasFailed(action.jobIds.last)) + assert(query.collect().length == 1) + } finally { + jobStart.release.countDown() + paused.release.countDown() + active.foreach(_.cancel()) + sc.setLocalProperties(previousProperties) + manager.pausedMapCompletion = None + manager.remoteFailureAfterLocalMap = None + sc.removeSparkListener(jobStart) + sc.removeSparkListener(completedJobs) + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + manager.localAttempts.asScala.drop(localBefore).foreach { attempt => + manager.unregisterShuffle(attempt.shuffleId) + } + } + } + } + + test("job-group cancellation while registering fallback does not submit a replacement job") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val sc = spark.sparkContext + val previousProperties = sc.getLocalProperties.clone().asInstanceOf[Properties] + val groupId = "comet-materialization-cancel-during-local-registration" + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val registration = new PausedLocalShuffleRegistration + val completedJobs = new CompletedMaterializationJobs + manager.pausedLocalRegistration = Some(registration) + // The materialization holds its lock during registration, so its own listener can block. + // An independent queue lets the test observe Spark's cancellation before releasing it. + sc.listenerBus.addToQueue(completedJobs, groupId) + var active: Option[FutureAction[MapOutputStatistics]] = None + try { + sc.setJobGroup(groupId, "Cancel while local shuffle registration is paused", true) + val query = spark + .range(0, 1, 1, 1) + .selectExpr("CAST(id AS INT) AS key", "repeat('x', 1048576) AS payload") + .repartition(2, $"key") + val exchange = collect(query.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val action = exchange.mapOutputStatisticsFuture + .asInstanceOf[FutureAction[MapOutputStatistics]] + active = Some(action) + assert(registration.entered.await(20, TimeUnit.SECONDS)) + val remoteJobId = completedJobs.awaitJobInGroup(groupId) + + // Fallback has won, but submitMapStage(local) has not happened. Spark's ordinary + // cancelJobGroup only cancels current jobs, so a later replacement must be fenced. + sc.cancelJobGroup(groupId) + completedJobs.awaitCompletion(Seq(remoteJobId)) + assert(completedJobs.hasFailed(remoteJobId)) + registration.release.countDown() + assert(registration.returned.await(20, TimeUnit.SECONDS)) + Await.ready(action, 20.seconds) + + assert(action.value.exists(_.isFailure)) + assert(original.materialization.get.completedDependency.isEmpty) + assert(action.jobIds == Seq(remoteJobId)) + assert(manager.localAttempts.size() == localBefore) + assert(manager.remoteAttempts.size() == remoteBefore + 1) + } finally { + registration.release.countDown() + active.foreach(_.cancel()) + sc.setLocalProperties(previousProperties) + manager.pausedLocalRegistration = None + sc.removeSparkListener(completedJobs) + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + if (registration.shuffleId.get() >= 0) { + manager.unregisterShuffle(registration.shuffleId.get()) + } + } + } + } + + for (cancelLocal <- Seq(false, true)) { + test(s"cancelling materialization ends every submitted job: local=$cancelLocal") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val remoteBefore = manager.remoteAttempts.size() + val localBefore = manager.localAttempts.size() + val paused = new PausedShuffleMapCompletion(cancelLocal) + val completedJobs = new CompletedMaterializationJobs + manager.pausedMapCompletion = Some(paused) + spark.sparkContext.addSparkListener(completedJobs) + var active: Option[FutureAction[MapOutputStatistics]] = None + try { + val payload = if (cancelLocal) "repeat('x', 1048576)" else "'small'" + val query = spark + .range(0, 1, 1, 1) + .selectExpr("CAST(id AS INT) AS key", s"$payload AS payload") + .repartition(2, $"key") + val exchange = collect(query.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val original = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val action = exchange.mapOutputStatisticsFuture + .asInstanceOf[FutureAction[MapOutputStatistics]] + active = Some(action) + assert(paused.stopped.await(20, TimeUnit.SECONDS)) + + // stop(true) has already produced a real MapStatus, but Spark has not received it. + // Cancelling here must prevent that delayed completion from publishing a dependency. + if (cancelLocal) { + original.materialization.get.cancel(Some("Cancelled during local materialization")) + } else { + action.cancel() + } + paused.release.countDown() + assert(paused.returned.await(20, TimeUnit.SECONDS)) + Await.ready(action, 20.seconds) + completedJobs.awaitCompletion(action.jobIds) + + assert(action.isCancelled) + assert(action.value.exists(_.isFailure)) + if (cancelLocal) { + assert( + action.value.get.failed.get.getMessage == "Cancelled during local materialization") + } + assert(original.materialization.get.completedDependency.isEmpty) + assert(action.jobIds.size == (if (cancelLocal) 2 else 1)) + val local = manager.localAttempts.asScala.drop(localBefore).toSeq + assert(local.size == (if (cancelLocal) 1 else 0)) + assert(local.forall(_.shuffleId != original.shuffleId)) + val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq + assert(remote.size == 1) + } finally { + paused.release.countDown() + active.foreach(_.cancel()) + manager.pausedMapCompletion = None + spark.sparkContext.removeSparkListener(completedJobs) + manager.remoteAttempts.asScala.drop(remoteBefore).foreach { case (attempt, _) => + manager.unregisterShuffle(attempt.shuffleId) + } + manager.localAttempts.asScala.drop(localBefore).foreach { attempt => + manager.unregisterShuffle(attempt.shuffleId) + } + } + } + } + } + +} + +private[shuffle] case class CometShuffleFallbackAttempt( + shuffleId: Int, + stageId: Int, + stageAttempt: Int, + partitionId: Int) + +/** Synchronizes one old remote map's validated completion with the new local stage. */ +private[shuffle] class DelayedRemoteMapCompletion { + val remoteCommitted = new CountDownLatch(1) + val localStageStarted = new CountDownLatch(1) + val remoteReturned = new CountDownLatch(1) + + def awaitRemoteCommit(): Unit = + require( + remoteCommitted.await(20, TimeUnit.SECONDS), + "The small remote map must commit before the oversized map starts") + + def awaitLocalStage(): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + var interrupted = false + var started = false + try { + while (!started && System.nanoTime() < deadline) { + try { + started = localStageStarted.await( + math.max(1L, deadline - System.nanoTime()), + TimeUnit.NANOSECONDS) + } catch { + case _: InterruptedException => interrupted = true + } + } + require(started, "A fresh local stage must start before the old map result is returned") + } finally { + if (interrupted) { + Thread.currentThread().interrupt() + } + } + } +} + +/** Holds one already computed map result so cancellation races with a concrete completion. */ +private[shuffle] class PausedShuffleMapCompletion(val local: Boolean) { + val stopped = new CountDownLatch(1) + val release = new CountDownLatch(1) + val returned = new CountDownLatch(1) + + def awaitRelease(): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + var interrupted = false + var released = false + try { + while (!released && System.nanoTime() < deadline) { + try { + released = + release.await(math.max(1L, deadline - System.nanoTime()), TimeUnit.NANOSECONDS) + } catch { + case _: InterruptedException => interrupted = true + } + } + require(released, "The test must release the completed map") + } finally { + if (interrupted) Thread.currentThread().interrupt() + } + } +} + +/** Holds listener delivery while Spark processes the old remote stage's failure. */ +private[shuffle] class PausedReplacementJobStart(groupId: String) extends SparkListener { + val entered = new CountDownLatch(1) + val release = new CountDownLatch(1) + private var startedJobs = 0 + + override def onJobStart(event: SparkListenerJobStart): Unit = { + if (Option(event.properties).exists(_.getProperty("spark.jobGroup.id") == groupId)) { + startedJobs += 1 + if (startedJobs == 2) { + entered.countDown() + require(release.await(20, TimeUnit.SECONDS), "The test must release local JobStart") + } + } + } +} + +/** Holds fallback registration before the materialization can submit its local map stage. */ +private[shuffle] class PausedLocalShuffleRegistration { + val shuffleId = new AtomicInteger(-1) + val entered = new CountDownLatch(1) + val release = new CountDownLatch(1) + val returned = new CountDownLatch(1) +} + +/** Waits for scheduler-confirmed job completion, including jobs cancelled during fallback. */ +private[shuffle] class CompletedMaterializationJobs extends SparkListener { + private val completed = mutable.Set.empty[Int] + private val failed = mutable.Set.empty[Int] + private val jobsByGroup = mutable.Map.empty[String, Int] + + override def onJobStart(event: SparkListenerJobStart): Unit = synchronized { + Option(event.properties) + .flatMap(properties => Option(properties.getProperty("spark.jobGroup.id"))) + .foreach(groupId => jobsByGroup(groupId) = event.jobId) + notifyAll() + } + + override def onJobEnd(event: SparkListenerJobEnd): Unit = synchronized { + completed += event.jobId + if (event.jobResult.isInstanceOf[JobFailed]) failed += event.jobId + notifyAll() + } + + def hasFailed(jobId: Int): Boolean = synchronized { failed.contains(jobId) } + + def awaitJobInGroup(groupId: String): Int = synchronized { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (!jobsByGroup.contains(groupId) && System.nanoTime() < deadline) { + wait(math.max(1L, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime()))) + } + require(jobsByGroup.contains(groupId), s"No materialization job started in group $groupId") + jobsByGroup(groupId) + } + + def awaitCompletion(jobIds: Seq[Int]): Unit = synchronized { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (!jobIds.forall(completed.contains) && System.nanoTime() < deadline) { + wait(math.max(1L, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime()))) + } + require(jobIds.forall(completed.contains), s"Materialization jobs still running: $jobIds") + } +} + +class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean) + extends CometCelebornShuffleManager( + conf, + isDriver, + (configuration, _) => new CometCelebornFallbackTestBackend(configuration), + planningSupportFactory = _ => CelebornNativeShufflePlanningSupport()) { + + private[shuffle] val remoteAttempts = + new ConcurrentLinkedQueue[(CometShuffleFallbackAttempt, CometCelebornFallbackTestClient)]() + private[shuffle] val localAttempts = new ConcurrentLinkedQueue[CometShuffleFallbackAttempt]() + val nativeReads = new AtomicInteger() + val remoteMapsBeforeFailure = new AtomicInteger() + @volatile var waitForFirstRemoteMap = false + @volatile private[shuffle] var delayedRemoteCompletion: Option[DelayedRemoteMapCompletion] = + None + @volatile private[shuffle] var pausedMapCompletion: Option[PausedShuffleMapCompletion] = None + @volatile private[shuffle] var remoteFailureAfterLocalMap: Option[PausedShuffleMapCompletion] = + None + @volatile private[shuffle] var pausedLocalRegistration: Option[PausedLocalShuffleRegistration] = + None + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { + val registration = dependency match { + case native: CometShuffleDependency[_, _, _] if native.useLocalShuffle => + pausedLocalRegistration + case _ => None + } + registration.foreach { paused => + paused.shuffleId.set(shuffleId) + paused.entered.countDown() + require( + paused.release.await(20, TimeUnit.SECONDS), + "The test must release local shuffle registration") + } + try super.registerShuffle(shuffleId, dependency) + finally registration.foreach(_.returned.countDown()) + } + + override protected[shuffle] def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean = + true + + override protected[shuffle] def createRemotePusher( + handle: ShuffleHandle, + context: TaskContext, + onGenerationResolved: (Int, Int) => Unit, + onGenerationInvalidated: (Int, Int) => Unit, + onInvalidationUnsafe: (Int, Int) => Boolean): ResolvedCelebornShufflePusher = { + val remoteHandle = handle.asInstanceOf[CelebornShuffleHandle[_, _, _]] + val dependency = remoteHandle.dependency + val client = new CometCelebornFallbackTestClient + val generation = handle.shuffleId + 1000 + val numMappers = remoteHandle.numMappers + onGenerationResolved(generation, numMappers) + remoteAttempts.add((recordAttempt(handle, context), client)) + ResolvedCelebornShufflePusher( + CelebornShufflePusherFactory.create( + conf, + client, + generation, + numMappers, + dependency.partitioner.numPartitions, + context), + client, + generation) + } + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + if (waitForFirstRemoteMap && handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] && + context.stageAttemptNumber() == 0 && context.partitionId() == 1) { + // This fixture runs Spark locally, so the executor can observe the driver's tracker. + // Waiting for a recorded MapStatus proves the retry discards already published output. + val tracker = SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (tracker.getNumAvailableOutputs(handle.shuffleId) == 0 && + System.nanoTime() < deadline) { + Thread.sleep(10) + } + val completed = tracker.getNumAvailableOutputs(handle.shuffleId) + require( + completed == 1, + s"Expected one completed remote map before size failure: $completed") + remoteMapsBeforeFailure.set(completed) + } + val delayed = delayedRemoteCompletion + if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] && context.partitionId() == 1) { + delayed.foreach(_.awaitRemoteCommit()) + } + val underlying = super.getWriter[K, V](handle, mapId, context, metrics) + val failureGate = remoteFailureAfterLocalMap.filter { _ => + handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] + } + val writer = if (failureGate.nonEmpty) { + new ShuffleWriter[K, V] { + override def write(records: Iterator[Product2[K, V]]): Unit = { + try underlying.write(records) + catch { + case failure if CometNativeShuffleWriter.isSizeLimitFailure(failure) => + // The replacement must be active when Spark processes the abandoned stage's + // exception; otherwise the shared-RDD bug depends on scheduler timing. + require(failureGate.get.stopped.await(20, TimeUnit.SECONDS)) + throw failure + } + } + + override def getPartitionLengths(): Array[Long] = underlying.getPartitionLengths() + + override def stop(success: Boolean): Option[MapStatus] = underlying.stop(success) + } + } else { + underlying + } + if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) { + localAttempts.add(recordAttempt(handle, context)) + delayed.foreach(_.localStageStarted.countDown()) + } + val paused = pausedMapCompletion.filter { completion => + completion.local == handle.isInstanceOf[CometNativeShuffleHandle[_, _]] + } + if (paused.nonEmpty) { + new ShuffleWriter[K, V] { + override def write(records: Iterator[Product2[K, V]]): Unit = writer.write(records) + + override def getPartitionLengths(): Array[Long] = writer.getPartitionLengths() + + override def stop(success: Boolean): Option[MapStatus] = { + val result = writer.stop(success) + if (success && result.nonEmpty) { + paused.get.stopped.countDown() + paused.get.awaitRelease() + paused.get.returned.countDown() + } + result + } + } + } else if (handle + .isInstanceOf[CelebornShuffleHandle[_, _, _]] && context.partitionId() == 0 && + delayed.nonEmpty) { + new ShuffleWriter[K, V] { + override def write(records: Iterator[Product2[K, V]]): Unit = writer.write(records) + + override def getPartitionLengths(): Array[Long] = writer.getPartitionLengths() + + override def stop(success: Boolean): Option[MapStatus] = { + val result = writer.stop(success) + if (success && result.nonEmpty) { + delayed.get.remoteCommitted.countDown() + // Keep the already validated result until local writers have started. Cancellation + // may interrupt this worker, but must not let this old result replace local output. + delayed.get.awaitLocalStage() + delayed.get.remoteReturned.countDown() + } + result + } + } + } else { + writer + } + } + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + val reader = super.getReader[K, C]( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + if (handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) { + nativeReads.incrementAndGet() + } + reader + } + + private def recordAttempt( + handle: ShuffleHandle, + context: TaskContext): CometShuffleFallbackAttempt = + CometShuffleFallbackAttempt( + handle.shuffleId, + context.stageId(), + context.stageAttemptNumber(), + context.partitionId()) +} + +class CometCelebornFallbackTestClient extends RecordingCelebornPushClient { + val fetchFailureReports = new AtomicInteger() + val shuffleCleanupCalls = new AtomicInteger() + + def reportShuffleFetchFailure( + shuffleId: Int, + celebornShuffleId: Int, + taskAttemptId: Long): Boolean = { + fetchFailureReports.incrementAndGet() + true + } + + def cleanupShuffle(shuffleId: Int): Unit = { + shuffleCleanupCalls.incrementAndGet() + } +} + +private[shuffle] class CometCelebornFallbackTestBackend(conf: SparkConf) extends ShuffleManager { + private val ordinaryShuffle = new CometShuffleManager(conf) + + override def registerShuffle[K, V, C]( + shuffleId: Int, + dependency: ShuffleDependency[K, V, C]): ShuffleHandle = dependency match { + case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle => + new CelebornShuffleHandle(shuffleId, dependency) + case _ => ordinaryShuffle.registerShuffle(shuffleId, dependency) + } + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + require( + !handle.isInstanceOf[CelebornShuffleHandle[_, _, _]], + "Unexpected ordinary native writer") + ordinaryShuffle.getWriter(handle, mapId, context, metrics) + } + + override def getReader[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { + require( + !handle.isInstanceOf[CelebornShuffleHandle[_, _, _]], + "Unexpected remote native reader") + ordinaryShuffle.getReader( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + } + + override def shuffleBlockResolver: ShuffleBlockResolver = ordinaryShuffle.shuffleBlockResolver + + override def unregisterShuffle(shuffleId: Int): Boolean = + ordinaryShuffle.unregisterShuffle(shuffleId) + + override def stop(): Unit = ordinaryShuffle.stop() +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala index 538531d7284..c22626a5208 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet.execution.shuffle +import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} + import scala.collection.mutable import org.scalatest.funsuite.AnyFunSuite @@ -360,6 +362,157 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)).authorized) } + test("an accepted size fallback abandons the remote shuffle and invalidates every owner") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val requestedFallbacks = mutable.ArrayBuffer.empty[Int] + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => true, + shuffleId => { + requestedFallbacks += shuffleId + true + }) + val shuffleId = 13 + val stageId = 77 + startStage(sparkCoordinator, stageId, numMappers = 2) + val generation = PrepareCelebornShuffleGeneration(shuffleId, 99, stageId, 0, 2) + assert(coordinator.prepareGeneration(generation)) + val validations = (0 until 2).map { mapId => + val owner = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, mapId, 0)) + ValidateCelebornMapAttempt(shuffleId, 99, stageId, 0, mapId, 0, owner.epoch) + } + assert(validations.forall(coordinator.validateMapAttempt)) + + assert(coordinator.requestLocalShuffle(RequestLocalCometShuffle(validations.head, 100L))) + assert(requestedFallbacks.toSeq == Seq(shuffleId)) + assert(validations.forall(validation => !coordinator.validateMapAttempt(validation))) + assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validations.last, 101L))) + assert( + !coordinator.prepareGeneration(generation.copy(celebornShuffleId = 100, stageAttempt = 1))) + assert( + !coordinator + .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 0, 1)) + .authorized) + assert( + !coordinator + .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 0)) + .authorized) + assert(requestedFallbacks.toSeq == Seq(shuffleId)) + + coordinator.unregisterShuffle(shuffleId) + assert( + coordinator.prepareGeneration(generation.copy(celebornShuffleId = 101, stageAttempt = 2))) + } + + test("a declined size fallback preserves the remote generation and can be requested again") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + var acceptFallback = false + var requests = 0 + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => true, + _ => { + requests += 1 + acceptFallback + }) + val stageId = 80 + startStage(sparkCoordinator, stageId, numMappers = 1) + assert( + coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(16, 105, stageId, 0, 1))) + val remote = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(16, stageId, 0, 0, 0)) + val validation = ValidateCelebornMapAttempt(16, 105, stageId, 0, 0, 0, remote.epoch) + val request = RequestLocalCometShuffle(validation, 105L) + + assert(!coordinator.requestLocalShuffle(request)) + assert(coordinator.validateMapAttempt(validation)) + acceptFallback = true + assert(coordinator.requestLocalShuffle(request)) + assert(!coordinator.validateMapAttempt(validation)) + assert(requests == 2) + } + + test("concurrent size reports invoke the active fallback once") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val requestedFallbacks = mutable.ArrayBuffer.empty[Int] + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => true, + shuffleId => { + requestedFallbacks += shuffleId + true + }) + val stageId = 78 + startStage(sparkCoordinator, stageId, numMappers = 2) + assert( + coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(14, 102, stageId, 0, 2))) + val requests = (0 until 2).map { mapId => + val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(14, stageId, 0, mapId, 0)) + RequestLocalCometShuffle( + ValidateCelebornMapAttempt(14, 102, stageId, 0, mapId, 0, owner.epoch), + 102L + mapId) + } + val executor = Executors.newFixedThreadPool(2) + val ready = new CountDownLatch(2) + val start = new CountDownLatch(1) + try { + val reports = requests.map { request => + executor.submit(new java.util.concurrent.Callable[Boolean] { + override def call(): Boolean = { + ready.countDown() + assert(start.await(5, TimeUnit.SECONDS)) + coordinator.requestLocalShuffle(request) + } + }) + } + assert(ready.await(5, TimeUnit.SECONDS)) + start.countDown() + assert(reports.map(_.get(5, TimeUnit.SECONDS)).count(identity) == 1) + assert(requestedFallbacks.toSeq == Seq(14)) + } finally { + start.countDown() + executor.shutdownNow() + } + } + + test("stale or unsafe size reports never invoke the active fallback") { + Seq(false, true).foreach { stale => + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => stale, + _ => fail("A stale or unsafe map attempt must not request fallback")) + val stageId = 79 + startStage(sparkCoordinator, stageId, numMappers = 1) + val generation = PrepareCelebornShuffleGeneration(15, 103, stageId, 0, 1) + assert(coordinator.prepareGeneration(generation)) + val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(15, stageId, 0, 0, 0)) + val validation = ValidateCelebornMapAttempt(15, 103, stageId, 0, 0, 0, owner.epoch) + if (stale) { + assert( + coordinator.prepareGeneration( + generation.copy(celebornShuffleId = 104, stageAttempt = 1))) + } + + assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validation, 104L))) + if (!stale) assert(coordinator.validateMapAttempt(validation)) + } + } + + test("a size report without an active materialization preserves the remote generation") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator, _ => true) + val stageId = 81 + startStage(sparkCoordinator, stageId, numMappers = 1) + assert( + coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(17, 106, stageId, 0, 1))) + val remote = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(17, stageId, 0, 0, 0)) + val validation = ValidateCelebornMapAttempt(17, 106, stageId, 0, 0, 0, remote.epoch) + + assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validation, 106L))) + assert(coordinator.validateMapAttempt(validation)) + } + test("manager preserves the application Spark configuration and driver identity") { val conf = new SparkConf(false).set("spark.app.name", "existing-celeborn-application") val backend = new RecordingShuffleManager @@ -414,6 +567,24 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(composite.registerShuffle[Any, Any, Any](31, null) eq backend.returnedHandle) } + test("native planning uses client compatibility without a scheduler-version gate") { + var configurationLoads = 0 + var completionProbes = 0 + val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( + new SparkConf(false), + _ => { + configurationLoads += 1 + new ReflectedPlanningConf + }, + () => { + completionProbes += 1 + None + }) + assert(support.fallbackReason(1).isEmpty) + assert(configurationLoads == 1) + assert(completionProbes == 1) + } + test("encrypted native planning falls back before loading the optional Celeborn API") { val conf = new SparkConf(false).set("spark.io.encryption.enabled", "true") var loads = 0 @@ -429,6 +600,75 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(loads == 0) } + test("native planning requires executor retention for local size-limit fallback") { + val cases = Seq( + ("fixed executors", false, false, false, true), + ("shuffle tracking", true, false, true, true), + ("external shuffle service", true, true, false, true), + ("both retention mechanisms", true, true, true, true), + ("remote storage alone", true, false, false, false)) + cases.foreach { case (name, dynamicAllocation, shuffleService, tracking, supported) => + val conf = new SparkConf(false) + .set("spark.dynamicAllocation.enabled", dynamicAllocation.toString) + .set("spark.shuffle.service.enabled", shuffleService.toString) + .set("spark.dynamicAllocation.shuffleTracking.enabled", tracking.toString) + var probes = 0 + val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( + conf, + _ => new ReflectedPlanningConf, + () => { + probes += 1 + None + }) + assert(support.fallbackReason(1).isEmpty == supported, name) + assert(probes == (if (supported) 1 else 0), name) + } + + // Read Spark's ConfigEntry so omitting the setting agrees with the executor allocation + // manager's effective behavior on each supported Spark version. + val defaultTracking = CometCelebornShuffleManager.nativeShufflePlanningSupport( + new SparkConf(false).set("spark.dynamicAllocation.enabled", "true"), + _ => new ReflectedPlanningConf, + () => None) + assert( + defaultTracking.fallbackReason(1).isEmpty == + new SparkConf(false).get( + org.apache.spark.internal.config.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED)) + } + + test("unprotected local fallback retains ordinary delegated shuffle") { + Seq(false, true).foreach { decommission => + val conf = new SparkConf(false) + .set("spark.dynamicAllocation.enabled", "true") + .set("spark.shuffle.service.enabled", "false") + .set("spark.dynamicAllocation.shuffleTracking.enabled", "false") + .set("spark.decommission.enabled", decommission.toString) + .set("spark.storage.decommission.enabled", decommission.toString) + .set("spark.storage.decommission.shuffleBlocks.enabled", decommission.toString) + val backend = new RecordingShuffleManager + val composite = new CometCelebornShuffleManager( + conf, + false, + (_, _) => backend, + planningSupportFactory = actualConf => + CometCelebornShuffleManager.nativeShufflePlanningSupport( + actualConf, + _ => fail("Unprotected local fallback must not load native Celeborn configuration"), + () => fail("Unprotected local fallback must not probe native push completion"))) + try { + assert(composite.nativeShuffleFallbackReason(1).exists(_.contains("local fallback"))) + val handle = composite.registerShuffle[Any, Any, Any](31, null) + assert(composite.getWriter[Any, Any](handle, 12L, null, null) == null) + assert(composite.getReader[Any, Any](handle, 0, 2, null, null) == null) + assert(backend.writerCall.contains((handle, 12L))) + assert(backend.rangedReaderCall.contains((handle, 0, Int.MaxValue, 0, 2))) + } finally { + composite.stop() + } + assert(backend.stopped) + } + } + test("native planning uses Celeborn's effective stage-rerun setting") { // Celeborn resolves JVM properties and legacy aliases itself. Spark SQL settings must not // override the resolved value used by the actual client and lifecycle manager. diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index db89c1840f0..c1e186e0eac 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -24,7 +24,7 @@ import java.lang.reflect.{InvocationHandler, Method, Proxy} import java.nio.{ByteBuffer, ByteOrder} import java.nio.channels.Channels import java.nio.charset.StandardCharsets -import java.nio.file.Files +import java.nio.file.{Files, Path} import java.util.{LinkedHashSet, Set => JSet} import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit} import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} @@ -81,11 +81,15 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { extends ShuffleManager { var readRange: Option[(Int, Int, Int, Int)] = None var unregistered: Option[Int] = None + var fileToRemove: Option[Path] = None + var stopped = false override def registerShuffle[K, V, C]( shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = - throw new UnsupportedOperationException("This fixture only reads existing native shuffles") + new CometNativeShuffleHandle[K, V]( + shuffleId, + dependency.asInstanceOf[ShuffleDependency[K, V, V]]) override def getWriter[K, V]( handle: ShuffleHandle, @@ -110,10 +114,11 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { override def unregisterShuffle(shuffleId: Int): Boolean = { unregistered = Some(shuffleId) + fileToRemove.foreach(Files.deleteIfExists(_)) true } - override def stop(): Unit = () + override def stop(): Unit = stopped = true } private final class RecordingReaderApi extends CelebornRawPartitionReader.Api { @@ -169,13 +174,17 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { context.taskMetrics().createTempShuffleReadMetrics(), retries) - private def simpleDependency(partitions: Int = 3, attributes: Seq[Attribute] = Seq.empty) + private def simpleDependency( + partitions: Int = 3, + attributes: Seq[Attribute] = Seq.empty, + useLocalShuffle: Boolean = false) : CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch] = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( spark.sparkContext.emptyRDD[(Int, ColumnarBatch)], new HashPartitioner(partitions), decodeTime = SQLMetrics.createMetric(spark.sparkContext, "Celeborn decode time"), - outputAttributes = attributes) + outputAttributes = attributes, + useLocalShuffle = useLocalShuffle) private def ownedClients(manager: CometCelebornShuffleManager) : java.util.concurrent.ConcurrentHashMap[AnyRef, java.lang.Boolean] = { @@ -298,6 +307,47 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { context.markTaskCompleted(None) } + test("local fallback shutdown preserves files until explicit shuffle unregister") { + Seq(false, true).foreach { unregister => + val context = TaskContext.empty() + val dependency = simpleDependency(useLocalShuffle = true) + val backend = new RecordingReaderBackend(new RecordingCelebornRawClient) + val local = new RecordingReaderBackend(new RecordingCelebornRawClient) + val file = Files.createTempFile("comet-local-fallback", ".data") + local.fileToRemove = Some(file) + val manager = new CometCelebornShuffleManager( + new SparkConf(false).set("spark.shuffle.service.enabled", "true"), + false, + (_, _) => backend, + localManagerFactory = _ => local) + val handle = manager.registerShuffle(17, dependency) + assert(handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) + try { + manager.getReader[Int, ColumnarBatch]( + handle, + 0, + 1, + context, + context.taskMetrics.createTempShuffleReadMetrics()) + assert(local.readRange.contains((0, Int.MaxValue, 0, 1))) + assert(backend.readRange.isEmpty) + if (unregister) { + assert(manager.unregisterShuffle(17)) + assert(local.unregistered.contains(17)) + assert(backend.unregistered.isEmpty) + } + manager.stop() + assert(local.stopped) + assert(backend.stopped) + assert(Files.exists(file) == !unregister) + if (!unregister) assert(local.unregistered.isEmpty) + } finally { + Files.deleteIfExists(file) + context.markTaskCompleted(None) + } + } + } + test("the manager routes native Celeborn handles through its reflected raw reader") { val context = TaskContext.empty() val dependency = simpleDependency() diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala new file mode 100644 index 00000000000..f4be2680366 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala @@ -0,0 +1,132 @@ +/* + * 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.sql.comet.execution.shuffle + +import scala.concurrent.Await +import scala.concurrent.duration.DurationInt + +import org.apache.spark.{SparkConf, SparkEnv} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter.{SHUFFLE_BYTES_WRITTEN, SHUFFLE_RECORDS_WRITTEN} + +import org.apache.comet.CometConf + +class CometCelebornShuffleStatisticsSuite extends CometTestBase { + import testImplicits._ + + override protected val shuffleManager: String = + classOf[CometCelebornFallbackTestShuffleManager].getName + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SHUFFLE_MODE.key, "native") + .set(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key, "false") + .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "32k") + .set(CometConf.COMET_EXCHANGE_SIZE_MULTIPLIER.key, "1") + .set("spark.shuffle.compress", "false") + .set("spark.celeborn.client.spark.stageRerun.enabled", "true") + + private def manager: CometCelebornFallbackTestShuffleManager = + SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornFallbackTestShuffleManager] + + for (adaptive <- Seq(false, true)) { + test(s"fallback publishes only the local destination's output statistics: AQE=$adaptive") { + withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { + manager.remoteMapsBeforeFailure.set(0) + manager.waitForFirstRemoteMap = true + try { + val shuffled = spark + .range(0, 2, 1, 2) + .selectExpr( + "CAST(id AS INT) AS key", + "CASE WHEN id = 0 THEN 'small' ELSE repeat('x', 1048576) END AS payload") + .repartition(3, $"key") + val actual = shuffled.collect().map(row => (row.getInt(0), row.getString(1))).toSeq + val exchange = collect(shuffled.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val remote = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + val local = remote.currentShuffleDependency + + assert(actual.sortBy(_._1) == Seq((0, "small"), (1, "x" * (1024 * 1024)))) + assert(manager.remoteMapsBeforeFailure.get() == 1) + assert(local.useLocalShuffle) + assert(remote.shuffleWriteMetrics(SHUFFLE_RECORDS_WRITTEN).value == 1L) + assert(exchange.runtimeStatistics.rowCount.contains(BigInt(2))) + assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L) + assert(remote.shuffleWriteMetrics("dataSize").value > 0L) + assert( + exchange.runtimeStatistics.sizeInBytes == + BigInt(local.shuffleWriteMetrics("dataSize").value)) + assert( + exchange.metrics("dataSize").value == + local.shuffleWriteMetrics("dataSize").value) + assert( + exchange.metrics(SHUFFLE_BYTES_WRITTEN).value == + local.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).value) + + // Simulate another accumulator merge from the abandoned remote stage after publication. + // Resetting the original counters would allow that update to inflate AQE statistics. + val published = exchange.runtimeStatistics + val publishedBytes = exchange.metrics(SHUFFLE_BYTES_WRITTEN).value + remote.shuffleWriteMetrics("dataSize").add(4096L) + remote.shuffleWriteMetrics(SHUFFLE_RECORDS_WRITTEN).add(1L) + remote.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).add(4096L) + assert(exchange.runtimeStatistics == published) + assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L) + assert(exchange.metrics("dataSize").value == published.sizeInBytes.toLong) + assert(exchange.metrics(SHUFFLE_BYTES_WRITTEN).value == publishedBytes) + + assert(manager.unregisterShuffle(remote.shuffleId)) + assert(manager.unregisterShuffle(local.shuffleId)) + } finally { + manager.waitForFirstRemoteMap = false + } + } + } + } + + test("successful remote materialization publishes remote output statistics") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val query = spark.range(0, 2, 1, 2).repartition(2, $"id") + val exchange = collect(query.queryExecution.executedPlan) { + case value: CometShuffleExchangeExec => value + }.headOption.getOrElse(fail("Expected a native Comet shuffle exchange")) + val remote = exchange.shuffleDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + try { + val statistics = Await.result(exchange.mapOutputStatisticsFuture, 20.seconds) + assert(remote.currentShuffleDependency eq remote) + assert(exchange.runtimeStatistics.rowCount.contains(BigInt(2))) + assert(exchange.runtimeStatistics.sizeInBytes > 0) + assert(exchange.metrics(SHUFFLE_RECORDS_WRITTEN).value == 2L) + assert(statistics.bytesByPartitionId.sum > 0L) + // MapOutputStatistics uses Spark's compressed MapStatus sizes, which are approximate. + assert( + exchange.metrics(SHUFFLE_BYTES_WRITTEN).value == + remote.shuffleWriteMetrics(SHUFFLE_BYTES_WRITTEN).value) + assert( + exchange.metrics("dataSize").value == + exchange.runtimeStatistics.sizeInBytes.toLong) + } finally { + manager.unregisterShuffle(remote.shuffleId) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index c3a19c8790b..99d9fe93e0c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark.{HashPartitioner, Partition, TaskContext} -import org.apache.spark.rdd.RDD +import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} @@ -45,6 +45,73 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ class CometNativeShuffleInputRDDSuite extends CometTestBase { + test("native shuffle input preserves its parents' determinism") { + Seq( + DeterministicLevel.DETERMINATE, + DeterministicLevel.UNORDERED, + DeterministicLevel.INDETERMINATE).foreach { level => + val parent = new RDD[AnyRef](spark.sparkContext, Nil) { + override protected def getOutputDeterministicLevel: DeterministicLevel.Value = level + override protected def getPartitions: Array[Partition] = Array.empty + override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] = + Iterator.empty + } + val input = new CometNativeShuffleInputRDD( + spark.sparkContext, + Seq(parent), + 0, + Set.empty, + CometMetricNode(Map.empty)) + assert(input.outputDeterministicLevel == level) + assert(input.copyForLocalShuffle().outputDeterministicLevel == level) + } + } + + test("local shuffle input is an independent sibling with the same partition inputs") { + val upstream = new RDD[AnyRef](spark.sparkContext, Nil) { + override protected def getPartitions: Array[Partition] = Array.tabulate(2) { i => + new Partition { + override def index: Int = i + } + } + + override def compute(split: Partition, context: TaskContext): Iterator[AnyRef] = + Iterator.single(null) + + override def getPreferredLocations(split: Partition): Seq[String] = + Seq(s"host-${split.index}") + } + val planData = Map("scan-0" -> Array(Array[Byte](1), Array[Byte](2))) + val remote = new CometNativeShuffleInputRDD( + spark.sparkContext, + Seq(upstream), + 2, + Set.empty, + CometMetricNode(Map.empty), + planData) + val local = remote.copyForLocalShuffle() + + assert(local.id != remote.id) + assert(local.dependencies.map(_.rdd) == Seq(upstream)) + assert(local.getNumPartitions == remote.getNumPartitions) + local.partitions.foreach { part => + val input = part.asInstanceOf[CometNativeShuffleInputPartition] + assert(input.inputPartitions.toSeq == Seq(upstream.partitions(part.index))) + assert(input.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index))) + assert(local.preferredLocations(part) == Seq(s"host-${part.index}")) + + val context = TaskContext.empty() + try { + val iterator = local.iterator(part, context).asInstanceOf[CometNativeShuffleInputIterator] + assert(iterator.partitionIndex == part.index) + assert(iterator.inputObjects.toSeq == Seq(null)) + assert(iterator.planDataByKey("scan-0").sameElements(planData("scan-0")(part.index))) + } finally { + context.markTaskCompleted(None) + } + } + } + test("spill reporting is registered before native shuffle input producers") { Seq(None, Some(new IllegalStateException("failed native shuffle"))).foreach { failure => val writerDisk = new SQLMetric("writerDisk")