akshaychitneni commented on code in PR #2657:
URL:
https://github.com/apache/datafusion-ballista/pull/2657#discussion_r4427015377
##########
ballista/scheduler/src/state/mod.rs:
##########
@@ -453,6 +591,24 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
) -> Result<()> {
let start = Instant::now();
+ // A pipelined job must be placed all at once, which the scheduler can
+ // only do when it pushes tasks to executors.
+ if !self.config.is_push_staged_scheduling()
+ && session_ctx
+ .state()
+ .config()
+ .ballista_config()
+ .execution_mode()?
+ == ExecutionMode::Pipelined
+ {
+ return Err(BallistaError::General(
+ "ballista.execution.mode = pipelined requires push-staged
scheduling, \
+ but this scheduler runs pull-staged (start it with \
+ --scheduler-policy push-staged)"
+ .to_owned(),
+ ));
+ }
+
self.task_manager
Review Comment:
**(1) Blocker.** Cache builds inherit the session's mode, so a pipelined
session builds them as pipelined jobs, which leave nothing to cache (see 2).
Suggest: if `self.cache_linker.key_of(job_id)` is set, submit with a copy of
the session that has `ballista.execution.mode = blocking`, made with
`session_manager.create_or_update_session` so the session id is kept. Read jobs
keep the user's mode; reading a blocking-built cache from a pipelined job
already works.
##########
ballista/scheduler/src/state/pipelined_graph.rs:
##########
@@ -0,0 +1,1734 @@
+// 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.
+
+//! The execution graph of a job run with `ballista.execution.mode =
+//! pipelined`, [`PipelinedExecutionGraph`].
+//!
+//! A blocking job runs stage by stage, each stage becoming runnable once the
+//! stages it reads have written their shuffle files. A pipelined job instead
+//! places every task of every stage at once, and stages stream to each other
+//! while they run (see `ballista_core::execution_plans::pipelined`). So:
+//!
+//! - Stages are planned with the static [`DistributedPlanner`], and each
+//! stage's root becomes a `PipelinedShuffleWriterExec`.
+//! - Placement is all or nothing: a partially placed job would deadlock, its
+//! consumers waiting on producers that never run. The scheduler asks for a
+//! [`Placement`] with [`plan_placement`], reserves the vcores, and commits
+//! it with [`PipelinedExecutionGraph::place`], which resolves every reader
+//! to the producer tasks it streams from.
+//! - The job reserves, on each executor, the most vcores any one of its
+//! stages uses there rather than the sum, since its stages share CPU and a
+//! task waiting on another is a parked future, not a thread. For the same
+//! reason a stage wider than the cluster, such as a `UNION ALL` of several
+//! scans, can run several partitions per vcore: it waits for every vcore
+//! to be free instead of failing. The reservation is released when the job
+//! finishes, not task by task.
+//! - Nothing is retried. A failed task or a lost executor fails the job.
+//! When the final stage completes, tasks still running — producers whose
+//! consumer stopped early, as under a `LIMIT` — are cancelled.
+
+use std::collections::{BTreeMap, HashMap, HashSet};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+use ballista_core::JobId;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::execution_plans::{
+ PipelinedShuffleReaderExec, PipelinedShuffleWriterExec, ShuffleWriter,
+ UnresolvedShuffleExec, compute_global_output_partition_ids,
+};
+use ballista_core::serde::protobuf::job_status::Status;
+use ballista_core::serde::protobuf::{
+ AvailableVcores, FailedJob, JobStatus, RunningJob, SuccessfulJob,
TaskStatus,
+ task_status,
+};
+use ballista_core::serde::scheduler::{
+ ExecutorMetadata, PartitionId, PartitionLocation,
PipelinedPartitionLocation, TaskKey,
+};
+use datafusion::physical_plan::{ExecutionPlan, replace_children_if_necessary};
+use datafusion::prelude::SessionConfig;
+use log::{info, warn};
+
+use crate::cluster::{BoundTask, ExecutorSlot, stage_has_input_collapse};
+use crate::planner::{DistributedPlanner, find_unresolved_shuffles};
+use crate::scheduler_server::event::QueryStageSchedulerEvent;
+use crate::scheduler_server::timestamp_millis;
+use crate::state::execution_graph::{
+ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ExecutionStageBuilder,
+ RunningTaskInfo, TaskDescription, create_task_info,
+};
+use crate::state::execution_stage::{
+ PendingPartitions, RunningStage, stage_input_partitions,
+};
+use crate::state::task_builder::stages_read_whole;
+
+/// What placing one stage of a pipelined job takes.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StageRequirement {
+ /// The stage.
+ pub stage_id: usize,
+ /// Input partitions the stage's tasks process between them.
+ pub partitions: usize,
+ /// The stage collapses its input to one output partition, so all of its
+ /// partitions must go to a single task.
+ pub collapse: bool,
+}
+
+/// Where every task of a pipelined job runs.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Placement {
+ /// For each stage, its tasks in task-id order: the executor each runs on
+ /// and the partitions it processes.
+ pub tasks: BTreeMap<usize, Vec<(String, Vec<usize>)>>,
+ /// Vcores to reserve on each executor: the most any one stage uses there.
+ pub reservation: Vec<ExecutorSlot>,
+}
+
+/// The outcome of [`plan_placement`].
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum PlacementDecision {
+ /// The job fits in the free vcores now.
+ Fits(Placement),
+ /// The job does not fit now, but will once vcores free up.
+ Wait,
+}
+
+/// Decides where every task of a pipelined job runs, given each executor's
+/// free vcores.
+///
+/// Each stage is planned against the full free capacity, because a job's
+/// stages share the vcores it reserves (see the [module documentation](self)).
+/// A collapse stage gets one task and one vcore. Any other stage gets one
+/// vcore per partition: its partitions are spread round-robin over the
+/// executors with room, in contiguous ranges, and each executor's range is
+/// cut into tasks of at most `max_partitions_per_task` partitions (`0` means
+/// unbounded).
+///
+/// `total_vcores` is the cluster's capacity when idle. A stage wider than
+/// that needs every vcore free, and then shares them: its partitions are
+/// dealt round-robin again once each free vcore has one.
+pub fn plan_placement(
+ requirements: &[StageRequirement],
+ available: &[AvailableVcores],
+ total_vcores: usize,
+ max_partitions_per_task: usize,
+) -> PlacementDecision {
+ // Most free first, then by id, so placement is deterministic.
+ let mut budgets: Vec<(&str, usize)> = available
+ .iter()
+ .filter(|a| a.vcores > 0)
+ .map(|a| (a.executor_id.as_str(), a.vcores as usize))
+ .collect();
+ budgets.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
+ if budgets.is_empty() {
+ return PlacementDecision::Wait;
+ }
+ let cap = if max_partitions_per_task == 0 {
+ usize::MAX
+ } else {
+ max_partitions_per_task
+ };
+
+ let mut reserved: HashMap<&str, usize> = HashMap::new();
+ let mut tasks = BTreeMap::new();
+ for requirement in requirements {
+ let mut stage_tasks = vec![];
+ if requirement.collapse {
+ let executor = budgets[0].0;
+ stage_tasks
+ .push((executor.to_owned(),
(0..requirement.partitions).collect()));
+ let entry = reserved.entry(executor).or_default();
+ *entry = (*entry).max(1);
+ } else {
+ let Some(counts) = spread(requirement.partitions, &budgets,
total_vcores)
+ else {
+ return PlacementDecision::Wait;
+ };
+ let mut next = 0;
+ for ((executor, free), count) in budgets.iter().zip(counts) {
+ if count == 0 {
+ continue;
+ }
+ let partitions: Vec<usize> = (next..next + count).collect();
+ next += count;
+ for chunk in partitions.chunks(cap) {
+ stage_tasks.push(((*executor).to_owned(), chunk.to_vec()));
+ }
+ let entry = reserved.entry(executor).or_default();
+ *entry = (*entry).max(count.min(*free));
+ }
+ }
+ tasks.insert(requirement.stage_id, stage_tasks);
+ }
+
+ let mut reservation: Vec<ExecutorSlot> = reserved
+ .into_iter()
+ .map(|(executor, vcores)| (executor.to_owned(), vcores as u32))
+ .collect();
+ reservation.sort();
+ PlacementDecision::Fits(Placement { tasks, reservation })
+}
+
+/// Deals `partitions` one at a time, round-robin, to the executors that still
+/// have a free vcore, and deals what is left the same way once every free
+/// vcore has a partition. Returns how many each executor got, or `None` if
+/// fewer vcores are free than the stage needs: one per partition, but no more
+/// than the cluster's `total_vcores`.
+fn spread(
+ partitions: usize,
+ budgets: &[(&str, usize)],
+ total_vcores: usize,
+) -> Option<Vec<usize>> {
+ let free: usize = budgets.iter().map(|(_, free)| free).sum();
+ if free < partitions.min(total_vcores) || free == 0 {
+ return None;
+ }
+ let mut counts = vec![0; budgets.len()];
+ let mut left = partitions;
+ while left > 0 {
+ let round = left.min(free);
+ let mut dealt = vec![0; budgets.len()];
+ let mut total = 0;
+ while total < round {
+ for (i, (_, free)) in budgets.iter().enumerate() {
+ if total == round {
+ break;
+ }
+ if dealt[i] < *free {
+ dealt[i] += 1;
+ total += 1;
+ }
+ }
+ }
+ for (count, dealt) in counts.iter_mut().zip(dealt) {
+ *count += dealt;
+ }
+ left -= round;
+ }
+ Some(counts)
+}
+
+/// The [`ExecutionGraph`] of a pipelined job. See the
+/// [module documentation](self).
+#[derive(Clone)]
+pub struct PipelinedExecutionGraph {
+ scheduler_id: Option<String>,
+ job_id: JobId,
+ job_name: String,
+ session_id: String,
+ status: JobStatus,
+ queued_at: u64,
+ start_time: u64,
+ end_time: u64,
+ stages: HashMap<usize, ExecutionStage>,
+ /// The stage whose output the client reads.
+ final_stage_id: usize,
+ /// Whether every task has been placed. A job is placed once.
+ placed: bool,
+ /// Vcores reserved on each executor, held until the job finishes.
+ reservation: Vec<ExecutorSlot>,
+ /// Executors running a task of this job. Losing one fails the job.
+ executors: HashSet<String>,
+ session_config: Arc<SessionConfig>,
+ logical_plan: Option<String>,
+ physical_plan: Arc<dyn ExecutionPlan>,
+}
+
+impl PipelinedExecutionGraph {
+ /// Plans `plan` into stages and turns each stage's root into a
+ /// `PipelinedShuffleWriterExec`. A stage read through a broadcast reader
+ /// retains its outputs for replay to every consumer task.
+ #[allow(clippy::too_many_arguments)]
+ pub fn new(
+ scheduler_id: &str,
+ job_id: &JobId,
+ job_name: &str,
+ session_id: &str,
+ plan: Arc<dyn ExecutionPlan>,
+ queued_at: u64,
+ session_config: Arc<SessionConfig>,
+ planner: &mut dyn DistributedPlanner,
+ logical_plan: Option<String>,
+ ) -> Result<Self> {
+ let shuffle_stages =
+ planner.plan_query_stages(job_id, plan.clone(),
session_config.options())?;
+
+ // Stages whose every output partition more than one task reads in
+ // full: a broadcast join's build side, and the input of a collapse
+ // inside a stage of several tasks, such as a cross join's left side.
+ // Their outputs are kept for each reader rather than queued once.
+ let mut replayed_stages = HashSet::new();
+ for stage in &shuffle_stages {
+ let stage_plan: Arc<dyn ExecutionPlan> = stage.clone();
+ for unresolved in find_unresolved_shuffles(&stage_plan)? {
+ if unresolved.broadcast {
+ replayed_stages.insert(unresolved.stage_id);
+ }
+ }
+ // A collapse stage runs as one task, so it reads its inputs once.
+ if !stage_has_input_collapse(&stage_plan) {
+ replayed_stages.extend(stages_read_whole(&stage_plan));
+ }
+ }
+
+ let pipelined_stages = shuffle_stages
+ .into_iter()
+ .map(|stage| {
+ let input = stage.children()[0].clone();
+ let writer = PipelinedShuffleWriterExec::try_new(
+ job_id.clone(),
+ stage.stage_id(),
+ input,
+ stage.shuffle_output_partitioning().cloned(),
+ replayed_stages.contains(&stage.stage_id()),
+ )?;
+ Ok(Arc::new(writer) as Arc<dyn ShuffleWriter>)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let stages =
+
ExecutionStageBuilder::new(session_config.clone()).build(pipelined_stages)?;
+ let final_stage_id = stages
+ .iter()
+ .filter(|(_, stage)| stage.output_links().is_empty())
+ .map(|(stage_id, _)| *stage_id)
+ .max()
+ .ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "pipelined job {job_id} has no final stage"
+ ))
+ })?;
+
+ let started_at = timestamp_millis();
+ Ok(Self {
+ scheduler_id: Some(scheduler_id.to_owned()),
+ job_id: job_id.clone(),
+ job_name: job_name.to_owned(),
+ session_id: session_id.to_owned(),
+ status: JobStatus {
+ job_id: job_id.to_string(),
+ job_name: job_name.to_owned(),
+ status: Some(Status::Running(RunningJob {
+ queued_at,
+ started_at,
+ scheduler: scheduler_id.to_owned(),
+ pipelined_output_locations: vec![],
+ })),
+ },
+ queued_at,
+ start_time: started_at,
+ end_time: 0,
+ stages,
+ final_stage_id,
+ placed: false,
+ reservation: vec![],
+ executors: HashSet::new(),
+ session_config,
+ logical_plan,
+ physical_plan: plan,
+ })
+ }
+
+ fn is_running(&self) -> bool {
+ 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.
+ pub fn awaits_placement(&self) -> bool {
+ !self.placed && self.is_running()
+ }
+
+ /// What placing each stage takes, or `None` once the job is placed or no
+ /// longer running.
+ pub fn placement_requirements(&self) -> Option<Vec<StageRequirement>> {
+ if !self.awaits_placement() {
+ return None;
+ }
+ let mut requirements: Vec<StageRequirement> = self
+ .stages
+ .iter()
+ .map(|(stage_id, stage)| {
+ let plan = stage.plan_arc();
+ StageRequirement {
+ stage_id: *stage_id,
+ partitions: stage_input_partitions(&plan),
+ collapse: stage_has_input_collapse(&plan),
+ }
+ })
+ .collect();
+ requirements.sort_by_key(|r| r.stage_id);
+ Some(requirements)
+ }
+
+ /// The scheduler's `max_partitions_per_task` for this job.
+ pub fn max_partitions_per_task(&self) -> usize {
+ self.session_config
+ .options()
+ .extensions
+ .get::<ballista_core::config::BallistaConfig>()
+ .map(|config| config.max_partitions_per_task())
+ .unwrap_or(0)
+ }
+
+ /// Commits a placement whose vcores the scheduler has reserved: resolves
+ /// every reader to the producer tasks it streams from, marks every stage
+ /// running, publishes the final stage's locations in the job status, and
+ /// returns the tasks to launch.
+ ///
+ /// `executors` must hold the metadata of every executor the placement
+ /// uses.
+ pub fn place(
+ &mut self,
+ placement: &Placement,
+ executors: &HashMap<String, ExecutorMetadata>,
+ ) -> Result<Vec<BoundTask>> {
+ if !self.awaits_placement() {
+ return Err(BallistaError::Internal(format!(
+ "pipelined job {} is already placed or no longer running",
+ self.job_id
+ )));
+ }
+ let missing: Vec<usize> = self
+ .stages
+ .keys()
+ .filter(|stage_id| !placement.tasks.contains_key(stage_id))
+ .copied()
+ .collect();
+ if !missing.is_empty() {
+ return Err(BallistaError::Internal(format!(
+ "placement of pipelined job {} leaves stages {missing:?}
unplaced",
+ self.job_id
+ )));
+ }
+
+ // Where every producer task serves each of its output partitions.
+ let mut outputs: HashMap<
+ usize,
+ BTreeMap<usize, Vec<PipelinedPartitionLocation>>,
+ > = HashMap::new();
+ for (stage_id, tasks) in &placement.tasks {
+ let plan = self.stage(*stage_id)?.plan_arc();
+ let stage_outputs = outputs.entry(*stage_id).or_default();
+ for (task_id, (executor_id, partitions)) in
tasks.iter().enumerate() {
+ let executor_meta = executors.get(executor_id).ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "no metadata for executor {executor_id} placed in
pipelined job {}",
+ self.job_id
+ ))
+ })?;
+ for partition in compute_global_output_partition_ids(&plan,
partitions) {
+ stage_outputs.entry(partition).or_default().push(
+ PipelinedPartitionLocation {
+ partition_id: PartitionId::new(
+ &self.job_id,
+ *stage_id,
+ partition,
+ ),
+ task_id,
+ executor_meta: executor_meta.clone(),
+ },
+ );
+ }
+ }
+ }
+
+ let reserved: HashMap<&str, u32> = placement
+ .reservation
+ .iter()
+ .map(|(executor_id, vcores)| (executor_id.as_str(), *vcores))
+ .collect();
+ let mut bound = vec![];
+ for (stage_id, tasks) in &placement.tasks {
+ let stage = self.stage(*stage_id)?;
+ let (output_links, inputs) = match stage {
+ ExecutionStage::UnResolved(stage) => {
+ (stage.output_links.clone(), stage.inputs.clone())
+ }
+ ExecutionStage::Resolved(stage) => {
+ (stage.output_links.clone(), stage.inputs.clone())
+ }
+ other => {
+ return Err(BallistaError::Internal(format!(
+ "stage {stage_id} of pipelined job {} is {} before
placement",
+ self.job_id,
+ other.variant_name()
+ )));
+ }
+ };
+ let plan = resolve_pipelined_readers(stage.plan_arc(), &outputs)?;
+ let partitions = stage_input_partitions(&plan);
+ let mut running = RunningStage::new(
+ *stage_id,
+ 0,
+ plan.clone(),
+ partitions,
+ output_links,
+ inputs,
+ self.session_config.clone(),
+ );
+ // Every partition is handed out here, so nothing is pending.
+ running.pending = PendingPartitions::empty(partitions);
+ for (task_id, (executor_id, slice)) in tasks.iter().enumerate() {
+ let mut task_info = create_task_info(executor_id.clone(),
task_id);
+ task_info.global_input_partition_ids = slice.clone();
+ // The job's reservation covers every task and is released
+ // when the job ends, so no task carries vcores of its own.
+ task_info.vcores_consumed = 0;
+ running.task_infos.push(task_info);
+ bound.push((
+ executor_id.clone(),
+ TaskDescription {
+ session_id: self.session_id.clone(),
+ key: TaskKey {
+ job_id: self.job_id.clone(),
+ stage_id: *stage_id,
+ task_id,
+ },
+ stage_attempt_num: 0,
+ task_attempt: 0,
+ global_input_partition_ids: slice.clone(),
+ // The job's reservation on the executor, which sizes
+ // the memory pool all its tasks there share.
+ vcores_consumed: reserved
+ .get(executor_id.as_str())
+ .copied()
+ .unwrap_or(1),
+ plan: plan.clone(),
+ session_config: self.session_config.clone(),
+ },
+ ));
+ }
+ self.stages
+ .insert(*stage_id, ExecutionStage::Running(running));
+ }
+
+ let final_locations: Vec<PipelinedPartitionLocation> = outputs
+ .remove(&self.final_stage_id)
+ .unwrap_or_default()
+ .into_values()
+ .flatten()
+ .collect();
+ if let Some(Status::Running(running)) = &mut self.status.status {
+ running.pipelined_output_locations =
+ final_locations.into_iter().map(Into::into).collect();
+ }
+ self.reservation = placement.reservation.clone();
+ self.executors = placement
+ .tasks
+ .values()
+ .flatten()
+ .map(|(executor_id, _)| executor_id.clone())
+ .collect();
+ self.placed = true;
+ info!(
+ "Placed pipelined job {}: {} tasks over {} executors, reserving
{:?}",
+ self.job_id,
+ bound.len(),
+ self.executors.len(),
+ self.reservation
+ );
+ Ok(bound)
+ }
+
+ /// Takes the vcores the job reserved, to be released. Empty once taken,
+ /// so they are released once.
+ pub fn take_reservation(&mut self) -> Vec<ExecutorSlot> {
+ std::mem::take(&mut self.reservation)
+ }
+
+ /// If losing `executor_id` breaks this job, the reason to fail it with.
+ pub fn lost_executor(&self, executor_id: &str) -> Option<String> {
+ (self.is_running() && self.executors.contains(executor_id)).then(|| {
+ format!(
+ "executor {executor_id}, which was running part of pipelined
job {}, was lost",
+ self.job_id
+ )
+ })
+ }
+
+ fn stage(&self, stage_id: usize) -> Result<&ExecutionStage> {
+ self.stages.get(&stage_id).ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "pipelined job {} has no stage {stage_id}",
+ self.job_id
+ ))
+ })
+ }
+
+ /// Marks every stage whose tasks have all succeeded as successful.
+ fn succeed_finished_stages(&mut self) {
+ let finished: Vec<usize> = self
+ .stages
+ .iter()
+ .filter_map(|(stage_id, stage)| match stage {
+ ExecutionStage::Running(running) if running.is_successful() =>
{
+ Some(*stage_id)
+ }
+ _ => None,
+ })
+ .collect();
+ for stage_id in finished {
+ self.succeed_stage(stage_id);
+ }
+ }
+}
+
+/// Replaces every `UnresolvedShuffleExec` in `plan` with a
+/// `PipelinedShuffleReaderExec` over the producer tasks in `outputs`.
+fn resolve_pipelined_readers(
+ plan: Arc<dyn ExecutionPlan>,
+ outputs: &HashMap<usize, BTreeMap<usize, Vec<PipelinedPartitionLocation>>>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ if let Some(unresolved) = plan.downcast_ref::<UnresolvedShuffleExec>() {
+ if unresolved.coalesce.is_some() {
+ return Err(BallistaError::Internal(format!(
+ "pipelined jobs do not support coalesced readers (stage {})",
+ unresolved.stage_id
+ )));
+ }
+ let empty = BTreeMap::new();
+ let produced = outputs.get(&unresolved.stage_id).unwrap_or(&empty);
+ let schema = unresolved.schema.clone();
+ let reader = if unresolved.broadcast {
+ PipelinedShuffleReaderExec::new_broadcast(
+ unresolved.stage_id,
+ produced.values().flatten().cloned().collect(),
+ schema,
+ )
+ } else {
+ let partitions = (0..unresolved.output_partition_count)
+ .map(|p| produced.get(&p).cloned().unwrap_or_default())
+ .collect();
+ PipelinedShuffleReaderExec::try_new(
+ unresolved.stage_id,
+ partitions,
+ schema,
+ unresolved.properties().output_partitioning().clone(),
+ )?
+ };
+ return Ok(Arc::new(reader));
+ }
+ let children = plan
+ .children()
+ .into_iter()
+ .map(|child| resolve_pipelined_readers(child.clone(), outputs))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(replace_children_if_necessary(plan, children)?)
+}
+
+impl ExecutionGraph for PipelinedExecutionGraph {
+ fn job_id(&self) -> &JobId {
+ &self.job_id
+ }
+
+ fn job_name(&self) -> &str {
+ &self.job_name
+ }
+
+ fn session_id(&self) -> &str {
+ &self.session_id
+ }
+
+ fn scheduler_id(&self) -> Option<&str> {
+ self.scheduler_id.as_deref()
+ }
+
+ fn session_config(&self) -> Arc<SessionConfig> {
+ self.session_config.clone()
+ }
+
+ fn status(&self) -> &JobStatus {
+ &self.status
+ }
+
+ fn logical_plan(&self) -> Option<&str> {
+ self.logical_plan.as_deref()
+ }
+
+ fn physical_plan(&self) -> Arc<dyn ExecutionPlan> {
+ self.physical_plan.clone()
+ }
+
+ fn start_time(&self) -> u64 {
+ self.start_time
+ }
+
+ fn end_time(&self) -> u64 {
+ self.end_time
+ }
+
+ fn completed_stages(&self) -> usize {
+ self.stages
+ .values()
+ .filter(|stage| matches!(stage, ExecutionStage::Successful(_)))
+ .count()
+ }
+
+ fn is_successful(&self) -> bool {
+ self.stages
+ .values()
+ .all(|stage| matches!(stage, ExecutionStage::Successful(_)))
+ }
+
+ /// Stages never become runnable one at a time: [`Self::place`] starts
+ /// them all.
+ fn revive(&mut self) -> bool {
+ false
+ }
+
+ /// Records task outcomes. The first failure fails the job; the final
+ /// stage completing finishes it and cancels whatever is still running.
+ fn update_task_status(
+ &mut self,
+ _executor: &ExecutorMetadata,
+ task_statuses: Vec<TaskStatus>,
+ _max_task_failures: usize,
+ _max_stage_failures: usize,
+ ) -> Result<Vec<QueryStageSchedulerEvent>> {
+ if !self.is_running() {
+ // Late reports from tasks cancelled when the job finished or
+ // failed.
+ return Ok(vec![]);
+ }
+ let mut failure = None;
+ for status in task_statuses {
+ let stage_id = status.stage_id as usize;
+ let task_id = status.task_id as usize;
+ let Some(ExecutionStage::Running(stage)) =
self.stages.get_mut(&stage_id)
+ else {
+ warn!(
+ "Ignoring status of task {task_id} in stage {stage_id} of
pipelined job {}: the stage is not running",
+ self.job_id
+ );
+ continue;
+ };
+ if task_id >= stage.task_infos.len() {
+ warn!(
+ "Ignoring status of unknown task {task_id} in stage
{stage_id} of pipelined job {}",
+ self.job_id
+ );
+ continue;
+ }
+ let outcome = status.status.clone();
+ let metrics = status.metrics.clone();
+ if !stage.update_task_info(task_id, status) {
+ continue;
+ }
+ match outcome {
+ Some(task_status::Status::Successful(_)) => {
+ stage.update_task_metrics(task_id, metrics)?;
+ }
+ Some(task_status::Status::Failed(failed)) => {
+ failure.get_or_insert_with(|| {
+ format!(
+ "task {task_id} of stage {stage_id} failed: {}",
+ failed.error
+ )
+ });
+ }
+ _ => {}
+ }
+ }
+
+ if let Some(reason) = failure {
+ let fail_message = format!(
+ "Pipelined job {} failed because {reason}. Pipelined jobs do
not retry tasks",
+ self.job_id
+ );
+ self.fail_job(fail_message.clone());
+ return Ok(vec![QueryStageSchedulerEvent::JobRunningFailed {
+ job_id: self.job_id.clone(),
+ fail_message,
+ queued_at: self.queued_at,
+ failed_at: timestamp_millis(),
+ }]);
+ }
+
+ let final_finished = matches!(
+ self.stages.get(&self.final_stage_id),
+ Some(ExecutionStage::Running(stage)) if stage.is_successful()
+ );
+ if !final_finished {
+ self.succeed_finished_stages();
+ return Ok(vec![]);
+ }
+
+ // The client has read the whole result. Anything still running is
+ // a producer whose output nobody needs any more.
+ let leftovers = self.running_tasks();
+ for stage in self.stages.values_mut() {
+ if let ExecutionStage::Running(running) = stage {
+ *stage = ExecutionStage::Successful(running.to_successful());
+ }
+ }
+ self.succeed_job()?;
+ let mut events = vec![];
+ if !leftovers.is_empty() {
+ info!(
+ "Pipelined job {} finished; cancelling {} tasks it no longer
needs",
+ self.job_id,
+ leftovers.len()
+ );
+ events.push(QueryStageSchedulerEvent::CancelTasks(leftovers));
+ }
+ events.push(QueryStageSchedulerEvent::JobFinished {
+ job_id: self.job_id.clone(),
+ queued_at: self.queued_at,
+ completed_at: timestamp_millis(),
+ });
+ Ok(events)
+ }
+
+ fn running_stages(&self) -> Vec<usize> {
+ self.stages
+ .iter()
+ .filter(|(_, stage)| matches!(stage, ExecutionStage::Running(_)))
+ .map(|(stage_id, _)| *stage_id)
+ .collect()
+ }
+
+ fn running_tasks(&self) -> Vec<RunningTaskInfo> {
+ self.stages
+ .values()
+ .flat_map(|stage| match stage {
+ ExecutionStage::Running(stage) => stage
+ .running_tasks()
+ .into_iter()
+ .map(|(task_id, stage_id, executor_id)| RunningTaskInfo {
+ task_id,
+ job_id: self.job_id.clone(),
+ stage_id,
+ executor_id,
+ })
+ .collect(),
+ _ => vec![],
+ })
+ .collect()
+ }
+
+ /// Nothing is ever pending: [`Self::place`] hands out every partition.
+ fn available_tasks(&self) -> usize {
+ 0
+ }
+
+ /// Stage-by-stage binding never applies to a pipelined job.
+ fn fetch_running_stage(
+ &mut self,
+ _black_list: &[usize],
+ ) -> Option<&mut RunningStage> {
+ None
+ }
+
+ fn update_status(&mut self, status: JobStatus) {
+ self.status = status;
+ }
+
+ /// A pipelined job has no shuffle files to read; its final stage streams
+ /// from the locations in the running status.
+ fn output_locations(&self) -> Vec<PartitionLocation> {
Review Comment:
**(2)** This empty list is what a pipelined cache build records: the cache
ends up empty and reads return no rows without an error (reproduced in the
scheduler test harness). A cache larger than its queue fails after
`registration_timeout_seconds` instead (from reading the code). With 1 this
can't happen, but `record_materialized_cache` could also refuse a pipelined
graph as a guard.
##########
ballista/scheduler/src/state/task_manager.rs:
##########
@@ -822,41 +890,45 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskManager<T, U>
/// Handles executor loss by resetting affected tasks and stages.
///
/// Returns the running tasks that need to be cancelled, and the jobs that
- /// lost `DataFrame::cache()` data with the executor, with when they were
- /// queued. Their graphs cannot re-run another job's output, so they are
- /// not reset; the caller must fail them.
- pub async fn executor_lost(
- &self,
- executor_id: &str,
- ) -> Result<(Vec<RunningTaskInfo>, Vec<(JobId, u64)>)> {
+ /// cannot survive the loss and must fail, with the reason. A pipelined
+ /// job is one: its tasks stream to each other, so there is nothing to
+ /// roll back to. So is a job that lost `DataFrame::cache()` data with the
+ /// executor: its graph cannot re-run another job's output.
+ pub async fn executor_lost(&self, executor_id: &str) ->
Result<ExecutorLostOutcome> {
// Its capacity has been removed, so there is nothing left to refund.
// Do this even for evicted jobs and before a graph reset can fail.
self.outstanding_task_reservations
.retain(|(executor, _, _), _| executor != executor_id);
- // Collect all the running task need to cancel when there are running
stages rolled back.
- let mut running_tasks_to_cancel: Vec<RunningTaskInfo> = vec![];
- let mut lost_cache = vec![];
-
- {
- for pairs in self.active_job_cache.iter() {
- let (job_id, job_info) = pairs.pair();
- let mut graph = job_info.execution_graph.write().await;
- if graph.reads_external_data_on(executor_id) {
- let queued_at = match &graph.status().status {
- Some(job_status::Status::Running(running)) =>
running.queued_at,
- _ => 0,
- };
- lost_cache.push((job_id.clone(), queued_at));
- continue;
- }
+ let mut outcome = ExecutorLostOutcome::default();
+
+ for pairs in self.active_job_cache.iter() {
+ let (job_id, job_info) = pairs.pair();
+ let mut graph = job_info.execution_graph.write().await;
+ let reason = if let Some(pipelined) = graph.as_pipelined_mut() {
Review Comment:
**(3)** This branch skips `reads_external_data_on` for pipelined jobs. Check
both, and use `CACHE_DATA_LOST` when cache data was lost, so the lost build is
dropped.
##########
ballista/scheduler/src/state/pipelined_graph.rs:
##########
@@ -0,0 +1,1734 @@
+// 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.
+
+//! The execution graph of a job run with `ballista.execution.mode =
+//! pipelined`, [`PipelinedExecutionGraph`].
+//!
+//! A blocking job runs stage by stage, each stage becoming runnable once the
+//! stages it reads have written their shuffle files. A pipelined job instead
+//! places every task of every stage at once, and stages stream to each other
+//! while they run (see `ballista_core::execution_plans::pipelined`). So:
+//!
+//! - Stages are planned with the static [`DistributedPlanner`], and each
+//! stage's root becomes a `PipelinedShuffleWriterExec`.
+//! - Placement is all or nothing: a partially placed job would deadlock, its
+//! consumers waiting on producers that never run. The scheduler asks for a
+//! [`Placement`] with [`plan_placement`], reserves the vcores, and commits
+//! it with [`PipelinedExecutionGraph::place`], which resolves every reader
+//! to the producer tasks it streams from.
+//! - The job reserves, on each executor, the most vcores any one of its
+//! stages uses there rather than the sum, since its stages share CPU and a
+//! task waiting on another is a parked future, not a thread. For the same
+//! reason a stage wider than the cluster, such as a `UNION ALL` of several
+//! scans, can run several partitions per vcore: it waits for every vcore
+//! to be free instead of failing. The reservation is released when the job
+//! finishes, not task by task.
+//! - Nothing is retried. A failed task or a lost executor fails the job.
+//! When the final stage completes, tasks still running — producers whose
+//! consumer stopped early, as under a `LIMIT` — are cancelled.
+
+use std::collections::{BTreeMap, HashMap, HashSet};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+use ballista_core::JobId;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::execution_plans::{
+ PipelinedShuffleReaderExec, PipelinedShuffleWriterExec, ShuffleWriter,
+ UnresolvedShuffleExec, compute_global_output_partition_ids,
+};
+use ballista_core::serde::protobuf::job_status::Status;
+use ballista_core::serde::protobuf::{
+ AvailableVcores, FailedJob, JobStatus, RunningJob, SuccessfulJob,
TaskStatus,
+ task_status,
+};
+use ballista_core::serde::scheduler::{
+ ExecutorMetadata, PartitionId, PartitionLocation,
PipelinedPartitionLocation, TaskKey,
+};
+use datafusion::physical_plan::{ExecutionPlan, replace_children_if_necessary};
+use datafusion::prelude::SessionConfig;
+use log::{info, warn};
+
+use crate::cluster::{BoundTask, ExecutorSlot, stage_has_input_collapse};
+use crate::planner::{DistributedPlanner, find_unresolved_shuffles};
+use crate::scheduler_server::event::QueryStageSchedulerEvent;
+use crate::scheduler_server::timestamp_millis;
+use crate::state::execution_graph::{
+ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ExecutionStageBuilder,
+ RunningTaskInfo, TaskDescription, create_task_info,
+};
+use crate::state::execution_stage::{
+ PendingPartitions, RunningStage, stage_input_partitions,
+};
+use crate::state::task_builder::stages_read_whole;
+
+/// What placing one stage of a pipelined job takes.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StageRequirement {
+ /// The stage.
+ pub stage_id: usize,
+ /// Input partitions the stage's tasks process between them.
+ pub partitions: usize,
+ /// The stage collapses its input to one output partition, so all of its
+ /// partitions must go to a single task.
+ pub collapse: bool,
+}
+
+/// Where every task of a pipelined job runs.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Placement {
+ /// For each stage, its tasks in task-id order: the executor each runs on
+ /// and the partitions it processes.
+ pub tasks: BTreeMap<usize, Vec<(String, Vec<usize>)>>,
+ /// Vcores to reserve on each executor: the most any one stage uses there.
+ pub reservation: Vec<ExecutorSlot>,
+}
+
+/// The outcome of [`plan_placement`].
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum PlacementDecision {
+ /// The job fits in the free vcores now.
+ Fits(Placement),
+ /// The job does not fit now, but will once vcores free up.
+ Wait,
+}
+
+/// Decides where every task of a pipelined job runs, given each executor's
+/// free vcores.
+///
+/// Each stage is planned against the full free capacity, because a job's
+/// stages share the vcores it reserves (see the [module documentation](self)).
+/// A collapse stage gets one task and one vcore. Any other stage gets one
+/// vcore per partition: its partitions are spread round-robin over the
+/// executors with room, in contiguous ranges, and each executor's range is
+/// cut into tasks of at most `max_partitions_per_task` partitions (`0` means
+/// unbounded).
+///
+/// `total_vcores` is the cluster's capacity when idle. A stage wider than
+/// that needs every vcore free, and then shares them: its partitions are
+/// dealt round-robin again once each free vcore has one.
+pub fn plan_placement(
+ requirements: &[StageRequirement],
+ available: &[AvailableVcores],
+ total_vcores: usize,
+ max_partitions_per_task: usize,
+) -> PlacementDecision {
+ // Most free first, then by id, so placement is deterministic.
+ let mut budgets: Vec<(&str, usize)> = available
+ .iter()
+ .filter(|a| a.vcores > 0)
+ .map(|a| (a.executor_id.as_str(), a.vcores as usize))
+ .collect();
+ budgets.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
+ if budgets.is_empty() {
+ return PlacementDecision::Wait;
+ }
+ let cap = if max_partitions_per_task == 0 {
+ usize::MAX
+ } else {
+ max_partitions_per_task
+ };
+
+ let mut reserved: HashMap<&str, usize> = HashMap::new();
+ let mut tasks = BTreeMap::new();
+ for requirement in requirements {
+ let mut stage_tasks = vec![];
+ if requirement.collapse {
+ let executor = budgets[0].0;
+ stage_tasks
+ .push((executor.to_owned(),
(0..requirement.partitions).collect()));
+ let entry = reserved.entry(executor).or_default();
+ *entry = (*entry).max(1);
+ } else {
+ let Some(counts) = spread(requirement.partitions, &budgets,
total_vcores)
+ else {
+ return PlacementDecision::Wait;
+ };
+ let mut next = 0;
+ for ((executor, free), count) in budgets.iter().zip(counts) {
+ if count == 0 {
+ continue;
+ }
+ let partitions: Vec<usize> = (next..next + count).collect();
+ next += count;
+ for chunk in partitions.chunks(cap) {
+ stage_tasks.push(((*executor).to_owned(), chunk.to_vec()));
+ }
+ let entry = reserved.entry(executor).or_default();
+ *entry = (*entry).max(count.min(*free));
+ }
+ }
+ tasks.insert(requirement.stage_id, stage_tasks);
+ }
+
+ let mut reservation: Vec<ExecutorSlot> = reserved
+ .into_iter()
+ .map(|(executor, vcores)| (executor.to_owned(), vcores as u32))
+ .collect();
+ reservation.sort();
+ PlacementDecision::Fits(Placement { tasks, reservation })
+}
+
+/// Deals `partitions` one at a time, round-robin, to the executors that still
+/// have a free vcore, and deals what is left the same way once every free
+/// vcore has a partition. Returns how many each executor got, or `None` if
+/// fewer vcores are free than the stage needs: one per partition, but no more
+/// than the cluster's `total_vcores`.
+fn spread(
+ partitions: usize,
+ budgets: &[(&str, usize)],
+ total_vcores: usize,
+) -> Option<Vec<usize>> {
+ let free: usize = budgets.iter().map(|(_, free)| free).sum();
+ if free < partitions.min(total_vcores) || free == 0 {
+ return None;
+ }
+ let mut counts = vec![0; budgets.len()];
+ let mut left = partitions;
+ while left > 0 {
+ let round = left.min(free);
+ let mut dealt = vec![0; budgets.len()];
+ let mut total = 0;
+ while total < round {
+ for (i, (_, free)) in budgets.iter().enumerate() {
+ if total == round {
+ break;
+ }
+ if dealt[i] < *free {
+ dealt[i] += 1;
+ total += 1;
+ }
+ }
+ }
+ for (count, dealt) in counts.iter_mut().zip(dealt) {
+ *count += dealt;
+ }
+ left -= round;
+ }
+ Some(counts)
+}
+
+/// The [`ExecutionGraph`] of a pipelined job. See the
+/// [module documentation](self).
+#[derive(Clone)]
+pub struct PipelinedExecutionGraph {
+ scheduler_id: Option<String>,
+ job_id: JobId,
+ job_name: String,
+ session_id: String,
+ status: JobStatus,
+ queued_at: u64,
+ start_time: u64,
+ end_time: u64,
+ stages: HashMap<usize, ExecutionStage>,
+ /// The stage whose output the client reads.
+ final_stage_id: usize,
+ /// Whether every task has been placed. A job is placed once.
+ placed: bool,
+ /// Vcores reserved on each executor, held until the job finishes.
+ reservation: Vec<ExecutorSlot>,
+ /// Executors running a task of this job. Losing one fails the job.
+ executors: HashSet<String>,
+ session_config: Arc<SessionConfig>,
+ logical_plan: Option<String>,
+ physical_plan: Arc<dyn ExecutionPlan>,
+}
+
+impl PipelinedExecutionGraph {
+ /// Plans `plan` into stages and turns each stage's root into a
+ /// `PipelinedShuffleWriterExec`. A stage read through a broadcast reader
+ /// retains its outputs for replay to every consumer task.
+ #[allow(clippy::too_many_arguments)]
+ pub fn new(
+ scheduler_id: &str,
+ job_id: &JobId,
+ job_name: &str,
+ session_id: &str,
+ plan: Arc<dyn ExecutionPlan>,
+ queued_at: u64,
+ session_config: Arc<SessionConfig>,
+ planner: &mut dyn DistributedPlanner,
+ logical_plan: Option<String>,
+ ) -> Result<Self> {
+ let shuffle_stages =
+ planner.plan_query_stages(job_id, plan.clone(),
session_config.options())?;
+
+ // Stages whose every output partition more than one task reads in
+ // full: a broadcast join's build side, and the input of a collapse
+ // inside a stage of several tasks, such as a cross join's left side.
+ // Their outputs are kept for each reader rather than queued once.
+ let mut replayed_stages = HashSet::new();
+ for stage in &shuffle_stages {
+ let stage_plan: Arc<dyn ExecutionPlan> = stage.clone();
+ for unresolved in find_unresolved_shuffles(&stage_plan)? {
+ if unresolved.broadcast {
+ replayed_stages.insert(unresolved.stage_id);
+ }
+ }
+ // A collapse stage runs as one task, so it reads its inputs once.
+ if !stage_has_input_collapse(&stage_plan) {
+ replayed_stages.extend(stages_read_whole(&stage_plan));
+ }
+ }
+
+ let pipelined_stages = shuffle_stages
+ .into_iter()
+ .map(|stage| {
+ let input = stage.children()[0].clone();
+ let writer = PipelinedShuffleWriterExec::try_new(
+ job_id.clone(),
+ stage.stage_id(),
+ input,
+ stage.shuffle_output_partitioning().cloned(),
+ replayed_stages.contains(&stage.stage_id()),
+ )?;
+ Ok(Arc::new(writer) as Arc<dyn ShuffleWriter>)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let stages =
+
ExecutionStageBuilder::new(session_config.clone()).build(pipelined_stages)?;
+ let final_stage_id = stages
+ .iter()
+ .filter(|(_, stage)| stage.output_links().is_empty())
+ .map(|(stage_id, _)| *stage_id)
+ .max()
+ .ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "pipelined job {job_id} has no final stage"
+ ))
+ })?;
+
+ let started_at = timestamp_millis();
+ Ok(Self {
+ scheduler_id: Some(scheduler_id.to_owned()),
+ job_id: job_id.clone(),
+ job_name: job_name.to_owned(),
+ session_id: session_id.to_owned(),
+ status: JobStatus {
+ job_id: job_id.to_string(),
+ job_name: job_name.to_owned(),
+ status: Some(Status::Running(RunningJob {
+ queued_at,
+ started_at,
+ scheduler: scheduler_id.to_owned(),
+ pipelined_output_locations: vec![],
+ })),
+ },
+ queued_at,
+ start_time: started_at,
+ end_time: 0,
+ stages,
+ final_stage_id,
+ placed: false,
+ reservation: vec![],
+ executors: HashSet::new(),
+ session_config,
+ logical_plan,
+ physical_plan: plan,
+ })
+ }
+
+ fn is_running(&self) -> bool {
+ 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.
+ pub fn awaits_placement(&self) -> bool {
+ !self.placed && self.is_running()
+ }
+
+ /// What placing each stage takes, or `None` once the job is placed or no
+ /// longer running.
+ pub fn placement_requirements(&self) -> Option<Vec<StageRequirement>> {
+ if !self.awaits_placement() {
+ return None;
+ }
+ let mut requirements: Vec<StageRequirement> = self
+ .stages
+ .iter()
+ .map(|(stage_id, stage)| {
+ let plan = stage.plan_arc();
+ StageRequirement {
+ stage_id: *stage_id,
+ partitions: stage_input_partitions(&plan),
+ collapse: stage_has_input_collapse(&plan),
+ }
+ })
+ .collect();
+ requirements.sort_by_key(|r| r.stage_id);
+ Some(requirements)
+ }
+
+ /// The scheduler's `max_partitions_per_task` for this job.
+ pub fn max_partitions_per_task(&self) -> usize {
+ self.session_config
+ .options()
+ .extensions
+ .get::<ballista_core::config::BallistaConfig>()
+ .map(|config| config.max_partitions_per_task())
+ .unwrap_or(0)
+ }
+
+ /// Commits a placement whose vcores the scheduler has reserved: resolves
+ /// every reader to the producer tasks it streams from, marks every stage
+ /// running, publishes the final stage's locations in the job status, and
+ /// returns the tasks to launch.
+ ///
+ /// `executors` must hold the metadata of every executor the placement
+ /// uses.
+ pub fn place(
+ &mut self,
+ placement: &Placement,
+ executors: &HashMap<String, ExecutorMetadata>,
+ ) -> Result<Vec<BoundTask>> {
+ if !self.awaits_placement() {
+ return Err(BallistaError::Internal(format!(
+ "pipelined job {} is already placed or no longer running",
+ self.job_id
+ )));
+ }
+ let missing: Vec<usize> = self
+ .stages
+ .keys()
+ .filter(|stage_id| !placement.tasks.contains_key(stage_id))
+ .copied()
+ .collect();
+ if !missing.is_empty() {
+ return Err(BallistaError::Internal(format!(
+ "placement of pipelined job {} leaves stages {missing:?}
unplaced",
+ self.job_id
+ )));
+ }
+
+ // Where every producer task serves each of its output partitions.
+ let mut outputs: HashMap<
+ usize,
+ BTreeMap<usize, Vec<PipelinedPartitionLocation>>,
+ > = HashMap::new();
+ for (stage_id, tasks) in &placement.tasks {
+ let plan = self.stage(*stage_id)?.plan_arc();
+ let stage_outputs = outputs.entry(*stage_id).or_default();
+ for (task_id, (executor_id, partitions)) in
tasks.iter().enumerate() {
+ let executor_meta = executors.get(executor_id).ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "no metadata for executor {executor_id} placed in
pipelined job {}",
+ self.job_id
+ ))
+ })?;
+ for partition in compute_global_output_partition_ids(&plan,
partitions) {
+ stage_outputs.entry(partition).or_default().push(
+ PipelinedPartitionLocation {
+ partition_id: PartitionId::new(
+ &self.job_id,
+ *stage_id,
+ partition,
+ ),
+ task_id,
+ executor_meta: executor_meta.clone(),
+ },
+ );
+ }
+ }
+ }
+
+ let reserved: HashMap<&str, u32> = placement
+ .reservation
+ .iter()
+ .map(|(executor_id, vcores)| (executor_id.as_str(), *vcores))
+ .collect();
+ let mut bound = vec![];
+ for (stage_id, tasks) in &placement.tasks {
+ let stage = self.stage(*stage_id)?;
+ let (output_links, inputs) = match stage {
+ ExecutionStage::UnResolved(stage) => {
+ (stage.output_links.clone(), stage.inputs.clone())
+ }
+ ExecutionStage::Resolved(stage) => {
+ (stage.output_links.clone(), stage.inputs.clone())
+ }
+ other => {
+ return Err(BallistaError::Internal(format!(
+ "stage {stage_id} of pipelined job {} is {} before
placement",
+ self.job_id,
+ other.variant_name()
+ )));
+ }
+ };
+ let plan = resolve_pipelined_readers(stage.plan_arc(), &outputs)?;
+ let partitions = stage_input_partitions(&plan);
+ let mut running = RunningStage::new(
+ *stage_id,
+ 0,
+ plan.clone(),
+ partitions,
+ output_links,
+ inputs,
+ self.session_config.clone(),
+ );
+ // Every partition is handed out here, so nothing is pending.
+ running.pending = PendingPartitions::empty(partitions);
+ for (task_id, (executor_id, slice)) in tasks.iter().enumerate() {
+ let mut task_info = create_task_info(executor_id.clone(),
task_id);
+ task_info.global_input_partition_ids = slice.clone();
+ // The job's reservation covers every task and is released
+ // when the job ends, so no task carries vcores of its own.
+ task_info.vcores_consumed = 0;
+ running.task_infos.push(task_info);
+ bound.push((
+ executor_id.clone(),
+ TaskDescription {
+ session_id: self.session_id.clone(),
+ key: TaskKey {
+ job_id: self.job_id.clone(),
+ stage_id: *stage_id,
+ task_id,
+ },
+ stage_attempt_num: 0,
+ task_attempt: 0,
+ global_input_partition_ids: slice.clone(),
+ // The job's reservation on the executor, which sizes
+ // the memory pool all its tasks there share.
+ vcores_consumed: reserved
+ .get(executor_id.as_str())
+ .copied()
+ .unwrap_or(1),
+ plan: plan.clone(),
+ session_config: self.session_config.clone(),
+ },
+ ));
+ }
+ self.stages
+ .insert(*stage_id, ExecutionStage::Running(running));
+ }
+
+ let final_locations: Vec<PipelinedPartitionLocation> = outputs
+ .remove(&self.final_stage_id)
+ .unwrap_or_default()
+ .into_values()
+ .flatten()
+ .collect();
+ if let Some(Status::Running(running)) = &mut self.status.status {
+ running.pipelined_output_locations =
+ final_locations.into_iter().map(Into::into).collect();
+ }
+ self.reservation = placement.reservation.clone();
+ self.executors = placement
+ .tasks
+ .values()
+ .flatten()
+ .map(|(executor_id, _)| executor_id.clone())
+ .collect();
+ self.placed = true;
+ info!(
+ "Placed pipelined job {}: {} tasks over {} executors, reserving
{:?}",
+ self.job_id,
+ bound.len(),
+ self.executors.len(),
+ self.reservation
+ );
+ Ok(bound)
+ }
+
+ /// Takes the vcores the job reserved, to be released. Empty once taken,
+ /// so they are released once.
+ pub fn take_reservation(&mut self) -> Vec<ExecutorSlot> {
+ std::mem::take(&mut self.reservation)
+ }
+
+ /// If losing `executor_id` breaks this job, the reason to fail it with.
+ pub fn lost_executor(&self, executor_id: &str) -> Option<String> {
+ (self.is_running() && self.executors.contains(executor_id)).then(|| {
+ format!(
+ "executor {executor_id}, which was running part of pipelined
job {}, was lost",
+ self.job_id
+ )
+ })
+ }
+
+ fn stage(&self, stage_id: usize) -> Result<&ExecutionStage> {
+ self.stages.get(&stage_id).ok_or_else(|| {
+ BallistaError::Internal(format!(
+ "pipelined job {} has no stage {stage_id}",
+ self.job_id
+ ))
+ })
+ }
+
+ /// Marks every stage whose tasks have all succeeded as successful.
+ fn succeed_finished_stages(&mut self) {
+ let finished: Vec<usize> = self
+ .stages
+ .iter()
+ .filter_map(|(stage_id, stage)| match stage {
+ ExecutionStage::Running(running) if running.is_successful() =>
{
+ Some(*stage_id)
+ }
+ _ => None,
+ })
+ .collect();
+ for stage_id in finished {
+ self.succeed_stage(stage_id);
+ }
+ }
+}
+
+/// Replaces every `UnresolvedShuffleExec` in `plan` with a
+/// `PipelinedShuffleReaderExec` over the producer tasks in `outputs`.
+fn resolve_pipelined_readers(
+ plan: Arc<dyn ExecutionPlan>,
+ outputs: &HashMap<usize, BTreeMap<usize, Vec<PipelinedPartitionLocation>>>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ if let Some(unresolved) = plan.downcast_ref::<UnresolvedShuffleExec>() {
+ if unresolved.coalesce.is_some() {
+ return Err(BallistaError::Internal(format!(
+ "pipelined jobs do not support coalesced readers (stage {})",
+ unresolved.stage_id
+ )));
+ }
+ let empty = BTreeMap::new();
+ let produced = outputs.get(&unresolved.stage_id).unwrap_or(&empty);
+ let schema = unresolved.schema.clone();
+ let reader = if unresolved.broadcast {
+ PipelinedShuffleReaderExec::new_broadcast(
+ unresolved.stage_id,
+ produced.values().flatten().cloned().collect(),
+ schema,
+ )
+ } else {
+ let partitions = (0..unresolved.output_partition_count)
+ .map(|p| produced.get(&p).cloned().unwrap_or_default())
+ .collect();
+ PipelinedShuffleReaderExec::try_new(
+ unresolved.stage_id,
+ partitions,
+ schema,
+ unresolved.properties().output_partitioning().clone(),
+ )?
+ };
+ return Ok(Arc::new(reader));
+ }
+ let children = plan
+ .children()
+ .into_iter()
+ .map(|child| resolve_pipelined_readers(child.clone(), outputs))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(replace_children_if_necessary(plan, children)?)
+}
+
+impl ExecutionGraph for PipelinedExecutionGraph {
+ fn job_id(&self) -> &JobId {
+ &self.job_id
+ }
+
+ fn job_name(&self) -> &str {
+ &self.job_name
+ }
+
+ fn session_id(&self) -> &str {
+ &self.session_id
+ }
+
+ fn scheduler_id(&self) -> Option<&str> {
+ self.scheduler_id.as_deref()
+ }
+
+ fn session_config(&self) -> Arc<SessionConfig> {
+ self.session_config.clone()
+ }
+
+ fn status(&self) -> &JobStatus {
+ &self.status
+ }
+
+ fn logical_plan(&self) -> Option<&str> {
+ self.logical_plan.as_deref()
+ }
+
+ fn physical_plan(&self) -> Arc<dyn ExecutionPlan> {
+ self.physical_plan.clone()
+ }
+
+ fn start_time(&self) -> u64 {
+ self.start_time
+ }
+
+ fn end_time(&self) -> u64 {
+ self.end_time
+ }
+
+ fn completed_stages(&self) -> usize {
+ self.stages
+ .values()
+ .filter(|stage| matches!(stage, ExecutionStage::Successful(_)))
+ .count()
+ }
+
+ fn is_successful(&self) -> bool {
+ self.stages
+ .values()
+ .all(|stage| matches!(stage, ExecutionStage::Successful(_)))
+ }
+
+ /// Stages never become runnable one at a time: [`Self::place`] starts
+ /// them all.
+ fn revive(&mut self) -> bool {
+ false
+ }
+
+ /// Records task outcomes. The first failure fails the job; the final
+ /// stage completing finishes it and cancels whatever is still running.
+ fn update_task_status(
+ &mut self,
+ _executor: &ExecutorMetadata,
+ task_statuses: Vec<TaskStatus>,
+ _max_task_failures: usize,
+ _max_stage_failures: usize,
+ ) -> Result<Vec<QueryStageSchedulerEvent>> {
+ if !self.is_running() {
+ // Late reports from tasks cancelled when the job finished or
+ // failed.
+ return Ok(vec![]);
+ }
+ let mut failure = None;
+ for status in task_statuses {
+ let stage_id = status.stage_id as usize;
+ let task_id = status.task_id as usize;
+ let Some(ExecutionStage::Running(stage)) =
self.stages.get_mut(&stage_id)
+ else {
+ warn!(
+ "Ignoring status of task {task_id} in stage {stage_id} of
pipelined job {}: the stage is not running",
+ self.job_id
+ );
+ continue;
+ };
+ if task_id >= stage.task_infos.len() {
+ warn!(
+ "Ignoring status of unknown task {task_id} in stage
{stage_id} of pipelined job {}",
+ self.job_id
+ );
+ continue;
+ }
+ let outcome = status.status.clone();
+ let metrics = status.metrics.clone();
+ if !stage.update_task_info(task_id, status) {
+ continue;
+ }
+ match outcome {
+ Some(task_status::Status::Successful(_)) => {
+ stage.update_task_metrics(task_id, metrics)?;
+ }
+ Some(task_status::Status::Failed(failed)) => {
+ failure.get_or_insert_with(|| {
Review Comment:
**(4)** A `FetchPartitionError` on cache data (files gone, executor alive)
should carry `CACHE_DATA_LOST`, detected with `is_external_fetch_failure` as in
the static graph. Otherwise every later pipelined query on that cache fails.
--
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]
