diff --git a/.cursor/rules.md b/.cursor/rules.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/.cursor/rules.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/.gemini/rules.md b/.gemini/rules.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/.gemini/rules.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..4ee2be4ad3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,5 @@ +When making a code review never link to any issue or a pull request +in any GitHub repository! + +Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md! + diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fd3060297d..84ec3a8bbe 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -212,7 +212,7 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| ConfigEntry::new(BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE.to_string(), "Configuration for max message size in gRPC clients".to_string(), DataType::UInt64, - Some((16 * 1024 * 1024).to_string())), + Some((128 * 1024 * 1024).to_string())), ConfigEntry::new(BALLISTA_CLIENT_GRPC_CONNECT_TIMEOUT_SECONDS.to_string(), "Connection timeout for gRPC client in seconds".to_string(), DataType::UInt64, @@ -934,10 +934,37 @@ mod tests { #[test] fn default_config() -> Result<()> { let config = BallistaConfig::default(); - assert_eq!(16777216, config.grpc_client_max_message_size()); + assert_eq!(134217728, config.grpc_client_max_message_size()); 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(), + 128 * 1024 * 1024, + "unexpected default (did the bump land?)", + ); + + 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/core/src/utils.rs b/ballista/core/src/utils.rs index 34b5b05f10..6ce842b82c 100644 --- a/ballista/core/src/utils.rs +++ b/ballista/core/src/utils.rs @@ -107,7 +107,7 @@ impl Default for GrpcClientConfig { tcp_keepalive_seconds: 3600, http2_keepalive_interval_seconds: 300, use_tls: false, - max_message_size: 16 * 1024 * 1024, + max_message_size: 128 * 1024 * 1024, io_retries_times: 3, io_retry_wait_time_ms: 3000, initial_connection_window_size: 67108864, diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs index 2b0ffe608d..1071b345d7 100644 --- a/ballista/executor/src/config.rs +++ b/ballista/executor/src/config.rs @@ -169,17 +169,17 @@ 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: 128MB). #[arg( long, - default_value_t = 16777216, + default_value_t = 134217728, 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: 128MB). #[arg( long, - default_value_t = 16777216, + default_value_t = 134217728, help = "The maximum size of an encoded message at the grpc server side." )] pub grpc_server_max_encoding_message_size: u32, diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 091a94c5f9..e70397ecb1 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -375,8 +375,8 @@ impl Default for ExecutorProcessConfig { log_rotation_policy: Default::default(), job_data_ttl_seconds: 604800, job_data_clean_up_interval_seconds: 0, - grpc_max_decoding_message_size: 16777216, - grpc_max_encoding_message_size: 16777216, + grpc_max_decoding_message_size: 134217728, + grpc_max_encoding_message_size: 134217728, grpc_server_config: Default::default(), executor_heartbeat_interval_seconds: 60, metric_collection_policy: ExecutorMetricCollectionPolicy::default(), diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 940888714b..459835427a 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -148,17 +148,27 @@ pub struct Config { /// Maximum size of decoded gRPC messages. #[arg( long, - default_value_t = 16777216, + default_value_t = 134217728, help = "The maximum size of a decoded message at the grpc server side." )] pub grpc_server_max_decoding_message_size: u32, /// Maximum size of encoded gRPC messages. #[arg( long, - default_value_t = 16777216, + default_value_t = 134217728, 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 = 134217728, + 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 @@ -305,8 +317,9 @@ impl Default for SchedulerConfig { job_resubmit_interval_ms: None, executor_termination_grace_period: 0, scheduler_event_expected_processing_duration: 0, - grpc_server_max_decoding_message_size: 16777216, - grpc_server_max_encoding_message_size: 16777216, + grpc_server_max_decoding_message_size: 134217728, + grpc_server_max_encoding_message_size: 134217728, + grpc_client_max_message_size: 134217728, 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 d8a8331778..910350cd40 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, diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c40f1a582c..8622b50b7a 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -109,7 +109,7 @@ standard DataFusion settings. | ballista.cache.noop | Boolean | true | Disable default cache node extension | | ballista.client.grpc_connect_timeout_seconds | UInt64 | 20 | Connection timeout for gRPC client in seconds | | ballista.client.grpc_http2_keepalive_interval_seconds | UInt64 | 300 | HTTP/2 keep-alive interval for gRPC client in seconds | -| ballista.client.grpc_max_message_size | UInt64 | 16777216 | Configuration for max message size in gRPC clients | +| ballista.client.grpc_max_message_size | UInt64 | 134217728 | Configuration for max message size in gRPC clients | | ballista.client.grpc_tcp_keepalive_seconds | UInt64 | 3600 | TCP keep-alive interval for gRPC client in seconds | | ballista.client.grpc_timeout_seconds | UInt64 | 20 | Request timeout for gRPC client in seconds | | ballista.client.initial_connection_window_size | UInt64 | 67108864 | HTTP/2 initial connection-level flow-control window for gRPC data-plane clients, in bytes. Should be >= the shuffle governor byte budget so the governor, not the transport window, is the binding backpressure. 0 leaves the tonic default. |