WIP: feat: treat ResourcesExhausted as a retriable task failure

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;

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<BallistaError> 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<BallistaError> 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(_))
+        ));
+    }
+}

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<PlanProperties>,
+        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<PlanProperties> {
+            &self.properties
+        }
+
+        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+            vec![]
+        }
+
+        fn with_new_children(
+            self: Arc<Self>,
+            _children: Vec<Arc<dyn ExecutionPlan>>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(self)
+        }
+
+        fn execute(
+            &self,
+            _partition: usize,
+            _context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            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<usize>,
+        ) -> Result<Arc<Statistics>> {
+            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<dyn ExecutionPlan> =
+            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<Arc<dyn ExecutionPlan>> {
         let schema = Arc::new(Schema::new(vec![
             Field::new("a", DataType::UInt32, true),