From fc6e54d6d2efd39e527f859b38f8c7699a683a42 Mon Sep 17 00:00:00 2001 From: Jeadie Date: Tue, 23 Jun 2026 13:58:54 +1000 Subject: [PATCH 1/3] feat: add poll_now_notify to poll_loop and on_work_available callback Adds an optional poll_now_notify: Option> parameter to the executor poll_loop so the scheduler can wake an idle executor immediately instead of waiting for the next poll interval. The idle wait becomes a tokio::select! over the poll interval and the notify. Adds an OnWorkAvailableFn callback to SchedulerConfig, invoked from the query stage scheduler when new work becomes available (after JobSubmitted and when new stages become runnable). Ported from spiceai/datafusion-ballista#12 for upstreaming. --- ballista/executor/src/execution_loop.rs | 21 +++++++++++++++++-- ballista/executor/src/executor_process.rs | 1 + ballista/executor/src/standalone.rs | 4 +++- ballista/scheduler/src/config.rs | 15 +++++++++++++ .../scheduler_server/query_stage_scheduler.rs | 12 +++++++++++ 5 files changed, 50 insertions(+), 3 deletions(-) diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index f64fae3014..9d5d82a8a9 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -46,7 +46,7 @@ use std::error::Error; use std::sync::mpsc::{Receiver, Sender, TryRecvError}; use std::time::{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}; /// Main execution loop that polls the scheduler for available tasks. @@ -57,10 +57,14 @@ 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, allows the scheduler to wake the poll loop +/// immediately instead of waiting for the next idle poll interval. pub async fn poll_loop( mut scheduler: SchedulerGrpcClient, executor: Arc, codec: BallistaCodec, + poll_now_notify: Option>, ) -> Result<(), BallistaError> where C: tonic::client::GrpcService, @@ -204,7 +208,20 @@ where } if !active_job { - tokio::time::sleep(Duration::from_millis(50)).await; + // Wait for either the poll interval or a poll_now notification + match &poll_now_notify { + Some(notify) => { + tokio::select! { + () = tokio::time::sleep(Duration::from_millis(50)) => {} + () = notify.notified() => { + debug!("Received poll_now notification, polling immediately"); + } + } + } + None => { + tokio::time::sleep(Duration::from_millis(50)).await; + } + } } } } diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 6050649ab9..43e2a448c8 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -505,6 +505,7 @@ pub async fn start_executor_process( scheduler.clone(), executor.clone(), default_codec, + None, // poll_now_notify ))); } }; diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index bd72dee4b4..f21e5c765f 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -142,7 +142,9 @@ pub async fn new_standalone_executor_from_builder( )), ); - tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec)); + tokio::spawn(execution_loop::poll_loop( + scheduler, executor, codec, None, + )); Ok(()) } diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 3c6cdb0074..8bb85b93fd 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -34,6 +34,16 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use std::fmt::Display; use std::sync::Arc; +/// Callback invoked when new work becomes available for executors. +/// +/// This is called after: +/// - A job is submitted and tasks are ready to be scheduled +/// - Tasks complete and new stages become runnable +/// +/// This allows external systems to notify executors to poll immediately +/// rather than waiting for their next poll interval. +pub type OnWorkAvailableFn = Arc; + /// Command-line configuration for the scheduler binary. #[cfg(feature = "build-binary")] #[derive(clap::Parser, Debug)] @@ -278,6 +288,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. + /// The string argument is a reason/description for debugging purposes. + pub on_work_available: Option, } impl Default for SchedulerConfig { @@ -314,6 +327,7 @@ impl Default for SchedulerConfig { cors_allowed_origins: String::default(), #[cfg(feature = "rest-api")] cors_allowed_methods: String::default(), + on_work_available: None, } } } @@ -546,6 +560,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 0d84950bc7..427b7b6164 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -180,6 +180,11 @@ impl .post_event(QueryStageSchedulerEvent::ReviveOffers) .await?; } + + // Notify external systems that new work is available + if let Some(ref callback) = self.config.on_work_available { + callback(&format!("job_submitted:{job_id}")); + } } QueryStageSchedulerEvent::JobPlanningFailed { job_id, @@ -295,6 +300,13 @@ impl .await?; } + // Notify external systems when new stages become runnable + if !stage_events.is_empty() + && let Some(ref callback) = self.config.on_work_available + { + callback("tasks_completed:new_stages_runnable"); + } + for stage_event in stage_events { event_sender.post_event(stage_event).await?; } From 683dfa77318e56751a8290fd3aa0aab254a04788 Mon Sep 17 00:00:00 2001 From: jeadie Date: Tue, 23 Jun 2026 15:21:14 +1000 Subject: [PATCH 2/3] fix(examples): pass new poll_loop arg in mtls-cluster and reformat standalone --- ballista/executor/src/standalone.rs | 4 +--- examples/examples/mtls-cluster.rs | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index f21e5c765f..a77ba1b868 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -142,9 +142,7 @@ pub async fn new_standalone_executor_from_builder( )), ); - tokio::spawn(execution_loop::poll_loop( - scheduler, executor, codec, None, - )); + tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec, None)); Ok(()) } diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs index 593248183e..128e2b64a7 100644 --- a/examples/examples/mtls-cluster.rs +++ b/examples/examples/mtls-cluster.rs @@ -430,7 +430,7 @@ async fn run_executor() -> Result<(), Box> { // This registers the executor and starts polling for tasks info!("Starting execution poll loop..."); let poll_handle = tokio::spawn(async move { - execution_loop::poll_loop(scheduler, executor, codec).await + execution_loop::poll_loop(scheduler, executor, codec, None).await }); tokio::select! { From 413bcbfc6e6c48997f3b17ebcd34cbf0e566bded Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:10:03 +0900 Subject: [PATCH 3/3] fix: address review feedback - Replace the callback's `&str` argument with a `WorkAvailableReason` enum - Fire the stage-resolution callback from the JobUpdated arm after `update_job` has revived the graph, so it only fires for genuinely new runnable work and never before the work is pollable - Add `SchedulerConfig::with_on_work_available` - Document that the callback runs synchronously in the scheduler event loop and must not block - Idle-poll less aggressively (1s fallback) when a poll_now notifier is wired --- ballista/executor/src/execution_loop.rs | 19 +- ballista/scheduler/src/config.rs | 49 ++++- .../scheduler_server/query_stage_scheduler.rs | 183 ++++++++++++++++-- 3 files changed, 221 insertions(+), 30 deletions(-) diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index 71ac742226..874aa8632e 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -49,6 +49,13 @@ use std::{sync::Arc, time::Duration}; 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 @@ -58,8 +65,11 @@ 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, allows the scheduler to wake the poll loop -/// immediately instead of waiting for the next idle poll interval. +/// `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, @@ -211,18 +221,17 @@ where } if !active_job { - // Wait for either the poll interval or a poll_now notification match &poll_now_notify { Some(notify) => { tokio::select! { - () = tokio::time::sleep(Duration::from_millis(50)) => {} + () = tokio::time::sleep(NOTIFIED_IDLE_POLL_INTERVAL) => {} () = notify.notified() => { debug!("Received poll_now notification, polling immediately"); } } } None => { - tokio::time::sleep(Duration::from_millis(50)).await; + tokio::time::sleep(IDLE_POLL_INTERVAL).await; } } } diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index f634632cf0..74f5185f69 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -28,21 +28,42 @@ 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; -/// Callback invoked when new work becomes available for executors. +/// 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. /// -/// This is called after: -/// - A job is submitted and tasks are ready to be scheduled -/// - Tasks complete and new stages become runnable +/// # Warning /// -/// This allows external systems to notify executors to poll immediately -/// rather than waiting for their next poll interval. -pub type OnWorkAvailableFn = Arc; +/// 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")] @@ -298,7 +319,7 @@ pub struct SchedulerConfig { /// Comma-separated list of allowed methods for CORS pub cors_allowed_methods: String, /// Callback invoked when new work becomes available for executors. - /// The string argument is a reason/description for debugging purposes. + /// See [`OnWorkAvailableFn`]. pub on_work_available: Option, } @@ -473,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 diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 13fe0e69a5..4fdfca736d 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; @@ -198,9 +198,10 @@ impl .await?; } - // Notify external systems that new work is available - if let Some(ref callback) = self.config.on_work_available { - callback(&format!("job_submitted:{job_id}")); + // 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 { @@ -263,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) => { @@ -310,13 +322,6 @@ impl .await?; } - // Notify external systems when new stages become runnable - if !stage_events.is_empty() - && let Some(ref callback) = self.config.on_work_available - { - callback("tasks_completed:new_stages_runnable"); - } - for stage_event in stage_events { event_sender.post_event(stage_event).await?; } @@ -388,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<()> { @@ -453,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),