Phase 2: expose ExecutionGraph for task_history; persist executor_id on terminal TaskInfo

Author: phillipleblanc

ballista/scheduler/src/state/aqe/mod.rs

@@ -1267,6 +1267,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph {
                 let task_attempt = stage.task_failure_numbers[partition_id];
                 let task_info = crate::state::execution_graph::TaskInfo {
                     task_id,
+                    executor_id: executor_id.to_owned(),
                     scheduled_time: SystemTime::now()
                         .duration_since(UNIX_EPOCH)
                         .unwrap()

ballista/scheduler/src/state/execution_graph.rs

@@ -1658,6 +1658,7 @@ impl ExecutionGraph for StaticExecutionGraph {
                 let task_attempt = stage.task_failure_numbers[partition_id];
                 let task_info = TaskInfo {
                     task_id,
+                    executor_id: executor_id.to_owned(),
                     scheduled_time: SystemTime::now()
                         .duration_since(UNIX_EPOCH)
                         .unwrap()
@@ -1742,6 +1743,7 @@ impl Debug for StaticExecutionGraph {
 pub fn create_task_info(executor_id: String, task_id: usize) -> TaskInfo {
     TaskInfo {
         task_id,
+        executor_id: executor_id.clone(),
         scheduled_time: SystemTime::now()
             .duration_since(UNIX_EPOCH)
             .unwrap()

ballista/scheduler/src/state/execution_stage.rs

@@ -247,6 +247,14 @@ pub struct FailedStage {
 pub struct TaskInfo {
     /// Unique task identifier within the execution graph.
     pub task_id: usize,
+    /// ID of the executor that ran (or is running) this task.
+    ///
+    /// Carried at the top level so it survives terminal status transitions
+    /// to `Failed` — `FailedTask` does not embed an `executor_id`, so without
+    /// this field the scheduler would lose the executor mapping for any
+    /// failed partition. Populated at task launch and preserved through
+    /// `update_task_info`.
+    pub executor_id: String,
     /// Timestamp when the task was scheduled (in milliseconds since epoch).
     pub scheduled_time: u128,
     /// Timestamp when the task was launched on an executor (in milliseconds since epoch).
@@ -661,9 +669,11 @@ impl RunningStage {
             return false;
         }
         let scheduled_time = task_info.scheduled_time;
+        let executor_id = task_info.executor_id.clone();
         let task_status = status.status.unwrap();
         let updated_task_info = TaskInfo {
             task_id,
+            executor_id,
             scheduled_time,
             launch_time: status.launch_time as u128,
             start_exec_time: status.start_exec_time as u128,
@@ -898,6 +908,7 @@ impl SuccessfulStage {
                 } if *executor == *executor_id => {
                     *task = TaskInfo {
                         task_id: *task_id,
+                        executor_id: executor_id.clone(),
                         scheduled_time: *scheduled_time,
                         launch_time: 0,
                         start_exec_time: 0,
@@ -1119,6 +1130,7 @@ mod tests {
         // Simulate scheduling the task: populate the task slot.
         stage.task_infos[0] = Some(TaskInfo {
             task_id: 0,
+            executor_id: "executor-1".to_string(),
             scheduled_time: 50,
             launch_time: 0,
             start_exec_time: 0,
@@ -1148,6 +1160,7 @@ mod tests {
         for i in 0..2 {
             stage.task_infos[i] = Some(TaskInfo {
                 task_id: i,
+                executor_id: "executor-1".to_string(),
                 scheduled_time: 50,
                 launch_time: 100,
                 start_exec_time: 200,

ballista/scheduler/src/state/task_manager.rs

@@ -408,8 +408,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
 
     /// Get the execution graph of of a job. First look in the active cache.
     /// If no one found, then in the Active/Completed jobs.
-    #[cfg(feature = "rest-api")]
-    pub(crate) async fn get_job_execution_graph(
+    ///
+    /// Exposed as `pub` so embedded callers (e.g., Spice's distributed
+    /// task_history writer) can walk per-stage and per-task state directly
+    /// without going through a gRPC method.
+    pub async fn get_job_execution_graph(
         &self,
         job_id: &str,
     ) -> Result<Option<ExecutionGraphBox>> {