feat: make gRPC max message size consistently configurable

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

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,

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,

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
+        );
+    }
+}