andygrove opened a new issue, #2150:
URL: https://github.com/apache/datafusion-ballista/issues/2150
### 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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
