andygrove commented on code in PR #2633:
URL:
https://github.com/apache/datafusion-ballista/pull/2633#discussion_r4372036765
##########
ballista/scheduler/src/state/execution_stage.rs:
##########
@@ -1379,6 +1380,48 @@ fn stage_input_partitions(plan: &Arc<dyn ExecutionPlan>)
-> usize {
.max(1)
}
+/// The inputs `plan` reads from outside its job, e.g. a cache read: a complete
+/// [`StageOutput`] per external [`ShuffleReaderExec`]. As stage inputs,
rollback
+/// and re-resolution restore them like any other input.
+pub(crate) fn external_inputs(plan: &dyn ExecutionPlan) -> HashMap<usize,
StageOutput> {
+ fn collect(plan: &dyn ExecutionPlan, inputs: &mut HashMap<usize,
StageOutput>) {
+ if let Some(reader) = plan.downcast_ref::<ShuffleReaderExec>()
+ && is_external_stage_id(reader.stage_id)
+ {
+ let output = inputs.entry(reader.stage_id).or_default();
+ for (partition, locations) in reader.partition.iter().enumerate() {
+ output
+ .partition_locations
+ .entry(partition)
+ .or_default()
+ .extend(locations.iter().cloned());
+ }
+ output.complete = true;
+ }
+ for child in plan.children() {
+ collect(child.as_ref(), inputs);
+ }
+ }
+ let mut inputs = HashMap::new();
+ collect(plan, &mut inputs);
+ inputs
+}
+
+/// Whether `inputs` hold external (cache) data on `executor_id`.
+pub(crate) fn reads_external_data_on(
Review Comment:
I think this check is a bit too broad. It matches any external input with
data on the failing executor, so if a stage's own shuffle fetch from E1 fails
for a transient reason while E1 also holds part of a cache, the job fails with
"lost DataFrame::cache() data" instead of rerunning the input stage.
`on_job_terminal` then drops every cache the job read too. On small clusters
the cache and the job's own output almost always share executors, so this could
hit most cache reads.
I reproduced it with the test below in `execution_graph.rs`. The cache is on
executor1 under build stage 7, and the fetch failure is for the job's own stage
1 on executor1. The job fails with `Stage 2 lost DataFrame::cache() data on
executor executor-id1` and stage 1 is never rerun.
Could it match `(executor_id, map_stage_id, map_partition_id)` against the
external locations instead, and fall through to the normal path when only an
own input matches? `FetchPartitionError` has no job id, so when both match it's
ambiguous, but treating that case as cache loss seems fine.
<details><summary>Repro test</summary>
```rust
#[tokio::test]
async fn own_fetch_failure_on_cache_executor_reruns_input() ->
Result<()> {
let executor1 = mock_executor("executor-id1".to_string());
let mut agg_graph = test_aggregation_plan(4).await;
revive_graph_and_complete_next_stage(&mut agg_graph)?;
agg_graph.revive();
let Some(ExecutionStage::Running(stage)) =
agg_graph.stages.get_mut(&2) else {
panic!("stage 2 should be running");
};
// The cache sits on executor1 too, under build stage 7.
let mut cache_loc = external_location(3, &executor1.id);
cache_loc.partition_id = PartitionId::new(&JobId::new("cache-job"),
7, 3);
let mut output = StageOutput::new();
output.partition_locations.insert(0, vec![cache_loc]);
output.complete = true;
stage.inputs.insert(EXTERNAL_STAGE_ID_BASE, output);
// Fetching the job's own stage 1 output from executor1 fails.
let task = agg_graph.pop_next_task(&executor1.id)?.unwrap();
let failed =
mock_failed_task(task, wrapped_fetch_failed_task(&executor1.id,
1, 0));
let events = agg_graph.update_task_status(&executor1, vec![failed],
4, 4)?;
assert!(
!events.iter().any(|event| matches!(
event,
QueryStageSchedulerEvent::JobRunningFailed { .. }
)),
"a fetch failure on the job's own output must not fail the job:
{events:?}"
);
Ok(())
}
```
</details>
##########
ballista/core/src/extension.rs:
##########
@@ -75,6 +76,25 @@ pub fn ballista_window_functions() -> Vec<Arc<WindowUDF>> {
all_default_window_functions()
}
+/// Lets a scheduler session plan `DataFrame::cache()` reads.
+pub trait SessionStateBuilderExt {
+ /// Installs [CacheReadQueryPlanner], so the session can plan a
+ /// `DataFrame::cache()` read, unless a query planner is already set: a
+ /// planner the builder's author chose is kept, and that author adds
+ /// [CacheReadPlanner](crate::cache::CacheReadPlanner) to it instead, for
+ /// example with [CacheReadQueryPlanner::with_extension_planners].
+ fn with_cache_read_planner(self) -> Self;
+}
+
+impl SessionStateBuilderExt for SessionStateBuilder {
+ fn with_cache_read_planner(mut self) -> Self {
Review Comment:
I think this means `standalone_with_state` and `remote_with_state` can't
read caches at all. `new_standalone_scheduler_from_state` gives the scheduler
the user's built state, and a built `SessionState` always has a query planner,
so `CacheReadPlanner` never gets added. The same goes for an
`override_session_builder` that builds its own state. The build job runs and
gets pinned, then the read fails, and because that error isn't treated as lost
data every retry fails the same way and the pinned data stays around.
I reproduced it with the test below in `context_cache.rs`. Both cases fail
on the first collect and again on the retry with `this scheduler's session
cannot plan DataFrame::cache() reads: add CacheReadPlanner to its query
planner`.
Could the scheduler lower `CacheReadNode` itself while resolving the cache
reads, so it doesn't depend on the session's planner? If that's too much for
this PR, checking that the planner can handle it before claiming the build
would at least avoid the wasted build and the leftover data. The design doc's
note about states that "have one [planner]" would need an update either way.
<details><summary>Repro test</summary>
```rust
#[rstest]
#[case::standalone_with_state(crate::common::standalone_context_with_state())]
#[case::remote_with_state(crate::common::remote_context_with_state())]
#[tokio::test]
async fn cache_with_state(
#[future(awt)]
#[case]
ctx: SessionContext,
test_data: String,
) {
enable_cache(&ctx).await;
let cached = alltypes(&ctx, &test_data)
.await
.select(vec![col("id")])
.unwrap()
.cache()
.await
.unwrap();
assert_eq!(cached.clone().count().await.unwrap(), 8);
assert_eq!(cached.count().await.unwrap(), 8);
}
```
</details>
##########
ballista/scheduler/src/state/cache_linker.rs:
##########
@@ -0,0 +1,1130 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Links the jobs behind a `DataFrame::cache()` query without blocking.
+//!
+//! A query reading a cache that is not built yet runs as a build (materialize)
+//! job per missing cache plus the client's read job, which cannot be planned
+//! until its builds finish. The read job is registered as queued, its id
+//! returned at once, and it is held here until the jobs it waits on finish.
+//!
+//! [`CacheLinker`](crate::state::cache_linker::CacheLinker) is the
bookkeeping;
+//! the [`SchedulerState`] methods below drive it.
+
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
+
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::event_loop::EventSender;
+use ballista_core::extension::BallistaCacheNode;
+use ballista_core::serde::protobuf::{FailedJob, JobStatus, job_status};
+use ballista_core::{JobId, JobStatusSubscriber};
+use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore};
+use datafusion::prelude::SessionContext;
+use datafusion_proto::logical_plan::AsLogicalPlan;
+use datafusion_proto::physical_plan::AsExecutionPlan;
+use futures::future::BoxFuture;
+use log::{debug, error, warn};
+use parking_lot::Mutex;
+
+use crate::cache::{
+ CACHE_DATA_LOST, OnMiss, ResolvedRead, all_locations_live, cache_node,
+ resolve_cache_reads, shape_locations,
+};
+use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan};
+use crate::scheduler_server::timestamp_millis;
+use crate::state::SchedulerState;
+use crate::state::cache_registry::{BeginOutcome, CacheKey, CompleteOutcome};
+
+/// What a held job is for.
+#[derive(Debug, Clone)]
+pub(crate) enum PendingKind {
+ /// The client's job, reading the caches.
+ Read,
+ /// Builds one cache, owning its registry claim on `key`.
+ Materialize { key: CacheKey },
+}
+
+/// A job held until the build jobs it reads have finished.
+pub(crate) struct PendingJob {
+ pub(crate) job_id: JobId,
+ pub(crate) job_name: String,
+ pub(crate) ctx: Arc<SessionContext>,
+ /// Still holds its `BallistaCacheNode`s; they become reads on release.
+ pub(crate) plan: LogicalPlan,
+ pub(crate) kind: PendingKind,
+ /// Build jobs not finished yet.
+ pub(crate) waiting_on: HashSet<JobId>,
+ pub(crate) queued_at: u64,
+ pub(crate) subscriber: Option<JobStatusSubscriber>,
+ /// Times a cache it reads went stale while held; bounded like stage
retries.
+ pub(crate) relinks: usize,
+}
+
+/// How a job ended, for [`SchedulerState::on_job_terminal`].
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum JobEnd<'a> {
+ Finished,
+ Failed(&'a str),
+ Cancelled,
+}
+
+/// A cache node's state when a submission is prepared.
+enum CacheClaim {
+ /// Built on live executors: read it.
+ Hit,
+ /// Claimed by this submission, for this build job.
+ Claimed(JobId),
+ /// Being built by another submission's job: wait on it.
+ InFlight(JobId),
+}
+
+/// A concurrent submission can change a key between lookup and claim.
+const CLAIM_ATTEMPTS: usize = 4;
+
+#[derive(Default)]
+struct Inner {
+ /// Build jobs claimed and not yet ended, with their cache. A job missing
+ /// here is done, or never was a build job.
+ inflight: HashMap<JobId, CacheKey>,
+ pending: HashMap<JobId, PendingJob>,
+ /// Build job -> the held jobs waiting on it.
+ dependents: HashMap<JobId, HashSet<JobId>>,
+ /// Running job -> the caches it reads.
+ reads: HashMap<JobId, Vec<ResolvedRead>>,
+}
+
+/// Tracks build jobs in flight, the jobs held for them, and running jobs'
cache
+/// reads. One lock guards it all; no operation awaits while holding it.
+#[derive(Default)]
+pub(crate) struct CacheLinker {
+ inner: Mutex<Inner>,
+}
+
+impl CacheLinker {
+ /// Marks `job` as building `key`. Called before the registry claim, so
+ /// whoever sees the claim also finds the job in flight.
+ pub(crate) fn track(&self, job: &JobId, key: CacheKey) {
+ self.inner.lock().inflight.insert(job.clone(), key);
+ }
+
+ /// The cache `job` builds, if it is in flight.
+ pub(crate) fn key_of(&self, job: &JobId) -> Option<CacheKey> {
+ self.inner.lock().inflight.get(job).cloned()
+ }
+
+ /// Undoes [`track`](Self::track) for a job that did not get its claim.
+ pub(crate) fn untrack(&self, job: &JobId) {
+ self.inner.lock().inflight.remove(job);
+ }
+
+ /// Holds `jobs` until the build jobs they wait on finish, ignoring builds
+ /// already finished. Returns the jobs with nothing to wait for.
+ pub(crate) fn register(&self, jobs: Vec<PendingJob>) -> Vec<PendingJob> {
+ let mut inner = self.inner.lock();
+ let mut ready = Vec::new();
+ for mut job in jobs {
+ job.waiting_on.retain(|m| inner.inflight.contains_key(m));
+ if job.waiting_on.is_empty() {
+ ready.push(job);
+ continue;
+ }
+ for m in &job.waiting_on {
+ inner
+ .dependents
+ .entry(m.clone())
+ .or_default()
+ .insert(job.job_id.clone());
+ }
+ inner.pending.insert(job.job_id.clone(), job);
+ }
+ ready
+ }
+
+ /// `finished` built its cache: returns the held jobs now ready.
+ pub(crate) fn complete(&self, finished: &JobId) -> Vec<PendingJob> {
+ let mut inner = self.inner.lock();
+ inner.inflight.remove(finished);
+ let mut ready = Vec::new();
+ for id in inner.dependents.remove(finished).unwrap_or_default() {
+ let done = inner.pending.get_mut(&id).map(|job| {
+ job.waiting_on.remove(finished);
+ job.waiting_on.is_empty()
+ });
+ if done == Some(true)
+ && let Some(job) = inner.pending.remove(&id)
+ {
+ ready.push(job);
+ }
+ }
+ ready
+ }
+
+ /// `failed` will not build its cache: removes and returns every held job
+ /// depending on it, directly or through another held build.
+ pub(crate) fn fail(&self, failed: &JobId) -> Vec<PendingJob> {
+ let mut inner = self.inner.lock();
+ inner.inflight.remove(failed);
+ let mut victims = Vec::new();
+ let mut frontier = vec![failed.clone()];
+ while let Some(id) = frontier.pop() {
+ for dependent in inner.dependents.remove(&id).unwrap_or_default() {
+ let Some(job) = inner.pending.remove(&dependent) else {
+ continue;
+ };
+ if matches!(job.kind, PendingKind::Materialize { .. }) {
+ inner.inflight.remove(&dependent);
+ frontier.push(dependent);
+ }
+ victims.push(job);
+ }
+ }
+ victims
+ }
+
+ /// Removes a held job the client cancelled.
+ pub(crate) fn cancel(&self, job: &JobId) -> Option<PendingJob> {
+ self.inner.lock().pending.remove(job)
+ }
+
+ /// Records the cache reads of a job being submitted.
+ pub(crate) fn record_reads(&self, job: &JobId, reads: Vec<ResolvedRead>) {
+ if !reads.is_empty() {
+ self.inner.lock().reads.insert(job.clone(), reads);
+ }
+ }
+
+ /// `job` ended: returns its cache reads, and whether it is still a build
in
+ /// flight, held, or waited on.
+ pub(crate) fn job_ended(&self, job: &JobId) -> (Vec<ResolvedRead>, bool) {
+ let mut inner = self.inner.lock();
+ let reads = inner.reads.remove(job).unwrap_or_default();
+ let linked = inner.inflight.contains_key(job)
+ || inner.pending.contains_key(job)
+ || inner.dependents.contains_key(job);
+ (reads, linked)
+ }
+}
+
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
SchedulerState<T, U> {
+ /// Submits a query holding `DataFrame::cache()` nodes without waiting:
claims
+ /// a build job per missing cache, holds the read job (queued) until they
+ /// finish, and returns its id.
+ pub(crate) async fn submit_cache_plan(
Review Comment:
This runs inside the `execute_query` handler, so if the client drops the
request after `prepare_caches` claims the build but before `submit_released`
posts it, the claim stays `Pending` forever. Every later read of that cache
then waits on a build that never runs. The awaits in that window are
`job_id_generator.next_id()` and `post_event` when the event channel is full.
I reproduced it with the test below in this file. A job id generator stalls
on the read job's id and the submit times out. After that the registry still
says `Pending(job-0)` and the next query over the cache stays `Queued`. The
`stall = false` case finishes with the same number of ticks.
Could everything after the first claim run in a spawned task, or roll the
claims back with a drop guard?
<details><summary>Repro test</summary>
```rust
/// Hands out the build job's id, then stalls on the read job's id.
struct StallSecondId {
calls: std::sync::atomic::AtomicUsize,
stall: std::sync::atomic::AtomicBool,
}
#[async_trait::async_trait]
impl crate::scheduler_server::JobIdGenerator for StallSecondId {
async fn next_id(&self) -> String {
let n = self.calls.fetch_add(1,
std::sync::atomic::Ordering::SeqCst);
if n == 1 &&
self.stall.load(std::sync::atomic::Ordering::SeqCst) {
std::future::pending::<()>().await;
}
format!("job-{n}")
}
}
#[rstest::rstest]
#[tokio::test]
async fn dropped_submit_does_not_strand_the_cache(#[values(true, false)]
stall: bool) {
let generator = Arc::new(StallSecondId {
calls: Default::default(),
stall: std::sync::atomic::AtomicBool::new(stall),
});
let config =
SchedulerConfig::default().with_job_id_generator(generator.clone());
let (mut test, _metrics) = scheduler_test_with(config, None).await;
let plan = cache_over("c1", scan());
// The client gives up while the read job's id is being generated.
let _ = tokio::time::timeout(Duration::from_millis(300),
test.submit("dropped", &plan))
.await;
generator.stall.store(false, std::sync::atomic::Ordering::SeqCst);
let again = test.submit("again", &plan).await.unwrap();
for _ in 0..30 {
let _ = tokio::time::timeout(Duration::from_millis(100),
test.tick()).await;
}
let status = test.job_status(&again).await.unwrap();
assert!(status.as_ref().is_some_and(is_terminal), "{status:?}");
}
```
</details>
##########
ballista/scheduler/src/state/mod.rs:
##########
@@ -433,20 +454,31 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
Ok(())
}
- /// Immediately reclaim intermediate shuffle data, then spawn a delayed
- /// future to clean up the remaining job data (final-stage output + job
- /// state) on both Scheduler and Executors.
- pub(crate) fn clean_up_successful_job(
+ /// Reclaim a finished job's shuffle data and schedule state cleanup. A job
+ /// backing a cache (recorded by [`on_job_terminal`](Self::on_job_terminal)
+ /// before this runs) keeps its final-stage data; unpinning is future work.
+ pub(crate) async fn clean_up_successful_job(
&self,
job_id: JobId,
intermediate_stage_ids: Vec<u32>,
) {
self.executor_manager
.clean_up_intermediate_job_data(job_id.clone(),
intermediate_stage_ids);
- self.executor_manager.clean_up_job_data_delayed(
- job_id.clone(),
- self.config.finished_job_data_clean_up_interval_seconds,
- );
+ let pinned = match self.cache_state.pinned_stages().await {
Review Comment:
Suggestion, fine for the eviction follow-up: a pinned job skips
`clean_up_job_data_delayed` entirely, and the entry that `invalidate_if`
removes is dropped without scheduling any cleanup. So every rebuild (after
`CACHE_DATA_LOST` or a dead-executor hit) leaves the old build's data on disk
unless the executor janitor is turned on, and it's off by default. Registry
entries also stay for the life of the scheduler. Could invalidation at least
schedule cleanup for the entry it removes? That would keep rebuilds from piling
up data until proper eviction lands.
--
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]
