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 f1bdc9f3b feat(scheduler): expose existing job state event stream
(#1891)
f1bdc9f3b is described below
commit f1bdc9f3b73adac1512fed10f64a0849b7fa61d0
Author: Jack Eadie <[email protected]>
AuthorDate: Wed Aug 5 20:04:46 2026 +1000
feat(scheduler): expose existing job state event stream (#1891)
* feat: add a broadcast channel for job state event notifications
Adds a tokio::sync::broadcast Sender to SchedulerServer, subscribable via
subscribe_job_updates(). On job state changes (queued/running/completed/
failed/cancelled) the scheduler broadcasts a JobStateEvent carrying the job
id and new state.
Ported from spiceai/datafusion-ballista#15 for upstreaming.
* fix(core): add From<&JobId> for String so borrowed ids satisfy
Into<String>
* fix(scheduler): address PR review feedback on job state broadcast channel
- Add job_state_channel_capacity to SchedulerConfig (default 256) so
operators with high job throughput can tune the buffer; removes the
hardcoded const from SchedulerServer
- Move all job state broadcasts to after their state transitions succeed,
so subscribers never see an event for a commit that failed; aligns
Queued/Failed/Completed/Cancelled with how Running already behaved
- Add From<job_status::Status> for JobState to keep the two enums in
sync; document that Cancelled has no protobuf counterpart
- Expand module doc with guidance on when to use JobStateEvent vs
JobStatusSubscriber, and reference job_state_channel_capacity for tuning
- Add integration tests covering the broadcast wiring: lifecycle order
(Queued -> Running -> Completed) and cancel path
* docs: link JobStateEvent by absolute path so rustdoc resolves it
* refactor(scheduler): expose existing job state event stream
* fix(scheduler): expose cluster state event stream
---------
Co-authored-by: Phillip LeBlanc
<[email protected]>
---
ballista/scheduler/src/cluster/memory.rs | 66 ++++++++++++++++------
ballista/scheduler/src/scheduler_server/mod.rs | 40 ++++++++++++-
.../src/scheduler_server/query_stage_scheduler.rs | 44 +++++++++++++++
ballista/scheduler/src/state/executor_manager.rs | 7 ++-
ballista/scheduler/src/state/task_manager.rs | 7 ++-
ballista/scheduler/src/test_utils.rs | 7 ++-
6 files changed, 151 insertions(+), 20 deletions(-)
diff --git a/ballista/scheduler/src/cluster/memory.rs
b/ballista/scheduler/src/cluster/memory.rs
index 51dfa37dc..3baef1ddd 100644
--- a/ballista/scheduler/src/cluster/memory.rs
+++ b/ballista/scheduler/src/cluster/memory.rs
@@ -493,22 +493,20 @@ impl JobState for InMemoryJobState {
async fn fail_unscheduled_job(&self, job_id: &JobId, reason: String) ->
Result<()> {
if let Some((job_id, (job_name, queued_at))) =
self.queued_jobs.remove(job_id) {
- self.completed_jobs.insert(
- job_id.clone(),
- (
- JobStatus {
- job_id: job_id.into(),
- job_name,
- status: Some(Status::Failed(FailedJob {
- error: reason,
- queued_at,
- started_at: 0,
- ended_at: timestamp_millis(),
- })),
- },
- None,
- ),
- );
+ let status = JobStatus {
+ job_id: job_id.clone().into(),
+ job_name,
+ status: Some(Status::Failed(FailedJob {
+ error: reason,
+ queued_at,
+ started_at: 0,
+ ended_at: timestamp_millis(),
+ })),
+ };
+ self.completed_jobs
+ .insert(job_id.clone(), (status.clone(), None));
+ self.job_event_sender
+ .send(&JobStateEvent::JobUpdated { job_id, status });
Ok(())
} else {
@@ -533,6 +531,7 @@ mod test {
use crate::test_utils::{
test_aggregation_plan, test_join_plan, test_two_aggregations_plan,
};
+ use ballista_core::JobId;
use ballista_core::error::Result;
use ballista_core::serde::protobuf::JobStatus;
use ballista_core::serde::scheduler::{
@@ -609,6 +608,41 @@ mod test {
Ok(())
}
+ #[tokio::test]
+ async fn test_in_memory_job_planning_failure_notification() -> Result<()> {
+ let state = InMemoryJobState::new(
+ "",
+ Arc::new(default_session_builder),
+ Arc::new(default_config_producer),
+ );
+ let mut events = state.job_state_events().await?;
+ let job_id = JobId::from("job-1");
+
+ state.accept_job(&job_id, "", 0)?;
+ state
+ .fail_unscheduled_job(&job_id, "failed planning".to_owned())
+ .await?;
+
+ let event =
+ tokio::time::timeout(std::time::Duration::from_secs(1),
events.next())
+ .await
+ .expect("job state event should arrive")
+ .expect("job state event stream should remain open");
+
+ assert!(matches!(
+ event,
+ JobStateEvent::JobUpdated {
+ job_id: event_job_id,
+ status: JobStatus {
+ status:
Some(ballista_core::serde::protobuf::job_status::Status::Failed(_)),
+ ..
+ },
+ } if event_job_id == job_id
+ ));
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_in_memory_job_notification() -> Result<()> {
let state = InMemoryJobState::new(
diff --git a/ballista/scheduler/src/scheduler_server/mod.rs
b/ballista/scheduler/src/scheduler_server/mod.rs
index 6fac8b0e4..ca439dfeb 100644
--- a/ballista/scheduler/src/scheduler_server/mod.rs
+++ b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -31,7 +31,7 @@ use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_proto::logical_plan::AsLogicalPlan;
use datafusion_proto::physical_plan::AsExecutionPlan;
-use crate::cluster::BallistaCluster;
+use crate::cluster::{BallistaCluster, ClusterStateEventStream,
JobStateEventStream};
use crate::config::SchedulerConfig;
use crate::metrics::SchedulerMetricsCollector;
use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata};
@@ -183,6 +183,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> SchedulerServer<T
self.state.task_manager.running_job_number()
}
+ /// Returns a stream of job state events from the configured state backend.
+ pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
+ self.state.task_manager.job_state_events().await
+ }
+
+ /// Returns a stream of cluster state events from the configured state
backend.
+ pub async fn cluster_state_events(&self) ->
Result<ClusterStateEventStream> {
+ self.state.executor_manager.cluster_state_events().await
+ }
+
/// True when at least `min_ready_executors` executors currently have
/// live heartbeats. Embedders can call this from their own
health/readiness
/// handler and AND it with whatever app-specific state they track. The
@@ -451,6 +461,7 @@ pub fn timestamp_millis() -> u64 {
#[cfg(test)]
mod test {
use std::sync::Arc;
+ use std::time::Duration;
use ballista_core::JobId;
use ballista_core::extension::SessionConfigExt;
@@ -463,7 +474,9 @@ mod test {
use datafusion::test_util::scan_empty_with_partitions;
use datafusion_proto::protobuf::LogicalPlanNode;
use datafusion_proto::protobuf::PhysicalPlanNode;
+ use futures::StreamExt;
+ use crate::cluster::ClusterStateEvent;
use crate::config::SchedulerConfig;
use ballista_core::config::TaskSchedulingPolicy;
use ballista_core::error::Result;
@@ -487,6 +500,31 @@ mod test {
assert_submitted_event, test_cluster_context,
};
+ #[tokio::test]
+ async fn test_scheduler_exposes_cluster_state_events() -> Result<()> {
+ let scheduler =
test_scheduler(TaskSchedulingPolicy::PushStaged).await?;
+ let mut events = scheduler.cluster_state_events().await?;
+ let (executor_metadata, executor_data) =
+ test_executors(2).into_iter().next().unwrap();
+ let executor_id = executor_metadata.id.clone();
+
+ scheduler
+ .state
+ .executor_manager
+ .register_executor(executor_metadata, executor_data)
+ .await?;
+
+ let event = tokio::time::timeout(Duration::from_secs(5), events.next())
+ .await
+ .expect("cluster state event should arrive");
+ assert_eq!(
+ event,
+ Some(ClusterStateEvent::RegisteredExecutor { executor_id })
+ );
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_pull_scheduling() -> Result<()> {
let plan = test_plan();
diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
index d4cb11fc2..ef3d118e8 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -482,6 +482,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
#[cfg(test)]
mod tests {
+ use crate::cluster::JobStateEvent;
use crate::config::{SchedulerConfig, WorkAvailableReason};
use crate::scheduler_server::SchedulerServer;
use crate::test_utils::{
@@ -507,6 +508,7 @@ mod tests {
use datafusion::prelude::SessionConfig;
use datafusion::test_util::scan_empty_with_partitions;
use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+ use futures::StreamExt;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tonic::Request;
@@ -565,6 +567,48 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_scheduler_exposes_job_state_events() -> Result<()> {
+ let plan = test_plan(1);
+ let metrics_collector = Arc::new(TestMetricsCollector::default());
+ let mut test = SchedulerTest::new(
+ SchedulerConfig::default()
+ .with_scheduler_policy(TaskSchedulingPolicy::PushStaged),
+ metrics_collector,
+ 1,
+ 1,
+ None,
+ )
+ .await?;
+
+ let mut events = test.job_state_events().await?;
+ let (_, job_id) = test.run("", &plan).await?;
+ let received_success = tokio::time::timeout(Duration::from_secs(5),
async {
+ while let Some(event) = events.next().await {
+ if matches!(
+ event,
+ JobStateEvent::JobUpdated {
+ job_id: event_job_id,
+ status,
+ } if event_job_id == job_id
+ && matches!(
+ status.status,
+ Some(job_status::Status::Successful(_))
+ )
+ ) {
+ return true;
+ }
+ }
+ false
+ })
+ .await
+ .expect("successful job state event should arrive");
+ assert!(
+ received_success,
+ "job state event stream closed unexpectedly"
+ );
+ Ok(())
+ }
#[tokio::test]
async fn test_on_work_available_callback() -> Result<()> {
let reasons: Arc<Mutex<Vec<WorkAvailableReason>>> = Arc::default();
diff --git a/ballista/scheduler/src/state/executor_manager.rs
b/ballista/scheduler/src/state/executor_manager.rs
index 2f0b162b0..1da756566 100644
--- a/ballista/scheduler/src/state/executor_manager.rs
+++ b/ballista/scheduler/src/state/executor_manager.rs
@@ -25,7 +25,7 @@ use ballista_core::serde::protobuf::ExecutorMetric;
use ballista_core::serde::protobuf::executor_metric::Metric;
use log::trace;
-use crate::cluster::{BoundTask, ClusterState, ExecutorSlot};
+use crate::cluster::{BoundTask, ClusterState, ClusterStateEventStream,
ExecutorSlot};
use crate::config::SchedulerConfig;
use crate::state::execution_graph::RunningTaskInfo;
@@ -112,6 +112,11 @@ impl ExecutorManager {
Ok(())
}
+ /// Returns a stream of cluster state events from the configured state
backend.
+ pub async fn cluster_state_events(&self) ->
Result<ClusterStateEventStream> {
+ self.cluster_state.cluster_state_events().await
+ }
+
/// Binds ready-to-run tasks from active jobs to available executor slots.
///
/// Returns a list of bound tasks that can be launched on executors.
diff --git a/ballista/scheduler/src/state/task_manager.rs
b/ballista/scheduler/src/state/task_manager.rs
index 7d46488c7..dea5af100 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::cluster::JobState;
+use crate::cluster::{JobState, JobStateEventStream};
use crate::config::SchedulerConfig;
use crate::planner::DefaultDistributedPlanner;
use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan};
@@ -243,6 +243,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskManager<T, U>
.count()
}
+ /// Returns a stream of job state events from the configured state backend.
+ pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
+ self.state.job_state_events().await
+ }
+
/// Generate an ExecutionGraph for the job and save it to the persistent
state.
/// By default, this job will be curated by the scheduler which receives
it.
/// Then we will also save it to the active execution graph
diff --git a/ballista/scheduler/src/test_utils.rs
b/ballista/scheduler/src/test_utils.rs
index c794a2a92..496aa8d58 100644
--- a/ballista/scheduler/src/test_utils.rs
+++ b/ballista/scheduler/src/test_utils.rs
@@ -53,7 +53,7 @@ use
datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::prelude::{CsvReadOptions, JoinType, col};
use datafusion::test_util::scan_empty_with_partitions;
-use crate::cluster::BallistaCluster;
+use crate::cluster::{BallistaCluster, JobStateEventStream};
use crate::scheduler_server::event::QueryStageSchedulerEvent;
use crate::state::execution_graph::{
@@ -479,6 +479,11 @@ impl SchedulerTest {
self.scheduler.running_job_number()
}
+ /// Returns job state events from the underlying scheduler.
+ pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
+ self.scheduler.job_state_events().await
+ }
+
/// Returns the session context for tests.
pub async fn ctx(&self) -> Result<Arc<SessionContext>> {
self.scheduler
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
