This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 9fecadae8 feat: make gRPC max message size consistently configurable
(#2173)
9fecadae8 is described below
commit 9fecadae8019b595cce2680f5c4964d76bf2c424
Author: Andy Grove <[email protected]>
AuthorDate: Tue Aug 4 17:09:17 2026 -0600
feat: make gRPC max message size consistently configurable (#2173)
---
ballista/core/src/config.rs | 26 ++++++++++
ballista/executor/src/config.rs | 4 +-
ballista/scheduler/src/config.rs | 20 ++++++++
ballista/scheduler/src/state/executor_manager.rs | 65 +++++++++++++++++++++++-
4 files changed, 111 insertions(+), 4 deletions(-)
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index b3f41ee32..b2d7ed662 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -966,6 +966,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 9fb01296c..ff935374b 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,
@@ -264,6 +274,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
@@ -321,6 +333,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,
no_executors_grace_period_seconds: 30,
@@ -448,6 +461,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,
@@ -561,6 +580,7 @@ impl TryFrom<Config> 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<dyn ClusterState>,
config: Arc<SchedulerConfig>,
) -> 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
+ );
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
