QuakeWang opened a new issue, #2214:
URL: https://github.com/apache/datafusion-ballista/issues/2214
# StaticExecutionGraph drops stages when plan transformations fail
**Describe the bug**
Several `StaticExecutionGraph` state transitions remove a stage from
`self.stages` before performing a fallible physical-plan transformation.
The affected transitions are:
- `resolve_stage`
- `rollback_running_stage`
- `rollback_resolved_stage`
For example, `resolve_stage` currently follows this pattern:
```rust
if let Some(ExecutionStage::UnResolved(stage)) =
self.stages.remove(&stage_id)
{
self.stages.insert(
stage_id,
ExecutionStage::Resolved(
stage.to_resolved(self.session_config.options())?,
),
);
}
```
The stage is removed before `to_resolved` is evaluated. If the
transformation returns an error, `?` returns from the transition before
`insert` is called, so the original stage is dropped and the stage ID is left
absent from the execution graph.
The two rollback transitions have the same ordering around
`to_unresolved(...)`.
These transformations can fail while replacing shuffle nodes, rebuilding
physical-plan nodes, constructing shuffle readers, or applying physical
optimizer rules.
The error is propagated through the scheduler, but callers operate directly
on the active execution graph and do not restore a previous snapshot. A failed
transition can therefore mutate the graph even though no replacement stage was
successfully produced.
**To Reproduce**
A deterministic unit-level reproduction can use a test execution-plan node
whose `with_new_children` implementation returns an error when the shuffle
child is rewritten.
For each affected transition:
1. Construct a `StaticExecutionGraph` containing a stage in the required
initial state.
2. Record the stage count and the stage's state.
3. Invoke the transition with a plan transformation that returns an error.
4. Verify that the transition returns `Err`.
5. Inspect `graph.stages()` for the original stage ID.
With the current implementation, the stage ID is no longer present because
the stage was removed before the transformation failed.
The same outcome follows directly from the current control flow: `remove`
completes before the fallible transformation is evaluated, while `insert` is
skipped by the early return.
**Expected behavior**
Stage transitions should be failure-atomic.
If a physical-plan transformation fails:
- the transition should return the original error;
- the original stage should remain in the execution graph;
- its state, plan, task metadata, and attempt information should remain as
they were at the start of the transition; and
- no replacement stage should be committed.
A successful transformation should retain the existing transition behavior.
**Additional context**
This issue was found by inspecting the transition control flow. The exact
user-visible outcome depends on which physical-plan transformation fails, so
this report does not claim a specific production failure mode or occurrence
frequency.
The proposed scope is limited to the three `StaticExecutionGraph`
transitions in `ballista/scheduler/src/state/execution_graph.rs`:
- compute the replacement stage before committing it to `self.stages`;
- commit the replacement only after the fallible transformation succeeds; and
- add regression coverage for the error path of each transition.
This should not require changes to public APIs, scheduling policy, retry
policy, stage-attempt behavior, or stage dependency semantics. It also does not
propose making the surrounding task-update or executor-loss operations fully
transactional.
--
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]
