From 8b26b2951576155f807212850b087fa26756deae Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Fri, 4 Sep 2026 02:16:42 +0000 Subject: [PATCH 1/4] fix: recover native Celeborn shuffle from oversized rows --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../comet/CometShuffleSizeLimitException.java | 27 ++ docs/source/user-guide/latest/configs.md | 17 + docs/source/user-guide/latest/tuning.md | 25 +- native/jni-bridge/src/errors.rs | 63 ++- native/shuffle/src/writers/rss/mod.rs | 151 ++++++- .../src/writers/rss/rss_partition_writer.rs | 80 +++- .../CelebornShufflePartitionPusher.java | 2 +- .../scala/org/apache/comet/CometConf.scala | 7 +- .../shuffle/CometCelebornShuffleManager.scala | 368 ++++++++++++++++-- .../shuffle/CometNativeShuffleInputRDD.scala | 16 +- .../shuffle/CometNativeShuffleWriter.scala | 18 +- .../shuffle/CometShuffleExchangeExec.scala | 6 +- .../org/apache/comet/CometConfSuite.scala | 2 +- .../celeborn/CelebornShuffleHandle.scala | 4 +- ...ometCelebornNativeShuffleWriterSuite.scala | 301 +++++++++++++- .../CometCelebornShuffleFallbackSuite.scala | 308 +++++++++++++++ .../CometCelebornShuffleManagerSuite.scala | 313 ++++++++++++++- .../CometCelebornShuffleReaderSuite.scala | 73 +++- 20 files changed, 1696 insertions(+), 87 deletions(-) create mode 100644 common/src/main/java/org/apache/comet/CometShuffleSizeLimitException.java create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0d62838cd27..7b097f31335 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -374,6 +374,7 @@ 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.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..3ea63b740ff 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -147,6 +147,7 @@ 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.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..cf6e6771112 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -45,6 +45,23 @@ 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. + +Native Celeborn shuffle requires Spark 3.5.1 or newer for safe whole-stage recovery. Earlier +Spark versions retain ordinary Spark/Celeborn shuffle. If a row cannot fit the remote limits +on a supported Spark version, Comet retries that shuffle using its local writer and reader. +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..efffb423076 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -257,7 +257,11 @@ 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 requires Spark 3.5.1 or newer, whose scheduler discards late map results from an +obsolete stage attempt during recovery. Earlier Spark versions retain ordinary Spark/Celeborn +shuffle, including when native mode is requested. + +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 +306,23 @@ 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 invalidates that shuffle's remote output and retries the +whole map stage using its local shuffle writer. All subsequent reads and retries for that +shuffle use local files and Spark's block transfer service. Native operators and Comet's Arrow +shuffle format are preserved, and remote admission limits remain enforced. This 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..fd40529de2a 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 @@ -27,14 +27,15 @@ import scala.collection.mutable import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.{MapOutputTrackerMaster, 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.spark.scheduler.{MapStatus, OutputCommitCoordinator} +import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.util.{RpcUtils, VersionUtils} 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 +52,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. */ @@ -77,6 +79,16 @@ class CometCelebornShuffleManager private[shuffle] ( 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, @@ -110,16 +122,23 @@ class CometCelebornShuffleManager private[shuffle] ( nativeDependency(handle) match { case Some(dependency) => val earlyClaim = claimNativeShuffleAttempt(handle.shuffleId, context) + if (earlyClaim.useLocalShuffle) { + if (earlyClaim.requiresStageRetry) { + throw localShuffleRetry(handle.shuffleId) + } + if (!earlyClaim.authorized) { + throw CelebornShufflePusherFactory.commitDenied(context) + } + return localWriter(handle.shuffleId, dependency, mapId, context, metrics, earlyClaim) + } if (!earlyClaim.authorized && context.attemptNumber() > 0 && !earlyClaim.requiresGenerationResolution) { 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 +190,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) @@ -190,6 +217,16 @@ class CometCelebornShuffleManager private[shuffle] ( metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = { nativeDependency(handle) match { case Some(dependency) => + if (usesLocalShuffle(handle.shuffleId)) { + return localShuffleManager.getReader( + localHandle[K, C](handle.shuffleId, dependency), + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics) + } if (startMapIndex > endMapIndex) { throw new UnsupportedOperationException( "Celeborn physical-skew chunk reads are not supported by native Comet shuffle") @@ -238,32 +275,55 @@ 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)) - } - 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() + } + 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 +333,11 @@ class CometCelebornShuffleManager private[shuffle] ( } val coordinator = new CelebornShuffleGenerationCoordinator( env.outputCommitCoordinator, - CelebornShufflePusherFactory.shouldReportShuffleFetchFailure) + shouldReportShuffleFetchFailure, + shuffleId => + env.mapOutputTracker + .asInstanceOf[MapOutputTrackerMaster] + .unregisterAllMapAndMergeOutput(shuffleId)) val endpoint = env.rpcEnv.setupEndpoint( CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT, new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator)) @@ -310,6 +374,124 @@ 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) + + protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = + generationEndpoint.askSync[Boolean](UsesLocalCometShuffle(shuffleId)) + + private def localHandle[K, V]( + shuffleId: Int, + dependency: CometShuffleDependency[_, _, _]): CometNativeShuffleHandle[K, V] = + new CometNativeShuffleHandle(shuffleId, dependency.asInstanceOf[ShuffleDependency[K, V, V]]) + + private def localWriter[K, V]( + shuffleId: Int, + dependency: CometShuffleDependency[_, _, _], + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter, + claim: CelebornMapAttemptClaim): ShuffleWriter[K, V] = { + val writer = localShuffleManager.getWriter[K, V]( + localHandle[K, V](shuffleId, dependency), + mapId, + context, + metrics) + 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] = { + if (success && !generationEndpoint.askSync[Boolean]( + ValidateLocalCometMapAttempt( + ClaimCelebornMapAttempt( + shuffleId, + context.stageId(), + context.stageAttemptNumber(), + context.partitionId(), + context.attemptNumber()), + claim.epoch))) { + val failure = CelebornShufflePusherFactory.commitDenied(context) + try writer.stop(false) + catch { + case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure) + } + throw failure + } + writer.stop(success) + } + } + } + + private def localShuffleRetry(shuffleId: Int, cause: Throwable = null): FetchFailedException = + new FetchFailedException( + null, + shuffleId, + -1L, + -1, + -1, + s"Native Celeborn shuffle $shuffleId exceeded its size limits; " + + "retrying the complete map stage with local Comet shuffle", + cause) + + 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) { + throw CelebornShufflePusherFactory.commitDenied(context) + } + + // The driver already cleared every map output before publishing the local decision. A lost + // executor or failed Celeborn RPC cannot leave a usable partial remote generation behind. + 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 localShuffleRetry(shuffleId, failure) + } + private def prepareNativeShuffleGeneration( shuffleId: Int, celebornShuffleId: Int, @@ -440,18 +622,40 @@ private[shuffle] object CometCelebornShuffleManager { .getMethod("fromSparkConf", classOf[SparkConf]) .invoke(null, conf) + private[shuffle] def supportsNativeShuffleStageRecovery(sparkVersion: String): Boolean = + VersionUtils.majorMinorPatchVersion(sparkVersion).exists { version => + implicitly[Ordering[(Int, Int, Int)]].gteq(version, (3, 5, 1)) + } + private[shuffle] def nativeShufflePlanningSupport( conf: SparkConf, loadCelebornConf: SparkConf => AnyRef = reflectedCelebornConf, pushCompletionUnavailableReason: () => Option[String] = () => Option( CelebornShufflePartitionPusher.nativePushCompletionUnavailableReason( - ClassLoaders.loadClass("org.apache.celeborn.client.ShuffleClientImpl")))) + ClassLoaders.loadClass("org.apache.celeborn.client.ShuffleClientImpl"))), + sparkVersion: String = org.apache.spark.SPARK_VERSION) : CelebornNativeShufflePlanningSupport = { + if (!supportsNativeShuffleStageRecovery(sparkVersion)) { + // Earlier schedulers accept late map successes from an obsolete indeterminate stage. + // Such a remote MapStatus could replace local output after the destination has changed. + return CelebornNativeShufflePlanningSupport( + Some("Native Celeborn shuffle requires Spark 3.5.1 or newer for safe stage recovery")) + } if (conf.getBoolean("spark.io.encryption.enabled", false)) { 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) { @@ -545,7 +749,9 @@ private[shuffle] final case class ClaimCelebornMapAttempt( private[shuffle] final case class CelebornMapAttemptClaim( authorized: Boolean, epoch: Long, - requiresGenerationResolution: Boolean = false) + requiresGenerationResolution: Boolean = false, + useLocalShuffle: Boolean = false, + requiresStageRetry: Boolean = false) extends Serializable private[shuffle] final case class ValidateCelebornMapAttempt( @@ -577,12 +783,25 @@ private[shuffle] final case class AbandonCelebornMapAttempt( taskAttemptId: Long) extends Serializable +private[shuffle] final case class RequestLocalCometShuffle( + validation: ValidateCelebornMapAttempt, + taskAttemptId: Long) + extends Serializable + +private[shuffle] final case class UsesLocalCometShuffle(shuffleId: Int) extends Serializable + +private[shuffle] final case class ValidateLocalCometMapAttempt( + claim: ClaimCelebornMapAttempt, + claimEpoch: 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, + unregisterMapOutputs: Int => Unit = _ => ()) { private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration] private val invalidatedGenerations = mutable.HashSet.empty[Int] @@ -590,6 +809,11 @@ 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]] + // A destination decision belongs to the Spark shuffle, not a task or Celeborn generation. + // Remember the failed stage until unregister so a retry cannot select the same size-limited + // writer again. The latest local stage also fences obsolete local commit claims. + private case class LocalShuffleFallback(failedStage: (Int, Int), latestStage: (Int, Int)) + private val localShuffles = mutable.HashMap.empty[Int, LocalShuffleFallback] private def currentEpoch(shuffleId: Int): Long = generationEpochs.getOrElse(shuffleId, 0L) @@ -664,6 +888,36 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( } def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized { + localShuffles.get(claim.shuffleId).foreach { local => + val stage = (claim.stageId, claim.stageAttempt) + val ordering = implicitly[Ordering[(Int, Int)]] + if (ordering.lt(stage, local.latestStage) || + !attemptCanRun(claim.stageId, claim.stageAttempt, claim.mapId, claim.taskAttempt)) { + return CelebornMapAttemptClaim( + false, + currentEpoch(claim.shuffleId), + useLocalShuffle = true) + } + if (ordering.lteq(stage, local.failedStage)) { + // The original reporter may die after the driver accepts fallback. A replacement task + // in that same stage must still trigger stage recovery, not write mixed map outputs. + return CelebornMapAttemptClaim( + false, + currentEpoch(claim.shuffleId), + useLocalShuffle = true, + requiresStageRetry = true) + } + if (ordering.gt(stage, local.latestStage)) { + invalidateOwners(claim.shuffleId) + localShuffles.update(claim.shuffleId, local.copy(latestStage = stage)) + } + return authorize( + claim.shuffleId, + claim.stageId, + claim.stageAttempt, + claim.mapId, + claim.taskAttempt).copy(useLocalShuffle = true) + } val previousGeneration = generations.get(claim.shuffleId) val stale = previousGeneration.exists { generation => generation.stageId == claim.stageId && @@ -724,6 +978,9 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized { require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive") + if (localShuffles.contains(generation.shuffleId)) { + return false + } generations.get(generation.shuffleId) match { case Some(previous) @@ -867,7 +1124,44 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( abandon } + def requestLocalShuffle(request: RequestLocalCometShuffle): Boolean = synchronized { + val validation = request.validation + if (localShuffles.contains(validation.shuffleId) || + !validateMapAttempt(validation) || + !shouldReportShuffleFetchFailure(request.taskAttemptId)) { + return false + } + + // Clear Spark's complete output set before exposing the destination change. The native + // input RDD is indeterminate, so Spark also rejects late success from the abandoned stage + // and starts every map again when it processes the accompanying FetchFailedException. + unregisterMapOutputs(validation.shuffleId) + val stage = (validation.stageId, validation.stageAttempt) + localShuffles.update(validation.shuffleId, LocalShuffleFallback(stage, stage)) + invalidatedGenerations.add(validation.shuffleId) + invalidateOwners(validation.shuffleId) + true + } + + def usesLocalShuffle(shuffleId: Int): Boolean = synchronized { + localShuffles.contains(shuffleId) + } + + def validateLocalMapAttempt(validation: ValidateLocalCometMapAttempt): Boolean = synchronized { + val claim = validation.claim + localShuffles.get(claim.shuffleId).exists { local => + local.latestStage == ((claim.stageId, claim.stageAttempt)) && + local.latestStage != local.failedStage && + currentEpoch(claim.shuffleId) == validation.claimEpoch && + attemptCanRun(claim.stageId, claim.stageAttempt, claim.mapId, claim.taskAttempt) && + claimOwners + .get(ownerKey(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId)) + .contains((claim.taskAttempt, validation.claimEpoch)) + } + } + def unregisterShuffle(shuffleId: Int): Unit = synchronized { + localShuffles.remove(shuffleId) generations.remove(shuffleId) invalidatedGenerations.remove(shuffleId) generationEpochs.remove(shuffleId) @@ -894,5 +1188,11 @@ 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)) + case UsesLocalCometShuffle(shuffleId) => + context.reply(coordinator.usesLocalShuffle(shuffleId)) + case validation: ValidateLocalCometMapAttempt => + context.reply(coordinator.validateLocalMapAttempt(validation)) } } 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..85b4723681d 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 @@ -20,7 +20,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ -import org.apache.spark.rdd.RDD +import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -40,11 +40,23 @@ private[shuffle] class CometNativeShuffleInputRDD( numPartitionsParam: Int, shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, - @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty) + @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, + private[shuffle] val requiresStageRetry: Boolean = false) extends RDD[Product2[Int, ColumnarBatch]]( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { + override protected def getOutputDeterministicLevel: DeterministicLevel.Value = { + if (requiresStageRetry) { + // A Celeborn generation can be replaced by local shuffle after a size-limit failure. + // Spark must recompute every map and ignore late successes from the previous stage + // attempt; otherwise remote MapStatus entries could be published as local file output. + DeterministicLevel.INDETERMINATE + } else { + super.getOutputDeterministicLevel + } + } + 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/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 46fe621e481..0f5f5597ff1 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 @@ -127,7 +127,8 @@ case class CometShuffleExchangeExec( ctx.numPartitions, ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), - ctx.perPartitionByKey) + ctx.perPartitionByKey, + requiresStageRetry = isCometCelebornShuffleManagerEnabled(conf)) case None => // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets // built via the convenience overload below; we just need a real RDD of batches. @@ -748,7 +749,8 @@ object CometShuffleExchangeExec Seq(streamRDD), rdd.getNumPartitions, shuffleScanIndices = Set.empty, - spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode))) + spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)), + requiresStageRetry = isCometCelebornShuffleManagerEnabled(conf)) val ctx = NativeExecContext( inputs = Seq(streamRDD), 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/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/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..201397a624c --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala @@ -0,0 +1,308 @@ +/* + * 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.{ConcurrentLinkedQueue, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{MapOutputTrackerMaster, ShuffleDependency, SparkConf, SparkEnv, TaskContext} +import org.apache.spark.rdd.DeterministicLevel +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.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 retries its complete map stage locally and succeeds: AQE=$adaptive") { + if (!CometCelebornShuffleManager.supportsNativeShuffleStageRecovery( + org.apache.spark.SPARK_VERSION)) { + cancel("Native Celeborn stage recovery requires Spark 3.5.1 or newer") + } + 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 writersBefore = manager.writerStages.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") + assert(remote.head._2 == 0) + val stages = manager.writerStages.asScala.drop(writersBefore).toSeq + assert(stages.sorted == Seq(0, 1), s"expected one remote and one local stage: $stages") + assert(manager.nativeReads.get() > readsBefore) + val client = remote.head._3 + assert(client.cleanupCalls.get() == 1) + assert(client.mapperEndCalls.get() == 0) + assert(client.fetchFailureReports.get() == 1) + + // Shuffle removal must clean both destinations after the successful replacement. + assert(manager.unregisterShuffle(remote.head._1)) + assert(client.shuffleCleanupCalls.get() == 1) + } + } + } + + test(s"fallback recomputes a previously completed remote map: AQE=$adaptive") { + if (!CometCelebornShuffleManager.supportsNativeShuffleStageRecovery( + org.apache.spark.SPARK_VERSION)) { + cancel("Native Celeborn stage recovery requires Spark 3.5.1 or newer") + } + withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { + val remoteBefore = manager.remoteAttempts.size() + val writersBefore = manager.writerAttempts.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).distinct.size == 1) + assert(remote.forall(_._2 == 0)) + assert(remote.count(_._3.mapperEndCalls.get() == 1) == 1) + assert(remote.map(_._3.fetchFailureReports.get()).sum == 1) + val attempts = manager.writerAttempts.asScala.drop(writersBefore).toSeq + assert( + attempts.sorted == Seq((0, 0), (0, 1), (1, 0), (1, 1)), + s"both maps must be recomputed exactly once after fallback: $attempts") + assert(manager.nativeReads.get() > readsBefore) + assert(manager.unregisterShuffle(remote.head._1)) + } finally { + manager.waitForFirstRemoteMap = false + } + } + } + } +} + +class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean) + extends CometCelebornShuffleManager( + conf, + isDriver, + (configuration, _) => new CometCelebornFallbackTestBackend(configuration), + planningSupportFactory = _ => CelebornNativeShufflePlanningSupport()) { + + val remoteAttempts = + new ConcurrentLinkedQueue[(Int, Int, CometCelebornFallbackTestClient)]() + val writerStages = new ConcurrentLinkedQueue[Int]() + val nativeReads = new AtomicInteger() + val writerAttempts = new ConcurrentLinkedQueue[(Int, Int)]() + val remoteMapsBeforeFailure = new AtomicInteger() + @volatile var waitForFirstRemoteMap = false + + 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((handle.shuffleId, context.stageAttemptNumber(), 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 writer = super.getWriter[K, V](handle, mapId, context, metrics) + if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]) { + writerStages.add(context.stageAttemptNumber()) + writerAttempts.add((context.stageAttemptNumber(), context.partitionId())) + } + 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[CelebornShuffleHandle[_, _, _]]) { + nativeReads.incrementAndGet() + } + reader + } +} + +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 => + require( + dependency.rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE, + "Remote shuffle input must invalidate every previous-stage map result during fallback") + 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..37bb8cef058 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,166 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)).authorized) } + test("a size failure invalidates every remote map and selects local shuffle until removal") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val removedOutputs = mutable.ArrayBuffer.empty[Int] + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => true, + shuffleId => removedOutputs += shuffleId) + val shuffleId = 13 + val stageId = 77 + startStage(sparkCoordinator, stageId, numMappers = 2) + val generation = PrepareCelebornShuffleGeneration(shuffleId, 99, stageId, 0, 2) + assert(coordinator.prepareGeneration(generation)) + val owners = (0 until 2).map { mapId => + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, mapId, 0)) + } + val validations = owners.zipWithIndex.map { case (owner, mapId) => + ValidateCelebornMapAttempt(shuffleId, 99, stageId, 0, mapId, 0, owner.epoch) + } + assert(validations.forall(coordinator.validateMapAttempt)) + + assert(coordinator.requestLocalShuffle(RequestLocalCometShuffle(validations.head, 100L))) + assert(coordinator.usesLocalShuffle(shuffleId)) + assert(removedOutputs.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))) + val sourceStageRetry = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 0, 1)) + assert(!sourceStageRetry.authorized) + assert(sourceStageRetry.useLocalShuffle) + assert(sourceStageRetry.requiresStageRetry) + + val localAttempt = ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 0) + val local = coordinator.claimMapAttempt(localAttempt) + assert(local.authorized) + assert(local.useLocalShuffle) + assert(local.epoch > owners.head.epoch) + assert( + coordinator.validateLocalMapAttempt( + ValidateLocalCometMapAttempt(localAttempt, local.epoch))) + completeFailedAttempt(sparkCoordinator, stageId, 1, 0, 0) + assert( + !coordinator.validateLocalMapAttempt( + ValidateLocalCometMapAttempt(localAttempt, local.epoch))) + val localRetry = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 1)) + assert(localRetry.authorized) + assert(localRetry.useLocalShuffle) + val staleSource = + coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 1, 1)) + assert(!staleSource.authorized) + assert(!staleSource.requiresStageRetry) + assert(removedOutputs.toSeq == Seq(shuffleId)) + + coordinator.unregisterShuffle(shuffleId) + assert(!coordinator.usesLocalShuffle(shuffleId)) + assert( + coordinator.prepareGeneration(generation.copy(celebornShuffleId = 101, stageAttempt = 2))) + } + + test("a newer local stage fences map completion from its previous local attempt") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + 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)) + assert( + coordinator.requestLocalShuffle( + RequestLocalCometShuffle( + ValidateCelebornMapAttempt(16, 105, stageId, 0, 0, 0, remote.epoch), + 105L))) + val firstAttempt = ClaimCelebornMapAttempt(16, stageId, 1, 0, 0) + val first = coordinator.claimMapAttempt(firstAttempt) + assert( + coordinator.validateLocalMapAttempt( + ValidateLocalCometMapAttempt(firstAttempt, first.epoch))) + + val replacementAttempt = firstAttempt.copy(stageAttempt = 2) + val replacement = coordinator.claimMapAttempt(replacementAttempt) + assert(replacement.authorized) + assert(replacement.useLocalShuffle) + assert( + !coordinator.validateLocalMapAttempt( + ValidateLocalCometMapAttempt(firstAttempt, first.epoch))) + assert( + coordinator.validateLocalMapAttempt( + ValidateLocalCometMapAttempt(replacementAttempt, replacement.epoch))) + } + + test("concurrent size reports select local shuffle and discard map outputs once") { + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val removedOutputs = mutable.ArrayBuffer.empty[Int] + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => true, + shuffleId => removedOutputs += shuffleId) + 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(coordinator.usesLocalShuffle(14)) + assert(removedOutputs.toSeq == Seq(14)) + } finally { + start.countDown() + executor.shutdownNow() + } + } + + test("stale or unsafe size reports preserve the current remote generation") { + Seq(false, true).foreach { stale => + val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) + val removedOutputs = mutable.ArrayBuffer.empty[Int] + val coordinator = new CelebornShuffleGenerationCoordinator( + sparkCoordinator, + _ => stale, + shuffleId => removedOutputs += shuffleId) + 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))) + assert(!coordinator.usesLocalShuffle(15)) + assert(removedOutputs.isEmpty) + if (!stale) 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 @@ -398,7 +560,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( actualConf, _ => effectiveConf, - () => None) + () => None, + sparkVersion = "3.5.1") }) // Mutating settings after manager construction cannot change its native capabilities. @@ -414,6 +577,55 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(composite.registerShuffle[Any, Any, Any](31, null) eq backend.returnedHandle) } + test("older Spark schedulers retain ordinary shuffle before probing native Celeborn APIs") { + Seq("3.4.0", "3.4.3", "3.4.4", "3.5.0", "3.5.0-SNAPSHOT", "unknown").foreach { version => + val backend = new RecordingShuffleManager + val composite = new CometCelebornShuffleManager( + new SparkConf(false), + false, + (_, _) => backend, + planningSupportFactory = conf => + CometCelebornShuffleManager.nativeShufflePlanningSupport( + conf, + _ => fail(s"Spark $version must not load native Celeborn configuration"), + () => fail(s"Spark $version must not probe native push completion"), + sparkVersion = version)) + try { + assert(composite.nativeShuffleFallbackReason(1).exists(_.contains("Spark 3.5.1"))) + val handle = composite.registerShuffle[Any, Any, Any](31, null) + assert(handle eq backend.returnedHandle) + 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("Spark versions with obsolete-stage filtering can select native Celeborn shuffle") { + Seq("3.5.1", "3.5.10", "3.5.1-SNAPSHOT", "4.0.0", "4.1.0").foreach { version => + var configurationLoads = 0 + var completionProbes = 0 + val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( + new SparkConf(false), + _ => { + configurationLoads += 1 + new ReflectedPlanningConf + }, + () => { + completionProbes += 1 + None + }, + sparkVersion = version) + assert(support.fallbackReason(1).isEmpty, s"Spark $version should support stage recovery") + 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 @@ -423,12 +635,82 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { loads += 1 new ReflectedPlanningConf }, - () => fail("Encrypted applications must not probe push completion")) + () => fail("Encrypted applications must not probe push completion"), + sparkVersion = "3.5.1") assert(support.fallbackReason(1).exists(_.contains("encryption"))) 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 + }, + sparkVersion = "3.5.1") + assert(support.fallbackReason(1).isEmpty == supported, name) + assert(probes == (if (supported) 1 else 0), name) + } + + // Spark 3.5.1 and later enable shuffle tracking by default. Read Spark's ConfigEntry so + // omitting the setting agrees with the executor allocation manager's effective behavior. + val defaultTracking = CometCelebornShuffleManager.nativeShufflePlanningSupport( + new SparkConf(false).set("spark.dynamicAllocation.enabled", "true"), + _ => new ReflectedPlanningConf, + () => None, + sparkVersion = "3.5.1") + assert(defaultTracking.fallbackReason(1).isEmpty) + } + + 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"), + sparkVersion = "3.5.1")) + 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. @@ -438,7 +720,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( conf, _ => new ReflectedPlanningConf(stageReruns = false), - () => None) + () => None, + sparkVersion = "3.5.1") assert(support.fallbackReason(1).nonEmpty) assert(support.fallbackReason(100).nonEmpty) @@ -463,7 +746,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { () => { completionProbes += 1 Some(reason) - })) + }, + sparkVersion = "3.5.1")) try { assert(composite.nativeShuffleFallbackReason(1).contains(reason)) assert(composite.nativeShuffleFallbackReason(Int.MaxValue).contains(reason)) @@ -487,7 +771,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(policy = policy, partitionThreshold = 4L), - () => None) + () => None, + sparkVersion = "3.5.1") assert(support.fallbackReason(1).isEmpty) assert(support.fallbackReason(3).isEmpty) assert(support.fallbackReason(4).nonEmpty) @@ -497,7 +782,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val always = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(policy = "ALWAYS", partitionThreshold = Long.MaxValue), - () => None) + () => None, + sparkVersion = "3.5.1") assert(always.fallbackReason(1).nonEmpty) // Celeborn's explicit NEVER policy takes precedence over the deprecated force-fallback flag. @@ -505,7 +791,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { new SparkConf(false) .set("spark.celeborn.client.spark.shuffle.forceFallback.enabled", "true"), _ => new ReflectedPlanningConf(policy = "NEVER", partitionThreshold = 1L), - () => None) + () => None, + sparkVersion = "3.5.1") assert(never.fallbackReason(1).isEmpty) assert(never.fallbackReason(4).isEmpty) } @@ -514,14 +801,16 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val largeThreshold = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(partitionThreshold = Int.MaxValue.toLong + 1L), - () => None) + () => None, + sparkVersion = "3.5.1") assert(largeThreshold.fallbackReason(Int.MaxValue).isEmpty) Seq(0L, -1L).foreach { threshold => val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(partitionThreshold = threshold), - () => None) + () => None, + sparkVersion = "3.5.1") assert(support.fallbackReason(1).nonEmpty) } } @@ -544,7 +833,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), loader, - () => None) + () => None, + sparkVersion = "3.5.1") assert(support.fallbackReason(1).nonEmpty) } } @@ -559,7 +849,8 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( conf, _ => throw new ClassNotFoundException("native planning API is unavailable"), - () => None)) + () => None, + sparkVersion = "3.5.1")) val handle = composite.registerShuffle[Any, Any, Any](31, null) assert(composite.nativeShuffleFallbackReason(1).nonEmpty) 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..15c196b081b 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,6 +81,8 @@ 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, @@ -110,10 +112,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 { @@ -279,7 +282,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { val api = new RecordingReaderApi val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val remote = manager .getReader[Int, ColumnarBatch]( @@ -298,6 +303,48 @@ 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() + val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + 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) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = true + } + 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.contains(17)) + } + 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() @@ -308,7 +355,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { val api = new RecordingReaderApi val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val remote = manager .getReader[Int, ColumnarBatch]( @@ -352,7 +401,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { dependency, stageRerunEnabled = unavailableApi) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val failure = intercept[IllegalStateException] { manager.getReader[Int, ColumnarBatch]( @@ -387,7 +438,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.retryLimitFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val failure = intercept[IllegalStateException] { manager.getReader[Int, ColumnarBatch]( @@ -417,7 +470,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.generationFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val failure = intercept[FetchFailedException] { manager.getReader[Int, ColumnarBatch]( @@ -445,7 +500,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.generationFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { + override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false + } val failure = intercept[ClassCastException] { manager.getReader[Int, ColumnarBatch]( From 6f2465a2f3d1cb2915a7c7f5551aa7032719467d Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Fri, 4 Sep 2026 03:20:31 +0000 Subject: [PATCH 2/4] fix: preserve fetch recovery after Celeborn fallback --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + docs/source/user-guide/latest/configs.md | 7 +- docs/source/user-guide/latest/tuning.md | 17 +- .../shuffle/CometCelebornShuffleManager.scala | 242 +++------ .../CometCelebornShuffleMaterialization.scala | 298 +++++++++++ .../shuffle/CometNativeShuffleInputRDD.scala | 16 +- .../shuffle/CometShuffleDependency.scala | 49 +- .../shuffle/CometShuffleExchangeExec.scala | 19 +- .../shuffle/CometShuffledRowRDD.scala | 5 + .../comet/exec/CometNativeShuffleSuite.scala | 9 +- .../CometCelebornLocalFetchFailureSuite.scala | 139 +++++ .../CometCelebornShuffleFallbackSuite.scala | 475 ++++++++++++++++-- .../CometCelebornShuffleManagerSuite.scala | 239 ++++----- .../CometCelebornShuffleReaderSuite.scala | 47 +- .../CometNativeShuffleInputRDDSuite.scala | 23 +- 16 files changed, 1162 insertions(+), 425 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornLocalFetchFailureSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 7b097f31335..db061fbaa0d 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -375,6 +375,7 @@ jobs: 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.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 3ea63b740ff..b9da5984d4d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -148,6 +148,7 @@ jobs: 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.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/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md index cf6e6771112..7b6001d8cac 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -54,9 +54,10 @@ reservation budget accommodates ordinary frames up to the default 64 MiB frame l Compressed frames still need workspace for their uncompressed data. Increase the reservation budget when larger rows or schemas need more workspace. -Native Celeborn shuffle requires Spark 3.5.1 or newer for safe whole-stage recovery. Earlier -Spark versions retain ordinary Spark/Celeborn shuffle. If a row cannot fit the remote limits -on a supported Spark version, Comet retries that shuffle using its local writer and reader. +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 identity, so late remote map results cannot replace local output. 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 diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index efffb423076..7d1a832c4c1 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -257,10 +257,6 @@ 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 Spark 3.5.1 or newer, whose scheduler discards late map results from an -obsolete stage attempt during recovery. Earlier Spark versions retain ordinary Spark/Celeborn -shuffle, including when native mode is requested. - 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 @@ -313,11 +309,14 @@ and codec overhead. Compression reduces the transmitted bytes but still needs un encoding workspace. Comet splits large batches between rows. If a single row, its schema, or its encoding workspace -cannot fit the remote limits, Comet invalidates that shuffle's remote output and retries the -whole map stage using its local shuffle writer. All subsequent reads and retries for that -shuffle use local files and Spark's block transfer service. Native operators and Comet's Arrow -shuffle format are preserved, and remote admission limits remain enforced. This fallback uses -executor disk. When `spark.dynamicAllocation.enabled=true`, native Celeborn shuffle requires +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 identity, so late remote map results cannot overwrite or skip local map +output. 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 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 fd40529de2a..39e73b3d158 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 @@ -27,12 +27,12 @@ import scala.collection.mutable import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.spark.{MapOutputTrackerMaster, ShuffleDependency, SparkConf, SparkEnv, TaskContext} +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.{MapStatus, OutputCommitCoordinator} -import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} -import org.apache.spark.util.{RpcUtils, VersionUtils} +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, ResolvedCelebornShufflePusher} @@ -76,6 +76,8 @@ class CometCelebornShuffleManager private[shuffle] ( 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 = _ @@ -94,6 +96,10 @@ class CometCelebornShuffleManager private[shuffle] ( 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)) { @@ -119,18 +125,13 @@ 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) - if (earlyClaim.useLocalShuffle) { - if (earlyClaim.requiresStageRetry) { - throw localShuffleRetry(handle.shuffleId) - } - if (!earlyClaim.authorized) { - throw CelebornShufflePusherFactory.commitDenied(context) - } - return localWriter(handle.shuffleId, dependency, mapId, context, metrics, earlyClaim) - } if (!earlyClaim.authorized && context.attemptNumber() > 0 && !earlyClaim.requiresGenerationResolution) { throw CelebornShufflePusherFactory.commitDenied(context) @@ -215,18 +216,19 @@ 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 (usesLocalShuffle(handle.shuffleId)) { - return localShuffleManager.getReader( - localHandle[K, C](handle.shuffleId, dependency), - startMapIndex, - endMapIndex, - startPartition, - endPartition, - context, - metrics) - } if (startMapIndex > endMapIndex) { throw new UnsupportedOperationException( "Celeborn physical-skew chunk reads are not supported by native Comet shuffle") @@ -275,6 +277,10 @@ class CometCelebornShuffleManager private[shuffle] ( override def shuffleBlockResolver: ShuffleBlockResolver = backend.shuffleBlockResolver override def unregisterShuffle(shuffleId: Int): Boolean = { + sizeLimitFallbacks.remove(shuffleId) + if (localShuffleIds.remove(shuffleId)) { + return localShuffleManager.unregisterShuffle(shuffleId) + } val generations = Option(nativeShuffleClients.remove(shuffleId)).toSeq .flatMap(_.asScala.toSeq) var removed = false @@ -299,6 +305,7 @@ class CometCelebornShuffleManager private[shuffle] ( val localCleanup = Option(localShuffleManagerInstance).toSeq.map { local => () => local.stop() } + sizeLimitFallbacks.clear() cleanupAll( localCleanup ++ Seq[() => Unit]( () => backend.stop(), @@ -334,10 +341,7 @@ class CometCelebornShuffleManager private[shuffle] ( val coordinator = new CelebornShuffleGenerationCoordinator( env.outputCommitCoordinator, shouldReportShuffleFetchFailure, - shuffleId => - env.mapOutputTracker - .asInstanceOf[MapOutputTrackerMaster] - .unregisterAllMapAndMergeOutput(shuffleId)) + shuffleId => Option(sizeLimitFallbacks.get(shuffleId)).exists(callback => callback())) val endpoint = env.rpcEnv.setupEndpoint( CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT, new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator)) @@ -392,63 +396,28 @@ class CometCelebornShuffleManager private[shuffle] ( protected[shuffle] def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean = CelebornShufflePusherFactory.shouldReportShuffleFetchFailure(taskAttemptId) - protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = - generationEndpoint.askSync[Boolean](UsesLocalCometShuffle(shuffleId)) - - private def localHandle[K, V]( + /** Registers the driver-owned materialization that can replace an unpublished shuffle. */ + private[shuffle] def registerSizeLimitFallback( shuffleId: Int, - dependency: CometShuffleDependency[_, _, _]): CometNativeShuffleHandle[K, V] = - new CometNativeShuffleHandle(shuffleId, dependency.asInstanceOf[ShuffleDependency[K, V, V]]) + 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 def localWriter[K, V]( - shuffleId: Int, - dependency: CometShuffleDependency[_, _, _], - mapId: Long, - context: TaskContext, - metrics: ShuffleWriteMetricsReporter, - claim: CelebornMapAttemptClaim): ShuffleWriter[K, V] = { - val writer = localShuffleManager.getWriter[K, V]( - localHandle[K, V](shuffleId, dependency), - mapId, - context, - metrics) - 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] = { - if (success && !generationEndpoint.askSync[Boolean]( - ValidateLocalCometMapAttempt( - ClaimCelebornMapAttempt( - shuffleId, - context.stageId(), - context.stageAttemptNumber(), - context.partitionId(), - context.attemptNumber()), - claim.epoch))) { - val failure = CelebornShufflePusherFactory.commitDenied(context) - try writer.stop(false) - catch { - case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure) - } - throw failure - } - writer.stop(success) - } - } + private[shuffle] def removeSizeLimitFallback(shuffleId: Int): Unit = { + sizeLimitFallbacks.remove(shuffleId) } - private def localShuffleRetry(shuffleId: Int, cause: Throwable = null): FetchFailedException = - new FetchFailedException( - null, - shuffleId, - -1L, - -1, - -1, - s"Native Celeborn shuffle $shuffleId exceeded its size limits; " + - "retrying the complete map stage with local Comet shuffle", - cause) + 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, @@ -469,11 +438,13 @@ class CometCelebornShuffleManager private[shuffle] ( claim.epoch), context.taskAttemptId())) if (!accepted) { - throw CelebornShufflePusherFactory.commitDenied(context) + // 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 already cleared every map output before publishing the local decision. A lost - // executor or failed Celeborn RPC cannot leave a usable partial remote generation behind. + // 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( @@ -489,7 +460,7 @@ class CometCelebornShuffleManager private[shuffle] ( } catch { case NonFatal(cleanupFailure) => failure.addSuppressed(cleanupFailure) } - throw localShuffleRetry(shuffleId, failure) + throw failure } private def prepareNativeShuffleGeneration( @@ -622,26 +593,14 @@ private[shuffle] object CometCelebornShuffleManager { .getMethod("fromSparkConf", classOf[SparkConf]) .invoke(null, conf) - private[shuffle] def supportsNativeShuffleStageRecovery(sparkVersion: String): Boolean = - VersionUtils.majorMinorPatchVersion(sparkVersion).exists { version => - implicitly[Ordering[(Int, Int, Int)]].gteq(version, (3, 5, 1)) - } - private[shuffle] def nativeShufflePlanningSupport( conf: SparkConf, loadCelebornConf: SparkConf => AnyRef = reflectedCelebornConf, pushCompletionUnavailableReason: () => Option[String] = () => Option( CelebornShufflePartitionPusher.nativePushCompletionUnavailableReason( - ClassLoaders.loadClass("org.apache.celeborn.client.ShuffleClientImpl"))), - sparkVersion: String = org.apache.spark.SPARK_VERSION) + ClassLoaders.loadClass("org.apache.celeborn.client.ShuffleClientImpl")))) : CelebornNativeShufflePlanningSupport = { - if (!supportsNativeShuffleStageRecovery(sparkVersion)) { - // Earlier schedulers accept late map successes from an obsolete indeterminate stage. - // Such a remote MapStatus could replace local output after the destination has changed. - return CelebornNativeShufflePlanningSupport( - Some("Native Celeborn shuffle requires Spark 3.5.1 or newer for safe stage recovery")) - } if (conf.getBoolean("spark.io.encryption.enabled", false)) { return CelebornNativeShufflePlanningSupport( Some("Native Celeborn shuffle does not support spark.io.encryption.enabled=true")) @@ -749,9 +708,7 @@ private[shuffle] final case class ClaimCelebornMapAttempt( private[shuffle] final case class CelebornMapAttemptClaim( authorized: Boolean, epoch: Long, - requiresGenerationResolution: Boolean = false, - useLocalShuffle: Boolean = false, - requiresStageRetry: Boolean = false) + requiresGenerationResolution: Boolean = false) extends Serializable private[shuffle] final case class ValidateCelebornMapAttempt( @@ -788,20 +745,13 @@ private[shuffle] final case class RequestLocalCometShuffle( taskAttemptId: Long) extends Serializable -private[shuffle] final case class UsesLocalCometShuffle(shuffleId: Int) extends Serializable - -private[shuffle] final case class ValidateLocalCometMapAttempt( - claim: ClaimCelebornMapAttempt, - claimEpoch: 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, - unregisterMapOutputs: Int => Unit = _ => ()) { + requestFallback: Int => Boolean = _ => false) { private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration] private val invalidatedGenerations = mutable.HashSet.empty[Int] @@ -809,11 +759,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]] - // A destination decision belongs to the Spark shuffle, not a task or Celeborn generation. - // Remember the failed stage until unregister so a retry cannot select the same size-limited - // writer again. The latest local stage also fences obsolete local commit claims. - private case class LocalShuffleFallback(failedStage: (Int, Int), latestStage: (Int, Int)) - private val localShuffles = mutable.HashMap.empty[Int, LocalShuffleFallback] + // 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) @@ -888,35 +836,8 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( } def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized { - localShuffles.get(claim.shuffleId).foreach { local => - val stage = (claim.stageId, claim.stageAttempt) - val ordering = implicitly[Ordering[(Int, Int)]] - if (ordering.lt(stage, local.latestStage) || - !attemptCanRun(claim.stageId, claim.stageAttempt, claim.mapId, claim.taskAttempt)) { - return CelebornMapAttemptClaim( - false, - currentEpoch(claim.shuffleId), - useLocalShuffle = true) - } - if (ordering.lteq(stage, local.failedStage)) { - // The original reporter may die after the driver accepts fallback. A replacement task - // in that same stage must still trigger stage recovery, not write mixed map outputs. - return CelebornMapAttemptClaim( - false, - currentEpoch(claim.shuffleId), - useLocalShuffle = true, - requiresStageRetry = true) - } - if (ordering.gt(stage, local.latestStage)) { - invalidateOwners(claim.shuffleId) - localShuffles.update(claim.shuffleId, local.copy(latestStage = stage)) - } - return authorize( - claim.shuffleId, - claim.stageId, - claim.stageAttempt, - claim.mapId, - claim.taskAttempt).copy(useLocalShuffle = true) + if (abandonedShuffles.contains(claim.shuffleId)) { + return CelebornMapAttemptClaim(false, currentEpoch(claim.shuffleId)) } val previousGeneration = generations.get(claim.shuffleId) val stale = previousGeneration.exists { generation => @@ -978,7 +899,7 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized { require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive") - if (localShuffles.contains(generation.shuffleId)) { + if (abandonedShuffles.contains(generation.shuffleId)) { return false } @@ -1126,42 +1047,21 @@ private[shuffle] final class CelebornShuffleGenerationCoordinator( def requestLocalShuffle(request: RequestLocalCometShuffle): Boolean = synchronized { val validation = request.validation - if (localShuffles.contains(validation.shuffleId) || + if (abandonedShuffles.contains(validation.shuffleId) || !validateMapAttempt(validation) || - !shouldReportShuffleFetchFailure(request.taskAttemptId)) { + !shouldReportShuffleFetchFailure(request.taskAttemptId) || + !requestFallback(validation.shuffleId)) { return false } - // Clear Spark's complete output set before exposing the destination change. The native - // input RDD is indeterminate, so Spark also rejects late success from the abandoned stage - // and starts every map again when it processes the accompanying FetchFailedException. - unregisterMapOutputs(validation.shuffleId) - val stage = (validation.stageId, validation.stageAttempt) - localShuffles.update(validation.shuffleId, LocalShuffleFallback(stage, stage)) + abandonedShuffles.add(validation.shuffleId) invalidatedGenerations.add(validation.shuffleId) invalidateOwners(validation.shuffleId) true } - def usesLocalShuffle(shuffleId: Int): Boolean = synchronized { - localShuffles.contains(shuffleId) - } - - def validateLocalMapAttempt(validation: ValidateLocalCometMapAttempt): Boolean = synchronized { - val claim = validation.claim - localShuffles.get(claim.shuffleId).exists { local => - local.latestStage == ((claim.stageId, claim.stageAttempt)) && - local.latestStage != local.failedStage && - currentEpoch(claim.shuffleId) == validation.claimEpoch && - attemptCanRun(claim.stageId, claim.stageAttempt, claim.mapId, claim.taskAttempt) && - claimOwners - .get(ownerKey(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId)) - .contains((claim.taskAttempt, validation.claimEpoch)) - } - } - def unregisterShuffle(shuffleId: Int): Unit = synchronized { - localShuffles.remove(shuffleId) + abandonedShuffles.remove(shuffleId) generations.remove(shuffleId) invalidatedGenerations.remove(shuffleId) generationEpochs.remove(shuffleId) @@ -1190,9 +1090,5 @@ private[shuffle] final class CelebornShuffleGenerationEndpoint( context.reply(coordinator.abandonMapAttempt(abandoned)) case request: RequestLocalCometShuffle => context.reply(coordinator.requestLocalShuffle(request)) - case UsesLocalCometShuffle(shuffleId) => - context.reply(coordinator.usesLocalShuffle(shuffleId)) - case validation: ValidateLocalCometMapAttempt => - context.reply(coordinator.validateLocalMapAttempt(validation)) } } 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..4a672f49878 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleMaterialization.scala @@ -0,0 +1,298 @@ +/* + * 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.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.scheduler.{SparkListener, SparkListenerJobStart} +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; its shuffle and stage IDs isolate all late remote + * 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] { + + import CometCelebornShuffleMaterialization._ + + private val sparkContext = remoteDependency.rdd.context + 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 = { + 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): Boolean = { + try { + val action = synchronized { + if (state != expected) return false + // Keep submission and activeAction assignment atomic with cancellation and a size report. + // submitMapStage queues the job without waiting for executor tasks to finish. + val submitted = withCapturedProperties(sparkContext.submitMapStage(dependency)) + activeAction = Some(submitted) + if (expected == RunningRemote) remoteAction = Some(submitted) + actions += submitted + submitted + } + action.onComplete(result => finish(dependency, expected, result))(ExecutionContext.global) + true + } catch { + case NonFatal(failure) => + fail(expected, failure) + cancelActions(failure) + false + } + } + + 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 + } + outcome match { + case Success(_) => selected = Some(dependency) + case _ => + } + state = Finished + completion.tryComplete(outcome) + outcome.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 input partitions were + // already computed for the remote submission, and submitMapStage only queues execution. + // 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 85b4723681d..41d1d6cc6a4 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 @@ -20,7 +20,7 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ -import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.rdd.RDD import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -40,23 +40,11 @@ private[shuffle] class CometNativeShuffleInputRDD( numPartitionsParam: Int, shuffleScanIndices: Set[Int], spillMetricNode: CometMetricNode, - @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty, - private[shuffle] val requiresStageRetry: Boolean = false) + @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty) extends RDD[Product2[Int, ColumnarBatch]]( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { - override protected def getOutputDeterministicLevel: DeterministicLevel.Value = { - if (requiresStageRetry) { - // A Celeborn generation can be replaced by local shuffle after a size-limit failure. - // Spark must recompute every map and ignore late successes from the previous stage - // attempt; otherwise remote MapStatus entries could be published as local file output. - DeterministicLevel.INDETERMINATE - } else { - super.getOutputDeterministicLevel - } - } - 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/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 2a058430074..02bb1ae59f6 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 @@ -68,7 +68,8 @@ 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) extends ShuffleDependency[K, V, C]( _rdd, partitioner, @@ -76,7 +77,51 @@ 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] = + new CometShuffleDependency[K, V, C]( + rdd, + partitioner, + serializer, + keyOrdering, + aggregator, + mapSideCombine, + shuffleWriterProcessor, + shuffleType, + schema, + decodeTime, + outputPartitioning, + outputAttributes, + shuffleWriteMetrics, + numParts, + rangePartitionBounds, + nativeShuffleSpec, + useLocalShuffle = true) +} /** 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 0f5f5597ff1..e8800c16003 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 @@ -127,8 +127,7 @@ case class CometShuffleExchangeExec( ctx.numPartitions, ctx.shuffleScanIndices, CometMetricNode(metrics, Seq(nativeChildMetricNode)), - ctx.perPartitionByKey, - requiresStageRetry = isCometCelebornShuffleManagerEnabled(conf)) + ctx.perPartitionByKey) case None => // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets // built via the convenience overload below; we just need a real RDD of batches. @@ -149,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) + } } } @@ -168,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 @@ -749,8 +757,7 @@ object CometShuffleExchangeExec Seq(streamRDD), rdd.getNumPartitions, shuffleScanIndices = Set.empty, - spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode)), - requiresStageRetry = isCometCelebornShuffleManagerEnabled(conf)) + spillMetricNode = CometMetricNode(metrics, Seq(childMetricNode))) val ctx = NativeExecContext( inputs = Seq(streamRDD), 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..d738c552cd6 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,6 +51,11 @@ class CometShuffledBatchRDD( SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY, SQLConf.get.fetchShuffleBlocksInBatch.toString) + // Materialize native Celeborn output before exposing a dependency to downstream stages. A + // size-limit fallback gets its own shuffle and stage IDs, so an obsolete remote completion + // cannot replace local output or tell Spark to skip a replacement map partition. + dependency = CometCelebornShuffleMaterialization.selectForRead(dependency) + override def getDependencies: Seq[Dependency[_]] = List(dependency) override val partitioner: Option[Partitioner] = 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/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/CometCelebornShuffleFallbackSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala index 201397a624c..842b4be4417 100644 --- 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 @@ -19,18 +19,24 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.util.concurrent.{ConcurrentLinkedQueue, TimeUnit} +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.{MapOutputTrackerMaster, ShuffleDependency, SparkConf, SparkEnv, TaskContext} -import org.apache.spark.rdd.DeterministicLevel +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} /** @@ -57,11 +63,7 @@ class CometCelebornShuffleFallbackSuite extends CometTestBase { SparkEnv.get.shuffleManager.asInstanceOf[CometCelebornFallbackTestShuffleManager] for (adaptive <- Seq(false, true)) { - test(s"an oversized row retries its complete map stage locally and succeeds: AQE=$adaptive") { - if (!CometCelebornShuffleManager.supportsNativeShuffleStageRecovery( - org.apache.spark.SPARK_VERSION)) { - cancel("Native Celeborn stage recovery requires Spark 3.5.1 or newer") - } + 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 => @@ -72,7 +74,7 @@ class CometCelebornShuffleFallbackSuite extends CometTestBase { .option("parquet.enable.dictionary", "false") .parquet(path.getCanonicalPath) val remoteBefore = manager.remoteAttempts.size() - val writersBefore = manager.writerStages.size() + val localBefore = manager.localAttempts.size() val readsBefore = manager.nativeReads.get() val shuffled = spark.read.parquet(path.getCanonicalPath).repartition(3, $"key") @@ -81,30 +83,29 @@ class CometCelebornShuffleFallbackSuite extends CometTestBase { 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") - assert(remote.head._2 == 0) - val stages = manager.writerStages.asScala.drop(writersBefore).toSeq - assert(stages.sorted == Seq(0, 1), s"expected one remote and one local stage: $stages") + 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._3 + val client = remote.head._2 assert(client.cleanupCalls.get() == 1) assert(client.mapperEndCalls.get() == 0) assert(client.fetchFailureReports.get() == 1) - // Shuffle removal must clean both destinations after the successful replacement. - assert(manager.unregisterShuffle(remote.head._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") { - if (!CometCelebornShuffleManager.supportsNativeShuffleStageRecovery( - org.apache.spark.SPARK_VERSION)) { - cancel("Native Celeborn stage recovery requires Spark 3.5.1 or newer") - } withSQLConf("spark.sql.adaptive.enabled" -> adaptive.toString) { val remoteBefore = manager.remoteAttempts.size() - val writersBefore = manager.writerAttempts.size() + val localBefore = manager.localAttempts.size() val readsBefore = manager.nativeReads.get() manager.remoteMapsBeforeFailure.set(0) manager.waitForFirstRemoteMap = true @@ -123,21 +124,337 @@ class CometCelebornShuffleFallbackSuite extends CometTestBase { assert(manager.remoteMapsBeforeFailure.get() == 1) val remote = manager.remoteAttempts.asScala.drop(remoteBefore).toSeq assert(remote.size == 2) - assert(remote.map(_._1).distinct.size == 1) - assert(remote.forall(_._2 == 0)) - assert(remote.count(_._3.mapperEndCalls.get() == 1) == 1) - assert(remote.map(_._3.fetchFailureReports.get()).sum == 1) - val attempts = manager.writerAttempts.asScala.drop(writersBefore).toSeq - assert( - attempts.sorted == Seq((0, 0), (0, 1), (1, 0), (1, 1)), - s"both maps must be recomputed exactly once after fallback: $attempts") + 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)) + 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("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 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") } } @@ -148,13 +465,36 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean (configuration, _) => new CometCelebornFallbackTestBackend(configuration), planningSupportFactory = _ => CelebornNativeShufflePlanningSupport()) { - val remoteAttempts = - new ConcurrentLinkedQueue[(Int, Int, CometCelebornFallbackTestClient)]() - val writerStages = new ConcurrentLinkedQueue[Int]() + private[shuffle] val remoteAttempts = + new ConcurrentLinkedQueue[(CometShuffleFallbackAttempt, CometCelebornFallbackTestClient)]() + private[shuffle] val localAttempts = new ConcurrentLinkedQueue[CometShuffleFallbackAttempt]() val nativeReads = new AtomicInteger() - val writerAttempts = new ConcurrentLinkedQueue[(Int, Int)]() 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 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 @@ -171,7 +511,7 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean val generation = handle.shuffleId + 1000 val numMappers = remoteHandle.numMappers onGenerationResolved(generation, numMappers) - remoteAttempts.add((handle.shuffleId, context.stageAttemptNumber(), client)) + remoteAttempts.add((recordAttempt(handle, context), client)) ResolvedCelebornShufflePusher( CelebornShufflePusherFactory.create( conf, @@ -205,12 +545,57 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean 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 writer = super.getWriter[K, V](handle, mapId, context, metrics) - if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]) { - writerStages.add(context.stageAttemptNumber()) - writerAttempts.add((context.stageAttemptNumber(), context.partitionId())) + 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 } - writer } override def getReader[K, C]( @@ -229,11 +614,20 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean endPartition, context, metrics) - if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]]) { + 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 { @@ -260,9 +654,6 @@ private[shuffle] class CometCelebornFallbackTestBackend(conf: SparkConf) extends shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = dependency match { case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle => - require( - dependency.rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE, - "Remote shuffle input must invalidate every previous-stage map result during fallback") new CelebornShuffleHandle(shuffleId, dependency) case _ => ordinaryShuffle.registerShuffle(shuffleId, dependency) } 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 37bb8cef058..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 @@ -362,105 +362,86 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)).authorized) } - test("a size failure invalidates every remote map and selects local shuffle until removal") { + test("an accepted size fallback abandons the remote shuffle and invalidates every owner") { val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) - val removedOutputs = mutable.ArrayBuffer.empty[Int] + val requestedFallbacks = mutable.ArrayBuffer.empty[Int] val coordinator = new CelebornShuffleGenerationCoordinator( sparkCoordinator, _ => true, - shuffleId => removedOutputs += shuffleId) + 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 owners = (0 until 2).map { mapId => - coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, mapId, 0)) - } - val validations = owners.zipWithIndex.map { case (owner, mapId) => + 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(coordinator.usesLocalShuffle(shuffleId)) - assert(removedOutputs.toSeq == Seq(shuffleId)) + 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))) - val sourceStageRetry = - coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 0, 1)) - assert(!sourceStageRetry.authorized) - assert(sourceStageRetry.useLocalShuffle) - assert(sourceStageRetry.requiresStageRetry) - - val localAttempt = ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 0) - val local = coordinator.claimMapAttempt(localAttempt) - assert(local.authorized) - assert(local.useLocalShuffle) - assert(local.epoch > owners.head.epoch) assert( - coordinator.validateLocalMapAttempt( - ValidateLocalCometMapAttempt(localAttempt, local.epoch))) - completeFailedAttempt(sparkCoordinator, stageId, 1, 0, 0) + !coordinator + .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 0, 1)) + .authorized) assert( - !coordinator.validateLocalMapAttempt( - ValidateLocalCometMapAttempt(localAttempt, local.epoch))) - val localRetry = - coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 1)) - assert(localRetry.authorized) - assert(localRetry.useLocalShuffle) - val staleSource = - coordinator.claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 0, 1, 1)) - assert(!staleSource.authorized) - assert(!staleSource.requiresStageRetry) - assert(removedOutputs.toSeq == Seq(shuffleId)) + !coordinator + .claimMapAttempt(ClaimCelebornMapAttempt(shuffleId, stageId, 1, 0, 0)) + .authorized) + assert(requestedFallbacks.toSeq == Seq(shuffleId)) coordinator.unregisterShuffle(shuffleId) - assert(!coordinator.usesLocalShuffle(shuffleId)) assert( coordinator.prepareGeneration(generation.copy(celebornShuffleId = 101, stageAttempt = 2))) } - test("a newer local stage fences map completion from its previous local attempt") { + test("a declined size fallback preserves the remote generation and can be requested again") { val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) - val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator) + 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)) - assert( - coordinator.requestLocalShuffle( - RequestLocalCometShuffle( - ValidateCelebornMapAttempt(16, 105, stageId, 0, 0, 0, remote.epoch), - 105L))) - val firstAttempt = ClaimCelebornMapAttempt(16, stageId, 1, 0, 0) - val first = coordinator.claimMapAttempt(firstAttempt) - assert( - coordinator.validateLocalMapAttempt( - ValidateLocalCometMapAttempt(firstAttempt, first.epoch))) + val validation = ValidateCelebornMapAttempt(16, 105, stageId, 0, 0, 0, remote.epoch) + val request = RequestLocalCometShuffle(validation, 105L) - val replacementAttempt = firstAttempt.copy(stageAttempt = 2) - val replacement = coordinator.claimMapAttempt(replacementAttempt) - assert(replacement.authorized) - assert(replacement.useLocalShuffle) - assert( - !coordinator.validateLocalMapAttempt( - ValidateLocalCometMapAttempt(firstAttempt, first.epoch))) - assert( - coordinator.validateLocalMapAttempt( - ValidateLocalCometMapAttempt(replacementAttempt, replacement.epoch))) + 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 select local shuffle and discard map outputs once") { + test("concurrent size reports invoke the active fallback once") { val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true) - val removedOutputs = mutable.ArrayBuffer.empty[Int] + val requestedFallbacks = mutable.ArrayBuffer.empty[Int] val coordinator = new CelebornShuffleGenerationCoordinator( sparkCoordinator, _ => true, - shuffleId => removedOutputs += shuffleId) + shuffleId => { + requestedFallbacks += shuffleId + true + }) val stageId = 78 startStage(sparkCoordinator, stageId, numMappers = 2) assert( @@ -487,22 +468,20 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(ready.await(5, TimeUnit.SECONDS)) start.countDown() assert(reports.map(_.get(5, TimeUnit.SECONDS)).count(identity) == 1) - assert(coordinator.usesLocalShuffle(14)) - assert(removedOutputs.toSeq == Seq(14)) + assert(requestedFallbacks.toSeq == Seq(14)) } finally { start.countDown() executor.shutdownNow() } } - test("stale or unsafe size reports preserve the current remote generation") { + 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 removedOutputs = mutable.ArrayBuffer.empty[Int] val coordinator = new CelebornShuffleGenerationCoordinator( sparkCoordinator, _ => stale, - shuffleId => removedOutputs += shuffleId) + _ => 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) @@ -516,12 +495,24 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { } assert(!coordinator.requestLocalShuffle(RequestLocalCometShuffle(validation, 104L))) - assert(!coordinator.usesLocalShuffle(15)) - assert(removedOutputs.isEmpty) 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 @@ -560,8 +551,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( actualConf, _ => effectiveConf, - () => None, - sparkVersion = "3.5.1") + () => None) }) // Mutating settings after manager construction cannot change its native capabilities. @@ -577,53 +567,22 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { assert(composite.registerShuffle[Any, Any, Any](31, null) eq backend.returnedHandle) } - test("older Spark schedulers retain ordinary shuffle before probing native Celeborn APIs") { - Seq("3.4.0", "3.4.3", "3.4.4", "3.5.0", "3.5.0-SNAPSHOT", "unknown").foreach { version => - val backend = new RecordingShuffleManager - val composite = new CometCelebornShuffleManager( - new SparkConf(false), - false, - (_, _) => backend, - planningSupportFactory = conf => - CometCelebornShuffleManager.nativeShufflePlanningSupport( - conf, - _ => fail(s"Spark $version must not load native Celeborn configuration"), - () => fail(s"Spark $version must not probe native push completion"), - sparkVersion = version)) - try { - assert(composite.nativeShuffleFallbackReason(1).exists(_.contains("Spark 3.5.1"))) - val handle = composite.registerShuffle[Any, Any, Any](31, null) - assert(handle eq backend.returnedHandle) - 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("Spark versions with obsolete-stage filtering can select native Celeborn shuffle") { - Seq("3.5.1", "3.5.10", "3.5.1-SNAPSHOT", "4.0.0", "4.1.0").foreach { version => - var configurationLoads = 0 - var completionProbes = 0 - val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( - new SparkConf(false), - _ => { - configurationLoads += 1 - new ReflectedPlanningConf - }, - () => { - completionProbes += 1 - None - }, - sparkVersion = version) - assert(support.fallbackReason(1).isEmpty, s"Spark $version should support stage recovery") - assert(configurationLoads == 1) - assert(completionProbes == 1) - } + 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") { @@ -635,8 +594,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { loads += 1 new ReflectedPlanningConf }, - () => fail("Encrypted applications must not probe push completion"), - sparkVersion = "3.5.1") + () => fail("Encrypted applications must not probe push completion")) assert(support.fallbackReason(1).exists(_.contains("encryption"))) assert(loads == 0) @@ -661,20 +619,21 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { () => { probes += 1 None - }, - sparkVersion = "3.5.1") + }) assert(support.fallbackReason(1).isEmpty == supported, name) assert(probes == (if (supported) 1 else 0), name) } - // Spark 3.5.1 and later enable shuffle tracking by default. Read Spark's ConfigEntry so - // omitting the setting agrees with the executor allocation manager's effective behavior. + // 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, - sparkVersion = "3.5.1") - assert(defaultTracking.fallbackReason(1).isEmpty) + () => 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") { @@ -695,8 +654,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( actualConf, _ => fail("Unprotected local fallback must not load native Celeborn configuration"), - () => fail("Unprotected local fallback must not probe native push completion"), - sparkVersion = "3.5.1")) + () => 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) @@ -720,8 +678,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( conf, _ => new ReflectedPlanningConf(stageReruns = false), - () => None, - sparkVersion = "3.5.1") + () => None) assert(support.fallbackReason(1).nonEmpty) assert(support.fallbackReason(100).nonEmpty) @@ -746,8 +703,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { () => { completionProbes += 1 Some(reason) - }, - sparkVersion = "3.5.1")) + })) try { assert(composite.nativeShuffleFallbackReason(1).contains(reason)) assert(composite.nativeShuffleFallbackReason(Int.MaxValue).contains(reason)) @@ -771,8 +727,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(policy = policy, partitionThreshold = 4L), - () => None, - sparkVersion = "3.5.1") + () => None) assert(support.fallbackReason(1).isEmpty) assert(support.fallbackReason(3).isEmpty) assert(support.fallbackReason(4).nonEmpty) @@ -782,8 +737,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val always = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(policy = "ALWAYS", partitionThreshold = Long.MaxValue), - () => None, - sparkVersion = "3.5.1") + () => None) assert(always.fallbackReason(1).nonEmpty) // Celeborn's explicit NEVER policy takes precedence over the deprecated force-fallback flag. @@ -791,8 +745,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { new SparkConf(false) .set("spark.celeborn.client.spark.shuffle.forceFallback.enabled", "true"), _ => new ReflectedPlanningConf(policy = "NEVER", partitionThreshold = 1L), - () => None, - sparkVersion = "3.5.1") + () => None) assert(never.fallbackReason(1).isEmpty) assert(never.fallbackReason(4).isEmpty) } @@ -801,16 +754,14 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { val largeThreshold = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(partitionThreshold = Int.MaxValue.toLong + 1L), - () => None, - sparkVersion = "3.5.1") + () => None) assert(largeThreshold.fallbackReason(Int.MaxValue).isEmpty) Seq(0L, -1L).foreach { threshold => val support = CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), _ => new ReflectedPlanningConf(partitionThreshold = threshold), - () => None, - sparkVersion = "3.5.1") + () => None) assert(support.fallbackReason(1).nonEmpty) } } @@ -833,8 +784,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( new SparkConf(false), loader, - () => None, - sparkVersion = "3.5.1") + () => None) assert(support.fallbackReason(1).nonEmpty) } } @@ -849,8 +799,7 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite { CometCelebornShuffleManager.nativeShufflePlanningSupport( conf, _ => throw new ClassNotFoundException("native planning API is unavailable"), - () => None, - sparkVersion = "3.5.1")) + () => None)) val handle = composite.registerShuffle[Any, Any, Any](31, null) assert(composite.nativeShuffleFallbackReason(1).nonEmpty) 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 15c196b081b..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 @@ -87,7 +87,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { 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, @@ -172,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] = { @@ -282,9 +288,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { val api = new RecordingReaderApi val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val remote = manager .getReader[Int, ColumnarBatch]( @@ -306,8 +310,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { test("local fallback shutdown preserves files until explicit shuffle unregister") { Seq(false, true).foreach { unregister => val context = TaskContext.empty() - val dependency = simpleDependency() - val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) + 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") @@ -316,9 +319,9 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { new SparkConf(false).set("spark.shuffle.service.enabled", "true"), false, (_, _) => backend, - localManagerFactory = _ => local) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = true - } + localManagerFactory = _ => local) + val handle = manager.registerShuffle(17, dependency) + assert(handle.isInstanceOf[CometNativeShuffleHandle[_, _]]) try { manager.getReader[Int, ColumnarBatch]( handle, @@ -331,7 +334,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { if (unregister) { assert(manager.unregisterShuffle(17)) assert(local.unregistered.contains(17)) - assert(backend.unregistered.contains(17)) + assert(backend.unregistered.isEmpty) } manager.stop() assert(local.stopped) @@ -355,9 +358,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { val api = new RecordingReaderApi val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val remote = manager .getReader[Int, ColumnarBatch]( @@ -401,9 +402,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { dependency, stageRerunEnabled = unavailableApi) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val failure = intercept[IllegalStateException] { manager.getReader[Int, ColumnarBatch]( @@ -438,9 +437,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.retryLimitFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val failure = intercept[IllegalStateException] { manager.getReader[Int, ColumnarBatch]( @@ -470,9 +467,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.generationFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val failure = intercept[FetchFailedException] { manager.getReader[Int, ColumnarBatch]( @@ -500,9 +495,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { api.generationFailure = expected val handle = new org.apache.spark.shuffle.celeborn.CelebornShuffleHandle(17, dependency) val manager = - new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) { - override protected[shuffle] def usesLocalShuffle(shuffleId: Int): Boolean = false - } + new CometCelebornShuffleManager(new SparkConf(false), false, (_, _) => backend, api) val failure = intercept[ClassCastException] { manager.getReader[Int, ColumnarBatch]( 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..990266cc65b 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,27 @@ 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) + } + } + 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") From f54d00bd0736f6a7fb13c800eee1e46711258c6b Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Fri, 4 Sep 2026 07:07:33 +0000 Subject: [PATCH 3/4] fix: preserve Celeborn recovery, concurrency, and statistics --- .github/workflows/pr_build_linux.yml | 2 + .github/workflows/pr_build_macos.yml | 2 + docs/source/user-guide/latest/configs.md | 5 +- docs/source/user-guide/latest/tuning.md | 7 +- .../CometCelebornShuffleMaterialization.scala | 92 ++++-- .../shuffle/CometNativeShuffleInputRDD.scala | 14 + .../shuffle/CometShuffleDependency.scala | 29 +- .../shuffle/CometShuffleExchangeExec.scala | 18 +- .../shuffle/CometShuffleOutputMetrics.scala | 69 +++++ .../shuffle/CometShuffledRowRDD.scala | 15 +- .../ShimCometShuffleMaterialization.scala | 44 +++ .../ShimCometShuffleMaterialization.scala | 47 ++++ .../ShimCometShuffleMaterialization.scala | 50 ++++ ...lebornConcurrentMaterializationSuite.scala | 265 ++++++++++++++++++ .../CometCelebornShuffleFallbackSuite.scala | 111 +++++++- .../CometCelebornShuffleStatisticsSuite.scala | 132 +++++++++ .../CometNativeShuffleInputRDDSuite.scala | 46 +++ 17 files changed, 900 insertions(+), 48 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleOutputMetrics.scala create mode 100644 spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala create mode 100644 spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala create mode 100644 spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometShuffleMaterialization.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleStatisticsSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index db061fbaa0d..645717e281f 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -375,6 +375,8 @@ jobs: 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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b9da5984d4d..673fac33a8d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -148,6 +148,8 @@ jobs: 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 diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md index 7b6001d8cac..49ee6dd48df 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -56,8 +56,9 @@ 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 identity, so late remote map results cannot replace local output. Subsequent fetch -failures retain Spark's normal recovery behavior. +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 diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 7d1a832c4c1..0ad8ec992c7 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -311,8 +311,11 @@ 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 identity, so late remote map results cannot overwrite or skip local map -output. All reads and retries for the replacement use local files and Spark's block transfer +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; 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 index 4a672f49878..eb62ff4e3a4 100644 --- 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 @@ -21,26 +21,31 @@ 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.{blocking, 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; its shuffle and stage IDs isolate all late remote - * task completions from the replacement. Once output is published, its destination is fixed. + * 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] { + extends FutureAction[MapOutputStatistics] + with ShimCometShuffleMaterialization { import CometCelebornShuffleMaterialization._ @@ -70,7 +75,7 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( throw failure } - private def withCapturedProperties[T](body: => T): T = { + private def withCapturedProperties[T](body: => T): T = withCapturedSession { val thread = Thread.currentThread() val previousProperties = sparkContext.getLocalProperties val previousClassLoader = thread.getContextClassLoader @@ -83,25 +88,52 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( } } - private def submit(dependency: CometShuffleDependency[K, V, C], expected: State): Boolean = { - try { - val action = synchronized { - if (state != expected) return false - // Keep submission and activeAction assignment atomic with cancellation and a size report. - // submitMapStage queues the job without waiting for executor tasks to finish. - val submitted = withCapturedProperties(sparkContext.submitMapStage(dependency)) - activeAction = Some(submitted) - if (expected == RunningRemote) remoteAction = Some(submitted) - actions += submitted - submitted + private def submit(dependency: CometShuffleDependency[K, V, C], expected: State): Unit = { + Future { + blocking { + try { + withCapturedProperties { + // submitMapStage eagerly resolves the input graph before queueing a job. An upstream + // Comet shuffle can still be choosing its destination, so resolve that graph on this + // worker, outside our lock. Independent branches can then be constructed and started, + // and cancellation can complete while we wait for an upstream materialization. + prepareInput(dependency.rdd) + val action = synchronized { + if (state != expected) { + None + } else { + if (expected == RunningLocal) remoteFailure.foreach(failure => throw failure) + // Dependencies and partitions are now cached. Keep job admission and action + // assignment atomic with cancellation and size-limit 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))( + ExecutionContext.global)) + } + } catch { + case NonFatal(failure) => + fail(expected, failure) + cancelActions(failure) + } + } + }(ExecutionContext.global) + } + + 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)) } - action.onComplete(result => finish(dependency, expected, result))(ExecutionContext.global) - true - } catch { - case NonFatal(failure) => - fail(expected, failure) - cancelActions(failure) - false } } @@ -128,13 +160,17 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( } else { result } - outcome match { + val published = outcome.map { statistics => + dependency.outputMetrics.foreach(_.publish(sparkContext)) + statistics + } + published match { case Success(_) => selected = Some(dependency) case _ => } state = Finished - completion.tryComplete(outcome) - outcome.failed.toOption.foreach(cancelActions) + completion.tryComplete(published) + published.failed.toOption.foreach(cancelActions) true } } @@ -194,8 +230,8 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( } else { state = RunningLocal try { - // Register the replacement before cancelling the old job. Its input partitions were - // already computed for the remote submission, and submitMapStage only queues execution. + // 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 = 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/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index 02bb1ae59f6..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 @@ -69,7 +70,8 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( val numParts: Int = 0, val rangePartitionBounds: Option[Seq[InternalRow]] = None, val nativeShuffleSpec: Option[NativeShuffleSpec] = None, - val useLocalShuffle: Boolean = false) + val useLocalShuffle: Boolean = false, + private[shuffle] val outputMetrics: Option[CometShuffleOutputMetrics] = None) extends ShuffleDependency[K, V, C]( _rdd, partitioner, @@ -102,25 +104,40 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( private[shuffle] def currentShuffleDependency: CometShuffleDependency[K, V, C] = Option(materializationInstance).flatMap(_.completedDependency).getOrElse(this) - private[shuffle] def createLocalShuffleDependency(): CometShuffleDependency[K, V, C] = + 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]( - rdd, + localRDD, partitioner, serializer, keyOrdering, aggregator, mapSideCombine, - shuffleWriterProcessor, + if (localOutputMetrics.nonEmpty) { + ShuffleExchangeExec.createShuffleWriteProcessor(localWriteMetrics) + } else { + shuffleWriterProcessor + }, shuffleType, schema, decodeTime, outputPartitioning, outputAttributes, - shuffleWriteMetrics, + localWriteMetrics, numParts, rangePartitionBounds, nativeShuffleSpec, - useLocalShuffle = true) + useLocalShuffle = true, + outputMetrics = localOutputMetrics) + } } /** Indicates shuffle type */ 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 e8800c16003..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 @@ -891,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 d738c552cd6..a8054209e79 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,12 +51,15 @@ class CometShuffledBatchRDD( SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY, SQLConf.get.fetchShuffleBlocksInBatch.toString) - // Materialize native Celeborn output before exposing a dependency to downstream stages. A - // size-limit fallback gets its own shuffle and stage IDs, so an obsolete remote completion - // cannot replace local output or tell Spark to skip a replacement map partition. - dependency = CometCelebornShuffleMaterialization.selectForRead(dependency) - - 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. + 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/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..53a87a6304a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornConcurrentMaterializationSuite.scala @@ -0,0 +1,265 @@ +/* + * 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.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) + } + } + } + + 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 + + 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/CometCelebornShuffleFallbackSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleFallbackSuite.scala index 842b4be4417..826ac7e86df 100644 --- 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 @@ -217,6 +217,73 @@ class CometCelebornShuffleFallbackSuite extends CometTestBase { } } + 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 @@ -411,6 +478,23 @@ private[shuffle] class PausedShuffleMapCompletion(val local: Boolean) { } } +/** 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) @@ -474,6 +558,8 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean @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 @@ -549,7 +635,30 @@ class CometCelebornFallbackTestShuffleManager(conf: SparkConf, isDriver: Boolean if (handle.isInstanceOf[CelebornShuffleHandle[_, _, _]] && context.partitionId() == 1) { delayed.foreach(_.awaitRemoteCommit()) } - val writer = super.getWriter[K, V](handle, mapId, context, metrics) + 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()) 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 990266cc65b..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 @@ -63,6 +63,52 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { 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) + } } } From f8d1aa08382340354aeed97c9964f1fe681d0250 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Fri, 4 Sep 2026 16:45:59 +0000 Subject: [PATCH 4/4] fix: avoid blocking shuffle materialization workers --- .../shuffle/CometCelebornShuffleManager.scala | 4 + .../CometCelebornShuffleMaterialization.scala | 55 ++++++++++---- .../shuffle/CometShuffledRowRDD.scala | 3 + ...lebornConcurrentMaterializationSuite.scala | 73 +++++++++++++++++++ 4 files changed, 120 insertions(+), 15 deletions(-) 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 39e73b3d158..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,6 +24,7 @@ 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 @@ -74,6 +75,9 @@ 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]() 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 index eb62ff4e3a4..f9f6770b407 100644 --- 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 @@ -23,7 +23,7 @@ import java.util.Properties import scala.collection.mutable import scala.collection.mutable.ArrayBuffer -import scala.concurrent.{blocking, CanAwait, ExecutionContext, Future, Promise} +import scala.concurrent.{CanAwait, ExecutionContext, Future, Promise} import scala.concurrent.duration.Duration import scala.util.{Failure, Success, Try} import scala.util.control.NonFatal @@ -50,6 +50,7 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( 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]() @@ -89,22 +90,27 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( } private def submit(dependency: CometShuffleDependency[K, V, C], expected: State): Unit = { - Future { - blocking { - try { + 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 { - // submitMapStage eagerly resolves the input graph before queueing a job. An upstream - // Comet shuffle can still be choosing its destination, so resolve that graph on this - // worker, outside our lock. Independent branches can then be constructed and started, - // and cancellation can complete while we wait for an upstream materialization. + // 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) - // Dependencies and partitions are now cached. Keep job admission and action - // assignment atomic with cancellation and size-limit reports. + // 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) @@ -112,17 +118,36 @@ private[shuffle] final class CometCelebornShuffleMaterialization[K, V, C]( Some(submitted) } } - action.foreach( - _.onComplete(result => finish(dependency, expected, result))( - ExecutionContext.global)) + action.foreach(_.onComplete(result => finish(dependency, expected, result))) } - } catch { + 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)) + } } - }(ExecutionContext.global) + } + upstream.toSeq } private def prepareInput(input: RDD[_]): Unit = { 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 a8054209e79..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 @@ -54,6 +54,9 @@ class CometShuffledBatchRDD( // 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[_]] = { 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 index 53a87a6304a..60132c984fd 100644 --- 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 @@ -31,6 +31,7 @@ import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriteMetricsReporter, Shu 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 @@ -174,6 +175,73 @@ class CometCelebornConcurrentMaterializationSuite extends CometTestBase { } } + 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 @@ -251,6 +319,11 @@ class CometCelebornConcurrentMaterializationTestManager(conf: SparkConf, isDrive 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,