This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new a9d9a7296 fix: classify wrapped FetchFailed as fetch partition error
(#2202)
a9d9a7296 is described below
commit a9d9a7296088598ce60ddfcdd1e00fc8542281ea
Author: Ville Brofeldt <[email protected]>
AuthorDate: Fri Aug 7 14:24:37 2026 -0700
fix: classify wrapped FetchFailed as fetch partition error (#2202)
---
ballista/core/src/error.rs | 221 +++++++++++++++++++--
.../core/src/execution_plans/shuffle_writer.rs | 119 +++++++----
.../src/execution_plans/sort_shuffle/writer.rs | 112 +++++++++--
ballista/executor/src/executor.rs | 4 +-
.../src/scheduler_server/query_stage_scheduler.rs | 5 +
ballista/scheduler/src/state/aqe/adapter.rs | 11 +-
ballista/scheduler/src/state/aqe/mod.rs | 23 ++-
ballista/scheduler/src/state/aqe/planner.rs | 14 +-
ballista/scheduler/src/state/execution_graph.rs | 44 ++--
ballista/scheduler/src/state/task_manager.rs | 12 ++
chaos-testing/README.md | 167 +++++++---------
chaos-testing/src/bin/chaos-executor.rs | 1 +
chaos-testing/src/cluster.rs | 102 +++++++++-
chaos-testing/src/fixture.rs | 15 ++
chaos-testing/tests/common/mod.rs | 23 ++-
chaos-testing/tests/ha.rs | 136 ++++++-------
16 files changed, 713 insertions(+), 296 deletions(-)
diff --git a/ballista/core/src/error.rs b/ballista/core/src/error.rs
index b5a7a82b0..33504fded 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<FetchFailedDetails> {
+ 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<FetchFailedDetails> {
+ match e.find_root() {
+ DataFusionError::ArrowError(e, _) => fetch_failed_in_arrow(e),
+ _ => None,
+ }
+}
+
+fn fetch_failed_in_arrow(e: &ArrowError) -> Option<FetchFailedDetails> {
+ let ArrowError::ExternalError(inner) = e else {
+ return None;
+ };
+ if let Some(e) = inner.downcast_ref::<BallistaError>() {
+ return find_fetch_failed(e);
+ }
+ if let Some(e) = inner.downcast_ref::<DataFusionError>() {
+ return fetch_failed_in_datafusion(e);
+ }
+ None
+}
+
impl From<BallistaError> 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 8854962d1..23257e633 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<ShuffleWriterExec>,
+ task_ctx: Arc<TaskContext>,
+ ) -> Vec<crate::error::Result<Vec<RecordBatch>>> {
+ 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::<BallistaError>().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<dyn ExecutionPlan> = 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 7802bcb87..4d8405d19 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<Vec<ShuffleWritePartition>> =
(0..k).map(|_| Vec::new()).collect();
- let mut first_error: Option<String> = None;
+ let mut first_error: Option<DataFusionError> = 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<SortShuffleWriterExec>,
+ task_ctx: Arc<TaskContext>,
+ ) -> Vec<crate::error::Result<Vec<RecordBatch>>> {
+ 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<dyn ExecutionPlan> = 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 dce01ae7f..b491cfa8a 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 2ccf5a9da..3086071ec 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<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
{
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 539b45c6b..b6ec54866 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<usize>,
+ inputs: HashMap<usize, StageOutput>,
}
///
@@ -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 0c133db64..dac848221 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<HashMap<usize,
ExecutionStage>> =
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<SessionConfig>,
- stage: Arc<dyn ShuffleWriter>,
+ 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 690a5f2bc..6dc08dab6 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<Vec<PartitionLocation>>` (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<dyn ShuffleWriter>,
- #[allow(dead_code)] // TODO: still not sure if this is needed
- pub(crate) inputs: Vec<usize>,
+ pub(crate) inputs: std::collections::HashMap<usize, StageOutput>,
}
diff --git a/ballista/scheduler/src/state/execution_graph.rs
b/ballista/scheduler/src/state/execution_graph.rs
index 49f3cab62..99b7c38d5 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 dea5af100..604b0a31e 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 88b198b57..b4afbfc63 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<BallistaError> 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<DataFusionError> 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(\"<executor-id>\", ...,
-\"...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 a2efe82a7..830d295ae 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 33fc8d040..d1972cd57 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<TestCluster, String> {
// 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<PathBuf> = 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<usize> {
+ 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<usize> {
+ 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 3871fce4e..f8c6d8784 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 1a2755d3d..831be3d16 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 fe9f93d11..4b5980aae 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()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
