Skip to content

[EPIC] Share execution graph logic between StaticExecutionGraph and AdaptiveExecutionGraph #2595

Open

Description

@andygrove

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

StaticExecutionGraph (state/execution_graph.rs) and AdaptiveExecutionGraph (state/aqe/mod.rs) carry near-identical copies of most of their logic. Both files are among the most frequently changed in the scheduler (29 and 25 commits in the last 12 months), and the copies have already drifted:

Duplicated code:

What Static graph Adaptive graph
update_task_status L760-L1195, 376 lines, cognitive complexity 164, nesting 9 L781-L1142, 361 lines, cognitive complexity 121
reset_stages_internal L532-L655 L548-L676 (adds planner.remove_exchange_locations)
accessors, revive, is_successful, completed_stages, lost-executor reset tail L609-L766 L630-L787 (identical apart from type names)
stage transitions (resolve_stage, succeed_stage, fail_stage, rollback_*, rerun_successful_stage, fail_job, succeed_job) L1307-L1438 L1254-L1391
processing_stages_update, get_running_stage_id L390-L455, L504-L530 copied
pop_next_task (test only) L1484+ L1440+

Describe the solution you'd like

Proposed sub-tasks, each a behaviour-preserving PR (apart from the bug fixes) relying on the existing execution-graph tests:

  • Fix the AQE stage-loss bug (AQE: resolve_stage and rollback_*_stage drop the stage from the graph when the conversion fails #2594) so both copies behave the same before they are merged.
  • Extract the shared state (job/session ids, status, timestamps, stages, failed_stage_attempts, ...) into a common struct embedded by both graphs, and implement the accessors, revive, is_successful, completed_stages and lost-executor reset once.
  • Share the stage-transition methods and processing_stages_update / get_running_stage_id as default trait methods or free functions over the stages map.
  • Share reset_stages_internal as a function returning a reset plan; the AQE graph adds its remove_exchange_locations step afterwards. Replace the nested for_each closures with plain loops.
  • Break update_task_status into handle_failed_task, handle_successful_task, handle_late_task_for_unresolved_stage and the resubmit/reset post-passes, then share those between the graphs. This also resolves the AQE gaps described in AQE: handle executor failure on the adaptive execution path #1986.
  • Share the test-only pop_next_task.

Describe alternatives you've considered

Keep the copies and port each fix by hand. The drift above shows this is already failing.

Additional context

Related: #1986, #2539, #2150. Line references are pinned to commit ea1126c.

Activity

  1. added
    enhancementNew feature or request
    help wantedExtra attention is needed
    rustPull requests that update Rust code
    tech debtCode complexity, duplication or structure that makes the codebase harder to change
    on Oct 14, 2026
  2. milenkovicm commented on Oct 14, 2026

    @milenkovicm
    Contributor

    i will take blame for this, meant to clean them up when AQE get mature but can't really catch up with ballista this days

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesthelp wantedExtra attention is neededrustPull requests that update Rust codetech debtCode complexity, duplication or structure that makes the codebase harder to change

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions