fix: Cancel prepare_task_definition fail job

Author: westhide

ballista/scheduler/src/cluster/mod.rs

@@ -44,7 +44,9 @@ use crate::cluster::memory::{InMemoryClusterState, InMemoryJobState};
 
 use crate::config::{ClusterStorageConfig, SchedulerConfig, TaskDistributionPolicy};
 use crate::scheduler_server::SessionBuilder;
-use crate::state::execution_graph::{create_task_info, ExecutionGraph, TaskDescription};
+use crate::state::execution_graph::{
+    create_task_info, ExecutionGraph, ExecutionStage, TaskDescription,
+};
 use crate::state::task_manager::JobInfoCache;
 
 pub mod event;
@@ -647,6 +649,29 @@ pub(crate) fn get_scan_files(
     Ok(collector)
 }
 
+pub(crate) async fn unbind_prepare_failed_tasks(
+    active_jobs: Arc<HashMap<String, JobInfoCache>>,
+    failed_jobs: &HashMap<String, Vec<TaskDescription>>,
+) {
+    for (job_id, failed_tasks) in failed_jobs.iter() {
+        if let Some(job_info) = active_jobs.get(job_id) {
+            let mut graph = job_info.execution_graph.write().await;
+            failed_tasks.iter().for_each(|task| {
+                if let Some(ExecutionStage::Running(running_stage)) =
+                    graph.stages_mut().get_mut(&task.partition.stage_id)
+                {
+                    if let Some(task_info) = running_stage
+                        .task_infos
+                        .get_mut(task.partition.partition_id)
+                    {
+                        *task_info = None;
+                    }
+                }
+            });
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct TopologyNode {
     pub id: String,

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -35,13 +35,17 @@ use ballista_core::serde::scheduler::ExecutorMetadata;
 use datafusion_proto::logical_plan::AsLogicalPlan;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 use log::{debug, error, info, trace, warn};
+use std::collections::HashMap;
 use std::net::SocketAddr;
 
 use std::ops::Deref;
 
-use crate::cluster::{bind_task_bias, bind_task_round_robin};
+use crate::cluster::{
+    bind_task_bias, bind_task_round_robin, unbind_prepare_failed_tasks,
+};
 use crate::config::TaskDistributionPolicy;
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
+use crate::state::execution_graph::TaskDescription;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tonic::{Request, Response, Status};
 
@@ -112,10 +116,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
             let active_jobs = self.state.task_manager.get_running_job_cache();
             let schedulable_tasks = match self.state.config.task_distribution {
                 TaskDistributionPolicy::Bias => {
-                    bind_task_bias(available_slots, active_jobs, |_| false).await
+                    bind_task_bias(available_slots, active_jobs.clone(), |_| false).await
                 }
                 TaskDistributionPolicy::RoundRobin => {
-                    bind_task_round_robin(available_slots, active_jobs, |_| false).await
+                    bind_task_round_robin(available_slots, active_jobs.clone(), |_| false).await
                 }
                 TaskDistributionPolicy::ConsistentHash{..} => {
                     return Err(Status::unimplemented(
@@ -124,14 +128,36 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
             };
 
             let mut tasks = vec![];
+            let mut prepare_failed_jobs = HashMap::<String, Vec<TaskDescription>>::new();
             for (_, task) in schedulable_tasks {
-                match self.state.task_manager.prepare_task_definition(task) {
+                let job_id = task.partition.job_id.clone();
+                if prepare_failed_jobs.contains_key(&job_id) {
+                    prepare_failed_jobs.entry(job_id).or_default().push(task);
+                    continue;
+                }
+                match self
+                    .state
+                    .task_manager
+                    .prepare_task_definition(task.clone())
+                {
                     Ok(task_definition) => tasks.push(task_definition),
                     Err(e) => {
                         error!("Error preparing task definition: {:?}", e);
+                        prepare_failed_jobs.entry(job_id).or_default().push(task);
                     }
                 }
             }
+
+            unbind_prepare_failed_tasks(active_jobs, &prepare_failed_jobs).await;
+            for job_id in prepare_failed_jobs.into_keys() {
+                info!("Cancel prepare task definition failed job: {}", job_id);
+                self.cancel_job(job_id).await.map_err(|e| {
+                    let msg = format!("Cancel job error due to {e:?}");
+                    error!("{}", msg);
+                    Status::internal(msg)
+                })?;
+            }
+
             Ok(Response::new(PollWorkResult { tasks }))
         } else {
             warn!("Received invalid executor poll_work request");
@@ -527,21 +553,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
     ) -> Result<Response<CancelJobResult>, Status> {
         let job_id = request.into_inner().job_id;
         info!("Received cancellation request for job {}", job_id);
-
-        self.query_stage_event_loop
-            .get_sender()
-            .map_err(|e| {
-                let msg = format!("Get query stage event loop error due to {e:?}");
-                error!("{}", msg);
-                Status::internal(msg)
-            })?
-            .post_event(QueryStageSchedulerEvent::JobCancel(job_id))
-            .await
-            .map_err(|e| {
-                let msg = format!("Post to query stage event loop error due to {e:?}");
-                error!("{}", msg);
-                Status::internal(msg)
-            })?;
+        self.cancel_job(job_id).await.map_err(|e| {
+            let msg = format!("Cancel job error due to {e:?}");
+            error!("{}", msg);
+            Status::internal(msg)
+        })?;
         Ok(Response::new(CancelJobResult { cancelled: true }))
     }
 

ballista/scheduler/src/scheduler_server/mod.rs

@@ -183,6 +183,13 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             .await
     }
 
+    pub(crate) async fn cancel_job(&self, job_id: String) -> Result<()> {
+        self.query_stage_event_loop
+            .get_sender()?
+            .post_event(QueryStageSchedulerEvent::JobCancel(job_id))
+            .await
+    }
+
     /// It just send task status update event to the channel,
     /// and will not guarantee the event processing completed after return
     pub(crate) async fn update_task_status(

ballista/scheduler/src/state/execution_graph.rs

@@ -225,6 +225,10 @@ impl ExecutionGraph {
         &self.stages
     }
 
+    pub(crate) fn stages_mut(&mut self) -> &mut HashMap<usize, ExecutionStage> {
+        &mut self.stages
+    }
+
     /// An ExecutionGraph is successful if all its stages are successful
     pub fn is_successful(&self) -> bool {
         self.stages