diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fd3060297..58e590163 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -938,6 +938,32 @@ mod tests { Ok(()) } + #[test] + fn cluster_config_plumbs_grpc_max_message_size() { + // Sanity check that setting `ballista.client.grpc_max_message_size` + // via `SessionConfig::options_mut().set(...)` — the path Python + // clients' `cluster_config` overrides use — actually reaches the + // reader consulted by `distributed_query`. + use crate::extension::SessionConfigExt; + use datafusion::prelude::SessionConfig; + + let mut config = SessionConfig::new_with_ballista(); + assert_eq!( + config.ballista_config().grpc_client_max_message_size(), + 16 * 1024 * 1024, + ); + + let r = config + .options_mut() + .set("ballista.client.grpc_max_message_size", "67108864"); + assert!(r.is_ok(), "set failed: {r:?}"); + + assert_eq!( + config.ballista_config().grpc_client_max_message_size(), + 64 * 1024 * 1024, + ); + } + // The default must stay non-zero: `0` disables the fit check, and since AQE // no longer consults `prefer_hash_join`, a disabled check would leave the // Partitioned arm unconditionally on hash join and make SortMergeJoin — the diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs index 2b0ffe608..9dbd4c78c 100644 --- a/ballista/executor/src/config.rs +++ b/ballista/executor/src/config.rs @@ -169,14 +169,14 @@ pub struct Config { help = "Tracing log rotation policy." )] pub log_rotation_policy: ballista_core::config::LogRotationPolicy, - /// Maximum size of incoming gRPC messages in bytes (default: 16MB). + /// Maximum size of incoming gRPC messages in bytes (default: 16MiB). #[arg( long, default_value_t = 16777216, help = "The maximum size of a decoded message at the grpc server side." )] pub grpc_server_max_decoding_message_size: u32, - /// Maximum size of outgoing gRPC messages in bytes (default: 16MB). + /// Maximum size of outgoing gRPC messages in bytes (default: 16MiB). #[arg( long, default_value_t = 16777216, diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 940888714..4bdea0283 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -159,6 +159,16 @@ pub struct Config { help = "The maximum size of an encoded message at the grpc server side." )] pub grpc_server_max_encoding_message_size: u32, + /// Maximum size of messages sent by the scheduler's outbound gRPC clients + /// (e.g. task assignment to executors). Should be at least as large as the + /// executor's `--grpc-server-max-decoding-message-size` for those RPCs to + /// succeed with big encoded plans. + #[arg( + long, + default_value_t = 16777216, + help = "The maximum size of a message sent by the scheduler's outbound gRPC clients (in bytes)." + )] + pub grpc_client_max_message_size: u32, /// Timeout in seconds before marking an executor as dead. #[arg( long, @@ -256,6 +266,8 @@ pub struct SchedulerConfig { pub grpc_server_max_decoding_message_size: u32, /// The maximum size of an encoded message at the grpc server side. pub grpc_server_max_encoding_message_size: u32, + /// The maximum size of a message sent by the scheduler's outbound gRPC clients. + pub grpc_client_max_message_size: u32, /// The executor timeout in seconds. It should be longer than executor's heartbeat intervals. pub executor_timeout_seconds: u64, /// The interval to check expired or dead executors @@ -307,6 +319,7 @@ impl Default for SchedulerConfig { scheduler_event_expected_processing_duration: 0, grpc_server_max_decoding_message_size: 16777216, grpc_server_max_encoding_message_size: 16777216, + grpc_client_max_message_size: 16777216, executor_timeout_seconds: 180, expire_dead_executor_interval_seconds: 15, override_config_producer: None, @@ -426,6 +439,12 @@ impl SchedulerConfig { self } + /// Sets the maximum message size for the scheduler's outbound gRPC clients. + pub fn with_grpc_client_max_message_size(mut self, value: u32) -> Self { + self.grpc_client_max_message_size = value; + self + } + /// Sets a custom config producer. pub fn with_override_config_producer( mut self, @@ -539,6 +558,7 @@ impl TryFrom for SchedulerConfig { .grpc_server_max_decoding_message_size, grpc_server_max_encoding_message_size: opt .grpc_server_max_encoding_message_size, + grpc_client_max_message_size: opt.grpc_client_max_message_size, executor_timeout_seconds: opt.executor_timeout_seconds, expire_dead_executor_interval_seconds: opt .expire_dead_executor_interval_seconds, diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index d8a833177..2f0b162b0 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -79,13 +79,22 @@ impl ExecutorManager { cluster_state: Arc, config: Arc, ) -> Self { + // Prefer an explicit override_config_producer if the embedder wired one, + // so a full BallistaConfig (with all its grpc-client knobs) still takes + // precedence. Otherwise, use `default()` but override + // `max_message_size` from the scheduler's `grpc_client_max_message_size` + // CLI flag so users can raise the ceiling for outbound task-assignment + // RPCs without having to write a config-producer in Rust. let grpc_client_config = if let Some(config_producer) = &config.override_config_producer { let session_config = config_producer(); let ballista_config = session_config.ballista_config(); GrpcClientConfig::from(&ballista_config) } else { - GrpcClientConfig::default() + GrpcClientConfig { + max_message_size: config.grpc_client_max_message_size as usize, + ..GrpcClientConfig::default() + } }; Self { cluster_state, @@ -552,7 +561,13 @@ impl ExecutorManager { } let connection = endpoint.connect().await?; - let client = ExecutorGrpcClient::new(connection); + // Message-size limits are tonic codec settings, not `Endpoint` + // settings, so `create_grpc_client_endpoint` cannot apply them. + // Without this the configured `max_message_size` is silently + // ignored and task assignment falls back to tonic's own defaults. + let client = ExecutorGrpcClient::new(connection) + .max_encoding_message_size(grpc_client_config.max_message_size) + .max_decoding_message_size(grpc_client_config.max_message_size); { self.clients.insert(executor_id.to_owned(), client.clone()); @@ -581,3 +596,49 @@ impl ExecutorManager { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::test_cluster_context; + use ballista_core::extension::SessionConfigExt; + use datafusion::prelude::SessionConfig; + + #[test] + fn grpc_client_max_message_size_flag_reaches_client_config() { + let config = Arc::new( + SchedulerConfig::default() + .with_grpc_client_max_message_size(64 * 1024 * 1024), + ); + let manager = + ExecutorManager::new(test_cluster_context().cluster_state(), config); + + assert_eq!( + manager.grpc_client_config.max_message_size, + 64 * 1024 * 1024 + ); + } + + #[test] + fn config_producer_still_wins_over_the_flag() { + let config = Arc::new( + SchedulerConfig::default() + .with_grpc_client_max_message_size(64 * 1024 * 1024) + .with_override_config_producer(Arc::new(|| { + let mut session_config = SessionConfig::new_with_ballista(); + session_config + .options_mut() + .set("ballista.client.grpc_max_message_size", "33554432") + .expect("valid setting"); + session_config + })), + ); + let manager = + ExecutorManager::new(test_cluster_context().cluster_state(), config); + + assert_eq!( + manager.grpc_client_config.max_message_size, + 32 * 1024 * 1024 + ); + } +}