mbutrovich commented on code in PR #2657:
URL: 
https://github.com/apache/datafusion-ballista/pull/2657#discussion_r4428021174


##########
ballista/scheduler/src/state/mod.rs:
##########
@@ -213,10 +227,14 @@ impl<T: 'static + AsLogicalPlan, U: 'static + 
AsExecutionPlan> SchedulerState<T,
         &self,
         sender: EventSender<QueryStageSchedulerEvent>,
     ) -> Result<()> {
-        let schedulable_tasks = self
-            .executor_manager
-            .bind_schedulable_tasks(self.task_manager.get_running_job_cache())
-            .await?;
+        // Pipelined jobs go first: each needs all of its vcores at once, and

Review Comment:
   Could we file issues for the admission follow-ups before this merges? As I 
read `revive_offers`, a pipelined job that needs more vcores than are ever free 
at the same moment can wait indefinitely while blocking jobs keep binding those 
vcores one task at a time. Pipelined jobs are also tried in `HashMap` order 
([`get_running_job_cache`](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/state/task_manager.rs#L466-L480)),
 so a wide one can wait behind narrower ones that keep fitting. #2652 lists 
fair admission as a follow-up and answers its question 3 with "wait". A client 
process that is killed outright also leaves its job holding its reservation 
until someone cancels it through the REST API, as you found in [this 
comment](https://github.com/apache/datafusion-ballista/pull/2657#issuecomment-6443057416).
 I couldn't find an issue for either.



##########
ballista/scheduler/src/state/pipelined_graph.rs:
##########
@@ -0,0 +1,1753 @@
+// 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::failed_task::FailedReason;
+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::cache::CACHE_DATA_LOST;
+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, is_external_fetch_failure, 
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
+            )
+        })
+    }

Review Comment:
   Should losing an executor also drop it from the job's reservation? 
`executor_lost` discards the blocking ledger's entries for a lost executor 
first, because its capacity is already gone 
([task_manager.rs#L898-L901](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/state/task_manager.rs#L898-L901)).
 The pipelined reservation keeps the lost executor's slot, and `abort_job` 
later hands it to `unbind_tasks`. If the executor has registered again under 
the same id by then, it gets those vcores on top of its fresh capacity. That is 
what a healthy executor does after a transient launch failure removes it: it 
"re-registers on its next heartbeat" 
([config.rs#L425-L427](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/config.rs#L425-L427)).
   
   I tried this with a test in this module. It places an 8-partition job on the 
two 4-vcore virtual executors with tasks that never report, calls 
`lose_executor("virtual-executor-0")`, registers `virtual-executor-0` again 
with 4 vcores, and waits for the job to fail. `free_vcores()` then returns 12. 
With the change below it returns 8. The call site in `executor_lost` already 
goes through `as_pipelined_mut()`. Could you add a test along those lines too?
   
   ```suggestion
       /// If losing `executor_id` breaks this job, the reason to fail it with.
       /// The lost executor's share of the reservation is dropped, since its
       /// capacity is already gone.
       pub fn lost_executor(&mut self, executor_id: &str) -> Option<String> {
           self.reservation.retain(|(id, _)| id != executor_id);
           (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
               )
           })
       }
   ```



##########
ballista/scheduler/src/state/task_manager.rs:
##########
@@ -681,27 +731,36 @@ impl<T: 'static + AsLogicalPlan, U: 'static + 
AsExecutionPlan> TaskManager<T, U>
     }
 
     /// Move a job from Active to Success and return the ids of its 
intermediate
-    /// (non-final) stages so their shuffle data can be reclaimed immediately.
-    /// Returns an empty vec if the job is not found or not successful.
-    pub(crate) async fn succeed_job(&self, job_id: &JobId) -> Result<Vec<u32>> 
{
+    /// (non-final) stages so their shuffle data can be reclaimed immediately,
+    /// along with any vcores the job held for its whole life (a pipelined
+    /// job's reservation), to be released.
+    /// Returns empty vecs if the job is not found or not successful.
+    pub(crate) async fn succeed_job(
+        &self,
+        job_id: &JobId,
+    ) -> Result<(Vec<u32>, Vec<ExecutorSlot>)> {
         debug!("Moving job {job_id} from Active to Success");
 
         if let Some(graph) = self.get_active_execution_graph(job_id) {
-            let (snapshot, intermediate_stage_ids) = {
-                let graph = graph.read().await;
+            let (snapshot, intermediate_stage_ids, reservation) = {
+                let mut graph = graph.write().await;
                 if !graph.is_successful() {
                     error!("Job {job_id} has not finished and cannot be 
completed");
-                    return Ok(vec![]);
+                    return Ok((vec![], vec![]));
                 }
+                let reservation = graph
+                    .as_pipelined_mut()
+                    .map(|graph| graph.take_reservation())
+                    .unwrap_or_default();
 
-                (graph.cloned(), graph.intermediate_stage_ids())
+                (graph.cloned(), graph.intermediate_stage_ids(), reservation)
             };
 
             self.persist_terminal_and_evict(job_id, &snapshot).await?;

Review Comment:
   What happens to the reservation if `persist_terminal_and_evict` fails here? 
By this point it has been taken from the graph. The `?` returns the error, and 
the `JobFinished` handler replaces it with empty vecs 
([query_stage_scheduler.rs#L470-L479](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs#L470-L479)),
 so those vcores are never released. A stage at least as wide as the cluster 
needs every vcore free 
([pipelined_graph.rs#L199](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/state/pipelined_graph.rs#L199)),
 so no such job could be placed after that. `abort_job` releases its slots 
before it persists 
([task_manager.rs#L838-L839](https://github.com/apache/datafusion-ballista/blob/81ff2e31f053fea26a293a79deb78bfe3de51ee7/ballista/scheduler/src/state/task_manager.rs#L838-L839)).
 Could `succeed_job` hand back th
 e reservation even when persisting fails?



##########
docs/source/user-guide/configs.md:
##########
@@ -170,15 +170,15 @@ explanation of the sort-based shuffle writer.
 
 These session-level keys select and tune the experimental pipelined execution
 mode, in which every task of a job is placed up front and stages stream data
-to each other instead of exchanging shuffle files. The scheduler does not run
-pipelined jobs yet.
+to each other instead of exchanging shuffle files. See the
+[pipelined execution guide](pipelined-execution.md).

Review Comment:
   `pipelined-execution.md` comes from #2713, which merges after this PR, so 
this link is broken on the published docs until then. Could the link move to 
#2713?
   
   ```suggestion
   These session-level keys select and tune the experimental pipelined execution
   mode, in which every task of a job is placed up front and stages stream data
   to each other instead of exchanging shuffle files.
   ```



-- 
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]

Reply via email to