diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 41c946365..8ecde3a01 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -473,6 +473,9 @@ message FailedTask { // A successful task's result is lost due to executor lost ResultLost result_lost = 8; TaskKilled task_killed = 9; + // The task ran out of memory. Retriable: the retry may land on a less-loaded + // executor, or run once its peers have drained. + ResourcesExhausted resources_exhausted = 10; } } @@ -504,6 +507,9 @@ message ResultLost { message TaskKilled { } +message ResourcesExhausted { +} + message ShuffleWritePartition { uint64 partition_id = 1; uint64 num_batches = 3; diff --git a/ballista/core/src/error.rs b/ballista/core/src/error.rs index b0ce594f3..19c4e37fc 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, ResourcesExhausted, +}; use datafusion::error::DataFusionError; use datafusion::{arrow::error::ArrowError, sql::sqlparser::parser}; use futures::future::Aborted; @@ -245,6 +247,25 @@ impl From for FailedTask { failed_reason: Some(FailedReason::IoError(IoError {})), } } + BallistaError::DataFusionError(ref e) + if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) => + { + FailedTask { + error: format!( + "Task failed due to exhausted resources (memory or spill capacity): {e}" + ), + // Retriable: the retry may land on a less-loaded executor, or run once + // this executor's other tasks have drained. Bounded by + // --task-max-failures, so a task that is simply too large still fails + // the job -- but with a clear message, rather than an OOM-killed + // executor and a FetchPartitionError cascade. + retryable: true, + count_to_failures: true, + failed_reason: Some(FailedReason::ResourcesExhausted( + ResourcesExhausted {}, + )), + } + } other => FailedTask { error: format!("Task failed due to runtime execution error: {other:?}"), retryable: false, @@ -256,3 +277,143 @@ impl From for FailedTask { } impl Error for BallistaError {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::serde::protobuf::failed_task::FailedReason; + + #[test] + fn resources_exhausted_maps_to_a_retriable_failed_task() { + let err = BallistaError::DataFusionError(Box::new( + DataFusionError::ResourcesExhausted("over budget".to_string()), + )); + let failed: FailedTask = err.into(); + + assert!(failed.retryable, "an OOM'd task must be retried"); + assert!(failed.count_to_failures, "retries must be bounded"); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + } + + #[test] + fn wrapped_resources_exhausted_is_still_recognized() { + // DataFusion routinely wraps errors in `Context`, so the mapping must look at + // the root cause rather than the outermost error. + let inner = DataFusionError::ResourcesExhausted("over budget".to_string()); + let err = BallistaError::DataFusionError(Box::new( + inner.context("while executing HashJoinExec"), + )); + let failed: FailedTask = err.into(); + + assert!(failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + )); + } + + /// `DataFusionError::Shared` is how DataFusion fans *one* stream error out to + /// *many* output partitions: `RepartitionExec` and `CoalescePartitionsExec` clone a + /// single error into an `Arc` and hand it to every consumer. A memory rejection + /// raised inside a join below a `RepartitionExec` therefore reaches the shuffle + /// writer wrapped in exactly this variant -- so the whole feature's retriability + /// depends on `find_root()` seeing through it. It does (`Shared`'s `Error::source` + /// yields the inner `DataFusionError`), and this is what holds that true. + #[test] + fn a_shared_resources_exhausted_is_still_recognized() { + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(DataFusionError::ResourcesExhausted( + "over budget".to_string(), + )), + ))); + let failed: FailedTask = err.into(); + + assert!( + failed.retryable, + "a ResourcesExhausted fanned out through RepartitionExec must still be retried" + ); + assert!(failed.count_to_failures, "retries must be bounded"); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + } + + /// The realistic shape: DataFusion adds a `Context` to the pool's rejection as it + /// unwinds out of the operator, and *then* the repartition shares it out. Both + /// wrappers have to be seen through. + #[test] + fn a_shared_context_wrapped_resources_exhausted_is_still_recognized() { + let inner = DataFusionError::ResourcesExhausted("over budget".to_string()); + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(inner.context("while executing HashJoinExec")), + ))); + let failed: FailedTask = err.into(); + + assert!(failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + )); + } + + /// The discriminator for the two tests above: `Shared` must not become a blanket + /// "retriable" arm. A shared error whose root is *not* `ResourcesExhausted` still + /// has to fall through to the non-retriable execution-error arm. + #[test] + fn a_shared_non_resource_error_remains_non_retriable() { + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(DataFusionError::Execution("boom".to_string())), + ))); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } + + #[test] + fn other_errors_remain_non_retriable_execution_errors() { + let err = BallistaError::General("boom".to_string()); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } + + #[test] + fn other_datafusion_errors_remain_non_retriable_execution_errors() { + // Sharper discriminator than a bare `General` error: this is a + // `DataFusionError` whose root is NOT `ResourcesExhausted`, so it must + // still fall through to the non-retriable arm. It also happens to be + // the exact shape the old `format!("{e:?}")` re-wrap at + // shuffle_writer.rs used to produce, so it guards against the + // resources-exhausted arm accidentally matching every DataFusionError. + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Execution( + "boom".to_string(), + ))); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.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 4a9919706..93aa8f7d4 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -35,6 +35,7 @@ use std::sync::Arc; use std::time::Instant; use crate::JobId; +use crate::error::BallistaError; use crate::execution_plans::create_shuffle_path; use crate::extension::SessionConfigExt; use crate::utils; @@ -242,7 +243,15 @@ impl ShuffleWriterExec { channel_capacity, ) .await - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + .map_err(|e| match e { + // Preserve the DataFusion error type: `find_root()` in the + // FailedTask mapping relies on it to classify a + // ResourcesExhausted as a retriable failure. A `format!` here + // would flatten it to an Execution error and silently make an + // OOM'd task non-retryable. + BallistaError::DataFusionError(e) => *e, + other => DataFusionError::Execution(format!("{other:?}")), + })?; write_metrics .input_rows @@ -600,10 +609,15 @@ fn result_schema() -> SchemaRef { #[allow(dead_code, unused_imports)] // clippy false positive with local imports mod tests { use super::*; + use crate::error::BallistaError; + use crate::serde::protobuf::FailedTask; + use crate::serde::protobuf::failed_task::FailedReason; use datafusion::arrow::array::{StringArray, StructArray, UInt32Array, UInt64Array}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::expressions::Column; use datafusion::prelude::SessionContext; use tempfile::TempDir; @@ -803,6 +817,138 @@ mod tests { Ok(()) } + /// Test-only plan whose single stream immediately yields a + /// `DataFusionError::ResourcesExhausted`, standing in for an OOM'd input + /// (e.g. a hash-join build side or a sort spilling past its budget) that + /// feeds a `None`-partitioned (unpartitioned) shuffle write stage. + #[derive(Debug)] + struct AlwaysResourcesExhaustedExec { + properties: Arc, + schema: SchemaRef, + } + + impl AlwaysResourcesExhaustedExec { + fn new(schema: SchemaRef) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { properties, schema } + } + } + + impl DisplayAs for AlwaysResourcesExhaustedExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => { + write!(f, "AlwaysResourcesExhaustedExec") + } + } + } + } + + impl ExecutionPlan for AlwaysResourcesExhaustedExec { + fn name(&self) -> &str { + "AlwaysResourcesExhaustedExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let schema = self.schema(); + let stream = futures::stream::once(async { + Err(DataFusionError::ResourcesExhausted( + "over budget".to_string(), + )) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + } + + #[tokio::test] + async fn test_no_repart_resources_exhausted_is_retryable() -> Result<()> { + // This pins the writer boundary, not just the FailedTask mapping in + // error.rs: a `ShuffleWriterExec` with `shuffle_output_partitioning: + // None` (the final stage, a broadcast-join build side, or a + // CoalescePartitionsExec/SortPreservingMergeExec input) must let a + // `ResourcesExhausted` from its input stream survive the disk-write + // path with its error type intact, so the scheduler retries the task + // instead of failing the whole job. + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])); + let input_plan: Arc = + Arc::new(AlwaysResourcesExhaustedExec::new(schema)); + let work_dir = TempDir::new()?; + let query_stage = ShuffleWriterExec::try_new( + JobId::new("jobOne"), + 1, + input_plan, + work_dir.path().to_str().unwrap().to_owned(), + None, + )?; + let mut stream = query_stage.execute(0, task_ctx)?; + let result = utils::collect_stream(&mut stream).await; + + let err: BallistaError = result.expect_err( + "AlwaysResourcesExhaustedExec's stream error must propagate as an error", + ); + let failed: FailedTask = err.into(); + + assert!( + failed.retryable, + "a ResourcesExhausted task on a None-partitioned shuffle write stage \ + must be retryable, got: {failed:?}" + ); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + + Ok(()) + } + fn create_input_plan() -> Result> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, true), diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 07dfa8a42..ca83fc896 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -717,7 +717,7 @@ pub struct FailedTask { /// Whether this task failure should be counted to the maximum number of times the task is allowed to retry #[prost(bool, tag = "3")] pub count_to_failures: bool, - #[prost(oneof = "failed_task::FailedReason", tags = "4, 5, 6, 7, 8, 9")] + #[prost(oneof = "failed_task::FailedReason", tags = "4, 5, 6, 7, 8, 9, 10")] pub failed_reason: ::core::option::Option, } /// Nested message and enum types in `FailedTask`. @@ -737,6 +737,10 @@ pub mod failed_task { ResultLost(super::ResultLost), #[prost(message, tag = "9")] TaskKilled(super::TaskKilled), + /// The task ran out of memory. Retriable: the retry may land on a less-loaded + /// executor, or run once its peers have drained. + #[prost(message, tag = "10")] + ResourcesExhausted(super::ResourcesExhausted), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -768,6 +772,8 @@ pub struct ResultLost {} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct TaskKilled {} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ResourcesExhausted {} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct ShuffleWritePartition { #[prost(uint64, tag = "1")] pub partition_id: u64, diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 8da872253..f1ab16c1d 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -23,7 +23,8 @@ use axum::{ response::{IntoResponse, Response}, }; use ballista_core::serde::protobuf::failed_task::FailedReason::{ - ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, + ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResourcesExhausted, + ResultLost, TaskKilled, }; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ @@ -823,6 +824,7 @@ fn failed_reason(failed: &FailedTask) -> String { Some(ExecutorLost(_)) => "ExecutorLost", Some(ResultLost(_)) => "ResultLost", Some(TaskKilled(_)) => "TaskKilled", + Some(ResourcesExhausted(_)) => "ResourcesExhausted", None => "Failed", } .to_string()