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 10d4d0877 feat: add `poll_now_notify` to `poll_loop` and
`on_work_available` callback (#1893)
10d4d0877 is described below
commit 10d4d0877c0bc705d9896e12bac26d88147fb83f
Author: Jack Eadie <[email protected]>
AuthorDate: Wed Aug 5 16:17:21 2026 +1000
feat: add `poll_now_notify` to `poll_loop` and `on_work_available` callback
(#1893)
---
ballista/executor/src/execution_loop.rs | 30 +++-
ballista/executor/src/executor_process.rs | 1 +
ballista/executor/src/standalone.rs | 1 +
ballista/scheduler/src/config.rs | 48 +++++-
.../src/scheduler_server/query_stage_scheduler.rs | 176 ++++++++++++++++++++-
5 files changed, 247 insertions(+), 9 deletions(-)
diff --git a/ballista/executor/src/execution_loop.rs
b/ballista/executor/src/execution_loop.rs
index b84c83b9d..2390364c1 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);
+
/// Maximum time the poll loop waits for a free vcore before polling the
/// scheduler anyway. `poll_work` doubles as the executor's heartbeat under
/// pull-based scheduling, so a fully-busy executor must keep polling
(reporting
@@ -68,6 +75,12 @@ const HEARTBEAT_POLL_INTERVAL: Duration =
Duration::from_secs(5);
/// Pass `None` to have the loop create a semaphore sized to the executor's
/// configured vcore count.
///
+/// `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.
+///
/// **Shared semaphores**: when one semaphore is shared across loops that
/// connect to different schedulers, each scheduler independently sees the
/// current free capacity and may dispatch up to that many tasks. The semaphore
@@ -89,6 +102,7 @@ pub async fn poll_loop<T: 'static + AsLogicalPlan, U:
'static + AsExecutionPlan,
mut scheduler: SchedulerGrpcClient<C>,
executor: Arc<Executor>,
codec: BallistaCodec<T, U>,
+ poll_now_notify: Option<Arc<Notify>>,
free_vcores: Option<Arc<Semaphore>>,
health: crate::health::ExecutorHealth,
) -> Result<(), BallistaError>
@@ -263,7 +277,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 091a94c5f..32c80afc9 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -728,6 +728,7 @@ pub async fn start_executor_process(
scheduler.clone(),
executor.clone(),
default_codec,
+ None, // poll_now_notify
None, // free_vcores: use internal semaphore
health,
)));
diff --git a/ballista/executor/src/standalone.rs
b/ballista/executor/src/standalone.rs
index fd9841a0e..78e704295 100644
--- a/ballista/executor/src/standalone.rs
+++ b/ballista/executor/src/standalone.rs
@@ -148,6 +148,7 @@ pub async fn new_standalone_executor_from_builder(
executor,
codec,
None,
+ None,
crate::health::ExecutorHealth::new(),
));
Ok(())
diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs
index ff935374b..28320783f 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<dyn Fn(WorkAvailableReason) + Send + Sync>;
+
/// Command-line configuration for the scheduler binary.
#[cfg(feature = "build-binary")]
#[derive(clap::Parser, Debug)]
@@ -313,6 +344,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<OnWorkAvailableFn>,
}
impl Default for SchedulerConfig {
@@ -352,6 +386,7 @@ impl Default for SchedulerConfig {
cors_allowed_origins: String::default(),
#[cfg(feature = "rest-api")]
cors_allowed_methods: String::default(),
+ on_work_available: None,
}
}
}
@@ -500,6 +535,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
@@ -600,6 +645,7 @@ impl TryFrom<Config> 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 4b64e205a..d4cb11fc2 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -27,7 +27,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;
@@ -204,6 +204,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
.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,
@@ -265,8 +271,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
}
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) => {
@@ -465,17 +482,34 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
#[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::job_status;
+ 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<()> {
@@ -531,6 +565,136 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_on_work_available_callback() -> Result<()> {
+ let reasons: Arc<Mutex<Vec<WorkAvailableReason>>> = Arc::default();
+ let captured = reasons.clone();
+
+ let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ 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(),
+ runtime_stats: vec![],
+ })),
+ })
+ .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(())
+ }
+
#[tokio::test]
async fn test_running_job_fails_when_all_executors_are_lost() ->
Result<()> {
let plan = test_plan(10);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
