diff --git a/native/native-executors/Cargo.lock b/native/native-executors/Cargo.lock index 4d21ac8a3e5..452fe6d373c 100644 --- a/native/native-executors/Cargo.lock +++ b/native/native-executors/Cargo.lock @@ -996,6 +996,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "ballista-api-types" +version = "55.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97d9eea12c759eda5f828b4f0eba2b42784d96a4e7858afe5624a9522507841" +dependencies = [ + "serde", + "utoipa", +] + [[package]] name = "ballista-core" version = "55.0.0" @@ -1023,6 +1033,7 @@ dependencies = [ "tonic-prost", "tonic-prost-build", "url", + "utoipa", "uuid", ] @@ -1056,6 +1067,19 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "55.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "926ed4783d3b4d24fc28e899c7aef2f0127a511238320eadfc06c65ba2772bf8" +dependencies = [ + "ballista-api-types", + "log", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "ballista-scheduler" version = "55.0.0" @@ -1065,7 +1089,9 @@ dependencies = [ "arrow-flight", "async-trait", "axum", + "ballista-api-types", "ballista-core", + "ballista-history", "dashmap", "datafusion", "datafusion-proto", @@ -1079,11 +1105,14 @@ dependencies = [ "prost", "rand 0.10.3", "serde", + "serde_json", "tokio", "tokio-stream", "tonic", "tower-http 0.7.1", + "tracing", "url-escape", + "utoipa", "uuid", ] @@ -3471,7 +3500,7 @@ dependencies = [ [[package]] name = "iceberg" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "aes-gcm", "anyhow", @@ -3529,7 +3558,7 @@ dependencies = [ [[package]] name = "iceberg-property-macro" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "proc-macro2", "quote", @@ -3539,7 +3568,7 @@ dependencies = [ [[package]] name = "iceberg-storage-opendal" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "anyhow", "async-trait", @@ -6880,6 +6909,29 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "utoipa" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8765fe27aeff71012a3f90fa474cf1df863a7d0dd81ed25de4e63fff2544994f" +dependencies = [ + "indexmap 2.14.2", + "serde", + "serde_json", + "utoipa-gen", +] + +[[package]] +name = "utoipa-gen" +version = "6.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d935f1c83fdf8b88f09bbe050d0cea78ea4afee6d7b0317403228c1200b2c8ab" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[package]] name = "uuid" version = "1.29.0" diff --git a/native/native-executors/Cargo.toml b/native/native-executors/Cargo.toml index 117cd7e3df4..388c06189aa 100644 --- a/native/native-executors/Cargo.toml +++ b/native/native-executors/Cargo.toml @@ -37,7 +37,9 @@ publish = false arrow-flight = "59.2.0" ballista-core = "55.0.0" ballista-executor = { version = "55.0.0", default-features = false, features = ["arrow-ipc-optimizations"] } -ballista-scheduler = { version = "55.0.0", default-features = false } +# `rest-api` compiles the scheduler's event-log writer (COMET_SCHEDULER_EVENT_LOG_DIR) and its REST +# API. In Ballista 55.0.0 it logs through `tracing` without enabling it, so that is enabled here. +ballista-scheduler = { version = "55.0.0", default-features = false, features = ["rest-api", "tracing"] } datafusion = { version = "55.2.0", default-features = false } # Comet's core crate, under the name of its library. Without the default hdfs-opendal feature, # which links libjvm through libhdfs. diff --git a/native/native-executors/src/bin/comet-executor.rs b/native/native-executors/src/bin/comet-executor.rs index f1b7cec6a99..9f0cfc8f78f 100644 --- a/native/native-executors/src/bin/comet-executor.rs +++ b/native/native-executors/src/bin/comet-executor.rs @@ -26,6 +26,17 @@ //! - `COMET_EXECUTOR_GRPC_PORT` (default `50052`) //! - `COMET_SCHEDULER_HOST` (default `localhost`) and `COMET_SCHEDULER_PORT` (default `50050`) //! - `COMET_EXECUTOR_TASKS`: concurrent tasks (default: the number of CPUs) +//! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest message, in bytes, the executor sends or accepts +//! on Ballista's gRPC connections with the scheduler, such as a task with many partitions of a +//! scan, and on its Flight service (default: Ballista's 16 MiB). Set it on the scheduler as +//! well. The executor's reads of other executors' shuffle data keep the limit of the query's +//! `ballista.client.grpc_max_message_size` (16 MiB by default), which this does not change. +//! - `COMET_EXECUTOR_MEMORY_POOL_SIZE`: the executor's memory budget in bytes, `0` for no limit. +//! It is not one pool that the tasks share: each task gets its own pool of the budget divided +//! by `COMET_EXECUTOR_TASKS`, times the task slots it takes, even when it runs alone. By +//! default the budget is 70% of the smaller of the container's memory limit and the host's +//! memory, or no limit when neither can be found. Memory outside the pools, such as batches in +//! flight between stages, comes on top of it. //! - `COMET_EXECUTOR_WORK_DIR`: where shuffle and spill files go (default: a temporary directory) use std::sync::Arc; @@ -61,6 +72,13 @@ async fn main() -> Result<(), Box> { config.vcores = tasks; } config.work_dir = env_var("COMET_EXECUTOR_WORK_DIR")?; + if let Some(size) = env_var("COMET_EXECUTOR_MEMORY_POOL_SIZE")? { + config.memory_pool_size = Some(size); + } + if let Some(size) = env_var("COMET_GRPC_MAX_MESSAGE_SIZE")? { + config.grpc_max_decoding_message_size = size; + config.grpc_max_encoding_message_size = size; + } start_executor_process(Arc::new(config)).await?; Ok(()) diff --git a/native/native-executors/src/bin/comet-scheduler.rs b/native/native-executors/src/bin/comet-scheduler.rs index d1506f335ac..ba68c19a1d6 100644 --- a/native/native-executors/src/bin/comet-scheduler.rs +++ b/native/native-executors/src/bin/comet-scheduler.rs @@ -20,12 +20,35 @@ //! //! Configured through environment variables, with Ballista's defaults otherwise: //! - `COMET_SCHEDULER_BIND_HOST` (default `127.0.0.1`) +//! - `COMET_SCHEDULER_EXTERNAL_HOST`: the host executors reach this scheduler at, which they +//! report task status to (default `localhost`) //! - `COMET_SCHEDULER_PORT`: Ballista's gRPC port, for executors (default `50050`) //! - `COMET_SCHEDULER_FLIGHT_PORT`: the Flight port, for drivers (default `50060`) //! - `COMET_SCHEDULER_RESULT_RETENTION_SECONDS`: how long a query's results stay readable after //! the query ends, whether or not a driver has read them (default `300`). `0` keeps them. +//! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest message, in bytes, of Ballista's gRPC service +//! between the scheduler and its executors, such as a task with many partitions of a scan +//! (default: Ballista's 16 MiB). Set it on the executors as well, for their side of the same +//! messages. The Flight service for drivers has no limit. Shuffle reads between executors keep +//! the limit of the query's `ballista.client.grpc_max_message_size` (16 MiB by default), which +//! this does not change. +//! - `COMET_SCHEDULER_EVENT_LOG_DIR`: a directory to write one event log per query to, as +//! Ballista's scheduler does (default: no event logs). The scheduler creates it if needed and +//! doesn't start if it can't write there. A log has the query's plan and a record for each +//! stage and task, with the task's times and the rows it wrote. The log sees a stage's Comet +//! block as one operator with no metrics, so the plan shows a block only by its size, input +//! rows count only the rows a stage reads from other stages (not those it reads from files), +//! and a task's compute time covers only that reading. A task's start and end times give how +//! long it ran. +//! The task records are best effort: when the writer falls behind, Ballista drops them, so a +//! log can have fewer records than the query had tasks. Logs stay in the directory until +//! something removes them; `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` applies to query +//! results, not to logs. +//! +//! Ballista's REST API, which lists and cancels jobs without authentication, is not served. use std::net::SocketAddr; +use std::path::Path; use std::sync::Arc; use arrow_flight::flight_service_server::FlightServiceServer; @@ -43,17 +66,38 @@ async fn main() -> Result<(), Box> { init_logging(); let mut config = SchedulerConfig { override_physical_codec: Some(Arc::new(CometPhysicalCodec::default())), + // Compiling the event-log writer compiles the REST API in with it, on the gRPC port. + disable_rest_api: true, ..Default::default() }; if let Some(host) = env_var("COMET_SCHEDULER_BIND_HOST")? { config.bind_host = host; } + if let Some(host) = env_var("COMET_SCHEDULER_EXTERNAL_HOST")? { + config.external_host = host; + } + if let Some(size) = env_var("COMET_GRPC_MAX_MESSAGE_SIZE")? { + // The scheduler is a gRPC server for the executors and also a client of theirs, and has + // a setting for the client. An executor has one pair of settings, for its server and for + // its client of the scheduler. + config.grpc_server_max_decoding_message_size = size; + config.grpc_server_max_encoding_message_size = size; + config.grpc_client_max_message_size = size; + } if let Some(port) = env_var("COMET_SCHEDULER_PORT")? { config.bind_port = port; } if let Some(seconds) = env_var("COMET_SCHEDULER_RESULT_RETENTION_SECONDS")? { config.finished_job_data_clean_up_interval_seconds = seconds; } + if let Some(dir) = env_var::("COMET_SCHEDULER_EVENT_LOG_DIR")? { + // Ballista creates the directory on a background task, and when it can't write there it + // only logs warnings. Fail at startup instead, as for a malformed setting. + check_writable(Path::new(&dir)).map_err(|e| { + format!("Cannot write event logs to COMET_SCHEDULER_EVENT_LOG_DIR {dir:?}: {e}") + })?; + config.event_log_dir = Some(dir); + } let flight_port: u16 = env_var("COMET_SCHEDULER_FLIGHT_PORT")?.unwrap_or(50060); let grpc_addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port).parse()?; let flight_addr = SocketAddr::new(grpc_addr.ip(), flight_port); @@ -79,3 +123,17 @@ async fn main() -> Result<(), Box> { } Ok(()) } + +/// Creates `dir` if it doesn't exist, and checks that the scheduler can create files in it: a +/// directory that already exists, such as a volume mounted with another owner, may not allow it. +fn check_writable(dir: &Path) -> std::io::Result<()> { + std::fs::create_dir_all(dir)?; + // Schedulers in containers are often PID 1, so the name also has the time, for two that share + // a log volume. + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()); + let probe = dir.join(format!(".comet-scheduler-{}-{nanos}", std::process::id())); + std::fs::File::create(&probe)?; + std::fs::remove_file(probe) +} diff --git a/native/native-executors/tests/cluster.rs b/native/native-executors/tests/cluster.rs index dee8dcaa7b5..2ed61c2cebf 100644 --- a/native/native-executors/tests/cluster.rs +++ b/native/native-executors/tests/cluster.rs @@ -21,6 +21,7 @@ mod common; use std::future::Future; +use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream}; use std::path::PathBuf; use std::process::{Child, Command, Stdio}; @@ -43,6 +44,8 @@ struct Cluster { processes: Vec, work_dir: PathBuf, flight_url: String, + /// The port of Ballista's gRPC service, for executors. + scheduler_port: u16, } /// Held while a cluster starts, so that two tests don't pick the same free port in between. @@ -67,6 +70,7 @@ impl Cluster { std::process::id() )), flight_url: format!("http://127.0.0.1:{flight_port}"), + scheduler_port, }; let mut env = vec![ ("COMET_SCHEDULER_BIND_HOST", "127.0.0.1".to_string()), @@ -166,6 +170,27 @@ fn free_port() -> u16 { .port() } +/// The HTTP status of a GET of `path` on a local port. +fn http_status(port: u16, path: &str) -> u16 { + let mut stream = TcpStream::connect(("127.0.0.1", port)).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .unwrap(); + write!( + stream, + "GET {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n" + ) + .unwrap(); + let mut response = vec![]; + stream.read_to_end(&mut response).unwrap(); + let response = String::from_utf8_lossy(&response); + response + .split_whitespace() + .nth(1) + .and_then(|status| status.parse().ok()) + .unwrap_or_else(|| panic!("Not an HTTP response: {response}")) +} + fn wait_for(port: u16) { let deadline = Instant::now() + Duration::from_secs(60); while TcpStream::connect(("127.0.0.1", port)).is_err() { @@ -263,3 +288,83 @@ async fn endpoints_say_when_their_results_are_deleted() { tokio::time::sleep(Duration::from_millis(200)).await; } } + +#[tokio::test(flavor = "multi_thread")] +async fn writes_event_logs_without_serving_the_rest_api() { + let log_dir = std::env::temp_dir().join(format!( + "comet-native-executors-event-logs-{}-{}", + std::process::id(), + free_port() + )); + let cluster = Cluster::start_with(&[( + "COMET_SCHEDULER_EVENT_LOG_DIR", + log_dir.to_string_lossy().into_owned(), + )]); + cluster.run(&query(vec![range_stage(100, 2)])).await; + + // The scheduler renames a query's log to `.eventlog` once it has written all of it. + let deadline = Instant::now() + Duration::from_secs(60); + let log = loop { + let finished = std::fs::read_dir(&log_dir).ok().and_then(|entries| { + entries + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .find(|path| { + path.extension() + .is_some_and(|extension| extension == "eventlog") + }) + }); + if let Some(log) = finished { + break log; + } + assert!(Instant::now() < deadline, "No finished event log"); + tokio::time::sleep(Duration::from_millis(100)).await; + }; + assert!(!std::fs::read_to_string(&log).unwrap().is_empty()); + + // Compiling the event-log writer compiles Ballista's REST API in, which lists and cancels + // jobs without authentication on the scheduler's gRPC port. The scheduler doesn't serve it. + assert_eq!(http_status(cluster.scheduler_port, "/api/jobs"), 404); + + drop(cluster); + let _ = std::fs::remove_dir_all(&log_dir); +} + +#[test] +fn refuses_an_event_log_dir_it_cannot_write() { + // No one can create a directory inside a file, not even root, as in CI's container. + let file = std::env::temp_dir().join(format!( + "comet-native-executors-not-a-dir-{}-{}", + std::process::id(), + free_port() + )); + std::fs::write(&file, "").unwrap(); + let mut scheduler = Command::new(env!("CARGO_BIN_EXE_comet-scheduler")) + .env("COMET_SCHEDULER_PORT", free_port().to_string()) + .env("COMET_SCHEDULER_FLIGHT_PORT", free_port().to_string()) + .env("COMET_SCHEDULER_EVENT_LOG_DIR", file.join("logs")) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(60); + let status = loop { + if let Some(status) = scheduler.try_wait().unwrap() { + break status; + } + if Instant::now() > deadline { + let _ = scheduler.kill(); + panic!("The scheduler started with an event-log directory it can't write to"); + } + std::thread::sleep(Duration::from_millis(100)); + }; + let mut stderr = String::new(); + scheduler + .stderr + .take() + .unwrap() + .read_to_string(&mut stderr) + .unwrap(); + let _ = std::fs::remove_file(&file); + assert!(!status.success()); + assert!(stderr.contains("COMET_SCHEDULER_EVENT_LOG_DIR"), "{stderr}"); +}