fix(scheduler): retire AQE stages cancelled by replan and discard late tasks

ballista/scheduler/src/state/aqe/mod.rs

@@ -283,7 +283,7 @@ impl AdaptiveExecutionGraph {
         Ok(events)
     }
 
-    /// 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,
@@ -318,10 +318,12 @@ impl AdaptiveExecutionGraph {
                 // we update output locations
                 self.output_locations = locations.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)
@@ -330,6 +332,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<usize> {
         let mut running_stage_id = self.stages.iter().find_map(|(stage_id, stage)| {
             if black_list.contains(stage_id) {
@@ -823,21 +869,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::<Vec<_>>(),
-                    );
+                    // 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::<Vec<_>>(),
+                        );
+                    } 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::<Vec<_>>(),
+                        );
+                    }
                 }
             } else {
                 return Err(BallistaError::Internal(format!(

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

ballista/scheduler/src/state/task_manager.rs

@@ -565,12 +565,23 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
                 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!(