From a4a4b4ca791cb2e6225960d75de8bde47ff8e072 Mon Sep 17 00:00:00 2001 From: goingforstudying-ctrl Date: Tue, 18 Aug 2026 23:51:02 -0400 Subject: [PATCH 1/3] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks When the adaptive query planner replans after a stage completes, it may decide a still-running stage is no longer needed (e.g. the build side of a join produced no data, so the probe side is moot). Previously the scheduler marked such stages cancelled in the planner but left them in the graph and only logged 'there are stages to be cancelled but its not implemented'. The running stage was never completed, so the job could sit waiting on it forever, and a late task completion from the cancelled stage hit 'Invalid stage ID ... is not running' and failed the whole status update. Now the scheduler retires the stage when a replan cancels it: - a Running stage is moved to Failed (TaskKilled) so the job no longer waits on it, and its in-flight tasks are reported for executor kill; - a Resolved/UnResolved stage is dropped outright; - a late task status for a retired stage is discarded as stale instead of failing the whole update. The task manager is also made resilient: an unexpected stage id in an update no longer aborts the entire batch, so the remaining stages are still processed. Adds a lifecycle test reproducing the orphan-stage scenario. Fixes #2150 --- ballista/scheduler/src/state/aqe/mod.rs | 92 ++++++++++++--- .../src/state/aqe/test/job_failure.rs | 111 +++++++++++++++++- ballista/scheduler/src/state/task_manager.rs | 15 ++- 3 files changed, 200 insertions(+), 18 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index c45236cc46..c769121d99 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -355,7 +355,7 @@ impl AdaptiveExecutionGraph { })) } - /// Return a Vec of stages to cancel + /// Return the set of stage ids the replan cancelled fn update_stage_progress( &mut self, stage_id: usize, @@ -439,10 +439,12 @@ impl AdaptiveExecutionGraph { // we update output locations self.output_locations = partitions.into_iter().flatten().collect(); } - // marking stages which need cancelling as canceled. - // stage ids are returned for task cancellation action + // Drop stages the replan cancelled: remove the planner entry and + // retire the graph stage (cancelling in-flight tasks). The + // returned ids are handed to the caller for executor-side kill. for stage_id in stages_to_cancel.iter() { self.planner.cancel_stage(*stage_id)?; + self.retire_cancelled_stage(*stage_id); } Ok(stages_to_cancel) @@ -451,6 +453,50 @@ impl AdaptiveExecutionGraph { } } + /// Retire a stage the replan cancelled. The stage is moved out of + /// `self.stages` so the job no longer waits on it. Any tasks still + /// running are cancelled; a late completion from one of those tasks is + /// then discarded as coming from a cancelled attempt. + fn retire_cancelled_stage(&mut self, stage_id: usize) { + match self.stages.remove(&stage_id) { + Some(ExecutionStage::Running(running)) => { + let inflight = running.running_tasks().len(); + let cancelled = running.to_failed(format!( + "Stage {stage_id} was cancelled by an AQE replan that made it redundant" + )); + self.stages + .insert(stage_id, ExecutionStage::Failed(cancelled)); + debug!( + "Job {} stage {stage_id} retired after AQE replan ({inflight} in-flight task(s) cancelled)", + self.job_id(), + ); + } + Some(stage) => { + // Resolved/UnResolved stages have no in-flight tasks; drop + // them outright. Successful/Failed stages are already + // terminal and keep their outputs/history. + if matches!( + stage, + ExecutionStage::Resolved(_) | ExecutionStage::UnResolved(_) + ) { + debug!( + "Job {} stage {stage_id} dropped after AQE replan", + self.job_id(), + ); + } else { + self.stages.insert(stage_id, stage); + } + } + None => { + warn!( + "Stage {}/{} to be cancelled was not found in the execution graph", + self.job_id(), + stage_id + ); + } + } + } + fn get_running_stage_id(&mut self, black_list: &[usize]) -> Option { let mut running_stage_id = self.stages.iter().find_map(|(stage_id, stage)| { if black_list.contains(stage_id) { @@ -965,21 +1011,37 @@ impl ExecutionGraph for AdaptiveExecutionGraph { )?; if !stages_to_cancel.is_empty() { - warn!( - "there are stages to be cancelled but its not implemented. stages to cancel: {:?}", - stages_to_cancel + debug!( + "retired stages after AQE replan for job {}: {:?}", + job_id, stages_to_cancel ); } } else { - warn!( - "Stage {}/{} is not in running when updating the status of tasks {:?}", - job_id, - stage_id, - stage_task_statuses - .into_iter() - .map(|task_status| task_status.task_id) - .collect::>(), - ); + // The stage was retired from the graph. A running stage is + // only retired by an AQE replan that made it redundant, so a + // late status from one of its tasks is stale: drop it instead + // of failing the whole update (which would wedge the job). + if matches!(stage, ExecutionStage::Failed(_)) { + warn!( + "Stage {}/{} was retired by an AQE replan; ignoring late status for task(s) {:?}", + job_id, + stage_id, + stage_task_statuses + .iter() + .map(|task_status| task_status.task_id) + .collect::>(), + ); + } else { + warn!( + "Stage {}/{} is not in running when updating the status of tasks {:?}", + job_id, + stage_id, + stage_task_statuses + .into_iter() + .map(|task_status| task_status.task_id) + .collect::>(), + ); + } } } else { return Err(BallistaError::Internal(format!( diff --git a/ballista/scheduler/src/state/aqe/test/job_failure.rs b/ballista/scheduler/src/state/aqe/test/job_failure.rs index f468e55b8e..324acfb248 100644 --- a/ballista/scheduler/src/state/aqe/test/job_failure.rs +++ b/ballista/scheduler/src/state/aqe/test/job_failure.rs @@ -23,7 +23,8 @@ use crate::state::execution_stage::ExecutionStage; use crate::test_utils::mock_executor; use ballista_core::error::Result; use ballista_core::serde::protobuf::{ - FailedTask, JobStatus, failed_task, job_status, task_status, + FailedTask, JobStatus, ShuffleWritePartition, SuccessfulTask, failed_task, + job_status, task_status, }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::execution::context::{SessionConfig, SessionContext}; @@ -135,3 +136,111 @@ async fn test_abort_running_cancels_stages_and_returns_inflight_tasks() -> Resul Ok(()) } + +// Reproduces the orphan-stage lifecycle: when the build side of a join +// produces no data, the replan cancels the (already running) probe stage. +// The cancelled stage must be retired from the graph so the job does not +// wait on it forever, and a late task completion from it must be discarded +// rather than fail the whole update. +#[tokio::test] +async fn test_replan_cancelled_stage_is_retired_and_late_task_discarded() -> Result<()> { + let executor = mock_executor("executor-id1".to_string()); + let mut graph = test_join_plan(2).await; + + // Move the two leaf stages to Running so tasks can be dispatched + graph.revive(); + let running = graph.running_stages(); + assert!( + running.len() >= 2, + "expected two leaf stages, found {running:?}" + ); + + // The two leaf stages run concurrently. Dispatch every available task, + // then hold back exactly one (the "probe" task) in flight; complete all + // the rest with empty output so the replan sees that side produced no + // data and cancels the stage the held task belongs to. The dispatch + // order is not deterministic (stages live in a HashMap), so pick the + // held task arbitrarily and derive its stage id from the task itself. + let mut held_task = None; + let mut complete_tasks = Vec::new(); + while let Some(task) = graph.pop_next_task(&executor.id)? { + if held_task.is_none() { + held_task = Some(task); + } else { + complete_tasks.push(task); + } + } + let held_task = held_task.expect("expected at least one dispatchable task"); + let held_stage_id = held_task.key.stage_id; + assert!( + !complete_tasks.is_empty(), + "expected tasks from the sibling stage to drive the replan" + ); + + // Complete every other (sibling-stage) task with empty output so the + // replan sees that side produced no data and cancels the held task's + // stage. + for task in complete_tasks { + let status = ballista_core::serde::protobuf::TaskStatus { + task_id: task.key.task_id as u32, + job_id: graph.job_id().clone().into(), + stage_id: task.key.stage_id as u32, + stage_attempt_num: 0, + launch_time: 0, + start_exec_time: 0, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask { + executor_id: executor.id.clone(), + partitions: vec![ShuffleWritePartition { + partition_id: task.key.task_id as u64, + num_batches: 0, + num_rows: 0, + num_bytes: 0, + file_id: None, + is_sort_shuffle: false, + }], + })), + }; + graph.update_task_status(&executor, vec![status], 4, 4)?; + } + + // The held task's stage must have been retired: the job no longer tracks + // it as running, so it cannot wedge the job waiting for it. + assert!( + !graph.running_stages().contains(&held_stage_id), + "cancelled stage must not remain running, running={:?}", + graph.running_stages() + ); + assert!( + matches!( + graph.stages.get(&held_stage_id), + Some(ExecutionStage::Failed(_)) | None + ), + "cancelled stage must be Failed or removed, found {:?}", + graph.stages.get(&held_stage_id) + ); + + // A late completion from the already-cancelled held task must be + // discarded instead of failing the whole status update (which used to + // error with "Invalid stage ID" and wedge the job). + let late = ballista_core::serde::protobuf::TaskStatus { + task_id: held_task.key.task_id as u32, + job_id: graph.job_id().clone().into(), + stage_id: held_stage_id as u32, + stage_attempt_num: 0, + launch_time: 0, + start_exec_time: 0, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask { + executor_id: executor.id.clone(), + partitions: vec![], + })), + }; + graph + .update_task_status(&executor, vec![late], 4, 4) + .expect("late task status from a replan-cancelled stage must be discarded"); + + Ok(()) +} diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 60e5330478..d26066f102 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -564,12 +564,23 @@ impl TaskManager self.get_active_execution_graph(&job_id.clone().into()) { let mut graph = cached.write().await; - graph.update_task_status( + // A failure updating one job must not abort the batch and + // drop the remaining jobs' task updates. Log the failure and + // continue with the next job. + match graph.update_task_status( executor, statuses, self.task_max_failures, self.stage_max_failures, - )? + ) { + Ok(events) => events, + Err(error) => { + warn!( + "Failed to update task statuses for job {job_id}, skipping its updates: {error}" + ); + vec![] + } + } } else { // TODO Deal with curator changed case error!( From 82b773ea1740e530dbf379c21dfb4ac1f51c9656 Mon Sep 17 00:00:00 2001 From: goingforstudying-ctrl Date: Wed, 19 Aug 2026 10:12:16 -0400 Subject: [PATCH 2/3] test: add runtime_stats to SuccessfulTask initializers after upstream field addition --- ballista/scheduler/src/state/aqe/test/job_failure.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ballista/scheduler/src/state/aqe/test/job_failure.rs b/ballista/scheduler/src/state/aqe/test/job_failure.rs index 324acfb248..d2307fcf70 100644 --- a/ballista/scheduler/src/state/aqe/test/job_failure.rs +++ b/ballista/scheduler/src/state/aqe/test/job_failure.rs @@ -200,6 +200,7 @@ async fn test_replan_cancelled_stage_is_retired_and_late_task_discarded() -> Res file_id: None, is_sort_shuffle: false, }], + runtime_stats: vec![], })), }; graph.update_task_status(&executor, vec![status], 4, 4)?; @@ -236,6 +237,7 @@ async fn test_replan_cancelled_stage_is_retired_and_late_task_discarded() -> Res status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: executor.id.clone(), partitions: vec![], + runtime_stats: vec![], })), }; graph From d58bd6b98fe02e26c31ea879a6cf2f9b2c67c6f7 Mon Sep 17 00:00:00 2001 From: goingforstudying-ctrl Date: Sat, 22 Aug 2026 08:45:15 -0400 Subject: [PATCH 3/3] fix(scheduler): drop replan-cancelled stages from the graph entirely Retiring a cancelled running stage as Failed kept it in the stage map, and is_successful() requires every mapped stage to be Successful, so the job could never finish once all real work completed. The TPC-DS suite hung on this until the job timeout cancelled the run. Remove the cancelled stage from the map on retire and track its id in a retired_stages set instead; late task statuses for retired stages are discarded via that set rather than matching on a Failed stage. --- ballista/scheduler/src/state/aqe/mod.rs | 74 ++++++++++--------- .../src/state/aqe/test/job_failure.rs | 11 +-- 2 files changed, 45 insertions(+), 40 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index c769121d99..36a95cbde0 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -116,6 +116,12 @@ pub(crate) struct AdaptiveExecutionGraph { end_time: u64, /// Map from Stage ID -> ExecutionStage stages: HashMap, + /// Stages retired after an AQE replan made them redundant. Tracked + /// separately from `stages` because `is_successful()` requires every + /// stage in the map to be Successful: a cancelled stage left behind in + /// a non-Successful state would block job completion forever. Late + /// task statuses for these stage ids are discarded instead of erroring. + retired_stages: HashSet, /// Locations of this `ExecutionGraph` final output locations output_locations: Vec, @@ -194,6 +200,7 @@ impl AdaptiveExecutionGraph { start_time: started_at, end_time: 0, stages, + retired_stages: HashSet::new(), output_locations: vec![], failed_stage_attempts: HashMap::new(), session_config, @@ -454,31 +461,30 @@ impl AdaptiveExecutionGraph { } /// Retire a stage the replan cancelled. The stage is moved out of - /// `self.stages` so the job no longer waits on it. Any tasks still - /// running are cancelled; a late completion from one of those tasks is - /// then discarded as coming from a cancelled attempt. + /// `self.stages` so the job no longer waits on it: `is_successful()` + /// requires every stage in the map to be Successful, so a cancelled + /// stage left behind in a non-Successful state would wedge the job. + /// The id is remembered in `retired_stages` so a late completion from + /// one of its tasks is discarded as coming from a cancelled attempt. fn retire_cancelled_stage(&mut self, stage_id: usize) { match self.stages.remove(&stage_id) { Some(ExecutionStage::Running(running)) => { let inflight = running.running_tasks().len(); - let cancelled = running.to_failed(format!( - "Stage {stage_id} was cancelled by an AQE replan that made it redundant" - )); - self.stages - .insert(stage_id, ExecutionStage::Failed(cancelled)); + self.retired_stages.insert(stage_id); debug!( "Job {} stage {stage_id} retired after AQE replan ({inflight} in-flight task(s) cancelled)", self.job_id(), ); } Some(stage) => { - // Resolved/UnResolved stages have no in-flight tasks; drop - // them outright. Successful/Failed stages are already + // Resolved/UnResolved stages have no in-flight tasks; retire + // them the same way. Successful/Failed stages are already // terminal and keep their outputs/history. if matches!( stage, ExecutionStage::Resolved(_) | ExecutionStage::UnResolved(_) ) { + self.retired_stages.insert(stage_id); debug!( "Job {} stage {stage_id} dropped after AQE replan", self.job_id(), @@ -1017,32 +1023,30 @@ impl ExecutionGraph for AdaptiveExecutionGraph { ); } } else { - // The stage was retired from the graph. A running stage is - // only retired by an AQE replan that made it redundant, so a - // late status from one of its tasks is stale: drop it instead - // of failing the whole update (which would wedge the job). - if matches!(stage, ExecutionStage::Failed(_)) { - warn!( - "Stage {}/{} was retired by an AQE replan; ignoring late status for task(s) {:?}", - job_id, - stage_id, - stage_task_statuses - .iter() - .map(|task_status| task_status.task_id) - .collect::>(), - ); - } else { - warn!( - "Stage {}/{} is not in running when updating the status of tasks {:?}", - job_id, - stage_id, - stage_task_statuses - .into_iter() - .map(|task_status| task_status.task_id) - .collect::>(), - ); - } + warn!( + "Stage {}/{} is not in running when updating the status of tasks {:?}", + job_id, + stage_id, + stage_task_statuses + .into_iter() + .map(|task_status| task_status.task_id) + .collect::>(), + ); } + } else if self.retired_stages.contains(&stage_id) { + // The stage was retired after an AQE replan made it + // redundant, so a late status from one of its tasks is + // stale: discard it instead of failing the whole update + // (which would wedge the job). + warn!( + "Stage {}/{} was retired by an AQE replan; ignoring late status for task(s) {:?}", + job_id, + stage_id, + stage_task_statuses + .iter() + .map(|task_status| task_status.task_id) + .collect::>(), + ); } else { return Err(BallistaError::Internal(format!( "Invalid stage ID {stage_id} for job {job_id}" diff --git a/ballista/scheduler/src/state/aqe/test/job_failure.rs b/ballista/scheduler/src/state/aqe/test/job_failure.rs index d2307fcf70..b0beb196c9 100644 --- a/ballista/scheduler/src/state/aqe/test/job_failure.rs +++ b/ballista/scheduler/src/state/aqe/test/job_failure.rs @@ -213,12 +213,13 @@ async fn test_replan_cancelled_stage_is_retired_and_late_task_discarded() -> Res "cancelled stage must not remain running, running={:?}", graph.running_stages() ); + // The cancelled stage must be gone from the stage map entirely. + // `is_successful()` requires every stage in the map to be Successful, so + // a cancelled stage kept as Failed would make the job hang forever once + // all real work completed (this hung the TPC-DS suite on CI). assert!( - matches!( - graph.stages.get(&held_stage_id), - Some(ExecutionStage::Failed(_)) | None - ), - "cancelled stage must be Failed or removed, found {:?}", + !graph.stages.contains_key(&held_stage_id), + "cancelled stage must be removed from the graph, found {:?}", graph.stages.get(&held_stage_id) );