diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index c45236cc4..36a95cbde 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, @@ -355,7 +362,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 +446,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 +460,49 @@ impl AdaptiveExecutionGraph { } } + /// Retire a stage the replan cancelled. The stage is moved out of + /// `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(); + 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; 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(), + ); + } 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,9 +1017,9 @@ 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 { @@ -981,6 +1033,20 @@ impl ExecutionGraph for AdaptiveExecutionGraph { .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 f468e55b8..b0beb196c 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,114 @@ 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, + }], + runtime_stats: vec![], + })), + }; + 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() + ); + // 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!( + !graph.stages.contains_key(&held_stage_id), + "cancelled stage must be removed from the graph, 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![], + runtime_stats: 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 60e533047..d26066f10 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!(