From 7ad9a93decb50b2713cdeed8d08df4a19f43d98d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 09:52:06 -0600 Subject: [PATCH 1/7] refactor(scheduler): extract REST DTOs into ballista-history and factor out dto_build Move the scheduler's REST response types into a new leaf crate, `ballista-history`, and pull the graph-to-DTO construction out of the axum handlers into a pure `api::dto_build` module. Behavior preserving: the same DTOs are produced from the same state, so live REST responses are byte-identical. The existing handler tests cover this, and the helper unit tests move alongside the functions they test. This is the first step toward a history server that replays completed jobs and serves the same `/api/*` responses without a live scheduler. Splitting the DTOs into a serde-only crate lets that server build the identical wire types without depending on the scheduler's live execution graph, and moving construction out of the handlers means it can run against state that did not come from a handler request. `JobResponse::job_id` becomes a `String` rather than `ballista_core::JobId` so the new crate stays serde-only. `JobId` is `#[serde(transparent)]` over `String`, so the JSON is unchanged. --- Cargo.lock | 9 + Cargo.toml | 1 + ballista/history/Cargo.toml | 30 ++ ballista/history/src/dto.rs | 114 ++++ ballista/history/src/lib.rs | 24 + ballista/scheduler/Cargo.toml | 4 +- ballista/scheduler/src/api/dto_build.rs | 581 ++++++++++++++++++++ ballista/scheduler/src/api/handlers.rs | 683 +----------------------- ballista/scheduler/src/api/mod.rs | 2 + 9 files changed, 771 insertions(+), 677 deletions(-) create mode 100644 ballista/history/Cargo.toml create mode 100644 ballista/history/src/dto.rs create mode 100644 ballista/history/src/lib.rs create mode 100644 ballista/scheduler/src/api/dto_build.rs diff --git a/Cargo.lock b/Cargo.lock index dac81a2616..498f742dfe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1200,6 +1200,13 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "54.0.0" +dependencies = [ + "serde", +] + [[package]] name = "ballista-scheduler" version = "54.0.0" @@ -1208,6 +1215,7 @@ dependencies = [ "async-trait", "axum", "ballista-core", + "ballista-history", "clap 4.6.3", "dashmap", "datafusion", @@ -1230,6 +1238,7 @@ dependencies = [ "regex", "rstest", "serde", + "serde_json", "tokio", "tokio-stream", "tonic", diff --git a/Cargo.toml b/Cargo.toml index 5a6b966dd8..50ffac073d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ members = [ "ballista/client", "ballista/core", "ballista/executor", + "ballista/history", "ballista/scheduler", "benchmarks", "chaos-testing", diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml new file mode 100644 index 0000000000..a4e56d709c --- /dev/null +++ b/ballista/history/Cargo.toml @@ -0,0 +1,30 @@ +# 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 = "Shared REST response types for the Ballista scheduler and history server" +license = "Apache-2.0" +version = "54.0.0" +homepage = "https://datafusion.apache.org/ballista/" +repository = "https://github.com/apache/datafusion-ballista" +authors = ["Apache DataFusion "] +edition = { workspace = true } +rust-version = { workspace = true } + +[dependencies] +serde = { workspace = true, features = ["derive"] } diff --git a/ballista/history/src/dto.rs b/ballista/history/src/dto.rs new file mode 100644 index 0000000000..74bbb7fe89 --- /dev/null +++ b/ballista/history/src/dto.rs @@ -0,0 +1,114 @@ +// 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 { + /// A `String` rather than `ballista_core::JobId` so this crate stays a + /// serde-only leaf. `JobId` is `#[serde(transparent)]` over `String`, so + /// the serialized JSON is unchanged. + 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 { reason: String, error: String }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskSummary { + /// task id + pub id: usize, + /// Task status + pub status: TaskStatus, + /// Global partition ids covered by this task. For a single-partition + /// task this is a one-element list — JSON-compatible with the old + /// scalar `partition_id` for callers that read `partition_id[0]`, and + /// honestly plural for multi-partition tasks. + pub partition_id: Vec, + /// 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, 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/lib.rs b/ballista/history/src/lib.rs new file mode 100644 index 0000000000..affddf8a53 --- /dev/null +++ b/ballista/history/src/lib.rs @@ -0,0 +1,24 @@ +// 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. + +//! Shared types for Ballista's scheduler REST API. +//! +//! These response types live outside `ballista-scheduler` so that a future +//! history server can serialize byte-identical JSON from stored state without +//! depending on the scheduler's live execution graph. + +pub mod dto; diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index b8b7127733..17ce991bd2 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -41,7 +41,7 @@ disable-stage-plan-cache = [] graphviz-support = ["dep:graphviz-rust"] keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"] prometheus-metrics = ["prometheus", "once_cell"] -rest-api = [] +rest-api = ["dep:ballista-history"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -50,6 +50,7 @@ arrow-flight = { workspace = true } async-trait = { workspace = true } axum = "0.8.9" ballista-core = { path = "../core", version = "54.0.0" } +ballista-history = { path = "../history", version = "54.0.0", optional = true } clap = { workspace = true, optional = true } dashmap = { workspace = true } datafusion = { workspace = true } @@ -88,6 +89,7 @@ path = "tests/tpch_plan_stability/main.rs" datafusion-functions-aggregate-common = { workspace = true } regex = "1" 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 0000000000..ce1ec45315 --- /dev/null +++ b/ballista/scheduler/src/api/dto_build.rs @@ -0,0 +1,581 @@ +// 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 REST response DTOs from scheduler execution state. +//! +//! These functions are the only place the scheduler's `ExecutionGraph` / +//! `JobOverview` shapes are translated into the wire types in +//! [`ballista_history::dto`]. Keeping the translation here (rather than inline +//! in the axum handlers) means the same DTOs can be produced from state that +//! did not come from a live handler request. +//! +//! Everything in this module is pure: state in, DTO out, no I/O. + +use crate::api::handlers::{JobQueryParams, PlanFormat}; +use crate::display::format_stage_metrics; +use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage}; +use crate::state::execution_stage::TaskInfo; +use crate::state::task_manager::JobOverview; +use ballista_core::serde::protobuf::failed_task::FailedReason::{ + ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, +}; +use ballista_core::serde::protobuf::job_status::Status; +use ballista_core::serde::protobuf::{FailedTask, task_status}; +use ballista_core::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::time::Duration; + +/// Build the `JobResponse` list entry served by `GET /api/jobs`. +/// +/// Plan fields are always `None` here: the job list is built from +/// [`JobOverview`] summaries, which do not carry plans. +pub fn job_overview_to_response(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, + } +} + +/// Build the `JobResponse` served by `GET /api/job/{job_id}`, including the +/// rendered logical, physical, and stage plans. +pub fn graph_to_job_response( + graph: &ExecutionGraphBox, + query: &JobQueryParams, +) -> JobResponse { + 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(), + }; + + 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: job.logical_plan().map(str::to_owned), + physical_plan: Some(physical_plan), + stage_plan: Some(stage_plan), + } +} + +/// Build the per-stage summaries served by `GET /api/job/{job_id}/stages`. +pub fn graph_to_query_stages( + graph: &ExecutionGraphBox, + query: &JobQueryParams, +) -> QueryStagesResponse { + let plan_format = query.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(render_stage_plan( + running_stage.plan.as_ref(), + metrics, + &plan_format, + )); + summary.input_rows = get_combined_count(metrics, "input_rows"); + summary.output_rows = get_combined_count(metrics, "output_rows"); + summary.elapsed_compute = get_running_stage_time( + &running_stage.task_infos, + get_current_time(), + ); + summary.tasks = task_summaries(&running_stage.task_infos, metrics); + } + ExecutionStage::Successful(completed_stage) => { + let metrics = completed_stage.stage_metrics.as_slice(); + summary.stage_plan = Some(render_stage_plan( + completed_stage.plan.as_ref(), + metrics, + &plan_format, + )); + summary.input_rows = get_combined_count(metrics, "input_rows"); + summary.output_rows = get_combined_count(metrics, "output_rows"); + summary.elapsed_compute = + get_finished_stage_time(&completed_stage.task_infos); + summary.tasks = task_summaries(&completed_stage.task_infos, metrics); + } + ExecutionStage::Failed(failed_stage) => { + let metrics = failed_stage.stage_metrics.as_deref().unwrap_or(&[]); + summary.stage_plan = Some(render_stage_plan( + failed_stage.plan.as_ref(), + metrics, + &plan_format, + )); + summary.input_rows = get_combined_count(metrics, "input_rows"); + summary.output_rows = get_combined_count(metrics, "output_rows"); + summary.elapsed_compute = + get_finished_stage_time(&failed_stage.task_infos); + summary.tasks = task_summaries(&failed_stage.task_infos, metrics); + } + _ => {} + } + + summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); + summary.task_input_percentiles = task_input_percentiles(&summary.tasks); + summary + }) + .collect(); + + QueryStagesResponse { stages } +} + +/// Render one stage's plan in the requested format. `Metrics` overlays the +/// stage's aggregated metrics onto the plan; the other formats ignore them. +fn render_stage_plan( + plan: &dyn datafusion::physical_plan::ExecutionPlan, + metrics: &[MetricsSet], + plan_format: &PlanFormat, +) -> String { + match plan_format { + PlanFormat::Default => displayable(plan).indent(false).to_string(), + PlanFormat::Tree => displayable(plan).tree_render().to_string(), + PlanFormat::Metrics => format_stage_metrics(plan, metrics), + } +} + +/// Build one [`TaskSummary`] per task in a stage. Row counts are summed over +/// the global partitions each task owns. +/// +/// The `Option` wrapper on each entry is part of the wire format and is always +/// `Some` here; it predates multi-partition tasks, when a stage's task list was +/// indexed by partition and could be sparse. +fn task_summaries( + task_infos: &[TaskInfo], + metrics: &[MetricsSet], +) -> Vec> { + task_infos + .iter() + .map(|info| { + let (input_rows, output_rows) = + get_partition_counts(metrics, &info.global_input_partition_ids); + + let start_exec_time = info.start_exec_time as u64; + let end_exec_time = info.end_exec_time as u64; + + Some(TaskSummary { + id: info.task_id, + partition_id: info + .global_input_partition_ids + .iter() + .map(|&p| p as u32) + .collect(), + 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_to_dto(&info.task_status), + }) + }) + .collect() +} + +/// Map a protobuf task status onto the wire enum. +/// +/// A free function rather than a `From` impl: both types are foreign to this +/// crate now that [`TaskStatus`] lives in `ballista-history`. +pub fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { + match value { + task_status::Status::Running(_) => TaskStatus::Running, + task_status::Status::Failed(failed) => TaskStatus::Failed { + reason: failed_reason(failed), + error: failed.error.clone(), + }, + task_status::Status::Successful(_) => TaskStatus::Successful, + } +} + +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 percentiles_of(mut values: Vec) -> Option { + if values.is_empty() { + return None; + } + + values.sort_unstable(); + + Some(Percentiles { + min: values[0], + p25: percentile_duration(&values, 25.0), + median: percentile_duration(&values, 50.0), + p75: percentile_duration(&values, 75.0), + max: *values.last().unwrap(), + }) +} + +fn task_input_percentiles(tasks: &[Option]) -> Option { + percentiles_of( + tasks + .iter() + .flatten() + .map(|t| t.input_rows as u64) + .collect(), + ) +} + +fn task_duration_percentiles(tasks: &[Option]) -> Option { + percentiles_of(tasks.iter().flatten().map(|t| t.exec_duration).collect()) +} + +/// 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: &[TaskInfo], current_time: u128) -> Option { + let min_start = task_infos + .iter() + .map(|t| t.start_exec_time) + .filter(|t| *t > 0) + .min(); + + match (min_start, current_time) { + (Some(start), end) if end >= start => Some(format_millis(end - start)), + _ => 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 => Some(format_millis(end - start)), + _ => None, + } +} + +/// Format a millisecond duration the way DataFusion renders elapsed-time +/// metrics, so stage timings match plan output. +fn format_millis(millis: u128) -> String { + let time = Time::new(); + time.add_duration(Duration::from_millis(millis as u64)); + time.to_string() +} + +fn failed_reason(failed: &FailedTask) -> String { + match &failed.failed_reason { + Some(ExecutionError(_)) => "ExecutionError", + Some(FetchPartitionError(_)) => "FetchPartitionError", + Some(IoError(_)) => "IoError", + Some(ExecutorLost(_)) => "ExecutorLost", + Some(ResultLost(_)) => "ResultLost", + Some(TaskKilled(_)) => "TaskKilled", + None => "Failed", + } + .to_string() +} + +/// Sum a task's `input_rows` / `output_rows` across the global partitions the +/// task owns. Metrics are keyed by global partition id — for single-partition +/// tasks `partitions` is a one-element slice; for multi-partition tasks it is +/// the task's `global_input_partition_ids`. +fn get_partition_counts(metrics: &[MetricsSet], partitions: &[usize]) -> (usize, usize) { + let input_rows = get_partition_count(metrics, partitions, "input_rows"); + let output_rows = get_partition_count(metrics, partitions, "output_rows"); + (input_rows, output_rows) +} + +fn get_partition_count( + metrics: &[MetricsSet], + partitions: &[usize], + name: &str, +) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + let owned_by_task = metric + .partition() + .map(|p| partitions.contains(&p)) + .unwrap_or(false); + if owned_by_task && 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_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()), + global_input_partition_ids: vec![], + vcores_consumed: 0, + } + } + + #[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_zero_start_returns_none() { + let tasks: Vec = vec![make_task_info(0, 0), make_task_info(0, 0)]; + 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![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![make_task_info(start, 0)]; + assert_eq!( + get_running_stage_time(&tasks, now), + Some("3.00s".to_string()) + ); + } + + #[test] + fn test_running_mixed_zero_start_uses_earliest_nonzero() { + let now = 3_000; + let earlier = 1_000; + let later = 2_000; + let tasks = vec![ + make_task_info(0, 0), + make_task_info(later, 0), + make_task_info(earlier, 0), + make_task_info(0, 0), + ]; + 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/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 5409f07801..f236dd37a6 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::response::Redirect; @@ -23,22 +21,14 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; -use ballista_core::serde::protobuf::failed_task::FailedReason::{ - ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, -}; +use ballista_core::BALLISTA_VERSION; use ballista_core::serde::protobuf::job_status::Status; -use ballista_core::serde::protobuf::{ - ExecutorMetric, FailedTask, 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")] @@ -48,10 +38,8 @@ use graphviz_rust::{ printer::PrinterContext, }; use http::{HeaderMap, StatusCode, header::CONTENT_TYPE}; -use serde::Serialize; use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; #[derive(Debug, serde::Serialize)] struct SchedulerStateResponse { @@ -112,27 +100,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, @@ -140,80 +107,6 @@ struct CancelJobResponse { pub reason: Option, } -#[derive(Debug, serde::Serialize)] -pub struct TaskSummary { - /// task id - pub id: usize, - /// Task status - pub status: TaskStatus, - /// Global partition ids covered by this task. For a single-partition - /// task this is a one-element list — JSON-compatible with the old - /// scalar `partition_id` for callers that read `partition_id[0]`, and - /// honestly plural for multi-partition tasks. - pub partition_id: Vec, - /// 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 { reason: String, error: String }, -} - -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(failed) => TaskStatus::Failed { - reason: failed_reason(failed), - error: failed.error.clone(), - }, - 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 @@ -395,34 +288,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::job_overview_to_response) .collect(); Ok(Json(jobs)) @@ -446,45 +312,8 @@ 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::graph_to_job_response(&graph, &query))) } pub async fn cancel_job< @@ -554,11 +383,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, @@ -567,8 +391,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 @@ -582,383 +404,12 @@ 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() - .map(|info| { - let (input_rows, output_rows) = running_stage - .stage_metrics - .as_deref() - .map(|metrics| { - get_partition_counts( - metrics, - &info.global_input_partition_ids, - ) - }) - .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(); - - Some(TaskSummary { - id: info.task_id, - partition_id: info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - 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() - .map(|task_info| { - let (input_rows, output_rows) = get_partition_counts( - &completed_stage.stage_metrics, - &task_info.global_input_partition_ids, - ); - - 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: task_info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - 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(); - } - ExecutionStage::Failed(failed_stage) => { - let metrics = failed_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(failed_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(failed_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(failed_stage.plan.as_ref(), metrics), - }); - summary.input_rows = get_combined_count(metrics, "input_rows"); - summary.output_rows = get_combined_count(metrics, "output_rows"); - summary.elapsed_compute = - get_finished_stage_time(&failed_stage.task_infos); - - summary.tasks = failed_stage - .task_infos - .iter() - .map(|info| { - let (input_rows, output_rows) = get_partition_counts( - metrics, - &info.global_input_partition_ids, - ); - - 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(); - - Some(TaskSummary { - id: info.task_id, - partition_id: info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - 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(); - } - _ => {} - } - 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::graph_to_query_stages(&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: &[TaskInfo], current_time: u128) -> Option { - let min_start = task_infos - .iter() - .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 failed_reason(failed: &FailedTask) -> String { - match &failed.failed_reason { - Some(ExecutionError(_)) => "ExecutionError", - Some(FetchPartitionError(_)) => "FetchPartitionError", - Some(IoError(_)) => "IoError", - Some(ExecutorLost(_)) => "ExecutorLost", - Some(ResultLost(_)) => "ResultLost", - Some(TaskKilled(_)) => "TaskKilled", - None => "Failed", - } - .to_string() -} - -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, - } -} - -/// Sum a task's `input_rows` / `output_rows` across the global partitions the -/// task owns. Metrics are keyed by global partition id — for single-partition -/// tasks `partitions` is a one-element slice; for multi-partition tasks it is -/// the task's `global_input_partition_ids`. -fn get_partition_counts(metrics: &[MetricsSet], partitions: &[usize]) -> (usize, usize) { - let input_rows = get_partition_count(metrics, partitions, "input_rows"); - let output_rows = get_partition_count(metrics, partitions, "output_rows"); - (input_rows, output_rows) -} - -fn get_partition_count( - metrics: &[MetricsSet], - partitions: &[usize], - name: &str, -) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - let owned_by_task = metric - .partition() - .map(|p| partitions.contains(&p)) - .unwrap_or(false); - if owned_by_task && 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, @@ -1081,126 +532,6 @@ pub async fn get_job_config< #[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()), - global_input_partition_ids: vec![], - vcores_consumed: 0, - } - } - - #[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_zero_start_returns_none() { - let tasks: Vec = vec![make_task_info(0, 0), make_task_info(0, 0)]; - 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![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![make_task_info(start, 0)]; - assert_eq!( - get_running_stage_time(&tasks, now), - Some("3.00s".to_string()) - ); - } - - #[test] - fn test_running_mixed_zero_start_uses_earliest_nonzero() { - let now = 3_000; - let earlier = 1_000; - let later = 2_000; - let tasks = vec![ - make_task_info(0, 0), - make_task_info(later, 0), - make_task_info(earlier, 0), - make_task_info(0, 0), - ]; - 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); - } mod get_webtui { use super::*; diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index 81d7177cfa..527e7610b8 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")] +mod dto_build; #[cfg(feature = "rest-api")] mod handlers; mod health; From 8c55ee35fd6b435bad713971f43d87009fe3827c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 10:05:30 -0600 Subject: [PATCH 2/7] refactor(scheduler): tighten the DTO extraction Follow-up cleanups on the extraction: - Collapse the three near-identical ExecutionStage arms in graph_to_query_stages into one destructuring match, dropping the mutable placeholder-zero summary. - Take PlanFormat by value instead of &JobQueryParams, and move PlanFormat into ballista-history. It is part of the wire contract, and the pure builder no longer imports an axum query-param type back out of the handler module. - Inject `now` into graph_to_query_stages rather than reading the clock, so replaying a stored log renders stable elapsed times. - Share percent_complete and min_start_time; use displayable() instead of the longhand DisplayableExecutionPlan::new(). - Drop the dead JobConfig alias and the unused serde_json dev-dependency, make task_status_to_dto private, and remove a duplicated test. - Enable #![warn(missing_docs)] on ballista-history and document the types, matching the other Ballista crates. - Register ballista-history with the release tooling: version bump script, publish order, and crate dependency graph. --- ballista/history/src/dto.rs | 62 ++++++++- ballista/history/src/lib.rs | 14 ++ ballista/scheduler/src/api/dto_build.rs | 169 ++++++++++-------------- ballista/scheduler/src/api/handlers.rs | 26 ++-- dev/release/README.md | 2 + dev/release/crate-deps.dot | 2 + dev/update_ballista_versions.py | 2 + 7 files changed, 157 insertions(+), 120 deletions(-) diff --git a/ballista/history/src/dto.rs b/ballista/history/src/dto.rs index 74bbb7fe89..272730af84 100644 --- a/ballista/history/src/dto.rs +++ b/ballista/history/src/dto.rs @@ -19,37 +19,58 @@ //! server, so both serialize byte-identical JSON. use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; +/// Summary of one job, served by `GET /api/jobs` and `GET /api/job/{job_id}`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct JobResponse { /// A `String` rather than `ballista_core::JobId` so this crate stays a /// serde-only leaf. `JobId` is `#[serde(transparent)]` over `String`, so /// the serialized JSON is unchanged. pub job_id: String, + /// Human-readable job name. pub job_name: String, + /// Verbose status, including the completion or failure detail. pub job_status: String, + /// Plain status word: `Queued`, `Running`, `Completed`, `Failed`, or `Invalid`. pub status: String, + /// Total number of stages in the job. pub num_stages: usize, + /// Number of stages that finished successfully. pub completed_stages: usize, + /// Progress as a percentage of completed stages. 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, + /// Rendered logical plan. Absent in the job list. #[serde(skip_serializing_if = "Option::is_none")] pub logical_plan: Option, + /// Rendered physical plan. Absent in the job list. #[serde(skip_serializing_if = "Option::is_none")] pub physical_plan: Option, + /// Rendered stage DAG. Absent in the job list. #[serde(skip_serializing_if = "Option::is_none")] pub stage_plan: Option, } +/// Terminal or in-flight state of a single task. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum TaskStatus { + /// Task is currently executing. Running, + /// Task completed successfully. Successful, - Failed { reason: String, error: String }, + /// Task failed, with the classified reason and the underlying error. + Failed { + /// Failure category, e.g. `ExecutionError` or `FetchPartitionError`. + reason: String, + /// Underlying error message. + error: String, + }, } +/// Per-task detail within a stage. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TaskSummary { /// task id @@ -79,36 +100,67 @@ pub struct TaskSummary { pub output_rows: usize, } +/// Five-number summary over a stage's tasks, used to spot skew. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Percentiles { + /// Smallest observed value. pub min: u64, + /// 25th percentile. pub p25: u64, + /// 50th percentile. pub median: u64, + /// 75th percentile. pub p75: u64, + /// Largest observed value. pub max: u64, } +/// Summary of one query stage, served by `GET /api/job/{job_id}/stages`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueryStageSummary { + /// Stage id, as a string. pub stage_id: String, + /// Stage state, e.g. `Running`, `Successful`, `Failed`. pub stage_status: String, + /// Rows read by the stage, summed across tasks. pub input_rows: usize, + /// Rows produced by the stage, summed across tasks. pub output_rows: usize, + /// Formatted wall time across the stage's tasks, if any have started. pub elapsed_compute: Option, + /// Rendered plan for this stage, in the requested [`PlanFormat`]. #[serde(skip_serializing_if = "Option::is_none")] pub stage_plan: Option, + /// Distribution of per-task execution time. #[serde(skip_serializing_if = "Option::is_none")] pub task_duration_percentiles: Option, + /// Distribution of per-task input row counts. #[serde(skip_serializing_if = "Option::is_none")] pub task_input_percentiles: Option, + /// One entry per task. Always `Some` today; the `Option` predates + /// multi-partition tasks, when this list was indexed by partition and + /// could be sparse. pub tasks: Vec>, } +/// Response body for `GET /api/job/{job_id}/stages`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueryStagesResponse { + /// One summary per stage, in stage order. pub stages: Vec, } -/// Session config as flat key/value pairs (from `SessionConfig::to_props()`), -/// sorted for stable output. -pub type JobConfig = BTreeMap; +/// How a plan should be rendered. Parsed from the `?plan_format=` query +/// parameter, and part of the REST contract, so it lives alongside the +/// response types rather than with the HTTP handlers. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PlanFormat { + /// `?plan_format=default` => plain indent, no metrics + #[default] + Default, + /// `?plan_format=tree` => tree render, no metrics + Tree, + /// `?plan_format=metrics` => indent with aggregated metrics + Metrics, +} diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs index affddf8a53..22c4d525d8 100644 --- a/ballista/history/src/lib.rs +++ b/ballista/history/src/lib.rs @@ -15,10 +15,24 @@ // specific language governing permissions and limitations // under the License. +#![warn(missing_docs)] + //! Shared types for Ballista's scheduler REST API. //! //! These response types live outside `ballista-scheduler` so that a future //! history server can serialize byte-identical JSON from stored state without //! depending on the scheduler's live execution graph. +//! +//! # What belongs here +//! +//! A type belongs in this crate when it is part of the `/api/*` wire contract +//! *and* can be reconstructed from a stored event log. Types describing live +//! scheduler state (`SchedulerStateResponse`, `CancelJobResponse`) stay in +//! `ballista-scheduler`, because there is nothing to replay. +//! +//! `ExecutorResponse` also stays behind, for a different reason: it embeds +//! `ballista-core` types, and this crate is deliberately serde-only so it +//! stays cheap to depend on. Sharing it would mean either taking a +//! `ballista-core` dependency here or duplicating those structs. pub mod dto; diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs index ce1ec45315..70bafd10e9 100644 --- a/ballista/scheduler/src/api/dto_build.rs +++ b/ballista/scheduler/src/api/dto_build.rs @@ -23,9 +23,10 @@ //! in the axum handlers) means the same DTOs can be produced from state that //! did not come from a live handler request. //! -//! Everything in this module is pure: state in, DTO out, no I/O. +//! Everything in this module is pure: state in, DTO out, no I/O and no +//! wall-clock reads. `graph_to_query_stages` takes `now` from its caller so a +//! replay of a stored log renders the same elapsed times every time. -use crate::api::handlers::{JobQueryParams, PlanFormat}; use crate::display::format_stage_metrics; use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage}; use crate::state::execution_stage::TaskInfo; @@ -35,12 +36,11 @@ use ballista_core::serde::protobuf::failed_task::FailedReason::{ }; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{FailedTask, task_status}; -use ballista_core::utils::get_current_time; use ballista_history::dto::{ - JobResponse, Percentiles, QueryStageSummary, QueryStagesResponse, TaskStatus, - TaskSummary, + JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse, + TaskStatus, TaskSummary, }; -use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::displayable; use datafusion::physical_plan::metrics::{MetricsSet, Time}; use std::time::Duration; @@ -55,14 +55,6 @@ pub fn job_overview_to_response(job: &JobOverview) -> JobResponse { 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(), @@ -72,7 +64,7 @@ pub fn job_overview_to_response(job: &JobOverview) -> JobResponse { end_time: job.end_time, num_stages: job.num_stages, completed_stages: job.completed_stages, - percent_complete, + percent_complete: percent_complete(job.completed_stages, job.num_stages), logical_plan: None, physical_plan: None, stage_plan: None, @@ -83,7 +75,7 @@ pub fn job_overview_to_response(job: &JobOverview) -> JobResponse { /// rendered logical, physical, and stage plans. pub fn graph_to_job_response( graph: &ExecutionGraphBox, - query: &JobQueryParams, + plan_format: PlanFormat, ) -> JobResponse { let stage_plan = format!("{graph:?}"); let job = graph.as_ref(); @@ -95,13 +87,10 @@ pub fn graph_to_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 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()) + displayable(job.physical_plan().as_ref()) .indent(false) .to_string() } @@ -119,7 +108,7 @@ pub fn graph_to_job_response( end_time: job.end_time(), num_stages, completed_stages, - percent_complete, + percent_complete: percent_complete(completed_stages, num_stages), logical_plan: job.logical_plan().map(str::to_owned), physical_plan: Some(physical_plan), stage_plan: Some(stage_plan), @@ -129,75 +118,53 @@ pub fn graph_to_job_response( /// Build the per-stage summaries served by `GET /api/job/{job_id}/stages`. pub fn graph_to_query_stages( graph: &ExecutionGraphBox, - query: &JobQueryParams, + plan_format: PlanFormat, + now: u128, ) -> QueryStagesResponse { - let plan_format = query.plan_format.clone().unwrap_or_default(); - let stages = graph .as_ref() .stages() .iter() .map(|(id, stage)| { - let mut summary = QueryStageSummary { + // Every started stage contributes the same three things; only how + // it reaches its metrics and which elapsed-time rule applies + // differ. Stages that have not started yet contribute nothing. + let started: Option<(&[MetricsSet], &[TaskInfo], Option)> = + match stage { + ExecutionStage::Running(s) => Some(( + s.stage_metrics.as_deref().unwrap_or(&[]), + &s.task_infos, + get_running_stage_time(&s.task_infos, now), + )), + ExecutionStage::Successful(s) => Some(( + &s.stage_metrics, + &s.task_infos, + get_finished_stage_time(&s.task_infos), + )), + ExecutionStage::Failed(s) => Some(( + s.stage_metrics.as_deref().unwrap_or(&[]), + &s.task_infos, + get_finished_stage_time(&s.task_infos), + )), + _ => None, + }; + + let has_started = started.is_some(); + let (metrics, task_infos, elapsed_compute) = started.unwrap_or_default(); + let tasks = task_summaries(task_infos, metrics); + + 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(render_stage_plan( - running_stage.plan.as_ref(), - metrics, - &plan_format, - )); - summary.input_rows = get_combined_count(metrics, "input_rows"); - summary.output_rows = get_combined_count(metrics, "output_rows"); - summary.elapsed_compute = get_running_stage_time( - &running_stage.task_infos, - get_current_time(), - ); - summary.tasks = task_summaries(&running_stage.task_infos, metrics); - } - ExecutionStage::Successful(completed_stage) => { - let metrics = completed_stage.stage_metrics.as_slice(); - summary.stage_plan = Some(render_stage_plan( - completed_stage.plan.as_ref(), - metrics, - &plan_format, - )); - summary.input_rows = get_combined_count(metrics, "input_rows"); - summary.output_rows = get_combined_count(metrics, "output_rows"); - summary.elapsed_compute = - get_finished_stage_time(&completed_stage.task_infos); - summary.tasks = task_summaries(&completed_stage.task_infos, metrics); - } - ExecutionStage::Failed(failed_stage) => { - let metrics = failed_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(render_stage_plan( - failed_stage.plan.as_ref(), - metrics, - &plan_format, - )); - summary.input_rows = get_combined_count(metrics, "input_rows"); - summary.output_rows = get_combined_count(metrics, "output_rows"); - summary.elapsed_compute = - get_finished_stage_time(&failed_stage.task_infos); - summary.tasks = task_summaries(&failed_stage.task_infos, metrics); - } - _ => {} + input_rows: get_combined_count(metrics, "input_rows"), + output_rows: get_combined_count(metrics, "output_rows"), + elapsed_compute, + stage_plan: has_started + .then(|| render_stage_plan(stage.plan(), metrics, plan_format)), + task_duration_percentiles: task_duration_percentiles(&tasks), + task_input_percentiles: task_input_percentiles(&tasks), + tasks, } - - summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); - summary.task_input_percentiles = task_input_percentiles(&summary.tasks); - summary }) .collect(); @@ -207,9 +174,9 @@ pub fn graph_to_query_stages( /// Render one stage's plan in the requested format. `Metrics` overlays the /// stage's aggregated metrics onto the plan; the other formats ignore them. fn render_stage_plan( - plan: &dyn datafusion::physical_plan::ExecutionPlan, + plan: &dyn ExecutionPlan, metrics: &[MetricsSet], - plan_format: &PlanFormat, + plan_format: PlanFormat, ) -> String { match plan_format { PlanFormat::Default => displayable(plan).indent(false).to_string(), @@ -262,7 +229,7 @@ fn task_summaries( /// /// A free function rather than a `From` impl: both types are foreign to this /// crate now that [`TaskStatus`] lives in `ballista-history`. -pub fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { +fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { match value { task_status::Status::Running(_) => TaskStatus::Running, task_status::Status::Failed(failed) => TaskStatus::Failed { @@ -273,6 +240,15 @@ pub fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { } } +/// Progress as a percentage of stages completed. Zero-stage jobs report 0 +/// rather than dividing by zero. +fn percent_complete(completed_stages: usize, num_stages: usize) -> u8 { + if num_stages == 0 { + return 0; + } + ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8 +} + 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)] @@ -352,33 +328,31 @@ fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, Strin } } -fn get_running_stage_time(task_infos: &[TaskInfo], current_time: u128) -> Option { - let min_start = task_infos +/// Earliest non-zero task start in a stage. Zero means "not started yet", so +/// those entries are ignored rather than dragging the minimum to 0. +fn min_start_time(task_infos: &[TaskInfo]) -> Option { + task_infos .iter() .map(|t| t.start_exec_time) .filter(|t| *t > 0) - .min(); + .min() +} - match (min_start, current_time) { +fn get_running_stage_time(task_infos: &[TaskInfo], current_time: u128) -> Option { + match (min_start_time(task_infos), current_time) { (Some(start), end) if end >= start => Some(format_millis(end - start)), _ => 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) { + match (min_start_time(task_infos), max_end) { (Some(start), Some(end)) if end >= start => Some(format_millis(end - start)), _ => None, } @@ -475,11 +449,6 @@ mod tests { } } - #[test] - fn test_job_elapsed_saturates_when_end_precedes_start() { - assert_eq!(job_elapsed_ms(900, 100), 0); - } - // --- get_finished_stage_time --- #[test] diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index f236dd37a6..1c79c7c49c 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -27,7 +27,8 @@ use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric}; use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; -use ballista_history::dto::JobResponse; +use ballista_core::utils::get_current_time; +use ballista_history::dto::{JobResponse, PlanFormat}; use datafusion::DATAFUSION_VERSION; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; @@ -113,18 +114,6 @@ pub struct JobQueryParams { pub plan_format: Option, } -#[derive(Debug, serde::Deserialize, Default, Clone)] -#[serde(rename_all = "snake_case")] -pub enum PlanFormat { - /// ?plan_format=default => plain indent, no metrics - #[default] - Default, - /// ?plan_format=tree => tree render, no metrics - Tree, - /// ?plan_format=metrics => indent with aggregated metrics - Metrics, -} - /// A handler for GET requests to the root (`/`). /// It redirects to `https://nightlies.apache.org/datafusion/ballista/tui//` /// forwarding any query parameters @@ -313,7 +302,10 @@ pub async fn get_job< })? .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; - Ok(Json(dto_build::graph_to_job_response(&graph, &query))) + Ok(Json(dto_build::graph_to_job_response( + &graph, + query.plan_format.unwrap_or_default(), + ))) } pub async fn cancel_job< @@ -404,7 +396,11 @@ pub async fn get_query_stages< ) })? { - Ok(Json(dto_build::graph_to_query_stages(&graph, &query))) + Ok(Json(dto_build::graph_to_query_stages( + &graph, + query.plan_format.unwrap_or_default(), + get_current_time(), + ))) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } diff --git a/dev/release/README.md b/dev/release/README.md index 3f34ba10e5..cdacee9970 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -287,6 +287,7 @@ of the following crates: - [ballista-cli](https://crates.io/crates/ballista-cli) - [ballista-core](https://crates.io/crates/ballista-core) - [ballista-executor](https://crates.io/crates/ballista-executor) +- [ballista-history](https://crates.io/crates/ballista-history) - [ballista-scheduler](https://crates.io/crates/ballista-scheduler) Download and unpack the official release tarball @@ -306,6 +307,7 @@ dot -Tsvg dev/release/crate-deps.dot > dev/release/crate-deps.svg ```shell (cd ballista/core && cargo publish) (cd ballista/executor && cargo publish) +(cd ballista/history && cargo publish) (cd ballista/scheduler && cargo publish) (cd ballista/client && cargo publish) (cd ballista-cli && cargo publish) diff --git a/dev/release/crate-deps.dot b/dev/release/crate-deps.dot index c3113e6a3d..2986ee5251 100644 --- a/dev/release/crate-deps.dot +++ b/dev/release/crate-deps.dot @@ -18,12 +18,14 @@ digraph G { ballista_core + ballista_history ballista_scheduler ballista_executor ballista ballista_cli ballista_scheduler -> ballista_core + ballista_scheduler -> ballista_history ballista_executor -> ballista_core diff --git a/dev/update_ballista_versions.py b/dev/update_ballista_versions.py index 8fedb21023..9d29f65482 100755 --- a/dev/update_ballista_versions.py +++ b/dev/update_ballista_versions.py @@ -43,6 +43,7 @@ def update_cargo_toml(cargo_toml: str, new_version: str): 'ballista', 'ballista-core', 'ballista-executor', + 'ballista-history', 'ballista-scheduler', 'ballista-cli', ) @@ -80,6 +81,7 @@ def main(): for rel_path in [ 'ballista-cli', 'ballista/core', + 'ballista/history', 'ballista/scheduler', 'ballista/executor', 'ballista/client', From b6b6a714dcb2b13f57bcee60e4ac4500c63bb840 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 10:33:38 -0600 Subject: [PATCH 3/7] docs(history): state the write-once, replay-verbatim data flow The previous wording left it ambiguous whether the history server re-derives responses from stored execution state or replays stored DTOs. It replays them: the scheduler builds each response once against the live graph and writes it to the event log, so byte-identical output is a structural property rather than two implementations agreeing. Also records the consequence, that anything not captured at write time cannot be recovered at replay time. --- ballista/history/src/lib.rs | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs index 22c4d525d8..2f814e0fdc 100644 --- a/ballista/history/src/lib.rs +++ b/ballista/history/src/lib.rs @@ -19,15 +19,37 @@ //! Shared types for Ballista's scheduler REST API. //! -//! These response types live outside `ballista-scheduler` so that a future -//! history server can serialize byte-identical JSON from stored state without -//! depending on the scheduler's live execution graph. +//! This crate is the serialization boundary between the live scheduler and the +//! history server: the scheduler builds these types, writes them to an event +//! log, and the history server reads them back and serves them. +//! +//! # Write once, replay verbatim +//! +//! The history server does **not** re-derive responses from stored execution +//! state. When a job finishes, the scheduler runs its DTO builders once against +//! the live execution graph and writes the finished [`dto::JobResponse`] and +//! [`dto::QueryStagesResponse`] into the log's terminal record. Replay +//! deserializes those values and re-serializes them unchanged: +//! +//! ```text +//! scheduler ──> dto builders ──> DTO ──> event log ──> history server ──> JSON +//! ``` +//! +//! So byte-identical output is a structural property, not a convention two +//! implementations have to keep agreeing on. There is one definition of each +//! type and one place that populates it. That is also why the builders +//! themselves stay in `ballista-scheduler`: they need the live execution graph, +//! and nothing on the replay path calls them. +//! +//! One consequence worth knowing: anything not captured when the record is +//! written cannot be recovered later. Plans, for instance, are rendered in a +//! single format at write time, so replay cannot re-render them differently. //! //! # What belongs here //! //! A type belongs in this crate when it is part of the `/api/*` wire contract -//! *and* can be reconstructed from a stored event log. Types describing live -//! scheduler state (`SchedulerStateResponse`, `CancelJobResponse`) stay in +//! *and* is worth persisting in the event log. Types describing live scheduler +//! state (`SchedulerStateResponse`, `CancelJobResponse`) stay in //! `ballista-scheduler`, because there is nothing to replay. //! //! `ExecutorResponse` also stays behind, for a different reason: it embeds From e2e187352e3a7586f773295666ede81d7f2f9253 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 10:40:05 -0600 Subject: [PATCH 4/7] refactor: rename ballista-history to ballista-api-types The crate holds the /api/* wire types, and it has three parties, not one: the scheduler serves them, the web TUI deserializes them, and a future history server will serve replayed copies. Naming it after the history server made it awkward for the TUI, which parses live scheduler responses and today keeps its own duplicate declarations. Renaming it after the contract it defines removes that friction. The event-log schema, writer, and reader can then land as a separate ballista-history crate that depends on this one. --- Cargo.lock | 16 +++--- Cargo.toml | 2 +- ballista/{history => api-types}/Cargo.toml | 4 +- ballista/{history => api-types}/src/dto.rs | 0 ballista/api-types/src/lib.rs | 56 ++++++++++++++++++++ ballista/history/src/lib.rs | 60 ---------------------- ballista/scheduler/Cargo.toml | 4 +- ballista/scheduler/src/api/dto_build.rs | 12 ++--- ballista/scheduler/src/api/handlers.rs | 2 +- dev/release/README.md | 4 +- dev/release/crate-deps.dot | 4 +- dev/update_ballista_versions.py | 4 +- 12 files changed, 82 insertions(+), 86 deletions(-) rename ballista/{history => api-types}/Cargo.toml (90%) rename ballista/{history => api-types}/src/dto.rs (100%) create mode 100644 ballista/api-types/src/lib.rs delete mode 100644 ballista/history/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 498f742dfe..697df0a10f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1024,6 +1024,13 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-api-types" +version = "54.0.0" +dependencies = [ + "serde", +] + [[package]] name = "ballista-benchmarks" version = "54.0.0" @@ -1200,13 +1207,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "ballista-history" -version = "54.0.0" -dependencies = [ - "serde", -] - [[package]] name = "ballista-scheduler" version = "54.0.0" @@ -1214,8 +1214,8 @@ dependencies = [ "arrow-flight", "async-trait", "axum", + "ballista-api-types", "ballista-core", - "ballista-history", "clap 4.6.3", "dashmap", "datafusion", diff --git a/Cargo.toml b/Cargo.toml index 50ffac073d..15817827b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,7 +22,7 @@ members = [ "ballista/client", "ballista/core", "ballista/executor", - "ballista/history", + "ballista/api-types", "ballista/scheduler", "benchmarks", "chaos-testing", diff --git a/ballista/history/Cargo.toml b/ballista/api-types/Cargo.toml similarity index 90% rename from ballista/history/Cargo.toml rename to ballista/api-types/Cargo.toml index a4e56d709c..9c9485e332 100644 --- a/ballista/history/Cargo.toml +++ b/ballista/api-types/Cargo.toml @@ -16,8 +16,8 @@ # under the License. [package] -name = "ballista-history" -description = "Shared REST response types for the Ballista scheduler and history server" +name = "ballista-api-types" +description = "Wire types for the Ballista scheduler REST API" license = "Apache-2.0" version = "54.0.0" homepage = "https://datafusion.apache.org/ballista/" diff --git a/ballista/history/src/dto.rs b/ballista/api-types/src/dto.rs similarity index 100% rename from ballista/history/src/dto.rs rename to ballista/api-types/src/dto.rs diff --git a/ballista/api-types/src/lib.rs b/ballista/api-types/src/lib.rs new file mode 100644 index 0000000000..08fab21c28 --- /dev/null +++ b/ballista/api-types/src/lib.rs @@ -0,0 +1,56 @@ +// 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. + +#![warn(missing_docs)] + +//! Wire types for Ballista's scheduler REST API. +//! +//! This crate is the single definition of what `/api/*` puts on the wire. It +//! exists because that contract has more than one party: +//! +//! - **`ballista-scheduler`** serves these responses for live jobs. +//! - **The web TUI** deserializes them to render the cluster. +//! - **A history server** ([#1923]) will serve them for completed jobs, replayed +//! from a durable event log. +//! +//! Each of those independently re-declaring the same structs is how the shapes +//! drift apart, so they share one definition instead. +//! +//! The crate is deliberately serde-only. That keeps it cheap enough for anyone +//! to depend on, and it is what lets the TUI use it from a `wasm32` build. +//! +//! # What belongs here +//! +//! A type belongs here when it is part of the `/api/*` contract and more than +//! one party needs it. Types describing live scheduler internals +//! (`SchedulerStateResponse`, `CancelJobResponse`) stay in `ballista-scheduler`. +//! +//! `ExecutorResponse` stays behind for a different reason: it embeds +//! `ballista-core` types, so sharing it would mean either taking a +//! `ballista-core` dependency here or duplicating those structs. Neither is +//! worth it until something outside the scheduler needs it. +//! +//! Note that the types are shared but the *construction* is not: building a +//! response from a live execution graph needs scheduler internals, so those +//! builders live in `ballista-scheduler`. A history server does not re-derive +//! responses — it replays ones the scheduler already built and stored, so +//! byte-identical output is a structural property rather than two +//! implementations agreeing. +//! +//! [#1923]: https://github.com/apache/datafusion-ballista/issues/1923 + +pub mod dto; diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs deleted file mode 100644 index 2f814e0fdc..0000000000 --- a/ballista/history/src/lib.rs +++ /dev/null @@ -1,60 +0,0 @@ -// 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. - -#![warn(missing_docs)] - -//! Shared types for Ballista's scheduler REST API. -//! -//! This crate is the serialization boundary between the live scheduler and the -//! history server: the scheduler builds these types, writes them to an event -//! log, and the history server reads them back and serves them. -//! -//! # Write once, replay verbatim -//! -//! The history server does **not** re-derive responses from stored execution -//! state. When a job finishes, the scheduler runs its DTO builders once against -//! the live execution graph and writes the finished [`dto::JobResponse`] and -//! [`dto::QueryStagesResponse`] into the log's terminal record. Replay -//! deserializes those values and re-serializes them unchanged: -//! -//! ```text -//! scheduler ──> dto builders ──> DTO ──> event log ──> history server ──> JSON -//! ``` -//! -//! So byte-identical output is a structural property, not a convention two -//! implementations have to keep agreeing on. There is one definition of each -//! type and one place that populates it. That is also why the builders -//! themselves stay in `ballista-scheduler`: they need the live execution graph, -//! and nothing on the replay path calls them. -//! -//! One consequence worth knowing: anything not captured when the record is -//! written cannot be recovered later. Plans, for instance, are rendered in a -//! single format at write time, so replay cannot re-render them differently. -//! -//! # What belongs here -//! -//! A type belongs in this crate when it is part of the `/api/*` wire contract -//! *and* is worth persisting in the event log. Types describing live scheduler -//! state (`SchedulerStateResponse`, `CancelJobResponse`) stay in -//! `ballista-scheduler`, because there is nothing to replay. -//! -//! `ExecutorResponse` also stays behind, for a different reason: it embeds -//! `ballista-core` types, and this crate is deliberately serde-only so it -//! stays cheap to depend on. Sharing it would mean either taking a -//! `ballista-core` dependency here or duplicating those structs. - -pub mod dto; diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index 17ce991bd2..cd26d8d420 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -41,7 +41,7 @@ disable-stage-plan-cache = [] graphviz-support = ["dep:graphviz-rust"] keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"] prometheus-metrics = ["prometheus", "once_cell"] -rest-api = ["dep:ballista-history"] +rest-api = ["dep:ballista-api-types"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -49,8 +49,8 @@ substrait = ["dep:datafusion-substrait"] arrow-flight = { workspace = true } async-trait = { workspace = true } axum = "0.8.9" +ballista-api-types = { path = "../api-types", version = "54.0.0", optional = true } ballista-core = { path = "../core", version = "54.0.0" } -ballista-history = { path = "../history", version = "54.0.0", optional = true } clap = { workspace = true, optional = true } dashmap = { workspace = true } datafusion = { workspace = true } diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs index 70bafd10e9..d1f343666b 100644 --- a/ballista/scheduler/src/api/dto_build.rs +++ b/ballista/scheduler/src/api/dto_build.rs @@ -19,7 +19,7 @@ //! //! These functions are the only place the scheduler's `ExecutionGraph` / //! `JobOverview` shapes are translated into the wire types in -//! [`ballista_history::dto`]. Keeping the translation here (rather than inline +//! [`ballista_api_types::dto`]. Keeping the translation here (rather than inline //! in the axum handlers) means the same DTOs can be produced from state that //! did not come from a live handler request. //! @@ -31,15 +31,15 @@ use crate::display::format_stage_metrics; use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage}; use crate::state::execution_stage::TaskInfo; use crate::state::task_manager::JobOverview; +use ballista_api_types::dto::{ + JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse, + TaskStatus, TaskSummary, +}; use ballista_core::serde::protobuf::failed_task::FailedReason::{ ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, }; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{FailedTask, task_status}; -use ballista_history::dto::{ - JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse, - TaskStatus, TaskSummary, -}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::displayable; use datafusion::physical_plan::metrics::{MetricsSet, Time}; @@ -228,7 +228,7 @@ fn task_summaries( /// Map a protobuf task status onto the wire enum. /// /// A free function rather than a `From` impl: both types are foreign to this -/// crate now that [`TaskStatus`] lives in `ballista-history`. +/// crate now that [`TaskStatus`] lives in `ballista-api-types`. fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { match value { task_status::Status::Running(_) => TaskStatus::Running, diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 1c79c7c49c..bd6d77a501 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -21,6 +21,7 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; +use ballista_api_types::dto::{JobResponse, PlanFormat}; use ballista_core::BALLISTA_VERSION; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric}; @@ -28,7 +29,6 @@ use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; use ballista_core::utils::get_current_time; -use ballista_history::dto::{JobResponse, PlanFormat}; use datafusion::DATAFUSION_VERSION; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; diff --git a/dev/release/README.md b/dev/release/README.md index cdacee9970..c6b0de31c0 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -287,7 +287,7 @@ of the following crates: - [ballista-cli](https://crates.io/crates/ballista-cli) - [ballista-core](https://crates.io/crates/ballista-core) - [ballista-executor](https://crates.io/crates/ballista-executor) -- [ballista-history](https://crates.io/crates/ballista-history) +- [ballista-api-types](https://crates.io/crates/ballista-api-types) - [ballista-scheduler](https://crates.io/crates/ballista-scheduler) Download and unpack the official release tarball @@ -307,7 +307,7 @@ dot -Tsvg dev/release/crate-deps.dot > dev/release/crate-deps.svg ```shell (cd ballista/core && cargo publish) (cd ballista/executor && cargo publish) -(cd ballista/history && cargo publish) +(cd ballista/api-types && cargo publish) (cd ballista/scheduler && cargo publish) (cd ballista/client && cargo publish) (cd ballista-cli && cargo publish) diff --git a/dev/release/crate-deps.dot b/dev/release/crate-deps.dot index 2986ee5251..27bee709ba 100644 --- a/dev/release/crate-deps.dot +++ b/dev/release/crate-deps.dot @@ -18,14 +18,14 @@ digraph G { ballista_core - ballista_history + ballista_api_types ballista_scheduler ballista_executor ballista ballista_cli ballista_scheduler -> ballista_core - ballista_scheduler -> ballista_history + ballista_scheduler -> ballista_api_types ballista_executor -> ballista_core diff --git a/dev/update_ballista_versions.py b/dev/update_ballista_versions.py index 9d29f65482..ca5bef3bf8 100755 --- a/dev/update_ballista_versions.py +++ b/dev/update_ballista_versions.py @@ -43,7 +43,7 @@ def update_cargo_toml(cargo_toml: str, new_version: str): 'ballista', 'ballista-core', 'ballista-executor', - 'ballista-history', + 'ballista-api-types', 'ballista-scheduler', 'ballista-cli', ) @@ -81,7 +81,7 @@ def main(): for rel_path in [ 'ballista-cli', 'ballista/core', - 'ballista/history', + 'ballista/api-types', 'ballista/scheduler', 'ballista/executor', 'ballista/client', From cca6f23d1d034b0d175164f08e88a287b60c433f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 11:35:42 -0600 Subject: [PATCH 5/7] feat(history): add the event-log schema, writer, and reader Second step toward the history server (#1923), after the wire-type extraction in #2256. Adds the durable format and the machinery to read and write it. Nothing in the scheduler calls this yet. - A versioned JSONL schema. JobStart / StageStart / StageEnd / TaskEnd form an incremental timeline; the terminal JobEnd embeds the finished REST responses, so replay re-serves what the scheduler already built rather than re-deriving it. - An async buffered EventLogWriter. All file I/O runs on a background task so the scheduler's event loop never waits on disk. Timeline events are dropped rather than allowed to block when the queue backs up, since losing a progress record beats stalling scheduling. JobEnd is the exception and waits for capacity, because a job missing it is invisible to the history server. - A reader that folds a completed log back into the served payload, skipping malformed lines rather than failing. Restores the JobConfig alias to ballista-api-types, which now has a real consumer in the JobEnd record. TaskEnd names a task rather than a partition: under the multi-partition task model a task owns a slice of partitions, and TaskStatus carries task_id, not partition_id. --- Cargo.lock | 12 ++ Cargo.toml | 1 + ballista/api-types/src/dto.rs | 8 + ballista/history/Cargo.toml | 38 ++++ ballista/history/src/event.rs | 174 ++++++++++++++++++ ballista/history/src/lib.rs | 52 ++++++ ballista/history/src/reader.rs | 139 ++++++++++++++ ballista/history/src/writer.rs | 313 ++++++++++++++++++++++++++++++++ dev/release/README.md | 2 + dev/release/crate-deps.dot | 3 + dev/update_ballista_versions.py | 2 + 11 files changed, 744 insertions(+) create mode 100644 ballista/history/Cargo.toml 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.lock b/Cargo.lock index 697df0a10f..cf2a092943 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1207,6 +1207,18 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "54.0.0" +dependencies = [ + "ballista-api-types", + "log", + "serde", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "ballista-scheduler" version = "54.0.0" diff --git a/Cargo.toml b/Cargo.toml index 15817827b1..3a7e1966cf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ members = [ "ballista/client", "ballista/core", "ballista/executor", + "ballista/history", "ballista/api-types", "ballista/scheduler", "benchmarks", diff --git a/ballista/api-types/src/dto.rs b/ballista/api-types/src/dto.rs index 272730af84..094d8b6ddc 100644 --- a/ballista/api-types/src/dto.rs +++ b/ballista/api-types/src/dto.rs @@ -19,6 +19,7 @@ //! server, so both serialize byte-identical JSON. use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; /// Summary of one job, served by `GET /api/jobs` and `GET /api/job/{job_id}`. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -164,3 +165,10 @@ pub enum PlanFormat { /// `?plan_format=metrics` => indent with aggregated metrics Metrics, } + +/// Session config for one job, as flat key/value pairs. Served by +/// `GET /api/job/{job_id}/config`. +/// +/// A `BTreeMap` so key order is deterministic: the same job must serialize +/// identically whether it is served live or replayed from a stored log. +pub type JobConfig = BTreeMap; diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml new file mode 100644 index 0000000000..ecf4713669 --- /dev/null +++ b/ballista/history/Cargo.toml @@ -0,0 +1,38 @@ +# 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 the Ballista history server" +license = "Apache-2.0" +version = "54.0.0" +homepage = "https://datafusion.apache.org/ballista/" +repository = "https://github.com/apache/datafusion-ballista" +authors = ["Apache DataFusion "] +edition = { workspace = true } +rust-version = { workspace = true } + +[dependencies] +ballista-api-types = { path = "../api-types", version = "54.0.0" } +log = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = "1" +tokio = { workspace = true, features = ["fs", "io-util", "macros", "rt", "sync"] } + +[dev-dependencies] +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs new file mode 100644 index 0000000000..3f848ce509 --- /dev/null +++ b/ballista/history/src/event.rs @@ -0,0 +1,174 @@ +// 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 ballista_api_types::dto::{JobConfig, JobResponse, QueryStagesResponse, TaskStatus}; +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 { + /// Job completed successfully. + Succeeded, + /// Job terminated with the given error. + Failed(String), + /// Job was cancelled before completing. + Cancelled, +} + +/// Metrics captured per finished task on the incremental timeline. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskEndMetrics { + /// Rows the task read. + pub input_rows: u64, + /// Rows the task produced. + pub output_rows: u64, + /// Time the task spent computing, in nanoseconds. + 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 { + /// First record in every log. Written when the scheduler accepts the job. + JobStart { + /// Schema version, see [`SCHEMA_VERSION`]. + version: u32, + /// Job identifier, matching the log's filename. + job_id: String, + /// Human-readable job name. + job_name: String, + /// When the job entered the queue. + queued_at: u64, + /// When the job was submitted for planning. + submitted_at: u64, + /// Rendered logical plan, if one was captured. + logical_plan: Option, + /// Rendered physical plan, if one was captured. + physical_plan: Option, + }, + /// A stage became runnable. + StageStart { + /// Stage identifier within the job. + stage_id: u32, + /// Number of partitions the stage will produce. + partitions: u32, + }, + /// A stage reached a terminal state. + StageEnd { + /// Stage identifier within the job. + stage_id: u32, + /// Terminal stage status. + status: String, + }, + /// A task finished, successfully or otherwise. + TaskEnd { + /// Stage the task belonged to. + stage_id: u32, + /// Task's slot within the stage. Under the multi-partition-task model a + /// task owns a slice of partitions, so this names the task rather than + /// a single partition. + task_id: u32, + /// Executor that ran the task. + executor_id: String, + /// Outcome of the task. + status: TaskStatus, + /// When the scheduler launched the task. + launch_time: u64, + /// When the executor began running it. + start_exec_time: u64, + /// When the executor finished it. + end_exec_time: u64, + /// Row counts and compute time for the task. + metrics: TaskEndMetrics, + }, + /// Terminal record, and the only one the history server serves from. It + /// carries the finished REST responses so replay never re-derives them. + JobEnd { + /// Schema version, see [`SCHEMA_VERSION`]. + version: u32, + /// How the job ended. + status: JobEndStatus, + /// When the job entered the queue. + queued_at: u64, + /// When the job started executing. + started_at: u64, + /// When the job reached its terminal state. + completed_at: u64, + /// Finished `GET /api/job/{job_id}` response. + job: Box, + /// Finished `GET /api/job/{job_id}/stages` response. + stages: Box, + /// Finished `GET /api/job/{job_id}/config` response. + config: JobConfig, + /// Rendered DOT graph of the stage DAG. + dot: String, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_api_types::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 0000000000..156935e825 --- /dev/null +++ b/ballista/history/src/lib.rs @@ -0,0 +1,52 @@ +// 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. + +#![warn(missing_docs)] + +//! Durable event logs for completed Ballista jobs. +//! +//! The scheduler appends one JSONL record per event to +//! `/.eventlog` while a job runs. The history server +//! reads those files back and serves the same `/api/*` responses the live +//! scheduler does, long after the scheduler has forgotten the job. +//! +//! # Write once, replay verbatim +//! +//! The terminal [`event::HistoryEvent::JobEnd`] record embeds the finished +//! [`ballista_api_types::dto::JobResponse`] and +//! [`ballista_api_types::dto::QueryStagesResponse`] the scheduler built from +//! its live execution graph. Replay deserializes those values and re-serializes +//! them unchanged, so the history server never re-derives a response and there +//! is no second implementation to drift. +//! +//! The earlier records ([`event::HistoryEvent::JobStart`], `StageStart`, +//! `StageEnd`, `TaskEnd`) form an incremental timeline. Nothing reads them yet; +//! they exist so a future UI can show a job progressing rather than only its +//! final state. +//! +//! # Durability +//! +//! [`writer::EventLogWriter`] does all file I/O on a background task, so the +//! scheduler's event loop never waits on a disk write. Timeline events are +//! dropped rather than allowed to block if the queue backs up, on the grounds +//! that losing a progress record is better than stalling scheduling. The +//! terminal `JobEnd` is the exception: it waits for queue capacity, because a +//! job missing its `JobEnd` is invisible to the history server. + +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 0000000000..bf7a87d805 --- /dev/null +++ b/ballista/history/src/reader.rs @@ -0,0 +1,139 @@ +// 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. + +//! Reads a completed `.eventlog` into the DTO bundle the history server +//! serves. A file is "completed" once it contains a `JobEnd` record. + +use crate::event::HistoryEvent; +use ballista_api_types::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use std::io::BufRead; +use std::path::Path; + +/// The served payload recovered from a completed job's event log. +#[derive(Debug, Clone)] +pub struct ReplayedJob { + /// `GET /api/job/{job_id}` response. + pub job: JobResponse, + /// `GET /api/job/{job_id}/stages` response. + pub stages: QueryStagesResponse, + /// `GET /api/job/{job_id}/config` response. + pub config: JobConfig, + /// Rendered DOT graph of the stage DAG. + pub dot: String, +} + +/// Read a completed job's payload out of its event log. +/// +/// Returns `Ok(None)` when the file has no `JobEnd` record, which means the job +/// is still running or the scheduler died before finishing it. Malformed lines +/// are skipped rather than treated as fatal, so a partially-written log still +/// yields its job if the terminal record survived. +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::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use ballista_api_types::dto::{JobResponse, QueryStagesResponse}; + 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()); + } +} diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs new file mode 100644 index 0000000000..de0e798fea --- /dev/null +++ b/ballista/history/src/writer.rs @@ -0,0 +1,313 @@ +// 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. + +//! 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<()>, + }, + Finish { + job_id: String, + done: oneshot::Sender<()>, + }, +} + +/// Handle to the background event-log writer task. +/// +/// Cloning is cheap and every clone feeds the same task, so the scheduler can +/// hand one to each component that needs to record events. +#[derive(Clone)] +pub struct EventLogWriter { + tx: mpsc::Sender, +} + +impl EventLogWriter { + /// Spawn the writer task, appending logs under `log_dir`. + /// + /// `buffer` bounds the in-flight event queue; beyond it, timeline events are + /// dropped rather than allowed to stall the caller. + 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() { + log::warn!( + "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; + } + } + + /// 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() { + log::warn!( + "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 { + log::warn!( + "event-log writer: failed to enqueue finish for {job_id} (channel closed)" + ); + } + } +} + +async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { + if let Err(e) = tokio::fs::create_dir_all(&log_dir).await { + log::warn!("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 { + log::warn!( + "event-log writer: write failed for {job_id}: {e}" + ); + } + } + Err(e) => log::warn!("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(()); + } + 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(()); + } + } + } +} + +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) => { + log::warn!("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, JobEndStatus, SCHEMA_VERSION}; + use ballista_api_types::dto::{JobResponse, QueryStagesResponse}; + 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() { + 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\"")); + } +} diff --git a/dev/release/README.md b/dev/release/README.md index c6b0de31c0..502b0a71af 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -288,6 +288,7 @@ of the following crates: - [ballista-core](https://crates.io/crates/ballista-core) - [ballista-executor](https://crates.io/crates/ballista-executor) - [ballista-api-types](https://crates.io/crates/ballista-api-types) +- [ballista-history](https://crates.io/crates/ballista-history) - [ballista-scheduler](https://crates.io/crates/ballista-scheduler) Download and unpack the official release tarball @@ -308,6 +309,7 @@ dot -Tsvg dev/release/crate-deps.dot > dev/release/crate-deps.svg (cd ballista/core && cargo publish) (cd ballista/executor && cargo publish) (cd ballista/api-types && cargo publish) +(cd ballista/history && cargo publish) (cd ballista/scheduler && cargo publish) (cd ballista/client && cargo publish) (cd ballista-cli && cargo publish) diff --git a/dev/release/crate-deps.dot b/dev/release/crate-deps.dot index 27bee709ba..637030aca5 100644 --- a/dev/release/crate-deps.dot +++ b/dev/release/crate-deps.dot @@ -19,6 +19,7 @@ digraph G { ballista_core ballista_api_types + ballista_history ballista_scheduler ballista_executor ballista @@ -27,6 +28,8 @@ digraph G { ballista_scheduler -> ballista_core ballista_scheduler -> ballista_api_types + ballista_history -> ballista_api_types + ballista_executor -> ballista_core ballista -> ballista_core diff --git a/dev/update_ballista_versions.py b/dev/update_ballista_versions.py index ca5bef3bf8..926520f42d 100755 --- a/dev/update_ballista_versions.py +++ b/dev/update_ballista_versions.py @@ -44,6 +44,7 @@ def update_cargo_toml(cargo_toml: str, new_version: str): 'ballista-core', 'ballista-executor', 'ballista-api-types', + 'ballista-history', 'ballista-scheduler', 'ballista-cli', ) @@ -82,6 +83,7 @@ def main(): 'ballista-cli', 'ballista/core', 'ballista/api-types', + 'ballista/history', 'ballista/scheduler', 'ballista/executor', 'ballista/client', From f60cacbe59de10268f76c53715703f54d6366df1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 8 Aug 2026 12:05:46 -0600 Subject: [PATCH 6/7] feat(history): make the event log survive future Ballista versions A log is written once and may be read years later by a much newer binary, so the guarantee has to run one way: a reader accepts any log whose version is not newer than its own. That is the opposite of BALLISTA_PROTOCOL_VERSION, the strict-equality handshake between scheduler and executor, where both ends are live and upgraded together. The format did not support that yet. Four changes: Self-describing envelope. Every line is now a LogRecord carrying `ev`, `version` and an opaque `data` payload, so a reader can route on the kind and check the version before committing to a shape it may not understand. Previously only JobStart and JobEnd carried a version at all. Stored responses are opaque. JobEnd holds the finished /api/* responses as raw JSON rather than typed structs, plus a small frozen JobIndex for listing. Those responses are ballista-api-types shapes, which change with the live REST contract: partition_id went from u32 to Vec and TaskStatus::Failed gained a field within one release cycle. Stored typed, either change would have made every older log unreadable, and the job would have silently disappeared. Stored raw, nothing ever parses the inner shape and replay relays the exact bytes. Unreadable is no longer indistinguishable from absent. read_completed_job returned Ok(None) both when a log had no JobEnd and when it had one that could not be parsed, and the loader treated that as "still running" and said nothing. It now returns a ReadError distinguishing an unsupported version from a malformed record. The version is actually checked. A record newer than SCHEMA_VERSION is reported as such rather than skipped. Adds testdata/schema-v1.eventlog, a frozen log replayed by CI on every build. It includes a record kind this build does not know and a stages payload carrying the multi-partition shape, so it exercises both forward-compatibility paths. Verified it fails as intended: a field rename in a stored type breaks all four compatibility tests. --- ballista/history/Cargo.toml | 2 +- ballista/history/src/event.rs | 410 ++++++++++++++----- ballista/history/src/lib.rs | 28 +- ballista/history/src/reader.rs | 365 ++++++++++++++--- ballista/history/src/writer.rs | 57 ++- ballista/history/testdata/schema-v1.eventlog | 7 + 6 files changed, 646 insertions(+), 223 deletions(-) create mode 100644 ballista/history/testdata/schema-v1.eventlog diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml index ecf4713669..009bf0172c 100644 --- a/ballista/history/Cargo.toml +++ b/ballista/history/Cargo.toml @@ -30,7 +30,7 @@ rust-version = { workspace = true } ballista-api-types = { path = "../api-types", version = "54.0.0" } log = { workspace = true } serde = { workspace = true, features = ["derive"] } -serde_json = "1" +serde_json = { version = "1", features = ["raw_value"] } tokio = { workspace = true, features = ["fs", "io-util", "macros", "rt", "sync"] } [dev-dependencies] diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs index 3f848ce509..b42bb504f2 100644 --- a/ballista/history/src/event.rs +++ b/ballista/history/src/event.rs @@ -15,16 +15,74 @@ // 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. +//! The on-disk event-log schema. One [`LogRecord`] is serialized per JSONL line. +//! +//! # Compatibility +//! +//! A log is written once and may be read years later by a much newer Ballista, +//! so the guarantee runs one way: **a reader accepts any record whose version is +//! less than or equal to its own**. There is no way to upgrade the writer of a +//! file that already exists. +//! +//! That is the opposite of `BALLISTA_PROTOCOL_VERSION`, the strict-equality +//! handshake between scheduler and executor. Both ends of that handshake are +//! live and upgraded together, so refusing to proceed is the safe move. Here +//! there is nothing to negotiate with. +//! +//! Three properties make the format survivable: +//! +//! 1. **Every line self-describes.** [`LogRecord`] carries `ev` and `version` +//! next to an opaque `data` payload, so a reader can decide whether it +//! understands a record before committing to its shape. +//! 2. **Unknown record types are skipped, not fatal.** A future scheduler can +//! add event kinds without breaking today's reader. +//! 3. **The served responses are stored verbatim**, as raw JSON rather than +//! typed structs. See [`JobEnd`]. -use ballista_api_types::dto::{JobConfig, JobResponse, QueryStagesResponse, TaskStatus}; +use ballista_api_types::dto::{JobConfig, TaskStatus}; use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; -/// Current on-disk schema version, stamped on `JobStart`/`JobEnd`. +/// Current on-disk schema version, stamped on every record. +/// +/// Bump this only for a **breaking** change: removing a field, changing a +/// field's type, or changing the meaning of an existing one. Additive changes +/// do not bump it and must remain readable from older logs, which in practice +/// means every new field carries `#[serde(default)]`. pub const SCHEMA_VERSION: u32 = 1; +/// One line of the log: a self-describing envelope around an opaque payload. +/// +/// Keeping the payload opaque at this level is what lets a reader route on +/// `ev` and check `version` *before* attempting to parse a shape it may not +/// understand. Deserializing the whole record up front would collapse "written +/// by a newer Ballista" and "corrupt" into the same indistinguishable failure. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LogRecord { + /// Record kind, e.g. `JobStart` or `JobEnd`. + pub ev: String, + /// Schema version this record was written with. See [`SCHEMA_VERSION`]. + pub version: u32, + /// Kind-specific payload, left unparsed. + pub data: Box, +} + +impl LogRecord { + /// Wrap a payload in an envelope stamped with the current schema version. + pub fn new(ev: &str, payload: &T) -> serde_json::Result { + Ok(LogRecord { + ev: ev.to_string(), + version: SCHEMA_VERSION, + data: serde_json::value::to_raw_value(payload)?, + }) + } + + /// Parse the payload as `T`. + pub fn decode Deserialize<'de>>(&self) -> serde_json::Result { + serde_json::from_str(self.data.get()) + } +} + /// How a job reached its terminal state. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum JobEndStatus { @@ -47,128 +105,256 @@ pub struct TaskEndMetrics { pub elapsed_compute_nanos: u64, } -/// A single record in a job's event log. +/// Written when the scheduler accepts a job. First record in every log. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobStart { + /// Job identifier, matching the log's filename. + pub job_id: String, + /// Human-readable job name. + pub job_name: String, + /// When the job entered the queue. + pub queued_at: u64, + /// When the job was submitted for planning. + pub submitted_at: u64, + /// Rendered logical plan, if one was captured. + pub logical_plan: Option, + /// Rendered physical plan, if one was captured. + pub physical_plan: Option, +} + +/// A stage became runnable. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StageStart { + /// Stage identifier within the job. + pub stage_id: u32, + /// Number of partitions the stage will produce. + pub partitions: u32, +} + +/// A stage reached a terminal state. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StageEnd { + /// Stage identifier within the job. + pub stage_id: u32, + /// Terminal stage status. + pub status: String, +} + +/// A task finished, successfully or otherwise. +/// +/// Identifiers are fixed-width `u32` 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. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskEnd { + /// Stage the task belonged to. + pub stage_id: u32, + /// Task's slot within the stage. Under the multi-partition-task model a + /// task owns a slice of partitions, so this names the task rather than a + /// single partition. + pub task_id: u32, + /// Executor that ran the task. + pub executor_id: String, + /// Outcome of the task. + pub status: TaskStatus, + /// When the scheduler launched the task. + pub launch_time: u64, + /// When the executor began running it. + pub start_exec_time: u64, + /// When the executor finished it. + pub end_exec_time: u64, + /// Row counts and compute time for the task. + pub metrics: TaskEndMetrics, +} + +/// Frozen summary of a completed job, owned by this crate rather than shared +/// with the REST API. +/// +/// The history server needs *some* structure to list and sort jobs, but it does +/// not need to understand the full responses. Keeping that structure minimal and +/// local decouples the part that must stay readable forever from +/// `ballista-api-types`, which evolves with the live REST contract. +/// +/// Adding fields here later is fine; each one needs `#[serde(default)]` so older +/// logs still parse. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobIndex { + /// Job identifier. + pub job_id: String, + /// Human-readable job name. + pub job_name: String, + /// Plain status word, e.g. `Completed` or `Failed`. + pub status: String, + /// Verbose status, including completion or failure detail. + pub job_status: String, + /// When the job started executing. + pub start_time: u64, + /// When the job reached its terminal state. + pub end_time: u64, +} + +/// Terminal record, and the only one the history server serves from. +/// +/// `job` and `stages` hold the finished REST responses **as raw JSON**, not as +/// typed structs. That is deliberate. Those responses are +/// `ballista-api-types` shapes, which change with the live REST contract: +/// `TaskSummary::partition_id` went from `u32` to `Vec`, and +/// `TaskStatus::Failed` gained a field, both within a single release cycle. /// -/// 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. +/// If this record stored them typed, a reader built after any such change would +/// fail to deserialize a log written before it, and the job would disappear. +/// Storing raw JSON means nothing ever parses the inner shape: the history +/// server relays the exact bytes the scheduler produced. The log is therefore +/// immune to REST type churn, and replayed output is byte-identical rather than +/// merely equivalent. #[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "ev")] +pub struct JobEnd { + /// How the job ended. + pub status: JobEndStatus, + /// When the job entered the queue. + pub queued_at: u64, + /// When the job started executing. + pub started_at: u64, + /// When the job reached its terminal state. + pub completed_at: u64, + /// Frozen summary, used to list and sort jobs without parsing the payloads. + pub index: JobIndex, + /// Finished `GET /api/job/{job_id}` response, stored verbatim. + pub job: Box, + /// Finished `GET /api/job/{job_id}/stages` response, stored verbatim. + pub stages: Box, + /// Finished `GET /api/job/{job_id}/config` response. + pub config: JobConfig, + /// Rendered DOT graph of the stage DAG. + pub dot: String, +} + +/// Record-kind discriminators, as written to the `ev` field. +pub mod kind { + /// [`super::JobStart`] + pub const JOB_START: &str = "JobStart"; + /// [`super::StageStart`] + pub const STAGE_START: &str = "StageStart"; + /// [`super::StageEnd`] + pub const STAGE_END: &str = "StageEnd"; + /// [`super::TaskEnd`] + pub const TASK_END: &str = "TaskEnd"; + /// [`super::JobEnd`] + pub const JOB_END: &str = "JobEnd"; +} + +/// An event to append to a job's log. +/// +/// This is the in-memory shape callers build. It is encoded to a [`LogRecord`] +/// on the way to disk rather than serialized directly, so the envelope stays the +/// only thing a reader has to understand unconditionally. +#[derive(Debug, Clone)] pub enum HistoryEvent { - /// First record in every log. Written when the scheduler accepts the job. - JobStart { - /// Schema version, see [`SCHEMA_VERSION`]. - version: u32, - /// Job identifier, matching the log's filename. - job_id: String, - /// Human-readable job name. - job_name: String, - /// When the job entered the queue. - queued_at: u64, - /// When the job was submitted for planning. - submitted_at: u64, - /// Rendered logical plan, if one was captured. - logical_plan: Option, - /// Rendered physical plan, if one was captured. - physical_plan: Option, - }, - /// A stage became runnable. - StageStart { - /// Stage identifier within the job. - stage_id: u32, - /// Number of partitions the stage will produce. - partitions: u32, - }, - /// A stage reached a terminal state. - StageEnd { - /// Stage identifier within the job. - stage_id: u32, - /// Terminal stage status. - status: String, - }, - /// A task finished, successfully or otherwise. - TaskEnd { - /// Stage the task belonged to. - stage_id: u32, - /// Task's slot within the stage. Under the multi-partition-task model a - /// task owns a slice of partitions, so this names the task rather than - /// a single partition. - task_id: u32, - /// Executor that ran the task. - executor_id: String, - /// Outcome of the task. - status: TaskStatus, - /// When the scheduler launched the task. - launch_time: u64, - /// When the executor began running it. - start_exec_time: u64, - /// When the executor finished it. - end_exec_time: u64, - /// Row counts and compute time for the task. - metrics: TaskEndMetrics, - }, - /// Terminal record, and the only one the history server serves from. It - /// carries the finished REST responses so replay never re-derives them. - JobEnd { - /// Schema version, see [`SCHEMA_VERSION`]. - version: u32, - /// How the job ended. - status: JobEndStatus, - /// When the job entered the queue. - queued_at: u64, - /// When the job started executing. - started_at: u64, - /// When the job reached its terminal state. - completed_at: u64, - /// Finished `GET /api/job/{job_id}` response. - job: Box, - /// Finished `GET /api/job/{job_id}/stages` response. - stages: Box, - /// Finished `GET /api/job/{job_id}/config` response. - config: JobConfig, - /// Rendered DOT graph of the stage DAG. - dot: String, - }, + /// See [`JobStart`]. + JobStart(JobStart), + /// See [`StageStart`]. + StageStart(StageStart), + /// See [`StageEnd`]. + StageEnd(StageEnd), + /// See [`TaskEnd`]. + TaskEnd(TaskEnd), + /// See [`JobEnd`]. + JobEnd(Box), +} + +impl HistoryEvent { + /// The `ev` discriminator this event is written with. + pub fn kind(&self) -> &'static str { + match self { + HistoryEvent::JobStart(_) => kind::JOB_START, + HistoryEvent::StageStart(_) => kind::STAGE_START, + HistoryEvent::StageEnd(_) => kind::STAGE_END, + HistoryEvent::TaskEnd(_) => kind::TASK_END, + HistoryEvent::JobEnd(_) => kind::JOB_END, + } + } + + /// Encode to the envelope written to disk. + pub fn to_record(&self) -> serde_json::Result { + match self { + HistoryEvent::JobStart(p) => LogRecord::new(self.kind(), p), + HistoryEvent::StageStart(p) => LogRecord::new(self.kind(), p), + HistoryEvent::StageEnd(p) => LogRecord::new(self.kind(), p), + HistoryEvent::TaskEnd(p) => LogRecord::new(self.kind(), p), + HistoryEvent::JobEnd(p) => LogRecord::new(self.kind(), p), + } + } } #[cfg(test)] mod tests { use super::*; - use ballista_api_types::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, + fn sample_job_end() -> JobEnd { + JobEnd { status: JobEndStatus::Succeeded, queued_at: 5, started_at: 10, completed_at: 20, - job: Box::new(job), - stages: Box::new(QueryStagesResponse { stages: vec![] }), + index: JobIndex { + job_id: "job-1".into(), + job_name: "q1".into(), + status: "Completed".into(), + job_status: "COMPLETED".into(), + start_time: 10, + end_time: 20, + }, + job: RawValue::from_string(r#"{"job_id":"job-1"}"#.to_string()).unwrap(), + stages: RawValue::from_string(r#"{"stages":[]}"#.to_string()).unwrap(), 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()); + } + } + + #[test] + fn job_end_round_trips_through_jsonl() { + let event = HistoryEvent::JobEnd(Box::new(sample_job_end())); + let line = serde_json::to_string(&event.to_record().unwrap()).unwrap(); + + assert!(line.contains(r#""ev":"JobEnd""#)); + assert!(line.contains(r#""version":1"#)); + + let record: LogRecord = serde_json::from_str(&line).unwrap(); + assert_eq!(record.ev, kind::JOB_END); + let back: JobEnd = record.decode().unwrap(); + assert_eq!(back.index.job_id, "job-1"); + } + + /// The stored responses are relayed as raw JSON and never round-tripped + /// through a typed struct, so a log stays readable when the REST types it + /// came from change shape. This is the property that keeps old logs alive + /// across releases. + #[test] + fn stored_payloads_survive_shapes_this_build_cannot_model() { + let mut end = sample_job_end(); + end.job = RawValue::from_string( + r#"{"a_field_from_the_future":[1,2,3],"partition_id":{"nested":true}}"# + .to_string(), + ) + .unwrap(); + let line = serde_json::to_string(&end).unwrap(); + + let back: JobEnd = serde_json::from_str(&line).unwrap(); + assert!(back.job.get().contains("a_field_from_the_future")); + // Byte-for-byte, not merely equivalent. + assert_eq!(back.job.get(), end.job.get()); + } + + /// The envelope is readable without knowing the payload's shape at all, + /// which is what lets a reader check the version before parsing. + #[test] + fn envelope_is_readable_without_understanding_the_payload() { + let line = + r#"{"ev":"SomeFutureEvent","version":99,"data":{"anything":[1,{"x":null}]}}"#; + let record: LogRecord = serde_json::from_str(line).unwrap(); + assert_eq!(record.ev, "SomeFutureEvent"); + assert_eq!(record.version, 99); + assert!(record.data.get().contains("anything")); } } diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs index 156935e825..814783a8c8 100644 --- a/ballista/history/src/lib.rs +++ b/ballista/history/src/lib.rs @@ -26,17 +26,23 @@ //! //! # Write once, replay verbatim //! -//! The terminal [`event::HistoryEvent::JobEnd`] record embeds the finished -//! [`ballista_api_types::dto::JobResponse`] and -//! [`ballista_api_types::dto::QueryStagesResponse`] the scheduler built from -//! its live execution graph. Replay deserializes those values and re-serializes -//! them unchanged, so the history server never re-derives a response and there -//! is no second implementation to drift. -//! -//! The earlier records ([`event::HistoryEvent::JobStart`], `StageStart`, -//! `StageEnd`, `TaskEnd`) form an incremental timeline. Nothing reads them yet; -//! they exist so a future UI can show a job progressing rather than only its -//! final state. +//! The terminal [`event::JobEnd`] record embeds the finished `/api/*` responses +//! the scheduler built from its live execution graph, stored as raw JSON. Replay +//! relays those bytes unchanged, so the history server never re-derives a +//! response and there is no second implementation to drift. +//! +//! The earlier records ([`event::JobStart`], [`event::StageStart`], +//! [`event::StageEnd`], [`event::TaskEnd`]) form an incremental timeline. +//! Nothing reads them yet; they exist so a future UI can show a job progressing +//! rather than only its final state. +//! +//! # Compatibility +//! +//! Logs outlive the binaries that wrote them, so a newer reader must keep +//! reading older logs indefinitely. [`event::SCHEMA_VERSION`] documents the +//! policy, [`event::LogRecord`] is the self-describing envelope that makes it +//! enforceable, and `testdata/schema-v1.eventlog` is a frozen log that CI +//! replays on every build to catch regressions. //! //! # Durability //! diff --git a/ballista/history/src/reader.rs b/ballista/history/src/reader.rs index bf7a87d805..803ca2a9e6 100644 --- a/ballista/history/src/reader.rs +++ b/ballista/history/src/reader.rs @@ -15,125 +15,356 @@ // specific language governing permissions and limitations // under the License. -//! Reads a completed `.eventlog` into the DTO bundle the history server +//! Reads a completed `.eventlog` into the payload the history server //! serves. A file is "completed" once it contains a `JobEnd` record. +//! +//! Reading is deliberately forgiving about lines it does not recognise and +//! deliberately loud about a `JobEnd` it cannot use. An unreadable terminal +//! record means a job that exists on disk will not appear in the UI, and that +//! must never be indistinguishable from a job that is simply still running. -use crate::event::HistoryEvent; -use ballista_api_types::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use crate::event::{JobEnd, JobIndex, LogRecord, SCHEMA_VERSION, kind}; +use ballista_api_types::dto::JobConfig; +use serde_json::value::RawValue; use std::io::BufRead; use std::path::Path; /// The served payload recovered from a completed job's event log. #[derive(Debug, Clone)] pub struct ReplayedJob { - /// `GET /api/job/{job_id}` response. - pub job: JobResponse, - /// `GET /api/job/{job_id}/stages` response. - pub stages: QueryStagesResponse, + /// Frozen summary, used to list and sort jobs. + pub index: JobIndex, + /// `GET /api/job/{job_id}` response, verbatim as the scheduler wrote it. + pub job: Box, + /// `GET /api/job/{job_id}/stages` response, verbatim. + pub stages: Box, /// `GET /api/job/{job_id}/config` response. pub config: JobConfig, /// Rendered DOT graph of the stage DAG. pub dot: String, } +/// Why a log that exists on disk yielded no servable job. +#[derive(Debug)] +pub enum ReadError { + /// The file could not be read at all. + Io(std::io::Error), + /// A `JobEnd` record is present but was written by a newer schema than + /// this build understands. + UnsupportedVersion { + /// Version stamped on the record. + found: u32, + /// Highest version this build can read. + supported: u32, + }, + /// A `JobEnd` record is present and claims a supported version, but could + /// not be deserialized. Distinct from a missing `JobEnd`: this is a real + /// problem worth surfacing, not a job that is still running. + Malformed(String), +} + +impl std::fmt::Display for ReadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ReadError::Io(e) => write!(f, "{e}"), + ReadError::UnsupportedVersion { found, supported } => write!( + f, + "event log schema version {found} is newer than the highest \ + supported version {supported}; upgrade the history server" + ), + ReadError::Malformed(e) => write!(f, "malformed JobEnd record: {e}"), + } + } +} + +impl std::error::Error for ReadError {} + +impl From for ReadError { + fn from(e: std::io::Error) -> Self { + ReadError::Io(e) + } +} + /// Read a completed job's payload out of its event log. /// -/// Returns `Ok(None)` when the file has no `JobEnd` record, which means the job -/// is still running or the scheduler died before finishing it. Malformed lines -/// are skipped rather than treated as fatal, so a partially-written log still -/// yields its job if the terminal record survived. -pub fn read_completed_job(path: &Path) -> std::io::Result> { +/// Returns `Ok(None)` only when the log genuinely has no `JobEnd` record, which +/// means the job is still running or the scheduler died before finishing it. +/// A `JobEnd` that is present but unusable is an `Err`, so callers can report it +/// rather than silently dropping the job. +/// +/// Lines that are not `JobEnd` are skipped without inspection, including ones +/// this build does not recognise: a future schema may add record types, and an +/// older reader must tolerate them rather than choke on the file. +pub fn read_completed_job(path: &Path) -> Result, ReadError> { 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() { + if line.trim().is_empty() { + continue; + } + + // Route on the envelope alone. A line we cannot even read an envelope + // from is treated as a record we do not understand, not as a failure. + let Ok(record) = serde_json::from_str::(&line) else { + continue; + }; + if record.ev != kind::JOB_END { 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, - })); + + if record.version > SCHEMA_VERSION { + return Err(ReadError::UnsupportedVersion { + found: record.version, + supported: SCHEMA_VERSION, + }); } + + return match record.decode::() { + Ok(end) => Ok(Some(ReplayedJob { + index: end.index, + job: end.job, + stages: end.stages, + config: end.config, + dot: end.dot, + })), + Err(e) => Err(ReadError::Malformed(e.to_string())), + }; } + Ok(None) } #[cfg(test)] mod tests { use super::*; - use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; - use ballista_api_types::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{HistoryEvent, JobEnd, JobEndStatus, SCHEMA_VERSION}; use std::io::Write; + fn sample_index() -> JobIndex { + JobIndex { + job_id: "job-1".into(), + job_name: "q1".into(), + status: "Completed".into(), + job_status: "COMPLETED".into(), + start_time: 2, + end_time: 3, + } + } + fn job_end_line() -> String { - let event = HistoryEvent::JobEnd { - version: SCHEMA_VERSION, + let event = HistoryEvent::JobEnd(Box::new(JobEnd { 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![] }), + index: sample_index(), + job: RawValue::from_string( + r#"{"job_id":"job-1","physical_plan":"ProjectionExec"}"#.to_string(), + ) + .unwrap(), + stages: RawValue::from_string(r#"{"stages":[]}"#.to_string()).unwrap(), config: Default::default(), dot: "digraph {}".into(), - }; - serde_json::to_string(&event).unwrap() + })); + serde_json::to_string(&event.to_record().unwrap()).unwrap() + } + + fn write_log( + dir: &tempfile::TempDir, + name: &str, + lines: &[&str], + ) -> std::path::PathBuf { + let path = dir.path().join(name); + let mut f = std::fs::File::create(&path).unwrap(); + for line in lines { + writeln!(f, "{line}").unwrap(); + } + path } #[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 path = write_log( + &dir, + "job-1.eventlog", + &[ + r#"{"ev":"StageStart","version":1,"data":{"stage_id":1,"partitions":4}}"#, + &job_end_line(), + ], + ); 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.index.job_id, "job-1"); assert_eq!(replayed.dot, "digraph {}"); + // The payload is re-served verbatim rather than round-tripped through a + // typed struct. + assert!(replayed.job.get().contains("ProjectionExec")); } #[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(); + let path = write_log( + &dir, + "job-2.eventlog", + &[r#"{"ev":"StageStart","version":1,"data":{"stage_id":1,"partitions":4}}"#], + ); assert!(read_completed_job(&path).unwrap().is_none()); } + + /// A record type this build has never heard of must not stop it reading the + /// rest of the file. This is the forward-compatibility guarantee: an older + /// history server keeps working against logs from a newer scheduler, as + /// long as the schema version itself has not been bumped. + #[test] + fn skips_unrecognized_record_types() { + let dir = tempfile::tempdir().unwrap(); + let path = write_log( + &dir, + "job-3.eventlog", + &[ + r#"{"ev":"SomeFutureEvent","version":1,"data":{"whatever":true}}"#, + &job_end_line(), + ], + ); + assert!(read_completed_job(&path).unwrap().is_some()); + } + + /// A newer schema version is reported as such rather than being silently + /// skipped, so an operator finds out their history server is too old. + #[test] + fn newer_schema_version_is_reported_not_skipped() { + let dir = tempfile::tempdir().unwrap(); + let line = format!( + r#"{{"ev":"JobEnd","version":{},"data":{{}}}}"#, + SCHEMA_VERSION + 1 + ); + let path = write_log(&dir, "job-4.eventlog", &[&line]); + + match read_completed_job(&path) { + Err(ReadError::UnsupportedVersion { found, supported }) => { + assert_eq!(found, SCHEMA_VERSION + 1); + assert_eq!(supported, SCHEMA_VERSION); + } + other => panic!("expected UnsupportedVersion, got {other:?}"), + } + } + + /// A corrupt terminal record must be distinguishable from a job that has + /// simply not finished. Both used to surface as `Ok(None)`, which meant a + /// job could vanish from the UI with nothing logged. + #[test] + fn malformed_job_end_is_an_error_not_a_missing_job() { + let dir = tempfile::tempdir().unwrap(); + let path = write_log( + &dir, + "job-5.eventlog", + &[r#"{"ev":"JobEnd","version":1,"data":{"status":"Succeeded"}}"#], + ); + + match read_completed_job(&path) { + Err(ReadError::Malformed(_)) => {} + other => panic!("expected Malformed, got {other:?}"), + } + } +} + +/// Compatibility tests against a checked-in log from an earlier schema version. +/// +/// `testdata/schema-v1.eventlog` is a **frozen artifact**. It is never +/// regenerated: the whole point is that it was written by an older Ballista and +/// must stay readable forever. If a change here makes this module fail, the +/// change breaks every event log already on disk in the field. +/// +/// Fixing such a failure by editing the fixture defeats the test. The options +/// are to make the change backward-compatible (usually `#[serde(default)]` on a +/// new field), or to bump [`SCHEMA_VERSION`] and keep a path that can still read +/// the old version. +#[cfg(test)] +mod compatibility { + use super::*; + + fn golden_v1() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("testdata") + .join("schema-v1.eventlog") + } + + #[test] + fn reads_a_v1_log_written_by_an_earlier_ballista() { + let replayed = read_completed_job(&golden_v1()) + .expect("a v1 log must remain readable") + .expect("the fixture contains a JobEnd record"); + + assert_eq!(replayed.index.job_id, "golden-v1"); + assert_eq!(replayed.index.job_name, "tpch-q1"); + assert_eq!(replayed.index.status, "Completed"); + assert_eq!(replayed.index.start_time, 1005); + assert_eq!(replayed.index.end_time, 1100); + assert_eq!( + replayed + .config + .get("datafusion.execution.target_partitions"), + Some(&"4".to_string()) + ); + assert!(replayed.dot.contains("digraph")); + } + + /// The stored responses must come back byte-for-byte, because that is what + /// lets the history server re-serve them without understanding them. + /// + /// Compared against the raw file text rather than a parsed + /// `serde_json::Value`: `Value` sorts object keys, so round-tripping + /// through it would reorder the payload and hide exactly the property + /// under test. + #[test] + fn v1_payloads_are_relayed_verbatim() { + let raw = std::fs::read_to_string(golden_v1()).unwrap(); + let replayed = read_completed_job(&golden_v1()).unwrap().unwrap(); + + assert!( + raw.contains(replayed.job.get()), + "the job payload must appear verbatim in the log, got: {}", + replayed.job.get() + ); + assert!( + raw.contains(replayed.stages.get()), + "the stages payload must appear verbatim in the log" + ); + // Key order survives, which a typed round trip would not preserve. + assert!( + replayed.job.get().starts_with(r#"{"job_id":"golden-v1""#), + "original key order must be preserved, got: {}", + replayed.job.get() + ); + } + + /// The v1 fixture deliberately contains a record kind this build has never + /// heard of, standing in for one a future scheduler might write. An older + /// reader must step over it and still find the terminal record. + #[test] + fn unknown_record_kinds_in_a_v1_log_do_not_break_reading() { + let raw = std::fs::read_to_string(golden_v1()).unwrap(); + assert!( + raw.contains("AnEventKindFromTheFuture"), + "fixture should exercise the unknown-record path" + ); + assert!(read_completed_job(&golden_v1()).unwrap().is_some()); + } + + /// The stages payload in the fixture carries `partition_id` as an array, + /// the shape that broke the TUI when it was declared as a scalar (#2257). + /// Storing payloads opaquely means a change like that can never make an + /// existing log unreadable. + #[test] + fn a_rest_shape_change_cannot_orphan_a_stored_log() { + let replayed = read_completed_job(&golden_v1()).unwrap().unwrap(); + assert!( + replayed.stages.get().contains(r#""partition_id":[0,1]"#), + "fixture should carry the multi-partition shape verbatim" + ); + } } diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs index de0e798fea..ff2e4d990a 100644 --- a/ballista/history/src/writer.rs +++ b/ballista/history/src/writer.rs @@ -146,7 +146,7 @@ async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { Some(f) => f, None => continue, }; - match serde_json::to_string(&*event) { + match event.to_record().and_then(|r| serde_json::to_string(&r)) { Ok(mut line) => { line.push('\n'); if let Err(e) = file.write_all(line.as_bytes()).await { @@ -203,8 +203,10 @@ async fn open_for<'a>( #[cfg(test)] mod tests { use super::*; - use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; - use ballista_api_types::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{ + HistoryEvent, JobEnd, JobEndStatus, JobIndex, JobStart, StageStart, + }; + use serde_json::value::RawValue; use std::collections::BTreeMap; #[tokio::test] @@ -216,51 +218,43 @@ mod tests { writer.append( "job-1", - HistoryEvent::JobStart { - version: SCHEMA_VERSION, + HistoryEvent::JobStart(JobStart { 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 { + HistoryEvent::StageStart(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, + let job_end = HistoryEvent::JobEnd(Box::new(JobEnd { status: JobEndStatus::Succeeded, queued_at: 1, started_at: 2, completed_at: 20, - job: Box::new(job), - stages: Box::new(QueryStagesResponse { stages: vec![] }), + index: JobIndex { + job_id: "job-1".into(), + job_name: "q1".into(), + status: "Completed".into(), + job_status: "COMPLETED".into(), + start_time: 10, + end_time: 20, + }, + job: RawValue::from_string(r#"{"job_id":"job-1"}"#.to_string()).unwrap(), + stages: RawValue::from_string(r#"{"stages":[]}"#.to_string()).unwrap(), config: BTreeMap::new(), dot: "digraph {}".into(), - }; + })); writer.append_final("job-1", job_end).await; writer.finish_job("job-1").await; @@ -284,22 +278,21 @@ mod tests { let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); writer.append( "job-1", - HistoryEvent::JobStart { - version: SCHEMA_VERSION, + HistoryEvent::JobStart(JobStart { 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 { + HistoryEvent::StageStart(StageStart { stage_id: 1, partitions: 4, - }, + }), ); writer.flush_job("job-1").await; diff --git a/ballista/history/testdata/schema-v1.eventlog b/ballista/history/testdata/schema-v1.eventlog new file mode 100644 index 0000000000..dd6c289a3c --- /dev/null +++ b/ballista/history/testdata/schema-v1.eventlog @@ -0,0 +1,7 @@ +{"ev":"JobStart","version":1,"data":{"job_id":"golden-v1","job_name":"tpch-q1","queued_at":1000,"submitted_at":1005,"logical_plan":"Projection: l_returnflag","physical_plan":"ProjectionExec: expr=[l_returnflag@0 as l_returnflag]"}} +{"ev":"StageStart","version":1,"data":{"stage_id":1,"partitions":4}} +{"ev":"TaskEnd","version":1,"data":{"stage_id":1,"task_id":0,"executor_id":"executor-a","status":"Successful","launch_time":1010,"start_exec_time":1012,"end_exec_time":1040,"metrics":{"input_rows":1000,"output_rows":250,"elapsed_compute_nanos":28000000}}} +{"ev":"TaskEnd","version":1,"data":{"stage_id":1,"task_id":1,"executor_id":"executor-b","status":{"Failed":{"reason":"ExecutionError","error":"divide by zero"}},"launch_time":1010,"start_exec_time":1013,"end_exec_time":1020,"metrics":{"input_rows":0,"output_rows":0,"elapsed_compute_nanos":0}}} +{"ev":"StageEnd","version":1,"data":{"stage_id":1,"status":"Successful"}} +{"ev":"AnEventKindFromTheFuture","version":1,"data":{"whatever":[1,2,{"nested":true}]}} +{"ev":"JobEnd","version":1,"data":{"status":"Succeeded","queued_at":1000,"started_at":1005,"completed_at":1100,"index":{"job_id":"golden-v1","job_name":"tpch-q1","status":"Completed","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","start_time":1005,"end_time":1100},"job":{"job_id":"golden-v1","job_name":"tpch-q1","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","status":"Completed","num_stages":1,"completed_stages":1,"percent_complete":100,"start_time":1005,"end_time":1100,"logical_plan":"Projection: l_returnflag","physical_plan":"ProjectionExec: expr=[l_returnflag@0 as l_returnflag]","stage_plan":"ShuffleWriterExec"},"stages":{"stages":[{"stage_id":"1","stage_status":"Successful","input_rows":1000,"output_rows":250,"elapsed_compute":"28ms","stage_plan":"ProjectionExec","task_duration_percentiles":{"min":7,"p25":7,"median":28,"p75":28,"max":28},"task_input_percentiles":{"min":0,"p25":0,"median":1000,"p75":1000,"max":1000},"tasks":[{"id":0,"status":"Successful","partition_id":[0,1],"scheduled_time":1008,"launch_time":1010,"start_exec_time":1012,"end_exec_time":1040,"exec_duration":28,"finish_time":1042,"input_rows":1000,"output_rows":250}]}]},"config":{"datafusion.execution.batch_size":"8192","datafusion.execution.target_partitions":"4"},"dot":"digraph G {\n stage_1;\n}"}} From 4e856b3fc7db6528803b8c46fc159ec8fc7deb1c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 9 Aug 2026 07:53:18 -0600 Subject: [PATCH 7/7] feat(history): carry the job-list fields in JobIndex The index exists so the history server can render GET /api/jobs without parsing the stored payloads, but it was missing num_stages, completed_stages and percent_complete, which that response includes. It could not actually serve the list it was there for. Adds the three fields and updates the v1 fixture to match. Doing this before release, while the schema is still unpublished, so the frozen fixture stays a faithful record of what v1 actually looks like. --- ballista/history/src/event.rs | 12 ++++++++++++ ballista/history/src/reader.rs | 3 +++ ballista/history/src/writer.rs | 3 +++ ballista/history/testdata/schema-v1.eventlog | 2 +- 4 files changed, 19 insertions(+), 1 deletion(-) diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs index b42bb504f2..dcd00af976 100644 --- a/ballista/history/src/event.rs +++ b/ballista/history/src/event.rs @@ -175,6 +175,9 @@ pub struct TaskEnd { /// local decouples the part that must stay readable forever from /// `ballista-api-types`, which evolves with the live REST contract. /// +/// These are exactly the fields `GET /api/jobs` renders, which is what lets the +/// history server build the job list without touching the payloads at all. +/// /// Adding fields here later is fine; each one needs `#[serde(default)]` so older /// logs still parse. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -191,6 +194,12 @@ pub struct JobIndex { pub start_time: u64, /// When the job reached its terminal state. pub end_time: u64, + /// Total number of stages in the job. + pub num_stages: usize, + /// Number of stages that finished successfully. + pub completed_stages: usize, + /// Progress as a percentage of completed stages. + pub percent_complete: u8, } /// Terminal record, and the only one the history server serves from. @@ -304,6 +313,9 @@ mod tests { job_status: "COMPLETED".into(), start_time: 10, end_time: 20, + num_stages: 1, + completed_stages: 1, + percent_complete: 100, }, job: RawValue::from_string(r#"{"job_id":"job-1"}"#.to_string()).unwrap(), stages: RawValue::from_string(r#"{"stages":[]}"#.to_string()).unwrap(), diff --git a/ballista/history/src/reader.rs b/ballista/history/src/reader.rs index 803ca2a9e6..36382e648e 100644 --- a/ballista/history/src/reader.rs +++ b/ballista/history/src/reader.rs @@ -150,6 +150,9 @@ mod tests { job_status: "COMPLETED".into(), start_time: 2, end_time: 3, + num_stages: 1, + completed_stages: 1, + percent_complete: 100, } } diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs index ff2e4d990a..ccf9939889 100644 --- a/ballista/history/src/writer.rs +++ b/ballista/history/src/writer.rs @@ -249,6 +249,9 @@ mod tests { job_status: "COMPLETED".into(), start_time: 10, end_time: 20, + num_stages: 1, + completed_stages: 1, + percent_complete: 100, }, job: RawValue::from_string(r#"{"job_id":"job-1"}"#.to_string()).unwrap(), stages: RawValue::from_string(r#"{"stages":[]}"#.to_string()).unwrap(), diff --git a/ballista/history/testdata/schema-v1.eventlog b/ballista/history/testdata/schema-v1.eventlog index dd6c289a3c..9d15a69e84 100644 --- a/ballista/history/testdata/schema-v1.eventlog +++ b/ballista/history/testdata/schema-v1.eventlog @@ -4,4 +4,4 @@ {"ev":"TaskEnd","version":1,"data":{"stage_id":1,"task_id":1,"executor_id":"executor-b","status":{"Failed":{"reason":"ExecutionError","error":"divide by zero"}},"launch_time":1010,"start_exec_time":1013,"end_exec_time":1020,"metrics":{"input_rows":0,"output_rows":0,"elapsed_compute_nanos":0}}} {"ev":"StageEnd","version":1,"data":{"stage_id":1,"status":"Successful"}} {"ev":"AnEventKindFromTheFuture","version":1,"data":{"whatever":[1,2,{"nested":true}]}} -{"ev":"JobEnd","version":1,"data":{"status":"Succeeded","queued_at":1000,"started_at":1005,"completed_at":1100,"index":{"job_id":"golden-v1","job_name":"tpch-q1","status":"Completed","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","start_time":1005,"end_time":1100},"job":{"job_id":"golden-v1","job_name":"tpch-q1","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","status":"Completed","num_stages":1,"completed_stages":1,"percent_complete":100,"start_time":1005,"end_time":1100,"logical_plan":"Projection: l_returnflag","physical_plan":"ProjectionExec: expr=[l_returnflag@0 as l_returnflag]","stage_plan":"ShuffleWriterExec"},"stages":{"stages":[{"stage_id":"1","stage_status":"Successful","input_rows":1000,"output_rows":250,"elapsed_compute":"28ms","stage_plan":"ProjectionExec","task_duration_percentiles":{"min":7,"p25":7,"median":28,"p75":28,"max":28},"task_input_percentiles":{"min":0,"p25":0,"median":1000,"p75":1000,"max":1000},"tasks":[{"id":0,"status":"Successful","partition_id":[0,1],"scheduled_time":1008,"launch_time":1010,"start_exec_time":1012,"end_exec_time":1040,"exec_duration":28,"finish_time":1042,"input_rows":1000,"output_rows":250}]}]},"config":{"datafusion.execution.batch_size":"8192","datafusion.execution.target_partitions":"4"},"dot":"digraph G {\n stage_1;\n}"}} +{"ev":"JobEnd","version":1,"data":{"status":"Succeeded","queued_at":1000,"started_at":1005,"completed_at":1100,"index":{"job_id":"golden-v1","job_name":"tpch-q1","status":"Completed","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","start_time":1005,"end_time":1100,"num_stages":1,"completed_stages":1,"percent_complete":100},"job":{"job_id":"golden-v1","job_name":"tpch-q1","job_status":"Completed. Produced 4 partitions containing 1000 rows. Elapsed time: 95 ms.","status":"Completed","num_stages":1,"completed_stages":1,"percent_complete":100,"start_time":1005,"end_time":1100,"logical_plan":"Projection: l_returnflag","physical_plan":"ProjectionExec: expr=[l_returnflag@0 as l_returnflag]","stage_plan":"ShuffleWriterExec"},"stages":{"stages":[{"stage_id":"1","stage_status":"Successful","input_rows":1000,"output_rows":250,"elapsed_compute":"28ms","stage_plan":"ProjectionExec","task_duration_percentiles":{"min":7,"p25":7,"median":28,"p75":28,"max":28},"task_input_percentiles":{"min":0,"p25":0,"median":1000,"p75":1000,"max":1000},"tasks":[{"id":0,"status":"Successful","partition_id":[0,1],"scheduled_time":1008,"launch_time":1010,"start_exec_time":1012,"end_exec_time":1040,"exec_duration":28,"finish_time":1042,"input_rows":1000,"output_rows":250}]}]},"config":{"datafusion.execution.batch_size":"8192","datafusion.execution.target_partitions":"4"},"dot":"digraph G {\n stage_1;\n}"}}