andygrove opened a new issue, #1996:
URL: https://github.com/apache/datafusion-ballista/issues/1996
### Background
#1982 reclaims a job's intermediate shuffle data immediately on success
(instead of waiting for `finished_job_data_clean_up_interval_seconds`), to
relieve executor disk pressure (#1981). It identifies intermediate stages via
`ExecutionGraph::intermediate_stage_ids()`, which selects stages whose
`output_links` is non-empty (the final stage has empty `output_links`).
### Problem
This works for the **static** planner (`StaticExecutionGraph`), where
`ExecutionStageBuilder` populates `output_links` up front. It does **not** work
under **AQE** (`AdaptiveExecutionGraph`): stages are created incrementally by
`create_resolved_stage`, which always sets `output_links` to `vec![]`
(`ballista/scheduler/src/state/aqe/mod.rs:212`, *"we do not know output links
at this moment"*), and no AQE code path ever back-fills it.
Consequently, under AQE every stage has empty `output_links`, so
`intermediate_stage_ids()` returns an empty set and the immediate reclaim does
nothing — the job falls back to the existing delayed whole-job cleanup. This is
**safe** (the "never delete a final stage" invariant holds trivially), but
delivers **no disk benefit under AQE**. That matters because the motivating
workload (the SF10 CI job in #1981) runs all 22 queries **both AQE off and
on**, so the AQE-on half gets no relief.
### Proposed fix
Override `intermediate_stage_ids()` for `AdaptiveExecutionGraph` to derive
final stages from `output_locations` rather than `output_links`: each
`PartitionLocation` in `output_locations` carries its terminal `stage_id`, so
`final = {stage_ids in output_locations}` and `intermediate = all_stages −
final`. Guard: if `output_locations` is empty, return an empty set (reclaim
nothing) to preserve the safety invariant.
### Tests
Add an AQE test that drives a multi-stage adaptive job to success and
asserts `intermediate_stage_ids()` returns exactly the non-terminal stages (and
never a stage present in `output_locations`).
Follow-up to #1982. Related: #1981.
--
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]
