You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "goingforstudying-ctrl (via GitHub)" <gi...@apache.org> on 2026/07/27 08:29:58 UTC
[PR] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks [datafusion-ballista]
goingforstudying-ctrl opened a new pull request, #2192:
URL: https://github.com/apache/datafusion-ballista/pull/2192
## Which issue does this PR close?
Closes #2150.
## Rationale for this change
When the adaptive query planner replans after a stage completes, it may decide that a stage which is **still running** is no longer needed — e.g. the build side of a join produced no data, so the probe side is moot. The planner correctly returns those stage ids as "cancelled", but the scheduler never acted on that list: it only logged
```
there are stages to be cancelled but its not implemented. stages to cancel: [...]
```
The `Running` stage stayed in the graph, so the job kept waiting on it and could sit forever, and a late task completion from the cancelled stage failed the whole status update with `Invalid stage ID ... is not running`.
## What changes are included in this PR?
- **`AdaptiveExecutionGraph::update_stage_progress`** now retires each stage the replan cancels, via a new `retire_cancelled_stage`:
- a `Running` stage is moved to `Failed(TaskKilled)` (cancelling in-flight tasks) so the job no longer waits on it;
- a `Resolved`/`UnResolved` stage is dropped outright (it has no in-flight tasks);
- a terminal (`Successful`/`Failed`) stage is left untouched.
- **`update_task_status`** treats a late task status for a retired (`Failed`) stage as stale and discards it with a warning instead of erroring, so the rest of the batch still processes.
- **`TaskManager`** no longer aborts the entire grouped update when one stage id is unexpected — it logs and continues with the remaining stages.
- Added a lifecycle test (`test_replan_cancelled_stage_is_retired_and_late_task_discarded`) reproducing the orphan-stage scenario end to end.
## Are these changes tested?
Yes. New unit test drives a join through `revive`, completes the build side with empty output (triggering the replan that cancels the probe stage), then asserts the probe stage is retired and that a late task completion from it is discarded rather than failing the update.
Full `state::` suite: **195 passed, 0 failed**. `cargo clippy -p ballista-scheduler --lib` and `cargo fmt` are clean.
## Are there any user-facing changes?
No API changes; the scheduler no longer wedges on AQE-cancelled stages.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2192:
URL: https://github.com/apache/datafusion-ballista/pull/2192#issuecomment-5969853566
**[P1] Preserve retired-task reservations after job completion**
The retained reservations live inside the execution graph, but retirement now allows the job to finish while redundant tasks are still running. `TaskManager::succeed_job()` then evicts the graph, and subsequent completions contribute **0** in `sum_vcores_for_statuses()`. Under push scheduling, those reserved vcores are never returned; repeated empty-side joins can eventually exhaust scheduler-visible executor capacity.
Although post-eviction refund handling is unchanged, this PR makes that path reachable for successfully completed jobs with retired tasks still in flight.
I reproduced this by holding a probe task, completing the empty sibling and remaining stages, calling `succeed_job()`, then checking the held completion’s refund: **1 before eviction, 0 afterward**. The current lifecycle test checks the reservation before eviction, so it misses this case.
Could we retain outstanding reservation accounting independently of the active graph until completion/cancellation is acknowledged, and extend the regression test to cover completion after successful-job eviction?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks [datafusion-ballista]
Posted by "goingforstudying-ctrl (via GitHub)" <gi...@apache.org>.
goingforstudying-ctrl commented on PR #2192:
URL: https://github.com/apache/datafusion-ballista/pull/2192#issuecomment-5973090333
Pushed b476fe16. Reservations are now tracked outside the execution graph, keyed by executor and stage attempt, and consumed exactly once across terminal reports, launch rejection, and cancellation. Reservations from a lost executor are discarded, so a job can complete and evict its graph while redundant tasks are still running without leaking vcores.
Added `reservations_tests.rs` covering the post-eviction completion case you reproduced, plus duplicate and weighted refunds, stale attempts, executor loss, and cancellation/rejection overlap. CI is green.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks [datafusion-ballista]
Posted by "goingforstudying-ctrl (via GitHub)" <gi...@apache.org>.
goingforstudying-ctrl commented on PR #2192:
URL: https://github.com/apache/datafusion-ballista/pull/2192#issuecomment-5609896884
Agreed that client-visible reporting of retired stages would be useful as a separate follow-up. This PR keeps that API unchanged. The current revision also preserves retired stages’ task-vcore metadata so late completions can refund reservations while the job graph is active. Already dispatched tasks may finish; their results are discarded rather than cancelled through executor RPCs.
At `838aaa99`, the scheduler state suite passed locally with 266 tests and one ignored, plus the launch-rejection regression and Clippy. The description now states the retirement behavior and its limits explicitly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix(scheduler): retire AQE stages cancelled by replan and discard late tasks [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #2192:
URL: https://github.com/apache/datafusion-ballista/pull/2192
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org