From 7cbcfea8ac984723dd70acf2b5792600f753a51c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 18:10:14 -0600 Subject: [PATCH 01/14] feat(history): add ballista-history crate with event-log schema --- Cargo.toml | 2 +- ballista/history/Cargo.toml | 16 +++++ ballista/history/src/dto.rs | 97 ++++++++++++++++++++++++++ ballista/history/src/event.rs | 124 +++++++++++++++++++++++++++++++++ ballista/history/src/lib.rs | 21 ++++++ ballista/history/src/reader.rs | 16 +++++ ballista/history/src/writer.rs | 16 +++++ 7 files changed, 291 insertions(+), 1 deletion(-) create mode 100644 ballista/history/Cargo.toml create mode 100644 ballista/history/src/dto.rs create mode 100644 ballista/history/src/event.rs create mode 100644 ballista/history/src/lib.rs create mode 100644 ballista/history/src/reader.rs create mode 100644 ballista/history/src/writer.rs diff --git a/Cargo.toml b/Cargo.toml index 5719b9b16..71e28483b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ [workspace] exclude = ["dev/msrvcheck", "python"] -members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/scheduler", "benchmarks", "examples"] +members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/history", "ballista/scheduler", "benchmarks", "examples"] resolver = "3" [workspace.package] diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml new file mode 100644 index 000000000..26751269b --- /dev/null +++ b/ballista/history/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "ballista-history" +description = "Event-log schema, writer, and reader for Ballista's history server" +license = "Apache-2.0" +version = "53.0.0" +edition = { workspace = true } +rust-version = { workspace = true } + +[dependencies] +serde = { workspace = true, features = ["derive"] } +serde_json = "1" +tokio = { workspace = true, features = ["rt", "sync", "fs", "io-util", "macros"] } + +[dev-dependencies] +tokio = { workspace = true, features = ["rt", "sync", "fs", "io-util", "macros", "rt-multi-thread"] } +tempfile = { workspace = true } diff --git a/ballista/history/src/dto.rs b/ballista/history/src/dto.rs new file mode 100644 index 000000000..3485bfc00 --- /dev/null +++ b/ballista/history/src/dto.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! REST response DTOs shared by the live scheduler API handlers and the history +//! server, so both serialize byte-identical JSON. + +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobResponse { + pub job_id: String, + pub job_name: String, + pub job_status: String, + pub status: String, + pub num_stages: usize, + pub completed_stages: usize, + pub percent_complete: u8, + pub start_time: u64, + pub end_time: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub logical_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub physical_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum TaskStatus { + Running, + Successful, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskSummary { + pub id: usize, + pub status: TaskStatus, + pub partition_id: u32, + pub scheduled_time: u64, + pub launch_time: u64, + pub start_exec_time: u64, + pub end_exec_time: u64, + pub exec_duration: u64, + pub finish_time: u64, + pub input_rows: usize, + pub output_rows: usize, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Percentiles { + pub min: u64, + pub p25: u64, + pub median: u64, + pub p75: u64, + pub max: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStageSummary { + pub stage_id: String, + pub stage_status: String, + pub input_rows: usize, + pub output_rows: usize, + pub elapsed_compute: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_duration_percentiles: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_input_percentiles: Option, + pub tasks: Vec>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStagesResponse { + pub stages: Vec, +} + +/// Session config as flat key/value pairs (from `SessionConfig::to_props()`), +/// sorted for stable output. +pub type JobConfig = BTreeMap; diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs new file mode 100644 index 000000000..8db1c05f4 --- /dev/null +++ b/ballista/history/src/event.rs @@ -0,0 +1,124 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The on-disk event-log schema. One `HistoryEvent` is serialized per JSONL line. +//! This is a frozen public projection of the scheduler's internal events; the +//! embedded DTOs are the stable contract the history server serves. + +use crate::dto::{JobConfig, JobResponse, QueryStagesResponse, TaskStatus}; +use serde::{Deserialize, Serialize}; + +/// Current on-disk schema version, stamped on `JobStart`/`JobEnd`. +pub const SCHEMA_VERSION: u32 = 1; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum JobEndStatus { + Succeeded, + Failed(String), +} + +/// Metrics captured per finished task on the incremental timeline. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskEndMetrics { + pub input_rows: u64, + pub output_rows: u64, + pub elapsed_compute_nanos: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "ev")] +pub enum HistoryEvent { + JobStart { + version: u32, + job_id: String, + job_name: String, + queued_at: u64, + submitted_at: u64, + logical_plan: Option, + physical_plan: Option, + }, + StageStart { + stage_id: usize, + partitions: usize, + }, + StageEnd { + stage_id: usize, + status: String, + }, + TaskEnd { + stage_id: u32, + partition: u32, + executor_id: String, + status: TaskStatus, + launch_time: u64, + start_exec_time: u64, + end_exec_time: u64, + metrics: TaskEndMetrics, + }, + JobEnd { + version: u32, + status: JobEndStatus, + queued_at: u64, + started_at: u64, + completed_at: u64, + job: Box, + stages: Box, + config: JobConfig, + dot: String, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::{JobResponse, QueryStagesResponse}; + use std::collections::BTreeMap; + + #[test] + fn job_end_round_trips_through_jsonl() { + let job = JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 2, + completed_stages: 2, + percent_complete: 100, + start_time: 10, + end_time: 20, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage plan".into()), + }; + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 5, + started_at: 10, + completed_at: 20, + job: Box::new(job), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: BTreeMap::from([("k".to_string(), "v".to_string())]), + dot: "digraph {}".into(), + }; + let line = serde_json::to_string(&event).unwrap(); + assert!(line.contains("\"ev\":\"JobEnd\"")); + let back: HistoryEvent = serde_json::from_str(&line).unwrap(); + // Re-serialize and compare strings (stable, discriminating). + assert_eq!(line, serde_json::to_string(&back).unwrap()); + } +} diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs new file mode 100644 index 000000000..0b17d91d6 --- /dev/null +++ b/ballista/history/src/lib.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod dto; +pub mod event; +pub mod reader; +pub mod writer; diff --git a/ballista/history/src/reader.rs b/ballista/history/src/reader.rs new file mode 100644 index 000000000..b248758bc --- /dev/null +++ b/ballista/history/src/reader.rs @@ -0,0 +1,16 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs new file mode 100644 index 000000000..b248758bc --- /dev/null +++ b/ballista/history/src/writer.rs @@ -0,0 +1,16 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. From 564a999d60a0d45b2110460e8e0035812cedd6fc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 18:14:10 -0600 Subject: [PATCH 02/14] feat(history): add async buffered event-log writer --- ballista/history/src/writer.rs | 163 +++++++++++++++++++++++++++++++++ 1 file changed, 163 insertions(+) diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs index b248758bc..b7b79132c 100644 --- a/ballista/history/src/writer.rs +++ b/ballista/history/src/writer.rs @@ -14,3 +14,166 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + +//! Async, buffered event-log writer. Each job's events append to +//! `/.eventlog` as JSONL. Appends are non-blocking; a background +//! task performs the file I/O so the scheduler hot path never waits on disk. + +use crate::event::HistoryEvent; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use tokio::io::AsyncWriteExt; +use tokio::sync::{mpsc, oneshot}; + +enum WriterMsg { + Event { + job_id: String, + event: Box, + }, + Flush { + job_id: String, + done: oneshot::Sender<()>, + }, +} + +#[derive(Clone)] +pub struct EventLogWriter { + tx: mpsc::Sender, +} + +impl EventLogWriter { + pub fn new(log_dir: PathBuf, buffer: usize) -> EventLogWriter { + let (tx, rx) = mpsc::channel(buffer.max(1)); + tokio::spawn(run(log_dir, rx)); + EventLogWriter { tx } + } + + /// Enqueue an event for `job_id`. Never blocks; drops (with a warning) if the + /// channel is full, so logging cannot stall scheduling. + pub fn append(&self, job_id: &str, event: HistoryEvent) { + let msg = WriterMsg::Event { + job_id: job_id.to_string(), + event: Box::new(event), + }; + if self.tx.try_send(msg).is_err() { + eprintln!( + "event-log writer: dropping event for {job_id} (channel full or closed)" + ); + } + } + + /// Await all currently-enqueued writes for `job_id` (best effort). + pub async fn flush_job(&self, job_id: &str) { + let (done, wait) = oneshot::channel(); + if self + .tx + .send(WriterMsg::Flush { + job_id: job_id.to_string(), + done, + }) + .await + .is_ok() + { + let _ = wait.await; + } + } +} + +async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { + if let Err(e) = tokio::fs::create_dir_all(&log_dir).await { + eprintln!("event-log writer: cannot create {}: {e}", log_dir.display()); + return; + } + // One open append handle per job for the life of the process. + let mut handles: HashMap = HashMap::new(); + + while let Some(msg) = rx.recv().await { + match msg { + WriterMsg::Event { job_id, event } => { + let file = match open_for(&log_dir, &mut handles, &job_id).await { + Some(f) => f, + None => continue, + }; + match serde_json::to_string(&*event) { + Ok(mut line) => { + line.push('\n'); + if let Err(e) = file.write_all(line.as_bytes()).await { + eprintln!("event-log writer: write failed for {job_id}: {e}"); + } + } + Err(e) => eprintln!("event-log writer: serialize failed: {e}"), + } + } + WriterMsg::Flush { job_id, done } => { + if let Some(file) = handles.get_mut(&job_id) { + let _ = file.flush().await; + } + let _ = done.send(()); + } + } + } +} + +async fn open_for<'a>( + log_dir: &Path, + handles: &'a mut HashMap, + job_id: &str, +) -> Option<&'a mut tokio::fs::File> { + if !handles.contains_key(job_id) { + let path = log_dir.join(format!("{job_id}.eventlog")); + match tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .await + { + Ok(f) => { + handles.insert(job_id.to_string(), f); + } + Err(e) => { + eprintln!("event-log writer: cannot open {}: {e}", path.display()); + return None; + } + } + } + handles.get_mut(job_id) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::event::{HistoryEvent, SCHEMA_VERSION}; + + #[tokio::test] + async fn append_writes_one_jsonl_line_per_event() { + let dir = tempfile::tempdir().unwrap(); + let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); + writer.append( + "job-1", + HistoryEvent::JobStart { + version: SCHEMA_VERSION, + job_id: "job-1".into(), + job_name: "q1".into(), + queued_at: 1, + submitted_at: 2, + logical_plan: None, + physical_plan: None, + }, + ); + writer.append( + "job-1", + HistoryEvent::StageStart { + stage_id: 1, + partitions: 4, + }, + ); + writer.flush_job("job-1").await; + + let path = dir.path().join("job-1.eventlog"); + let contents = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + assert_eq!(lines.len(), 2); + assert!(lines[0].contains("\"ev\":\"JobStart\"")); + assert!(lines[1].contains("\"ev\":\"StageStart\"")); + } +} From 20daac9ee4e90d05f6fc08dd7f5a9b0ea68253d3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 18:37:10 -0600 Subject: [PATCH 03/14] refactor(scheduler): share REST DTOs and extract graph-to-DTO builders Move the scheduler REST API's JobResponse/TaskSummary/TaskStatus/Percentiles/ QueryStageSummary/QueryStagesResponse types onto the ballista-history crate's shared DTOs, and extract the graph-to-DTO construction logic that previously lived inline in the handlers into pub(crate) builder functions in a new api::dto_build module. Handlers become thin wrappers over these builders, which a later event-log writer will also call. Behavior is unchanged; the existing handler unit tests move to dto_build.rs alongside the moved helper functions they exercise, plus one new plan-string assertion test. --- Cargo.lock | 12 + ballista/scheduler/Cargo.toml | 2 + ballista/scheduler/src/api/dto_build.rs | 658 ++++++++++++++++++ ballista/scheduler/src/api/handlers.rs | 617 +--------------- ballista/scheduler/src/api/mod.rs | 2 + .../src/state/execution_graph_dot.rs | 4 +- 6 files changed, 693 insertions(+), 602 deletions(-) create mode 100644 ballista/scheduler/src/api/dto_build.rs diff --git a/Cargo.lock b/Cargo.lock index 369a96718..b2973825b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1168,6 +1168,16 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "53.0.0" +dependencies = [ + "serde", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "ballista-scheduler" version = "53.0.0" @@ -1176,6 +1186,7 @@ dependencies = [ "async-trait", "axum", "ballista-core", + "ballista-history", "clap 4.6.1", "dashmap", "datafusion", @@ -1196,6 +1207,7 @@ dependencies = [ "rand 0.10.1", "rstest", "serde", + "serde_json", "tokio", "tokio-stream", "tonic", diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index f882d26f3..1ba88fa13 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -51,6 +51,7 @@ async-trait = { workspace = true } axum = "0.8.9" tower-http = { version = "0.7", features = ["cors"] } ballista-core = { path = "../core", version = "53.0.0" } +ballista-history = { path = "../history", version = "53.0.0" } clap = { workspace = true, optional = true } dashmap = { workspace = true } datafusion = { workspace = true } @@ -81,6 +82,7 @@ uuid = { workspace = true } [dev-dependencies] rstest = { workspace = true } +serde_json = "1" [build-dependencies] tonic-prost-build = { workspace = true, optional = true } diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs new file mode 100644 index 000000000..8d4d18fd7 --- /dev/null +++ b/ballista/scheduler/src/api/dto_build.rs @@ -0,0 +1,658 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Builds the shared REST DTOs (`ballista_history::dto`) from the scheduler's +//! internal execution-graph state. These builders back both the live REST API +//! handlers (`api::handlers`) and, eventually, the event-log writer, so both +//! serialize byte-identical JSON for the same graph state. + +use crate::api::handlers::{JobQueryParams, PlanFormat}; +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_core::serde::protobuf::job_status::Status; +use ballista_core::serde::protobuf::{OperatorMetricsSet, task_status}; +use ballista_core::utils::get_current_time; +use ballista_history::dto::{ + JobResponse, Percentiles, QueryStageSummary, QueryStagesResponse, TaskStatus, + TaskSummary, +}; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::displayable; +use datafusion::physical_plan::metrics::{MetricsSet, Time}; +use std::fmt; +use std::time::Duration; + +/// Builds the `JobResponse` DTO for a job's execution graph. +/// +/// `with_plans` controls whether `logical_plan`/`physical_plan`/`stage_plan` are +/// populated (the single-job detail view, and eventually the history detail +/// view) or left `None` (list views, which use +/// [`build_job_response_from_overview`] instead). `plan_format` selects how the +/// physical plan is rendered when `with_plans` is true; it is ignored otherwise. +pub(crate) fn build_job_response( + graph: &ExecutionGraphBox, + with_plans: bool, + plan_format: PlanFormat, +) -> JobResponse { + let job = graph.as_ref(); + let (plain_status, job_status) = format_job_status( + &job.status().status, + job_elapsed_ms(job.start_time(), job.end_time()), + ); + + let num_stages = job.stage_count(); + let completed_stages = job.completed_stages(); + let percent_complete = + ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; + + let (logical_plan, physical_plan, stage_plan) = if with_plans { + let physical_plan = match plan_format { + PlanFormat::Default | PlanFormat::Metrics => { + DisplayableExecutionPlan::new(job.physical_plan().as_ref()) + .indent(false) + .to_string() + } + PlanFormat::Tree => displayable(job.physical_plan().as_ref()) + .tree_render() + .to_string(), + }; + ( + job.logical_plan().map(str::to_owned), + Some(physical_plan), + Some(format!("{:?}", graph)), + ) + } else { + (None, None, None) + }; + + JobResponse { + job_id: job.job_id().to_string(), + job_name: job.job_name().to_owned(), + job_status, + status: plain_status, + start_time: job.start_time(), + end_time: job.end_time(), + num_stages, + completed_stages, + percent_complete, + logical_plan, + physical_plan, + stage_plan, + } +} + +/// Builds the `JobResponse` DTO for a job's cached [`JobOverview`] (no +/// execution graph, and hence no plans available). +pub(crate) fn build_job_response_from_overview(job: &JobOverview) -> JobResponse { + let (plain_status, job_status) = format_job_status( + &job.status.status, + job_elapsed_ms(job.start_time, job.end_time), + ); + + // calculate progress based on completed stages for now, but we could use completed + // tasks in the future to make this more accurate + let percent_complete = if job.num_stages == 0 { + 0 + } else { + ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 + }; + JobResponse { + job_id: job.job_id.to_string(), + job_name: job.job_name.to_owned(), + job_status, + status: plain_status, + start_time: job.start_time, + end_time: job.end_time, + num_stages: job.num_stages, + completed_stages: job.completed_stages, + percent_complete, + logical_plan: None, + physical_plan: None, + stage_plan: None, + } +} + +/// Builds the `QueryStagesResponse` DTO for a job's execution graph. +pub(crate) fn build_query_stages_response( + graph: &ExecutionGraphBox, + params: &JobQueryParams, +) -> QueryStagesResponse { + let plan_format = params.plan_format.clone().unwrap_or_default(); + + let stages = graph + .as_ref() + .stages() + .iter() + .map(|(id, stage)| { + let mut summary = QueryStageSummary { + stage_id: id.to_string(), + stage_status: stage.variant_name().to_string(), + input_rows: 0, + output_rows: 0, + elapsed_compute: None, + tasks: vec![], + task_duration_percentiles: None, + task_input_percentiles: None, + stage_plan: None, + }; + match stage { + ExecutionStage::Running(running_stage) => { + let metrics = running_stage.stage_metrics.as_deref().unwrap_or(&[]); + summary.stage_plan = Some(match plan_format { + PlanFormat::Default => displayable(running_stage.plan.as_ref()) + .indent(false) + .to_string(), + PlanFormat::Tree => displayable(running_stage.plan.as_ref()) + .tree_render() + .to_string(), + PlanFormat::Metrics => { + format_stage_metrics(running_stage.plan.as_ref(), metrics) + } + }); + summary.input_rows = running_stage + .stage_metrics + .as_ref() + .map(|m| get_combined_count(m.as_slice(), "input_rows")) + .unwrap_or(0); + summary.output_rows = running_stage + .stage_metrics + .as_ref() + .map(|m| get_combined_count(m.as_slice(), "output_rows")) + .unwrap_or(0); + summary.elapsed_compute = get_running_stage_time( + &running_stage.task_infos, + get_current_time(), + ); + summary.tasks = running_stage + .task_infos + .iter() + .enumerate() + .map(|(partition_id, task_info)| { + task_info.as_ref().map(|info| { + let (input_rows, output_rows) = running_stage + .stage_metrics + .as_deref() + .map(|metrics| { + get_partition_counts(metrics, partition_id) + }) + .unwrap_or((0, 0)); + + let start_exec_time = info.start_exec_time as u64; + let end_exec_time = info.end_exec_time as u64; + + let task_status = to_api_task_status(&info.task_status); + + TaskSummary { + id: info.task_id, + partition_id: partition_id as u32, + scheduled_time: info.scheduled_time as u64, + launch_time: info.launch_time as u64, + start_exec_time, + end_exec_time, + exec_duration: end_exec_time + .saturating_sub(start_exec_time), + finish_time: info.finish_time as u64, + input_rows, + output_rows, + status: task_status, + } + }) + }) + .collect(); + } + ExecutionStage::Successful(completed_stage) => { + summary.stage_plan = Some(match plan_format { + PlanFormat::Default => displayable(completed_stage.plan.as_ref()) + .indent(false) + .to_string(), + PlanFormat::Tree => displayable(completed_stage.plan.as_ref()) + .tree_render() + .to_string(), + PlanFormat::Metrics => format_stage_metrics( + completed_stage.plan.as_ref(), + &completed_stage.stage_metrics, + ), + }); + summary.input_rows = + get_combined_count(&completed_stage.stage_metrics, "input_rows"); + summary.output_rows = + get_combined_count(&completed_stage.stage_metrics, "output_rows"); + summary.elapsed_compute = + get_finished_stage_time(&completed_stage.task_infos); + + summary.tasks = completed_stage + .task_infos + .iter() + .enumerate() + .map(|(partition_id, task_info)| { + let (input_rows, output_rows) = get_partition_counts( + &completed_stage.stage_metrics, + partition_id, + ); + + let start_exec_time = task_info.start_exec_time as u64; + let end_exec_time = task_info.end_exec_time as u64; + let task_status = to_api_task_status(&task_info.task_status); + Some(TaskSummary { + id: task_info.task_id, + partition_id: partition_id as u32, + scheduled_time: task_info.scheduled_time as u64, + launch_time: task_info.launch_time as u64, + start_exec_time, + end_exec_time, + exec_duration: end_exec_time + .saturating_sub(start_exec_time), + finish_time: task_info.finish_time as u64, + input_rows, + output_rows, + status: task_status, + }) + }) + .collect(); + } + _ => {} + } + summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); + summary.task_input_percentiles = task_input_percentiles(&summary.tasks); + summary + }) + .collect(); + + QueryStagesResponse { stages } +} + +/// Builds the DOT-format graph for a job's execution graph. +pub(crate) fn build_job_dot(graph: &ExecutionGraphBox) -> Result { + ExecutionGraphDot::generate(graph.as_ref()) +} + +/// Converts a task's protobuf status into the shared `TaskStatus` DTO. +pub(crate) fn to_api_task_status(status: &task_status::Status) -> TaskStatus { + match status { + task_status::Status::Running(_) => TaskStatus::Running, + task_status::Status::Failed(_) => TaskStatus::Failed, + task_status::Status::Successful(_) => TaskStatus::Successful, + } +} + +/// Sums a single task's raw operator metrics into +/// `(input_rows, output_rows, elapsed_compute_nanos)`. +/// +/// Note: this is *not* a literal extraction of `get_partition_counts`/ +/// `get_combined_count` below. Those operate on the stage's already-merged +/// per-partition [`MetricsSet`]s (DataFusion's own metrics type) and filter by +/// partition id, returning `usize` row counts only. `task_row_counts` instead +/// converts a single task's raw protobuf [`OperatorMetricsSet`]s (as reported +/// by an executor for one task, e.g. `TaskStatus::metrics`) directly, with no +/// partition filtering (a task already corresponds to exactly one partition), +/// and additionally sums `elapsed_compute` (nanoseconds) alongside the row +/// counts, matching `ballista_history::event::TaskEndMetrics` for Task 5's +/// timeline events. +/// +/// Not yet called from production code — the event-log writer that will +/// consume it lands in a later task. Silence `dead_code` until then. +#[allow(dead_code)] +pub(crate) 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::::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) +} + +fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { + let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; + sorted[idx.min(sorted.len() - 1)] +} + +fn task_input_percentiles(tasks: &[Option]) -> Option { + let mut durations: Vec = tasks + .iter() + .flatten() + .map(|t| t.input_rows as u64) + .collect(); + + if durations.is_empty() { + return None; + } + + durations.sort_unstable(); + + Some(Percentiles { + min: durations[0], + p25: percentile_duration(&durations, 25.0), + median: percentile_duration(&durations, 50.0), + p75: percentile_duration(&durations, 75.0), + max: *durations.last().unwrap(), + }) +} + +fn task_duration_percentiles(tasks: &[Option]) -> Option { + let mut durations: Vec = + tasks.iter().flatten().map(|t| t.exec_duration).collect(); + + if durations.is_empty() { + return None; + } + + durations.sort_unstable(); + + Some(Percentiles { + min: durations[0], + p25: percentile_duration(&durations, 25.0), + median: percentile_duration(&durations, 50.0), + p75: percentile_duration(&durations, 75.0), + max: *durations.last().unwrap(), + }) +} + +/// Returns elapsed wall time in milliseconds for API formatting. +/// +/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or +/// `end_time` still zero while `start_time` is set) do not panic on subtract. +fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { + end_time.saturating_sub(start_time) +} + +fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { + match status { + Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), + Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), + Some(Status::Failed(error)) => { + ("Failed".to_string(), format!("Failed: {}", error.error)) + } + Some(Status::Successful(completed)) => { + let num_rows = completed + .partition_location + .iter() + .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) + .sum::(); + let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; + let num_partitions = completed.partition_location.len(); + let num_partitions_term = if num_partitions == 1 { + "partition" + } else { + "partitions" + }; + ( + "Completed".to_string(), + format!( + "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", + num_partitions, + num_partitions_term, + num_rows, + num_rows_term, + elapsed_ms + ), + ) + } + _ => ("Invalid".to_string(), "Invalid State".to_string()), + } +} + +fn get_running_stage_time( + task_infos: &[Option], + current_time: u128, +) -> Option { + let min_start = task_infos + .iter() + .flat_map(|t| t.as_ref().map(|t| t.start_exec_time)) + .filter(|t| *t > 0) + .min(); + + match (min_start, current_time) { + (Some(start), end) if end >= start => { + let time = Time::new(); + time.add_duration(Duration::from_millis((end - start) as u64)); + Some(time.to_string()) + } + _ => None, + } +} + +fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { + let min_start = task_infos + .iter() + .map(|t| t.start_exec_time) + .filter(|t| *t > 0) + .min(); + + let max_end = task_infos + .iter() + .map(|t| t.end_exec_time) + .filter(|t| *t > 0) + .max(); + + match (min_start, max_end) { + (Some(start), Some(end)) if end >= start => { + let time = Time::new(); + time.add_duration(Duration::from_millis((end - start) as u64)); + Some(time.to_string()) + } + _ => None, + } +} + +fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) { + let input_rows = get_partition_count(metrics, partition_id, "input_rows"); + let output_rows = get_partition_count(metrics, partition_id, "output_rows"); + (input_rows, output_rows) +} + +fn get_partition_count(metrics: &[MetricsSet], partition_id: usize, name: &str) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + if metric.partition() == Some(partition_id) && metric_value.name() == name + { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + if metric_value.name() == name { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::execution_graph_dot::tests::test_graph; + use ballista_core::serde::protobuf::task_status; + + fn make_task_info(start: u128, end: u128) -> TaskInfo { + TaskInfo { + task_id: 0, + scheduled_time: 0, + launch_time: 0, + start_exec_time: start, + end_exec_time: end, + finish_time: 0, + task_status: task_status::Status::Running(Default::default()), + } + } + + #[test] + fn test_job_elapsed_saturates_when_end_precedes_start() { + assert_eq!(job_elapsed_ms(900, 100), 0); + } + + // --- get_finished_stage_time --- + + #[test] + fn test_finished_empty_slice_returns_none() { + assert_eq!(get_finished_stage_time(&[]), None); + } + + #[test] + fn test_finished_all_zero_timestamps_returns_none() { + let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + #[test] + fn test_finished_single_task_elapsed() { + // 600 - 100 = 500 ms → "500.00ms" + let tasks = vec![make_task_info(100, 600)]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("500.00ms".to_string()) + ); + } + + #[test] + fn test_finished_picks_earliest_start_and_latest_end() { + // min start = 100, max end = 900 → 800 ms + let tasks = vec![ + make_task_info(100, 500), + make_task_info(200, 900), + make_task_info(300, 700), + ]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("800.00ms".to_string()) + ); + } + + #[test] + fn test_finished_end_before_start_returns_none() { + let tasks = vec![make_task_info(900, 100)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + // --- get_running_stage_time --- + + #[test] + fn test_running_empty_slice_returns_none() { + assert_eq!(get_running_stage_time(&[], 1000), None); + } + + #[test] + fn test_running_all_none_returns_none() { + let tasks: Vec> = vec![None, None]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_future_start_returns_none() { + // start_exec_time beyond current time → elapsed clamped to 0 + let tasks = vec![Some(make_task_info(u128::MAX, 0))]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_past_start_returns_some() { + let now = 4_000; + let start = 1_000; + let tasks = vec![Some(make_task_info(start, 0))]; + assert_eq!( + get_running_stage_time(&tasks, now), + Some("3.00s".to_string()) + ); + } + + #[test] + fn test_running_mixed_some_none_uses_earliest_some() { + let now = 3_000; + let earlier = 1_000; + let later = 2_000; + let tasks = vec![ + None, + Some(make_task_info(later, 0)), + Some(make_task_info(earlier, 0)), + None, + ]; + let result = get_running_stage_time(&tasks, now); + assert_eq!(result, Some("2.00s".to_string())); + } + + #[test] + fn test_job_elapsed_ms_normal() { + assert_eq!(super::job_elapsed_ms(100, 500), 400); + } + + #[test] + fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { + assert_eq!(super::job_elapsed_ms(500, 100), 0); + } + + #[test] + fn test_task_row_counts_empty_is_zero() { + assert_eq!(task_row_counts(&[]), (0, 0, 0)); + } + + #[tokio::test] + async fn test_dto_builders_plan_string() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + // build_query_stages_response: sanity-check the stage topology survives + // the DTO builder + JSON serialization round trip. test_graph() is a + // freshly built (never executed) graph, so every stage is + // Resolved/Unresolved and `stage_plan` is intentionally left `None` + // (only Running/Successful stages populate it) -- so this builder alone + // can't be used for the plan-string assertion below. + let params = JobQueryParams::default(); + let stages_response = build_query_stages_response(&graph, ¶ms); + let stages_json = serde_json::to_string(&stages_response).unwrap(); + assert!( + stages_json.contains(r#""stage_status":"Resolved""#), + "expected a resolved stage in response, got: {stages_json}" + ); + + // build_job_dot: assert on the embedded physical-plan string, following + // the repo's plan-string assertion convention (mirrors + // `execution_graph_dot::tests::dot()`, which renders the same graph). + let dot = build_job_dot(&graph).unwrap(); + assert!( + dot.contains("DataSourceExec: (Memory)"), + "expected plan string in dot graph, got: {dot}" + ); + } +} diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index ff886a69f..4aae6ea9d 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -10,11 +10,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::display::format_stage_metrics; +use crate::api::dto_build; use crate::scheduler_server::event::QueryStageSchedulerEvent; -use crate::state::execution_graph::ExecutionStage; use crate::state::execution_graph_dot::ExecutionGraphDot; -use crate::state::execution_stage::TaskInfo; use crate::{api::SchedulerErrorResponse, scheduler_server::SchedulerServer}; use axum::extract::Query; use axum::{ @@ -22,19 +20,14 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; +use ballista_core::BALLISTA_VERSION; use ballista_core::serde::protobuf::job_status::Status; -use ballista_core::serde::protobuf::{ - ExecutorMetric, executor_metric::Metric, task_status, -}; +use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric}; use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; -use ballista_core::utils::get_current_time; -use ballista_core::{BALLISTA_VERSION, JobId}; +use ballista_history::dto::JobResponse; use datafusion::DATAFUSION_VERSION; -use datafusion::physical_plan::display::DisplayableExecutionPlan; -use datafusion::physical_plan::displayable; -use datafusion::physical_plan::metrics::{MetricsSet, Time}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; #[cfg(feature = "graphviz-support")] @@ -44,9 +37,7 @@ use graphviz_rust::{ printer::PrinterContext, }; use http::{StatusCode, header::CONTENT_TYPE}; -use serde::Serialize; use std::sync::Arc; -use std::time::Duration; #[derive(Debug, serde::Serialize)] struct SchedulerStateResponse { @@ -107,27 +98,6 @@ impl ExecutorMetricResponse { } } -#[derive(Debug, serde::Serialize)] -pub struct JobResponse { - pub job_id: JobId, - pub job_name: String, - pub job_status: String, - pub status: String, - pub num_stages: usize, - pub completed_stages: usize, - pub percent_complete: u8, - /// Timestamp when the job started. - pub start_time: u64, - /// Timestamp when the job ended (0 if still running). - pub end_time: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub logical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub physical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, -} - #[derive(Debug, serde::Serialize)] struct CancelJobResponse { pub cancelled: bool, @@ -135,74 +105,6 @@ struct CancelJobResponse { pub reason: Option, } -#[derive(Debug, serde::Serialize)] -pub struct TaskSummary { - /// task id - pub id: usize, - /// Task status - pub status: TaskStatus, - /// partition id - pub partition_id: u32, - /// Scheduler schedule time - pub scheduled_time: u64, - /// Scheduler launch time (ms since epoch) - pub launch_time: u64, - /// The time the Executor start to run the task (ms since epoch) - pub start_exec_time: u64, - /// The time the Executor finish the task (ms since epoch) - pub end_exec_time: u64, - /// total execution time (ms) - pub exec_duration: u64, - /// Scheduler side finish time (ms since epoch) - pub finish_time: u64, - /// Number of input rows - pub input_rows: usize, - /// Number of output rows - pub output_rows: usize, -} - -#[derive(Debug, Clone, Serialize)] -pub enum TaskStatus { - Running, - Successful, - Failed, -} - -impl From<&task_status::Status> for TaskStatus { - fn from(value: &task_status::Status) -> Self { - match value { - task_status::Status::Running(_) => TaskStatus::Running, - task_status::Status::Failed(_) => TaskStatus::Failed, - task_status::Status::Successful(_) => TaskStatus::Successful, - } - } -} - -#[derive(Debug, serde::Serialize)] -pub struct Percentiles { - pub min: u64, - pub p25: u64, - pub median: u64, - pub p75: u64, - pub max: u64, -} - -#[derive(Debug, serde::Serialize)] -pub struct QueryStageSummary { - pub stage_id: String, - pub stage_status: String, - pub input_rows: usize, - pub output_rows: usize, - pub elapsed_compute: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_duration_percentiles: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_input_percentiles: Option, - pub tasks: Vec>, -} - #[derive(Debug, serde::Deserialize, Default)] pub struct JobQueryParams { /// Controls plan format @@ -332,34 +234,7 @@ pub async fn get_jobs< let jobs: Vec = jobs .iter() - .map(|job| { - let (plain_status, job_status) = format_job_status( - &job.status.status, - job_elapsed_ms(job.start_time, job.end_time), - ); - - // calculate progress based on completed stages for now, but we could use completed - // tasks in the future to make this more accurate - let percent_complete = if job.num_stages == 0 { - 0 - } else { - ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 - }; - JobResponse { - job_id: job.job_id.to_owned(), - job_name: job.job_name.to_owned(), - job_status, - status: plain_status, - start_time: job.start_time, - end_time: job.end_time, - num_stages: job.num_stages, - completed_stages: job.completed_stages, - percent_complete, - logical_plan: None, - physical_plan: None, - stage_plan: None, - } - }) + .map(dto_build::build_job_response_from_overview) .collect(); Ok(Json(jobs)) @@ -383,45 +258,14 @@ pub async fn get_job< SchedulerErrorResponse::with_error(StatusCode::INTERNAL_SERVER_ERROR, format!("Error occurred while getting the execution graph for job '{job_id}'")) })? .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; - let stage_plan = format!("{:?}", graph); - let job = graph.as_ref(); - let (plain_status, job_status) = format_job_status( - &job.status().status, - job_elapsed_ms(job.start_time(), job.end_time()), - ); - - let num_stages = job.stage_count(); - let completed_stages = job.completed_stages(); - let percent_complete = - ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; let plan_format = query.plan_format.clone().unwrap_or_default(); - let physical_plan = match plan_format { - PlanFormat::Default | PlanFormat::Metrics => { - DisplayableExecutionPlan::new(job.physical_plan().as_ref()) - .indent(false) - .to_string() - } - PlanFormat::Tree => displayable(job.physical_plan().as_ref()) - .tree_render() - .to_string(), - }; - - Ok(Json(JobResponse { - job_id: job.job_id().to_owned(), - job_name: job.job_name().to_owned(), - job_status, - status: plain_status, - start_time: job.start_time(), - end_time: job.end_time(), - num_stages, - completed_stages, - percent_complete, - logical_plan: job.logical_plan().map(str::to_owned), - physical_plan: Some(physical_plan), - stage_plan: Some(stage_plan), - })) + Ok(Json(dto_build::build_job_response( + &graph, + true, + plan_format, + ))) } pub async fn cancel_job< @@ -491,11 +335,6 @@ pub async fn cancel_job< } } -#[derive(Debug, serde::Serialize)] -pub struct QueryStagesResponse { - pub stages: Vec, -} - pub async fn get_query_stages< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -504,8 +343,6 @@ pub async fn get_query_stages< Path(job_id): Path, query: Query, ) -> Result { - let plan_format = query.plan_format.clone().unwrap_or_default(); - if let Some(graph) = data_server .state .task_manager @@ -519,310 +356,14 @@ pub async fn get_query_stages< ) })? { - let stages = graph - .as_ref() - .stages() - .iter() - .map(|(id, stage)| { - let mut summary = QueryStageSummary { - stage_id: id.to_string(), - stage_status: stage.variant_name().to_string(), - input_rows: 0, - output_rows: 0, - elapsed_compute: None, - tasks: vec![], - task_duration_percentiles: None, - task_input_percentiles: None, - stage_plan: None, - }; - match stage { - ExecutionStage::Running(running_stage) => { - let metrics = running_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(running_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(running_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(running_stage.plan.as_ref(), metrics), - }); - summary.input_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "input_rows")) - .unwrap_or(0); - summary.output_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "output_rows")) - .unwrap_or(0); - summary.elapsed_compute = get_running_stage_time(&running_stage - .task_infos, get_current_time()); - summary.tasks = running_stage - .task_infos - .iter() - .enumerate() - .map(|(partition_id, task_info)| { - task_info.as_ref().map(|info| { - let (input_rows, output_rows) = running_stage - .stage_metrics - .as_deref() - .map(|metrics| { - get_partition_counts(metrics, partition_id) - }) - .unwrap_or((0, 0)); - - let start_exec_time = info.start_exec_time as u64; - let end_exec_time = info.end_exec_time as u64; - - let task_status: TaskStatus = (&info.task_status).into(); - - TaskSummary { - id: info.task_id, - partition_id: partition_id as u32, - scheduled_time: info.scheduled_time as u64, - launch_time: info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: info.finish_time as u64, - input_rows, - output_rows, - status: task_status - } - }) - }) - .collect(); - } - ExecutionStage::Successful(completed_stage) => { - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(completed_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(completed_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(completed_stage.plan.as_ref(), &completed_stage.stage_metrics), - }); - summary.input_rows = get_combined_count( - &completed_stage.stage_metrics, - "input_rows", - ); - summary.output_rows = get_combined_count( - &completed_stage.stage_metrics, - "output_rows", - ); - summary.elapsed_compute = - get_finished_stage_time(&completed_stage.task_infos); - - summary.tasks = completed_stage - .task_infos - .iter() - .enumerate() - .map(|(partition_id, task_info)| { - let (input_rows, output_rows) = get_partition_counts( - &completed_stage.stage_metrics, - partition_id, - ); - - let start_exec_time = task_info.start_exec_time as u64; - let end_exec_time = task_info.end_exec_time as u64; - let task_status = (&task_info.task_status).into(); - Some(TaskSummary { - id: task_info.task_id, - partition_id: partition_id as u32, - scheduled_time: task_info.scheduled_time as u64, - launch_time: task_info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: task_info.finish_time as u64, - input_rows, - output_rows, - status: task_status - }) - }) - .collect(); - } - _ => {} - } - summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); - summary.task_input_percentiles = task_input_percentiles(&summary.tasks); - summary - }) - .collect(); - - Ok(Json(QueryStagesResponse { stages })) + Ok(Json(dto_build::build_query_stages_response( + &graph, &query, + ))) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } } -fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { - let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn task_input_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = tasks - .iter() - .flatten() - .map(|t| t.input_rows as u64) - .collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -fn task_duration_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = - tasks.iter().flatten().map(|t| t.exec_duration).collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -/// Returns elapsed wall time in milliseconds for API formatting. -/// -/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or -/// `end_time` still zero while `start_time` is set) do not panic on subtract. -fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { - end_time.saturating_sub(start_time) -} - -fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { - match status { - Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), - Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), - Some(Status::Failed(error)) => { - ("Failed".to_string(), format!("Failed: {}", error.error)) - } - Some(Status::Successful(completed)) => { - let num_rows = completed - .partition_location - .iter() - .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) - .sum::(); - let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; - let num_partitions = completed.partition_location.len(); - let num_partitions_term = if num_partitions == 1 { - "partition" - } else { - "partitions" - }; - ( - "Completed".to_string(), - format!( - "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", - num_partitions, - num_partitions_term, - num_rows, - num_rows_term, - elapsed_ms - ), - ) - } - _ => ("Invalid".to_string(), "Invalid State".to_string()), - } -} - -fn get_running_stage_time( - task_infos: &[Option], - current_time: u128, -) -> Option { - let min_start = task_infos - .iter() - .flat_map(|t| t.as_ref().map(|t| t.start_exec_time)) - .filter(|t| *t > 0) - .min(); - - match (min_start, current_time) { - (Some(start), end) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { - let min_start = task_infos - .iter() - .map(|t| t.start_exec_time) - .filter(|t| *t > 0) - .min(); - - let max_end = task_infos - .iter() - .map(|t| t.end_exec_time) - .filter(|t| *t > 0) - .max(); - - match (min_start, max_end) { - (Some(start), Some(end)) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) { - let input_rows = get_partition_count(metrics, partition_id, "input_rows"); - let output_rows = get_partition_count(metrics, partition_id, "output_rows"); - (input_rows, output_rows) -} - -fn get_partition_count(metrics: &[MetricsSet], partition_id: usize, name: &str) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - if metric.partition() == Some(partition_id) && metric_value.name() == name - { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - -fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - if metric_value.name() == name { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - pub async fn get_job_dot_graph< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -840,11 +381,10 @@ pub async fn get_job_dot_graph< SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) })? { - ExecutionGraphDot::generate(graph.as_ref()) - .map_err(|e| { - tracing::error!("Error occurred while getting the dot graph for job '{job_id}' reason: {e:?}"); - SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) - }) + dto_build::build_job_dot(&graph).map_err(|e| { + tracing::error!("Error occurred while getting the dot graph for job '{job_id}' reason: {e:?}"); + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) + }) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } @@ -941,126 +481,3 @@ pub async fn get_job_config< .map(|e| Json(e.to_props())) .map_err(|_| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::state::execution_stage::TaskInfo; - use ballista_core::serde::protobuf::task_status; - - fn make_task_info(start: u128, end: u128) -> TaskInfo { - TaskInfo { - task_id: 0, - scheduled_time: 0, - launch_time: 0, - start_exec_time: start, - end_exec_time: end, - finish_time: 0, - task_status: task_status::Status::Running(Default::default()), - } - } - - #[test] - fn test_job_elapsed_saturates_when_end_precedes_start() { - assert_eq!(job_elapsed_ms(900, 100), 0); - } - - // --- get_finished_stage_time --- - - #[test] - fn test_finished_empty_slice_returns_none() { - assert_eq!(get_finished_stage_time(&[]), None); - } - - #[test] - fn test_finished_all_zero_timestamps_returns_none() { - let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - #[test] - fn test_finished_single_task_elapsed() { - // 600 - 100 = 500 ms → "500.00ms" - let tasks = vec![make_task_info(100, 600)]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("500.00ms".to_string()) - ); - } - - #[test] - fn test_finished_picks_earliest_start_and_latest_end() { - // min start = 100, max end = 900 → 800 ms - let tasks = vec![ - make_task_info(100, 500), - make_task_info(200, 900), - make_task_info(300, 700), - ]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("800.00ms".to_string()) - ); - } - - #[test] - fn test_finished_end_before_start_returns_none() { - let tasks = vec![make_task_info(900, 100)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - // --- get_running_stage_time --- - - #[test] - fn test_running_empty_slice_returns_none() { - assert_eq!(get_running_stage_time(&[], 1000), None); - } - - #[test] - fn test_running_all_none_returns_none() { - let tasks: Vec> = vec![None, None]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_future_start_returns_none() { - // start_exec_time beyond current time → elapsed clamped to 0 - let tasks = vec![Some(make_task_info(u128::MAX, 0))]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_past_start_returns_some() { - let now = 4_000; - let start = 1_000; - let tasks = vec![Some(make_task_info(start, 0))]; - assert_eq!( - get_running_stage_time(&tasks, now), - Some("3.00s".to_string()) - ); - } - - #[test] - fn test_running_mixed_some_none_uses_earliest_some() { - let now = 3_000; - let earlier = 1_000; - let later = 2_000; - let tasks = vec![ - None, - Some(make_task_info(later, 0)), - Some(make_task_info(earlier, 0)), - None, - ]; - let result = get_running_stage_time(&tasks, now); - assert_eq!(result, Some("2.00s".to_string())); - } - - #[test] - fn test_job_elapsed_ms_normal() { - assert_eq!(super::job_elapsed_ms(100, 500), 400); - } - - #[test] - fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { - assert_eq!(super::job_elapsed_ms(500, 100), 0); - } -} diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index c02b47fd4..90ba23768 100644 --- a/ballista/scheduler/src/api/mod.rs +++ b/ballista/scheduler/src/api/mod.rs @@ -10,6 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "rest-api")] +pub(crate) mod dto_build; #[cfg(feature = "rest-api")] mod handlers; #[cfg(feature = "rest-api")] diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index eaef4dc5e..c9a622aab 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -406,7 +406,7 @@ fn get_file_scan(scan: &FileScanConfig) -> String { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use crate::planner::DefaultDistributedPlanner; use crate::state::execution_graph::StaticExecutionGraph; use crate::state::execution_graph_dot::ExecutionGraphDot; @@ -578,7 +578,7 @@ filter_expr="] Ok(()) } - async fn test_graph() -> Result { + pub(crate) async fn test_graph() -> Result { let mut config = SessionConfig::new() .with_target_partitions(48) .with_batch_size(4096); From c542809648bdc62f3d95097d7bca985acc0af278 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 18:43:17 -0600 Subject: [PATCH 04/14] feat(scheduler): add event_log_dir config option --- ballista/scheduler/src/config.rs | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 3c6cdb007..807cff1e6 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -103,6 +103,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, + /// Directory to write per-job event logs to (enables the history server). + #[arg(long)] + pub event_log_dir: Option, /// Whether to print thread IDs and names in log files. #[arg( long, @@ -278,6 +281,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, } impl Default for SchedulerConfig { @@ -314,6 +319,7 @@ impl Default for SchedulerConfig { cors_allowed_origins: String::default(), #[cfg(feature = "rest-api")] cors_allowed_methods: String::default(), + event_log_dir: None, } } } @@ -449,6 +455,12 @@ impl SchedulerConfig { self.use_tls = use_tls; self } + + /// Sets the directory to write per-job event logs to. + pub fn with_event_log_dir(mut self, event_log_dir: Option) -> Self { + self.event_log_dir = event_log_dir; + self + } } /// Policy of distributing tasks to available executor slots @@ -546,8 +558,22 @@ impl TryFrom 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, }; Ok(config) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn event_log_dir_defaults_to_none_and_is_settable() { + let config = SchedulerConfig::default(); + assert!(config.event_log_dir.is_none()); + let config = config.with_event_log_dir(Some("/tmp/history".to_string())); + assert_eq!(config.event_log_dir.as_deref(), Some("/tmp/history")); + } +} From da381927853a3b7ea88e1f04b31c3bbaa0e7ea43 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:00:20 -0600 Subject: [PATCH 05/14] feat(scheduler): tee lifecycle events into the event log Emit HistoryEvents from the QueryStageScheduler event loop: JobSubmitted -> JobStart, TaskUpdating -> TaskEnd (one per finished task), and JobFinished/JobRunningFailed -> JobEnd. Event builders live in a new scheduler_server::event_log module and reuse the same dto_build DTOs the REST API serves, so a job's JobEnd event matches its live GET /api/job/{id} response. The EventLogWriter is constructed from SchedulerConfig::event_log_dir and threaded into QueryStageScheduler; emission is best-effort and a no-op when event_log_dir is unset. --- Cargo.lock | 1 + ballista/scheduler/Cargo.toml | 1 + ballista/scheduler/src/api/dto_build.rs | 8 +- ballista/scheduler/src/api/mod.rs | 5 +- .../src/scheduler_server/event_log.rs | 283 ++++++++++++++++++ .../scheduler/src/scheduler_server/mod.rs | 25 ++ .../scheduler_server/query_stage_scheduler.rs | 108 +++++++ 7 files changed, 424 insertions(+), 7 deletions(-) create mode 100644 ballista/scheduler/src/scheduler_server/event_log.rs diff --git a/Cargo.lock b/Cargo.lock index b2973825b..a27f7c30e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1208,6 +1208,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "tempfile", "tokio", "tokio-stream", "tonic", diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index 1ba88fa13..ad1711c79 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -83,6 +83,7 @@ uuid = { workspace = true } [dev-dependencies] rstest = { workspace = true } serde_json = "1" +tempfile = { workspace = true } [build-dependencies] tonic-prost-build = { workspace = true, optional = true } diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs index 8d4d18fd7..179aa47f3 100644 --- a/ballista/scheduler/src/api/dto_build.rs +++ b/ballista/scheduler/src/api/dto_build.rs @@ -298,12 +298,8 @@ pub(crate) fn to_api_task_status(status: &task_status::Status) -> TaskStatus { /// by an executor for one task, e.g. `TaskStatus::metrics`) directly, with no /// partition filtering (a task already corresponds to exactly one partition), /// and additionally sums `elapsed_compute` (nanoseconds) alongside the row -/// counts, matching `ballista_history::event::TaskEndMetrics` for Task 5's -/// timeline events. -/// -/// Not yet called from production code — the event-log writer that will -/// consume it lands in a later task. Silence `dead_code` until then. -#[allow(dead_code)] +/// counts, matching `ballista_history::event::TaskEndMetrics` for the +/// event-log's timeline events (see `scheduler_server::event_log`). pub(crate) fn task_row_counts(metrics: &[OperatorMetricsSet]) -> (u64, u64, u64) { let mut input_rows: u64 = 0; let mut output_rows: u64 = 0; diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index 90ba23768..bdcfc2e52 100644 --- a/ballista/scheduler/src/api/mod.rs +++ b/ballista/scheduler/src/api/mod.rs @@ -13,7 +13,10 @@ #[cfg(feature = "rest-api")] pub(crate) mod dto_build; #[cfg(feature = "rest-api")] -mod handlers; +// `pub(crate)` (rather than private) so `scheduler_server::event_log` can reuse +// `JobQueryParams`/`PlanFormat` to build event-log DTOs identical to the live +// REST responses. +pub(crate) mod handlers; #[cfg(feature = "rest-api")] mod routes; #[cfg(feature = "rest-api")] diff --git a/ballista/scheduler/src/scheduler_server/event_log.rs b/ballista/scheduler/src/scheduler_server/event_log.rs new file mode 100644 index 000000000..9514afcd0 --- /dev/null +++ b/ballista/scheduler/src/scheduler_server/event_log.rs @@ -0,0 +1,283 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Builds `HistoryEvent`s from scheduler event-loop state and appends them to +//! the [`ballista_history::writer::EventLogWriter`]. These are the pure, +//! synchronous event-builder functions; the actual emission -- deciding +//! *when* to call them -- lives at the top of +//! `query_stage_scheduler::QueryStageScheduler::on_receive`. +//! +//! The builders reuse `crate::api::dto_build`, the same DTO builders backing +//! the live REST API, so a job's `JobEnd` event and its `GET /api/job/{id}` +//! response serialize identically for the same graph state. + +use crate::api::dto_build::{ + build_job_dot, build_job_response, build_query_stages_response, task_row_counts, + to_api_task_status, +}; +use crate::api::handlers::{JobQueryParams, PlanFormat}; +use crate::state::execution_graph::ExecutionGraphBox; +use ballista_core::serde::protobuf::{TaskStatus, task_status}; +use ballista_history::dto::TaskStatus as ApiTaskStatus; +use ballista_history::event::{ + HistoryEvent, JobEndStatus, SCHEMA_VERSION, TaskEndMetrics, +}; +use datafusion::physical_plan::display::DisplayableExecutionPlan; + +/// Builds the `JobStart` event for a job that has just been submitted. +pub(crate) fn job_start_event( + graph: &ExecutionGraphBox, + queued_at: u64, + submitted_at: u64, +) -> HistoryEvent { + HistoryEvent::JobStart { + version: SCHEMA_VERSION, + job_id: graph.job_id().to_string(), + job_name: graph.job_name().to_string(), + queued_at, + submitted_at, + logical_plan: graph.logical_plan().map(|p| p.to_string()), + // Same rendering the `get_job` handler uses for `PlanFormat::Default`. + physical_plan: Some( + DisplayableExecutionPlan::new(graph.physical_plan().as_ref()) + .indent(false) + .to_string(), + ), + } +} + +/// Builds one `TaskEnd` event per finished task in `statuses` (skips tasks +/// still `Running`, which are transient/in-flight status updates rather than +/// terminal states). +pub(crate) fn task_end_events( + executor_id: &str, + statuses: &[TaskStatus], +) -> Vec { + statuses + .iter() + .filter(|s| !matches!(s.status, Some(task_status::Status::Running(_)))) + .map(|s| { + let status = s + .status + .as_ref() + .map(to_api_task_status) + .unwrap_or(ApiTaskStatus::Running); + HistoryEvent::TaskEnd { + stage_id: s.stage_id, + partition: s.partition_id, + executor_id: executor_id.to_string(), + status, + launch_time: s.launch_time, + start_exec_time: s.start_exec_time, + end_exec_time: s.end_exec_time, + metrics: task_end_metrics(s), + } + }) + .collect() +} + +/// Sums a task's raw operator metrics into the timeline's `TaskEndMetrics`, +/// using the same extraction the stage-summary REST DTO performs so the +/// timeline and stage views agree. Absent metrics yield zeros. +fn task_end_metrics(status: &TaskStatus) -> TaskEndMetrics { + let (input_rows, output_rows, elapsed_compute_nanos) = + task_row_counts(&status.metrics); + TaskEndMetrics { + input_rows, + output_rows, + elapsed_compute_nanos, + } +} + +/// Builds the `JobEnd` event for a job that has just finished (successfully +/// or not), embedding the same DTOs the REST API would serve for this job at +/// this moment: the job summary, per-stage summaries, the session config, and +/// the DOT-format graph. +pub(crate) fn job_end_event( + graph: &ExecutionGraphBox, + status: JobEndStatus, + queued_at: u64, + completed_at: u64, +) -> HistoryEvent { + // `PlanFormat::Default` -- the same default the `get_job` handler falls + // back to when no `?plan_format=` query param is supplied -- so the + // stored JSON matches what a live `GET /api/job/{id}` would have + // returned at this point. + let job = Box::new(build_job_response(graph, true, PlanFormat::Default)); + let stages = Box::new(build_query_stages_response( + graph, + &JobQueryParams::default(), + )); + let dot = build_job_dot(graph).unwrap_or_default(); + let config = graph.session_config().to_props().into_iter().collect(); + + HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status, + queued_at, + started_at: graph.start_time(), + completed_at, + job, + stages, + config, + dot, + } +} + +#[cfg(test)] +mod tests { + use crate::scheduler_server::event_log::*; + use crate::state::execution_graph::ExecutionGraphBox; + use crate::state::execution_graph_dot::tests::test_graph; + use ballista_core::serde::protobuf::{SuccessfulTask, TaskStatus, task_status}; + use ballista_history::writer::EventLogWriter; + + #[test] + fn task_end_events_map_one_per_finished_task() { + let statuses = vec![TaskStatus { + task_id: 0, + job_id: "job-1".into(), + stage_id: 1, + stage_attempt_num: 0, + partition_id: 3, + launch_time: 100, + start_exec_time: 110, + end_exec_time: 200, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask::default())), + }]; + let events = task_end_events("exec-1", &statuses); + assert_eq!(events.len(), 1); + let line = serde_json::to_string(&events[0]).unwrap(); + assert!(line.contains("\"ev\":\"TaskEnd\"")); + assert!(line.contains("\"partition\":3")); + assert!(line.contains("\"executor_id\":\"exec-1\"")); + } + + // `task_end_events` skips `Running` statuses -- they are transient + // in-flight updates, not the terminal states the timeline records. + #[test] + fn task_end_events_skips_running_tasks() { + let statuses = vec![TaskStatus { + task_id: 0, + job_id: "job-1".into(), + stage_id: 1, + stage_attempt_num: 0, + partition_id: 0, + launch_time: 100, + start_exec_time: 110, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Running(Default::default())), + }]; + assert!(task_end_events("exec-1", &statuses).is_empty()); + } + + /// `job_end_event`'s embedded DTOs (`job`, `stages`, `dot`) are built by + /// the same `dto_build` functions backing the live REST API, so a direct + /// builder test on the serialized event is an adequate substitute for + /// wiring a full `QueryStageScheduler` through `on_receive` here (which + /// would additionally require standing up executor/task-manager state + /// just to reach a finished job). + #[tokio::test] + async fn job_end_event_embeds_job_and_stage_plan_strings() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + let event = job_end_event( + &graph, + JobEndStatus::Succeeded, + /* queued_at */ 1, + /* completed_at */ 2, + ); + let json = serde_json::to_string(&event).unwrap(); + + assert!(json.contains("\"ev\":\"JobEnd\"")); + assert!(json.contains("\"status\":\"Succeeded\"")); + assert!(json.contains("\"job_id\":\"job_id\"")); + // The job's top-level physical plan, embedded via `build_job_response`. + assert!( + json.contains("DataSourceExec: (Memory)"), + "expected job physical plan in JobEnd event, got: {json}" + ); + // The per-stage summaries, embedded via `build_query_stages_response`. + assert!( + json.contains("\"stage_status\":\"Resolved\""), + "expected a resolved stage summary in JobEnd event, got: {json}" + ); + // The DOT graph, embedded via `build_job_dot`. + assert!( + json.contains("digraph"), + "expected a dot graph in JobEnd event, got: {json}" + ); + } + + #[tokio::test] + async fn job_start_event_captures_plan_and_timestamps() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + let event = + job_start_event(&graph, /* queued_at */ 5, /* submitted_at */ 10); + let json = serde_json::to_string(&event).unwrap(); + + assert!(json.contains("\"ev\":\"JobStart\"")); + assert!(json.contains("\"job_id\":\"job_id\"")); + assert!(json.contains("\"job_name\":\"job_name\"")); + assert!(json.contains("\"queued_at\":5")); + assert!(json.contains("\"submitted_at\":10")); + // `job_start_event` renders the graph's pre-staging physical plan + // (`graph.physical_plan()`), unlike `job_end_event`'s embedded job + // DTO which shows the plan reconstructed from stages -- assert on the + // join at its root rather than the leaf scan string. + assert!( + json.contains("HashJoinExec"), + "expected job physical plan in JobStart event, got: {json}" + ); + } + + /// End-to-end: drive the real `EventLogWriter` over a temp dir with a + /// `JobStart` followed by a `JobEnd`, then read the `.eventlog` file back + /// and assert on line order and content -- proving the builders here + /// produce events the writer can actually persist and that a consumer + /// (the eventual history server) can read back as ordered JSONL. + #[tokio::test] + async fn writer_round_trip_orders_job_start_before_job_end() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + let job_id = graph.job_id().to_string(); + + let dir = tempfile::tempdir().unwrap(); + let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); + + writer.append(&job_id, job_start_event(&graph, 1, 2)); + writer.append( + &job_id, + job_end_event(&graph, JobEndStatus::Succeeded, 1, 20), + ); + writer.flush_job(&job_id).await; + + let path = dir.path().join(format!("{job_id}.eventlog")); + let contents = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + + assert_eq!(lines.len(), 2); + assert!(lines[0].contains("\"ev\":\"JobStart\"")); + assert!(lines[1].contains("\"ev\":\"JobEnd\"")); + assert!(lines[1].contains("DataSourceExec: (Memory)")); + } +} diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 17b955a98..77d9b38f6 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -54,6 +54,13 @@ pub mod externalscaler { /// Events for the scheduler event loop. pub mod event; +/// Builds `HistoryEvent`s emitted from the event loop into the event log. +/// +/// Depends on `crate::api::dto_build`, which is only compiled with the +/// `rest-api` feature (the default), so this module -- and the event-log +/// wiring into `QueryStageScheduler` below -- is gated the same way. +#[cfg(feature = "rest-api")] +mod event_log; #[cfg(feature = "keda-scaler")] mod external_scaler; mod grpc; @@ -102,10 +109,19 @@ impl SchedulerServer SchedulerServer>, metrics_collector: Arc, config: Arc, + #[cfg(feature = "rest-api")] + event_log: Option, } impl QueryStageScheduler { @@ -51,11 +55,16 @@ impl QueryStageSchedul state: Arc>, metrics_collector: Arc, config: Arc, + #[cfg(feature = "rest-api")] event_log: Option< + ballista_history::writer::EventLogWriter, + >, ) -> Self { Self { state, metrics_collector, config, + #[cfg(feature = "rest-api")] + event_log, } } #[cfg(feature = "rest-api")] @@ -64,6 +73,25 @@ impl QueryStageSchedul } } +/// Groups task status updates by job id, so a single `TaskUpdating` batch +/// (which can span multiple jobs) can be appended to each job's own event log. +#[cfg(feature = "rest-api")] +fn group_by_job( + statuses: &[ballista_core::serde::protobuf::TaskStatus], +) -> std::collections::HashMap> { + let mut by_job: std::collections::HashMap< + String, + Vec, + > = std::collections::HashMap::new(); + for status in statuses { + by_job + .entry(status.job_id.clone()) + .or_default() + .push(status.clone()); + } + by_job +} + #[async_trait::async_trait] impl EventAction for QueryStageScheduler @@ -82,6 +110,86 @@ impl tx_event: &mpsc::Sender, _rx_event: &mpsc::Receiver, ) -> Result<()> { + #[cfg(feature = "rest-api")] + if let Some(log) = &self.event_log { + match &event { + QueryStageSchedulerEvent::JobSubmitted { + job_id, + queued_at, + submitted_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append( + job_id.as_str(), + event_log::job_start_event(&graph, *queued_at, *submitted_at), + ); + } + } + QueryStageSchedulerEvent::TaskUpdating(executor_id, statuses) => { + for (job_id, group) in group_by_job(statuses) { + for ev in event_log::task_end_events(executor_id, &group) { + log.append(&job_id, ev); + } + } + } + QueryStageSchedulerEvent::JobFinished { + job_id, + queued_at, + completed_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append( + job_id.as_str(), + event_log::job_end_event( + &graph, + ballista_history::event::JobEndStatus::Succeeded, + *queued_at, + *completed_at, + ), + ); + } + log.flush_job(job_id.as_str()).await; + } + QueryStageSchedulerEvent::JobRunningFailed { + job_id, + fail_message, + queued_at, + failed_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append( + job_id.as_str(), + event_log::job_end_event( + &graph, + ballista_history::event::JobEndStatus::Failed( + fail_message.clone(), + ), + *queued_at, + *failed_at, + ), + ); + } + log.flush_job(job_id.as_str()).await; + } + _ => {} + } + } + let mut time_recorder = None; if self.config.scheduler_event_expected_processing_duration > 0 { time_recorder = Some((Instant::now(), event.clone())); From 1eda5177eee0d9574825fab8429083b790ddb3c3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:05:29 -0600 Subject: [PATCH 06/14] feat(history): add event-log reader producing ReplayedJob --- ballista/history/src/reader.rs | 112 +++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/ballista/history/src/reader.rs b/ballista/history/src/reader.rs index b248758bc..d4b60f210 100644 --- a/ballista/history/src/reader.rs +++ b/ballista/history/src/reader.rs @@ -14,3 +14,115 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. + +//! Reads a completed `.eventlog` into the DTO bundle the history server +//! serves. A file is "completed" once it contains a `JobEnd` record. + +use crate::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use crate::event::HistoryEvent; +use std::io::BufRead; +use std::path::Path; + +#[derive(Debug, Clone)] +pub struct ReplayedJob { + pub job: JobResponse, + pub stages: QueryStagesResponse, + pub config: JobConfig, + pub dot: String, +} + +pub fn read_completed_job(path: &Path) -> std::io::Result> { + let file = std::fs::File::open(path)?; + let reader = std::io::BufReader::new(file); + for line in reader.lines() { + let line = line?; + if line.is_empty() { + continue; + } + // Only JobEnd carries the served payload; other lines are the timeline + // and are ignored here. Unknown/garbled lines are skipped, not fatal. + if let Ok(HistoryEvent::JobEnd { + job, + stages, + config, + dot, + .. + }) = serde_json::from_str::(&line) + { + return Ok(Some(ReplayedJob { + job: *job, + stages: *stages, + config, + dot, + })); + } + } + Ok(None) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::io::Write; + + fn job_end_line() -> String { + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 1, + started_at: 2, + completed_at: 3, + job: Box::new(JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 1, + completed_stages: 1, + percent_complete: 100, + start_time: 2, + end_time: 3, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage".into()), + }), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: Default::default(), + dot: "digraph {}".into(), + }; + serde_json::to_string(&event).unwrap() + } + + #[test] + fn reads_job_end_and_ignores_unknown_timeline_lines() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("job-1.eventlog"); + let mut f = std::fs::File::create(&path).unwrap(); + // Unknown/other event lines before the JobEnd must be tolerated. + writeln!(f, r#"{{"ev":"StageStart","stage_id":1,"partitions":4}}"#).unwrap(); + writeln!(f, "{}", job_end_line()).unwrap(); + drop(f); + + let replayed = read_completed_job(&path).unwrap().expect("completed"); + assert_eq!(replayed.job.job_id, "job-1"); + assert_eq!( + replayed.job.physical_plan.as_deref(), + Some("ProjectionExec") + ); + assert_eq!(replayed.dot, "digraph {}"); + } + + #[test] + fn returns_none_when_no_job_end() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("job-2.eventlog"); + std::fs::write( + &path, + "{\"ev\":\"StageStart\",\"stage_id\":1,\"partitions\":4}\n", + ) + .unwrap(); + assert!(read_completed_job(&path).unwrap().is_none()); + } +} From af575ef0ad4017db5e074b6988a7ddc806aa3977 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:11:00 -0600 Subject: [PATCH 07/14] feat(scheduler): add ballista-history-server binary and router --- Cargo.lock | 1 + ballista/scheduler/Cargo.toml | 6 + ballista/scheduler/src/bin/history_server.rs | 94 +++++++++ ballista/scheduler/src/history/mod.rs | 193 +++++++++++++++++++ ballista/scheduler/src/lib.rs | 3 + 5 files changed, 297 insertions(+) create mode 100644 ballista/scheduler/src/bin/history_server.rs create mode 100644 ballista/scheduler/src/history/mod.rs diff --git a/Cargo.lock b/Cargo.lock index a27f7c30e..a432314ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1214,6 +1214,7 @@ dependencies = [ "tonic", "tonic-prost", "tonic-prost-build", + "tower", "tower-http 0.7.0", "tracing", "tracing-appender", diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index ad1711c79..3ddbad532 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -32,6 +32,11 @@ name = "ballista-scheduler" path = "src/bin/main.rs" required-features = ["build-binary"] +[[bin]] +name = "ballista-history-server" +path = "src/bin/history_server.rs" +required-features = ["build-binary", "rest-api"] + [features] build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"] default = ["build-binary", "rest-api"] @@ -84,6 +89,7 @@ uuid = { workspace = true } rstest = { workspace = true } serde_json = "1" tempfile = { workspace = true } +tower = "0.5" [build-dependencies] tonic-prost-build = { workspace = true, optional = true } diff --git a/ballista/scheduler/src/bin/history_server.rs b/ballista/scheduler/src/bin/history_server.rs new file mode 100644 index 000000000..631c7f5cc --- /dev/null +++ b/ballista/scheduler/src/bin/history_server.rs @@ -0,0 +1,94 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Standalone Ballista history server binary: loads completed event logs from +//! a directory and serves the same `/api/*` responses the live scheduler does, +//! so the existing TUI can connect to it unchanged. + +use ballista_core::error::{BallistaError, Result}; +use ballista_scheduler::history::{HistoryStore, history_router}; +use clap::Parser; +use std::env; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use tracing_subscriber::EnvFilter; + +#[derive(Debug, clap::Parser)] +#[command( + name = "ballista-history-server", + version, + about = "Ballista history server" +)] +struct Args { + /// Directory containing per-job event logs. + #[arg(long)] + event_log_dir: PathBuf, + /// Host to bind the HTTP server to. + #[arg(long, default_value = "0.0.0.0")] + bind_host: String, + /// Port to bind the HTTP server to. + #[arg(long, default_value_t = 50060)] + bind_port: u16, +} + +fn main() -> Result<()> { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .enable_time() + .build() + .map_err(BallistaError::IoError)?; + + runtime.block_on(inner()) +} + +async fn inner() -> Result<()> { + let rust_log = env::var(EnvFilter::DEFAULT_ENV); + let log_filter = EnvFilter::new(rust_log.unwrap_or_else(|_| "info".to_string())); + tracing_subscriber::fmt() + .with_ansi(false) + .with_writer(std::io::stdout) + .with_env_filter(log_filter) + .init(); + + let args = Args::parse(); + + let store = Arc::new(HistoryStore::load(&args.event_log_dir)?); + tracing::info!( + "Loaded {} completed job(s) from {}", + store.jobs.len(), + args.event_log_dir.display() + ); + let app = history_router(store); + + let addr: SocketAddr = format!("{}:{}", args.bind_host, args.bind_port) + .parse() + .map_err(|e: std::net::AddrParseError| { + BallistaError::Configuration(e.to_string()) + })?; + + let listener = tokio::net::TcpListener::bind(&addr) + .await + .map_err(BallistaError::IoError)?; + tracing::info!("History server listening on http://{addr}"); + + axum::serve(listener, app.into_make_service()) + .await + .map_err(BallistaError::IoError)?; + + Ok(()) +} diff --git a/ballista/scheduler/src/history/mod.rs b/ballista/scheduler/src/history/mod.rs new file mode 100644 index 000000000..b026307ab --- /dev/null +++ b/ballista/scheduler/src/history/mod.rs @@ -0,0 +1,193 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Standalone history server: loads completed event logs and serves the same +//! `/api/*` responses the live scheduler does, from stored DTOs. + +use axum::{ + Json, Router, + extract::{Path as AxumPath, State}, + routing::get, +}; +use ballista_history::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use ballista_history::reader::{ReplayedJob, read_completed_job}; +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; + +/// In-memory store of completed jobs, loaded once at startup from a directory +/// of `.eventlog` files. +#[derive(Default)] +pub struct HistoryStore { + /// Completed jobs keyed by job id. + pub jobs: HashMap, +} + +impl HistoryStore { + /// Load every completed job found under `dir`. Missing directories yield + /// an empty store rather than an error. + pub fn load(dir: &Path) -> std::io::Result { + let mut jobs = HashMap::new(); + if dir.exists() { + for entry in std::fs::read_dir(dir)? { + let path = entry?.path(); + if path.extension().and_then(|e| e.to_str()) == Some("eventlog") + && let Some(replayed) = read_completed_job(&path)? + { + jobs.insert(replayed.job.job_id.clone(), replayed); + } + } + } + Ok(HistoryStore { jobs }) + } +} + +/// Build the axum router serving `/api/*` from a loaded [`HistoryStore`]. +pub fn history_router(store: Arc) -> Router { + Router::new() + .route("/api/jobs", get(get_jobs)) + .route("/api/job/{job_id}", get(get_job)) + .route("/api/job/{job_id}/stages", get(get_stages)) + .route("/api/job/{job_id}/config", get(get_config)) + .route("/api/job/{job_id}/dot", get(get_dot)) + .route("/api/executors", get(get_executors_empty)) + .with_state(store) +} + +async fn get_jobs(State(store): State>) -> Json> { + // The live list endpoint omits plans; null them for byte-identical output. + let mut jobs: Vec = store + .jobs + .values() + .map(|j| { + let mut r = j.job.clone(); + r.logical_plan = None; + r.physical_plan = None; + r.stage_plan = None; + r + }) + .collect(); + jobs.sort_by(|a, b| a.job_id.cmp(&b.job_id)); + Json(jobs) +} + +async fn get_job( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Json> { + Json(store.jobs.get(&job_id).map(|j| j.job.clone())) +} + +async fn get_stages( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Json> { + Json(store.jobs.get(&job_id).map(|j| j.stages.clone())) +} + +async fn get_config( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Json> { + Json(store.jobs.get(&job_id).map(|j| j.config.clone())) +} + +async fn get_dot( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> String { + store + .jobs + .get(&job_id) + .map(|j| j.dot.clone()) + .unwrap_or_default() +} + +async fn get_executors_empty() -> Json> { + Json(vec![]) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use tower::ServiceExt; // oneshot + + fn store_with_one_job() -> Arc { + let mut jobs = HashMap::new(); + jobs.insert( + "job-1".to_string(), + ReplayedJob { + job: JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 1, + completed_stages: 1, + percent_complete: 100, + start_time: 2, + end_time: 3, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage".into()), + }, + stages: QueryStagesResponse { stages: vec![] }, + config: Default::default(), + dot: "digraph {}".into(), + }, + ); + Arc::new(HistoryStore { jobs }) + } + + #[tokio::test] + async fn jobs_endpoint_nulls_plan_fields() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/jobs") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + assert!(body.contains("\"job_id\":\"job-1\"")); + assert!(!body.contains("physical_plan")); // nulled + skip_serializing_if + } + + #[tokio::test] + async fn stages_endpoint_returns_stored_dto() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/job/job-1/stages") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } +} diff --git a/ballista/scheduler/src/lib.rs b/ballista/scheduler/src/lib.rs index 533711ecb..8338a6537 100644 --- a/ballista/scheduler/src/lib.rs +++ b/ballista/scheduler/src/lib.rs @@ -25,6 +25,9 @@ pub mod cluster; pub mod config; /// Display utilities for execution plans and state. pub mod display; +/// Standalone history server: serves `/api/*` from stored event logs. +#[cfg(feature = "rest-api")] +pub mod history; /// Metrics collection and reporting. pub mod metrics; /// Physical query plan optimizers. From dcbd2bbca1d05d1e02c794f3b5cc680761cc742c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:22:29 -0600 Subject: [PATCH 08/14] fix(scheduler): harden history-server router against corrupt logs and missing jobs - HistoryStore::load now skips unreadable/corrupt .eventlog files with a warning instead of failing the whole load, so one bad log can't hide every other completed job. - Job endpoints (/api/job/{id}, /stages, /config, /dot) now return 404 for unknown job ids instead of 200 with a null body, matching the live scheduler's behavior. - Add /api/state with a static payload matching the shape the TUI deserializes at startup. - Strengthen router tests: assert stage-response body contents, cover the corrupt-eventlog skip path, and cover the 404 behavior. --- ballista/scheduler/Cargo.toml | 1 + ballista/scheduler/src/history/mod.rs | 245 ++++++++++++++++++++++---- 2 files changed, 212 insertions(+), 34 deletions(-) diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index 3ddbad532..5a90b28f1 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -76,6 +76,7 @@ prost = { workspace = true } prost-types = { workspace = true } rand = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = "1" tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true, features = ["net"] } tonic = { workspace = true, features = ["router"] } diff --git a/ballista/scheduler/src/history/mod.rs b/ballista/scheduler/src/history/mod.rs index b026307ab..8f6dbccd0 100644 --- a/ballista/scheduler/src/history/mod.rs +++ b/ballista/scheduler/src/history/mod.rs @@ -18,13 +18,17 @@ //! Standalone history server: loads completed event logs and serves the same //! `/api/*` responses the live scheduler does, from stored DTOs. +use crate::api::SchedulerErrorResponse; use axum::{ Json, Router, extract::{Path as AxumPath, State}, routing::get, }; +use ballista_core::BALLISTA_VERSION; use ballista_history::dto::{JobConfig, JobResponse, QueryStagesResponse}; use ballista_history::reader::{ReplayedJob, read_completed_job}; +use datafusion::DATAFUSION_VERSION; +use http::StatusCode; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; @@ -40,15 +44,30 @@ pub struct HistoryStore { impl HistoryStore { /// Load every completed job found under `dir`. Missing directories yield /// an empty store rather than an error. + /// + /// A single unreadable/corrupt `.eventlog` file (e.g. truncated by a + /// crash mid-write) is logged and skipped rather than failing the whole + /// load — one bad log must not hide every other completed job. Only a + /// failure to read the directory itself is propagated. pub fn load(dir: &Path) -> std::io::Result { let mut jobs = HashMap::new(); if dir.exists() { for entry in std::fs::read_dir(dir)? { let path = entry?.path(); - if path.extension().and_then(|e| e.to_str()) == Some("eventlog") - && let Some(replayed) = read_completed_job(&path)? - { - jobs.insert(replayed.job.job_id.clone(), replayed); + if path.extension().and_then(|e| e.to_str()) != Some("eventlog") { + continue; + } + match read_completed_job(&path) { + Ok(Some(replayed)) => { + jobs.insert(replayed.job.job_id.clone(), replayed); + } + Ok(None) => {} + Err(err) => { + tracing::warn!( + "skipping unreadable event log {}: {err}", + path.display() + ); + } } } } @@ -65,6 +84,7 @@ pub fn history_router(store: Arc) -> Router { .route("/api/job/{job_id}/config", get(get_config)) .route("/api/job/{job_id}/dot", get(get_dot)) .route("/api/executors", get(get_executors_empty)) + .route("/api/state", get(get_state)) .with_state(store) } @@ -88,70 +108,122 @@ async fn get_jobs(State(store): State>) -> Json>, AxumPath(job_id): AxumPath, -) -> Json> { - Json(store.jobs.get(&job_id).map(|j| j.job.clone())) +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.job.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } async fn get_stages( State(store): State>, AxumPath(job_id): AxumPath, -) -> Json> { - Json(store.jobs.get(&job_id).map(|j| j.stages.clone())) +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.stages.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } async fn get_config( State(store): State>, AxumPath(job_id): AxumPath, -) -> Json> { - Json(store.jobs.get(&job_id).map(|j| j.config.clone())) +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.config.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } async fn get_dot( State(store): State>, AxumPath(job_id): AxumPath, -) -> String { +) -> Result { store .jobs .get(&job_id) .map(|j| j.dot.clone()) - .unwrap_or_default() + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } async fn get_executors_empty() -> Json> { Json(vec![]) } +/// Static `/api/state` payload. The history server has no live scheduler +/// process behind it, so every field that would normally reflect runtime +/// state (uptime, feature flags, scheduling policy) is a fixed placeholder. +/// Field names/types match the live `/api/state` response +/// (`SchedulerStateResponse` in `api/handlers.rs`) and what the TUI +/// deserializes into (`ballista-cli/src/tui/domain/mod.rs::SchedulerState`), +/// so the TUI's startup call succeeds instead of erroring out. +async fn get_state() -> Json { + Json(serde_json::json!({ + "started": 0, + "version": BALLISTA_VERSION, + "datafusion_version": DATAFUSION_VERSION, + "substrait_support": false, + "keda_support": false, + "prometheus_support": false, + "graphviz_support": false, + "spark_support": false, + "scheduling_policy": "history-server", + })) +} + #[cfg(test)] mod tests { use super::*; use axum::body::Body; use axum::http::{Request, StatusCode}; + use ballista_history::dto::QueryStageSummary; + use ballista_history::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::io::Write; + use tempfile::tempdir; use tower::ServiceExt; // oneshot + const STAGE_ID_MARKER: &str = "stage-42"; + + fn sample_replayed_job(job_id: &str) -> ReplayedJob { + ReplayedJob { + job: JobResponse { + job_id: job_id.into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 1, + completed_stages: 1, + percent_complete: 100, + start_time: 2, + end_time: 3, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage".into()), + }, + stages: QueryStagesResponse { + stages: vec![QueryStageSummary { + stage_id: STAGE_ID_MARKER.into(), + stage_status: "Completed".into(), + input_rows: 10, + output_rows: 5, + elapsed_compute: Some("1ms".into()), + stage_plan: None, + task_duration_percentiles: None, + task_input_percentiles: None, + tasks: vec![], + }], + }, + config: Default::default(), + dot: "digraph {}".into(), + } + } + fn store_with_one_job() -> Arc { let mut jobs = HashMap::new(); - jobs.insert( - "job-1".to_string(), - ReplayedJob { - job: JobResponse { - job_id: "job-1".into(), - job_name: "q1".into(), - job_status: "COMPLETED".into(), - status: "Successful".into(), - num_stages: 1, - completed_stages: 1, - percent_complete: 100, - start_time: 2, - end_time: 3, - logical_plan: Some("Projection".into()), - physical_plan: Some("ProjectionExec".into()), - stage_plan: Some("stage".into()), - }, - stages: QueryStagesResponse { stages: vec![] }, - config: Default::default(), - dot: "digraph {}".into(), - }, - ); + jobs.insert("job-1".to_string(), sample_replayed_job("job-1")); Arc::new(HistoryStore { jobs }) } @@ -189,5 +261,110 @@ mod tests { .await .unwrap(); assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body: QueryStagesResponse = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(body.stages.len(), 1); + assert_eq!(body.stages[0].stage_id, STAGE_ID_MARKER); + assert_eq!(body.stages[0].input_rows, 10); + assert_eq!(body.stages[0].output_rows, 5); + } + + #[tokio::test] + async fn missing_job_returns_404_on_job_and_stages() { + let app = history_router(store_with_one_job()); + + let resp = app + .clone() + .oneshot( + Request::builder() + .uri("/api/job/does-not-exist") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let resp = app + .oneshot( + Request::builder() + .uri("/api/job/does-not-exist/stages") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn state_endpoint_returns_static_payload() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/state") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let value: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + for field in [ + "started", + "version", + "datafusion_version", + "substrait_support", + "keda_support", + "prometheus_support", + "graphviz_support", + "spark_support", + "scheduling_policy", + ] { + assert!(value.get(field).is_some(), "missing field: {field}"); + } + } + + fn write_job_end_log(path: &Path, job_id: &str) { + let replayed = sample_replayed_job(job_id); + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 0, + started_at: 2, + completed_at: 3, + job: Box::new(replayed.job), + stages: Box::new(replayed.stages), + config: replayed.config, + dot: replayed.dot, + }; + let line = serde_json::to_string(&event).unwrap(); + std::fs::write(path, format!("{line}\n")).unwrap(); + } + + #[test] + fn load_skips_corrupt_eventlog_and_keeps_good_one() { + let dir = tempdir().unwrap(); + + // A good, readable event log. + write_job_end_log(&dir.path().join("job-good.eventlog"), "job-good"); + + // A corrupt file: invalid UTF-8, as if a crash truncated a write + // mid-multibyte-character. + let mut corrupt = + std::fs::File::create(dir.path().join("job-bad.eventlog")).unwrap(); + corrupt.write_all(&[0xff, 0xfe, 0xfd]).unwrap(); + drop(corrupt); + + let store = HistoryStore::load(dir.path()).unwrap(); + assert_eq!(store.jobs.len(), 1); + assert!(store.jobs.contains_key("job-good")); + assert!(!store.jobs.contains_key("job-bad")); } } From ed0065c191549b11a01bf9a6165f93233ca57f09 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:32:54 -0600 Subject: [PATCH 09/14] test(history): verify history store serves byte-identical JSON to live scheduler Add an end-to-end DTO parity test that builds live JobResponse/ QueryStagesResponse DTOs, writes the same JobEnd event through the real EventLogWriter, reloads it via HistoryStore::load, and asserts the serialized JSON matches exactly -- proving the history server would serve the same data a running scheduler produced. --- .../src/scheduler_server/event_log.rs | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/ballista/scheduler/src/scheduler_server/event_log.rs b/ballista/scheduler/src/scheduler_server/event_log.rs index 9514afcd0..dbb867be4 100644 --- a/ballista/scheduler/src/scheduler_server/event_log.rs +++ b/ballista/scheduler/src/scheduler_server/event_log.rs @@ -280,4 +280,56 @@ mod tests { assert!(lines[1].contains("\"ev\":\"JobEnd\"")); assert!(lines[1].contains("DataSourceExec: (Memory)")); } + + /// End-to-end DTO parity: the history server, replaying a real event log + /// through `EventLogWriter` + `HistoryStore::load`, must serve + /// byte-identical JSON to what the live scheduler would return for the + /// same job via `dto_build::{build_job_response, build_query_stages_response}`. + /// + /// This is `job_end_event`'s only round-trip coverage through the real + /// writer + `HistoryStore` (as opposed to the direct builder-output + /// assertions above), so it lives here as a crate-internal `#[cfg(test)]` + /// module rather than an integration test under `tests/`: `dto_build`'s + /// builders, `job_end_event` itself, and + /// `execution_graph_dot::tests::test_graph` are all `pub(crate)` / + /// cfg(test)-gated and therefore unreachable from a separate `tests/` + /// integration-test crate, which only sees the scheduler crate's public API. + #[tokio::test] + async fn history_store_serves_byte_identical_json_to_live_scheduler() { + use crate::api::handlers::JobQueryParams; + use crate::history::HistoryStore; + + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + let job_id = graph.job_id().to_string(); + + // The live DTOs a running scheduler would serve for this job right now, + // via the same builders the REST handlers call. + let live_job = build_job_response(&graph, true, PlanFormat::Default); + let live_stages = build_query_stages_response(&graph, &JobQueryParams::default()); + + // Produce the same `JobEnd` event the scheduler emits on completion, + // write it through the real async writer, then load it back via the + // history server's own `HistoryStore::load` -- exercising the full + // write -> read -> serve path, not just the event builder. + let event = job_end_event(&graph, JobEndStatus::Succeeded, 1, 2); + let dir = tempfile::tempdir().unwrap(); + let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); + writer.append(&job_id, event); + writer.flush_job(&job_id).await; + + let store = HistoryStore::load(dir.path()).unwrap(); + let replayed = store.jobs.get(&job_id).expect("job should be replayed"); + + // The core fidelity guarantee: the history server would return + // byte-identical JSON to what the live scheduler returned for this job. + assert_eq!( + serde_json::to_string(&replayed.job).unwrap(), + serde_json::to_string(&live_job).unwrap(), + ); + assert_eq!( + serde_json::to_string(&replayed.stages).unwrap(), + serde_json::to_string(&live_stages).unwrap(), + ); + } } From 9b545d71020add6f1474cb4263e5154f2a436418 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 19:45:57 -0600 Subject: [PATCH 10/14] fix(history): guarantee terminal JobEnd is not dropped on a full channel EventLogWriter::append enqueues via non-blocking try_send, which drops the event when the channel is saturated. Using it for the terminal JobEnd event meant a completed job's log could end up with JobStart and TaskEnds but no JobEnd, making read_completed_job return None and the job silently disappear from the history server. Add append_final, which awaits channel capacity instead of dropping, and finish_job, which flushes and then closes the per-job file handle (also fixing fd accumulation). Wire both into the JobFinished and JobRunningFailed on_receive arms in place of append + flush_job. --- ballista/history/src/writer.rs | 126 +++++++++++++++++- .../scheduler_server/query_stage_scheduler.rs | 14 +- 2 files changed, 133 insertions(+), 7 deletions(-) diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs index b7b79132c..61ae0a825 100644 --- a/ballista/history/src/writer.rs +++ b/ballista/history/src/writer.rs @@ -34,6 +34,10 @@ enum WriterMsg { job_id: String, done: oneshot::Sender<()>, }, + Finish { + job_id: String, + done: oneshot::Sender<()>, + }, } #[derive(Clone)] @@ -77,6 +81,46 @@ impl EventLogWriter { let _ = wait.await; } } + + /// Enqueue a terminal event (e.g. `JobEnd`) for `job_id`. Unlike `append`, this + /// awaits channel capacity instead of dropping the event when the channel is + /// full, so the terminal record is never silently lost. Still best-effort at + /// the process boundary: if the channel is closed (background task gone) this + /// logs and returns rather than panicking. + pub async fn append_final(&self, job_id: &str, event: HistoryEvent) { + let msg = WriterMsg::Event { + job_id: job_id.to_string(), + event: Box::new(event), + }; + if self.tx.send(msg).await.is_err() { + eprintln!( + "event-log writer: failed to enqueue terminal event for {job_id} (channel closed)" + ); + } + } + + /// Flush and close the per-job file handle for `job_id`. Must be called after + /// the terminal event has been enqueued (e.g. via `append_final`) so it is + /// ordered after it on the single-consumer FIFO channel. Best effort: if the + /// channel is closed this logs and returns. + pub async fn finish_job(&self, job_id: &str) { + let (done, wait) = oneshot::channel(); + if self + .tx + .send(WriterMsg::Finish { + job_id: job_id.to_string(), + done, + }) + .await + .is_ok() + { + let _ = wait.await; + } else { + eprintln!( + "event-log writer: failed to enqueue finish for {job_id} (channel closed)" + ); + } + } } async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { @@ -110,6 +154,13 @@ async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { } let _ = done.send(()); } + WriterMsg::Finish { job_id, done } => { + if let Some(mut file) = handles.remove(&job_id) { + let _ = file.flush().await; + // Dropping `file` here closes the fd. + } + let _ = done.send(()); + } } } } @@ -142,7 +193,80 @@ async fn open_for<'a>( #[cfg(test)] mod tests { use super::*; - use crate::event::{HistoryEvent, SCHEMA_VERSION}; + use crate::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::collections::BTreeMap; + + #[tokio::test] + async fn terminal_job_end_is_not_dropped_on_a_saturated_channel() { + let dir = tempfile::tempdir().unwrap(); + // Tiny buffer so the non-blocking `append` path would readily drop events + // under load; `append_final` must still guarantee delivery. + let writer = EventLogWriter::new(dir.path().to_path_buf(), 1); + + writer.append( + "job-1", + HistoryEvent::JobStart { + version: SCHEMA_VERSION, + job_id: "job-1".into(), + job_name: "q1".into(), + queued_at: 1, + submitted_at: 2, + logical_plan: None, + physical_plan: None, + }, + ); + for stage_id in 0..10 { + writer.append( + "job-1", + HistoryEvent::StageStart { + stage_id, + partitions: 4, + }, + ); + } + + let job = JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 2, + completed_stages: 2, + percent_complete: 100, + start_time: 10, + end_time: 20, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage plan".into()), + }; + let job_end = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 1, + started_at: 2, + completed_at: 20, + job: Box::new(job), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: BTreeMap::new(), + dot: "digraph {}".into(), + }; + writer.append_final("job-1", job_end).await; + writer.finish_job("job-1").await; + + let path = dir.path().join("job-1.eventlog"); + let contents = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + assert!( + lines.iter().any(|l| l.contains("\"ev\":\"JobEnd\"")), + "expected a JobEnd line in the event log, got: {contents}" + ); + assert_eq!( + lines.last().map(|l| l.contains("\"ev\":\"JobEnd\"")), + Some(true), + "JobEnd should be the last line written" + ); + } #[tokio::test] async fn append_writes_one_jsonl_line_per_event() { diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 62a8efb3b..320343f1b 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -148,7 +148,7 @@ impl .get_job_execution_graph(job_id) .await { - log.append( + log.append_final( job_id.as_str(), event_log::job_end_event( &graph, @@ -156,9 +156,10 @@ impl *queued_at, *completed_at, ), - ); + ) + .await; } - log.flush_job(job_id.as_str()).await; + log.finish_job(job_id.as_str()).await; } QueryStageSchedulerEvent::JobRunningFailed { job_id, @@ -172,7 +173,7 @@ impl .get_job_execution_graph(job_id) .await { - log.append( + log.append_final( job_id.as_str(), event_log::job_end_event( &graph, @@ -182,9 +183,10 @@ impl *queued_at, *failed_at, ), - ); + ) + .await; } - log.flush_job(job_id.as_str()).await; + log.finish_job(job_id.as_str()).await; } _ => {} } From d919d1b679438215c4410cee4c7abdd11f097974 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 22:26:19 -0600 Subject: [PATCH 11/14] chore(history): add Apache license header to Cargo.toml --- ballista/history/Cargo.toml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml index 26751269b..873f4b44a 100644 --- a/ballista/history/Cargo.toml +++ b/ballista/history/Cargo.toml @@ -1,3 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + [package] name = "ballista-history" description = "Event-log schema, writer, and reader for Ballista's history server" From 7852a91666349494d87c39650fc36ac1c8aaeba5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 22:47:02 -0600 Subject: [PATCH 12/14] chore(scheduler): make ballista-history an optional dep gated by rest-api All ballista_history usage in the scheduler is behind the rest-api feature (DTO builders, event-log wiring, history module). Gating the dependency on rest-api keeps it out of the graph for consumers that build the scheduler with default-features = false (e.g. pyballista), and out of non-rest-api builds. --- ballista/scheduler/Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index 5a90b28f1..c366772b6 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -46,7 +46,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 = [] +rest-api = ["dep:ballista-history"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -56,7 +56,7 @@ async-trait = { workspace = true } axum = "0.8.9" tower-http = { version = "0.7", features = ["cors"] } ballista-core = { path = "../core", version = "53.0.0" } -ballista-history = { path = "../history", version = "53.0.0" } +ballista-history = { path = "../history", version = "53.0.0", optional = true } clap = { workspace = true, optional = true } dashmap = { workspace = true } datafusion = { workspace = true } From 09686e39043a7ff4eb87edfb354a8d1fdb0543c4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 23:05:09 -0600 Subject: [PATCH 13/14] chore(scheduler): gate serde_json behind rest-api feature serde_json is only used in production by the rest-api-gated history module (all other uses are in tests, covered by the dev-dependency). Gating it keeps it out of the dependency graph for default-features = false consumers such as pyballista, so python/Cargo.lock stays in sync. --- ballista/scheduler/Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index c366772b6..2b4421818 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -46,7 +46,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-history"] +rest-api = ["dep:ballista-history", "dep:serde_json"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -76,7 +76,7 @@ prost = { workspace = true } prost-types = { workspace = true } rand = { workspace = true } serde = { workspace = true, features = ["derive"] } -serde_json = "1" +serde_json = { version = "1", optional = true } tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true, features = ["net"] } tonic = { workspace = true, features = ["router"] } From 970d72ec5602b34bd99cf73b69fe9c5067f2af4c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 5 Aug 2026 07:24:16 -0600 Subject: [PATCH 14/14] fix(scheduler): address event-log review feedback - record cancelled jobs: `JobCancel` is terminal for a job, so without a `JobEnd` the log was left without a terminal record (invisible to the history server) and its file handle stayed open for the life of the process. The tee runs before the cancel is applied to the graph, so the event stamps the cancelled status onto the DTO rather than recording the job as perpetually `Running`. - guard `percent_complete` against a zero stage count, matching `build_job_response_from_overview`. - skip status-less task updates in `task_end_events` instead of recording them as `Running`, which contradicted both the filter and the doc. - report (rather than swallow) a failed or missing execution-graph lookup on the event-log path, and document why `finish_job` still runs. - use fixed-width stage/partition ids throughout the on-disk event schema. - load the history store before the tokio runtime starts, so its blocking directory walk does not park a runtime worker. --- ballista/history/src/event.rs | 14 ++- ballista/scheduler/src/api/dto_build.rs | 7 +- ballista/scheduler/src/bin/history_server.rs | 24 +++-- .../src/scheduler_server/event_log.rs | 95 ++++++++++++++++--- .../scheduler_server/query_stage_scheduler.rs | 74 +++++++++++---- 5 files changed, 168 insertions(+), 46 deletions(-) diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs index 8db1c05f4..59b43b470 100644 --- a/ballista/history/src/event.rs +++ b/ballista/history/src/event.rs @@ -25,10 +25,12 @@ use serde::{Deserialize, Serialize}; /// Current on-disk schema version, stamped on `JobStart`/`JobEnd`. pub const SCHEMA_VERSION: u32 = 1; +/// How a job reached its terminal state. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum JobEndStatus { Succeeded, Failed(String), + Cancelled, } /// Metrics captured per finished task on the incremental timeline. @@ -39,6 +41,12 @@ pub struct TaskEndMetrics { pub elapsed_compute_nanos: u64, } +/// A single record in a job's event log. +/// +/// Stage and partition identifiers are fixed-width (`u32`) throughout rather +/// than `usize`: this is a durable, cross-machine format, so a log written by a +/// 64-bit scheduler must mean the same thing to any reader. Callers holding the +/// scheduler's own `usize` stage ids cast on the way in. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "ev")] pub enum HistoryEvent { @@ -52,11 +60,11 @@ pub enum HistoryEvent { physical_plan: Option, }, StageStart { - stage_id: usize, - partitions: usize, + stage_id: u32, + partitions: u32, }, StageEnd { - stage_id: usize, + stage_id: u32, status: String, }, TaskEnd { diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs index 179aa47f3..d9305cafa 100644 --- a/ballista/scheduler/src/api/dto_build.rs +++ b/ballista/scheduler/src/api/dto_build.rs @@ -54,8 +54,11 @@ pub(crate) fn build_job_response( let num_stages = job.stage_count(); let completed_stages = job.completed_stages(); - let percent_complete = - ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; + let percent_complete = if num_stages == 0 { + 0 + } else { + ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8 + }; let (logical_plan, physical_plan, stage_plan) = if with_plans { let physical_plan = match plan_format { diff --git a/ballista/scheduler/src/bin/history_server.rs b/ballista/scheduler/src/bin/history_server.rs index 631c7f5cc..212df5817 100644 --- a/ballista/scheduler/src/bin/history_server.rs +++ b/ballista/scheduler/src/bin/history_server.rs @@ -47,16 +47,6 @@ struct Args { } fn main() -> Result<()> { - let runtime = tokio::runtime::Builder::new_multi_thread() - .enable_io() - .enable_time() - .build() - .map_err(BallistaError::IoError)?; - - runtime.block_on(inner()) -} - -async fn inner() -> Result<()> { let rust_log = env::var(EnvFilter::DEFAULT_ENV); let log_filter = EnvFilter::new(rust_log.unwrap_or_else(|_| "info".to_string())); tracing_subscriber::fmt() @@ -67,12 +57,26 @@ async fn inner() -> Result<()> { let args = Args::parse(); + // `HistoryStore::load` walks the log directory with blocking file I/O, and + // how long it takes scales with the number of stored jobs. Run it here, + // before the runtime exists, rather than parking a runtime worker on it. let store = Arc::new(HistoryStore::load(&args.event_log_dir)?); tracing::info!( "Loaded {} completed job(s) from {}", store.jobs.len(), args.event_log_dir.display() ); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .enable_time() + .build() + .map_err(BallistaError::IoError)?; + + runtime.block_on(serve(args, store)) +} + +async fn serve(args: Args, store: Arc) -> Result<()> { let app = history_router(store); let addr: SocketAddr = format!("{}:{}", args.bind_host, args.bind_port) diff --git a/ballista/scheduler/src/scheduler_server/event_log.rs b/ballista/scheduler/src/scheduler_server/event_log.rs index dbb867be4..8b924ebe9 100644 --- a/ballista/scheduler/src/scheduler_server/event_log.rs +++ b/ballista/scheduler/src/scheduler_server/event_log.rs @@ -32,7 +32,6 @@ use crate::api::dto_build::{ use crate::api::handlers::{JobQueryParams, PlanFormat}; use crate::state::execution_graph::ExecutionGraphBox; use ballista_core::serde::protobuf::{TaskStatus, task_status}; -use ballista_history::dto::TaskStatus as ApiTaskStatus; use ballista_history::event::{ HistoryEvent, JobEndStatus, SCHEMA_VERSION, TaskEndMetrics, }; @@ -60,23 +59,23 @@ pub(crate) fn job_start_event( } } -/// Builds one `TaskEnd` event per finished task in `statuses` (skips tasks -/// still `Running`, which are transient/in-flight status updates rather than -/// terminal states). +/// Builds one `TaskEnd` event per finished task in `statuses`. +/// +/// Only terminal statuses are recorded: `Running` updates are transient +/// in-flight reports, and a status-less update carries no outcome to record at +/// all, so both are skipped. pub(crate) fn task_end_events( executor_id: &str, statuses: &[TaskStatus], ) -> Vec { statuses .iter() - .filter(|s| !matches!(s.status, Some(task_status::Status::Running(_)))) - .map(|s| { - let status = s - .status - .as_ref() - .map(to_api_task_status) - .unwrap_or(ApiTaskStatus::Running); - HistoryEvent::TaskEnd { + .filter_map(|s| { + let status = match s.status.as_ref()? { + task_status::Status::Running(_) => return None, + terminal => to_api_task_status(terminal), + }; + Some(HistoryEvent::TaskEnd { stage_id: s.stage_id, partition: s.partition_id, executor_id: executor_id.to_string(), @@ -85,7 +84,7 @@ pub(crate) fn task_end_events( start_exec_time: s.start_exec_time, end_exec_time: s.end_exec_time, metrics: task_end_metrics(s), - } + }) }) .collect() } @@ -138,6 +137,37 @@ pub(crate) fn job_end_event( } } +/// Builds the terminal `JobEnd` event for a cancelled job. +/// +/// Cancellation is recorded from the event-log tee, which runs *before* the +/// scheduler applies the cancel to the graph, so the graph still reports the +/// job as running here. The embedded job DTO's status fields are therefore +/// overwritten, otherwise the history server would show a cancelled job as +/// perpetually `Running`. +/// +/// `queued_at` is not carried on `JobCancel`, so it is recorded as 0. +pub(crate) fn job_cancel_event( + graph: &ExecutionGraphBox, + cancelled_at: u64, +) -> HistoryEvent { + let mut event = job_end_event( + graph, + JobEndStatus::Cancelled, + /* queued_at */ 0, + cancelled_at, + ); + if let HistoryEvent::JobEnd { job, .. } = &mut event { + job.job_status = CANCELLED_STATUS.to_string(); + job.status = CANCELLED_STATUS.to_string(); + } + event +} + +/// The status string a cancelled job is recorded (and served) with. The live +/// REST API has no equivalent -- a cancelled job's graph is dropped before it +/// can be observed -- so this is history-only. +const CANCELLED_STATUS: &str = "Cancelled"; + #[cfg(test)] mod tests { use crate::scheduler_server::event_log::*; @@ -187,6 +217,45 @@ mod tests { assert!(task_end_events("exec-1", &statuses).is_empty()); } + // A status update with no status at all reports no outcome, so there is + // nothing terminal to record for it either. + #[test] + fn task_end_events_skips_tasks_with_no_status() { + let statuses = vec![TaskStatus { + task_id: 0, + job_id: "job-1".into(), + stage_id: 1, + stage_attempt_num: 0, + partition_id: 0, + launch_time: 100, + start_exec_time: 110, + end_exec_time: 0, + metrics: vec![], + status: None, + }]; + assert!(task_end_events("exec-1", &statuses).is_empty()); + } + + /// A cancelled job is recorded from a graph that has not yet been marked + /// failed, so the event must carry the cancelled status itself rather than + /// the graph's in-flight `Running`. + #[tokio::test] + async fn job_cancel_event_reports_cancelled_not_running() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + let event = job_cancel_event(&graph, /* cancelled_at */ 42); + let json = serde_json::to_string(&event).unwrap(); + + assert!(json.contains("\"ev\":\"JobEnd\"")); + assert!(json.contains("\"status\":\"Cancelled\"")); + assert!(json.contains("\"job_status\":\"Cancelled\"")); + assert!( + !json.contains("\"Running\""), + "cancelled job should not be recorded as running, got: {json}" + ); + } + /// `job_end_event`'s embedded DTOs (`job`, `stages`, `dot`) are built by /// the same `dto_build` functions backing the live REST API, so a direct /// builder test on the serialized event is an adequate substitute for diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 320343f1b..1afd5bc4b 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -71,6 +71,35 @@ impl QueryStageSchedul pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector { self.metrics_collector.as_ref() } + + /// Fetches a job's execution graph for the event log, reporting (rather + /// than swallowing) the cases where it is unavailable. Returning `None` + /// costs the job its event, not its execution -- event logging is never + /// allowed to fail scheduling. + #[cfg(feature = "rest-api")] + async fn event_log_graph( + &self, + job_id: &ballista_core::JobId, + ) -> Option { + match self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + Ok(Some(graph)) => Some(graph), + Ok(None) => { + warn!("event log: no execution graph for job {job_id}, skipping event"); + None + } + Err(e) => { + warn!( + "event log: failed to read execution graph for job {job_id}: {e:?}" + ); + None + } + } + } } /// Groups task status updates by job id, so a single `TaskUpdating` batch @@ -118,12 +147,7 @@ impl queued_at, submitted_at, } => { - if let Ok(Some(graph)) = self - .state - .task_manager - .get_job_execution_graph(job_id) - .await - { + if let Some(graph) = self.event_log_graph(job_id).await { log.append( job_id.as_str(), event_log::job_start_event(&graph, *queued_at, *submitted_at), @@ -137,17 +161,17 @@ impl } } } + // The three terminal events below each close out the job's log. + // `finish_job` runs even when no `JobEnd` could be built, so a + // job that cannot be recorded still releases its file handle; + // such a log has no `JobEnd` line and the history server skips + // it rather than serving a half-written job. QueryStageSchedulerEvent::JobFinished { job_id, queued_at, completed_at, } => { - if let Ok(Some(graph)) = self - .state - .task_manager - .get_job_execution_graph(job_id) - .await - { + if let Some(graph) = self.event_log_graph(job_id).await { log.append_final( job_id.as_str(), event_log::job_end_event( @@ -167,12 +191,7 @@ impl queued_at, failed_at, } => { - if let Ok(Some(graph)) = self - .state - .task_manager - .get_job_execution_graph(job_id) - .await - { + if let Some(graph) = self.event_log_graph(job_id).await { log.append_final( job_id.as_str(), event_log::job_end_event( @@ -188,6 +207,25 @@ impl } log.finish_job(job_id.as_str()).await; } + QueryStageSchedulerEvent::JobCancel(job_id) => { + // Cancellation is terminal: the handler below drops the + // graph, so this is the last point at which the job can be + // recorded. Without this the log would never receive a + // `JobEnd` -- leaving the cancelled job invisible to the + // history server and its file handle open for the life of + // the process. + if let Some(graph) = self.event_log_graph(job_id).await { + log.append_final( + job_id.as_str(), + event_log::job_cancel_event(&graph, timestamp_millis()), + ) + .await; + } + log.finish_job(job_id.as_str()).await; + } + // `JobPlanningFailed` is deliberately absent: it is only posted + // when `submit_job` fails, i.e. instead of `JobSubmitted`, so + // the job has neither an execution graph nor an open log file. _ => {} } }