Repository navigation
fix(scheduler): keep AQE stage in the graph when a stage transition fails - #2613
Open
DeviousCardi wants to merge 1 commit into
Open
DeviousCardi wants to merge 1 commit into
DeviousCardi wants to merge 1 commit into
Conversation
…ails AdaptiveExecutionGraph::resolve_stage, rollback_running_stage and rollback_resolved_stage removed the stage from `stages` before running the fallible plan rewrite, so an error dropped the stage from the graph. Look the stage up, build the new stage, and only then insert it, as StaticExecutionGraph already does. The FailingPlanRewriteExec test helper moves from the static graph's tests to test_utils so the adaptive graph's tests can use it too. Closes apache#2594 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y1b55sEL1gA1c3FCMcQCbm
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #2594.
Rationale for this change
AdaptiveExecutionGraph::resolve_stage,rollback_running_stageandrollback_resolved_stageremoved the stage fromstagesbefore running the fallible plan rewrite (to_resolved/to_unresolved). If the rewrite failed,?returned early and the stage was never put back, so later lookups failed with "stage not found" instead of the original error.StaticExecutionGraphalready avoids this.What changes are included in this PR?
getthe stage, build the new one, andinsertit only once that succeeds, the same wayStaticExecutionGraphdoes it.succeed_stage,fail_stageandrerun_successful_stagekeepremovebecause their conversions can't fail.FailingPlanRewriteExec/fail_plan_rewritesmove from the static graph's test module totest_utils.rsso both graphs can use them. This is a move only; the static graph's tests are otherwise unchanged.state/aqe/test/stage_transitions.rswith three tests that mirror the static graph's tests. Each one forces the plan rewrite to fail and checks that the stage is still in the graph in its original state (same plan, attempt number, and pending and running tasks). All three fail onmainand pass with this change.cargo test -p ballista-schedulerpasses, and so doescargo clippy -p ballista-scheduler --all-targets --all-features -- -D warnings.Are there any user-facing changes?
No.
AI disclosure: Claude Code prepared the implementation and tests. I reviewed them and will reply to review comments myself.
🤖 Generated with Claude Code
https://claude.ai/code/session_01Y1b55sEL1gA1c3FCMcQCbm