diff --git a/Cargo.lock b/Cargo.lock index dac81a2616..697df0a10f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1024,6 +1024,13 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-api-types" +version = "54.0.0" +dependencies = [ + "serde", +] + [[package]] name = "ballista-benchmarks" version = "54.0.0" @@ -1207,6 +1214,7 @@ dependencies = [ "arrow-flight", "async-trait", "axum", + "ballista-api-types", "ballista-core", "clap 4.6.3", "dashmap", @@ -1230,6 +1238,7 @@ dependencies = [ "regex", "rstest", "serde", + "serde_json", "tokio", "tokio-stream", "tonic", diff --git a/Cargo.toml b/Cargo.toml index 5a6b966dd8..15817827b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ members = [ "ballista/client", "ballista/core", "ballista/executor", + "ballista/api-types", "ballista/scheduler", "benchmarks", "chaos-testing", diff --git a/ballista/api-types/Cargo.toml b/ballista/api-types/Cargo.toml new file mode 100644 index 0000000000..9c9485e332 --- /dev/null +++ b/ballista/api-types/Cargo.toml @@ -0,0 +1,30 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "ballista-api-types" +description = "Wire types for the Ballista scheduler REST API" +license = "Apache-2.0" +version = "54.0.0" +homepage = "https://datafusion.apache.org/ballista/" +repository = "https://github.com/apache/datafusion-ballista" +authors = ["Apache DataFusion "] +edition = { workspace = true } +rust-version = { workspace = true } + +[dependencies] +serde = { workspace = true, features = ["derive"] } diff --git a/ballista/api-types/src/dto.rs b/ballista/api-types/src/dto.rs new file mode 100644 index 0000000000..272730af84 --- /dev/null +++ b/ballista/api-types/src/dto.rs @@ -0,0 +1,166 @@ +// 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}; + +/// Summary of one job, served by `GET /api/jobs` and `GET /api/job/{job_id}`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobResponse { + /// A `String` rather than `ballista_core::JobId` so this crate stays a + /// serde-only leaf. `JobId` is `#[serde(transparent)]` over `String`, so + /// the serialized JSON is unchanged. + pub job_id: String, + /// Human-readable job name. + pub job_name: String, + /// Verbose status, including the completion or failure detail. + pub job_status: String, + /// Plain status word: `Queued`, `Running`, `Completed`, `Failed`, or `Invalid`. + pub status: String, + /// Total number of stages in the job. + pub num_stages: usize, + /// Number of stages that finished successfully. + pub completed_stages: usize, + /// Progress as a percentage of completed stages. + pub percent_complete: u8, + /// Timestamp when the job started. + pub start_time: u64, + /// Timestamp when the job ended (0 if still running). + pub end_time: u64, + /// Rendered logical plan. Absent in the job list. + #[serde(skip_serializing_if = "Option::is_none")] + pub logical_plan: Option, + /// Rendered physical plan. Absent in the job list. + #[serde(skip_serializing_if = "Option::is_none")] + pub physical_plan: Option, + /// Rendered stage DAG. Absent in the job list. + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, +} + +/// Terminal or in-flight state of a single task. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum TaskStatus { + /// Task is currently executing. + Running, + /// Task completed successfully. + Successful, + /// Task failed, with the classified reason and the underlying error. + Failed { + /// Failure category, e.g. `ExecutionError` or `FetchPartitionError`. + reason: String, + /// Underlying error message. + error: String, + }, +} + +/// Per-task detail within a stage. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskSummary { + /// task id + pub id: usize, + /// Task status + pub status: TaskStatus, + /// Global partition ids covered by this task. For a single-partition + /// task this is a one-element list — JSON-compatible with the old + /// scalar `partition_id` for callers that read `partition_id[0]`, and + /// honestly plural for multi-partition tasks. + pub partition_id: Vec, + /// Scheduler schedule time + pub scheduled_time: u64, + /// Scheduler launch time (ms since epoch) + pub launch_time: u64, + /// The time the Executor start to run the task (ms since epoch) + pub start_exec_time: u64, + /// The time the Executor finish the task (ms since epoch) + pub end_exec_time: u64, + /// total execution time (ms) + pub exec_duration: u64, + /// Scheduler side finish time (ms since epoch) + pub finish_time: u64, + /// Number of input rows + pub input_rows: usize, + /// Number of output rows + pub output_rows: usize, +} + +/// Five-number summary over a stage's tasks, used to spot skew. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Percentiles { + /// Smallest observed value. + pub min: u64, + /// 25th percentile. + pub p25: u64, + /// 50th percentile. + pub median: u64, + /// 75th percentile. + pub p75: u64, + /// Largest observed value. + pub max: u64, +} + +/// Summary of one query stage, served by `GET /api/job/{job_id}/stages`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStageSummary { + /// Stage id, as a string. + pub stage_id: String, + /// Stage state, e.g. `Running`, `Successful`, `Failed`. + pub stage_status: String, + /// Rows read by the stage, summed across tasks. + pub input_rows: usize, + /// Rows produced by the stage, summed across tasks. + pub output_rows: usize, + /// Formatted wall time across the stage's tasks, if any have started. + pub elapsed_compute: Option, + /// Rendered plan for this stage, in the requested [`PlanFormat`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub stage_plan: Option, + /// Distribution of per-task execution time. + #[serde(skip_serializing_if = "Option::is_none")] + pub task_duration_percentiles: Option, + /// Distribution of per-task input row counts. + #[serde(skip_serializing_if = "Option::is_none")] + pub task_input_percentiles: Option, + /// One entry per task. Always `Some` today; the `Option` predates + /// multi-partition tasks, when this list was indexed by partition and + /// could be sparse. + pub tasks: Vec>, +} + +/// Response body for `GET /api/job/{job_id}/stages`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryStagesResponse { + /// One summary per stage, in stage order. + pub stages: Vec, +} + +/// How a plan should be rendered. Parsed from the `?plan_format=` query +/// parameter, and part of the REST contract, so it lives alongside the +/// response types rather than with the HTTP handlers. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PlanFormat { + /// `?plan_format=default` => plain indent, no metrics + #[default] + Default, + /// `?plan_format=tree` => tree render, no metrics + Tree, + /// `?plan_format=metrics` => indent with aggregated metrics + Metrics, +} diff --git a/ballista/api-types/src/lib.rs b/ballista/api-types/src/lib.rs new file mode 100644 index 0000000000..08fab21c28 --- /dev/null +++ b/ballista/api-types/src/lib.rs @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![warn(missing_docs)] + +//! Wire types for Ballista's scheduler REST API. +//! +//! This crate is the single definition of what `/api/*` puts on the wire. It +//! exists because that contract has more than one party: +//! +//! - **`ballista-scheduler`** serves these responses for live jobs. +//! - **The web TUI** deserializes them to render the cluster. +//! - **A history server** ([#1923]) will serve them for completed jobs, replayed +//! from a durable event log. +//! +//! Each of those independently re-declaring the same structs is how the shapes +//! drift apart, so they share one definition instead. +//! +//! The crate is deliberately serde-only. That keeps it cheap enough for anyone +//! to depend on, and it is what lets the TUI use it from a `wasm32` build. +//! +//! # What belongs here +//! +//! A type belongs here when it is part of the `/api/*` contract and more than +//! one party needs it. Types describing live scheduler internals +//! (`SchedulerStateResponse`, `CancelJobResponse`) stay in `ballista-scheduler`. +//! +//! `ExecutorResponse` stays behind for a different reason: it embeds +//! `ballista-core` types, so sharing it would mean either taking a +//! `ballista-core` dependency here or duplicating those structs. Neither is +//! worth it until something outside the scheduler needs it. +//! +//! Note that the types are shared but the *construction* is not: building a +//! response from a live execution graph needs scheduler internals, so those +//! builders live in `ballista-scheduler`. A history server does not re-derive +//! responses — it replays ones the scheduler already built and stored, so +//! byte-identical output is a structural property rather than two +//! implementations agreeing. +//! +//! [#1923]: https://github.com/apache/datafusion-ballista/issues/1923 + +pub mod dto; diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index b8b7127733..cd26d8d420 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -41,7 +41,7 @@ disable-stage-plan-cache = [] graphviz-support = ["dep:graphviz-rust"] keda-scaler = ["dep:tonic-prost-build", "dep:tonic-prost"] prometheus-metrics = ["prometheus", "once_cell"] -rest-api = [] +rest-api = ["dep:ballista-api-types"] spark-compat = ["ballista-core/spark-compat"] substrait = ["dep:datafusion-substrait"] @@ -49,6 +49,7 @@ substrait = ["dep:datafusion-substrait"] arrow-flight = { workspace = true } async-trait = { workspace = true } axum = "0.8.9" +ballista-api-types = { path = "../api-types", version = "54.0.0", optional = true } ballista-core = { path = "../core", version = "54.0.0" } clap = { workspace = true, optional = true } dashmap = { workspace = true } @@ -88,6 +89,7 @@ path = "tests/tpch_plan_stability/main.rs" datafusion-functions-aggregate-common = { workspace = true } regex = "1" rstest = { workspace = true } +serde_json = "1" [build-dependencies] tonic-prost-build = { workspace = true, optional = true } diff --git a/ballista/scheduler/src/api/dto_build.rs b/ballista/scheduler/src/api/dto_build.rs new file mode 100644 index 0000000000..d1f343666b --- /dev/null +++ b/ballista/scheduler/src/api/dto_build.rs @@ -0,0 +1,550 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Builds REST response DTOs from scheduler execution state. +//! +//! These functions are the only place the scheduler's `ExecutionGraph` / +//! `JobOverview` shapes are translated into the wire types in +//! [`ballista_api_types::dto`]. Keeping the translation here (rather than inline +//! in the axum handlers) means the same DTOs can be produced from state that +//! did not come from a live handler request. +//! +//! Everything in this module is pure: state in, DTO out, no I/O and no +//! wall-clock reads. `graph_to_query_stages` takes `now` from its caller so a +//! replay of a stored log renders the same elapsed times every time. + +use crate::display::format_stage_metrics; +use crate::state::execution_graph::{ExecutionGraphBox, ExecutionStage}; +use crate::state::execution_stage::TaskInfo; +use crate::state::task_manager::JobOverview; +use ballista_api_types::dto::{ + JobResponse, Percentiles, PlanFormat, QueryStageSummary, QueryStagesResponse, + TaskStatus, TaskSummary, +}; +use ballista_core::serde::protobuf::failed_task::FailedReason::{ + ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, +}; +use ballista_core::serde::protobuf::job_status::Status; +use ballista_core::serde::protobuf::{FailedTask, task_status}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::displayable; +use datafusion::physical_plan::metrics::{MetricsSet, Time}; +use std::time::Duration; + +/// Build the `JobResponse` list entry served by `GET /api/jobs`. +/// +/// Plan fields are always `None` here: the job list is built from +/// [`JobOverview`] summaries, which do not carry plans. +pub fn job_overview_to_response(job: &JobOverview) -> JobResponse { + let (plain_status, job_status) = format_job_status( + &job.status.status, + job_elapsed_ms(job.start_time, job.end_time), + ); + + 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: percent_complete(job.completed_stages, job.num_stages), + logical_plan: None, + physical_plan: None, + stage_plan: None, + } +} + +/// Build the `JobResponse` served by `GET /api/job/{job_id}`, including the +/// rendered logical, physical, and stage plans. +pub fn graph_to_job_response( + graph: &ExecutionGraphBox, + plan_format: PlanFormat, +) -> JobResponse { + let stage_plan = format!("{graph:?}"); + let job = graph.as_ref(); + + let (plain_status, job_status) = format_job_status( + &job.status().status, + job_elapsed_ms(job.start_time(), job.end_time()), + ); + + let num_stages = job.stage_count(); + let completed_stages = job.completed_stages(); + + let physical_plan = match plan_format { + PlanFormat::Default | PlanFormat::Metrics => { + displayable(job.physical_plan().as_ref()) + .indent(false) + .to_string() + } + PlanFormat::Tree => displayable(job.physical_plan().as_ref()) + .tree_render() + .to_string(), + }; + + JobResponse { + job_id: job.job_id().to_string(), + job_name: job.job_name().to_owned(), + job_status, + status: plain_status, + start_time: job.start_time(), + end_time: job.end_time(), + num_stages, + completed_stages, + percent_complete: percent_complete(completed_stages, num_stages), + logical_plan: job.logical_plan().map(str::to_owned), + physical_plan: Some(physical_plan), + stage_plan: Some(stage_plan), + } +} + +/// Build the per-stage summaries served by `GET /api/job/{job_id}/stages`. +pub fn graph_to_query_stages( + graph: &ExecutionGraphBox, + plan_format: PlanFormat, + now: u128, +) -> QueryStagesResponse { + let stages = graph + .as_ref() + .stages() + .iter() + .map(|(id, stage)| { + // Every started stage contributes the same three things; only how + // it reaches its metrics and which elapsed-time rule applies + // differ. Stages that have not started yet contribute nothing. + let started: Option<(&[MetricsSet], &[TaskInfo], Option)> = + match stage { + ExecutionStage::Running(s) => Some(( + s.stage_metrics.as_deref().unwrap_or(&[]), + &s.task_infos, + get_running_stage_time(&s.task_infos, now), + )), + ExecutionStage::Successful(s) => Some(( + &s.stage_metrics, + &s.task_infos, + get_finished_stage_time(&s.task_infos), + )), + ExecutionStage::Failed(s) => Some(( + s.stage_metrics.as_deref().unwrap_or(&[]), + &s.task_infos, + get_finished_stage_time(&s.task_infos), + )), + _ => None, + }; + + let has_started = started.is_some(); + let (metrics, task_infos, elapsed_compute) = started.unwrap_or_default(); + let tasks = task_summaries(task_infos, metrics); + + QueryStageSummary { + stage_id: id.to_string(), + stage_status: stage.variant_name().to_string(), + input_rows: get_combined_count(metrics, "input_rows"), + output_rows: get_combined_count(metrics, "output_rows"), + elapsed_compute, + stage_plan: has_started + .then(|| render_stage_plan(stage.plan(), metrics, plan_format)), + task_duration_percentiles: task_duration_percentiles(&tasks), + task_input_percentiles: task_input_percentiles(&tasks), + tasks, + } + }) + .collect(); + + QueryStagesResponse { stages } +} + +/// Render one stage's plan in the requested format. `Metrics` overlays the +/// stage's aggregated metrics onto the plan; the other formats ignore them. +fn render_stage_plan( + plan: &dyn ExecutionPlan, + metrics: &[MetricsSet], + plan_format: PlanFormat, +) -> String { + match plan_format { + PlanFormat::Default => displayable(plan).indent(false).to_string(), + PlanFormat::Tree => displayable(plan).tree_render().to_string(), + PlanFormat::Metrics => format_stage_metrics(plan, metrics), + } +} + +/// Build one [`TaskSummary`] per task in a stage. Row counts are summed over +/// the global partitions each task owns. +/// +/// The `Option` wrapper on each entry is part of the wire format and is always +/// `Some` here; it predates multi-partition tasks, when a stage's task list was +/// indexed by partition and could be sparse. +fn task_summaries( + task_infos: &[TaskInfo], + metrics: &[MetricsSet], +) -> Vec> { + task_infos + .iter() + .map(|info| { + let (input_rows, output_rows) = + get_partition_counts(metrics, &info.global_input_partition_ids); + + let start_exec_time = info.start_exec_time as u64; + let end_exec_time = info.end_exec_time as u64; + + Some(TaskSummary { + id: info.task_id, + partition_id: info + .global_input_partition_ids + .iter() + .map(|&p| p as u32) + .collect(), + scheduled_time: info.scheduled_time as u64, + launch_time: info.launch_time as u64, + start_exec_time, + end_exec_time, + exec_duration: end_exec_time.saturating_sub(start_exec_time), + finish_time: info.finish_time as u64, + input_rows, + output_rows, + status: task_status_to_dto(&info.task_status), + }) + }) + .collect() +} + +/// Map a protobuf task status onto the wire enum. +/// +/// A free function rather than a `From` impl: both types are foreign to this +/// crate now that [`TaskStatus`] lives in `ballista-api-types`. +fn task_status_to_dto(value: &task_status::Status) -> TaskStatus { + match value { + task_status::Status::Running(_) => TaskStatus::Running, + task_status::Status::Failed(failed) => TaskStatus::Failed { + reason: failed_reason(failed), + error: failed.error.clone(), + }, + task_status::Status::Successful(_) => TaskStatus::Successful, + } +} + +/// Progress as a percentage of stages completed. Zero-stage jobs report 0 +/// rather than dividing by zero. +fn percent_complete(completed_stages: usize, num_stages: usize) -> u8 { + if num_stages == 0 { + return 0; + } + ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8 +} + +fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { + let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; + sorted[idx.min(sorted.len() - 1)] +} + +fn percentiles_of(mut values: Vec) -> Option { + if values.is_empty() { + return None; + } + + values.sort_unstable(); + + Some(Percentiles { + min: values[0], + p25: percentile_duration(&values, 25.0), + median: percentile_duration(&values, 50.0), + p75: percentile_duration(&values, 75.0), + max: *values.last().unwrap(), + }) +} + +fn task_input_percentiles(tasks: &[Option]) -> Option { + percentiles_of( + tasks + .iter() + .flatten() + .map(|t| t.input_rows as u64) + .collect(), + ) +} + +fn task_duration_percentiles(tasks: &[Option]) -> Option { + percentiles_of(tasks.iter().flatten().map(|t| t.exec_duration).collect()) +} + +/// Returns elapsed wall time in milliseconds for API formatting. +/// +/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or +/// `end_time` still zero while `start_time` is set) do not panic on subtract. +fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { + end_time.saturating_sub(start_time) +} + +fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { + match status { + Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), + Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), + Some(Status::Failed(error)) => { + ("Failed".to_string(), format!("Failed: {}", error.error)) + } + Some(Status::Successful(completed)) => { + let num_rows = completed + .partition_location + .iter() + .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) + .sum::(); + let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; + let num_partitions = completed.partition_location.len(); + let num_partitions_term = if num_partitions == 1 { + "partition" + } else { + "partitions" + }; + ( + "Completed".to_string(), + format!( + "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", + num_partitions, + num_partitions_term, + num_rows, + num_rows_term, + elapsed_ms + ), + ) + } + _ => ("Invalid".to_string(), "Invalid State".to_string()), + } +} + +/// Earliest non-zero task start in a stage. Zero means "not started yet", so +/// those entries are ignored rather than dragging the minimum to 0. +fn min_start_time(task_infos: &[TaskInfo]) -> Option { + task_infos + .iter() + .map(|t| t.start_exec_time) + .filter(|t| *t > 0) + .min() +} + +fn get_running_stage_time(task_infos: &[TaskInfo], current_time: u128) -> Option { + match (min_start_time(task_infos), current_time) { + (Some(start), end) if end >= start => Some(format_millis(end - start)), + _ => None, + } +} + +fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { + let max_end = task_infos + .iter() + .map(|t| t.end_exec_time) + .filter(|t| *t > 0) + .max(); + + match (min_start_time(task_infos), max_end) { + (Some(start), Some(end)) if end >= start => Some(format_millis(end - start)), + _ => None, + } +} + +/// Format a millisecond duration the way DataFusion renders elapsed-time +/// metrics, so stage timings match plan output. +fn format_millis(millis: u128) -> String { + let time = Time::new(); + time.add_duration(Duration::from_millis(millis as u64)); + time.to_string() +} + +fn failed_reason(failed: &FailedTask) -> String { + match &failed.failed_reason { + Some(ExecutionError(_)) => "ExecutionError", + Some(FetchPartitionError(_)) => "FetchPartitionError", + Some(IoError(_)) => "IoError", + Some(ExecutorLost(_)) => "ExecutorLost", + Some(ResultLost(_)) => "ResultLost", + Some(TaskKilled(_)) => "TaskKilled", + None => "Failed", + } + .to_string() +} + +/// Sum a task's `input_rows` / `output_rows` across the global partitions the +/// task owns. Metrics are keyed by global partition id — for single-partition +/// tasks `partitions` is a one-element slice; for multi-partition tasks it is +/// the task's `global_input_partition_ids`. +fn get_partition_counts(metrics: &[MetricsSet], partitions: &[usize]) -> (usize, usize) { + let input_rows = get_partition_count(metrics, partitions, "input_rows"); + let output_rows = get_partition_count(metrics, partitions, "output_rows"); + (input_rows, output_rows) +} + +fn get_partition_count( + metrics: &[MetricsSet], + partitions: &[usize], + name: &str, +) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + let owned_by_task = metric + .partition() + .map(|p| partitions.contains(&p)) + .unwrap_or(false); + if owned_by_task && metric_value.name() == name { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { + metrics + .iter() + .flat_map(|vec| { + vec.iter().map(|metric| { + let metric_value = metric.value(); + if metric_value.name() == name { + metric_value.as_usize() + } else { + 0 + } + }) + }) + .sum() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::state::execution_stage::TaskInfo; + use ballista_core::serde::protobuf::task_status; + + fn make_task_info(start: u128, end: u128) -> TaskInfo { + TaskInfo { + task_id: 0, + scheduled_time: 0, + launch_time: 0, + start_exec_time: start, + end_exec_time: end, + finish_time: 0, + task_status: task_status::Status::Running(Default::default()), + global_input_partition_ids: vec![], + vcores_consumed: 0, + } + } + + // --- get_finished_stage_time --- + + #[test] + fn test_finished_empty_slice_returns_none() { + assert_eq!(get_finished_stage_time(&[]), None); + } + + #[test] + fn test_finished_all_zero_timestamps_returns_none() { + let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + #[test] + fn test_finished_single_task_elapsed() { + // 600 - 100 = 500 ms → "500.00ms" + let tasks = vec![make_task_info(100, 600)]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("500.00ms".to_string()) + ); + } + + #[test] + fn test_finished_picks_earliest_start_and_latest_end() { + // min start = 100, max end = 900 → 800 ms + let tasks = vec![ + make_task_info(100, 500), + make_task_info(200, 900), + make_task_info(300, 700), + ]; + assert_eq!( + get_finished_stage_time(&tasks), + Some("800.00ms".to_string()) + ); + } + + #[test] + fn test_finished_end_before_start_returns_none() { + let tasks = vec![make_task_info(900, 100)]; + assert_eq!(get_finished_stage_time(&tasks), None); + } + + // --- get_running_stage_time --- + + #[test] + fn test_running_empty_slice_returns_none() { + assert_eq!(get_running_stage_time(&[], 1000), None); + } + + #[test] + fn test_running_all_zero_start_returns_none() { + let tasks: Vec = vec![make_task_info(0, 0), make_task_info(0, 0)]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_future_start_returns_none() { + // start_exec_time beyond current time → elapsed clamped to 0 + let tasks = vec![make_task_info(u128::MAX, 0)]; + assert_eq!(get_running_stage_time(&tasks, 1000), None); + } + + #[test] + fn test_running_past_start_returns_some() { + let now = 4_000; + let start = 1_000; + let tasks = vec![make_task_info(start, 0)]; + assert_eq!( + get_running_stage_time(&tasks, now), + Some("3.00s".to_string()) + ); + } + + #[test] + fn test_running_mixed_zero_start_uses_earliest_nonzero() { + let now = 3_000; + let earlier = 1_000; + let later = 2_000; + let tasks = vec![ + make_task_info(0, 0), + make_task_info(later, 0), + make_task_info(earlier, 0), + make_task_info(0, 0), + ]; + let result = get_running_stage_time(&tasks, now); + assert_eq!(result, Some("2.00s".to_string())); + } + + #[test] + fn test_job_elapsed_ms_normal() { + assert_eq!(super::job_elapsed_ms(100, 500), 400); + } + + #[test] + fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { + assert_eq!(super::job_elapsed_ms(500, 100), 0); + } +} diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 5409f07801..bd6d77a501 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -10,11 +10,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::display::format_stage_metrics; +use crate::api::dto_build; use crate::scheduler_server::event::QueryStageSchedulerEvent; -use crate::state::execution_graph::ExecutionStage; use crate::state::execution_graph_dot::ExecutionGraphDot; -use crate::state::execution_stage::TaskInfo; use crate::{api::SchedulerErrorResponse, scheduler_server::SchedulerServer}; use axum::extract::Query; use axum::response::Redirect; @@ -23,22 +21,15 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; -use ballista_core::serde::protobuf::failed_task::FailedReason::{ - ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, -}; +use ballista_api_types::dto::{JobResponse, PlanFormat}; +use ballista_core::BALLISTA_VERSION; use ballista_core::serde::protobuf::job_status::Status; -use ballista_core::serde::protobuf::{ - ExecutorMetric, FailedTask, executor_metric::Metric, task_status, -}; +use ballista_core::serde::protobuf::{ExecutorMetric, executor_metric::Metric}; use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; use ballista_core::utils::get_current_time; -use ballista_core::{BALLISTA_VERSION, JobId}; use datafusion::DATAFUSION_VERSION; -use datafusion::physical_plan::display::DisplayableExecutionPlan; -use datafusion::physical_plan::displayable; -use datafusion::physical_plan::metrics::{MetricsSet, Time}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; #[cfg(feature = "graphviz-support")] @@ -48,10 +39,8 @@ use graphviz_rust::{ printer::PrinterContext, }; use http::{HeaderMap, StatusCode, header::CONTENT_TYPE}; -use serde::Serialize; use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; #[derive(Debug, serde::Serialize)] struct SchedulerStateResponse { @@ -112,27 +101,6 @@ impl ExecutorMetricResponse { } } -#[derive(Debug, serde::Serialize)] -pub struct JobResponse { - pub job_id: JobId, - pub job_name: String, - pub job_status: String, - pub status: String, - pub num_stages: usize, - pub completed_stages: usize, - pub percent_complete: u8, - /// Timestamp when the job started. - pub start_time: u64, - /// Timestamp when the job ended (0 if still running). - pub end_time: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub logical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub physical_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, -} - #[derive(Debug, serde::Serialize)] struct CancelJobResponse { pub cancelled: bool, @@ -140,98 +108,12 @@ struct CancelJobResponse { pub reason: Option, } -#[derive(Debug, serde::Serialize)] -pub struct TaskSummary { - /// task id - pub id: usize, - /// Task status - pub status: TaskStatus, - /// Global partition ids covered by this task. For a single-partition - /// task this is a one-element list — JSON-compatible with the old - /// scalar `partition_id` for callers that read `partition_id[0]`, and - /// honestly plural for multi-partition tasks. - pub partition_id: Vec, - /// Scheduler schedule time - pub scheduled_time: u64, - /// Scheduler launch time (ms since epoch) - pub launch_time: u64, - /// The time the Executor start to run the task (ms since epoch) - pub start_exec_time: u64, - /// The time the Executor finish the task (ms since epoch) - pub end_exec_time: u64, - /// total execution time (ms) - pub exec_duration: u64, - /// Scheduler side finish time (ms since epoch) - pub finish_time: u64, - /// Number of input rows - pub input_rows: usize, - /// Number of output rows - pub output_rows: usize, -} - -#[derive(Debug, Clone, Serialize)] -pub enum TaskStatus { - Running, - Successful, - Failed { reason: String, error: String }, -} - -impl From<&task_status::Status> for TaskStatus { - fn from(value: &task_status::Status) -> Self { - match value { - task_status::Status::Running(_) => TaskStatus::Running, - task_status::Status::Failed(failed) => TaskStatus::Failed { - reason: failed_reason(failed), - error: failed.error.clone(), - }, - task_status::Status::Successful(_) => TaskStatus::Successful, - } - } -} - -#[derive(Debug, serde::Serialize)] -pub struct Percentiles { - pub min: u64, - pub p25: u64, - pub median: u64, - pub p75: u64, - pub max: u64, -} - -#[derive(Debug, serde::Serialize)] -pub struct QueryStageSummary { - pub stage_id: String, - pub stage_status: String, - pub input_rows: usize, - pub output_rows: usize, - pub elapsed_compute: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stage_plan: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_duration_percentiles: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub task_input_percentiles: Option, - pub tasks: Vec>, -} - #[derive(Debug, serde::Deserialize, Default)] pub struct JobQueryParams { /// Controls plan format pub plan_format: Option, } -#[derive(Debug, serde::Deserialize, Default, Clone)] -#[serde(rename_all = "snake_case")] -pub enum PlanFormat { - /// ?plan_format=default => plain indent, no metrics - #[default] - Default, - /// ?plan_format=tree => tree render, no metrics - Tree, - /// ?plan_format=metrics => indent with aggregated metrics - Metrics, -} - /// A handler for GET requests to the root (`/`). /// It redirects to `https://nightlies.apache.org/datafusion/ballista/tui//` /// forwarding any query parameters @@ -395,34 +277,7 @@ pub async fn get_jobs< let jobs: Vec = jobs .iter() - .map(|job| { - let (plain_status, job_status) = format_job_status( - &job.status.status, - job_elapsed_ms(job.start_time, job.end_time), - ); - - // calculate progress based on completed stages for now, but we could use completed - // tasks in the future to make this more accurate - let percent_complete = if job.num_stages == 0 { - 0 - } else { - ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 - }; - JobResponse { - job_id: job.job_id.to_owned(), - job_name: job.job_name.to_owned(), - job_status, - status: plain_status, - start_time: job.start_time, - end_time: job.end_time, - num_stages: job.num_stages, - completed_stages: job.completed_stages, - percent_complete, - logical_plan: None, - physical_plan: None, - stage_plan: None, - } - }) + .map(dto_build::job_overview_to_response) .collect(); Ok(Json(jobs)) @@ -446,45 +301,11 @@ pub async fn get_job< SchedulerErrorResponse::with_error(StatusCode::INTERNAL_SERVER_ERROR, format!("Error occurred while getting the execution graph for job '{job_id}'")) })? .ok_or_else(|| SchedulerErrorResponse::new(StatusCode::NOT_FOUND))?; - let stage_plan = format!("{:?}", graph); - let job = graph.as_ref(); - let (plain_status, job_status) = format_job_status( - &job.status().status, - job_elapsed_ms(job.start_time(), job.end_time()), - ); - - let num_stages = job.stage_count(); - let completed_stages = job.completed_stages(); - let percent_complete = - ((completed_stages as f32 / num_stages as f32) * 100_f32) as u8; - - let plan_format = query.plan_format.clone().unwrap_or_default(); - - let physical_plan = match plan_format { - PlanFormat::Default | PlanFormat::Metrics => { - DisplayableExecutionPlan::new(job.physical_plan().as_ref()) - .indent(false) - .to_string() - } - PlanFormat::Tree => displayable(job.physical_plan().as_ref()) - .tree_render() - .to_string(), - }; - Ok(Json(JobResponse { - job_id: job.job_id().to_owned(), - job_name: job.job_name().to_owned(), - job_status, - status: plain_status, - start_time: job.start_time(), - end_time: job.end_time(), - num_stages, - completed_stages, - percent_complete, - logical_plan: job.logical_plan().map(str::to_owned), - physical_plan: Some(physical_plan), - stage_plan: Some(stage_plan), - })) + Ok(Json(dto_build::graph_to_job_response( + &graph, + query.plan_format.unwrap_or_default(), + ))) } pub async fn cancel_job< @@ -554,11 +375,6 @@ pub async fn cancel_job< } } -#[derive(Debug, serde::Serialize)] -pub struct QueryStagesResponse { - pub stages: Vec, -} - pub async fn get_query_stages< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -567,8 +383,6 @@ pub async fn get_query_stages< Path(job_id): Path, query: Query, ) -> Result { - let plan_format = query.plan_format.clone().unwrap_or_default(); - if let Some(graph) = data_server .state .task_manager @@ -582,383 +396,16 @@ pub async fn get_query_stages< ) })? { - let stages = graph - .as_ref() - .stages() - .iter() - .map(|(id, stage)| { - let mut summary = QueryStageSummary { - stage_id: id.to_string(), - stage_status: stage.variant_name().to_string(), - input_rows: 0, - output_rows: 0, - elapsed_compute: None, - tasks: vec![], - task_duration_percentiles: None, - task_input_percentiles: None, - stage_plan: None, - }; - match stage { - ExecutionStage::Running(running_stage) => { - let metrics = running_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(running_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(running_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(running_stage.plan.as_ref(), metrics), - }); - summary.input_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "input_rows")) - .unwrap_or(0); - summary.output_rows = running_stage - .stage_metrics - .as_ref() - .map(|m| get_combined_count(m.as_slice(), "output_rows")) - .unwrap_or(0); - summary.elapsed_compute = get_running_stage_time(&running_stage - .task_infos, get_current_time()); - summary.tasks = running_stage - .task_infos - .iter() - .map(|info| { - let (input_rows, output_rows) = running_stage - .stage_metrics - .as_deref() - .map(|metrics| { - get_partition_counts( - metrics, - &info.global_input_partition_ids, - ) - }) - .unwrap_or((0, 0)); - - let start_exec_time = info.start_exec_time as u64; - let end_exec_time = info.end_exec_time as u64; - - let task_status: TaskStatus = (&info.task_status).into(); - - Some(TaskSummary { - id: info.task_id, - partition_id: info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - scheduled_time: info.scheduled_time as u64, - launch_time: info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: info.finish_time as u64, - input_rows, - output_rows, - status: task_status, - }) - }) - .collect(); - } - ExecutionStage::Successful(completed_stage) => { - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(completed_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(completed_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(completed_stage.plan.as_ref(), &completed_stage.stage_metrics), - }); - summary.input_rows = get_combined_count( - &completed_stage.stage_metrics, - "input_rows", - ); - summary.output_rows = get_combined_count( - &completed_stage.stage_metrics, - "output_rows", - ); - summary.elapsed_compute = - get_finished_stage_time(&completed_stage.task_infos); - - summary.tasks = completed_stage - .task_infos - .iter() - .map(|task_info| { - let (input_rows, output_rows) = get_partition_counts( - &completed_stage.stage_metrics, - &task_info.global_input_partition_ids, - ); - - let start_exec_time = task_info.start_exec_time as u64; - let end_exec_time = task_info.end_exec_time as u64; - let task_status = (&task_info.task_status).into(); - Some(TaskSummary { - id: task_info.task_id, - partition_id: task_info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - scheduled_time: task_info.scheduled_time as u64, - launch_time: task_info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: task_info.finish_time as u64, - input_rows, - output_rows, - status: task_status, - }) - }) - .collect(); - } - ExecutionStage::Failed(failed_stage) => { - let metrics = failed_stage.stage_metrics.as_deref().unwrap_or(&[]); - summary.stage_plan = Some(match plan_format { - PlanFormat::Default => displayable(failed_stage.plan.as_ref()).indent(false).to_string(), - PlanFormat::Tree => displayable(failed_stage.plan.as_ref()).tree_render().to_string(), - PlanFormat::Metrics => format_stage_metrics(failed_stage.plan.as_ref(), metrics), - }); - summary.input_rows = get_combined_count(metrics, "input_rows"); - summary.output_rows = get_combined_count(metrics, "output_rows"); - summary.elapsed_compute = - get_finished_stage_time(&failed_stage.task_infos); - - summary.tasks = failed_stage - .task_infos - .iter() - .map(|info| { - let (input_rows, output_rows) = get_partition_counts( - metrics, - &info.global_input_partition_ids, - ); - - let start_exec_time = info.start_exec_time as u64; - let end_exec_time = info.end_exec_time as u64; - let task_status: TaskStatus = (&info.task_status).into(); - - Some(TaskSummary { - id: info.task_id, - partition_id: info - .global_input_partition_ids - .iter() - .map(|&p| p as u32) - .collect(), - scheduled_time: info.scheduled_time as u64, - launch_time: info.launch_time as u64, - start_exec_time, - end_exec_time, - exec_duration: end_exec_time.saturating_sub(start_exec_time), - finish_time: info.finish_time as u64, - input_rows, - output_rows, - status: task_status, - }) - }) - .collect(); - } - _ => {} - } - summary.task_duration_percentiles = task_duration_percentiles(&summary.tasks); - summary.task_input_percentiles = task_input_percentiles(&summary.tasks); - summary - }) - .collect(); - - Ok(Json(QueryStagesResponse { stages })) + Ok(Json(dto_build::graph_to_query_stages( + &graph, + query.plan_format.unwrap_or_default(), + get_current_time(), + ))) } else { Err(SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) } } -fn percentile_duration(sorted: &[u64], pct: f64) -> u64 { - let idx = ((pct / 100.0) * (sorted.len() - 1) as f64).round() as usize; - sorted[idx.min(sorted.len() - 1)] -} - -fn task_input_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = tasks - .iter() - .flatten() - .map(|t| t.input_rows as u64) - .collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -fn task_duration_percentiles(tasks: &[Option]) -> Option { - let mut durations: Vec = - tasks.iter().flatten().map(|t| t.exec_duration).collect(); - - if durations.is_empty() { - return None; - } - - durations.sort_unstable(); - - Some(Percentiles { - min: durations[0], - p25: percentile_duration(&durations, 25.0), - median: percentile_duration(&durations, 50.0), - p75: percentile_duration(&durations, 75.0), - max: *durations.last().unwrap(), - }) -} - -/// Returns elapsed wall time in milliseconds for API formatting. -/// -/// Uses saturating subtraction so inconsistent timestamps (e.g. failed jobs, or -/// `end_time` still zero while `start_time` is set) do not panic on subtract. -fn job_elapsed_ms(start_time: u64, end_time: u64) -> u64 { - end_time.saturating_sub(start_time) -} - -fn format_job_status(status: &Option, elapsed_ms: u64) -> (String, String) { - match status { - Some(Status::Queued(_)) => ("Queued".to_string(), "Queued".to_string()), - Some(Status::Running(_)) => ("Running".to_string(), "Running".to_string()), - Some(Status::Failed(error)) => { - ("Failed".to_string(), format!("Failed: {}", error.error)) - } - Some(Status::Successful(completed)) => { - let num_rows = completed - .partition_location - .iter() - .map(|p| p.partition_stats.as_ref().map(|s| s.num_rows).unwrap_or(0)) - .sum::(); - let num_rows_term = if num_rows == 1 { "row" } else { "rows" }; - let num_partitions = completed.partition_location.len(); - let num_partitions_term = if num_partitions == 1 { - "partition" - } else { - "partitions" - }; - ( - "Completed".to_string(), - format!( - "Completed. Produced {} {} containing {} {}. Elapsed time: {} ms.", - num_partitions, - num_partitions_term, - num_rows, - num_rows_term, - elapsed_ms - ), - ) - } - _ => ("Invalid".to_string(), "Invalid State".to_string()), - } -} - -fn get_running_stage_time(task_infos: &[TaskInfo], current_time: u128) -> Option { - let min_start = task_infos - .iter() - .map(|t| t.start_exec_time) - .filter(|t| *t > 0) - .min(); - - match (min_start, current_time) { - (Some(start), end) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -fn failed_reason(failed: &FailedTask) -> String { - match &failed.failed_reason { - Some(ExecutionError(_)) => "ExecutionError", - Some(FetchPartitionError(_)) => "FetchPartitionError", - Some(IoError(_)) => "IoError", - Some(ExecutorLost(_)) => "ExecutorLost", - Some(ResultLost(_)) => "ResultLost", - Some(TaskKilled(_)) => "TaskKilled", - None => "Failed", - } - .to_string() -} - -fn get_finished_stage_time(task_infos: &[TaskInfo]) -> Option { - let min_start = task_infos - .iter() - .map(|t| t.start_exec_time) - .filter(|t| *t > 0) - .min(); - - let max_end = task_infos - .iter() - .map(|t| t.end_exec_time) - .filter(|t| *t > 0) - .max(); - - match (min_start, max_end) { - (Some(start), Some(end)) if end >= start => { - let time = Time::new(); - time.add_duration(Duration::from_millis((end - start) as u64)); - Some(time.to_string()) - } - _ => None, - } -} - -/// Sum a task's `input_rows` / `output_rows` across the global partitions the -/// task owns. Metrics are keyed by global partition id — for single-partition -/// tasks `partitions` is a one-element slice; for multi-partition tasks it is -/// the task's `global_input_partition_ids`. -fn get_partition_counts(metrics: &[MetricsSet], partitions: &[usize]) -> (usize, usize) { - let input_rows = get_partition_count(metrics, partitions, "input_rows"); - let output_rows = get_partition_count(metrics, partitions, "output_rows"); - (input_rows, output_rows) -} - -fn get_partition_count( - metrics: &[MetricsSet], - partitions: &[usize], - name: &str, -) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - let owned_by_task = metric - .partition() - .map(|p| partitions.contains(&p)) - .unwrap_or(false); - if owned_by_task && metric_value.name() == name { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - -fn get_combined_count(metrics: &[MetricsSet], name: &str) -> usize { - metrics - .iter() - .flat_map(|vec| { - vec.iter().map(|metric| { - let metric_value = metric.value(); - if metric_value.name() == name { - metric_value.as_usize() - } else { - 0 - } - }) - }) - .sum() -} - pub async fn get_job_dot_graph< T: AsLogicalPlan + Clone + Send + Sync + 'static, U: AsExecutionPlan + Send + Sync + 'static, @@ -1081,126 +528,6 @@ pub async fn get_job_config< #[cfg(test)] mod tests { use super::*; - use crate::state::execution_stage::TaskInfo; - use ballista_core::serde::protobuf::task_status; - - fn make_task_info(start: u128, end: u128) -> TaskInfo { - TaskInfo { - task_id: 0, - scheduled_time: 0, - launch_time: 0, - start_exec_time: start, - end_exec_time: end, - finish_time: 0, - task_status: task_status::Status::Running(Default::default()), - global_input_partition_ids: vec![], - vcores_consumed: 0, - } - } - - #[test] - fn test_job_elapsed_saturates_when_end_precedes_start() { - assert_eq!(job_elapsed_ms(900, 100), 0); - } - - // --- get_finished_stage_time --- - - #[test] - fn test_finished_empty_slice_returns_none() { - assert_eq!(get_finished_stage_time(&[]), None); - } - - #[test] - fn test_finished_all_zero_timestamps_returns_none() { - let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - #[test] - fn test_finished_single_task_elapsed() { - // 600 - 100 = 500 ms → "500.00ms" - let tasks = vec![make_task_info(100, 600)]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("500.00ms".to_string()) - ); - } - - #[test] - fn test_finished_picks_earliest_start_and_latest_end() { - // min start = 100, max end = 900 → 800 ms - let tasks = vec![ - make_task_info(100, 500), - make_task_info(200, 900), - make_task_info(300, 700), - ]; - assert_eq!( - get_finished_stage_time(&tasks), - Some("800.00ms".to_string()) - ); - } - - #[test] - fn test_finished_end_before_start_returns_none() { - let tasks = vec![make_task_info(900, 100)]; - assert_eq!(get_finished_stage_time(&tasks), None); - } - - // --- get_running_stage_time --- - - #[test] - fn test_running_empty_slice_returns_none() { - assert_eq!(get_running_stage_time(&[], 1000), None); - } - - #[test] - fn test_running_all_zero_start_returns_none() { - let tasks: Vec = vec![make_task_info(0, 0), make_task_info(0, 0)]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_future_start_returns_none() { - // start_exec_time beyond current time → elapsed clamped to 0 - let tasks = vec![make_task_info(u128::MAX, 0)]; - assert_eq!(get_running_stage_time(&tasks, 1000), None); - } - - #[test] - fn test_running_past_start_returns_some() { - let now = 4_000; - let start = 1_000; - let tasks = vec![make_task_info(start, 0)]; - assert_eq!( - get_running_stage_time(&tasks, now), - Some("3.00s".to_string()) - ); - } - - #[test] - fn test_running_mixed_zero_start_uses_earliest_nonzero() { - let now = 3_000; - let earlier = 1_000; - let later = 2_000; - let tasks = vec![ - make_task_info(0, 0), - make_task_info(later, 0), - make_task_info(earlier, 0), - make_task_info(0, 0), - ]; - let result = get_running_stage_time(&tasks, now); - assert_eq!(result, Some("2.00s".to_string())); - } - - #[test] - fn test_job_elapsed_ms_normal() { - assert_eq!(super::job_elapsed_ms(100, 500), 400); - } - - #[test] - fn test_job_elapsed_ms_end_before_start_saturates_to_zero() { - assert_eq!(super::job_elapsed_ms(500, 100), 0); - } mod get_webtui { use super::*; diff --git a/ballista/scheduler/src/api/mod.rs b/ballista/scheduler/src/api/mod.rs index 81d7177cfa..527e7610b8 100644 --- a/ballista/scheduler/src/api/mod.rs +++ b/ballista/scheduler/src/api/mod.rs @@ -10,6 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "rest-api")] +mod dto_build; #[cfg(feature = "rest-api")] mod handlers; mod health; diff --git a/dev/release/README.md b/dev/release/README.md index 3f34ba10e5..c6b0de31c0 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -287,6 +287,7 @@ of the following crates: - [ballista-cli](https://crates.io/crates/ballista-cli) - [ballista-core](https://crates.io/crates/ballista-core) - [ballista-executor](https://crates.io/crates/ballista-executor) +- [ballista-api-types](https://crates.io/crates/ballista-api-types) - [ballista-scheduler](https://crates.io/crates/ballista-scheduler) Download and unpack the official release tarball @@ -306,6 +307,7 @@ dot -Tsvg dev/release/crate-deps.dot > dev/release/crate-deps.svg ```shell (cd ballista/core && cargo publish) (cd ballista/executor && cargo publish) +(cd ballista/api-types && cargo publish) (cd ballista/scheduler && cargo publish) (cd ballista/client && cargo publish) (cd ballista-cli && cargo publish) diff --git a/dev/release/crate-deps.dot b/dev/release/crate-deps.dot index c3113e6a3d..27bee709ba 100644 --- a/dev/release/crate-deps.dot +++ b/dev/release/crate-deps.dot @@ -18,12 +18,14 @@ digraph G { ballista_core + ballista_api_types ballista_scheduler ballista_executor ballista ballista_cli ballista_scheduler -> ballista_core + ballista_scheduler -> ballista_api_types ballista_executor -> ballista_core diff --git a/dev/update_ballista_versions.py b/dev/update_ballista_versions.py index 8fedb21023..ca5bef3bf8 100755 --- a/dev/update_ballista_versions.py +++ b/dev/update_ballista_versions.py @@ -43,6 +43,7 @@ def update_cargo_toml(cargo_toml: str, new_version: str): 'ballista', 'ballista-core', 'ballista-executor', + 'ballista-api-types', 'ballista-scheduler', 'ballista-cli', ) @@ -80,6 +81,7 @@ def main(): for rel_path in [ 'ballista-cli', 'ballista/core', + 'ballista/api-types', 'ballista/scheduler', 'ballista/executor', 'ballista/client',