bug: Prevent exec hang when task decode fails

ballista/scheduler/src/scheduler_server/mod.rs

@@ -482,9 +482,9 @@ mod test {
     use crate::scheduler_server::{SchedulerServer, timestamp_millis};
 
     use crate::test_utils::{
-        ExplodingTableProvider, SchedulerTest, TaskRunnerFn, TestMetricsCollector,
-        assert_completed_event, assert_failed_event, assert_no_submitted_event,
-        assert_submitted_event, test_cluster_context,
+        ExplodingTableProvider, RejectingTaskLauncher, SchedulerTest, TaskRunnerFn,
+        TestMetricsCollector, assert_completed_event, assert_failed_event,
+        assert_no_submitted_event, assert_submitted_event, test_cluster_context,
     };
 
     #[tokio::test]
@@ -1013,6 +1013,49 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn deterministic_launch_rejection_fails_job() -> Result<()> {
+        // A launcher that always rejects with gRPC InvalidArgument (an executor that
+        // cannot decode the task). The job must fail fast instead of hanging (#1908).
+        let metrics_collector = Arc::new(TestMetricsCollector::default());
+        let mut test = SchedulerTest::new_with_launcher(
+            SchedulerConfig::default()
+                .with_scheduler_policy(TaskSchedulingPolicy::PushStaged),
+            metrics_collector,
+            1,
+            1,
+            None,
+            Arc::new(RejectingTaskLauncher::default()),
+        )
+        .await?;
+
+        let plan = test_plan();
+        let job_id = test.submit("", &plan).await?;
+
+        // Hard wall-clock bound so a stuck job fails the test instead of hanging.
+        let status = tokio::time::timeout(
+            std::time::Duration::from_secs(10),
+            test.await_completion(&job_id),
+        )
+        .await
+        .expect(
+            "job did not reach a terminal state within 10s — likely not being failed",
+        )?;
+
+        assert!(
+            matches!(
+                status,
+                JobStatus {
+                    status: Some(job_status::Status::Failed(_)),
+                    ..
+                }
+            ),
+            "expected job to fail on task rejection, got {status:?}"
+        );
+
+        Ok(())
+    }
+
     async fn test_scheduler(
         scheduling_policy: TaskSchedulingPolicy,
     ) -> Result<SchedulerServer<LogicalPlanNode, PhysicalPlanNode>> {

ballista/scheduler/src/state/executor_manager.rs

@@ -417,12 +417,7 @@ impl ExecutorManager {
                 multi_tasks,
                 scheduler_id,
             })
-            .await
-            .map_err(|e| {
-                BallistaError::Internal(format!(
-                    "Failed to connect to executor {executor_id}: {e:?}"
-                ))
-            })?;
+            .await?;
 
         Ok(())
     }

ballista/scheduler/src/state/mod.rs

@@ -18,6 +18,7 @@
 use crate::cluster::{BallistaCluster, BoundTask, ExecutorSlot};
 use crate::config::SchedulerConfig;
 use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan};
+use crate::scheduler_server::timestamp_millis;
 use crate::state::execution_graph::TaskDescription;
 use crate::state::executor_manager::ExecutorManager;
 use crate::state::session_manager::SessionManager;
@@ -36,6 +37,7 @@ use std::any::type_name;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Instant;
+use tonic::Code;
 
 mod aqe;
 mod distributed_explain;
@@ -191,7 +193,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         tokio::spawn(async move {
             let mut if_revive = false;
             match state.launch_tasks(schedulable_tasks).await {
-                Ok(unassigned_executor_slots) => {
+                Ok((unassigned_executor_slots, failed_jobs)) => {
                     if !unassigned_executor_slots.is_empty() {
                         if let Err(e) = state
                             .executor_manager
@@ -202,6 +204,20 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
                         }
                         if_revive = true;
                     }
+                    for job in failed_jobs {
+                        if let Err(e) = sender
+                            .post_event(QueryStageSchedulerEvent::JobRunningFailed {
+                                job_id: job,
+                                fail_message: "task serialization failed by executor"
+                                    .to_string(),
+                                queued_at: timestamp_millis(),
+                                failed_at: timestamp_millis(),
+                            })
+                            .await
+                        {
+                            error!("Fail to post JobRunningFailed: {e:?}");
+                        }
+                    }
                 }
                 Err(e) => {
                     error!("Fail to launch tasks: {e}");
@@ -260,7 +276,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
     async fn launch_tasks(
         &self,
         bound_tasks: Vec<BoundTask>,
-    ) -> Result<Vec<ExecutorSlot>> {
+    ) -> Result<(Vec<ExecutorSlot>, Vec<JobId>)> {
         // Put tasks to the same executor together
         // And put tasks belonging to the same stage together for creating MultiTaskDefinition
         let mut executor_stage_assignments: HashMap<
@@ -284,7 +300,6 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
                 executor_stage_assignments.insert(executor_id, executor_stage_tasks);
             }
         }
-
         let mut join_handles = vec![];
         for (executor_id, tasks) in executor_stage_assignments.into_iter() {
             let tasks: Vec<Vec<TaskDescription>> = tasks.into_values().collect();
@@ -293,6 +308,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
 
             let state = self.clone();
             let join_handle = tokio::spawn(async move {
+                let mut failed_jobs: Vec<JobId> = Vec::new();
+                let job_ids: Vec<JobId> = tasks
+                    .iter()
+                    .flatten()
+                    .map(|t| t.key.job_id.clone())
+                    .collect();
                 let success = match state
                     .executor_manager
                     .get_executor_metadata(&executor_id)
@@ -304,12 +325,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
                             .launch_multi_task(&executor, tasks, &state.executor_manager)
                             .await
                         {
-                            let err_msg = format!("Failed to launch new task: {e}");
-                            error!("{}", err_msg.clone());
+                            let is_task_rejected = matches!(&e,BallistaError::GrpcError(status)
+                                if status.code() == Code::InvalidArgument);
+                            failed_jobs = if is_task_rejected {
+                                job_ids
+                            } else {
+                                let err_msg = format!("Failed to launch new task: {e}");
+                                error!("{}", err_msg.clone());
 
-                            // It's OK to remove executor aggressively,
-                            // since if the executor is in healthy state, it will be registered again.
-                            state.remove_executor(&executor_id, Some(err_msg)).await;
+                                // It's OK to remove executor aggressively,
+                                // since if the executor is in healthy state, it will be registered again.
+                                state.remove_executor(&executor_id, Some(err_msg)).await;
+                                Vec::new()
+                            };
 
                             false
                         } else {
@@ -324,27 +352,30 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
                     }
                 };
                 if success {
-                    vec![]
+                    (vec![], vec![])
                 } else {
-                    vec![(executor_id.clone(), n_tasks as u32)]
+                    (vec![(executor_id.clone(), n_tasks as u32)], failed_jobs)
                 }
             });
             join_handles.push(join_handle);
         }
 
-        let unassigned_executor_slots =
-            futures::future::join_all(join_handles)
-                .await
-                .into_iter()
-                .collect::<std::result::Result<
-                    Vec<Vec<ExecutorSlot>>,
-                    tokio::task::JoinError,
-                >>()?;
-
-        Ok(unassigned_executor_slots
+        let results = futures::future::join_all(join_handles)
+            .await
             .into_iter()
-            .flatten()
-            .collect::<Vec<ExecutorSlot>>())
+            .collect::<std::result::Result<
+            Vec<(Vec<ExecutorSlot>, Vec<JobId>)>,
+            tokio::task::JoinError,
+        >>()?;
+
+        let mut unassigned_executor_slots = Vec::new();
+        let mut failed_jobs = Vec::new();
+        for (slots, jobs) in results {
+            unassigned_executor_slots.extend(slots);
+            failed_jobs.extend(jobs);
+        }
+
+        Ok((unassigned_executor_slots, failed_jobs))
     }
 
     pub(crate) async fn update_task_statuses(

ballista/scheduler/src/test_utils.rs

@@ -380,6 +380,25 @@ impl TaskLauncher for VirtualTaskLauncher {
     }
 }
 
+/// Launcher that rejects every launch with a deterministic gRPC `InvalidArgument`,
+/// simulating an executor that cannot decode/validate the task (see issue #1908).
+#[derive(Default)]
+pub struct RejectingTaskLauncher {}
+
+#[async_trait::async_trait]
+impl TaskLauncher for RejectingTaskLauncher {
+    async fn launch_tasks(
+        &self,
+        _executor: &ExecutorMetadata,
+        _tasks: Vec<MultiTaskDefinition>,
+        _executor_manager: &ExecutorManager,
+    ) -> Result<()> {
+        Err(BallistaError::GrpcError(Box::new(
+            tonic::Status::invalid_argument("undecodable task"),
+        )))
+    }
+}
+
 /// Test harness for scheduler integration tests.
 pub struct SchedulerTest {
     scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode>,
@@ -468,6 +487,83 @@ impl SchedulerTest {
         })
     }
 
+    /// Like [`SchedulerTest::new`] but injects a custom [`TaskLauncher`].
+    pub async fn new_with_launcher(
+        config: SchedulerConfig,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
+        num_executors: usize,
+        task_slots_per_executor: usize,
+        runner: Option<Arc<dyn TaskRunner>>,
+        launcher: Arc<dyn TaskLauncher>,
+    ) -> Result<Self> {
+        let cluster = BallistaCluster::new_from_config(&config).await?;
+
+        let session_config = if num_executors > 0 && task_slots_per_executor > 0 {
+            SessionConfig::new_with_ballista()
+                .with_target_partitions(num_executors * task_slots_per_executor)
+        } else {
+            SessionConfig::new_with_ballista()
+        };
+
+        let runner = runner.unwrap_or_else(|| Arc::new(default_task_runner()));
+
+        let executors: HashMap<String, VirtualExecutor> = (0..num_executors)
+            .map(|i| {
+                let id = format!("virtual-executor-{i}");
+                let executor = VirtualExecutor {
+                    executor_id: id.clone(),
+                    task_slots: task_slots_per_executor,
+                    runner: runner.clone(),
+                };
+                (id, executor)
+            })
+            .collect();
+
+        // This launcher does not report task statuses back, so no receiver is needed.
+        let (_status_sender, status_receiver) = channel(1000);
+
+        let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+            SchedulerServer::new_with_task_launcher(
+                "localhost:50050".to_owned(),
+                cluster,
+                BallistaCodec::default(),
+                Arc::new(config),
+                metrics_collector,
+                launcher,
+            );
+        scheduler.init().await?;
+
+        for (executor_id, VirtualExecutor { task_slots, .. }) in executors {
+            let metadata = ExecutorMetadata {
+                id: executor_id.clone(),
+                host: String::default(),
+                port: 0,
+                grpc_port: 0,
+                specification: ExecutorSpecification::default()
+                    .with_task_slots(task_slots as u32),
+                os_info: ExecutorOperatingSystemSpecification::default(),
+            };
+
+            let executor_data = ExecutorData {
+                executor_id,
+                total_task_slots: task_slots as u32,
+                available_task_slots: task_slots as u32,
+            };
+
+            scheduler
+                .state
+                .executor_manager
+                .register_executor(metadata, executor_data)
+                .await?;
+        }
+
+        Ok(Self {
+            scheduler,
+            session_config,
+            status_receiver: Some(status_receiver),
+        })
+    }
+
     /// Returns the number of pending jobs.
     pub fn pending_job_number(&self) -> usize {
         self.scheduler.pending_job_number()