apache / apache/datafusion-ballista
AQE: a replan that reassigns stage ids orphans the sibling leaf stage, and its task completion errors out
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
### Background
Under `AdaptiveExecutionGraph`, finalising a stage can trigger a replan that changes the set of stages. `AdaptivePlanner::cancel_stage` drops the stage from `runnable_stage_cache` (`ballista/scheduler/src/state/aqe/planner.rs:197`), but nothing on the graph side transitions that stage out of `Running` or cancels its in-flight tasks. `update_stage_progress` just logs and moves on:
```rust
if !stages_to_cancel.is_empty() {
warn!(
"there are stages to be cancelled but its not implemented. stages to cancel: {:?}",
stages_to_cancel
);
}
```
(`ballista/scheduler/src/state/aqe/mod.rs:825`)
### Problem
When the orphaned stage's tasks later report success, `update_stage_progress` calls `AdaptivePlanner::finalise_stage`, which no longer finds the stage in `runnable_stage_cache` and fails:
```
Err(DataFusionError(Execution("Can't find active cache resolve")))
```
(`ballista/scheduler/src/state/aqe/planner.rs:270`)
Two consequences:
1. **The job can never complete.** The error propagates out of `update_task_status` before the stage is transitioned, so the orphaned stage is left `Running` with every one of its tasks already reported successful and zero available tasks. `is_successful()` requires *all* stages to be `Successful`, so it stays false forever. The job neither finishes nor fails.
2. **Unrelated task updates are dropped.** `TaskManager::update_task_statuses` groups the incoming batch by job and uses `?` inside that loop (`ballista/scheduler/src/state/task_manager.rs:568`), so one bad job aborts the whole batch. `query_stage_scheduler` only logs `Failed to update N task statuses for Executor X` and discards the rest, including statuses belonging to other jobs and any events already accumulated.
### Reproduction
A join over two leaf stages. Completing stage 0 triggers a replan that creates stages 2 and 3 and drops stage 1 from the planner cache, while stage 1 remains `Running` in the graph. Reporting stage 1's tasks as successful then errors on the task that completes the stage:
```
initial stages: [0, 1]
popped tasks for stages: [0, 0, 1, 1]
stages after stage 0 finished: [0, 1, 2, 3]
stage 0: Successful
stage 1: Running
stage 2: Resolved
stage 3: Resolved
stage 1 task completion result: Ok([])
stage 1 task completion result: Err(DataFusionError(Execution("Can't find active cache resolve")))
is_successful after error: false
final stage 0: Successful
final stage 1: Running
final stage 2: Running
final stage 3: Running
running_stages: [2, 3, 1]
available_tasks: 4
```
The scratch test that produced this builds an `AdaptiveExecutionGraph` for `left JOIN right ... GROUP BY left.id` with `target_partitions = 2`, pops every task of both leaf stages up front, completes only stage 0's tasks, then reports stage 1's tasks as successful.
### Proposed fix
Handle `stages_to_cancel` on the graph side rather than logging it: transition the cancelled stage out of `Running` (or remove it from `stages` entirely so `is_successful()` is not blocked), and emit `CancelTasks` for its in-flight tasks. Late task statuses arriving for a cancelled stage should be ignored rather than routed into `finalise_stage`.
Separately, `TaskManager::update_task_statuses` should isolate per-job failures so one job's error cannot discard another job's task updates in the same batch.
### Notes
Found while writing an AQE test for #1996. Not caused by that change, and #2149 sidesteps it by testing with an aggregation chain instead of a join.
Contributor guide
Research direction
Start with ballista/scheduler/src/state/aqe/mod.rs around update_stage_progress and ballista/scheduler/src/state/aqe/planner.rs around cancel_stage and finalise_stage. Reproduce the two-leaf join scenario described in the issue, then inspect ballista/scheduler/src/state/task_manager.rs around update_task_statuses. Done means cancelled stages no longer block completion, late statuses are ignored, cancellation events are emitted, and one job's update error does not discard other jobs' updates.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100