andygrove commented on code in PR #2730:
URL:
https://github.com/apache/datafusion-ballista/pull/2730#discussion_r4486156891
##########
ballista/scheduler/src/state/execution_graph.rs:
##########
@@ -893,69 +935,96 @@ impl ExecutionGraph for StaticExecutionGraph {
.entry(stage_id)
.or_default();
failed_attempts.insert(task_stage_attempt_num);
+ let map_stage_id =
+ fetch_partiton_error.map_stage_id as
usize;
+ let map_partition_id =
+ fetch_partiton_error.map_partition_id
as usize;
+ let executor_id =
fetch_partiton_error.executor_id;
+ // A cache's build stage id can match one
of this
+ // job's: the cache inputs the failure
names, and
+ // whether an own input on that executor
does too.
+ let lost_cache = external_fetch_failures(
+ &running_stage.inputs,
+ &executor_id,
+ map_stage_id,
+ map_partition_id,
+ );
+ let lost_own = lost_cache.is_empty()
Review Comment:
I think there's a problem when two tasks fail fetching the same lost cache
data in one status batch.
Say two tasks of the stage reading cache input `EXT` both fail fetching from
dead executor E, and their statuses arrive together. The first failure removes
every E location from `EXT`. For the second one, `external_fetch_failures` now
finds nothing, because those locations are already gone, so `lost_own =
lost_cache.is_empty()` comes out true.
We then call `running_stage.remove_input_partitions(map_stage_id, ...)`, but
`map_stage_id` here is the cache build job's stage id, not one of this graph's
inputs. That returns `Internal("... is not a valid child stage ID")` and the
`?` propagates it. As far as I can tell `TaskManager` logs it and drops the
whole batch, so there's no rollback and no `ExternalInputLost`, and the stage
input is left half updated. The job never recovers. If the build's stage id
happens to match one of the job's own input ids, we'd instead re-run the job's
own producer for no reason.
Before this PR the second failure hit the "Stages was marked failed, ignore
FetchPartitionError" branch, so this looks like a regression. Maybe `lost_own`
should only be true when `map_stage_id` is actually one of
`running_stage.inputs`, and a repeat failure on an already-removed cache
location should just be ignored? A graph test with two fetch failures from the
same executor in one `update_task_status` call would pin it down.
##########
ballista/scheduler/src/state/cache_linker.rs:
##########
@@ -734,6 +1005,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
}
}
+/// Whether a job with `reads` reads cache `key` more than once: directly, or
+/// through another cache built from it.
+fn read_more_than_once(reads: &[ResolvedRead], key: &CacheKey) -> bool {
Review Comment:
I think this guard misses one direction of nesting. It counts reads of the
lost key and caches built from it, but not the case where the lost cache was
itself built from another cache the job also reads directly.
Example: a job reads `c1` directly (that stage already finished) and also
reads `c2`, which was cached from a plan over `c1`. E holds parts of both and
dies, but only `c2`'s input is reported lost, so this counts one read and the
rebuild goes ahead. If I'm reading `prepare_caches` right, rebuilding `c2`
recurses into its subplan, `claim_or_hit` finds `c1`'s entry on dead E,
invalidates it, and rebuilds `c1` too. Then `c2` is rebuilt from the new `c1`
while the job already consumed the old one. With a non-reproducible subplan
(`random()`, or `LIMIT` without `ORDER BY`) the query ends up mixing two
versions of `c1`, which is exactly what this guard is meant to prevent.
Would it work to also fail when any cache the lost one was built from is
read by the job, or to check after preparing the rebuild whether it rebuilt
anything the job already read?
##########
ballista/scheduler/src/state/cache_linker.rs:
##########
@@ -575,6 +686,156 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
}
}
+ /// A running job lost the cache data it reads under `inputs`, and its
stages
+ /// wait for it: rebuilds those caches and provides them again, or fails
the
+ /// job. On its own task, off the event loop.
+ pub(crate) fn on_external_input_lost(
+ self: &Arc<Self>,
+ job_id: JobId,
+ inputs: Vec<usize>,
+ sender: &EventSender<QueryStageSchedulerEvent>,
+ ) {
+ let state = Arc::clone(self);
+ let sender = sender.clone();
+ tokio::spawn(async move {
+ let Some(job_reads) = state.cache_linker.reads_of(&job_id) else {
+ return; // the job ended meanwhile
+ };
+ if let Err(reason) = state
+ .rebuild_lost_inputs(&job_id, &job_reads, &inputs, &sender)
+ .await
+ {
+ state
+ .post_running_failure(&job_id, job_reads.queued_at,
&reason, &sender)
+ .await;
+ }
+ });
+ }
+
+ async fn rebuild_lost_inputs(
+ &self,
+ job_id: &JobId,
+ job_reads: &JobReads,
+ inputs: &[usize],
+ sender: &EventSender<QueryStageSchedulerEvent>,
+ ) -> std::result::Result<(), String> {
+ for input in inputs {
+ let Some(read) = job_reads.read(*input) else {
+ return Err(format!(
+ "external input {input} is not a cache this scheduler can
rebuild"
+ ));
+ };
+ // A rebuild would give the query a second version of the cache,
next
+ // to the one it already read elsewhere.
+ if read_more_than_once(&job_reads.reads, &read.key) {
Review Comment:
Small one: the read-more-than-once check and the missing-read check run per
input inside the loop, after earlier inputs have already been invalidated and
their rebuilds claimed and submitted.
So with inputs `[A, B]`, A passes, its build is dropped and a full rebuild
starts under the job's admission, then B fails the guard and the job fails. The
rebuild of A keeps running for a job that's gone, and A's entry was dropped
even though it was healthy on the other executors. Validating every input
first, before invalidating or claiming anything, would avoid that.
##########
ballista/scheduler/src/state/cache_linker.rs:
##########
@@ -98,16 +98,59 @@ enum CacheClaim {
/// A concurrent submission can change a key between lookup and claim.
const CLAIM_ATTEMPTS: usize = 4;
+/// What waits on a build job.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+enum Waiter {
+ /// A held job.
+ Held(JobId),
+ /// A running job's lost cache read, by external stage id.
+ Input(JobId, usize),
+}
+
+/// What a build job's end releases, or fails.
+#[derive(Default)]
+pub(crate) struct Released {
+ /// Held jobs.
+ pub(crate) jobs: Vec<PendingJob>,
+ /// Running jobs' lost reads, by external stage id.
+ pub(crate) inputs: Vec<(JobId, usize)>,
+}
+
+/// A running job's cache reads, kept to rebuild a cache whose data it loses.
+#[derive(Clone)]
+pub(crate) struct JobReads {
+ pub(crate) ctx: Arc<SessionContext>,
+ pub(crate) queued_at: u64,
+ /// The submission's admission, for rebuilds; dropped when the job ends.
+ pub(crate) admission: Admission,
Review Comment:
I want to check the memory and lifetime side of this. `JobReads` keeps the
job's `SessionContext`, its `Admission`, and each cache's full `LogicalPlan`
subplan (via `ResolvedRead.cache`) alive for the job's whole run. Before this,
the subplan was dropped once the read was resolved. That's extra scheduler
memory for every running cache-reading job, which we've been strict about.
On `Admission`, I see the design doc wants shutdown to wait for rebuilds,
which makes sense. But `Admission` is documented in `event.rs` as held until
the job is running or fails to plan, so that doc needs updating. Also, now any
end path that skips `on_job_terminal` would leak it and the scheduler would
never report idle. Could we keep only what a rebuild actually needs (the cache
key and enough to re-plan it), and double check that every terminal path drops
the record?
##########
ballista/scheduler/src/state/execution_stage.rs:
##########
@@ -1491,28 +1493,109 @@ pub(crate) fn reads_external_data_on(
})
}
-/// Whether a failed fetch was for external (cache) data in `inputs`: one of
-/// their locations is on `executor_id` with the stage and partition the
executor
-/// reported. The report has no job id, so when one of the job's own locations
-/// matches too, the failure counts as external.
-pub(crate) fn is_external_fetch_failure(
+/// The external (cache) inputs in `inputs` a failed fetch was for: those with
a
+/// location on `executor_id` at the stage and partition the executor reported.
+/// The report has no job id, so one of the job's own inputs may match too.
+pub(crate) fn external_fetch_failures(
inputs: &HashMap<usize, StageOutput>,
executor_id: &str,
stage_id: usize,
partition_id: usize,
-) -> bool {
- inputs.iter().any(|(id, output)| {
- is_external_stage_id(*id)
- && output
- .partition_locations
- .values()
- .flatten()
- .any(|location| {
- location.executor_meta.id == executor_id
- && location.partition_id.stage_id == stage_id
- && location.partition_id.partition_id == partition_id
- })
- })
+) -> Vec<usize> {
+ let mut lost: Vec<usize> = inputs
+ .iter()
+ .filter(|(id, output)| {
+ is_external_stage_id(**id)
+ && output
+ .partition_locations
+ .values()
+ .flatten()
+ .any(|location| {
+ location.executor_meta.id == executor_id
+ && location.partition_id.stage_id == stage_id
+ && location.partition_id.partition_id ==
partition_id
+ })
+ })
+ .map(|(id, _)| *id)
+ .collect();
+ lost.sort_unstable();
+ lost
+}
+
+/// The partitions `plan`'s reader of external input `input_id` declares.
+pub(crate) fn external_reader_partitions(
+ plan: &dyn ExecutionPlan,
+ input_id: usize,
+) -> Option<usize> {
+ if let Some(reader) = plan.downcast_ref::<ShuffleReaderExec>()
+ && reader.stage_id == input_id
+ {
+ return Some(reader.partition.len());
+ }
+ if let Some(unresolved) = plan.downcast_ref::<UnresolvedShuffleExec>()
+ && unresolved.stage_id == input_id
+ {
+ return Some(unresolved.output_partition_count);
+ }
+ plan.children()
+ .into_iter()
+ .find_map(|child| external_reader_partitions(child.as_ref(), input_id))
+}
+
+/// `plan` with its reader of external input `input_id` reading `locations`,
+/// for a stage whose plan keeps its readers rather than resolving them from
+/// its inputs.
+pub(crate) fn with_external_locations(
+ plan: Arc<dyn ExecutionPlan>,
+ input_id: usize,
+ locations: &StageOutput,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let rewritten = plan.transform(|node| {
+ let Some(reader) = node.downcast_ref::<ShuffleReaderExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ if reader.stage_id != input_id {
+ return Ok(Transformed::no(node));
+ }
+ let read = ShuffleReaderExec::try_new(
Review Comment:
Not a bug today as far as I can tell, since the cache linker also builds
these readers with plain `try_new`. But this drops whatever else the replaced
reader had (broadcast, coalesced, skew join, `work_dir`, client pool). If a
lowering ever starts touching external readers, a broadcast reader would
quietly become per-partition here and give wrong join results. Maybe add an
assert or a debug check that the reader being replaced is a plain one, so it
fails loudly if that assumption ever changes?
##########
ballista/scheduler/src/state/pipelined_graph.rs:
##########
@@ -867,15 +870,102 @@ impl ExecutionGraph for PipelinedExecutionGraph {
vec![]
}
- /// Nothing is rolled back: losing an executor fails the job instead (see
- /// [`PipelinedExecutionGraph::lost_executor`]).
+ /// Nothing is rolled back: once placed, losing an executor fails the job
+ /// instead (see [`PipelinedExecutionGraph::lost_executor`]). Before then
+ /// no task has run, so a cache read loses only its locations on the
+ /// executor, and the job waits for the cache to be rebuilt.
fn reset_stages_on_lost_executor(
&mut self,
- _executor_id: &str,
+ executor_id: &str,
) -> Result<(HashSet<usize>, Vec<RunningTaskInfo>)> {
+ if !self.placed {
+ for stage in self.stages.values_mut() {
+ let inputs = match stage {
+ ExecutionStage::UnResolved(stage) => &mut stage.inputs,
+ ExecutionStage::Resolved(stage) => &mut stage.inputs,
+ _ => continue,
+ };
+ for (_, output) in inputs
+ .iter_mut()
+ .filter(|(id, _)| is_external_stage_id(**id))
+ {
+ for locations in output.partition_locations.values_mut() {
+ let before = locations.len();
+ locations.retain(|loc| loc.executor_meta.id !=
executor_id);
+ if locations.len() < before {
+ output.complete = false;
+ }
+ }
+ }
+ }
+ }
Ok((HashSet::new(), vec![]))
}
+ /// Stages are resolved before placement, so a lost cache read can leave a
+ /// resolved stage waiting too.
+ fn waiting_external_inputs(&self) -> Vec<usize> {
Review Comment:
Nit: this is a copy of the trait default with a `Resolved` arm added, so the
two filter/sort/dedup pipelines have to be kept in sync. Since a static graph
never has an incomplete external input on a `Resolved` stage, I think the
default could just cover `Resolved` too and this override could go. A shared
helper would also work.
##########
ballista/scheduler/src/state/cache_linker.rs:
##########
@@ -575,6 +686,156 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
}
}
+ /// A running job lost the cache data it reads under `inputs`, and its
stages
+ /// wait for it: rebuilds those caches and provides them again, or fails
the
+ /// job. On its own task, off the event loop.
+ pub(crate) fn on_external_input_lost(
+ self: &Arc<Self>,
+ job_id: JobId,
+ inputs: Vec<usize>,
+ sender: &EventSender<QueryStageSchedulerEvent>,
+ ) {
+ let state = Arc::clone(self);
+ let sender = sender.clone();
+ tokio::spawn(async move {
+ let Some(job_reads) = state.cache_linker.reads_of(&job_id) else {
+ return; // the job ended meanwhile
Review Comment:
If `reads_of` comes back `None` here, we return silently. By this point the
graph has already rolled the stage back and marked the external input
incomplete, so if the job is actually still running, nothing rebuilds it,
resolves it, or fails it, and the query hangs with no error to the client.
I'm not sure there's a path today where a running cache reader has no read
record, so this may be theoretical. It still feels safer to check whether the
job is still active and post `JobRunningFailed` if so, rather than assume it
ended. Same thing in `provide_input` below.
##########
chaos-testing/src/cluster.rs:
##########
@@ -914,12 +899,6 @@ impl TestCluster {
}
/// Stage ids come back from the REST API as strings.
Review Comment:
Looks like this doc comment got left behind when `find_stage` moved to
`rest.rs`, so it now documents `dir_has_entries`. Should move with `find_stage`
or go.
##########
ballista/scheduler/src/state/pipelined_graph.rs:
##########
@@ -349,11 +350,12 @@ impl PipelinedExecutionGraph {
matches!(self.status.status, Some(Status::Running(_)))
}
- /// Whether the job still waits to be placed: running, and not placed yet.
- /// A job can stop waiting between planning its placement and committing
- /// it, if it is cancelled or fails meanwhile.
+ /// Whether the job still waits to be placed: running, not placed yet, and
+ /// not waiting for lost `DataFrame::cache()` data to be rebuilt. A job can
+ /// stop waiting between planning its placement and committing it, if it
+ /// is cancelled or fails meanwhile.
pub fn awaits_placement(&self) -> bool {
- !self.placed && self.is_running()
+ !self.placed && self.is_running() &&
self.waiting_external_inputs().is_empty()
Review Comment:
Nit: `awaits_placement` is called for every unplaced pipelined job on each
offer round (`placement_requirements`, `place`, and the offer loop in
`state/mod.rs`), and `waiting_external_inputs()` allocates, sorts and dedups a
Vec just so we can check it's empty. An `.any(|(id, out)|
is_external_stage_id(*id) && !out.is_complete())` over the stage inputs would
short-circuit without allocating.
##########
docs/source/contributors-guide/dataframe-cache.md:
##########
@@ -141,17 +133,47 @@ instead of failing, and the scheduler rebuilds the cache
and resumes the stage:
stage resolves and re-runs → results to the client
```
-- A failed fetch is reported under the reader's stage id — the external id for
a
- cache read — so the graph knows exactly which input was lost.
-- Readers that lose the same cache share one rebuild. A rebuild may have a
- different partition count; the reader declares unknown partitioning, so the
- new partitions are folded into the old count.
-- Retries count towards `stage_max_failures`, like any stage's. If the rebuild
- fails, the job fails with the reason.
-- The cache is rebuilt in full and the reading stage re-runs; the job keeps its
- other stages. See [Per-partition recovery](#per-partition-recovery).
-- The adaptive planner cannot yet re-resolve a stage it rolls back (for any
- job), so under it a job that loses cache data still fails cleanly.
+- **Which data was lost.** A failed fetch is cache data when it names one of
+ the stage's cache locations (executor, build stage, partition). The report
+ carries no job id, so when one of the job's own inputs on that executor
+ matches too, both are recovered: the own producer re-runs its tasks there,
and
+ the cache is rebuilt. On executor loss, Ballista's reset already drops the
+ executor's locations from every stage input, cache reads included, and rolls
+ back their readers; the cache inputs it leaves incomplete are reported.
+- **No producer re-run.** For an own input the graph re-runs the producer's
lost
+ tasks; a cache's producer is another job, so the graph reports the input lost
+ and the stage waits until it is provided again.
+- **Shared rebuilds.** Readers that lose the same cache share one rebuild, and
+ one that finds a newer build resumes on it at once. A rebuild may have
+ another partition count; the reader declares unknown partitioning, so the new
+ partitions are folded into the old count. A rebuild must write the schema the
+ read was planned with, as the read picks its columns by position; otherwise
+ the job fails.
+- **Retries and failure.** Retries count towards `stage_max_failures`, like any
+ stage's; a stage out of attempts, or a rebuild that fails, fails the job with
Review Comment:
I don't think this matches the code for a failed rebuild. `fail_cache_job`
reports "a cache this query reads could not be built: ..." without
`CACHE_DATA_LOST`, so anyone matching on "lost DataFrame::cache() data" won't
see it. It also means `on_job_terminal` doesn't take its `CACHE_DATA_LOST`
invalidation branch for that failure. Either the message or the doc should
change so they agree.
--
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]
