diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index ae8098bce9..3935ce0174 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -1028,8 +1028,12 @@ message LaunchTaskResult { } message LaunchMultiTaskResult { - bool success = 1; - // TODO when part of the task set are scheduled successfully + reserved 1; + reserved "success"; + // Job IDs the executor could not decode/validate. These jobs are failed + // individually while the rest of the batch still runs; an empty list means + // the whole batch was accepted. A successful RPC does not imply every job ran. + repeated string failed_jobs = 2; } message CancelTasksParams { diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 9e05f2ed28..25b92def4f 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -1483,11 +1483,13 @@ pub struct LaunchTaskResult { #[prost(bool, tag = "1")] pub success: bool, } -#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct LaunchMultiTaskResult { - /// TODO when part of the task set are scheduled successfully - #[prost(bool, tag = "1")] - pub success: bool, + /// Job IDs the executor could not decode/validate. These jobs are failed + /// individually while the rest of the batch still runs; an empty list means + /// the whole batch was accepted. A successful RPC does not imply every job ran. + #[prost(string, repeated, tag = "2")] + pub failed_jobs: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct CancelTasksParams { diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs index 0d56417e8e..acb509405f 100644 --- a/ballista/executor/src/executor_server.rs +++ b/ballista/executor/src/executor_server.rs @@ -23,7 +23,7 @@ use ballista_core::BALLISTA_VERSION; use memory_stats::memory_stats; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::convert::TryInto; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -897,8 +897,10 @@ impl ExecutorGrpc scheduler_id, } = request.into_inner(); let task_sender = self.executor_env.tx_task.clone(); + let mut failed_jobs: HashSet = HashSet::new(); for multi_task in multi_tasks { - let multi_task: Vec = get_task_definition_vec( + let job_id = multi_task.job_id.clone(); + let multi_task: Vec = match get_task_definition_vec( multi_task, self.executor.runtime_producer.clone(), self.executor.produce_config(), @@ -910,8 +912,15 @@ impl ExecutorGrpc .higher_order_functions .clone(), self.codec.clone(), - ) - .map_err(|e| Status::invalid_argument(format!("{e}")))?; + ) { + Ok(tasks) => tasks, + Err(e) => { + error!("failed to decode tasks for {job_id} : {e}"); + failed_jobs.insert(job_id); + continue; + } + }; + for task in multi_task { task_sender .send(CuratorTaskDefinition { @@ -922,7 +931,9 @@ impl ExecutorGrpc .unwrap(); } } - Ok(Response::new(LaunchMultiTaskResult { success: true })) + Ok(Response::new(LaunchMultiTaskResult { + failed_jobs: failed_jobs.into_iter().collect(), + })) } async fn stop_executor( diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 2e8684dbff..c5c293c464 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -573,9 +573,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] @@ -1134,6 +1134,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> { diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index da8d7470b7..5ba270e1ef 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -419,28 +419,28 @@ impl ExecutorManager { } /// Launches multiple tasks on the specified executor. + /// + /// `Ok` means the RPC was dispatched; the returned set holds job IDs the + /// executor rejected (could not decode) and failed individually. `Err` is + /// only returned for a transport-level failure of the whole RPC. pub async fn launch_multi_task( &self, executor_id: &str, multi_tasks: Vec, scheduler_id: String, - ) -> Result<()> { + ) -> Result> { let mut client = self .get_client(executor_id, &self.grpc_client_config) .await?; - client + let res = client .launch_multi_task(protobuf::LaunchMultiTaskParams { multi_tasks, scheduler_id, }) - .await - .map_err(|e| { - BallistaError::Internal(format!( - "Failed to connect to executor {executor_id}: {e:?}" - )) - })?; + .await? + .into_inner(); - Ok(()) + Ok(res.failed_jobs.into_iter().map(JobId::from).collect()) } pub(crate) fn drain_pending_cleanup_jobs( diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index f4a6ebec4c..a23db7af2f 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/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; @@ -33,7 +34,7 @@ use datafusion_proto::physical_plan::AsExecutionPlan; use log::{debug, error, info, warn}; use prost::Message; use std::any::type_name; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Instant; @@ -191,7 +192,7 @@ impl SchedulerState { + Ok((unassigned_executor_slots, failed_jobs)) => { if !unassigned_executor_slots.is_empty() { if let Err(e) = state .executor_manager @@ -202,6 +203,20 @@ impl SchedulerState { error!("Fail to launch tasks: {e}"); @@ -261,11 +276,14 @@ impl SchedulerState, sender: &EventSender, - ) -> Result> { + ) -> Result<(Vec, HashSet)> { // 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< @@ -289,70 +307,82 @@ impl SchedulerState> = tasks.into_values().collect(); // Total number of tasks to be launched for one executor let n_tasks: usize = tasks.iter().map(|stage_tasks| stage_tasks.len()).sum(); - let state = self.clone(); let sender = sender.clone(); let join_handle = tokio::spawn(async move { - let success = match state + let job_ids: Vec = tasks + .iter() + .flatten() + .map(|t| t.key.job_id.clone()) + .collect(); + match state .executor_manager .get_executor_metadata(&executor_id) .await { Ok(executor) => { - if let Err(e) = state + match state .task_manager .launch_multi_task(&executor, tasks, &state.executor_manager) .await { - 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), &sender) - .await; - - false - } else { - true + Ok(rejected) => { + let freed = job_ids + .iter() + .filter(|j| rejected.contains(*j)) + .count() + as u32; + (vec![(executor_id.clone(), freed)], rejected) + } + Err(e) => { + 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), &sender) + .await; + + ( + vec![(executor_id.clone(), n_tasks as u32)], + HashSet::new(), + ) + } } } Err(e) => { error!( "Failed to launch new task, could not get executor metadata: {e}" ); - false + (vec![(executor_id.clone(), n_tasks as u32)], HashSet::new()) } - }; - if success { - vec![] - } else { - vec![(executor_id.clone(), n_tasks as u32)] } }); join_handles.push(join_handle); } - let unassigned_executor_slots = - futures::future::join_all(join_handles) - .await - .into_iter() - .collect::>, - tokio::task::JoinError, - >>()?; - - Ok(unassigned_executor_slots + let results = futures::future::join_all(join_handles) + .await .into_iter() - .flatten() - .collect::>()) + .collect::, HashSet)>, + tokio::task::JoinError, + >>()?; + + let mut unassigned_executor_slots = Vec::new(); + let mut failed_jobs = HashSet::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( @@ -421,3 +451,153 @@ impl SchedulerState, + _executor_manager: &ExecutorManager, + ) -> Result> { + Ok(tasks + .iter() + .map(|t| JobId::from(t.job_id.clone())) + .filter(|j| j == &self.reject) + .take(1) + .collect()) + } + } + + fn agg_plan() -> LogicalPlan { + let schema = Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("gmv", DataType::UInt64, false), + ]); + scan_empty_with_partitions(None, &schema, Some(vec![0, 1]), 2) + .unwrap() + .aggregate(vec![col("id")], vec![sum(col("gmv"))]) + .unwrap() + .build() + .unwrap() + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn launch_tasks_isolates_single_rejected_job() -> Result<()> { + let bad_job = JobId::from("job-bad"); + let good_job = JobId::from("job-good"); + + let state: SchedulerState = + SchedulerState::new_with_task_launcher( + test_cluster_context(), + BallistaCodec::default(), + "localhost:50050".to_owned(), + Arc::new(SchedulerConfig::default()), + Arc::new(RejectOne { + reject: bad_job.clone(), + }), + ); + + let vcores = 8; + state + .executor_manager + .register_executor( + ExecutorMetadata { + id: "executor-1".to_string(), + host: String::default(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default().with_vcores(vcores), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + ExecutorData { + executor_id: "executor-1".to_string(), + total_vcores: vcores, + available_vcores: vcores, + }, + ) + .await?; + + let ctx = state + .session_manager + .create_or_update_session("session", &SessionConfig::new_with_ballista()) + .await?; + + for job_id in [&good_job, &bad_job] { + state + .task_manager + .queue_job(job_id, "", timestamp_millis())?; + state + .task_manager + .submit_job( + job_id, + "", + ctx.clone(), + &agg_plan(), + timestamp_millis(), + None, + ) + .await?; + } + + let bound = state + .executor_manager + .bind_schedulable_tasks(state.task_manager.get_running_job_cache()) + .await?; + + let bad_task_count = bound + .iter() + .filter(|(_, t)| t.key.job_id == bad_job) + .count() as u32; + let good_task_count = bound + .iter() + .filter(|(_, t)| t.key.job_id == good_job) + .count() as u32; + assert!(bad_task_count > 0 && good_task_count > 0); + + let (tx_event, _rx_event) = tokio::sync::mpsc::channel(100); + let sender = EventSender::new(tx_event); + let (unassigned_slots, failed_jobs) = state.launch_tasks(bound, &sender).await?; + + assert_eq!(failed_jobs, HashSet::from([bad_job.clone()])); + + let freed: u32 = unassigned_slots.iter().map(|(_, n)| *n).sum(); + assert_eq!(freed, bad_task_count); + + assert!( + state + .executor_manager + .get_executor_metadata("executor-1") + .await + .is_ok() + ); + + Ok(()) + } +} diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 60e5330478..8a0a3b81db 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -61,12 +61,16 @@ type ActiveJobCache = Arc>; #[async_trait::async_trait] pub trait TaskLauncher: Send + Sync + 'static { /// Launches the given tasks on the specified executor. + /// + /// `Ok` means the RPC was dispatched; the returned set holds job IDs the + /// executor rejected and failed individually. `Err` is only for a + /// transport-level failure of the whole RPC. async fn launch_tasks( &self, executor: &ExecutorMetadata, tasks: Vec, executor_manager: &ExecutorManager, - ) -> Result<()>; + ) -> Result>; } struct DefaultTaskLauncher { @@ -86,7 +90,7 @@ impl TaskLauncher for DefaultTaskLauncher { executor: &ExecutorMetadata, tasks: Vec, executor_manager: &ExecutorManager, - ) -> Result<()> { + ) -> Result> { if log::max_level() >= log::Level::Info { let tasks_ids: Vec = tasks .iter() @@ -104,10 +108,10 @@ impl TaskLauncher for DefaultTaskLauncher { executor.id, tasks_ids ); } - executor_manager + let res = executor_manager .launch_multi_task(&executor.id, tasks, self.scheduler_id.clone()) .await?; - Ok(()) + Ok(res) } } @@ -820,7 +824,7 @@ impl TaskManager executor: &ExecutorMetadata, tasks: Vec>, executor_manager: &ExecutorManager, - ) -> Result<()> { + ) -> Result> { let mut multi_tasks = vec![]; for stage_tasks in tasks { match self.prepare_multi_task_definition(stage_tasks) { @@ -834,7 +838,7 @@ impl TaskManager .launch_tasks(executor, multi_tasks, executor_manager) .await } else { - Ok(()) + Ok(HashSet::new()) } } diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 25604673b4..e0bd34fc69 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -341,8 +341,8 @@ impl TaskLauncher for BlackholeTaskLauncher { _executor: &ExecutorMetadata, _tasks: Vec, _executor_manager: &ExecutorManager, - ) -> Result<()> { - Ok(()) + ) -> Result> { + Ok(HashSet::new()) } } @@ -363,7 +363,7 @@ impl TaskLauncher for VirtualTaskLauncher { executor: &ExecutorMetadata, tasks: Vec, _executor_manager: &ExecutorManager, - ) -> Result<()> { + ) -> Result> { if self.unreachable.lock().contains(&executor.id) { return Err(BallistaError::Internal(format!( "test: executor {} is unreachable", @@ -388,7 +388,30 @@ impl TaskLauncher for VirtualTaskLauncher { .await .map_err(|e| { BallistaError::Internal(format!("Error sending task status: {e:?}")) - }) + })?; + Ok(HashSet::new()) + } +} + +/// Launcher that reports every job in the batch as rejected via the +/// `failed_jobs` channel, simulating an executor that cannot decode/validate +/// the task (see issue #1908). The RPC itself succeeds; the jobs are failed +/// individually rather than the whole batch. +#[derive(Default)] +pub struct RejectingTaskLauncher {} + +#[async_trait::async_trait] +impl TaskLauncher for RejectingTaskLauncher { + async fn launch_tasks( + &self, + _executor: &ExecutorMetadata, + tasks: Vec, + _executor_manager: &ExecutorManager, + ) -> Result> { + Ok(tasks + .iter() + .map(|t| JobId::from(t.job_id.clone())) + .collect()) } } @@ -489,6 +512,84 @@ impl SchedulerTest { }) } + /// Like [`SchedulerTest::new`] but injects a custom [`TaskLauncher`]. + pub async fn new_with_launcher( + config: SchedulerConfig, + metrics_collector: Arc, + num_executors: usize, + task_slots_per_executor: usize, + runner: Option>, + launcher: Arc, + ) -> Result { + 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 = (0..num_executors) + .map(|i| { + let id = format!("virtual-executor-{i}"); + let executor = VirtualExecutor { + executor_id: id.clone(), + vcores: 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 = + 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 { vcores, .. }) in executors { + let metadata = ExecutorMetadata { + id: executor_id.clone(), + host: String::default(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default() + .with_vcores(vcores as u32), + os_info: ExecutorOperatingSystemSpecification::default(), + }; + + let executor_data = ExecutorData { + executor_id, + total_vcores: vcores as u32, + available_vcores: vcores as u32, + }; + + scheduler + .state + .executor_manager + .register_executor(metadata, executor_data) + .await?; + } + + Ok(Self { + scheduler, + session_config, + status_receiver: Some(status_receiver), + unreachable_executors: Arc::default(), + }) + } + /// Returns the number of pending jobs. pub fn pending_job_number(&self) -> usize { self.scheduler.pending_job_number()