You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/09/26 21:08:03 UTC
(datafusion-ballista) branch main updated: fix(scheduler): refund the vcores of running tasks when a job is aborted (#2477)
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 2fedccdcf fix(scheduler): refund the vcores of running tasks when a job is aborted (#2477)
2fedccdcf is described below
commit 2fedccdcf62f59b0cb480047c94b5867caf371c0
Author: Amit Vijapur <vi...@gmail.com>
AuthorDate: Sat Sep 26 22:07:58 2026 +0100
fix(scheduler): refund the vcores of running tasks when a job is aborted (#2477)
* fix(scheduler): refund the vcores of running tasks when a job is aborted
Under PushStaged the only refund happens in TaskUpdating, from the status
an executor reports. abort_job cancels the running tasks and then evicts
the job from the active cache, so that later report is dropped with
"Fail to find job ... in the active cache" and the vcores stay deducted
until the scheduler can no longer bind anything.
abort_job now computes the vcores its running tasks hold from the graph
snapshot and hands them to the cancel closure; the scheduler refunds them
once the executors have acknowledged the cancel.
Closes #2418
* Say in abort_job that the refund is made when the cancel is sent
---
.../src/scheduler_server/query_stage_scheduler.rs | 19 +++++++++++------
ballista/scheduler/src/state/task_manager.rs | 24 +++++++++++++++++-----
2 files changed, 32 insertions(+), 11 deletions(-)
diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
index e2273a777..6f4e0be22 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -109,15 +109,22 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryStageSchedul
async fn abort_job(&self, job_id: &JobId, failure_reason: String) -> Result<()> {
let executor_manager = self.state.executor_manager.clone();
+ let refund_vcores = self.state.config.is_push_staged_scheduling();
self.state
.task_manager
- .abort_job(job_id, failure_reason, move |running_tasks| async move {
- if running_tasks.is_empty() {
+ .abort_job(
+ job_id,
+ failure_reason,
+ move |running_tasks, freed_slots| async move {
+ if !running_tasks.is_empty() {
+ executor_manager.cancel_running_tasks(running_tasks).await?;
+ }
+ if refund_vcores && !freed_slots.is_empty() {
+ executor_manager.unbind_tasks(freed_slots).await?;
+ }
Ok(())
- } else {
- executor_manager.cancel_running_tasks(running_tasks).await
- }
- })
+ },
+ )
.await?;
Ok(())
}
diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs
index 0ed587f4e..e0ef644bf 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::cluster::{JobState, JobStateEventStream};
+use crate::cluster::{ExecutorSlot, JobState, JobStateEventStream};
use crate::config::SchedulerConfig;
use crate::planner::DefaultDistributedPlanner;
use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan};
@@ -647,7 +647,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
cancel_tasks: F,
) -> Result<usize>
where
- F: FnOnce(Vec<RunningTaskInfo>) -> Fut,
+ F: FnOnce(Vec<RunningTaskInfo>, Vec<ExecutorSlot>) -> Fut,
Fut: Future<Output = Result<()>>,
{
let Some(graph) = self.get_active_execution_graph(job_id) else {
@@ -669,13 +669,26 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
(running_tasks, pending_tasks, snapshot)
};
+ // The vcores those tasks hold on their executors. Their terminal
+ // status arrives after the job has left the active cache, so the
+ // `TaskUpdating` refund never sees them (#2418). They are refunded
+ // as soon as the cancel is sent, so an executor can be briefly
+ // oversubscribed while the cancelled tasks wind down.
+ let mut freed: HashMap<String, u32> = HashMap::new();
+ for task in &running_tasks {
+ if let Some(vcores) = snapshot.task_vcores(task.stage_id, task.task_id) {
+ *freed.entry(task.executor_id.clone()).or_default() += vcores;
+ }
+ }
+ let freed_slots: Vec<ExecutorSlot> = freed.into_iter().collect();
+
info!(
"Cancelling {} running tasks for job {}",
running_tasks.len(),
job_id
);
- let cancel_result = cancel_tasks(running_tasks).await;
+ let cancel_result = cancel_tasks(running_tasks, freed_slots).await;
let persist_result = self.persist_terminal_and_evict(job_id, &snapshot).await;
match (cancel_result, persist_result) {
(Ok(()), Ok(())) => Ok(pending_tasks),
@@ -1376,7 +1389,7 @@ mod tests {
.abort_job(
&job_id_for_abort,
"test failure".to_string(),
- move |tasks| async move {
+ move |tasks, slots| async move {
let status = manager_for_cancel
.get_job_status(&job_id_for_cancel)
.await?
@@ -1388,6 +1401,7 @@ mod tests {
"job must be terminal before executor tasks are cancelled"
);
cancelled_tasks_for_abort.store(tasks.len(), Ordering::SeqCst);
+ assert_eq!(slots, vec![("executor-1".to_string(), 1)]);
Ok(())
},
)
@@ -1437,7 +1451,7 @@ mod tests {
.abort_job(
&job_id_for_abort,
"test failure".to_string(),
- move |tasks| async move {
+ move |tasks, _slots| async move {
cancelled_tasks_for_abort.store(tasks.len(), Ordering::SeqCst);
Ok(())
},
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org