fix(scheduler): fail running jobs when all executors are lost

ballista/scheduler/src/config.rs

@@ -173,6 +173,14 @@ pub struct Config {
         help = "Interval, in seconds, to check expired or dead executors."
     )]
     pub expire_dead_executor_interval_seconds: u64,
+    /// Grace period in seconds to wait for an executor to (re)appear after the
+    /// cluster has lost its last executor before failing the running jobs.
+    #[arg(
+        long,
+        default_value_t = 30,
+        help = "Grace period, in seconds, to wait for an executor to (re)register after the last executor is lost before failing running jobs. Prevents jobs from hanging forever when every executor dies, while still tolerating a transient total loss (e.g. a rolling restart). Set to 0 to fail as soon as the loss is observed."
+    )]
+    pub no_executors_grace_period_seconds: u64,
     /// Minimum number of registered executors before /readyz returns 200
     #[arg(
         long,
@@ -260,6 +268,12 @@ pub struct SchedulerConfig {
     pub executor_timeout_seconds: u64,
     /// The interval to check expired or dead executors
     pub expire_dead_executor_interval_seconds: u64,
+    /// Grace period in seconds to wait for an executor to (re)register after the
+    /// cluster has lost its last executor before failing the running jobs. This
+    /// bounds the otherwise-unbounded wait so that a total executor loss fails
+    /// the affected jobs instead of hanging forever. Set to 0 to fail as soon as
+    /// the loss is observed.
+    pub no_executors_grace_period_seconds: u64,
     /// [ConfigProducer] override option
     pub override_config_producer: Option<ConfigProducer>,
     /// [SessionBuilder] override option
@@ -309,6 +323,7 @@ impl Default for SchedulerConfig {
             grpc_server_max_encoding_message_size: 16777216,
             executor_timeout_seconds: 180,
             expire_dead_executor_interval_seconds: 15,
+            no_executors_grace_period_seconds: 30,
             override_config_producer: None,
             override_session_builder: None,
             override_logical_codec: None,
@@ -414,6 +429,13 @@ impl SchedulerConfig {
         self
     }
 
+    /// Sets the grace period, in seconds, to wait for an executor to (re)register
+    /// after the last executor is lost before failing running jobs.
+    pub fn with_no_executors_grace_period_seconds(mut self, value: u64) -> Self {
+        self.no_executors_grace_period_seconds = value;
+        self
+    }
+
     /// Sets the maximum gRPC server decoding message size.
     pub fn with_grpc_server_max_decoding_message_size(mut self, value: u32) -> Self {
         self.grpc_server_max_decoding_message_size = value;
@@ -542,6 +564,7 @@ impl TryFrom<Config> for SchedulerConfig {
             executor_timeout_seconds: opt.executor_timeout_seconds,
             expire_dead_executor_interval_seconds: opt
                 .expire_dead_executor_interval_seconds,
+            no_executors_grace_period_seconds: opt.no_executors_grace_period_seconds,
             override_config_producer: None,
             override_logical_codec: None,
             override_physical_codec: None,

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -16,10 +16,11 @@
 // under the License.
 
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 
 use ballista_core::JobId;
-use ballista_core::serde::protobuf::{FailedJob, JobStatus};
+use ballista_core::serde::protobuf::{FailedJob, JobStatus, job_status};
 use log::{debug, error, info, trace, warn};
 
 use ballista_core::error::{BallistaError, Result};
@@ -45,6 +46,11 @@ pub(crate) struct QueryStageScheduler<
     state: Arc<SchedulerState<T, U>>,
     metrics_collector: Arc<dyn SchedulerMetricsCollector>,
     config: Arc<SchedulerConfig>,
+    /// Guards against arming more than one "all executors lost" grace timer at a
+    /// time. When a whole cluster dies at once the reaper posts an `ExecutorLost`
+    /// per executor, and each would otherwise arm its own timer and fail every
+    /// running job again. See <https://github.com/apache/datafusion-ballista/issues/2029>
+    no_executor_check_pending: Arc<AtomicBool>,
 }
 
 impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryStageScheduler<T, U> {
@@ -57,6 +63,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryStageSchedul
             state,
             metrics_collector,
             config,
+            no_executor_check_pending: Arc::new(AtomicBool::new(false)),
         }
     }
 
@@ -339,6 +346,88 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                         error!("{msg}");
                     }
                 }
+
+                // If that was the last executor, the running jobs whose tasks were
+                // just reset can no longer make progress — there is nothing to
+                // schedule them onto. Rather than hang forever, wait a bounded
+                // grace period for an executor to (re)register (e.g. a rolling
+                // restart) and then fail any job still running on an empty cluster.
+                // Only fires for executors that were actually present, so jobs
+                // merely queued waiting for their first executor (autoscaling cold
+                // start) are never affected.
+                // See https://github.com/apache/datafusion-ballista/issues/2029
+                //
+                // `no_executor_check_pending` collapses the burst of `ExecutorLost`
+                // events produced when a whole cluster dies at once into a single
+                // timer, so each running job is failed at most once.
+                if self.state.executor_manager.get_alive_executors().is_empty()
+                    && !self.no_executor_check_pending.swap(true, Ordering::SeqCst)
+                {
+                    let state = self.state.clone();
+                    let sender = event_sender.clone();
+                    let pending = self.no_executor_check_pending.clone();
+                    let grace = Duration::from_secs(
+                        state.config.no_executors_grace_period_seconds,
+                    );
+                    let lost_at = timestamp_millis();
+                    tokio::spawn(async move {
+                        tokio::time::sleep(grace).await;
+
+                        // An executor may have (re)registered during the grace
+                        // window; if so the reset tasks will be scheduled onto it
+                        // and there is nothing to fail.
+                        if state.executor_manager.get_alive_executors().is_empty() {
+                            for job_id in
+                                state.task_manager.get_running_job_cache().keys()
+                            {
+                                // Re-read the live status right before failing: a
+                                // job that finished during the grace window must
+                                // not be failed, and a job planned *after* the
+                                // cluster went empty (started_at > lost_at) has its
+                                // own window and must not inherit this one.
+                                let queued_at = match state
+                                    .task_manager
+                                    .get_job_status(job_id)
+                                    .await
+                                {
+                                    Ok(Some(JobStatus {
+                                        status: Some(job_status::Status::Running(running)),
+                                        ..
+                                    })) if running.started_at <= lost_at => {
+                                        running.queued_at
+                                    }
+                                    _ => continue,
+                                };
+
+                                let fail_message = format!(
+                                    "all executors were lost and no executor re-registered within {}s; no executors remain to run the tasks for this job",
+                                    grace.as_secs()
+                                );
+                                warn!("Failing job {job_id}: {fail_message}");
+                                if let Err(e) = sender
+                                    .post_event(
+                                        QueryStageSchedulerEvent::JobRunningFailed {
+                                            job_id: job_id.clone(),
+                                            fail_message,
+                                            queued_at,
+                                            failed_at: timestamp_millis(),
+                                        },
+                                    )
+                                    .await
+                                {
+                                    error!(
+                                        "Fail to post JobRunningFailed for job {job_id}: {e:?}"
+                                    );
+                                }
+                            }
+                        }
+
+                        // Cleared last, so the whole burst of `ExecutorLost` events
+                        // that a simultaneous cluster death produces collapses into
+                        // this single check — even when the grace period is 0.
+                        pending.store(false, Ordering::SeqCst);
+                    });
+                }
             }
             QueryStageSchedulerEvent::CancelTasks(tasks) => {
                 if let Err(e) = self
@@ -380,6 +469,7 @@ mod tests {
     use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition};
     use ballista_core::config::TaskSchedulingPolicy;
     use ballista_core::error::Result;
+    use ballista_core::serde::protobuf::job_status;
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::functions_aggregate::sum::sum;
     use datafusion::logical_expr::{LogicalPlan, col};
@@ -441,6 +531,113 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_running_job_fails_when_all_executors_are_lost() -> Result<()> {
+        let plan = test_plan(10);
+
+        let metrics_collector = Arc::new(TestMetricsCollector::default());
+
+        // Grace period of 0 so the job is failed as soon as the loss is observed,
+        // keeping the test fast.
+        let mut test = SchedulerTest::new(
+            SchedulerConfig::default()
+                .with_scheduler_policy(TaskSchedulingPolicy::PushStaged)
+                .with_no_executors_grace_period_seconds(0),
+            metrics_collector.clone(),
+            1,
+            1,
+            None,
+        )
+        .await?;
+
+        let job_id = test.submit("", &plan).await?;
+
+        // Wait until the job is actually running with tasks in flight. We
+        // deliberately never `tick()`, so its tasks never complete.
+        let job_id_ref = &job_id;
+        let test_ref = &test;
+        let running = await_condition(Duration::from_millis(50), 40, || async move {
+            let status = test_ref.job_status(job_id_ref).await?;
+            Ok(matches!(
+                status.and_then(|s| s.status),
+                Some(job_status::Status::Running(_))
+            ))
+        })
+        .await?;
+        assert!(running, "job should reach the running state");
+
+        // The only executor is lost. With no executors left, the reset tasks can
+        // never be scheduled, so the job must fail rather than hang forever
+        // (#2029).
+        test.lose_executor("virtual-executor-0").await?;
+
+        let failed = await_condition(Duration::from_millis(100), 50, || async move {
+            let status = test_ref.job_status(job_id_ref).await?;
+            Ok(matches!(
+                status.and_then(|s| s.status),
+                Some(job_status::Status::Failed(_))
+            ))
+        })
+        .await?;
+        assert!(
+            failed,
+            "job should be failed after all executors were lost, but status was {:?}",
+            test.job_status(&job_id).await?
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_running_job_survives_partial_executor_loss() -> Result<()> {
+        let plan = test_plan(10);
+
+        let metrics_collector = Arc::new(TestMetricsCollector::default());
+
+        let mut test = SchedulerTest::new(
+            SchedulerConfig::default()
+                .with_scheduler_policy(TaskSchedulingPolicy::PushStaged)
+                .with_no_executors_grace_period_seconds(0),
+            metrics_collector.clone(),
+            2,
+            1,
+            None,
+        )
+        .await?;
+
+        let job_id = test.submit("", &plan).await?;
+
+        let job_id_ref = &job_id;
+        let test_ref = &test;
+        let running = await_condition(Duration::from_millis(50), 40, || async move {
+            let status = test_ref.job_status(job_id_ref).await?;
+            Ok(matches!(
+                status.and_then(|s| s.status),
+                Some(job_status::Status::Running(_))
+            ))
+        })
+        .await?;
+        assert!(running, "job should reach the running state");
+
+        // Lose only one of two executors. One remains alive, so the job must not
+        // be failed by the total-loss guard.
+        test.lose_executor("virtual-executor-0").await?;
+
+        // Give the (grace-0) failure path ample time to fire if it were going to.
+        tokio::time::sleep(Duration::from_millis(500)).await;
+
+        let status = test.job_status(&job_id).await?;
+        assert!(
+            !matches!(
+                status.as_ref().and_then(|s| s.status.clone()),
+                Some(job_status::Status::Failed(_))
+            ),
+            "job must not be failed while an executor remains, but status was {status:?}"
+        );
+
+        Ok(())
+    }
+
     fn test_plan(partitions: usize) -> LogicalPlan {
         let schema = Schema::new(vec![
             Field::new("id", DataType::Utf8, false),

ballista/scheduler/src/test_utils.rs

@@ -543,6 +543,32 @@ impl SchedulerTest {
             .await
     }
 
+    /// Simulates the loss of an executor: deregisters it from the executor
+    /// manager and posts the `ExecutorLost` event. This mirrors the reaper's
+    /// `remove_executor` path without waiting out the heartbeat timeout.
+    pub async fn lose_executor(&self, executor_id: &str) -> Result<()> {
+        let reason = Some("test: executor lost".to_owned());
+        self.scheduler
+            .state
+            .executor_manager
+            .remove_executor(executor_id, reason.clone())
+            .await?;
+        self.post_scheduler_event(QueryStageSchedulerEvent::ExecutorLost(
+            executor_id.to_owned(),
+            reason,
+        ))
+        .await
+    }
+
+    /// Returns the current status of a job, if known.
+    pub async fn job_status(&self, job_id: &JobId) -> Result<Option<JobStatus>> {
+        self.scheduler
+            .state
+            .task_manager
+            .get_job_status(job_id)
+            .await
+    }
+
     /// Waits for job completion with a timeout in milliseconds.
     pub async fn await_completion_timeout(
         &self,

chaos-testing/README.md

@@ -188,7 +188,7 @@ sanity check that every other scenario's assertions depend on.
 | 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.**                                 |
 | 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; the only requirement is that the job _terminates_ within 120s rather than hanging.                                                                                        | **Ignored (both), reproduces [#2029](https://github.com/apache/datafusion-ballista/issues/2029) — Finding 3.**                                 |
+| 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
@@ -325,19 +325,20 @@ describes for `FetchFailed`, which is why both scenarios are ignored against
 shuffle output — which is every non-final stage — is classified non-retryable,
 turning what should be a retried task into an immediate job failure.
 
-### Finding 3 — Killing every executor hangs the job instead of failing it
+### Finding 3 — Killing every executor hung the job instead of failing it (fixed)
 
-Tracked by [#2029](https://github.com/apache/datafusion-ballista/issues/2029).
+Tracked by [#2029](https://github.com/apache/datafusion-ballista/issues/2029), fixed in that issue's PR.
 
-**Proven by:** Scenario G (`killing_every_executor_terminates_the_job`), both
-AQE settings, `#[ignore]`d against that issue.
+**Regression test:** Scenario G (`killing_every_executor_terminates_the_job`),
+both AQE settings, now enabled (no longer `#[ignore]`d).
 
 With every executor dead mid-query, there is nothing left to schedule tasks
-onto. The job does not terminate within the scenario's 120s timeout — the
-scheduler waits rather than failing the query once it can determine no
-executor can ever satisfy the remaining tasks. Note that this scenario
-deliberately asserts only _termination_, not success or a particular error; it
-is a hang detector, and what it detects is the hang itself.
+onto. Previously the job never terminated — the scheduler waited forever rather
+than failing the query. The fix makes the scheduler wait a bounded grace period
+(`no_executors_grace_period_seconds`) after losing its last executor and then
+fail the job with a clear error. The scenario turns that grace down via the
+cluster builder and asserts the query fails with an error naming the executor
+loss.
 
 ### For comparison: the heartbeat-expiry path does recover
 

chaos-testing/src/bin/chaos-scheduler.rs

@@ -54,6 +54,12 @@ async fn main() -> ballista_core::error::Result<()> {
         ),
         task_max_failures: env_parsed("CHAOS_TASK_MAX_FAILURES", 4),
         stage_max_failures: env_parsed("CHAOS_STAGE_MAX_FAILURES", 4),
+        // The default is 30s. Turn it down so the total-executor-loss scenario
+        // fails the job a second or so after the reap instead of waiting it out.
+        no_executors_grace_period_seconds: env_parsed(
+            "CHAOS_NO_EXECUTORS_GRACE_SECONDS",
+            1,
+        ),
         override_session_builder: Some(Arc::new(chaos_session_state)),
         ..Default::default()
     };