diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index deec10063..df027e6e2 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -1208,6 +1208,34 @@ impl ExecutionGraph for AdaptiveExecutionGraph { &self.stages } + /// The adaptive planner creates stages incrementally and never knows a + /// stage's consumers up front, so `output_links` is always empty here and + /// the default `output_links`-based implementation would classify every + /// stage as final (reclaiming nothing). Derive the final stage(s) from + /// `output_locations` instead: each `PartitionLocation` carries the + /// terminal stage that produced it, so everything else is intermediate. + /// + /// `output_locations` is only populated once the job has no more stages to + /// run. Until then it is empty and we reclaim nothing, which keeps the + /// "never delete a final stage" invariant. + fn intermediate_stage_ids(&self) -> Vec { + if self.output_locations.is_empty() { + return vec![]; + } + + let final_stage_ids: HashSet = self + .output_locations + .iter() + .map(|location| location.partition_id.stage_id) + .collect(); + + self.stages + .keys() + .filter(|stage_id| !final_stage_ids.contains(stage_id)) + .map(|stage_id| *stage_id as u32) + .collect() + } + fn stage_count(&self) -> usize { self.stages.len() } diff --git a/ballista/scheduler/src/state/aqe/test/intermediate_stages.rs b/ballista/scheduler/src/state/aqe/test/intermediate_stages.rs new file mode 100644 index 000000000..66480a065 --- /dev/null +++ b/ballista/scheduler/src/state/aqe/test/intermediate_stages.rs @@ -0,0 +1,111 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Intermediate-stage identification for the adaptive execution graph, which +//! drives the immediate reclaim of shuffle data when a job succeeds. + +use crate::state::aqe::AdaptiveExecutionGraph; +use crate::state::aqe::test::{mock_batch, mock_context}; +use crate::state::execution_graph::ExecutionGraph; +use crate::test_utils::revive_graph_and_complete_next_stage; +use ballista_core::error::Result; +use std::collections::HashSet; + +/// Builds an adaptive graph for two chained aggregations, which the adaptive +/// planner runs as a chain of stages. +async fn test_two_aggregations_plan() -> AdaptiveExecutionGraph { + let ctx = mock_context(); + ctx.register_batch("t", mock_batch().unwrap()).unwrap(); + + let logical_plan = ctx + .sql("select c0, count(*) from (select min(a) as c0, c as c1 from t group by c) group by c0") + .await + .unwrap() + .into_unoptimized_plan(); + + AdaptiveExecutionGraph::try_new( + "localhost:50050", + &"job".into(), + "", + &ctx, + &logical_plan, + 0, + ) + .await + .unwrap() +} + +/// Drives the graph until every stage is complete, guarding against a stuck +/// loop if no stage makes progress. +fn run_to_success(graph: &mut AdaptiveExecutionGraph) -> Result<()> { + for _ in 0..16 { + if graph.is_successful() { + return Ok(()); + } + revive_graph_and_complete_next_stage(graph)?; + } + panic!("adaptive job did not complete: {:?}", graph.status()); +} + +// The adaptive planner never populates `output_links`, so the default +// implementation would treat every stage as final and reclaim nothing. The +// override derives the final stage from `output_locations` instead. +#[tokio::test] +async fn test_intermediate_stage_ids_excludes_final_stage() -> Result<()> { + let mut graph = test_two_aggregations_plan().await; + + // Nothing is reclaimable before the job produces its output locations. + assert!( + graph.intermediate_stage_ids().is_empty(), + "a running job must not have its stage data reclaimed" + ); + + run_to_success(&mut graph)?; + + let final_stage_ids: HashSet = graph + .output_locations() + .iter() + .map(|location| location.partition_id.stage_id) + .collect(); + assert!( + !final_stage_ids.is_empty(), + "a successful job must report its output locations" + ); + + let all_stage_ids: HashSet = graph.stages().keys().copied().collect(); + assert!( + all_stage_ids.len() > final_stage_ids.len(), + "expected a multi-stage job, found stages {all_stage_ids:?}" + ); + + let intermediate: HashSet = graph + .intermediate_stage_ids() + .into_iter() + .map(|stage_id| stage_id as usize) + .collect(); + + // Exactly the non-terminal stages, and never a stage that is still + // referenced by the job's output locations. + let expected: HashSet = all_stage_ids + .difference(&final_stage_ids) + .copied() + .collect(); + assert_eq!(intermediate, expected); + assert!(intermediate.is_disjoint(&final_stage_ids)); + + Ok(()) +} diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index 52f173d19..4fcaf194e 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -21,6 +21,8 @@ mod alter_stages; mod broadcast_thresholds; /// Functional tests for the CoalescePartitionsRule end-to-end through the planner mod coalesce_rule; +/// Intermediate-stage identification used by the on-success shuffle reclaim +mod intermediate_stages; /// Job-failure lifecycle tests for the adaptive graph mod job_failure; /// covers join selection tests