diff --git a/ballista/core/src/error.rs b/ballista/core/src/error.rs index b5a7a82b02..33504fdedf 100644 --- a/ballista/core/src/error.rs +++ b/ballista/core/src/error.rs @@ -24,7 +24,9 @@ use std::{ }; use crate::serde::protobuf::failed_task::FailedReason; -use crate::serde::protobuf::{ExecutionError, FailedTask, FetchPartitionError, IoError}; +use crate::serde::protobuf::{ + ExecutionError, FailedTask, FetchPartitionError, IoError, TaskKilled, +}; use datafusion::error::DataFusionError; use datafusion::{arrow::error::ArrowError, sql::sqlparser::parser}; use futures::future::Aborted; @@ -202,29 +204,75 @@ impl Display for BallistaError { } } +struct FetchFailedDetails { + executor_id: String, + map_stage_id: usize, + map_partition_id: usize, + desc: String, +} + +/// Recovers a shuffle fetch failure that crossed DataFusion as +/// `ArrowError::ExternalError(FetchFailed)` and may now be wrapped in +/// `DataFusionError` layers. +fn find_fetch_failed(e: &BallistaError) -> Option { + match e { + BallistaError::FetchFailed(executor_id, map_stage_id, map_partition_id, desc) => { + Some(FetchFailedDetails { + executor_id: executor_id.clone(), + map_stage_id: *map_stage_id, + map_partition_id: *map_partition_id, + desc: desc.clone(), + }) + } + BallistaError::ArrowError(e) => fetch_failed_in_arrow(e), + BallistaError::DataFusionError(e) => fetch_failed_in_datafusion(e), + _ => None, + } +} + +fn fetch_failed_in_datafusion(e: &DataFusionError) -> Option { + match e.find_root() { + DataFusionError::ArrowError(e, _) => fetch_failed_in_arrow(e), + _ => None, + } +} + +fn fetch_failed_in_arrow(e: &ArrowError) -> Option { + let ArrowError::ExternalError(inner) = e else { + return None; + }; + if let Some(e) = inner.downcast_ref::() { + return find_fetch_failed(e); + } + if let Some(e) = inner.downcast_ref::() { + return fetch_failed_in_datafusion(e); + } + None +} + impl From for FailedTask { fn from(e: BallistaError) -> Self { + if let Some(fetch_failed) = find_fetch_failed(&e) { + return FailedTask { + error: fetch_failed.desc, + retryable: false, + count_to_failures: false, + failed_reason: Some(FailedReason::FetchPartitionError( + FetchPartitionError { + executor_id: fetch_failed.executor_id, + map_stage_id: fetch_failed.map_stage_id as u32, + map_partition_id: fetch_failed.map_partition_id as u32, + }, + )), + }; + } match e { - BallistaError::FetchFailed( - executor_id, - map_stage_id, - map_partition_id, - desc, - ) => { - FailedTask { - error: desc, - // fetch partition error is considered to be non-retryable - retryable: false, - count_to_failures: false, - failed_reason: Some(FailedReason::FetchPartitionError( - FetchPartitionError { - executor_id, - map_stage_id: map_stage_id as u32, - map_partition_id: map_partition_id as u32, - }, - )), - } - } + BallistaError::Cancelled => FailedTask { + error: "Task cancelled".to_string(), + retryable: true, + count_to_failures: false, + failed_reason: Some(FailedReason::TaskKilled(TaskKilled {})), + }, BallistaError::IoError(io) => { FailedTask { error: format!("Task failed due to Ballista IO error: {io:?}"), @@ -266,6 +314,20 @@ mod tests { FailedTask::from(e) } + fn fetch_failed( + executor_id: &str, + map_stage_id: usize, + map_partition_id: usize, + desc: &str, + ) -> BallistaError { + BallistaError::FetchFailed( + executor_id.to_owned(), + map_stage_id, + map_partition_id, + desc.to_owned(), + ) + } + #[test] fn bare_datafusion_io_error_is_retryable() { let e = BallistaError::DataFusionError(Box::new(DataFusionError::IoError( @@ -316,6 +378,17 @@ mod tests { )); } + #[test] + fn cancelled_task_is_retryable_without_counting_to_failures() { + let task = FailedTask::from(BallistaError::Cancelled); + assert!(task.retryable); + assert!(!task.count_to_failures); + assert!(matches!( + task.failed_reason, + Some(FailedReason::TaskKilled(_)) + )); + } + #[test] fn shared_wrapped_non_io_error_stays_non_retryable() { let inner = DataFusionError::Plan("bad plan".to_string()); @@ -328,4 +401,110 @@ mod tests { Some(FailedReason::ExecutionError(_)) )); } + + /// Builds the wrapped shape that can reach task failure classification. + fn wrap_in_arrow_external(inner: BallistaError) -> BallistaError { + BallistaError::DataFusionError(Box::new(datafusion_arrow_external(inner))) + } + + fn datafusion_arrow_external(inner: BallistaError) -> DataFusionError { + DataFusionError::ArrowError( + Box::new(ArrowError::ExternalError(Box::new(inner))), + None, + ) + } + + fn wrap_in_shared_arrow_external(inner: BallistaError) -> BallistaError { + let df = DataFusionError::Shared(Arc::new(datafusion_arrow_external(inner))); + BallistaError::DataFusionError(Box::new(df)) + } + + fn wrap_in_context_arrow_external(inner: BallistaError) -> BallistaError { + let df = datafusion_arrow_external(inner).context("reading shuffle partition"); + BallistaError::DataFusionError(Box::new(df)) + } + + fn assert_fetch_partition_error( + task: FailedTask, + executor_id: &str, + map_stage_id: u32, + map_partition_id: u32, + error: &str, + ) { + assert!(!task.retryable); + assert!(!task.count_to_failures); + assert_eq!(task.error, error); + match task.failed_reason { + Some(FailedReason::FetchPartitionError(fp)) => { + assert_eq!(fp.executor_id, executor_id); + assert_eq!(fp.map_stage_id, map_stage_id); + assert_eq!(fp.map_partition_id, map_partition_id); + } + other => panic!("expected FetchPartitionError, got {other:?}"), + } + } + + #[test] + fn bare_fetch_failed_maps_to_fetch_partition_error() { + let task = FailedTask::from(fetch_failed("exec-1", 3, 7, "boom")); + assert_fetch_partition_error(task, "exec-1", 3, 7, "boom"); + } + + #[test] + fn datafusion_arrow_external_fetch_failed_converts_to_bare_fetch_failed() { + let e = BallistaError::from(datafusion_arrow_external(fetch_failed( + "exec-1", + 3, + 7, + "connection reset", + ))); + + match e { + BallistaError::FetchFailed( + executor_id, + map_stage_id, + map_partition_id, + desc, + ) => { + assert_eq!(executor_id, "exec-1"); + assert_eq!(map_stage_id, 3); + assert_eq!(map_partition_id, 7); + assert_eq!(desc, "connection reset"); + } + other => panic!("expected bare FetchFailed, got {other:?}"), + } + } + + #[test] + fn wrapped_fetch_failed_is_recovered_as_fetch_partition_error() { + let e = wrap_in_arrow_external(fetch_failed("exec-1", 3, 7, "connection reset")); + let task = FailedTask::from(e); + assert_fetch_partition_error(task, "exec-1", 3, 7, "connection reset"); + } + + #[test] + fn shared_wrapped_fetch_failed_is_recovered() { + let e = + wrap_in_shared_arrow_external(fetch_failed("exec-2", 1, 2, "peer closed")); + let task = FailedTask::from(e); + assert_fetch_partition_error(task, "exec-2", 1, 2, "peer closed"); + } + + #[test] + fn context_wrapped_fetch_failed_is_recovered() { + let e = wrap_in_context_arrow_external(fetch_failed("exec-3", 5, 9, "timeout")); + let task = FailedTask::from(e); + assert_fetch_partition_error(task, "exec-3", 5, 9, "timeout"); + } + + #[test] + fn wrapped_non_fetch_error_stays_execution_error() { + let e = wrap_in_arrow_external(BallistaError::General("boom".to_string())); + let task = FailedTask::from(e); + assert!(!task.retryable); + assert!(matches!( + task.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } } diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 8854962d1b..23257e633a 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -575,7 +575,12 @@ impl ShuffleWriterExec { compression_type, ) .await - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + .map_err(|e| { + DataFusionError::ArrowError( + Box::new(ArrowError::ExternalError(Box::new(e))), + None, + ) + })?; let rows = stats.num_rows.unwrap_or(0) as usize; write_metrics.input_rows.add(rows); write_metrics.output_rows.add(rows); @@ -738,15 +743,12 @@ impl ExecutionPlan for ShuffleWriterExec { Ok(Box::pin(RecordBatchStreamAdapter::new( schema, futures::stream::once(async move { - let summaries = rx - .await - .map_err(|_| { - DataFusionError::Internal( - "ShuffleWriterExec coordinator dropped without sending" - .to_owned(), - ) - })? - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + let summaries = rx.await.map_err(|_| { + DataFusionError::Internal( + "ShuffleWriterExec coordinator dropped without sending" + .to_owned(), + ) + })??; summaries_to_batch( summaries, schema_captured, @@ -847,15 +849,13 @@ async fn run_coordinator( } } Err(e) => { - let msg = format!("{e:?}"); - for (i, slot) in senders.iter_mut().enumerate() { + // Share the original error with every output handoff so classification + // preserves FetchFailed/IO details regardless of stream completion order. + let shared = Arc::new(e); + for slot in senders.iter_mut() { if let Some(sender) = slot.take() { - let err_msg = if i == 0 { - msg.clone() - } else { - format!("shuffle writer failed: {msg}") - }; - let _ = sender.send(Err(DataFusionError::Execution(err_msg))); + let err = DataFusionError::from(&shared); + let _ = sender.send(Err(err)); } } } @@ -936,6 +936,8 @@ pub(crate) fn summaries_to_batch( #[allow(dead_code, unused_imports)] // clippy false positive with local imports mod tests { use super::*; + use crate::error::BallistaError; + use crate::execution_plans::ChaosExec; use datafusion::arrow::array::{StringArray, StructArray, UInt32Array, UInt64Array}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -975,6 +977,50 @@ mod tests { Ok(all) } + async fn drive_partition_results( + plan: Arc, + task_ctx: Arc, + ) -> Vec>> { + let k = plan.properties().output_partitioning().partition_count(); + let mut handles = Vec::with_capacity(k); + for n in 0..k { + let plan = plan.clone(); + let ctx = task_ctx.clone(); + handles.push(tokio::spawn(async move { + let mut stream = plan.execute(n, ctx).map_err(BallistaError::from)?; + utils::collect_stream(&mut stream).await + })); + } + let mut results = Vec::new(); + for h in handles { + results.push( + h.await + .expect("drive_partition_results task should not panic"), + ); + } + results + } + + fn assert_shared_structural_error(err: &BallistaError) { + let BallistaError::DataFusionError(e) = err else { + panic!("expected DataFusionError, got {err:?}"); + }; + assert!( + matches!(e.as_ref(), DataFusionError::Shared(_)), + "expected shared DataFusionError, got {e:?}" + ); + let DataFusionError::ArrowError(arrow, _) = e.find_root() else { + panic!("expected ArrowError root, got {:?}", e.find_root()); + }; + let ArrowError::ExternalError(inner) = arrow.as_ref() else { + panic!("expected Arrow external error, got {arrow:?}"); + }; + assert!( + inner.downcast_ref::().is_some(), + "expected Arrow external BallistaError, got {inner:?}" + ); + } + #[tokio::test] async fn test() -> Result<()> { let session_ctx = SessionContext::new(); @@ -1062,36 +1108,29 @@ mod tests { } #[tokio::test] - async fn test_no_repart_write_failure_propagates() -> Result<()> { + async fn write_failure_is_shared_to_all_output_partitions() -> Result<()> { let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); - // Place a directory at the data-{file_id}.arrow path so File::create - // fails inside write_stream_to_disk, not at create_dir_all. - // Path: work_dir / job_id / stage_id / partition / data-{file_id}.arrow - // task_slot defaults to 0 in try_new, so file_id is 0. - let tmp = tempfile::TempDir::new().unwrap(); - let data_arrow_dir = tmp - .path() - .join("jobOne") - .join("1") - .join("0") - .join("data-0.arrow"); - std::fs::create_dir_all(&data_arrow_dir).unwrap(); - let work_dir = tmp.path().to_str().unwrap().to_owned(); - - let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); + let input_plan: Arc = Arc::new(ChaosExec::new( + create_input_plan()?, + 1.0, + "transient", + Some(42), + )?); + let work_dir = TempDir::new()?; let query_stage = Arc::new(ShuffleWriterExec::try_new( "jobOne".into(), 1, input_plan, - work_dir, + work_dir.path().to_str().unwrap().to_owned(), )?); - let result = drive_all_partitions(query_stage, task_ctx).await; - assert!( - result.is_err(), - "expected File::create failure in write_stream_to_disk to propagate" - ); + let results = drive_partition_results(query_stage, task_ctx).await; + assert_eq!(2, results.len()); + for result in results { + let err = result.expect_err("expected injected write failure"); + assert_shared_structural_error(&err); + } Ok(()) } diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 7802bcb876..4d8405d196 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -958,15 +958,12 @@ impl ExecutionPlan for SortShuffleWriterExec { Ok(Box::pin(RecordBatchStreamAdapter::new( schema, futures::stream::once(async move { - let summaries = rx - .await - .map_err(|_| { - DataFusionError::Internal( - "SortShuffleWriterExec coordinator dropped without sending" - .to_owned(), - ) - })? - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + let summaries = rx.await.map_err(|_| { + DataFusionError::Internal( + "SortShuffleWriterExec coordinator dropped without sending" + .to_owned(), + ) + })??; summaries_to_batch( summaries, schema_captured, @@ -1049,7 +1046,7 @@ async fn run_coordinator( let mut grouped: Vec> = (0..k).map(|_| Vec::new()).collect(); - let mut first_error: Option = None; + let mut first_error: Option = None; for handle in handles { match handle.await { Ok(Ok(summaries)) => { @@ -1061,24 +1058,24 @@ async fn run_coordinator( } } Ok(Err(e)) => { - first_error.get_or_insert_with(|| format!("{e:?}")); + first_error.get_or_insert(e); } Err(join_err) => { - first_error - .get_or_insert_with(|| format!("write task panicked: {join_err}")); + first_error.get_or_insert_with(|| { + DataFusionError::Execution(format!("write task panicked: {join_err}")) + }); } } } - if let Some(msg) = first_error { - for (i, slot) in senders.iter_mut().enumerate() { + if let Some(e) = first_error { + // Share the original error with every output handoff so classification + // preserves FetchFailed/IO details regardless of stream completion order. + let shared = Arc::new(e); + for slot in senders.iter_mut() { if let Some(sender) = slot.take() { - let err_msg = if i == 0 { - msg.clone() - } else { - format!("sort shuffle writer failed: {msg}") - }; - let _ = sender.send(Err(DataFusionError::Execution(err_msg))); + let err = DataFusionError::from(&shared); + let _ = sender.send(Err(err)); } } return; @@ -1141,6 +1138,8 @@ fn compute_partition_indices( #[cfg(test)] mod tests { use super::*; + use crate::error::BallistaError; + use crate::execution_plans::ChaosExec; use datafusion::arrow::array::{StringArray, UInt32Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; @@ -1173,6 +1172,45 @@ mod tests { Ok(Arc::new(DataSourceExec::new(memory_data_source))) } + async fn drive_partition_results( + plan: Arc, + task_ctx: Arc, + ) -> Vec>> { + let k = plan.properties().output_partitioning().partition_count(); + let mut handles = Vec::with_capacity(k); + for n in 0..k { + let plan = plan.clone(); + let ctx = task_ctx.clone(); + handles.push(tokio::spawn(async move { + let mut stream = plan.execute(n, ctx).map_err(BallistaError::from)?; + crate::utils::collect_stream(&mut stream).await + })); + } + let mut results = Vec::new(); + for h in handles { + results.push( + h.await + .expect("drive_partition_results task should not panic"), + ); + } + results + } + + fn assert_shared_io_error(err: &BallistaError) { + let BallistaError::DataFusionError(e) = err else { + panic!("expected DataFusionError, got {err:?}"); + }; + assert!( + matches!(e.as_ref(), DataFusionError::Shared(_)), + "expected shared DataFusionError, got {e:?}" + ); + assert!( + matches!(e.find_root(), DataFusionError::IoError(_)), + "expected IO root, got {:?}", + e.find_root() + ); + } + #[test] fn compute_partition_indices_distributes_rows_by_hash() { use datafusion::arrow::array::{Int64Array, StringArray}; @@ -1298,6 +1336,38 @@ mod tests { Ok(()) } + #[tokio::test] + async fn write_failure_is_shared_to_all_output_partitions() -> Result<()> { + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + + let input_plan: Arc = Arc::new(ChaosExec::new( + create_test_input()?, + 1.0, + "transient", + Some(42), + )?); + let work_dir = TempDir::new()?; + + let writer = Arc::new(SortShuffleWriterExec::try_new( + "job1".into(), + 1, + input_plan, + work_dir.path().to_str().unwrap().to_string(), + Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2), + SortShuffleConfig::default(), + )?); + + let results = drive_partition_results(writer, task_ctx).await; + assert_eq!(2, results.len()); + for result in results { + let err = result.expect_err("expected injected write failure"); + assert_shared_io_error(&err); + } + + Ok(()) + } + /// Shared helper for round-trip tests. Builds `num_batches` batches of /// `rows_per_batch` rows each with schema `(k: Int64, v: Int64)`, writes /// them through `SortShuffleWriterExec` into `num_partitions` output diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs index dce01ae7f5..b491cfa8a6 100644 --- a/ballista/executor/src/executor.rs +++ b/ballista/executor/src/executor.rs @@ -236,9 +236,7 @@ impl Executor { self.abort_handles.insert(key.clone(), abort_handle); let partitions = match std::panic::AssertUnwindSafe(task).catch_unwind().await { - Ok(Ok(result)) => { - result.map_err(|e| BallistaError::DataFusionError(Box::new(e))) - } + Ok(Ok(result)) => result.map_err(BallistaError::from), Ok(Err(_)) => { warn!("Task has been aborted!"); Err(BallistaError::Cancelled) diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 2ccf5a9daf..3086071ec4 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -355,6 +355,11 @@ impl { warn!("Fail to cancel running tasks due to {e:?}"); } + if self.state.config.is_push_staged_scheduling() { + event_sender + .post_event(QueryStageSchedulerEvent::ReviveOffers) + .await?; + } } Err(e) => { let msg = format!( diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 539b45c6b8..b6ec54866b 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -18,6 +18,7 @@ use crate::planner::create_shuffle_writer_with_config; use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use crate::state::aqe::planner::AdaptiveStageInfo; +use crate::state::execution_graph::StageOutput; use ballista_core::JobId; use ballista_core::execution_plans::{ PerPartitionFilterExec, ShuffleReaderExec, range_partition_predicates, @@ -31,11 +32,12 @@ use datafusion::{ physical_plan::ExecutionPlan, }; use log::debug; +use std::collections::HashMap; use std::sync::Arc; #[derive(Debug, Clone, Default)] pub(crate) struct BallistaAdapter { - inputs: Vec, + inputs: HashMap, } /// @@ -61,7 +63,12 @@ impl BallistaAdapter { "stage ID has to be generated at this point".to_string(), ) })?; - self.inputs.push(stage_id); + let mut stage_output = StageOutput::new(); + for partition in partitions.iter().flatten().cloned() { + stage_output.add_partition(partition); + } + stage_output.complete = true; + self.inputs.insert(stage_id, stage_output); let partitioning = exchange.properties().partitioning.clone(); let reader = match (exchange.coalesce(), exchange.broadcast) { diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 0c133db640..dac8482214 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -19,7 +19,7 @@ use crate::display::print_stage_metrics; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; use crate::state::aqe::execution_plan::RangeRepartitionRouting; -use crate::state::aqe::planner::AdaptivePlanner; +use crate::state::aqe::planner::{AdaptivePlanner, AdaptiveStageInfo}; use crate::state::execution_graph::{ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ResolvedStage, RunningTaskInfo, StageOutput, @@ -29,7 +29,7 @@ use crate::state::task_manager::UpdatedStages; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{ - ShuffleWriter, cut_partitions, merge_runtime_stats_reports, repartition_routing_expr, + cut_partitions, merge_runtime_stats_reports, repartition_routing_expr, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; @@ -167,7 +167,7 @@ impl AdaptiveExecutionGraph { let stages: ballista_core::error::Result> = runnable .into_iter() - .map(|s| Self::create_resolved_stage(session_config.clone(), s.plan)) + .map(|s| Self::create_resolved_stage(session_config.clone(), s)) .collect(); let stages = stages?; @@ -202,15 +202,15 @@ impl AdaptiveExecutionGraph { impl AdaptiveExecutionGraph { fn create_resolved_stage( session_config: Arc, - stage: Arc, + stage: AdaptiveStageInfo, ) -> ballista_core::error::Result<(usize, ExecutionStage)> { - let stage_id = stage.stage_id(); + let stage_id = stage.plan.stage_id(); let stage = ExecutionStage::Resolved(ResolvedStage::new( stage_id, 0, - stage, - vec![], // we do not know output links at this moment - HashMap::new(), // we do not keep inputs at the moment + stage.plan, + vec![], // we do not know output links at this moment + stage.inputs, HashSet::new(), session_config, )); @@ -399,7 +399,7 @@ impl AdaptiveExecutionGraph { if !self.stages.contains_key(&stage.plan.stage_id()) { let (stage_id, stage) = Self::create_resolved_stage( self.session_config.clone(), - stage.plan, + stage, )?; self.stages.insert(stage_id, stage); } @@ -525,6 +525,11 @@ impl AdaptiveExecutionGraph { } }); + for stage_id in &reset_running_stage { + self.planner + .remove_exchange_locations(*stage_id, executor_id); + } + // check and reset the successful stages if !resubmit_inputs.is_empty() { self.stages diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 690a5f2bcd..6dc08dab6a 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -286,6 +286,17 @@ impl AdaptivePlanner { } } + pub fn remove_exchange_locations(&mut self, stage_id: usize, executor_id: &str) { + if let Some(stage_output) = self.runnable_stage_output.get_mut(&stage_id) { + stage_output + .partition_locations + .iter_mut() + .for_each(|(_partition, locs)| { + locs.retain(|loc| loc.executor_meta.id != executor_id); + }); + } + } + /// Once all tasks have completed, pop the accumulated stage output as a /// K-shaped `Vec>` (or the broadcast-shape variant) /// *without* parking it on the ExchangeExec. Caller can post-process @@ -650,6 +661,5 @@ impl AdaptivePlanner { /// Wraps stage plan with addition of references to previous stages pub(crate) struct AdaptiveStageInfo { pub(crate) plan: Arc, - #[allow(dead_code)] // TODO: still not sure if this is needed - pub(crate) inputs: Vec, + pub(crate) inputs: std::collections::HashMap, } diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 49f3cab627..99b7c38d5d 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -1819,11 +1819,12 @@ mod test { use std::sync::Arc; use crate::scheduler_server::event::QueryStageSchedulerEvent; - use ballista_core::error::Result; + use ballista_core::error::{BallistaError, Result}; use ballista_core::serde::protobuf::{ self, ExecutionError, FailedTask, FetchPartitionError, IoError, JobStatus, TaskKilled, failed_task, job_status, task_status, }; + use datafusion::arrow::error::ArrowError; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::TaskContext; use datafusion::physical_plan::{ @@ -2572,23 +2573,13 @@ mod test { let task1 = agg_graph.pop_next_task(&executor2.id)?.unwrap(); let task_status1 = mock_completed_task(task1, &executor2.id); - // 2nd task in the Stage 2, failed due to FetchPartitionError let task2 = agg_graph.pop_next_task(&executor2.id)?.unwrap(); - let task_status2 = mock_failed_task( - task2, - FailedTask { - error: "FetchPartitionError".to_string(), - retryable: false, - count_to_failures: false, - failed_reason: Some(failed_task::FailedReason::FetchPartitionError( - FetchPartitionError { - executor_id: executor1.id.clone(), - map_stage_id: 1, - map_partition_id: 0, - }, - )), - }, - ); + let failed_task = wrapped_fetch_failed_task(&executor1.id, 1, 0); + assert!(matches!( + failed_task.failed_reason, + Some(failed_task::FailedReason::FetchPartitionError(_)) + )); + let task_status2 = mock_failed_task(task2, failed_task); let mut running_task_count = 0; while let Some(_task) = agg_graph.pop_next_task(&executor2.id)? { @@ -3227,6 +3218,25 @@ mod test { // todo!() // } + fn wrapped_fetch_failed_task( + executor_id: &str, + map_stage_id: usize, + map_partition_id: usize, + ) -> FailedTask { + let err = BallistaError::DataFusionError(Box::new(DataFusionError::ArrowError( + Box::new(ArrowError::ExternalError(Box::new( + BallistaError::FetchFailed( + executor_id.to_owned(), + map_stage_id, + map_partition_id, + "FetchPartitionError".to_owned(), + ), + ))), + None, + ))); + FailedTask::from(err) + } + fn drain_tasks(graph: &mut dyn ExecutionGraph) -> Result<()> { let executor = mock_executor("executor-id1".to_string()); while let Some(task) = graph.pop_next_task(&executor.id)? { diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index dea5af1004..604b0a31e8 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -1338,11 +1338,23 @@ mod tests { let manager_for_abort = manager.clone(); let job_id_for_abort = job_id.clone(); let abort = tokio::spawn(async move { + let manager_for_cancel = manager_for_abort.clone(); + let job_id_for_cancel = job_id_for_abort.clone(); manager_for_abort .abort_job( &job_id_for_abort, "test failure".to_string(), move |tasks| async move { + let status = manager_for_cancel + .get_job_status(&job_id_for_cancel) + .await? + .expect( + "aborted job should remain visible during cancellation", + ); + assert!( + matches!(status.status, Some(Status::Failed(_))), + "job must be terminal before executor tasks are cancelled" + ); cancelled_tasks_for_abort.store(tasks.len(), Ordering::SeqCst); Ok(()) }, diff --git a/chaos-testing/README.md b/chaos-testing/README.md index 88b198b571..b4afbfc634 100644 --- a/chaos-testing/README.md +++ b/chaos-testing/README.md @@ -26,10 +26,11 @@ injects faults into real queries, to exercise Ballista's high-availability **This is a bug-hunting harness, not a regression suite in the usual sense.** Its job is to surface real defects in Ballista's HA behavior. Where it finds one, the corresponding test reproduces the bug rather than working around it. -Such a test is marked `#[ignore]` with the issue it reproduces, so that it -does not hold CI red on a bug it did not introduce, and is un-ignored — not -rewritten — when that issue is fixed, at which point it becomes the regression -test for the fix. Run them with `cargo test -p ballista-chaos -- --ignored`. +Such a test is marked `#[ignore]` with the issue or follow-up path it +reproduces, so that it does not hold CI red on a bug it did not introduce, and +is un-ignored — not rewritten — when that issue is fixed, at which point it +becomes the regression test for the fix. Run them with +`cargo test -p ballista-chaos -- --ignored`. See [Findings](#findings) below for the confirmed bugs this harness has found so far, each with the test that reproduces it. @@ -133,8 +134,8 @@ non-null without invoking the (volatile) UDF. Two regression tests in ## How to run ```sh -cargo test -p ballista-chaos # everything except the known-bug scenarios -cargo test -p ballista-chaos -- --ignored # the known-bug scenarios; these fail, on purpose +cargo test -p ballista-chaos # active regression scenarios +cargo test -p ballista-chaos -- --ignored # any currently ignored known-bug scenarios ``` Every test that spawns a cluster does so through `TestCluster`, which holds @@ -182,19 +183,17 @@ sanity check that every other scenario's assertions depend on. | Scenario | Test | What it does | Expected result | | -------- | ------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| A | `retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}` | Injects one retryable IO fault (budget 1); the retry must succeed and match baseline. | **Ignored (both), reproduces the error flattening tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 2.** | +| A | `retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}` | Injects one retryable IO fault (budget 1); the retry must succeed and match baseline. | Pass (both). | | B | `exhausted_retries_fail_the_job_and_leave_the_cluster_healthy` | Injects an inexhaustible IO fault (budget 99 ≫ `task_max_failures`); job must fail, cluster must stay usable after. | Pass (both). | | C | `panicking_task_fails_the_job_but_the_executor_survives` | Injects a task panic; job must fail non-retryably, both executor processes must survive, cluster must stay usable after. | Pass (both). | -| D | `executor_killed_mid_stage_is_recovered` | SIGKILLs an executor while its tasks are genuinely running (held open by `chaos_delay`); scheduler must reschedule onto the survivor and return the correct result. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | -| E | `executor_killed_after_shuffle_write_is_recovered` | SIGKILLs the map-side executor _after_ it wrote shuffle output, with a long executor timeout to bias toward the fetch-failure path rather than heartbeat expiry; downstream stage must re-run the map stage. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | +| D | `executor_killed_mid_stage_is_recovered` | SIGKILLs an executor while its tasks are genuinely running (held open by `chaos_delay`); scheduler must reschedule onto the survivor and return the correct result. | Pass (both), after stale task-attempt cancellation is classified as retryable cleanup. | +| E | `executor_killed_after_shuffle_write_is_recovered` | SIGKILLs the map-side executor _after_ it wrote shuffle output, with a long executor timeout to bias toward the fetch-failure path rather than heartbeat expiry; downstream stage must re-run the map stage. | Pass (both). | | F | `restarted_executor_rejoins_and_serves_queries` | Kills an executor, waits for the scheduler to reap it, restarts it, asserts the registered count returns to 2 and the cluster still serves the baseline query. | Pass (both), after the race fix in this crate (see below). | | G | `killing_every_executor_terminates_the_job` | SIGKILLs every executor mid-query; asserts the job fails with an error naming the executor loss rather than hanging. | Regression test for [#2029](https://github.com/apache/datafusion-ballista/issues/2029) — Finding 3. | -An ignored scenario above is not a defect in this harness, and its assertions -have not been weakened to make it pass — it reproduces a real Ballista bug and -is ignored only so that CI is not red on a bug this crate did not introduce. -Run the ignored scenarios with `-- --ignored --test-threads=1` to see the -failures; see [Findings](#findings) for what each one proves. +If a future scenario is ignored, it should stay tied to its tracking issue or +follow-up path, and its assertions should keep reproducing the underlying bug +rather than being weakened to pass. ### A note on Scenario F: the harness race that was fixed here @@ -216,78 +215,55 @@ actually promises. ## Findings -Four scenarios above fail because they have found real bugs in Ballista, not -because the harness is broken. Each is `#[ignore]`d against the issue it -reproduces so CI stays green on a tree whose bugs predate this crate, and each -keeps its original assertions: nothing is relaxed to manufacture a pass. When -the issue is fixed, delete the `#[ignore]` — the scenario is then the -regression test for it. One of the findings (#2028) has already gone through -that cycle in reverse: it was fixed on main, but a concurrent refactor -re-broke the same scenario through a different mechanism — see Finding 2. +The findings below are bugs this harness exposed. Fixed findings remain here +as context for the active regression scenarios; unfixed findings stay ignored +against their tracking issue or follow-up path so CI is not red on a known bug. ### Finding 1 — Shuffle-fetch failures lose their type, so the map-stage resubmit never fires Tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027). -**Proven by:** Scenario D (`executor_killed_mid_stage_is_recovered`) and -Scenario E (`executor_killed_after_shuffle_write_is_recovered`), both AQE -settings, `#[ignore]`d against that issue. - -Scenarios D and E are **races**, not deterministic reproducers, and this is -the one place in the crate where that is true. Killing an executor can be -noticed by the scheduler in either of two ways, and they are in a footrace: if -the heartbeat expires first, the `ExecutorLost` path recovers the job -correctly and the scenario passes in a few seconds; if a downstream task tries -to fetch shuffle output from the dead executor first, the bug below bites and -the job fails (or hangs until the scenario's timeout). Scenario D failed -locally on two of three runs; Scenario E passed locally but failed in CI -(`aqe_off`, with the reduce stage reporting a `FetchFailed` flattened inside -`DataFusionError::Shared` — see the log excerpt in #2027). The CI failure also -settles which way `executor_timeout_seconds` biases the race: raising it to -60s (as Scenario E does) _delays_ heartbeat expiry, so the downstream fetch -hits the dead executor first and the scenario exercises the broken -fetch-failure path; it only passes when the kill happens to land after the -reduce tasks already fetched their input. Un-ignoring these scenarios once -#2027 is fixed therefore also means pinning which of the two paths each -exercises — `executor_timeout_seconds` is the knob that decides the race, and -Scenario D currently leaves it at the harness default — otherwise they will be -flaky regression tests. - -The shuffle reader (`ballista/core/src/execution_plans/shuffle_reader.rs`) -correctly produces a typed `BallistaError::FetchFailed(executor_id, -map_stage_id, map_partition_id, desc)` when it cannot reach a dead executor, -and `ballista/core/src/error.rs`'s `impl From for FailedTask` -has a dedicated arm for exactly that variant (around line 205) which produces -`FailedReason::FetchPartitionError` — the signal -`ballista/scheduler/src/state/execution_graph.rs` (around line 826) uses to -resubmit the lost map stage rather than simply failing the job. - -The type does not survive to that point, however. Two real, non-test code -paths erase it before the executor reports its `TaskStatus`: - -- `ballista/executor/src/executor.rs:238-239`, in - `Executor::execute_query_stage`, converts the stage's result with - `result.map_err(|e| BallistaError::DataFusionError(Box::new(e)))` instead of - `BallistaError::from(e)` / `e.into()`. That bypasses the very unwrapping - logic `error.rs`'s `impl From for BallistaError` exists to - provide (`DataFusionError::ArrowError(e, _) => Self::from(*e)`, which would - otherwise recover a `FetchFailed` wrapped inside an `ArrowError::ExternalError`). -- `ballista/core/src/execution_plans/shuffle_writer.rs:245`, in - `ShuffleWriterExec`'s unpartitioned write branch, Debug-formats a - `BallistaError` into an opaque `DataFusionError::Execution(format!("{e:?}"))` - — a conversion that can never be undone by any later `.into()`, because the - original variant no longer exists, only its printed form. - -Either path leaves the executor reporting something like -`BallistaError::DataFusionError(Execution("FetchFailed(\"\", ..., -\"...Connection refused...\")"))` — the `FetchFailed` information is present -only as inert text inside a string. `error.rs`'s `FetchFailed` arm cannot match -a `DataFusionError::Execution`, so the task falls to the catch-all arm -(around line 248) and is marked `retryable: false`, `FailedReason::ExecutionError`. - -**Net effect:** when an executor dies after producing shuffle output that a -downstream stage still needs, Ballista fails the whole query instead of -re-running the map stage that produced it. +**Regression coverage:** Scenario D +(`executor_killed_mid_stage_is_recovered`) and Scenario E +(`executor_killed_after_shuffle_write_is_recovered`), both AQE settings. + +The shuffle reader produces a typed `BallistaError::FetchFailed(executor_id, +map_stage_id, map_partition_id, desc)` when it cannot reach a dead executor. +That type has to survive until task-failure classification, because the +scheduler's map-stage resubmit path is keyed on `FetchPartitionError`. + +Before the fix, production code could turn that structured error into inert +text inside `DataFusionError::Execution`. Once that happened, the classifier +saw only a generic execution error and the scheduler never received the +`FetchPartitionError` signal. + +The fix keeps the error structural across both places where it was being lost: +the executor now uses the existing `BallistaError` conversion when a stage +fails, and the shuffle-writer coordinator no longer Debug-formats the first +child error before handing it back through the output stream. The classifier +also looks through the DataFusion/Arrow wrapper stack so a wrapped +`FetchFailed` still becomes `FetchPartitionError`. + +Scenario E is the direct fetch-failure regression: it kills the map-side +executor after shuffle output is written and uses a long executor timeout so a +downstream fetch is likely to hit the dead executor before heartbeat expiry. +Scenario D covers the adjacent executor-loss race while a stage is still +running. In that path, heartbeat expiry can win first, so executor-loss task +resets need to wake push scheduling with fresh offers. + +### Scenario D note — Mid-stage executor loss can cancel stale task attempts + +Scenario D also exposed an adjacent recovery edge. When executor-loss recovery +rolls back or resets stages, the scheduler can ask surviving executors to +cancel tasks from stale stage attempts. Those aborted task futures are reported +back as `BallistaError::Cancelled`; this is task attempt cleanup, not a +client-cancelled job. + +The fix maps executor-reported `Cancelled` task attempts to retryable, +non-counting `TaskKilled` failures, so stale work cleanup does not fail a job +that executor-loss recovery is already rescheduling. Explicit job cancellation +is still handled by the scheduler by marking the job terminal before cancelling +running executor tasks. ### Finding 2 — Retryable IO errors are misclassified because the shuffle writer flattens them @@ -296,9 +272,8 @@ Originally tracked by fixed); the surviving flattening mechanism is the one [#2027](https://github.com/apache/datafusion-ballista/issues/2027) tracks. -**Proven by:** Scenario A, both cases -(`retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}`), -`#[ignore]`d against #2027. +**Regression coverage:** Scenario A, both cases +(`retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}`). The history matters here because the failure mode moved underneath the harness. As first found, only the `aqe_on` case failed: an `IoError` raised on @@ -317,13 +292,9 @@ fault now reaches the classifier as `Execution("IoError(Custom { .. })")` (`aqe_off`) or `Execution("Shared(IoError(Custom { .. }))")` (`aqe_on`) — the variant exists only as printed text, so `find_root()` has nothing to unwrap and the task is -marked non-retryable. This is the same type-erasing conversion Finding 1 -describes for `FetchFailed`, which is why both scenarios are ignored against -#2027 rather than the fixed #2028. - -**Net effect:** any genuine transient IO error inside a stage that writes -shuffle output — which is every non-final stage — is classified non-retryable, -turning what should be a retried task into an immediate job failure. +marked non-retryable. Keeping the original `DataFusionError` through the +shuffle-writer handoff fixes this by giving the existing `find_root()` based +classifier the real IO error again. ### Finding 3 — Killing every executor hung the job instead of failing it (fixed) @@ -355,12 +326,10 @@ the same `ExecutorLost` event; the unit test `test_running_job_fails_when_launch_failure_loses_last_executor` covers the launch-failure path deterministically. -### For comparison: the heartbeat-expiry path does recover +### For comparison: the heartbeat-expiry path also recovers -Ballista's HA recovery is not uniformly broken. Whenever the kill in Scenario -D or E happens to be noticed by heartbeat expiry (`ExecutorLost`) before any -downstream task fetches from the dead executor, the job recovers and matches -the baseline — that is why both scenarios pass on some runs. What Finding 1 -breaks is specifically the fetch-failure path, and both scenarios are ignored -against #2027 because whether a given run exercises that path is a timing -race, not something the harness currently controls. +Killing an executor can be noticed in two ways: heartbeat expiry +(`ExecutorLost`) or a downstream shuffle fetch from the dead executor. Scenario +E biases toward the fetch-failure path. Scenario D exercises the broader +mid-stage executor-loss path, including stale task-attempt cancellation during +executor-loss recovery. Both paths now recover and return the baseline result. diff --git a/chaos-testing/src/bin/chaos-executor.rs b/chaos-testing/src/bin/chaos-executor.rs index a2efe82a72..830d295ae3 100644 --- a/chaos-testing/src/bin/chaos-executor.rs +++ b/chaos-testing/src/bin/chaos-executor.rs @@ -43,6 +43,7 @@ async fn main() -> ballista_core::error::Result<()> { grpc_port: env_u16("CHAOS_EXECUTOR_GRPC_PORT"), scheduler_host: "127.0.0.1".to_string(), scheduler_port: env_u16("CHAOS_SCHEDULER_PORT"), + scheduler_connect_timeout_seconds: 10, vcores: std::env::var("CHAOS_CONCURRENT_TASKS") .ok() .and_then(|v| v.parse().ok()) diff --git a/chaos-testing/src/cluster.rs b/chaos-testing/src/cluster.rs index 33fc8d040f..d1972cd578 100644 --- a/chaos-testing/src/cluster.rs +++ b/chaos-testing/src/cluster.rs @@ -69,19 +69,13 @@ fn binary(name: &str) -> PathBuf { /// One supervised executor process. /// -/// `child` is used by this task's `kill_executor`/`executor_is_alive`. `port`, -/// `grpc_port`, and `work_dir` are still not read anywhere yet — no code in -/// this task needed to target an executor by its network address or inspect -/// its working directory — so they keep a narrower `#[allow(dead_code)]` than -/// the struct-wide one Task 5 left; a later scenario that needs to address a -/// specific executor's port or inspect its shuffle files can drop it then. +/// `child` is used by this task's `kill_executor`/`executor_is_alive`. pub(crate) struct ExecutorHandle { pub(crate) child: Child, #[allow(dead_code)] pub(crate) port: u16, #[allow(dead_code)] pub(crate) grpc_port: u16, - #[allow(dead_code)] pub(crate) work_dir: PathBuf, } @@ -139,6 +133,11 @@ impl TestClusterBuilder { self } + pub fn concurrent_tasks(mut self, n: usize) -> Self { + self.concurrent_tasks = n; + self + } + pub async fn start(self) -> Result { // Held for the cluster's whole lifetime, so only one cluster exists in // this process at a time. `--test-threads=1` gives the same guarantee, @@ -414,7 +413,7 @@ impl TestCluster { } /// The last lines of every child process log, for timeout diagnostics. - fn log_tails(&self) -> String { + pub fn log_tails(&self) -> String { let mut out = String::new(); let mut paths: Vec = std::fs::read_dir(&self.log_dir) .map(|d| d.filter_map(|e| e.ok().map(|e| e.path())).collect()) @@ -436,6 +435,16 @@ impl TestCluster { out } + pub async fn diagnostics(&self, job_id: &str) -> String { + let stages = self + .stages(job_id) + .await + .ok() + .and_then(|stages| serde_json::to_string_pretty(&stages).ok()) + .unwrap_or_else(|| "stage summary unavailable".to_string()); + format!("--- stages ---\n{stages}\n{}", self.log_tails()) + } + /// Block until the scheduler considers exactly `n` executors registered. /// /// Unlike `await_executors` (which waits for *at least* `n`, the right @@ -657,6 +666,50 @@ impl TestCluster { Ok(()) } + pub async fn await_successful_shuffle_output( + &self, + job_id: &str, + ) -> Result<(usize, usize), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + for stage_id in running_stage_ids_with_successful_tasks(&stages) { + if let Some(executor_index) = + self.executor_with_shuffle_output(job_id, stage_id) + { + return Ok((executor_index, stage_id)); + } + } + if self.job_status(job_id).await.unwrap_or_default() == "Completed" { + return Err(format!( + "job {job_id} completed before a running shuffle-writing stage could be targeted\n{}", + self.log_tails() + )); + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for successful shuffle output for job {job_id}\n{}", + self.log_tails() + )); + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + fn executor_with_shuffle_output( + &self, + job_id: &str, + stage_id: usize, + ) -> Option { + self.executors + .iter() + .enumerate() + .find_map(|(index, executor)| { + let stage_dir = executor.work_dir.join(job_id).join(stage_id.to_string()); + dir_has_entries(&stage_dir).then_some(index) + }) + } + /// Start a fresh executor process in the given slot and wait for it to register. pub async fn restart_executor(&mut self, index: usize) -> Result<(), String> { let expected = self.executors.len(); @@ -677,6 +730,34 @@ fn find_stage(stages: &serde_json::Value, stage_id: usize) -> Option<&serde_json }) } +fn dir_has_entries(path: &Path) -> bool { + std::fs::read_dir(path) + .map(|mut entries| entries.next().is_some()) + .unwrap_or(false) +} + +fn running_stage_ids_with_successful_tasks(stages: &serde_json::Value) -> Vec { + stages + .get("stages") + .and_then(|s| s.as_array()) + .into_iter() + .flatten() + .filter(|stage| { + stage.get("stage_status").and_then(|v| v.as_str()) == Some("Running") + && stage + .get("tasks") + .and_then(|tasks| tasks.as_array()) + .into_iter() + .flatten() + .any(|task| { + task.get("status").and_then(|status| status.as_str()) + == Some("Successful") + }) + }) + .filter_map(|stage| stage.get("stage_id")?.as_str()?.parse().ok()) + .collect() +} + /// If `child` already exited with a non-zero status, log the path of its /// output so a human investigating a failed scenario knows where to look. /// Then make sure it is actually gone. @@ -712,6 +793,11 @@ mod tests { async fn cluster_starts_with_the_requested_executors_registered() { let cluster = TestCluster::builder() .executors(2) + // This canary only checks that registration happened; it kills + // nothing, so the default short reap timeout buys it nothing and + // lets a CPU-starved executor be reaped mid-startup on a loaded CI + // runner before the snapshot below. + .executor_timeout_seconds(30) .start() .await .expect("cluster must start"); diff --git a/chaos-testing/src/fixture.rs b/chaos-testing/src/fixture.rs index 3871fce4ef..f8c6d8784d 100644 --- a/chaos-testing/src/fixture.rs +++ b/chaos-testing/src/fixture.rs @@ -128,6 +128,12 @@ impl Fixture { GROUP BY d.name ORDER BY d.name" } + pub fn shuffle_loss_query() -> &'static str { + "SELECT f.value, d.name \ + FROM facts f JOIN dims d ON f.key = d.key \ + ORDER BY f.value" + } + /// The same query with a chaos UDF spliced into the WHERE clause. /// /// `injection` is a complete `chaos_*(...)` call returning BOOLEAN. Its own @@ -155,6 +161,15 @@ impl Fixture { GROUP BY d.name ORDER BY d.name" ) } + + pub fn shuffle_loss_chaos_query(injection: &str) -> String { + format!( + "SELECT f.value, d.name \ + FROM facts f JOIN dims d ON f.key = d.key \ + WHERE {injection} IS NOT NULL \ + ORDER BY f.value" + ) + } } #[cfg(test)] diff --git a/chaos-testing/tests/common/mod.rs b/chaos-testing/tests/common/mod.rs index 1a2755d3d2..831be3d166 100644 --- a/chaos-testing/tests/common/mod.rs +++ b/chaos-testing/tests/common/mod.rs @@ -43,12 +43,23 @@ impl ChaosRun { aqe: bool, executors: usize, executor_timeout_seconds: u64, + ) -> Self { + Self::start_with_concurrent_tasks(aqe, executors, executor_timeout_seconds, 4) + .await + } + + pub async fn start_with_concurrent_tasks( + aqe: bool, + executors: usize, + executor_timeout_seconds: u64, + concurrent_tasks: usize, ) -> Self { let _ = env_logger::builder().is_test(true).try_init(); let cluster = TestCluster::builder() .executors(executors) .executor_timeout_seconds(executor_timeout_seconds) + .concurrent_tasks(concurrent_tasks) .start() .await .expect("cluster must start"); @@ -105,17 +116,15 @@ impl ChaosRun { /// The expected answer, computed by plain local DataFusion. pub async fn local_baseline(&self) -> String { + self.local_sql(Fixture::baseline_query()).await + } + + pub async fn local_sql(&self, sql: &str) -> String { let ctx = SessionContext::new(); for stmt in self.fixture.register_sql() { ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); } - let batches = ctx - .sql(Fixture::baseline_query()) - .await - .unwrap() - .collect() - .await - .unwrap(); + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); pretty_format_batches(&batches).unwrap().to_string() } diff --git a/chaos-testing/tests/ha.rs b/chaos-testing/tests/ha.rs index fe9f93d11f..4b5980aae8 100644 --- a/chaos-testing/tests/ha.rs +++ b/chaos-testing/tests/ha.rs @@ -22,9 +22,9 @@ //! statistics, so a re-run map stage can come back with a different plan than the //! one whose output was lost. //! -//! Every test in this file spawns a whole multi-process cluster, so this file -//! must always be run with `--test-threads=1` or ports and CPU will be -//! exhausted by concurrent clusters. +//! Every test in this file spawns a whole multi-process cluster. `TestCluster` +//! serializes those clusters, so `--test-threads=1` is useful for readable +//! local output but is not required for correctness. mod common; @@ -66,37 +66,12 @@ use chaos_testing::fixture::Fixture; /// result still equals the baseline: a retried stage is exactly where duplicated /// or dropped partitions would show up. /// -/// This scenario is split into one test per AQE setting rather than being an -/// `rstest` over both, because historically only the AQE-on case failed and -/// `rstest` cannot ignore an individual case. -/// -/// Both cases originally passed or reproduced #2028 (the AQE-on case: -/// `Shared(IoError)` misses the classifier's shallow match). #2028 was fixed -/// by classifying on `find_root()` (#2119), but the sort-shuffle writer -/// refactor (#2038/#2106) then made both cases fail the same way: the shuffle -/// write coordinator flattens any task error into -/// `DataFusionError::Execution(format!("{e:?}"))` -/// (`ballista/core/src/execution_plans/shuffle_writer.rs`, error arm of the -/// coordinator fan-out), so the injected `IoError` reaches the classifier as -/// inert text inside an `Execution` string, `find_root()` has nothing to see -/// through, and the task is marked non-retryable. That is the same -/// type-erasing mechanism tracked for fetch failures by #2027. -/// -/// Ignored, not deleted or weakened. Un-ignore both as the regression tests -/// when #2027's error-flattening is fixed; see chaos-testing/README.md, -/// Finding 2. #[tokio::test] -#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] async fn retryable_fault_is_retried_and_result_is_correct_aqe_off() { retryable_fault_is_retried_and_result_is_correct(false).await; } -/// See `retryable_fault_is_retried_and_result_is_correct_aqe_off` above; the -/// AQE-on case fails identically (the error arrives as -/// `Execution("Shared(IoError(..))")` — stringified before the classifier, -/// which is what distinguishes this from the fixed #2028). #[tokio::test] -#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] async fn retryable_fault_is_retried_and_result_is_correct_aqe_on() { retryable_fault_is_retried_and_result_is_correct(true).await; } @@ -210,23 +185,15 @@ use std::time::Duration; /// Scenario D: SIGKILL an executor while it is running tasks. /// -/// `chaos_delay` holds stage 1 open so the kill lands while tasks are genuinely -/// in flight. The scheduler must detect the loss, reschedule the dead executor's -/// tasks onto the survivor, and still return the correct result. +/// `chaos_delay` holds the delayed scan stage open so the kill lands while +/// tasks are genuinely in flight. The scheduler must detect the loss, +/// reschedule the dead executor's tasks onto the survivor, and still return the +/// correct result. /// -/// Both cases reproduce #2027: the shuffle reader's typed -/// `BallistaError::FetchFailed` is flattened into an opaque -/// `DataFusionError::Execution` before the executor reports its `TaskStatus`, -/// so the scheduler never sees the `FetchPartitionError` that would make it -/// resubmit the lost map stage, and fails the query instead of recovering it. -/// -/// Ignored, not deleted or weakened. Un-ignore it as the regression test when -/// #2027 is fixed; see chaos-testing/README.md, Finding 1. #[rstest] #[case::aqe_off(false)] #[case::aqe_on(true)] #[tokio::test] -#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] async fn executor_killed_mid_stage_is_recovered(#[case] aqe: bool) { let mut run = ChaosRun::start(aqe, 2).await; let expected = run.local_baseline().await; @@ -235,18 +202,21 @@ async fn executor_killed_mid_stage_is_recovered(#[case] aqe: bool) { // enough to kill an executor inside it. let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 300)"); - // Submit the query concurrently, then kill executor 0 once stage 1 is running. + // Submit the query concurrently, then kill executor 0 once the delayed + // stage is running. AQE and non-AQE planning assign different stage ids + // to that work. let query = tokio::spawn({ let ctx = run.clone_ctx(); let sql = sql.clone(); async move { ctx.sql(&sql).await?.collect().await } }); + let delayed_stage_id = if aqe { 0 } else { 1 }; let job_id = run.cluster.running_job_id().await.expect("job must appear"); run.cluster - .await_stage_running(&job_id, 1) + .await_stage_running(&job_id, delayed_stage_id) .await - .expect("stage 1 must start running"); + .expect("delayed stage must start running"); run.cluster.kill_executor(0).expect("kill executor 0"); let batches = tokio::time::timeout(Duration::from_secs(120), query) @@ -273,27 +243,22 @@ async fn executor_killed_mid_stage_is_recovered(#[case] aqe: bool) { /// recoveries, so the assertion is on correctness, and the path that actually /// fired is only recorded. /// -/// Reproduces #2027: whenever the reduce stage genuinely has to fetch from the -/// dead executor, the typed `FetchFailed` arrives flattened inside -/// `DataFusionError::Shared`, the scheduler never resubmits the map stage, and -/// the job fails. The scenario only passes when the kill happens to land after -/// the reduce tasks have already fetched their partitions, which makes it -/// timing-dependent (it failed in CI, aqe_off case). -/// -/// Ignored, not deleted or weakened. Un-ignore it as the regression test when -/// #2027 is fixed; see chaos-testing/README.md, Finding 1. #[rstest] #[case::aqe_off(false)] #[case::aqe_on(true)] #[tokio::test] -#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] async fn executor_killed_after_shuffle_write_is_recovered(#[case] aqe: bool) { - let mut run = ChaosRun::start_with(aqe, 2, 60).await; - let expected = run.local_baseline().await; + let mut run = ChaosRun::start_with_concurrent_tasks(aqe, 2, 60, 1).await; + let expected = run.local_sql(Fixture::shuffle_loss_query()).await; - // Delay the *aggregate* side so the reduce stage is slow, giving us a window - // between "stage 1 succeeded" and "stage 2 has finished fetching". - let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 50)"); + run.sql("SET ballista.shuffle.force_remote_read = true") + .await + .expect("force remote shuffle reads"); + run.sql("SET ballista.shuffle.max_concurrent_read_requests = 1") + .await + .expect("serialize shuffle reads"); + + let sql = Fixture::shuffle_loss_chaos_query("chaos_delay(d.name IS NOT NULL, 300)"); let query = tokio::spawn({ let ctx = run.clone_ctx(); @@ -302,17 +267,54 @@ async fn executor_killed_after_shuffle_write_is_recovered(#[case] aqe: bool) { }); let job_id = run.cluster.running_job_id().await.expect("job must appear"); - run.cluster - .await_stage_successful(&job_id, 1) + let (executor_index, shuffle_stage_id) = run + .cluster + .await_successful_shuffle_output(&job_id) .await - .expect("stage 1 must complete before we kill its executor"); - run.cluster.kill_executor(0).expect("kill executor 0"); - - let batches = tokio::time::timeout(Duration::from_secs(180), query) + .expect("a shuffle-writing stage must complete before we kill its executor"); + if run + .cluster + .job_status(&job_id) .await - .expect("query must not hang after shuffle output is lost") - .expect("query task must not panic") - .expect("query must recover by re-running the map stage"); + .expect("job status must be readable") + == "Completed" + { + panic!( + "job completed before executor {executor_index} could be killed after stage {shuffle_stage_id} wrote shuffle output\n{}", + run.cluster.diagnostics(&job_id).await + ); + } + if query.is_finished() { + panic!( + "query completed before executor {executor_index} could be killed after stage {shuffle_stage_id} wrote shuffle output\n{}", + run.cluster.diagnostics(&job_id).await + ); + } + run.cluster + .kill_executor(executor_index) + .unwrap_or_else(|e| { + panic!( + "kill executor {executor_index} with stage {shuffle_stage_id} shuffle output: {e}" + ) + }); + + let batches = match tokio::time::timeout(Duration::from_secs(180), query).await { + Ok(result) => match result.expect("query task must not panic") { + Ok(batches) => batches, + Err(e) => { + panic!( + "query must recover by re-running the map stage: {e}\n{}", + run.cluster.diagnostics(&job_id).await + ) + } + }, + Err(e) => { + panic!( + "query must not hang after shuffle output is lost: {e}\n{}", + run.cluster.diagnostics(&job_id).await + ) + } + }; let actual = datafusion::arrow::util::pretty::pretty_format_batches(&batches) .unwrap()