diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index da4ff19cd..874aa8632 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -46,9 +46,16 @@ use std::error::Error; use std::sync::mpsc::{Receiver, Sender, TryRecvError}; use std::time::{Instant, SystemTime, UNIX_EPOCH}; use std::{sync::Arc, time::Duration}; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; use tonic::codegen::{Body, Bytes, StdError}; +/// Idle sleep between polls when polling is the only way to learn of new work. +const IDLE_POLL_INTERVAL: Duration = Duration::from_millis(50); + +/// Idle sleep when a `poll_now_notify` wake-up is wired and the timer is only +/// a fallback. +const NOTIFIED_IDLE_POLL_INTERVAL: Duration = Duration::from_secs(1); + /// Main execution loop that polls the scheduler for available tasks. /// /// This function runs indefinitely, periodically asking the scheduler for @@ -57,10 +64,17 @@ use tonic::codegen::{Body, Bytes, StdError}; /// /// The loop respects the executor's concurrent task limit via a semaphore, /// ensuring no more than the configured number of tasks run simultaneously. +/// +/// `poll_now_notify`, when provided, wakes an idle poll loop immediately +/// (typically wired to the scheduler's `on_work_available` callback) instead +/// of waiting out the idle interval. A notification sent mid-poll is not +/// lost: `Notify` stores the permit and the next `notified().await` returns +/// immediately. pub async fn poll_loop( mut scheduler: SchedulerGrpcClient, executor: Arc, codec: BallistaCodec, + poll_now_notify: Option>, health: crate::health::ExecutorHealth, ) -> Result<(), BallistaError> where @@ -207,7 +221,19 @@ where } if !active_job { - tokio::time::sleep(Duration::from_millis(50)).await; + match &poll_now_notify { + Some(notify) => { + tokio::select! { + () = tokio::time::sleep(NOTIFIED_IDLE_POLL_INTERVAL) => {} + () = notify.notified() => { + debug!("Received poll_now notification, polling immediately"); + } + } + } + None => { + tokio::time::sleep(IDLE_POLL_INTERVAL).await; + } + } } } } diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 9bce0d549..c8c5b0d94 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -564,6 +564,7 @@ pub async fn start_executor_process( scheduler.clone(), executor.clone(), default_codec, + None, // poll_now_notify health, ))); } diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index d39841e8b..fd9841a0e 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -147,6 +147,7 @@ pub async fn new_standalone_executor_from_builder( scheduler, executor, codec, + None, crate::health::ExecutorHealth::new(), )); Ok(()) diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 940888714..74f5185f6 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -28,12 +28,43 @@ use crate::SessionBuilder; use crate::cluster::DistributionPolicy; use ballista_core::extension::EndpointOverrideFn; -use ballista_core::{ConfigProducer, config::TaskSchedulingPolicy}; +use ballista_core::{ConfigProducer, JobId, config::TaskSchedulingPolicy}; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use std::fmt::Display; use std::sync::Arc; +/// Why the scheduler believes new work has become available for executors. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkAvailableReason { + /// A job was submitted and its initial tasks are ready to be scheduled. + JobSubmitted { + /// Identifier of the submitted job. + job_id: JobId, + }, + /// Completed tasks resolved downstream stages of a job, and the tasks of + /// those stages are now schedulable. + NewStagesRunnable { + /// Identifier of the job that gained schedulable tasks. + job_id: JobId, + }, +} + +/// Callback invoked when new work becomes available for executors, e.g. to +/// wake idle pull-based executors via the poll loop's `poll_now_notify`. +/// +/// It fires only after the work is visible to a polling executor, so waking +/// one cannot race the scheduler's internal event processing. +/// +/// # Warning +/// +/// The callback runs synchronously inside the scheduler's main event loop. +/// Implementations **must be non-blocking**; offload blocking or long-running +/// work (such as network I/O) to a separate task or thread. +/// +/// `Arc` rather than `Box` so [`SchedulerConfig`] remains [`Clone`]. +pub type OnWorkAvailableFn = Arc; + /// Command-line configuration for the scheduler binary. #[cfg(feature = "build-binary")] #[derive(clap::Parser, Debug)] @@ -287,6 +318,9 @@ pub struct SchedulerConfig { #[cfg(feature = "rest-api")] /// Comma-separated list of allowed methods for CORS pub cors_allowed_methods: String, + /// Callback invoked when new work becomes available for executors. + /// See [`OnWorkAvailableFn`]. + pub on_work_available: Option, } impl Default for SchedulerConfig { @@ -324,6 +358,7 @@ impl Default for SchedulerConfig { cors_allowed_origins: String::default(), #[cfg(feature = "rest-api")] cors_allowed_methods: String::default(), + on_work_available: None, } } } @@ -459,6 +494,16 @@ impl SchedulerConfig { self.use_tls = use_tls; self } + + /// Sets the callback invoked when new work becomes available for + /// executors. See [`OnWorkAvailableFn`]. + pub fn with_on_work_available( + mut self, + on_work_available: OnWorkAvailableFn, + ) -> Self { + self.on_work_available = Some(on_work_available); + self + } } /// Policy of distributing tasks to available executor slots @@ -557,6 +602,7 @@ impl TryFrom for SchedulerConfig { cors_allowed_origins: opt.cors_allowed_origins, #[cfg(feature = "rest-api")] cors_allowed_methods: opt.cors_allowed_methods, + on_work_available: None, }; Ok(config) diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index ee5174853..4fdfca736 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -26,7 +26,7 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::{EventAction, EventSender}; use tokio::sync::mpsc::error::TrySendError; -use crate::config::SchedulerConfig; +use crate::config::{SchedulerConfig, WorkAvailableReason}; use crate::metrics::SchedulerMetricsCollector; use crate::scheduler_server::timestamp_millis; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -197,6 +197,12 @@ impl .post_event(QueryStageSchedulerEvent::ReviveOffers) .await?; } + + // The graph was revived before caching, so the job's tasks + // are already visible to polling executors. + if let Some(callback) = &self.config.on_work_available { + callback(WorkAvailableReason::JobSubmitted { job_id }); + } } QueryStageSchedulerEvent::JobPlanningFailed { job_id, @@ -258,8 +264,19 @@ impl } QueryStageSchedulerEvent::JobUpdated(job_id) => { debug!("Job updated, job_id: [{job_id}]"); - if let Err(e) = self.state.task_manager.update_job(&job_id).await { - error!("Fail to invoke update_job for job {job_id} due to {e:?}"); + match self.state.task_manager.update_job(&job_id).await { + Ok(new_tasks) => { + // update_job revived the graph: the new tasks are + // already visible to polling executors. + if new_tasks > 0 + && let Some(callback) = &self.config.on_work_available + { + callback(WorkAvailableReason::NewStagesRunnable { job_id }); + } + } + Err(e) => { + error!("Fail to invoke update_job for job {job_id} due to {e:?}"); + } } } QueryStageSchedulerEvent::JobCancel(job_id) => { @@ -376,16 +393,33 @@ impl #[cfg(test)] mod tests { - use crate::config::SchedulerConfig; - use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; + use crate::config::{SchedulerConfig, WorkAvailableReason}; + use crate::scheduler_server::SchedulerServer; + use crate::test_utils::{ + SchedulerTest, TestMetricsCollector, await_condition, test_cluster_context, + }; + use ballista_core::BALLISTA_PROTOCOL_VERSION; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; + use ballista_core::extension::SessionConfigExt; + use ballista_core::serde::BallistaCodec; + use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpc; + use ballista_core::serde::protobuf::{ + ExecutorRegistration, PollWorkParams, ShuffleWritePartition, SuccessfulTask, + TaskStatus, task_status, + }; + use ballista_core::serde::scheduler::{ + ExecutorOperatingSystemSpecification, ExecutorSpecification, + }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; use datafusion::logical_expr::{LogicalPlan, col}; + use datafusion::prelude::SessionConfig; use datafusion::test_util::scan_empty_with_partitions; - use std::sync::Arc; + use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; + use std::sync::{Arc, Mutex}; use std::time::Duration; + use tonic::Request; #[tokio::test] async fn test_pending_job_metric() -> Result<()> { @@ -441,6 +475,135 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_on_work_available_callback() -> Result<()> { + let reasons: Arc>> = Arc::default(); + let captured = reasons.clone(); + + let mut scheduler: SchedulerServer = + SchedulerServer::new( + "localhost:50050".to_owned(), + test_cluster_context(), + BallistaCodec::default(), + Arc::new( + SchedulerConfig { + scheduling_policy: TaskSchedulingPolicy::PullStaged, + ..Default::default() + } + .with_on_work_available(Arc::new( + move |reason| { + captured.lock().unwrap().push(reason); + }, + )), + ), + Arc::new(TestMetricsCollector::default()), + ); + scheduler.init().await?; + + let ctx = scheduler + .state + .session_manager + .create_or_update_session( + "session", + &SessionConfig::new_with_ballista().with_target_partitions(2), + ) + .await?; + let job_id = scheduler.submit_job("", ctx, &test_plan(2), None).await?; + + // Job submission runs asynchronously through the event loop. + let submitted = await_condition(Duration::from_millis(10), 100, || { + futures::future::ready(Ok(!reasons.lock().unwrap().is_empty())) + }) + .await?; + assert!(submitted, "JobSubmitted callback never fired"); + assert_eq!( + reasons.lock().unwrap().first(), + Some(&WorkAvailableReason::JobSubmitted { + job_id: job_id.clone() + }) + ); + + // Pull the shuffle stage's tasks; the callback promised they are + // visible by the time it fired. + let exec_meta = ExecutorRegistration { + id: "executor-1".to_owned(), + host: Some("localhost".to_owned()), + port: 50051, + grpc_port: 50052, + specification: Some(ExecutorSpecification::default().with_vcores(2).into()), + os_info: Some(ExecutorOperatingSystemSpecification::default().into()), + ballista_protocol_version: BALLISTA_PROTOCOL_VERSION, + }; + let polled = scheduler + .poll_work(Request::new(PollWorkParams { + metadata: Some(exec_meta.clone()), + num_free_vcores: 2, + task_status: vec![], + })) + .await + .expect("poll_work failed") + .into_inner(); + assert!( + !polled.tasks.is_empty(), + "expected tasks after JobSubmitted" + ); + + // Report the pulled tasks as successful; each writes the plan's two + // shuffle output partitions. + let task_status = polled + .tasks + .iter() + .map(|task| TaskStatus { + task_id: task.task_id, + job_id: task.job_id.clone(), + stage_id: task.stage_id, + stage_attempt_num: task.stage_attempt_num, + launch_time: 0, + start_exec_time: 0, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask { + executor_id: exec_meta.id.clone(), + partitions: (0..2) + .map(|partition_id| ShuffleWritePartition { + partition_id, + num_batches: 1, + num_rows: 1, + num_bytes: 1, + file_id: None, + is_sort_shuffle: false, + }) + .collect(), + })), + }) + .collect(); + scheduler + .poll_work(Request::new(PollWorkParams { + metadata: Some(exec_meta), + num_free_vcores: 2, + task_status, + })) + .await + .expect("poll_work with task status failed"); + + // Completing the shuffle stage resolves the final stage. + let resolved = await_condition(Duration::from_millis(10), 100, || { + futures::future::ready(Ok(reasons.lock().unwrap().contains( + &WorkAvailableReason::NewStagesRunnable { + job_id: job_id.clone(), + }, + ))) + }) + .await?; + assert!( + resolved, + "expected NewStagesRunnable, got {:?}", + reasons.lock().unwrap() + ); + + Ok(()) + } + fn test_plan(partitions: usize) -> LogicalPlan { let schema = Schema::new(vec![ Field::new("id", DataType::Utf8, false), diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs index 3533019b9..609935d6b 100644 --- a/examples/examples/mtls-cluster.rs +++ b/examples/examples/mtls-cluster.rs @@ -433,7 +433,7 @@ async fn run_executor() -> Result<(), Box> { info!("Starting execution poll loop..."); let health = ballista_executor::health::ExecutorHealth::new(); let poll_handle = tokio::spawn(async move { - execution_loop::poll_loop(scheduler, executor, codec, health).await + execution_loop::poll_loop(scheduler, executor, codec, None, health).await }); tokio::select! {