feat(scheduler): record a per-job event log behind --event-log-dir

ballista/scheduler/Cargo.toml

@@ -41,7 +41,7 @@ disable-stage-plan-cache = []
 graphviz-support = ["dep:graphviz-rust"]
 keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"]
 prometheus-metrics = ["prometheus", "once_cell"]
-rest-api = ["dep:ballista-api-types"]
+rest-api = ["dep:ballista-api-types", "dep:ballista-history", "dep:serde_json"]
 spark-compat = ["ballista-core/spark-compat"]
 substrait = ["dep:datafusion-substrait"]
 
@@ -51,6 +51,7 @@ async-trait = { workspace = true }
 axum = "0.8.9"
 ballista-api-types = { path = "../api-types", version = "54.0.0", optional = true }
 ballista-core = { path = "../core", version = "54.0.0" }
+ballista-history = { path = "../history", version = "54.0.0", optional = true }
 clap = { workspace = true, optional = true }
 dashmap = { workspace = true }
 datafusion = { workspace = true }
@@ -70,6 +71,7 @@ prost = { workspace = true }
 prost-types = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+serde_json = { version = "1", optional = true }
 tokio = { workspace = true, features = ["full"] }
 tokio-stream = { workspace = true, features = ["net"] }
 tonic = { workspace = true, features = ["router"] }
@@ -90,6 +92,7 @@ datafusion-functions-aggregate-common = { workspace = true }
 regex = "1"
 rstest = { workspace = true }
 serde_json = "1"
+tempfile = { workspace = true }
 
 [build-dependencies]
 tonic-prost-build = { workspace = true, optional = true }

ballista/scheduler/src/api/dto_build.rs

@@ -29,17 +29,19 @@
 
 use crate::display::format_stage_metrics;
 use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage};
+use crate::state::execution_graph_dot::ExecutionGraphDot;
 use crate::state::execution_stage::TaskInfo;
 use crate::state::task_manager::JobOverview;
 use ballista_api_types::dto::{
-    JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse,
-    TaskStatus, TaskSummary,
+    JobConfig, JobResponse, Percentiles, PlanFormat, QueryStageSummary,
+    QueryStagesResponse, TaskStatus, TaskSummary,
 };
 use ballista_core::serde::protobuf::failed_task::FailedReason::{
     ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled,
 };
 use ballista_core::serde::protobuf::job_status::Status;
-use ballista_core::serde::protobuf::{FailedTask, task_status};
+use ballista_core::serde::protobuf::{FailedTask, OperatorMetricsSet, task_status};
+use datafusion::execution::context::SessionConfig;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::displayable;
 use datafusion::physical_plan::metrics::{MetricsSet, Time};
@@ -115,6 +117,12 @@ pub fn graph_to_job_response(
     }
 }
 
+/// Flatten a session config into the sorted key/value map served by
+/// `GET /api/job/{job_id}/config`.
+pub fn session_config_to_job_config(config: &SessionConfig) -> JobConfig {
+    config.to_props().into_iter().collect()
+}
+
 /// Build the per-stage summaries served by `GET /api/job/{job_id}/stages`.
 pub fn graph_to_query_stages(
     graph: &ExecutionGraphBox,
@@ -225,11 +233,48 @@ fn task_summaries(
         .collect()
 }
 
+/// Render a job's stage DAG in DOT format.
+pub fn build_job_dot(graph: &ExecutionGraphBox) -> Result<String, std::fmt::Error> {
+    ExecutionGraphDot::generate(graph.as_ref())
+}
+
+/// Sum one task's raw operator metrics into
+/// `(input_rows, output_rows, elapsed_compute_nanos)`.
+///
+/// Distinct from [`get_partition_counts`], which reads a stage's already-merged
+/// [`MetricsSet`]s and filters by partition. This takes the raw protobuf
+/// [`OperatorMetricsSet`]s an executor reports for a single task, so there is no
+/// partition to filter on, and it also sums `elapsed_compute` for the event
+/// log's per-task timeline records.
+pub fn task_row_counts(metrics: &[OperatorMetricsSet]) -> (u64, u64, u64) {
+    let mut input_rows: u64 = 0;
+    let mut output_rows: u64 = 0;
+    let mut elapsed_compute_nanos: u64 = 0;
+
+    for operator_metrics in metrics {
+        let Ok(metrics_set) = TryInto::<MetricsSet>::try_into(operator_metrics.clone())
+        else {
+            continue;
+        };
+        for metric in metrics_set.iter() {
+            let value = metric.value();
+            match value.name() {
+                "input_rows" => input_rows += value.as_usize() as u64,
+                "output_rows" => output_rows += value.as_usize() as u64,
+                "elapsed_compute" => elapsed_compute_nanos += value.as_usize() as u64,
+                _ => {}
+            }
+        }
+    }
+
+    (input_rows, output_rows, elapsed_compute_nanos)
+}
+
 /// Map a protobuf task status onto the wire enum.
 ///
 /// A free function rather than a `From` impl: both types are foreign to this
 /// crate now that [`TaskStatus`] lives in `ballista-api-types`.
-fn task_status_to_dto(value: &task_status::Status) -> TaskStatus {
+pub fn task_status_to_dto(value: &task_status::Status) -> TaskStatus {
     match value {
         task_status::Status::Running(_) => TaskStatus::Running,
         task_status::Status::Failed(failed) => TaskStatus::Failed {

ballista/scheduler/src/api/mod.rs

@@ -11,7 +11,9 @@
 // limitations under the License.
 
 #[cfg(feature = "rest-api")]
-mod dto_build;
+// `pub(crate)` so `scheduler_server::event_log` can build the same DTOs the
+// live REST API serves.
+pub(crate) mod dto_build;
 #[cfg(feature = "rest-api")]
 mod handlers;
 mod health;

ballista/scheduler/src/config.rs

@@ -147,6 +147,9 @@ pub struct Config {
         help = "Log dir: a path to save log. This will create a new storage directory at the specified path if it does not already exist."
     )]
     pub log_dir: Option<String>,
+    /// Directory to write per-job event logs to. Enables the history server.
+    #[arg(long)]
+    pub event_log_dir: Option<String>,
     /// Whether to print thread IDs and names in log files.
     #[arg(
         long,
@@ -365,6 +368,8 @@ pub struct SchedulerConfig {
     #[cfg(feature = "rest-api")]
     /// Comma-separated list of allowed methods for CORS
     pub cors_allowed_methods: String,
+    /// Directory to write per-job event logs to. `None` disables event logging.
+    pub event_log_dir: Option<String>,
     #[cfg(feature = "rest-api")]
     /// The HTTP path that will redirect to the WebTUI app at `https://nightlies.apache.org`
     pub web_tui_route: String,
@@ -410,6 +415,7 @@ impl Default for SchedulerConfig {
             cors_allowed_origins: String::default(),
             #[cfg(feature = "rest-api")]
             cors_allowed_methods: String::default(),
+            event_log_dir: None,
             #[cfg(feature = "rest-api")]
             web_tui_route: String::from("/"),
             on_work_available: None,
@@ -556,6 +562,12 @@ impl SchedulerConfig {
         self
     }
 
+    /// Sets the directory to write per-job event logs to.
+    pub fn with_event_log_dir(mut self, event_log_dir: Option<String>) -> Self {
+        self.event_log_dir = event_log_dir;
+        self
+    }
+
     /// Sets whether TLS should be used when connecting to executors (for flight proxy).
     pub fn with_use_tls(mut self, use_tls: bool) -> Self {
         self.use_tls = use_tls;
@@ -671,6 +683,7 @@ impl TryFrom<Config> for SchedulerConfig {
             cors_allowed_origins: opt.cors_allowed_origins,
             #[cfg(feature = "rest-api")]
             cors_allowed_methods: opt.cors_allowed_methods,
+            event_log_dir: opt.event_log_dir,
             #[cfg(feature = "rest-api")]
             web_tui_route: opt.web_tui_route,
             on_work_available: None,