diff --git a/.cursor/rules.md b/.cursor/rules.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/.cursor/rules.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/.gemini/rules.md b/.gemini/rules.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/.gemini/rules.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/Cargo.lock b/Cargo.lock index 369a967188..a432314ef0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1168,6 +1168,16 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "53.0.0" +dependencies = [ + "serde", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "ballista-scheduler" version = "53.0.0" @@ -1176,6 +1186,7 @@ dependencies = [ "async-trait", "axum", "ballista-core", + "ballista-history", "clap 4.6.1", "dashmap", "datafusion", @@ -1196,11 +1207,14 @@ dependencies = [ "rand 0.10.1", "rstest", "serde", + "serde_json", + "tempfile", "tokio", "tokio-stream", "tonic", "tonic-prost", "tonic-prost-build", + "tower", "tower-http 0.7.0", "tracing", "tracing-appender", diff --git a/Cargo.toml b/Cargo.toml index 5719b9b16c..71e28483b2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ [workspace] exclude = ["dev/msrvcheck", "python"] -members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/scheduler", "benchmarks", "examples"] +members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/history", "ballista/scheduler", "benchmarks", "examples"] resolver = "3" [workspace.package] diff --git a/ballista/history/Cargo.toml b/ballista/history/Cargo.toml new file mode 100644 index 0000000000..873f4b44a2 --- /dev/null +++ b/ballista/history/Cargo.toml @@ -0,0 +1,33 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "ballista-history" +description = "Event-log schema, writer, and reader for Ballista's history server" +license = "Apache-2.0" +version = "53.0.0" +edition = { workspace = true } +rust-version = { workspace = true } + +[dependencies] +serde = { workspace = true, features = ["derive"] } +serde_json = "1" +tokio = { workspace = true, features = ["rt", "sync", "fs", "io-util", "macros"] } + +[dev-dependencies] +tokio = { workspace = true, features = ["rt", "sync", "fs", "io-util", "macros", "rt-multi-thread"] } +tempfile = { workspace = true } diff --git a/ballista/history/src/dto.rs b/ballista/history/src/dto.rs new file mode 100644 index 0000000000..3485bfc00f --- /dev/null +++ b/ballista/history/src/dto.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! REST response DTOs shared by the live scheduler API handlers and the history +//! server, so both serialize byte-identical JSON. + +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobResponse { + pub job_id: String, + pub job_name: String, + pub job_status: String, + pub status: String, + pub num_stages: usize, + pub completed_stages: usize, + pub percent_complete: u8, + pub start_time: u64, + pub end_time: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub logical_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub physical_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum TaskStatus { + Running, + Successful, + Failed, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskSummary { + pub id: usize, + pub status: TaskStatus, + pub partition_id: u32, + pub scheduled_time: u64, + pub launch_time: u64, + pub start_exec_time: u64, + pub end_exec_time: u64, + pub exec_duration: u64, + pub finish_time: u64, + pub input_rows: usize, + pub output_rows: usize, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Percentiles { + pub min: u64, + pub p25: u64, + pub median: u64, + pub p75: u64, + pub max: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStageSummary { + pub stage_id: String, + pub stage_status: String, + pub input_rows: usize, + pub output_rows: usize, + pub elapsed_compute: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_duration_percentiles: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_input_percentiles: Option, + pub tasks: Vec>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStagesResponse { + pub stages: Vec, +} + +/// Session config as flat key/value pairs (from `SessionConfig::to_props()`), +/// sorted for stable output. +pub type JobConfig = BTreeMap; diff --git a/ballista/history/src/event.rs b/ballista/history/src/event.rs new file mode 100644 index 0000000000..8db1c05f49 --- /dev/null +++ b/ballista/history/src/event.rs @@ -0,0 +1,124 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The on-disk event-log schema. One `HistoryEvent` is serialized per JSONL line. +//! This is a frozen public projection of the scheduler's internal events; the +//! embedded DTOs are the stable contract the history server serves. + +use crate::dto::{JobConfig, JobResponse, QueryStagesResponse, TaskStatus}; +use serde::{Deserialize, Serialize}; + +/// Current on-disk schema version, stamped on `JobStart`/`JobEnd`. +pub const SCHEMA_VERSION: u32 = 1; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum JobEndStatus { + Succeeded, + Failed(String), +} + +/// Metrics captured per finished task on the incremental timeline. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskEndMetrics { + pub input_rows: u64, + pub output_rows: u64, + pub elapsed_compute_nanos: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "ev")] +pub enum HistoryEvent { + JobStart { + version: u32, + job_id: String, + job_name: String, + queued_at: u64, + submitted_at: u64, + logical_plan: Option, + physical_plan: Option, + }, + StageStart { + stage_id: usize, + partitions: usize, + }, + StageEnd { + stage_id: usize, + status: String, + }, + TaskEnd { + stage_id: u32, + partition: u32, + executor_id: String, + status: TaskStatus, + launch_time: u64, + start_exec_time: u64, + end_exec_time: u64, + metrics: TaskEndMetrics, + }, + JobEnd { + version: u32, + status: JobEndStatus, + queued_at: u64, + started_at: u64, + completed_at: u64, + job: Box, + stages: Box, + config: JobConfig, + dot: String, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::{JobResponse, QueryStagesResponse}; + use std::collections::BTreeMap; + + #[test] + fn job_end_round_trips_through_jsonl() { + let job = JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 2, + completed_stages: 2, + percent_complete: 100, + start_time: 10, + end_time: 20, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage plan".into()), + }; + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 5, + started_at: 10, + completed_at: 20, + job: Box::new(job), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: BTreeMap::from([("k".to_string(), "v".to_string())]), + dot: "digraph {}".into(), + }; + let line = serde_json::to_string(&event).unwrap(); + assert!(line.contains("\"ev\":\"JobEnd\"")); + let back: HistoryEvent = serde_json::from_str(&line).unwrap(); + // Re-serialize and compare strings (stable, discriminating). + assert_eq!(line, serde_json::to_string(&back).unwrap()); + } +} diff --git a/ballista/history/src/lib.rs b/ballista/history/src/lib.rs new file mode 100644 index 0000000000..0b17d91d6b --- /dev/null +++ b/ballista/history/src/lib.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod dto; +pub mod event; +pub mod reader; +pub mod writer; diff --git a/ballista/history/src/reader.rs b/ballista/history/src/reader.rs new file mode 100644 index 0000000000..d4b60f210e --- /dev/null +++ b/ballista/history/src/reader.rs @@ -0,0 +1,128 @@ +// 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::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use crate::event::HistoryEvent; +use std::io::BufRead; +use std::path::Path; + +#[derive(Debug, Clone)] +pub struct ReplayedJob { + pub job: JobResponse, + pub stages: QueryStagesResponse, + pub config: JobConfig, + pub dot: String, +} + +pub fn read_completed_job(path: &Path) -> std::io::Result> { + let file = std::fs::File::open(path)?; + let reader = std::io::BufReader::new(file); + for line in reader.lines() { + let line = line?; + if line.is_empty() { + continue; + } + // Only JobEnd carries the served payload; other lines are the timeline + // and are ignored here. Unknown/garbled lines are skipped, not fatal. + if let Ok(HistoryEvent::JobEnd { + job, + stages, + config, + dot, + .. + }) = serde_json::from_str::(&line) + { + return Ok(Some(ReplayedJob { + job: *job, + stages: *stages, + config, + dot, + })); + } + } + Ok(None) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::io::Write; + + fn job_end_line() -> String { + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 1, + started_at: 2, + completed_at: 3, + job: Box::new(JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 1, + completed_stages: 1, + percent_complete: 100, + start_time: 2, + end_time: 3, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage".into()), + }), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: Default::default(), + dot: "digraph {}".into(), + }; + serde_json::to_string(&event).unwrap() + } + + #[test] + fn reads_job_end_and_ignores_unknown_timeline_lines() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("job-1.eventlog"); + let mut f = std::fs::File::create(&path).unwrap(); + // Unknown/other event lines before the JobEnd must be tolerated. + writeln!(f, r#"{{"ev":"StageStart","stage_id":1,"partitions":4}}"#).unwrap(); + writeln!(f, "{}", job_end_line()).unwrap(); + drop(f); + + let replayed = read_completed_job(&path).unwrap().expect("completed"); + assert_eq!(replayed.job.job_id, "job-1"); + assert_eq!( + replayed.job.physical_plan.as_deref(), + Some("ProjectionExec") + ); + assert_eq!(replayed.dot, "digraph {}"); + } + + #[test] + fn returns_none_when_no_job_end() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("job-2.eventlog"); + std::fs::write( + &path, + "{\"ev\":\"StageStart\",\"stage_id\":1,\"partitions\":4}\n", + ) + .unwrap(); + assert!(read_completed_job(&path).unwrap().is_none()); + } +} diff --git a/ballista/history/src/writer.rs b/ballista/history/src/writer.rs new file mode 100644 index 0000000000..61ae0a8259 --- /dev/null +++ b/ballista/history/src/writer.rs @@ -0,0 +1,303 @@ +// 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<()>, + }, +} + +#[derive(Clone)] +pub struct EventLogWriter { + tx: mpsc::Sender, +} + +impl EventLogWriter { + pub fn new(log_dir: PathBuf, buffer: usize) -> EventLogWriter { + let (tx, rx) = mpsc::channel(buffer.max(1)); + tokio::spawn(run(log_dir, rx)); + EventLogWriter { tx } + } + + /// Enqueue an event for `job_id`. Never blocks; drops (with a warning) if the + /// channel is full, so logging cannot stall scheduling. + pub fn append(&self, job_id: &str, event: HistoryEvent) { + let msg = WriterMsg::Event { + job_id: job_id.to_string(), + event: Box::new(event), + }; + if self.tx.try_send(msg).is_err() { + eprintln!( + "event-log writer: dropping event for {job_id} (channel full or closed)" + ); + } + } + + /// Await all currently-enqueued writes for `job_id` (best effort). + pub async fn flush_job(&self, job_id: &str) { + let (done, wait) = oneshot::channel(); + if self + .tx + .send(WriterMsg::Flush { + job_id: job_id.to_string(), + done, + }) + .await + .is_ok() + { + let _ = wait.await; + } + } + + /// Enqueue a terminal event (e.g. `JobEnd`) for `job_id`. Unlike `append`, this + /// awaits channel capacity instead of dropping the event when the channel is + /// full, so the terminal record is never silently lost. Still best-effort at + /// the process boundary: if the channel is closed (background task gone) this + /// logs and returns rather than panicking. + pub async fn append_final(&self, job_id: &str, event: HistoryEvent) { + let msg = WriterMsg::Event { + job_id: job_id.to_string(), + event: Box::new(event), + }; + if self.tx.send(msg).await.is_err() { + eprintln!( + "event-log writer: failed to enqueue terminal event for {job_id} (channel closed)" + ); + } + } + + /// Flush and close the per-job file handle for `job_id`. Must be called after + /// the terminal event has been enqueued (e.g. via `append_final`) so it is + /// ordered after it on the single-consumer FIFO channel. Best effort: if the + /// channel is closed this logs and returns. + pub async fn finish_job(&self, job_id: &str) { + let (done, wait) = oneshot::channel(); + if self + .tx + .send(WriterMsg::Finish { + job_id: job_id.to_string(), + done, + }) + .await + .is_ok() + { + let _ = wait.await; + } else { + eprintln!( + "event-log writer: failed to enqueue finish for {job_id} (channel closed)" + ); + } + } +} + +async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver) { + if let Err(e) = tokio::fs::create_dir_all(&log_dir).await { + eprintln!("event-log writer: cannot create {}: {e}", log_dir.display()); + return; + } + // One open append handle per job for the life of the process. + let mut handles: HashMap = HashMap::new(); + + while let Some(msg) = rx.recv().await { + match msg { + WriterMsg::Event { job_id, event } => { + let file = match open_for(&log_dir, &mut handles, &job_id).await { + Some(f) => f, + None => continue, + }; + match serde_json::to_string(&*event) { + Ok(mut line) => { + line.push('\n'); + if let Err(e) = file.write_all(line.as_bytes()).await { + eprintln!("event-log writer: write failed for {job_id}: {e}"); + } + } + Err(e) => eprintln!("event-log writer: serialize failed: {e}"), + } + } + WriterMsg::Flush { job_id, done } => { + if let Some(file) = handles.get_mut(&job_id) { + let _ = file.flush().await; + } + let _ = done.send(()); + } + 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) => { + eprintln!("event-log writer: cannot open {}: {e}", path.display()); + return None; + } + } + } + handles.get_mut(job_id) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::{JobResponse, QueryStagesResponse}; + use crate::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::collections::BTreeMap; + + #[tokio::test] + async fn terminal_job_end_is_not_dropped_on_a_saturated_channel() { + let dir = tempfile::tempdir().unwrap(); + // Tiny buffer so the non-blocking `append` path would readily drop events + // under load; `append_final` must still guarantee delivery. + let writer = EventLogWriter::new(dir.path().to_path_buf(), 1); + + writer.append( + "job-1", + HistoryEvent::JobStart { + version: SCHEMA_VERSION, + job_id: "job-1".into(), + job_name: "q1".into(), + queued_at: 1, + submitted_at: 2, + logical_plan: None, + physical_plan: None, + }, + ); + for stage_id in 0..10 { + writer.append( + "job-1", + HistoryEvent::StageStart { + stage_id, + partitions: 4, + }, + ); + } + + let job = JobResponse { + job_id: "job-1".into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 2, + completed_stages: 2, + percent_complete: 100, + start_time: 10, + end_time: 20, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage plan".into()), + }; + let job_end = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 1, + started_at: 2, + completed_at: 20, + job: Box::new(job), + stages: Box::new(QueryStagesResponse { stages: vec![] }), + config: BTreeMap::new(), + dot: "digraph {}".into(), + }; + writer.append_final("job-1", job_end).await; + writer.finish_job("job-1").await; + + let path = dir.path().join("job-1.eventlog"); + let contents = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + assert!( + lines.iter().any(|l| l.contains("\"ev\":\"JobEnd\"")), + "expected a JobEnd line in the event log, got: {contents}" + ); + assert_eq!( + lines.last().map(|l| l.contains("\"ev\":\"JobEnd\"")), + Some(true), + "JobEnd should be the last line written" + ); + } + + #[tokio::test] + async fn append_writes_one_jsonl_line_per_event() { + 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/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index f882d26f38..2b44218189 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -32,6 +32,11 @@ name = "ballista-scheduler" path = "src/bin/main.rs" required-features = ["build-binary"] +[[bin]] +name = "ballista-history-server" +path = "src/bin/history_server.rs" +required-features = ["build-binary", "rest-api"] + [features] build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"] default = ["build-binary", "rest-api"] @@ -41,7 +46,7 @@ disable-stage-plan-cache = [] graphviz-support = ["dep:graphviz-rust"] keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"] prometheus-metrics = ["prometheus", "once_cell"] -rest-api = [] +rest-api = ["dep:ballista-history", "dep:serde_json"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -51,6 +56,7 @@ async-trait = { workspace = true } axum = "0.8.9" tower-http = { version = "0.7", features = ["cors"] } ballista-core = { path = "../core", version = "53.0.0" } +ballista-history = { path = "../history", version = "53.0.0", optional = true } clap = { workspace = true, optional = true } dashmap = { workspace = true } datafusion = { workspace = true } @@ -70,6 +76,7 @@ prost = { workspace = true } prost-types = { workspace = true } rand = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = { version = "1", optional = true } tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true, features = ["net"] } tonic = { workspace = true, features = ["router"] } @@ -81,6 +88,9 @@ uuid = { workspace = true } [dev-dependencies] rstest = { workspace = true } +serde_json = "1" +tempfile = { workspace = true } +tower = "0.5" [build-dependencies] tonic-prost-build = { workspace = true, optional = true } diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs new file mode 100644 index 0000000000..179aa47f3e --- /dev/null +++ b/ballista/scheduler/src/api/dto_build.rs @@ -0,0 +1,654 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Builds the shared REST DTOs (`ballista_history::dto`) from the scheduler's +//! internal execution-graph state. These builders back both the live REST API +//! handlers (`api::handlers`) and, eventually, the event-log writer, so both +//! serialize byte-identical JSON for the same graph state. + +use crate::api::handlers::{JobQueryParams, PlanFormat}; +use crate::display::format_stage_metrics; +use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage}; +use crate::state::execution_graph_dot::ExecutionGraphDot; +use crate::state::execution_stage::TaskInfo; +use crate::state::task_manager::JobOverview; +use ballista_core::serde::protobuf::job_status::Status; +use ballista_core::serde::protobuf::{OperatorMetricsSet, task_status}; +use ballista_core::utils::get_current_time; +use ballista_history::dto::{ + JobResponse, Percentiles, QueryStageSummary, QueryStagesResponse, TaskStatus, + TaskSummary, +}; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::displayable; +use datafusion::physical_plan::metrics::{MetricsSet, Time}; +use std::fmt; +use std::time::Duration; + +/// Builds the `JobResponse` DTO for a job's execution graph. +/// +/// `with_plans` controls whether `logical_plan`/`physical_plan`/`stage_plan` are +/// populated (the single-job detail view, and eventually the history detail +/// view) or left `None` (list views, which use +/// [`build_job_response_from_overview`] instead). `plan_format` selects how the +/// physical plan is rendered when `with_plans` is true; it is ignored otherwise. +pub(crate) fn build_job_response( + graph: &ExecutionGraphBox, + with_plans: bool, + plan_format: PlanFormat, +) -> JobResponse { + let job = graph.as_ref(); + let (plain_status, job_status) = format_job_status( + &job.status().status, + job_elapsed_ms(job.start_time(), job.end_time()), + ); + + let num_stages = job.stage_count(); + let completed_stages = job.completed_stages(); + let percent_complete = + ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; + + let (logical_plan, physical_plan, stage_plan) = if with_plans { + let physical_plan = match plan_format { + PlanFormat::Default | PlanFormat::Metrics => { + DisplayableExecutionPlan::new(job.physical_plan().as_ref()) + .indent(false) + .to_string() + } + PlanFormat::Tree => displayable(job.physical_plan().as_ref()) + .tree_render() + .to_string(), + }; + ( + job.logical_plan().map(str::to_owned), + Some(physical_plan), + Some(format!("{:?}", graph)), + ) + } else { + (None, None, None) + }; + + JobResponse { + job_id: job.job_id().to_string(), + job_name: job.job_name().to_owned(), + job_status, + status: plain_status, + start_time: job.start_time(), + end_time: job.end_time(), + num_stages, + completed_stages, + percent_complete, + logical_plan, + physical_plan, + stage_plan, + } +} + +/// Builds the `JobResponse` DTO for a job's cached [`JobOverview`] (no +/// execution graph, and hence no plans available). +pub(crate) fn build_job_response_from_overview(job: &JobOverview) -> JobResponse { + let (plain_status, job_status) = format_job_status( + &job.status.status, + job_elapsed_ms(job.start_time, job.end_time), + ); + + // calculate progress based on completed stages for now, but we could use completed + // tasks in the future to make this more accurate + let percent_complete = if job.num_stages == 0 { + 0 + } else { + ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 + }; + JobResponse { + job_id: job.job_id.to_string(), + job_name: job.job_name.to_owned(), + job_status, + status: plain_status, + start_time: job.start_time, + end_time: job.end_time, + num_stages: job.num_stages, + completed_stages: job.completed_stages, + percent_complete, + logical_plan: None, + physical_plan: None, + stage_plan: None, + } +} + +/// Builds the `QueryStagesResponse` DTO for a job's execution graph. +pub(crate) fn build_query_stages_response( + graph: &ExecutionGraphBox, + params: &JobQueryParams, +) -> QueryStagesResponse { + let plan_format = params.plan_format.clone().unwrap_or_default(); + + let stages = graph + .as_ref() + .stages() + .iter() + .map(|(id, stage)| { + let mut summary = QueryStageSummary { + stage_id: id.to_string(), + stage_status: stage.variant_name().to_string(), + input_rows: 0, + output_rows: 0, + elapsed_compute: None, + tasks: vec![], + task_duration_percentiles: None, + task_input_percentiles: None, + stage_plan: None, + }; + match stage { + ExecutionStage::Running(running_stage) => { + let metrics = running_stage.stage_metrics.as_deref().unwrap_or(&[]); + summary.stage_plan = Some(match plan_format { + PlanFormat::Default => displayable(running_stage.plan.as_ref()) + .indent(false) + .to_string(), + PlanFormat::Tree => displayable(running_stage.plan.as_ref()) + .tree_render() + .to_string(), + PlanFormat::Metrics => { + format_stage_metrics(running_stage.plan.as_ref(), metrics) + } + }); + summary.input_rows = running_stage + .stage_metrics + .as_ref() + .map(|m| get_combined_count(m.as_slice(), "input_rows")) + .unwrap_or(0); + summary.output_rows = running_stage + .stage_metrics + .as_ref() + .map(|m| get_combined_count(m.as_slice(), "output_rows")) + .unwrap_or(0); + summary.elapsed_compute = get_running_stage_time( + &running_stage.task_infos, + get_current_time(), + ); + summary.tasks = running_stage + .task_infos + .iter() + .enumerate() + .map(|(partition_id, task_info)| { + task_info.as_ref().map(|info| { + let (input_rows, output_rows) = running_stage + .stage_metrics + .as_deref() + .map(|metrics| { + get_partition_counts(metrics, partition_id) + }) + .unwrap_or((0, 0)); + + let start_exec_time = info.start_exec_time as u64; + let end_exec_time = info.end_exec_time as u64; + + let task_status = to_api_task_status(&info.task_status); + + TaskSummary { + id: info.task_id, + partition_id: partition_id as u32, + scheduled_time: info.scheduled_time as u64, + launch_time: info.launch_time as u64, + start_exec_time, + end_exec_time, + exec_duration: end_exec_time + .saturating_sub(start_exec_time), + finish_time: info.finish_time as u64, + input_rows, + output_rows, + status: task_status, + } + }) + }) + .collect(); + } + ExecutionStage::Successful(completed_stage) => { + summary.stage_plan = Some(match plan_format { + PlanFormat::Default => displayable(completed_stage.plan.as_ref()) + .indent(false) + .to_string(), + PlanFormat::Tree => displayable(completed_stage.plan.as_ref()) + .tree_render() + .to_string(), + PlanFormat::Metrics => format_stage_metrics( + completed_stage.plan.as_ref(), + &completed_stage.stage_metrics, + ), + }); + summary.input_rows = + get_combined_count(&completed_stage.stage_metrics, "input_rows"); + summary.output_rows = + get_combined_count(&completed_stage.stage_metrics, "output_rows"); + summary.elapsed_compute = + get_finished_stage_time(&completed_stage.task_infos); + + summary.tasks = completed_stage + .task_infos + .iter() + .enumerate() + .map(|(partition_id, task_info)| { + let (input_rows, output_rows) = get_partition_counts( + &completed_stage.stage_metrics, + partition_id, + ); + + let start_exec_time = task_info.start_exec_time as u64; + let end_exec_time = task_info.end_exec_time as u64; + let task_status = to_api_task_status(&task_info.task_status); + Some(TaskSummary { + id: task_info.task_id, + partition_id: partition_id as u32, + scheduled_time: task_info.scheduled_time as u64, + launch_time: task_info.launch_time as u64, + start_exec_time, + end_exec_time, + exec_duration: end_exec_time + .saturating_sub(start_exec_time), + finish_time: task_info.finish_time as u64, + input_rows, + output_rows, + status: task_status, + }) + }) + .collect(); + } + _ => {} + } + summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); + summary.task_input_percentiles = task_input_percentiles(&summary.tasks); + summary + }) + .collect(); + + QueryStagesResponse { stages } +} + +/// Builds the DOT-format graph for a job's execution graph. +pub(crate) fn build_job_dot(graph: &ExecutionGraphBox) -> Result { + ExecutionGraphDot::generate(graph.as_ref()) +} + +/// Converts a task's protobuf status into the shared `TaskStatus` DTO. +pub(crate) fn to_api_task_status(status: &task_status::Status) -> TaskStatus { + match status { + task_status::Status::Running(_) => TaskStatus::Running, + task_status::Status::Failed(_) => TaskStatus::Failed, + task_status::Status::Successful(_) => TaskStatus::Successful, + } +} + +/// Sums a single task's raw operator metrics into +/// `(input_rows, output_rows, elapsed_compute_nanos)`. +/// +/// Note: this is *not* a literal extraction of `get_partition_counts`/ +/// `get_combined_count` below. Those operate on the stage's already-merged +/// per-partition [`MetricsSet`]s (DataFusion's own metrics type) and filter by +/// partition id, returning `usize` row counts only. `task_row_counts` instead +/// converts a single task's raw protobuf [`OperatorMetricsSet`]s (as reported +/// by an executor for one task, e.g. `TaskStatus::metrics`) directly, with no +/// partition filtering (a task already corresponds to exactly one partition), +/// and additionally sums `elapsed_compute` (nanoseconds) alongside the row +/// counts, matching `ballista_history::event::TaskEndMetrics` for the +/// event-log's timeline events (see `scheduler_server::event_log`). +pub(crate) fn task_row_counts(metrics: &[OperatorMetricsSet]) -> (u64, u64, u64) { + let mut input_rows: u64 = 0; + let mut output_rows: u64 = 0; + let mut elapsed_compute_nanos: u64 = 0; + + for operator_metrics in metrics { + let Ok(metrics_set) = TryInto::::try_into(operator_metrics.clone()) + else { + continue; + }; + for metric in metrics_set.iter() { + let value = metric.value(); + match value.name() { + "input_rows" => input_rows += value.as_usize() as u64, + "output_rows" => output_rows += value.as_usize() as u64, + "elapsed_compute" => elapsed_compute_nanos += value.as_usize() as u64, + _ => {} + } + } + } + + (input_rows, output_rows, elapsed_compute_nanos) +} + +fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { + let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; + sorted[idx.min(sorted.len() - 1)] +} + +fn task_input_percentiles(tasks: &[Option]) -> Option { + let mut durations: Vec = tasks + .iter() + .flatten() + .map(|t| t.input_rows as u64) + .collect(); + + if durations.is_empty() { + return None; + } + + durations.sort_unstable(); + + Some(Percentiles { + min: durations[0], + p25: percentile_duration(&durations, 25.0), + median: percentile_duration(&durations, 50.0), + p75: percentile_duration(&durations, 75.0), + max: *durations.last().unwrap(), + }) +} + +fn task_duration_percentiles(tasks: &[Option]) -> Option { + let mut durations: Vec = + tasks.iter().flatten().map(|t| t.exec_duration).collect(); + + if durations.is_empty() { + return None; + } + + durations.sort_unstable(); + + Some(Percentiles { + min: durations[0], + p25: percentile_duration(&durations, 25.0), + median: percentile_duration(&durations, 50.0), + p75: percentile_duration(&durations, 75.0), + max: *durations.last().unwrap(), + }) +} + +/// Returns elapsed wall time in milliseconds for API formatting. +/// +/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or +/// `end_time` still zero while `start_time` is set) do not panic on subtract. +fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { + end_time.saturating_sub(start_time) +} + +fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { + match status { + Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), + Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), + Some(Status::Failed(error)) => { + ("Failed".to_string(), format!("Failed: {}", error.error)) + } + Some(Status::Successful(completed)) => { + let num_rows = completed + .partition_location + .iter() + .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) + .sum::(); + let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; + let num_partitions = completed.partition_location.len(); + let num_partitions_term = if num_partitions == 1 { + "partition" + } else { + "partitions" + }; + ( + "Completed".to_string(), + format!( + "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", + num_partitions, + num_partitions_term, + num_rows, + num_rows_term, + elapsed_ms + ), + ) + } + _ => ("Invalid".to_string(), "Invalid State".to_string()), + } +} + +fn get_running_stage_time( + task_infos: &[Option], + current_time: u128, +) -> Option { + let min_start = task_infos + .iter() + .flat_map(|t| t.as_ref().map(|t| t.start_exec_time)) + .filter(|t| *t > 0) + .min(); + + match (min_start, current_time) { + (Some(start), end) if end >= start => { + let time = Time::new(); + time.add_duration(Duration::from_millis((end - start) as u64)); + Some(time.to_string()) + } + _ => None, + } +} + +fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { + let min_start = task_infos + .iter() + .map(|t| t.start_exec_time) + .filter(|t| *t > 0) + .min(); + + let max_end = task_infos + .iter() + .map(|t| t.end_exec_time) + .filter(|t| *t > 0) + .max(); + + match (min_start, max_end) { + (Some(start), Some(end)) if end >= start => { + let time = Time::new(); + time.add_duration(Duration::from_millis((end - start) as u64)); + Some(time.to_string()) + } + _ => None, + } +} + +fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) { + let input_rows = get_partition_count(metrics, partition_id, "input_rows"); + let output_rows = get_partition_count(metrics, partition_id, "output_rows"); + (input_rows, output_rows) +} + +fn get_partition_count(metrics: &[MetricsSet], partition_id: usize, name: &str) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + if metric.partition() == Some(partition_id) && metric_value.name() == name + { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + if metric_value.name() == name { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::execution_graph_dot::tests::test_graph; + use ballista_core::serde::protobuf::task_status; + + fn make_task_info(start: u128, end: u128) -> TaskInfo { + TaskInfo { + task_id: 0, + scheduled_time: 0, + launch_time: 0, + start_exec_time: start, + end_exec_time: end, + finish_time: 0, + task_status: task_status::Status::Running(Default::default()), + } + } + + #[test] + fn test_job_elapsed_saturates_when_end_precedes_start() { + assert_eq!(job_elapsed_ms(900, 100), 0); + } + + // --- get_finished_stage_time --- + + #[test] + fn test_finished_empty_slice_returns_none() { + assert_eq!(get_finished_stage_time(&[]), None); + } + + #[test] + fn test_finished_all_zero_timestamps_returns_none() { + let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + #[test] + fn test_finished_single_task_elapsed() { + // 600 - 100 = 500 ms → "500.00ms" + let tasks = vec![make_task_info(100, 600)]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("500.00ms".to_string()) + ); + } + + #[test] + fn test_finished_picks_earliest_start_and_latest_end() { + // min start = 100, max end = 900 → 800 ms + let tasks = vec![ + make_task_info(100, 500), + make_task_info(200, 900), + make_task_info(300, 700), + ]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("800.00ms".to_string()) + ); + } + + #[test] + fn test_finished_end_before_start_returns_none() { + let tasks = vec![make_task_info(900, 100)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + // --- get_running_stage_time --- + + #[test] + fn test_running_empty_slice_returns_none() { + assert_eq!(get_running_stage_time(&[], 1000), None); + } + + #[test] + fn test_running_all_none_returns_none() { + let tasks: Vec> = vec![None, None]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_future_start_returns_none() { + // start_exec_time beyond current time → elapsed clamped to 0 + let tasks = vec![Some(make_task_info(u128::MAX, 0))]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_past_start_returns_some() { + let now = 4_000; + let start = 1_000; + let tasks = vec![Some(make_task_info(start, 0))]; + assert_eq!( + get_running_stage_time(&tasks, now), + Some("3.00s".to_string()) + ); + } + + #[test] + fn test_running_mixed_some_none_uses_earliest_some() { + let now = 3_000; + let earlier = 1_000; + let later = 2_000; + let tasks = vec![ + None, + Some(make_task_info(later, 0)), + Some(make_task_info(earlier, 0)), + None, + ]; + let result = get_running_stage_time(&tasks, now); + assert_eq!(result, Some("2.00s".to_string())); + } + + #[test] + fn test_job_elapsed_ms_normal() { + assert_eq!(super::job_elapsed_ms(100, 500), 400); + } + + #[test] + fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { + assert_eq!(super::job_elapsed_ms(500, 100), 0); + } + + #[test] + fn test_task_row_counts_empty_is_zero() { + assert_eq!(task_row_counts(&[]), (0, 0, 0)); + } + + #[tokio::test] + async fn test_dto_builders_plan_string() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + // build_query_stages_response: sanity-check the stage topology survives + // the DTO builder + JSON serialization round trip. test_graph() is a + // freshly built (never executed) graph, so every stage is + // Resolved/Unresolved and `stage_plan` is intentionally left `None` + // (only Running/Successful stages populate it) -- so this builder alone + // can't be used for the plan-string assertion below. + let params = JobQueryParams::default(); + let stages_response = build_query_stages_response(&graph, ¶ms); + let stages_json = serde_json::to_string(&stages_response).unwrap(); + assert!( + stages_json.contains(r#""stage_status":"Resolved""#), + "expected a resolved stage in response, got: {stages_json}" + ); + + // build_job_dot: assert on the embedded physical-plan string, following + // the repo's plan-string assertion convention (mirrors + // `execution_graph_dot::tests::dot()`, which renders the same graph). + let dot = build_job_dot(&graph).unwrap(); + assert!( + dot.contains("DataSourceExec: (Memory)"), + "expected plan string in dot graph, got: {dot}" + ); + } +} diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index ff886a69f8..4aae6ea9d5 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -10,11 +10,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::display::format_stage_metrics; +use crate::api::dto_build; use crate::scheduler_server::event::QueryStageSchedulerEvent; -use crate::state::execution_graph::ExecutionStage; use crate::state::execution_graph_dot::ExecutionGraphDot; -use crate::state::execution_stage::TaskInfo; use crate::{api::SchedulerErrorResponse, scheduler_server::SchedulerServer}; use axum::extract::Query; use axum::{ @@ -22,19 +20,14 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; +use ballista_core::BALLISTA_VERSION; use ballista_core::serde::protobuf::job_status::Status; -use ballista_core::serde::protobuf::{ - ExecutorMetric, executor_metric::Metric, task_status, -}; +use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric}; use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; -use ballista_core::utils::get_current_time; -use ballista_core::{BALLISTA_VERSION, JobId}; +use ballista_history::dto::JobResponse; use datafusion::DATAFUSION_VERSION; -use datafusion::physical_plan::display::DisplayableExecutionPlan; -use datafusion::physical_plan::displayable; -use datafusion::physical_plan::metrics::{MetricsSet, Time}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; #[cfg(feature = "graphviz-support")] @@ -44,9 +37,7 @@ use graphviz_rust::{ printer::PrinterContext, }; use http::{StatusCode, header::CONTENT_TYPE}; -use serde::Serialize; use std::sync::Arc; -use std::time::Duration; #[derive(Debug, serde::Serialize)] struct SchedulerStateResponse { @@ -107,27 +98,6 @@ impl ExecutorMetricResponse { } } -#[derive(Debug, serde::Serialize)] -pub struct JobResponse { - pub job_id: JobId, - pub job_name: String, - pub job_status: String, - pub status: String, - pub num_stages: usize, - pub completed_stages: usize, - pub percent_complete: u8, - /// Timestamp when the job started. - pub start_time: u64, - /// Timestamp when the job ended (0 if still running). - pub end_time: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub logical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub physical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, -} - #[derive(Debug, serde::Serialize)] struct CancelJobResponse { pub cancelled: bool, @@ -135,74 +105,6 @@ struct CancelJobResponse { pub reason: Option, } -#[derive(Debug, serde::Serialize)] -pub struct TaskSummary { - /// task id - pub id: usize, - /// Task status - pub status: TaskStatus, - /// partition id - pub partition_id: u32, - /// Scheduler schedule time - pub scheduled_time: u64, - /// Scheduler launch time (ms since epoch) - pub launch_time: u64, - /// The time the Executor start to run the task (ms since epoch) - pub start_exec_time: u64, - /// The time the Executor finish the task (ms since epoch) - pub end_exec_time: u64, - /// total execution time (ms) - pub exec_duration: u64, - /// Scheduler side finish time (ms since epoch) - pub finish_time: u64, - /// Number of input rows - pub input_rows: usize, - /// Number of output rows - pub output_rows: usize, -} - -#[derive(Debug, Clone, Serialize)] -pub enum TaskStatus { - Running, - Successful, - Failed, -} - -impl From<&task_status::Status> for TaskStatus { - fn from(value: &task_status::Status) -> Self { - match value { - task_status::Status::Running(_) => TaskStatus::Running, - task_status::Status::Failed(_) => TaskStatus::Failed, - task_status::Status::Successful(_) => TaskStatus::Successful, - } - } -} - -#[derive(Debug, serde::Serialize)] -pub struct Percentiles { - pub min: u64, - pub p25: u64, - pub median: u64, - pub p75: u64, - pub max: u64, -} - -#[derive(Debug, serde::Serialize)] -pub struct QueryStageSummary { - pub stage_id: String, - pub stage_status: String, - pub input_rows: usize, - pub output_rows: usize, - pub elapsed_compute: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_duration_percentiles: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_input_percentiles: Option, - pub tasks: Vec>, -} - #[derive(Debug, serde::Deserialize, Default)] pub struct JobQueryParams { /// Controls plan format @@ -332,34 +234,7 @@ pub async fn get_jobs< let jobs: Vec = jobs .iter() - .map(|job| { - let (plain_status, job_status) = format_job_status( - &job.status.status, - job_elapsed_ms(job.start_time, job.end_time), - ); - - // calculate progress based on completed stages for now, but we could use completed - // tasks in the future to make this more accurate - let percent_complete = if job.num_stages == 0 { - 0 - } else { - ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 - }; - JobResponse { - job_id: job.job_id.to_owned(), - job_name: job.job_name.to_owned(), - job_status, - status: plain_status, - start_time: job.start_time, - end_time: job.end_time, - num_stages: job.num_stages, - completed_stages: job.completed_stages, - percent_complete, - logical_plan: None, - physical_plan: None, - stage_plan: None, - } - }) + .map(dto_build::build_job_response_from_overview) .collect(); Ok(Json(jobs)) @@ -383,45 +258,14 @@ pub async fn get_job< SchedulerErrorResponse::with_error(StatusCode::INTERNAL_SERVER_ERROR, format!("Error occurred while getting the execution graph for job '{job_id}'")) })? .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; - let stage_plan = format!("{:?}", graph); - let job = graph.as_ref(); - let (plain_status, job_status) = format_job_status( - &job.status().status, - job_elapsed_ms(job.start_time(), job.end_time()), - ); - - let num_stages = job.stage_count(); - let completed_stages = job.completed_stages(); - let percent_complete = - ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; let plan_format = query.plan_format.clone().unwrap_or_default(); - let physical_plan = match plan_format { - PlanFormat::Default | PlanFormat::Metrics => { - DisplayableExecutionPlan::new(job.physical_plan().as_ref()) - .indent(false) - .to_string() - } - PlanFormat::Tree => displayable(job.physical_plan().as_ref()) - .tree_render() - .to_string(), - }; - - Ok(Json(JobResponse { - job_id: job.job_id().to_owned(), - job_name: job.job_name().to_owned(), - job_status, - status: plain_status, - start_time: job.start_time(), - end_time: job.end_time(), - num_stages, - completed_stages, - percent_complete, - logical_plan: job.logical_plan().map(str::to_owned), - physical_plan: Some(physical_plan), - stage_plan: Some(stage_plan), - })) + Ok(Json(dto_build::build_job_response( + &graph, + true, + plan_format, + ))) } pub async fn cancel_job< @@ -491,11 +335,6 @@ pub async fn cancel_job< } } -#[derive(Debug, serde::Serialize)] -pub struct QueryStagesResponse { - pub stages: Vec, -} - pub async fn get_query_stages< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -504,8 +343,6 @@ pub async fn get_query_stages< Path(job_id): Path, query: Query, ) -> Result { - let plan_format = query.plan_format.clone().unwrap_or_default(); - if let Some(graph) = data_server .state .task_manager @@ -519,310 +356,14 @@ pub async fn get_query_stages< ) })? { - let stages = graph - .as_ref() - .stages() - .iter() - .map(|(id, stage)| { - let mut summary = QueryStageSummary { - stage_id: id.to_string(), - stage_status: stage.variant_name().to_string(), - input_rows: 0, - output_rows: 0, - elapsed_compute: None, - tasks: vec![], - task_duration_percentiles: None, - task_input_percentiles: None, - stage_plan: None, - }; - match stage { - ExecutionStage::Running(running_stage) => { - let metrics = running_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(running_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(running_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(running_stage.plan.as_ref(), metrics), - }); - summary.input_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "input_rows")) - .unwrap_or(0); - summary.output_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "output_rows")) - .unwrap_or(0); - summary.elapsed_compute = get_running_stage_time(&running_stage - .task_infos, get_current_time()); - summary.tasks = running_stage - .task_infos - .iter() - .enumerate() - .map(|(partition_id, task_info)| { - task_info.as_ref().map(|info| { - let (input_rows, output_rows) = running_stage - .stage_metrics - .as_deref() - .map(|metrics| { - get_partition_counts(metrics, partition_id) - }) - .unwrap_or((0, 0)); - - let start_exec_time = info.start_exec_time as u64; - let end_exec_time = info.end_exec_time as u64; - - let task_status: TaskStatus = (&info.task_status).into(); - - TaskSummary { - id: info.task_id, - partition_id: partition_id as u32, - scheduled_time: info.scheduled_time as u64, - launch_time: info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: info.finish_time as u64, - input_rows, - output_rows, - status: task_status - } - }) - }) - .collect(); - } - ExecutionStage::Successful(completed_stage) => { - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(completed_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(completed_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(completed_stage.plan.as_ref(), &completed_stage.stage_metrics), - }); - summary.input_rows = get_combined_count( - &completed_stage.stage_metrics, - "input_rows", - ); - summary.output_rows = get_combined_count( - &completed_stage.stage_metrics, - "output_rows", - ); - summary.elapsed_compute = - get_finished_stage_time(&completed_stage.task_infos); - - summary.tasks = completed_stage - .task_infos - .iter() - .enumerate() - .map(|(partition_id, task_info)| { - let (input_rows, output_rows) = get_partition_counts( - &completed_stage.stage_metrics, - partition_id, - ); - - let start_exec_time = task_info.start_exec_time as u64; - let end_exec_time = task_info.end_exec_time as u64; - let task_status = (&task_info.task_status).into(); - Some(TaskSummary { - id: task_info.task_id, - partition_id: partition_id as u32, - scheduled_time: task_info.scheduled_time as u64, - launch_time: task_info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: task_info.finish_time as u64, - input_rows, - output_rows, - status: task_status - }) - }) - .collect(); - } - _ => {} - } - summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); - summary.task_input_percentiles = task_input_percentiles(&summary.tasks); - summary - }) - .collect(); - - Ok(Json(QueryStagesResponse { stages })) + Ok(Json(dto_build::build_query_stages_response( + &graph, &query, + ))) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } } -fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { - let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn task_input_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = tasks - .iter() - .flatten() - .map(|t| t.input_rows as u64) - .collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -fn task_duration_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = - tasks.iter().flatten().map(|t| t.exec_duration).collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -/// Returns elapsed wall time in milliseconds for API formatting. -/// -/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or -/// `end_time` still zero while `start_time` is set) do not panic on subtract. -fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { - end_time.saturating_sub(start_time) -} - -fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { - match status { - Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), - Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), - Some(Status::Failed(error)) => { - ("Failed".to_string(), format!("Failed: {}", error.error)) - } - Some(Status::Successful(completed)) => { - let num_rows = completed - .partition_location - .iter() - .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) - .sum::(); - let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; - let num_partitions = completed.partition_location.len(); - let num_partitions_term = if num_partitions == 1 { - "partition" - } else { - "partitions" - }; - ( - "Completed".to_string(), - format!( - "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", - num_partitions, - num_partitions_term, - num_rows, - num_rows_term, - elapsed_ms - ), - ) - } - _ => ("Invalid".to_string(), "Invalid State".to_string()), - } -} - -fn get_running_stage_time( - task_infos: &[Option], - current_time: u128, -) -> Option { - let min_start = task_infos - .iter() - .flat_map(|t| t.as_ref().map(|t| t.start_exec_time)) - .filter(|t| *t > 0) - .min(); - - match (min_start, current_time) { - (Some(start), end) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { - let min_start = task_infos - .iter() - .map(|t| t.start_exec_time) - .filter(|t| *t > 0) - .min(); - - let max_end = task_infos - .iter() - .map(|t| t.end_exec_time) - .filter(|t| *t > 0) - .max(); - - match (min_start, max_end) { - (Some(start), Some(end)) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) { - let input_rows = get_partition_count(metrics, partition_id, "input_rows"); - let output_rows = get_partition_count(metrics, partition_id, "output_rows"); - (input_rows, output_rows) -} - -fn get_partition_count(metrics: &[MetricsSet], partition_id: usize, name: &str) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - if metric.partition() == Some(partition_id) && metric_value.name() == name - { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - -fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - if metric_value.name() == name { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - pub async fn get_job_dot_graph< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -840,11 +381,10 @@ pub async fn get_job_dot_graph< SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) })? { - ExecutionGraphDot::generate(graph.as_ref()) - .map_err(|e| { - tracing::error!("Error occurred while getting the dot graph for job '{job_id}' reason: {e:?}"); - SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) - }) + dto_build::build_job_dot(&graph).map_err(|e| { + tracing::error!("Error occurred while getting the dot graph for job '{job_id}' reason: {e:?}"); + SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) + }) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } @@ -941,126 +481,3 @@ pub async fn get_job_config< .map(|e| Json(e.to_props())) .map_err(|_| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::state::execution_stage::TaskInfo; - use ballista_core::serde::protobuf::task_status; - - fn make_task_info(start: u128, end: u128) -> TaskInfo { - TaskInfo { - task_id: 0, - scheduled_time: 0, - launch_time: 0, - start_exec_time: start, - end_exec_time: end, - finish_time: 0, - task_status: task_status::Status::Running(Default::default()), - } - } - - #[test] - fn test_job_elapsed_saturates_when_end_precedes_start() { - assert_eq!(job_elapsed_ms(900, 100), 0); - } - - // --- get_finished_stage_time --- - - #[test] - fn test_finished_empty_slice_returns_none() { - assert_eq!(get_finished_stage_time(&[]), None); - } - - #[test] - fn test_finished_all_zero_timestamps_returns_none() { - let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - #[test] - fn test_finished_single_task_elapsed() { - // 600 - 100 = 500 ms → "500.00ms" - let tasks = vec![make_task_info(100, 600)]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("500.00ms".to_string()) - ); - } - - #[test] - fn test_finished_picks_earliest_start_and_latest_end() { - // min start = 100, max end = 900 → 800 ms - let tasks = vec![ - make_task_info(100, 500), - make_task_info(200, 900), - make_task_info(300, 700), - ]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("800.00ms".to_string()) - ); - } - - #[test] - fn test_finished_end_before_start_returns_none() { - let tasks = vec![make_task_info(900, 100)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - // --- get_running_stage_time --- - - #[test] - fn test_running_empty_slice_returns_none() { - assert_eq!(get_running_stage_time(&[], 1000), None); - } - - #[test] - fn test_running_all_none_returns_none() { - let tasks: Vec> = vec![None, None]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_future_start_returns_none() { - // start_exec_time beyond current time → elapsed clamped to 0 - let tasks = vec![Some(make_task_info(u128::MAX, 0))]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_past_start_returns_some() { - let now = 4_000; - let start = 1_000; - let tasks = vec![Some(make_task_info(start, 0))]; - assert_eq!( - get_running_stage_time(&tasks, now), - Some("3.00s".to_string()) - ); - } - - #[test] - fn test_running_mixed_some_none_uses_earliest_some() { - let now = 3_000; - let earlier = 1_000; - let later = 2_000; - let tasks = vec![ - None, - Some(make_task_info(later, 0)), - Some(make_task_info(earlier, 0)), - None, - ]; - let result = get_running_stage_time(&tasks, now); - assert_eq!(result, Some("2.00s".to_string())); - } - - #[test] - fn test_job_elapsed_ms_normal() { - assert_eq!(super::job_elapsed_ms(100, 500), 400); - } - - #[test] - fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { - assert_eq!(super::job_elapsed_ms(500, 100), 0); - } -} diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index c02b47fd4c..bdcfc2e527 100644 --- a/ballista/scheduler/src/api/mod.rs +++ b/ballista/scheduler/src/api/mod.rs @@ -11,7 +11,12 @@ // limitations under the License. #[cfg(feature = "rest-api")] -mod handlers; +pub(crate) mod dto_build; +#[cfg(feature = "rest-api")] +// `pub(crate)` (rather than private) so `scheduler_server::event_log` can reuse +// `JobQueryParams`/`PlanFormat` to build event-log DTOs identical to the live +// REST responses. +pub(crate) mod handlers; #[cfg(feature = "rest-api")] mod routes; #[cfg(feature = "rest-api")] diff --git a/ballista/scheduler/src/bin/history_server.rs b/ballista/scheduler/src/bin/history_server.rs new file mode 100644 index 0000000000..631c7f5cc0 --- /dev/null +++ b/ballista/scheduler/src/bin/history_server.rs @@ -0,0 +1,94 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Standalone Ballista history server binary: loads completed event logs from +//! a directory and serves the same `/api/*` responses the live scheduler does, +//! so the existing TUI can connect to it unchanged. + +use ballista_core::error::{BallistaError, Result}; +use ballista_scheduler::history::{HistoryStore, history_router}; +use clap::Parser; +use std::env; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use tracing_subscriber::EnvFilter; + +#[derive(Debug, clap::Parser)] +#[command( + name = "ballista-history-server", + version, + about = "Ballista history server" +)] +struct Args { + /// Directory containing per-job event logs. + #[arg(long)] + event_log_dir: PathBuf, + /// Host to bind the HTTP server to. + #[arg(long, default_value = "0.0.0.0")] + bind_host: String, + /// Port to bind the HTTP server to. + #[arg(long, default_value_t = 50060)] + bind_port: u16, +} + +fn main() -> Result<()> { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .enable_time() + .build() + .map_err(BallistaError::IoError)?; + + runtime.block_on(inner()) +} + +async fn inner() -> Result<()> { + let rust_log = env::var(EnvFilter::DEFAULT_ENV); + let log_filter = EnvFilter::new(rust_log.unwrap_or_else(|_| "info".to_string())); + tracing_subscriber::fmt() + .with_ansi(false) + .with_writer(std::io::stdout) + .with_env_filter(log_filter) + .init(); + + let args = Args::parse(); + + let store = Arc::new(HistoryStore::load(&args.event_log_dir)?); + tracing::info!( + "Loaded {} completed job(s) from {}", + store.jobs.len(), + args.event_log_dir.display() + ); + let app = history_router(store); + + let addr: SocketAddr = format!("{}:{}", args.bind_host, args.bind_port) + .parse() + .map_err(|e: std::net::AddrParseError| { + BallistaError::Configuration(e.to_string()) + })?; + + let listener = tokio::net::TcpListener::bind(&addr) + .await + .map_err(BallistaError::IoError)?; + tracing::info!("History server listening on http://{addr}"); + + axum::serve(listener, app.into_make_service()) + .await + .map_err(BallistaError::IoError)?; + + Ok(()) +} diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 3c6cdb0074..807cff1e69 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -103,6 +103,9 @@ pub struct Config { help = "Log dir: a path to save log. This will create a new storage directory at the specified path if it does not already exist." )] pub log_dir: Option, + /// Directory to write per-job event logs to (enables the history server). + #[arg(long)] + pub event_log_dir: Option, /// Whether to print thread IDs and names in log files. #[arg( long, @@ -278,6 +281,8 @@ pub struct SchedulerConfig { #[cfg(feature = "rest-api")] /// Comma-separated list of allowed methods for CORS pub cors_allowed_methods: String, + /// Directory to write per-job event logs to. `None` disables event logging. + pub event_log_dir: Option, } impl Default for SchedulerConfig { @@ -314,6 +319,7 @@ impl Default for SchedulerConfig { cors_allowed_origins: String::default(), #[cfg(feature = "rest-api")] cors_allowed_methods: String::default(), + event_log_dir: None, } } } @@ -449,6 +455,12 @@ impl SchedulerConfig { self.use_tls = use_tls; self } + + /// Sets the directory to write per-job event logs to. + pub fn with_event_log_dir(mut self, event_log_dir: Option) -> Self { + self.event_log_dir = event_log_dir; + self + } } /// Policy of distributing tasks to available executor slots @@ -546,8 +558,22 @@ impl TryFrom for SchedulerConfig { cors_allowed_origins: opt.cors_allowed_origins, #[cfg(feature = "rest-api")] cors_allowed_methods: opt.cors_allowed_methods, + event_log_dir: opt.event_log_dir, }; Ok(config) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn event_log_dir_defaults_to_none_and_is_settable() { + let config = SchedulerConfig::default(); + assert!(config.event_log_dir.is_none()); + let config = config.with_event_log_dir(Some("/tmp/history".to_string())); + assert_eq!(config.event_log_dir.as_deref(), Some("/tmp/history")); + } +} diff --git a/ballista/scheduler/src/history/mod.rs b/ballista/scheduler/src/history/mod.rs new file mode 100644 index 0000000000..8f6dbccd0e --- /dev/null +++ b/ballista/scheduler/src/history/mod.rs @@ -0,0 +1,370 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Standalone history server: loads completed event logs and serves the same +//! `/api/*` responses the live scheduler does, from stored DTOs. + +use crate::api::SchedulerErrorResponse; +use axum::{ + Json, Router, + extract::{Path as AxumPath, State}, + routing::get, +}; +use ballista_core::BALLISTA_VERSION; +use ballista_history::dto::{JobConfig, JobResponse, QueryStagesResponse}; +use ballista_history::reader::{ReplayedJob, read_completed_job}; +use datafusion::DATAFUSION_VERSION; +use http::StatusCode; +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; + +/// In-memory store of completed jobs, loaded once at startup from a directory +/// of `.eventlog` files. +#[derive(Default)] +pub struct HistoryStore { + /// Completed jobs keyed by job id. + pub jobs: HashMap, +} + +impl HistoryStore { + /// Load every completed job found under `dir`. Missing directories yield + /// an empty store rather than an error. + /// + /// A single unreadable/corrupt `.eventlog` file (e.g. truncated by a + /// crash mid-write) is logged and skipped rather than failing the whole + /// load — one bad log must not hide every other completed job. Only a + /// failure to read the directory itself is propagated. + pub fn load(dir: &Path) -> std::io::Result { + let mut jobs = HashMap::new(); + if dir.exists() { + for entry in std::fs::read_dir(dir)? { + let path = entry?.path(); + if path.extension().and_then(|e| e.to_str()) != Some("eventlog") { + continue; + } + match read_completed_job(&path) { + Ok(Some(replayed)) => { + jobs.insert(replayed.job.job_id.clone(), replayed); + } + Ok(None) => {} + Err(err) => { + tracing::warn!( + "skipping unreadable event log {}: {err}", + path.display() + ); + } + } + } + } + Ok(HistoryStore { jobs }) + } +} + +/// Build the axum router serving `/api/*` from a loaded [`HistoryStore`]. +pub fn history_router(store: Arc) -> Router { + Router::new() + .route("/api/jobs", get(get_jobs)) + .route("/api/job/{job_id}", get(get_job)) + .route("/api/job/{job_id}/stages", get(get_stages)) + .route("/api/job/{job_id}/config", get(get_config)) + .route("/api/job/{job_id}/dot", get(get_dot)) + .route("/api/executors", get(get_executors_empty)) + .route("/api/state", get(get_state)) + .with_state(store) +} + +async fn get_jobs(State(store): State>) -> Json> { + // The live list endpoint omits plans; null them for byte-identical output. + let mut jobs: Vec = store + .jobs + .values() + .map(|j| { + let mut r = j.job.clone(); + r.logical_plan = None; + r.physical_plan = None; + r.stage_plan = None; + r + }) + .collect(); + jobs.sort_by(|a, b| a.job_id.cmp(&b.job_id)); + Json(jobs) +} + +async fn get_job( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.job.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) +} + +async fn get_stages( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.stages.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) +} + +async fn get_config( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Result, SchedulerErrorResponse> { + store + .jobs + .get(&job_id) + .map(|j| Json(j.config.clone())) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) +} + +async fn get_dot( + State(store): State>, + AxumPath(job_id): AxumPath, +) -> Result { + store + .jobs + .get(&job_id) + .map(|j| j.dot.clone()) + .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) +} + +async fn get_executors_empty() -> Json> { + Json(vec![]) +} + +/// Static `/api/state` payload. The history server has no live scheduler +/// process behind it, so every field that would normally reflect runtime +/// state (uptime, feature flags, scheduling policy) is a fixed placeholder. +/// Field names/types match the live `/api/state` response +/// (`SchedulerStateResponse` in `api/handlers.rs`) and what the TUI +/// deserializes into (`ballista-cli/src/tui/domain/mod.rs::SchedulerState`), +/// so the TUI's startup call succeeds instead of erroring out. +async fn get_state() -> Json { + Json(serde_json::json!({ + "started": 0, + "version": BALLISTA_VERSION, + "datafusion_version": DATAFUSION_VERSION, + "substrait_support": false, + "keda_support": false, + "prometheus_support": false, + "graphviz_support": false, + "spark_support": false, + "scheduling_policy": "history-server", + })) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use ballista_history::dto::QueryStageSummary; + use ballista_history::event::{HistoryEvent, JobEndStatus, SCHEMA_VERSION}; + use std::io::Write; + use tempfile::tempdir; + use tower::ServiceExt; // oneshot + + const STAGE_ID_MARKER: &str = "stage-42"; + + fn sample_replayed_job(job_id: &str) -> ReplayedJob { + ReplayedJob { + job: JobResponse { + job_id: job_id.into(), + job_name: "q1".into(), + job_status: "COMPLETED".into(), + status: "Successful".into(), + num_stages: 1, + completed_stages: 1, + percent_complete: 100, + start_time: 2, + end_time: 3, + logical_plan: Some("Projection".into()), + physical_plan: Some("ProjectionExec".into()), + stage_plan: Some("stage".into()), + }, + stages: QueryStagesResponse { + stages: vec![QueryStageSummary { + stage_id: STAGE_ID_MARKER.into(), + stage_status: "Completed".into(), + input_rows: 10, + output_rows: 5, + elapsed_compute: Some("1ms".into()), + stage_plan: None, + task_duration_percentiles: None, + task_input_percentiles: None, + tasks: vec![], + }], + }, + config: Default::default(), + dot: "digraph {}".into(), + } + } + + fn store_with_one_job() -> Arc { + let mut jobs = HashMap::new(); + jobs.insert("job-1".to_string(), sample_replayed_job("job-1")); + Arc::new(HistoryStore { jobs }) + } + + #[tokio::test] + async fn jobs_endpoint_nulls_plan_fields() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/jobs") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + assert!(body.contains("\"job_id\":\"job-1\"")); + assert!(!body.contains("physical_plan")); // nulled + skip_serializing_if + } + + #[tokio::test] + async fn stages_endpoint_returns_stored_dto() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/job/job-1/stages") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body: QueryStagesResponse = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(body.stages.len(), 1); + assert_eq!(body.stages[0].stage_id, STAGE_ID_MARKER); + assert_eq!(body.stages[0].input_rows, 10); + assert_eq!(body.stages[0].output_rows, 5); + } + + #[tokio::test] + async fn missing_job_returns_404_on_job_and_stages() { + let app = history_router(store_with_one_job()); + + let resp = app + .clone() + .oneshot( + Request::builder() + .uri("/api/job/does-not-exist") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let resp = app + .oneshot( + Request::builder() + .uri("/api/job/does-not-exist/stages") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn state_endpoint_returns_static_payload() { + let app = history_router(store_with_one_job()); + let resp = app + .oneshot( + Request::builder() + .uri("/api/state") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let value: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + for field in [ + "started", + "version", + "datafusion_version", + "substrait_support", + "keda_support", + "prometheus_support", + "graphviz_support", + "spark_support", + "scheduling_policy", + ] { + assert!(value.get(field).is_some(), "missing field: {field}"); + } + } + + fn write_job_end_log(path: &Path, job_id: &str) { + let replayed = sample_replayed_job(job_id); + let event = HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status: JobEndStatus::Succeeded, + queued_at: 0, + started_at: 2, + completed_at: 3, + job: Box::new(replayed.job), + stages: Box::new(replayed.stages), + config: replayed.config, + dot: replayed.dot, + }; + let line = serde_json::to_string(&event).unwrap(); + std::fs::write(path, format!("{line}\n")).unwrap(); + } + + #[test] + fn load_skips_corrupt_eventlog_and_keeps_good_one() { + let dir = tempdir().unwrap(); + + // A good, readable event log. + write_job_end_log(&dir.path().join("job-good.eventlog"), "job-good"); + + // A corrupt file: invalid UTF-8, as if a crash truncated a write + // mid-multibyte-character. + let mut corrupt = + std::fs::File::create(dir.path().join("job-bad.eventlog")).unwrap(); + corrupt.write_all(&[0xff, 0xfe, 0xfd]).unwrap(); + drop(corrupt); + + let store = HistoryStore::load(dir.path()).unwrap(); + assert_eq!(store.jobs.len(), 1); + assert!(store.jobs.contains_key("job-good")); + assert!(!store.jobs.contains_key("job-bad")); + } +} diff --git a/ballista/scheduler/src/lib.rs b/ballista/scheduler/src/lib.rs index 533711ecb0..8338a65372 100644 --- a/ballista/scheduler/src/lib.rs +++ b/ballista/scheduler/src/lib.rs @@ -25,6 +25,9 @@ pub mod cluster; pub mod config; /// Display utilities for execution plans and state. pub mod display; +/// Standalone history server: serves `/api/*` from stored event logs. +#[cfg(feature = "rest-api")] +pub mod history; /// Metrics collection and reporting. pub mod metrics; /// Physical query plan optimizers. diff --git a/ballista/scheduler/src/scheduler_server/event_log.rs b/ballista/scheduler/src/scheduler_server/event_log.rs new file mode 100644 index 0000000000..dbb867be46 --- /dev/null +++ b/ballista/scheduler/src/scheduler_server/event_log.rs @@ -0,0 +1,335 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Builds `HistoryEvent`s from scheduler event-loop state and appends them to +//! the [`ballista_history::writer::EventLogWriter`]. These are the pure, +//! synchronous event-builder functions; the actual emission -- deciding +//! *when* to call them -- lives at the top of +//! `query_stage_scheduler::QueryStageScheduler::on_receive`. +//! +//! The builders reuse `crate::api::dto_build`, the same DTO builders backing +//! the live REST API, so a job's `JobEnd` event and its `GET /api/job/{id}` +//! response serialize identically for the same graph state. + +use crate::api::dto_build::{ + build_job_dot, build_job_response, build_query_stages_response, task_row_counts, + to_api_task_status, +}; +use crate::api::handlers::{JobQueryParams, PlanFormat}; +use crate::state::execution_graph::ExecutionGraphBox; +use ballista_core::serde::protobuf::{TaskStatus, task_status}; +use ballista_history::dto::TaskStatus as ApiTaskStatus; +use ballista_history::event::{ + HistoryEvent, JobEndStatus, SCHEMA_VERSION, TaskEndMetrics, +}; +use datafusion::physical_plan::display::DisplayableExecutionPlan; + +/// Builds the `JobStart` event for a job that has just been submitted. +pub(crate) fn job_start_event( + graph: &ExecutionGraphBox, + queued_at: u64, + submitted_at: u64, +) -> HistoryEvent { + HistoryEvent::JobStart { + version: SCHEMA_VERSION, + job_id: graph.job_id().to_string(), + job_name: graph.job_name().to_string(), + queued_at, + submitted_at, + logical_plan: graph.logical_plan().map(|p| p.to_string()), + // Same rendering the `get_job` handler uses for `PlanFormat::Default`. + physical_plan: Some( + DisplayableExecutionPlan::new(graph.physical_plan().as_ref()) + .indent(false) + .to_string(), + ), + } +} + +/// Builds one `TaskEnd` event per finished task in `statuses` (skips tasks +/// still `Running`, which are transient/in-flight status updates rather than +/// terminal states). +pub(crate) fn task_end_events( + executor_id: &str, + statuses: &[TaskStatus], +) -> Vec { + statuses + .iter() + .filter(|s| !matches!(s.status, Some(task_status::Status::Running(_)))) + .map(|s| { + let status = s + .status + .as_ref() + .map(to_api_task_status) + .unwrap_or(ApiTaskStatus::Running); + HistoryEvent::TaskEnd { + stage_id: s.stage_id, + partition: s.partition_id, + executor_id: executor_id.to_string(), + status, + launch_time: s.launch_time, + start_exec_time: s.start_exec_time, + end_exec_time: s.end_exec_time, + metrics: task_end_metrics(s), + } + }) + .collect() +} + +/// Sums a task's raw operator metrics into the timeline's `TaskEndMetrics`, +/// using the same extraction the stage-summary REST DTO performs so the +/// timeline and stage views agree. Absent metrics yield zeros. +fn task_end_metrics(status: &TaskStatus) -> TaskEndMetrics { + let (input_rows, output_rows, elapsed_compute_nanos) = + task_row_counts(&status.metrics); + TaskEndMetrics { + input_rows, + output_rows, + elapsed_compute_nanos, + } +} + +/// Builds the `JobEnd` event for a job that has just finished (successfully +/// or not), embedding the same DTOs the REST API would serve for this job at +/// this moment: the job summary, per-stage summaries, the session config, and +/// the DOT-format graph. +pub(crate) fn job_end_event( + graph: &ExecutionGraphBox, + status: JobEndStatus, + queued_at: u64, + completed_at: u64, +) -> HistoryEvent { + // `PlanFormat::Default` -- the same default the `get_job` handler falls + // back to when no `?plan_format=` query param is supplied -- so the + // stored JSON matches what a live `GET /api/job/{id}` would have + // returned at this point. + let job = Box::new(build_job_response(graph, true, PlanFormat::Default)); + let stages = Box::new(build_query_stages_response( + graph, + &JobQueryParams::default(), + )); + let dot = build_job_dot(graph).unwrap_or_default(); + let config = graph.session_config().to_props().into_iter().collect(); + + HistoryEvent::JobEnd { + version: SCHEMA_VERSION, + status, + queued_at, + started_at: graph.start_time(), + completed_at, + job, + stages, + config, + dot, + } +} + +#[cfg(test)] +mod tests { + use crate::scheduler_server::event_log::*; + use crate::state::execution_graph::ExecutionGraphBox; + use crate::state::execution_graph_dot::tests::test_graph; + use ballista_core::serde::protobuf::{SuccessfulTask, TaskStatus, task_status}; + use ballista_history::writer::EventLogWriter; + + #[test] + fn task_end_events_map_one_per_finished_task() { + let statuses = vec![TaskStatus { + task_id: 0, + job_id: "job-1".into(), + stage_id: 1, + stage_attempt_num: 0, + partition_id: 3, + launch_time: 100, + start_exec_time: 110, + end_exec_time: 200, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask::default())), + }]; + let events = task_end_events("exec-1", &statuses); + assert_eq!(events.len(), 1); + let line = serde_json::to_string(&events[0]).unwrap(); + assert!(line.contains("\"ev\":\"TaskEnd\"")); + assert!(line.contains("\"partition\":3")); + assert!(line.contains("\"executor_id\":\"exec-1\"")); + } + + // `task_end_events` skips `Running` statuses -- they are transient + // in-flight updates, not the terminal states the timeline records. + #[test] + fn task_end_events_skips_running_tasks() { + let statuses = vec![TaskStatus { + task_id: 0, + job_id: "job-1".into(), + stage_id: 1, + stage_attempt_num: 0, + partition_id: 0, + launch_time: 100, + start_exec_time: 110, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Running(Default::default())), + }]; + assert!(task_end_events("exec-1", &statuses).is_empty()); + } + + /// `job_end_event`'s embedded DTOs (`job`, `stages`, `dot`) are built by + /// the same `dto_build` functions backing the live REST API, so a direct + /// builder test on the serialized event is an adequate substitute for + /// wiring a full `QueryStageScheduler` through `on_receive` here (which + /// would additionally require standing up executor/task-manager state + /// just to reach a finished job). + #[tokio::test] + async fn job_end_event_embeds_job_and_stage_plan_strings() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + let event = job_end_event( + &graph, + JobEndStatus::Succeeded, + /* queued_at */ 1, + /* completed_at */ 2, + ); + let json = serde_json::to_string(&event).unwrap(); + + assert!(json.contains("\"ev\":\"JobEnd\"")); + assert!(json.contains("\"status\":\"Succeeded\"")); + assert!(json.contains("\"job_id\":\"job_id\"")); + // The job's top-level physical plan, embedded via `build_job_response`. + assert!( + json.contains("DataSourceExec: (Memory)"), + "expected job physical plan in JobEnd event, got: {json}" + ); + // The per-stage summaries, embedded via `build_query_stages_response`. + assert!( + json.contains("\"stage_status\":\"Resolved\""), + "expected a resolved stage summary in JobEnd event, got: {json}" + ); + // The DOT graph, embedded via `build_job_dot`. + assert!( + json.contains("digraph"), + "expected a dot graph in JobEnd event, got: {json}" + ); + } + + #[tokio::test] + async fn job_start_event_captures_plan_and_timestamps() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + + let event = + job_start_event(&graph, /* queued_at */ 5, /* submitted_at */ 10); + let json = serde_json::to_string(&event).unwrap(); + + assert!(json.contains("\"ev\":\"JobStart\"")); + assert!(json.contains("\"job_id\":\"job_id\"")); + assert!(json.contains("\"job_name\":\"job_name\"")); + assert!(json.contains("\"queued_at\":5")); + assert!(json.contains("\"submitted_at\":10")); + // `job_start_event` renders the graph's pre-staging physical plan + // (`graph.physical_plan()`), unlike `job_end_event`'s embedded job + // DTO which shows the plan reconstructed from stages -- assert on the + // join at its root rather than the leaf scan string. + assert!( + json.contains("HashJoinExec"), + "expected job physical plan in JobStart event, got: {json}" + ); + } + + /// End-to-end: drive the real `EventLogWriter` over a temp dir with a + /// `JobStart` followed by a `JobEnd`, then read the `.eventlog` file back + /// and assert on line order and content -- proving the builders here + /// produce events the writer can actually persist and that a consumer + /// (the eventual history server) can read back as ordered JSONL. + #[tokio::test] + async fn writer_round_trip_orders_job_start_before_job_end() { + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + let job_id = graph.job_id().to_string(); + + let dir = tempfile::tempdir().unwrap(); + let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); + + writer.append(&job_id, job_start_event(&graph, 1, 2)); + writer.append( + &job_id, + job_end_event(&graph, JobEndStatus::Succeeded, 1, 20), + ); + writer.flush_job(&job_id).await; + + let path = dir.path().join(format!("{job_id}.eventlog")); + let contents = tokio::fs::read_to_string(&path).await.unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + + assert_eq!(lines.len(), 2); + assert!(lines[0].contains("\"ev\":\"JobStart\"")); + assert!(lines[1].contains("\"ev\":\"JobEnd\"")); + assert!(lines[1].contains("DataSourceExec: (Memory)")); + } + + /// End-to-end DTO parity: the history server, replaying a real event log + /// through `EventLogWriter` + `HistoryStore::load`, must serve + /// byte-identical JSON to what the live scheduler would return for the + /// same job via `dto_build::{build_job_response, build_query_stages_response}`. + /// + /// This is `job_end_event`'s only round-trip coverage through the real + /// writer + `HistoryStore` (as opposed to the direct builder-output + /// assertions above), so it lives here as a crate-internal `#[cfg(test)]` + /// module rather than an integration test under `tests/`: `dto_build`'s + /// builders, `job_end_event` itself, and + /// `execution_graph_dot::tests::test_graph` are all `pub(crate)` / + /// cfg(test)-gated and therefore unreachable from a separate `tests/` + /// integration-test crate, which only sees the scheduler crate's public API. + #[tokio::test] + async fn history_store_serves_byte_identical_json_to_live_scheduler() { + use crate::api::handlers::JobQueryParams; + use crate::history::HistoryStore; + + let graph = test_graph().await.unwrap(); + let graph: ExecutionGraphBox = Box::new(graph); + let job_id = graph.job_id().to_string(); + + // The live DTOs a running scheduler would serve for this job right now, + // via the same builders the REST handlers call. + let live_job = build_job_response(&graph, true, PlanFormat::Default); + let live_stages = build_query_stages_response(&graph, &JobQueryParams::default()); + + // Produce the same `JobEnd` event the scheduler emits on completion, + // write it through the real async writer, then load it back via the + // history server's own `HistoryStore::load` -- exercising the full + // write -> read -> serve path, not just the event builder. + let event = job_end_event(&graph, JobEndStatus::Succeeded, 1, 2); + let dir = tempfile::tempdir().unwrap(); + let writer = EventLogWriter::new(dir.path().to_path_buf(), 16); + writer.append(&job_id, event); + writer.flush_job(&job_id).await; + + let store = HistoryStore::load(dir.path()).unwrap(); + let replayed = store.jobs.get(&job_id).expect("job should be replayed"); + + // The core fidelity guarantee: the history server would return + // byte-identical JSON to what the live scheduler returned for this job. + assert_eq!( + serde_json::to_string(&replayed.job).unwrap(), + serde_json::to_string(&live_job).unwrap(), + ); + assert_eq!( + serde_json::to_string(&replayed.stages).unwrap(), + serde_json::to_string(&live_stages).unwrap(), + ); + } +} diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 17b955a98c..77d9b38f6f 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -54,6 +54,13 @@ pub mod externalscaler { /// Events for the scheduler event loop. pub mod event; +/// Builds `HistoryEvent`s emitted from the event loop into the event log. +/// +/// Depends on `crate::api::dto_build`, which is only compiled with the +/// `rest-api` feature (the default), so this module -- and the event-log +/// wiring into `QueryStageScheduler` below -- is gated the same way. +#[cfg(feature = "rest-api")] +mod event_log; #[cfg(feature = "keda-scaler")] mod external_scaler; mod grpc; @@ -102,10 +109,19 @@ impl SchedulerServer SchedulerServer>, metrics_collector: Arc, config: Arc, + #[cfg(feature = "rest-api")] + event_log: Option, } impl QueryStageScheduler { @@ -51,11 +55,16 @@ impl QueryStageSchedul state: Arc>, metrics_collector: Arc, config: Arc, + #[cfg(feature = "rest-api")] event_log: Option< + ballista_history::writer::EventLogWriter, + >, ) -> Self { Self { state, metrics_collector, config, + #[cfg(feature = "rest-api")] + event_log, } } #[cfg(feature = "rest-api")] @@ -64,6 +73,25 @@ impl QueryStageSchedul } } +/// Groups task status updates by job id, so a single `TaskUpdating` batch +/// (which can span multiple jobs) can be appended to each job's own event log. +#[cfg(feature = "rest-api")] +fn group_by_job( + statuses: &[ballista_core::serde::protobuf::TaskStatus], +) -> std::collections::HashMap> { + let mut by_job: std::collections::HashMap< + String, + Vec, + > = std::collections::HashMap::new(); + for status in statuses { + by_job + .entry(status.job_id.clone()) + .or_default() + .push(status.clone()); + } + by_job +} + #[async_trait::async_trait] impl EventAction for QueryStageScheduler @@ -82,6 +110,88 @@ impl tx_event: &mpsc::Sender, _rx_event: &mpsc::Receiver, ) -> Result<()> { + #[cfg(feature = "rest-api")] + if let Some(log) = &self.event_log { + match &event { + QueryStageSchedulerEvent::JobSubmitted { + job_id, + queued_at, + submitted_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append( + job_id.as_str(), + event_log::job_start_event(&graph, *queued_at, *submitted_at), + ); + } + } + QueryStageSchedulerEvent::TaskUpdating(executor_id, statuses) => { + for (job_id, group) in group_by_job(statuses) { + for ev in event_log::task_end_events(executor_id, &group) { + log.append(&job_id, ev); + } + } + } + QueryStageSchedulerEvent::JobFinished { + job_id, + queued_at, + completed_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append_final( + job_id.as_str(), + event_log::job_end_event( + &graph, + ballista_history::event::JobEndStatus::Succeeded, + *queued_at, + *completed_at, + ), + ) + .await; + } + log.finish_job(job_id.as_str()).await; + } + QueryStageSchedulerEvent::JobRunningFailed { + job_id, + fail_message, + queued_at, + failed_at, + } => { + if let Ok(Some(graph)) = self + .state + .task_manager + .get_job_execution_graph(job_id) + .await + { + log.append_final( + job_id.as_str(), + event_log::job_end_event( + &graph, + ballista_history::event::JobEndStatus::Failed( + fail_message.clone(), + ), + *queued_at, + *failed_at, + ), + ) + .await; + } + log.finish_job(job_id.as_str()).await; + } + _ => {} + } + } + let mut time_recorder = None; if self.config.scheduler_event_expected_processing_duration > 0 { time_recorder = Some((Instant::now(), event.clone())); diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index eaef4dc5e9..c9a622aaba 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -406,7 +406,7 @@ fn get_file_scan(scan: &FileScanConfig) -> String { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use crate::planner::DefaultDistributedPlanner; use crate::state::execution_graph::StaticExecutionGraph; use crate::state::execution_graph_dot::ExecutionGraphDot; @@ -578,7 +578,7 @@ filter_expr="] Ok(()) } - async fn test_graph() -> Result { + pub(crate) async fn test_graph() -> Result { let mut config = SessionConfig::new() .with_target_partitions(48) .with_batch_size(4096);