This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new e1e9514ba feat: delete intermediate shuffle files on job success
(#1982)
e1e9514ba is described below
commit e1e9514ba1fd89b59074b60a3b1e1ded334fb0a9
Author: Andy Grove <[email protected]>
AuthorDate: Sat Jul 11 09:01:50 2026 -0600
feat: delete intermediate shuffle files on job success (#1982)
---
ballista/core/proto/ballista.proto | 6 ++
ballista/core/src/serde/generated/ballista.rs | 8 +++
ballista/executor/src/execution_loop.rs | 9 ++-
ballista/executor/src/executor_process.rs | 74 ++++++++++++++++++++--
ballista/executor/src/executor_server.rs | 7 +-
ballista/scheduler/src/scheduler_server/grpc.rs | 3 +-
.../src/scheduler_server/query_stage_scheduler.rs | 16 +++--
ballista/scheduler/src/state/execution_graph.rs | 37 +++++++++++
ballista/scheduler/src/state/execution_stage.rs | 12 ++++
ballista/scheduler/src/state/executor_manager.rs | 60 +++++++++++++++---
ballista/scheduler/src/state/mod.rs | 12 +++-
ballista/scheduler/src/state/task_manager.rs | 13 ++--
12 files changed, 224 insertions(+), 33 deletions(-)
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index 50a3e25e7..41c946365 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -783,6 +783,9 @@ message CancelJobResult {
message CleanJobDataParams {
string job_id = 1;
+ // Specific stage ids to remove within the job dir.
+ // Empty means remove the whole job dir (legacy behavior).
+ repeated uint32 remove_stage_ids = 2;
}
message CleanJobDataResult {
@@ -820,6 +823,9 @@ message CancelTasksResult {
message RemoveJobDataParams {
string job_id = 1;
+ // Specific stage ids to remove within the job dir.
+ // Empty means remove the whole job dir (legacy behavior).
+ repeated uint32 remove_stage_ids = 2;
}
message RemoveJobDataResult {
diff --git a/ballista/core/src/serde/generated/ballista.rs
b/ballista/core/src/serde/generated/ballista.rs
index d07154666..07dfa8a42 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -1188,6 +1188,10 @@ pub struct CancelJobResult {
pub struct CleanJobDataParams {
#[prost(string, tag = "1")]
pub job_id: ::prost::alloc::string::String,
+ /// Specific stage ids to remove within the job dir.
+ /// Empty means remove the whole job dir (legacy behavior).
+ #[prost(uint32, repeated, tag = "2")]
+ pub remove_stage_ids: ::prost::alloc::vec::Vec<u32>,
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
pub struct CleanJobDataResult {}
@@ -1233,6 +1237,10 @@ pub struct CancelTasksResult {
pub struct RemoveJobDataParams {
#[prost(string, tag = "1")]
pub job_id: ::prost::alloc::string::String,
+ /// Specific stage ids to remove within the job dir.
+ /// Empty means remove the whole job dir (legacy behavior).
+ #[prost(uint32, repeated, tag = "2")]
+ pub remove_stage_ids: ::prost::alloc::vec::Vec<u32>,
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
pub struct RemoveJobDataResult {}
diff --git a/ballista/executor/src/execution_loop.rs
b/ballista/executor/src/execution_loop.rs
index a8b08efd7..77cd4c7c5 100644
--- a/ballista/executor/src/execution_loop.rs
+++ b/ballista/executor/src/execution_loop.rs
@@ -23,7 +23,7 @@
use crate::cpu_bound_executor::DedicatedExecutor;
use crate::executor::Executor;
-use crate::executor_process::remove_job_dir;
+use crate::executor_process::remove_job_data;
use crate::{TaskExecutionTimes, as_task_status};
use ballista_core::JobId;
use ballista_core::error::BallistaError;
@@ -119,12 +119,15 @@ where
for cleanup in jobs_to_clean {
let job_id = cleanup.job_id.clone().into();
let work_dir = executor.work_dir.clone();
+ let remove_stage_ids = cleanup.remove_stage_ids.clone();
// In poll-based cleanup, removing job data is
fire-and-forget.
// Failures here do not affect task execution and are only
logged.
tokio::spawn(async move {
- if let Err(e) = remove_job_dir(&work_dir,
&job_id).await {
- error!("failed to remove job dir {job_id}: {e}");
+ if let Err(e) =
+ remove_job_data(&work_dir, &job_id,
&remove_stage_ids).await
+ {
+ error!("failed to remove job data {job_id}: {e}");
}
});
}
diff --git a/ballista/executor/src/executor_process.rs
b/ballista/executor/src/executor_process.rs
index 22f526edc..4e37687e8 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -737,11 +737,16 @@ async fn clean_all_shuffle_data(work_dir: &str) ->
ballista_core::error::Result<
Ok(())
}
-/// Remove a job directory under work_dir.
+/// Remove job data under work_dir.
/// Used by both push-based (gRPC handler) and pull-based (poll loop) cleanup.
-pub(crate) async fn remove_job_dir(
+///
+/// `remove_stage_ids` empty ⇒ remove the whole `work_dir/{job_id}` dir (legacy
+/// behavior). Non-empty ⇒ remove only the listed
`work_dir/{job_id}/{stage_id}`
+/// subdirs, retaining the rest of the job dir (e.g. the final-stage output).
+pub(crate) async fn remove_job_data(
work_dir: &str,
job_id: &JobId,
+ remove_stage_ids: &[u32],
) -> ballista_core::error::Result<()> {
let work_path = PathBuf::from(&work_dir);
let job_path = work_path.join(job_id.as_str());
@@ -761,11 +766,30 @@ pub(crate) async fn remove_job_dir(
)));
}
- info!("Remove data for job {:?}", job_id);
+ // Empty stage list ⇒ remove the whole job dir (legacy behavior).
+ if remove_stage_ids.is_empty() {
+ info!("Remove data for job {:?}", job_id);
+ return tokio::fs::remove_dir_all(&job_path).await.map_err(|e| {
+ BallistaError::General(format!("Failed to remove {job_path:?} due
to {e}"))
+ });
+ }
- tokio::fs::remove_dir_all(&job_path).await.map_err(|e| {
- BallistaError::General(format!("Failed to remove {job_path:?} due to
{e}"))
- })?;
+ // Otherwise remove only the given (intermediate) stage subdirs.
+ for stage_id in remove_stage_ids {
+ let stage_path = job_path.join(stage_id.to_string());
+ if !tokio::fs::try_exists(&stage_path).await.unwrap_or(false) {
+ continue;
+ }
+ if !is_subdirectory(stage_path.as_path(), job_path.as_path()) {
+ return Err(BallistaError::General(format!(
+ "Path {stage_path:?} is not a subdirectory of {job_path:?}"
+ )));
+ }
+ info!("Remove intermediate data for job {job_id:?} stage {stage_id}");
+ tokio::fs::remove_dir_all(&stage_path).await.map_err(|e| {
+ BallistaError::General(format!("Failed to remove {stage_path:?}
due to {e}"))
+ })?;
+ }
Ok(())
}
@@ -886,6 +910,7 @@ mod tests {
use std::path::{Path, PathBuf};
use super::clean_shuffle_data_loop;
+ use super::remove_job_data;
use ballista_core::JobId;
use std::fs;
use std::fs::File;
@@ -984,6 +1009,43 @@ mod tests {
assert!(!is_subdirectory(&job_path, base_dir));
}
}
+ #[tokio::test]
+ async fn test_remove_intermediate_stage_data() {
+ let work_dir = TempDir::new().unwrap();
+ let work = work_dir.path();
+ let job_id: JobId = "job".into();
+
+ // Create job/1, job/2, job/3, each with a data file.
+ for stage in [1u32, 2, 3] {
+ let stage_dir = work.join("job").join(stage.to_string());
+ fs::create_dir_all(&stage_dir).unwrap();
+ File::create(stage_dir.join("data.arrow"))
+ .unwrap()
+ .write_all(b"x")
+ .unwrap();
+ }
+
+ // Remove intermediate stages 1 and 2; stage 3 (final) is retained.
+ remove_job_data(work.to_str().unwrap(), &job_id, &[1, 2])
+ .await
+ .unwrap();
+ assert!(!work.join("job").join("1").exists());
+ assert!(!work.join("job").join("2").exists());
+ assert!(work.join("job").join("3").exists());
+
+ // Removing a missing stage id is a no-op (idempotent).
+ remove_job_data(work.to_str().unwrap(), &job_id, &[99])
+ .await
+ .unwrap();
+ assert!(work.join("job").join("3").exists());
+
+ // Empty list removes the whole job dir.
+ remove_job_data(work.to_str().unwrap(), &job_id, &[])
+ .await
+ .unwrap();
+ assert!(!work.join("job").exists());
+ }
+
fn prepare_testing_job_directory(base_dir: &Path, job_id: &JobId) ->
PathBuf {
let mut path = base_dir.to_path_buf();
path.push(job_id.as_str());
diff --git a/ballista/executor/src/executor_server.rs
b/ballista/executor/src/executor_server.rs
index c7c564fe3..8e7b02354 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -63,7 +63,7 @@ use tokio::task::JoinHandle;
use crate::cpu_bound_executor::DedicatedExecutor;
use crate::executor::Executor;
-use crate::executor_process::{ExecutorProcessConfig, remove_job_dir};
+use crate::executor_process::{ExecutorProcessConfig, remove_job_data};
use crate::metrics::ExecutorMetricCollectionPolicy;
use crate::shutdown::ShutdownNotifier;
use crate::{TaskExecutionTimes, as_task_status};
@@ -933,9 +933,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorGrpc
&self,
request: Request<RemoveJobDataParams>,
) -> Result<Response<RemoveJobDataResult>, Status> {
- let job_id = request.into_inner().job_id.into();
+ let params = request.into_inner();
+ let job_id = params.job_id.into();
- remove_job_dir(&self.executor.work_dir, &job_id)
+ remove_job_data(&self.executor.work_dir, &job_id,
¶ms.remove_stage_ids)
.await
.map_err(|e| Status::invalid_argument(e.to_string()))?;
diff --git a/ballista/scheduler/src/scheduler_server/grpc.rs
b/ballista/scheduler/src/scheduler_server/grpc.rs
index ef1a93136..5c21fa806 100644
--- a/ballista/scheduler/src/scheduler_server/grpc.rs
+++ b/ballista/scheduler/src/scheduler_server/grpc.rs
@@ -161,8 +161,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerGrpc
.executor_manager
.drain_pending_cleanup_jobs(&executor_id)
.into_iter()
- .map(|job_id| CleanJobDataParams {
+ .map(|(job_id, remove_stage_ids)| CleanJobDataParams {
job_id: job_id.into_inner(),
+ remove_stage_ids,
})
.collect();
Ok(Response::new(PollWorkResult {
diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
index 9bf33307c..cd8e21db0 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -211,10 +211,18 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
.record_completed(&job_id, queued_at, completed_at);
info!("Job finished successfully: [{job_id}]");
- if let Err(e) =
self.state.task_manager.succeed_job(&job_id).await {
- error!("Fail to invoke succeed_job for job {job_id} due to
{e:?}");
- }
- self.state.clean_up_successful_job(job_id);
+ let intermediate_stage_ids =
+ match self.state.task_manager.succeed_job(&job_id).await {
+ Ok(ids) => ids,
+ Err(e) => {
+ error!(
+ "Fail to invoke succeed_job for job {job_id}
due to {e:?}"
+ );
+ vec![]
+ }
+ };
+ self.state
+ .clean_up_successful_job(job_id, intermediate_stage_ids);
}
QueryStageSchedulerEvent::JobRunningFailed {
job_id,
diff --git a/ballista/scheduler/src/state/execution_graph.rs
b/ballista/scheduler/src/state/execution_graph.rs
index 879a17296..68ab2f4a8 100644
--- a/ballista/scheduler/src/state/execution_graph.rs
+++ b/ballista/scheduler/src/state/execution_graph.rs
@@ -240,6 +240,16 @@ pub trait ExecutionGraph: Debug {
/// Exposes executions stages and stage id's
fn stages(&self) -> &HashMap<usize, ExecutionStage>;
+ /// Stage ids of all non-final (intermediate) stages — those whose
+ /// `output_links` is non-empty. The final stage(s) are excluded.
+ fn intermediate_stage_ids(&self) -> Vec<u32> {
+ self.stages()
+ .iter()
+ .filter(|(_, stage)| !stage.output_links().is_empty())
+ .map(|(stage_id, _)| *stage_id as u32)
+ .collect()
+ }
+
/// returns next task to run
/// (used for testing only)
#[cfg(test)]
@@ -1786,6 +1796,33 @@ mod test {
test_union_all_plan, test_union_plan,
};
+ #[tokio::test]
+ async fn test_intermediate_stage_ids() {
+ // A simple aggregation produces a 2-stage graph: one intermediate
+ // stage (non-empty output_links) feeding one final stage (empty
+ // output_links).
+ let graph = test_aggregation_plan(4).await;
+
+ assert_eq!(graph.stages().len(), 2);
+
+ // Exactly one final stage.
+ let final_count = graph
+ .stages()
+ .values()
+ .filter(|s| s.output_links().is_empty())
+ .count();
+ assert_eq!(final_count, 1);
+
+ // Intermediate = all - final = exactly one stage, and none of the
+ // returned ids is a final stage.
+ let intermediate = graph.intermediate_stage_ids();
+ assert_eq!(intermediate.len(), 1);
+ for id in &intermediate {
+ let stage = graph.stages().get(&(*id as usize)).unwrap();
+ assert!(!stage.output_links().is_empty());
+ }
+ }
+
#[tokio::test]
async fn test_fail_job_sets_end_time_and_failed_metadata() -> Result<()> {
let mut graph = test_aggregation_plan(4).await;
diff --git a/ballista/scheduler/src/state/execution_stage.rs
b/ballista/scheduler/src/state/execution_stage.rs
index 10440f357..eafc13af9 100644
--- a/ballista/scheduler/src/state/execution_stage.rs
+++ b/ballista/scheduler/src/state/execution_stage.rs
@@ -107,6 +107,18 @@ impl ExecutionStage {
ExecutionStage::Failed(stage) => stage.plan.as_ref(),
}
}
+
+ /// Get the output links for this stage. An empty slice means this is a
+ /// final stage in the `ExecutionGraph`.
+ pub fn output_links(&self) -> &[usize] {
+ match self {
+ ExecutionStage::UnResolved(stage) => &stage.output_links,
+ ExecutionStage::Resolved(stage) => &stage.output_links,
+ ExecutionStage::Running(stage) => &stage.output_links,
+ ExecutionStage::Successful(stage) => &stage.output_links,
+ ExecutionStage::Failed(stage) => &stage.output_links,
+ }
+ }
}
/// For a stage whose input stages are not all completed, we say it's a
unresolved stage
diff --git a/ballista/scheduler/src/state/executor_manager.rs
b/ballista/scheduler/src/state/executor_manager.rs
index 1871514a7..febb4084d 100644
--- a/ballista/scheduler/src/state/executor_manager.rs
+++ b/ballista/scheduler/src/state/executor_manager.rs
@@ -66,8 +66,9 @@ pub struct ExecutorManager {
config: Arc<SchedulerConfig>,
/// Cached gRPC clients for communicating with executors.
clients: ExecutorClients,
- /// Jobs pending cleanup on each executor.
- pending_cleanup_jobs: Arc<DashMap<String, HashSet<JobId>>>,
+ /// Per-executor pending cleanups: job id -> stage ids to remove
+ /// (empty stage ids ⇒ remove the whole job dir).
+ pending_cleanup_jobs: Arc<DashMap<String, HashMap<JobId, Vec<u32>>>>,
/// Configuration for gRPC client connections.
grpc_client_config: GrpcClientConfig,
}
@@ -191,7 +192,9 @@ impl ExecutorManager {
let executor_manager = self.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(clean_up_interval)).await;
- executor_manager.clean_up_job_data_inner(job_id).await;
+ executor_manager
+ .clean_up_job_data_inner(job_id, vec![])
+ .await;
});
}
@@ -199,13 +202,49 @@ impl ExecutorManager {
pub fn clean_up_job_data(&self, job_id: JobId) {
let executor_manager = self.clone();
tokio::spawn(async move {
- executor_manager.clean_up_job_data_inner(job_id).await;
+ executor_manager
+ .clean_up_job_data_inner(job_id, vec![])
+ .await;
+ });
+ }
+
+ /// Immediately reclaim intermediate-stage shuffle data for a successful
job,
+ /// retaining the rest of the job dir (e.g. the final-stage output) for the
+ /// existing delayed whole-job cleanup. No-op if `remove_stage_ids` is
empty.
+ ///
+ /// Contract: call this AT MOST ONCE per job, passing the complete set of
+ /// intermediate stage ids. It is not designed for per-stage invocation.
+ /// In poll-based scheduling the pending-cleanup value is a `Vec<u32>`
keyed
+ /// by job, and `HashMap::insert` REPLACES on key collision (it does not
+ /// merge). One intermediate call followed by the delayed whole-job cleanup
+ /// (empty `remove_stage_ids`) is correct: the empty/whole-job entry
+ /// supersedes and removes the entire job dir. But two DISTINCT partial
calls
+ /// for the same job before the executor polls would drop the earlier stage
+ /// ids. Do not "fix" this by merging/extending: an empty vec means "remove
+ /// the whole job dir", so extending could never represent whole-job
+ /// supersession.
+ pub(crate) fn clean_up_intermediate_job_data(
+ &self,
+ job_id: JobId,
+ remove_stage_ids: Vec<u32>,
+ ) {
+ if remove_stage_ids.is_empty() {
+ return;
+ }
+ let executor_manager = self.clone();
+ tokio::spawn(async move {
+ executor_manager
+ .clean_up_job_data_inner(job_id, remove_stage_ids)
+ .await;
});
}
/// 1. Push strategy: Send rpc to Executors to clean up the job data
/// 2. Poll strategy: Save cleanup job ids and send them to executors
- async fn clean_up_job_data_inner(&self, job_id: JobId) {
+ ///
+ /// `remove_stage_ids` empty ⇒ remove the whole job dir; non-empty ⇒ remove
+ /// only those stage subdirs.
+ async fn clean_up_job_data_inner(&self, job_id: JobId, remove_stage_ids:
Vec<u32>) {
let alive_executors = self.get_alive_executors();
for executor in alive_executors {
@@ -215,10 +254,12 @@ impl ExecutorManager {
if let Ok(mut client) =
self.get_client(&executor, &self.grpc_client_config).await
{
+ let remove_stage_ids = remove_stage_ids.clone();
tokio::spawn(async move {
if let Err(err) = client
.remove_job_data(RemoveJobDataParams {
job_id: job_id_clone,
+ remove_stage_ids,
})
.await
{
@@ -234,7 +275,7 @@ impl ExecutorManager {
self.pending_cleanup_jobs
.entry(executor)
.or_default()
- .insert(job_id.clone());
+ .insert(job_id.clone(), remove_stage_ids.clone());
}
}
}
@@ -387,10 +428,13 @@ impl ExecutorManager {
Ok(())
}
- pub(crate) fn drain_pending_cleanup_jobs(&self, executor_id: &str) ->
HashSet<JobId> {
+ pub(crate) fn drain_pending_cleanup_jobs(
+ &self,
+ executor_id: &str,
+ ) -> Vec<(JobId, Vec<u32>)> {
self.pending_cleanup_jobs
.remove(executor_id)
- .map(|(_, jobs)| jobs)
+ .map(|(_, jobs)| jobs.into_iter().collect())
.unwrap_or_default()
}
diff --git a/ballista/scheduler/src/state/mod.rs
b/ballista/scheduler/src/state/mod.rs
index ce007330e..4b8b4385d 100644
--- a/ballista/scheduler/src/state/mod.rs
+++ b/ballista/scheduler/src/state/mod.rs
@@ -382,8 +382,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
Ok(())
}
- /// Spawn a delayed future to clean up job data on both Scheduler and
Executors
- pub(crate) fn clean_up_successful_job(&self, job_id: JobId) {
+ /// Immediately reclaim intermediate shuffle data, then spawn a delayed
+ /// future to clean up the remaining job data (final-stage output + job
+ /// state) on both Scheduler and Executors.
+ pub(crate) fn clean_up_successful_job(
+ &self,
+ job_id: JobId,
+ intermediate_stage_ids: Vec<u32>,
+ ) {
+ self.executor_manager
+ .clean_up_intermediate_job_data(job_id.clone(),
intermediate_stage_ids);
self.executor_manager.clean_up_job_data_delayed(
job_id.clone(),
self.config.finished_job_data_clean_up_interval_seconds,
diff --git a/ballista/scheduler/src/state/task_manager.rs
b/ballista/scheduler/src/state/task_manager.rs
index 7f18d9581..5c4d90789 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -585,24 +585,25 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskManager<T, U>
Ok(events)
}
- /// Mark a job to success. This will create a key under the CompletedJobs
keyspace
- /// and remove the job from ActiveJobs
- pub(crate) async fn succeed_job(&self, job_id: &JobId) -> Result<()> {
+ /// Move a job from Active to Success and return the ids of its
intermediate
+ /// (non-final) stages so their shuffle data can be reclaimed immediately.
+ /// Returns an empty vec if the job is not found or not successful.
+ pub(crate) async fn succeed_job(&self, job_id: &JobId) -> Result<Vec<u32>>
{
debug!("Moving job {job_id} from Active to Success");
if let Some(graph) = self.remove_active_execution_graph(job_id) {
let graph = graph.read().await;
if graph.is_successful() {
self.state.save_job(job_id, &graph).await?;
+ Ok(graph.intermediate_stage_ids())
} else {
error!("Job {job_id} has not finished and cannot be
completed");
- return Ok(());
+ Ok(vec![])
}
} else {
warn!("Fail to find job {job_id} in the cache");
+ Ok(vec![])
}
-
- Ok(())
}
/// Cancel the job and return a Vec of running tasks need to cancel
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
