Add metrics collector to SchedulerState and instrument executor and planning metrics

phillipleblanc · phillipleblanc · commit 0a4174a63001 · 2026-01-21T20:00:34.000-08:00
- Add metrics_collector field to SchedulerState struct
- Instrument record_planning_duration in submit_job
- Instrument record_executor_registered/deregistered and set_active_executor_count
- Update all SchedulerState constructors and call sites
diff --git a/ballista/scheduler/src/scheduler_server/grpc.rs b/ballista/scheduler/src/scheduler_server/grpc.rs
@@ -48,10 +48,10 @@ use std::ops::Deref;
 
 use crate::cluster::{bind_task_bias, bind_task_round_robin};
 use crate::config::TaskDistributionPolicy;
-use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use crate::scheduler_server::SchedulerServer;
-use ballista_core::remote_catalog::remote_function_serialize_ext::RemoteFunctionSerializeExt;
+use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use ballista_core::remote_catalog::catalog_serialize_ext::CatalogSerializeExt;
+use ballista_core::remote_catalog::remote_function_serialize_ext::RemoteFunctionSerializeExt;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tonic::{Request, Response, Status};
 
@@ -490,6 +490,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
         );
 
         let executor_manager = self.state.executor_manager.clone();
+        let metrics_collector = self.state.metrics_collector.clone();
         let event_sender = self.query_stage_event_loop.get_sender().map_err(|e| {
             let msg = format!("Get query stage event loop error due to {e:?}");
             error!("{msg}");
@@ -499,6 +500,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
         Self::remove_executor(
             executor_manager,
             event_sender,
+            metrics_collector,
             &executor_id,
             Some(reason),
             self.config.executor_termination_grace_period,
@@ -665,6 +667,7 @@ mod test {
             SchedulerState::new_with_default_scheduler_name(
                 cluster.clone(),
                 BallistaCodec::default(),
+                default_metrics_collector().unwrap(),
             );
         state.init().await?;
 
@@ -697,6 +700,7 @@ mod test {
             SchedulerState::new_with_default_scheduler_name(
                 cluster.clone(),
                 BallistaCodec::default(),
+                default_metrics_collector().unwrap(),
             );
         state.init().await?;
 
diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -99,6 +99,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             codec,
             scheduler_name.clone(),
             config.clone(),
+            metrics_collector.clone(),
         ));
         let query_stage_scheduler = Arc::new(QueryStageScheduler::new(
             state.clone(),
@@ -137,6 +138,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             codec,
             scheduler_name.clone(),
             config.clone(),
+            metrics_collector.clone(),
             task_launcher,
         ));
         let query_stage_scheduler = Arc::new(QueryStageScheduler::new(
@@ -312,6 +314,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
                     Self::remove_executor(
                         state.executor_manager.clone(),
                         sender_clone,
+                        state.metrics_collector.clone(),
                         &executor_id,
                         Some(stop_reason.clone()),
                         0,
@@ -338,6 +341,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     pub(crate) fn remove_executor(
         executor_manager: ExecutorManager,
         event_sender: EventSender<QueryStageSchedulerEvent>,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
         executor_id: &str,
         reason: Option<String>,
         wait_secs: u64,
@@ -356,6 +360,13 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
                 error!("error removing executor {executor_id}: {e:?}");
             }
 
+            // Record executor deregistration metric
+            metrics_collector.record_executor_deregistered(&executor_id);
+
+            // Update active executor count
+            let count = executor_manager.get_alive_executors().len();
+            metrics_collector.set_active_executor_count(count);
+
             if let Err(e) = event_sender
                 .post_event(QueryStageSchedulerEvent::ExecutorLost(executor_id, reason))
                 .await
@@ -366,8 +377,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     }
 
     async fn do_register_executor(&self, metadata: ExecutorMetadata) -> Result<()> {
+        let executor_id = metadata.id.clone();
         let executor_data = ExecutorData {
-            executor_id: metadata.id.clone(),
+            executor_id: executor_id.clone(),
             total_task_slots: metadata.specification.task_slots,
             available_task_slots: metadata.specification.task_slots,
         };
@@ -378,6 +390,17 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             .register_executor(metadata, executor_data)
             .await?;
 
+        // Record executor registration metric
+        self.state
+            .metrics_collector
+            .record_executor_registered(&executor_id);
+
+        // Update active executor count
+        let count = self.state.executor_manager.get_alive_executors().len();
+        self.state
+            .metrics_collector
+            .set_active_executor_count(count);
+
         // If we are using push-based scheduling then reserve this executors slots and send
         // them for scheduling tasks.
         if self.state.config.is_push_staged_scheduling() {
diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs
@@ -36,10 +36,10 @@ use ballista_core::serde::protobuf::{
 use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata};
 
 use ballista_core::utils::{
-    GrpcClientConfig, create_grpc_client_connection, create_grpc_client_endpoint, get_time_before,
+    GrpcClientConfig, create_grpc_client_connection, create_grpc_client_endpoint,
+    get_time_before,
 };
 
-
 use dashmap::DashMap;
 use log::{debug, error, info, warn};
 use std::collections::{HashMap, HashSet};
diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs
@@ -37,6 +37,7 @@ use crate::state::task_manager::{TaskLauncher, TaskManager};
 
 use crate::cluster::{BallistaCluster, BoundTask, ExecutorSlot};
 use crate::config::SchedulerConfig;
+use crate::metrics::SchedulerMetricsCollector;
 use crate::state::execution_graph::TaskDescription;
 use ballista_core::error::{BallistaError, Result};
 use ballista_core::event_loop::EventSender;
@@ -117,6 +118,8 @@ pub struct SchedulerState<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPl
     pub codec: BallistaCodec<T, U>,
     /// Scheduler configuration.
     pub config: Arc<SchedulerConfig>,
+    /// Metrics collector for recording scheduler metrics.
+    pub metrics_collector: Arc<dyn SchedulerMetricsCollector>,
 }
 
 impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T, U> {
@@ -126,6 +129,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         codec: BallistaCodec<T, U>,
         scheduler_name: String,
         config: Arc<SchedulerConfig>,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
     ) -> Self {
         Self {
             executor_manager: ExecutorManager::new(
@@ -140,6 +144,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
             session_manager: SessionManager::new(cluster.job_state()),
             codec,
             config,
+            metrics_collector,
         }
     }
 
@@ -148,9 +153,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
     pub fn new_with_default_scheduler_name(
         cluster: BallistaCluster,
         codec: BallistaCodec<T, U>,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
     ) -> Self {
         let config = Arc::new(SchedulerConfig::default());
-        SchedulerState::new(cluster, codec, "localhost:50050".to_owned(), config)
+        SchedulerState::new(
+            cluster,
+            codec,
+            "localhost:50050".to_owned(),
+            config,
+            metrics_collector,
+        )
     }
 
     #[allow(dead_code)]
@@ -159,6 +171,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         codec: BallistaCodec<T, U>,
         scheduler_name: String,
         config: Arc<SchedulerConfig>,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
         dispatcher: Arc<dyn TaskLauncher>,
     ) -> Self {
         Self {
@@ -175,6 +188,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
             session_manager: SessionManager::new(cluster.job_state()),
             codec,
             config,
+            metrics_collector,
         }
     }
 
@@ -491,6 +505,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
 
         let elapsed = start.elapsed();
 
+        // Record planning duration metric
+        self.metrics_collector
+            .record_planning_duration(job_id, elapsed.as_millis() as u64);
+
         info!("Planned job {job_id} in {elapsed:?}");
 
         Ok(())
