Add REST endpoints for getting list of jobs and job summary

ballista/rust/scheduler/src/api/handlers.rs

@@ -19,6 +19,9 @@ use warp::Rejection;
 #[derive(Debug, serde::Serialize)]
 struct StateResponse {
     executors: Vec<ExecutorMetaResponse>,
+    active_jobs: Vec<String>,
+    completed_jobs: Vec<String>,
+    failed_jobs: Vec<String>,
     started: u128,
     version: &'static str,
 }
@@ -31,12 +34,14 @@ pub struct ExecutorMetaResponse {
     pub last_seen: u128,
 }
 
-pub(crate) async fn scheduler_state<T: AsLogicalPlan, U: AsExecutionPlan>(
+/// Return current scheduler state, including list of executors and active, completed, and failed
+/// job ids.
+pub(crate) async fn get_state<T: AsLogicalPlan, U: AsExecutionPlan>(
     data_server: SchedulerServer<T, U>,
 ) -> Result<impl warp::Reply, Rejection> {
     // TODO: Display last seen information in UI
-    let executors: Vec<ExecutorMetaResponse> = data_server
-        .state
+    let state = data_server.state;
+    let executors: Vec<ExecutorMetaResponse> = state
         .executor_manager
         .get_executor_state()
         .await
@@ -51,8 +56,52 @@ pub(crate) async fn scheduler_state<T: AsLogicalPlan, U: AsExecutionPlan>(
         .collect();
     let response = StateResponse {
         executors,
+        active_jobs: state
+            .task_manager
+            .get_active_job_ids()
+            .await
+            .map_err(|_| warp::reject())?,
+        completed_jobs: state
+            .task_manager
+            .get_completed_job_ids()
+            .await
+            .map_err(|_| warp::reject())?,
+        failed_jobs: state
+            .task_manager
+            .get_failed_job_ids()
+            .await
+            .map_err(|_| warp::reject())?,
         started: data_server.start_time,
         version: BALLISTA_VERSION,
     };
     Ok(warp::reply::json(&response))
 }
+
+#[derive(Debug, serde::Serialize)]
+pub struct JobSummaryResponse {
+    /// Just show debug output for now but what we really want here is a list of stages with
+    /// plans and metrics and the relationship between them
+    pub summary: String,
+}
+
+/// Get the execution graph for the specified job id
+pub(crate) async fn get_job_summary<T: AsLogicalPlan, U: AsExecutionPlan>(
+    data_server: SchedulerServer<T, U>,
+    job_id: String,
+) -> Result<impl warp::Reply, Rejection> {
+    let graph = data_server
+        .state
+        .task_manager
+        .get_job_execution_graph(&job_id)
+        .await
+        .map_err(|_| warp::reject())?;
+
+    match graph {
+        Some(x) => Ok(warp::reply::json(&JobSummaryResponse {
+            summary: format!("{:?}", x),
+        })),
+        _ => Ok(warp::reply::json(&JobSummaryResponse {
+            summary: "Not Found".to_string(),
+        })),
+    }
+}

ballista/rust/scheduler/src/api/mod.rs

@@ -85,8 +85,19 @@ fn with_data_server<T: AsLogicalPlan + Clone, U: 'static + AsExecutionPlan>(
 pub fn get_routes<T: AsLogicalPlan + Clone, U: 'static + AsExecutionPlan>(
     scheduler_server: SchedulerServer<T, U>,
 ) -> BoxedFilter<(impl Reply,)> {
-    let routes = warp::path("state")
-        .and(with_data_server(scheduler_server))
-        .and_then(handlers::scheduler_state);
+    let route_state = warp::path("state")
+        .and(with_data_server(scheduler_server.clone()))
+        .and_then(handlers::get_state);
+
+    let route_job_summary = warp::path!("jobs" / String)
+        .and(with_data_server(scheduler_server.clone()))
+        .and_then(|job_id, data_server| handlers::get_job_summary(data_server, job_id));
+
+    // let route_graph = warp::path("jobs" / job_id / "graph" )
+    //     .and(with_data_server(scheduler_server.clone()))
+    //     .map(|x, job_id| handlers::get_execution_graph(x, job_id).await);
+    // .and(with_data_server(scheduler_server))
+    // .and_then();
+    let routes = route_state.or(route_job_summary);
     routes.boxed()
 }

ballista/rust/scheduler/src/state/task_manager.rs

@@ -106,6 +106,24 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
         Ok(())
     }
 
+    /// Get a list of active job ids
+    pub async fn get_active_job_ids(&self) -> Result<Vec<String>> {
+        let keys = self.state.scan_keys(Keyspace::ActiveJobs).await?;
+        Ok(keys.iter().cloned().collect::<Vec<_>>())
+    }
+
+    /// Get a list of completed job ids
+    pub async fn get_completed_job_ids(&self) -> Result<Vec<String>> {
+        let keys = self.state.scan_keys(Keyspace::CompletedJobs).await?;
+        Ok(keys.iter().cloned().collect::<Vec<_>>())
+    }
+
+    /// Get a list of failed job ids
+    pub async fn get_failed_job_ids(&self) -> Result<Vec<String>> {
+        let keys = self.state.scan_keys(Keyspace::FailedJobs).await?;
+        Ok(keys.iter().cloned().collect::<Vec<_>>())
+    }
+
     /// Get the status of of a job. First look in the active cache.
     /// If no one found, then in the Active/Completed jobs, and then in Failed jobs
     pub async fn get_job_status(&self, job_id: &str) -> Result<Option<JobStatus>> {
@@ -126,6 +144,22 @@ 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.
+    pub async fn get_job_execution_graph(
+        &self,
+        job_id: &str,
+    ) -> Result<Option<Arc<ExecutionGraph>>> {
+        if let Some(graph) = self.get_active_execution_graph(job_id).await {
+            Ok(Some(Arc::new(graph.read().await.clone())))
+        } else if let Ok(graph) = self.get_execution_graph(job_id).await {
+            Ok(Some(Arc::new(graph.clone())))
+        } else {
+            // if the job failed then we return no graph for now
+            Ok(None)
+        }
+    }
+
     /// Update given task statuses in the respective job and return a tuple containing:
     /// 1. A list of QueryStageSchedulerEvent to publish.
     /// 2. A list of reservations that can now be offered.