apache / apache/datafusion-ballista

AQE: a replan that reassigns stage ids orphans the sibling leaf stage, and its task completion errors out

Open
#2,150 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.