feat: add cluster settings to the Comet Native Executor binaries

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.

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<dyn std::error::Error>> {
         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(())

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<dyn std::error::Error>> {
     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::<String>("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<dyn std::error::Error>> {
     }
     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)
+}

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<Child>,
     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 `<job id>.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}");
+}