Add remaining scheduler metrics instrumentation (stage_started, task_scheduled, queue sizes)

Author: phillipleblanc

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -140,6 +140,17 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
             let mut tasks = vec![];
             for (_, task) in schedulable_tasks {
                 let job_id = task.partition.job_id.clone();
+                let stage_id = task.partition.stage_id;
+
+                // Record task scheduling metric
+                // Note: latency_ms is 0 since we don't currently track when tasks became schedulable.
+                self.state.metrics_collector.record_task_scheduled(
+                    &job_id,
+                    stage_id,
+                    &executor_id,
+                    0,
+                );
+
                 match self.state.task_manager.prepare_task_definition(task) {
                     Ok(task_definition) => tasks.push(task_definition),
                     Err(e) => {

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -323,6 +323,18 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                 );
             }
         }
+
+        // Update queue size metrics after processing each event
+        let pending_jobs = self.state.task_manager.pending_job_number();
+        #[expect(clippy::cast_possible_truncation)]
+        self.metrics_collector
+            .set_pending_jobs_queue_size(pending_jobs as u64);
+
+        let pending_tasks = self.state.task_manager.total_pending_tasks().await;
+        #[expect(clippy::cast_possible_truncation)]
+        self.metrics_collector
+            .set_pending_tasks_queue_size(pending_tasks as u64);
+
         Ok(())
     }
 

ballista/scheduler/src/state/execution_graph.rs

@@ -305,6 +305,15 @@ impl ExecutionGraph {
     /// Revive the execution graph by converting the resolved stages to running stages
     /// If any stages are converted, return true; else false.
     pub fn revive(&mut self) -> bool {
+        self.revive_with_metrics().0
+    }
+
+    /// Revive the execution graph by converting the resolved stages to running stages.
+    ///
+    /// Returns a tuple of (stages_converted, stages_started_info) where:
+    /// - `stages_converted`: true if any stages were converted
+    /// - `stages_started_info`: list of (stage_id, task_count, started_at_ms) for each started stage
+    pub fn revive_with_metrics(&mut self) -> (bool, Vec<(usize, usize, u64)>) {
         let running_stages = self
             .stages
             .values()
@@ -318,15 +327,23 @@ impl ExecutionGraph {
             .collect::<Vec<_>>();
 
         if running_stages.is_empty() {
-            false
+            (false, vec![])
         } else {
+            let mut stages_started = Vec::with_capacity(running_stages.len());
             for running_stage in running_stages {
+                let stage_id = running_stage.stage_id;
+                let task_count = running_stage.partitions;
+                #[expect(clippy::cast_possible_truncation)]
+                let started_at_ms = running_stage.stage_running_time as u64;
+
+                stages_started.push((stage_id, task_count, started_at_ms));
+
                 self.stages.insert(
                     running_stage.stage_id,
                     ExecutionStage::Running(running_stage),
                 );
             }
-            true
+            (true, stages_started)
         }
     }
 
@@ -354,7 +371,15 @@ impl ExecutionGraph {
 
         // Revive before updating due to some updates not saved
         // It will be refined later
-        self.revive();
+        let (_, stages_started) = self.revive_with_metrics();
+        for (stage_id, task_count, started_at_ms) in stages_started {
+            metrics_info.stages_started.push((
+                job_id.clone(),
+                stage_id,
+                task_count,
+                started_at_ms,
+            ));
+        }
 
         let current_running_stages: HashSet<usize> =
             HashSet::from_iter(self.running_stages());
@@ -775,17 +800,16 @@ impl ExecutionGraph {
             }
         }
 
-        let (events, mut stage_metrics) =
-            self.processing_stages_update(UpdatedStages {
-                resolved_stages,
-                successful_stages,
-                failed_stages,
-                rollback_running_stages,
-                resubmit_successful_stages: resubmit_successful_stages
-                    .keys()
-                    .cloned()
-                    .collect(),
-            })?;
+        let (events, stage_metrics) = self.processing_stages_update(UpdatedStages {
+            resolved_stages,
+            successful_stages,
+            failed_stages,
+            rollback_running_stages,
+            resubmit_successful_stages: resubmit_successful_stages
+                .keys()
+                .cloned()
+                .collect(),
+        })?;
 
         // Combine task metrics collected during processing with stage metrics
         metrics_info

ballista/scheduler/src/state/mod.rs

@@ -284,6 +284,18 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         &self,
         bound_tasks: Vec<BoundTask>,
     ) -> Result<Vec<ExecutorSlot>> {
+        // Record task scheduling metrics for each task
+        for (executor_id, task) in &bound_tasks {
+            // Note: latency_ms is 0 since we don't currently track when tasks became schedulable.
+            // This could be enhanced by adding a schedulable_time field to TaskDescription.
+            self.metrics_collector.record_task_scheduled(
+                &task.partition.job_id,
+                task.partition.stage_id,
+                executor_id,
+                0, // latency_ms placeholder
+            );
+        }
+
         // Put tasks to the same executor together
         // And put tasks belonging to the same stage together for creating MultiTaskDefinition
         let mut executor_stage_assignments: HashMap<

ballista/scheduler/src/state/task_manager.rs

@@ -263,6 +263,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
         self.active_job_cache.len()
     }
 
+    /// Get the total number of pending tasks across all active jobs.
+    ///
+    /// A pending task is a task that is available to schedule on an executor
+    /// but cannot be scheduled because no resources are available.
+    pub async fn total_pending_tasks(&self) -> usize {
+        let mut total = 0;
+        for entry in self.active_job_cache.iter() {
+            let graph = entry.value().execution_graph.read().await;
+            total += graph.available_tasks();
+        }
+        total
+    }
+
     /// 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