This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 30c28b177 bug: Prevent exec hang when task decode fails (#2016)
30c28b177 is described below
commit 30c28b17758140602bec3c72a121674d6680ab1c
Author: Bhargava Vadlamani <[email protected]>
AuthorDate: Wed Sep 9 13:47:27 2026 -0700
bug: Prevent exec hang when task decode fails (#2016)
---
ballista/core/proto/ballista.proto | 8 +-
ballista/core/src/serde/generated/ballista.rs | 10 +-
ballista/executor/src/executor_server.rs | 21 +-
ballista/scheduler/src/scheduler_server/mod.rs | 49 ++++-
ballista/scheduler/src/state/executor_manager.rs | 18 +-
ballista/scheduler/src/state/mod.rs | 254 +++++++++++++++++++----
ballista/scheduler/src/state/task_manager.rs | 16 +-
ballista/scheduler/src/test_utils.rs | 109 +++++++++-
8 files changed, 415 insertions(+), 70 deletions(-)
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index 3fdfbdccb..f311e8b08 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -1262,8 +1262,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 313c56632..af5491b91 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -1741,11 +1741,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 6c009ce2c..c92caba6e 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};
@@ -910,8 +910,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorGrpc
scheduler_id,
} = request.into_inner();
let task_sender = self.executor_env.tx_task.clone();
+ let mut failed_jobs: HashSet<String> = HashSet::new();
for multi_task in multi_tasks {
- let multi_task: Vec<TaskDefinition> = get_task_definition_vec(
+ let job_id = multi_task.job_id.clone();
+ let multi_task: Vec<TaskDefinition> = match
get_task_definition_vec(
multi_task,
self.executor.runtime_producer.clone(),
self.executor.produce_config(),
@@ -923,8 +925,15 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> 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 {
@@ -935,7 +944,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> 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 d267bb725..c03a5ead9 100644
--- a/ballista/scheduler/src/scheduler_server/mod.rs
+++ b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -550,9 +550,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]
@@ -1112,6 +1112,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>> {
diff --git a/ballista/scheduler/src/state/executor_manager.rs
b/ballista/scheduler/src/state/executor_manager.rs
index da8d7470b..5ba270e1e 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<MultiTaskDefinition>,
scheduler_id: String,
- ) -> Result<()> {
+ ) -> Result<HashSet<JobId>> {
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 f4a6ebec4..a23db7af2 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<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
tokio::spawn(async move {
let mut if_revive = false;
match state.launch_tasks(schedulable_tasks, &sender).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 +203,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}");
@@ -261,11 +276,14 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
/// 2. Then launch the task set vector to each executor one by one.
///
/// If it fails to launch a task set, the related [`ExecutorSlot`] will be
returned.
+ ///
+ /// Returns the freed executor slots and the set of job IDs the executors
+ /// rejected (failed individually while the rest of the batch ran).
async fn launch_tasks(
&self,
bound_tasks: Vec<BoundTask>,
sender: &EventSender<QueryStageSchedulerEvent>,
- ) -> Result<Vec<ExecutorSlot>> {
+ ) -> Result<(Vec<ExecutorSlot>, HashSet<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<
@@ -289,70 +307,82 @@ 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();
// 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<JobId> = 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::<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>, HashSet<JobId>)>,
+ 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<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerState<T,
);
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::config::SchedulerConfig;
+ use crate::scheduler_server::timestamp_millis;
+ use crate::state::executor_manager::ExecutorManager;
+ use crate::state::task_manager::TaskLauncher;
+ use crate::test_utils::test_cluster_context;
+ use ballista_core::extension::SessionConfigExt;
+ use ballista_core::serde::BallistaCodec;
+ use ballista_core::serde::protobuf::MultiTaskDefinition;
+ use ballista_core::serde::scheduler::{
+ ExecutorData, ExecutorMetadata, ExecutorOperatingSystemSpecification,
+ ExecutorSpecification,
+ };
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::execution::context::SessionConfig;
+ use datafusion::functions_aggregate::sum::sum;
+ use datafusion::logical_expr::{LogicalPlan, col};
+ use datafusion::test_util::scan_empty_with_partitions;
+ use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+
+ struct RejectOne {
+ reject: JobId,
+ }
+
+ #[async_trait::async_trait]
+ impl TaskLauncher for RejectOne {
+ async fn launch_tasks(
+ &self,
+ _executor: &ExecutorMetadata,
+ tasks: Vec<MultiTaskDefinition>,
+ _executor_manager: &ExecutorManager,
+ ) -> Result<HashSet<JobId>> {
+ 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<LogicalPlanNode, PhysicalPlanNode> =
+ 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 2c7d9beee..7d67b8d8b 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -61,12 +61,16 @@ type ActiveJobCache = Arc<DashMap<JobId, JobInfoCache>>;
#[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<MultiTaskDefinition>,
executor_manager: &ExecutorManager,
- ) -> Result<()>;
+ ) -> Result<HashSet<JobId>>;
}
struct DefaultTaskLauncher {
@@ -86,7 +90,7 @@ impl TaskLauncher for DefaultTaskLauncher {
executor: &ExecutorMetadata,
tasks: Vec<MultiTaskDefinition>,
executor_manager: &ExecutorManager,
- ) -> Result<()> {
+ ) -> Result<HashSet<JobId>> {
if log::max_level() >= log::Level::Info {
let tasks_ids: Vec<String> = 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<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskManager<T, U>
executor: &ExecutorMetadata,
tasks: Vec<Vec<TaskDescription>>,
executor_manager: &ExecutorManager,
- ) -> Result<()> {
+ ) -> Result<HashSet<JobId>> {
let mut multi_tasks = vec![];
for stage_tasks in tasks {
match self.prepare_multi_task_definition(stage_tasks) {
@@ -834,7 +838,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskManager<T, U>
.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 c2de2b300..50b810e73 100644
--- a/ballista/scheduler/src/test_utils.rs
+++ b/ballista/scheduler/src/test_utils.rs
@@ -342,8 +342,8 @@ impl TaskLauncher for BlackholeTaskLauncher {
_executor: &ExecutorMetadata,
_tasks: Vec<MultiTaskDefinition>,
_executor_manager: &ExecutorManager,
- ) -> Result<()> {
- Ok(())
+ ) -> Result<HashSet<JobId>> {
+ Ok(HashSet::new())
}
}
@@ -364,7 +364,7 @@ impl TaskLauncher for VirtualTaskLauncher {
executor: &ExecutorMetadata,
tasks: Vec<MultiTaskDefinition>,
_executor_manager: &ExecutorManager,
- ) -> Result<()> {
+ ) -> Result<HashSet<JobId>> {
if self.unreachable.lock().contains(&executor.id) {
return Err(BallistaError::Internal(format!(
"test: executor {} is unreachable",
@@ -389,7 +389,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<MultiTaskDefinition>,
+ _executor_manager: &ExecutorManager,
+ ) -> Result<HashSet<JobId>> {
+ Ok(tasks
+ .iter()
+ .map(|t| JobId::from(t.job_id.clone()))
+ .collect())
}
}
@@ -490,6 +513,84 @@ 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(),
+ 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<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 { 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()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
