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 90f8096ed feat(executor): default to a bounded, auto-sized memory pool
(#2160)
90f8096ed is described below
commit 90f8096ed340d7e035efdc6e12eff566a8111e6c
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 24 05:46:33 2026 -0600
feat(executor): default to a bounded, auto-sized memory pool (#2160)
---
ballista/executor/src/config.rs | 45 +++-
ballista/executor/src/executor_process.rs | 342 +++++++++++++++++++++++++++++-
docs/source/user-guide/tuning-guide.md | 43 ++--
3 files changed, 401 insertions(+), 29 deletions(-)
diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs
index 4a7b638d3..2b0ffe608 100644
--- a/ballista/executor/src/config.rs
+++ b/ballista/executor/src/config.rs
@@ -42,6 +42,21 @@ fn parse_memory_pool_size(s: &str) -> Result<u64, String> {
.map_err(|e| format!("invalid byte size '{s}': {e}"))
}
+/// Parse and validate the `--memory-pool-fraction` value, which must be in
+/// the range `(0.0, 1.0]`.
+fn parse_memory_pool_fraction(s: &str) -> Result<f64, String> {
+ let f = s
+ .parse::<f64>()
+ .map_err(|e| format!("invalid fraction '{s}': {e}"))?;
+ if f > 0.0 && f <= 1.0 {
+ Ok(f)
+ } else {
+ Err(format!(
+ "memory pool fraction must be in (0.0, 1.0], got {f}"
+ ))
+ }
+}
+
/// Command-line arguments for configuring a Ballista executor.
///
/// This struct is parsed from command-line arguments using clap and contains
@@ -183,15 +198,26 @@ pub struct Config {
help = "Metric collection policy of this executor instance"
)]
pub metric_collection_policy: ExecutorMetricCollectionPolicy,
- /// Optional total memory budget for the executor. Accepts human-readable
- /// values like "8GB", "512MiB", or a plain byte count. When set, each
- /// per-vcore FairSpillPool gets `memory_pool_size / vcores`.
+ /// Total memory budget for the executor. When unset, defaults to
+ /// `--memory-pool-fraction` of the detected host/cgroup memory limit. `0`
+ /// disables the pool (unbounded, no spilling). Accepts human-readable
values
+ /// like "8GB", "512MiB", or a plain byte count. Split evenly across
vcores.
#[arg(
long,
value_parser = parse_memory_pool_size,
- help = "Optional total executor memory budget (e.g. \"8GB\",
\"512MiB\"). Split evenly across vcores."
+ help = "Total executor memory budget (e.g. \"8GB\"). Unset = fraction
of detected memory; 0 = unbounded. Split across vcores."
)]
pub memory_pool_size: Option<u64>,
+ /// Fraction (0.0, 1.0] of the detected host/cgroup memory limit used for
the
+ /// auto memory pool when `--memory-pool-size` is unset. Ignored when
+ /// `--memory-pool-size` is given.
+ #[arg(
+ long,
+ default_value_t =
crate::executor_process::DEFAULT_MEMORY_POOL_FRACTION,
+ value_parser = parse_memory_pool_fraction,
+ help = "Fraction (0-1] of detected memory for the auto pool. Default
0.70. Ignored if --memory-pool-size is set."
+ )]
+ pub memory_pool_fraction: f64,
/// Maximum number of sessions whose shared runtime state (object-store
/// clients, Parquet footer cache) is retained on the executor (LRU). `0`
/// disables caching.
@@ -247,6 +273,7 @@ impl TryFrom<Config> for ExecutorProcessConfig {
executor_heartbeat_interval_seconds:
opt.executor_heartbeat_interval_seconds,
metric_collection_policy: opt.metric_collection_policy,
memory_pool_size: opt.memory_pool_size,
+ memory_pool_fraction: opt.memory_pool_fraction,
session_runtime_cache_capacity: opt.session_runtime_cache_capacity,
override_execution_engine: None,
override_function_registry: None,
@@ -264,8 +291,18 @@ impl TryFrom<Config> for ExecutorProcessConfig {
#[cfg(test)]
mod tests {
+ use super::parse_memory_pool_fraction;
use super::parse_memory_pool_size;
+ #[test]
+ fn parse_fraction_accepts_valid_and_rejects_out_of_range() {
+ assert_eq!(parse_memory_pool_fraction("0.7").unwrap(), 0.7);
+ assert_eq!(parse_memory_pool_fraction("1.0").unwrap(), 1.0);
+ assert!(parse_memory_pool_fraction("0").is_err()); // must be > 0
+ assert!(parse_memory_pool_fraction("1.5").is_err()); // must be <= 1
+ assert!(parse_memory_pool_fraction("banana").is_err());
+ }
+
#[test]
fn parse_decimal_suffix() {
assert_eq!(parse_memory_pool_size("8GB").unwrap(), 8_000_000_000);
diff --git a/ballista/executor/src/executor_process.rs
b/ballista/executor/src/executor_process.rs
index a48216132..091a94c5f 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -32,7 +32,7 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use log::{error, info, warn};
-use sysinfo::{Disks, System};
+use sysinfo::{Disks, MemoryRefreshKind, System};
use tempfile::TempDir;
use tokio::fs::DirEntry;
use tokio::signal;
@@ -80,6 +80,139 @@ use crate::shutdown::ShutdownNotifier;
use crate::{ArrowFlightServerProvider, terminate};
use crate::{execution_loop, executor_server};
+/// Default fraction of the detected memory limit handed to the pool when
+/// `--memory-pool-fraction` is not set. The `FairSpillPool` only bounds memory
+/// operators register with it, so the remaining budget absorbs untracked
+/// overhead (in-flight Arrow batches, shuffle writer buffers, fragmentation).
+pub(crate) const DEFAULT_MEMORY_POOL_FRACTION: f64 = 0.70;
+
+/// How the operator interpreted `--memory-pool-size`.
+#[derive(Debug, PartialEq)]
+enum MemoryBudget {
+ /// Flag unset: auto-size to `fraction` of detected host/cgroup memory.
+ Auto { fraction: f64 },
+ /// `--memory-pool-size 0`: run DataFusion's unbounded pool.
+ Unbounded,
+ /// `--memory-pool-size N`: bound the pool to exactly `N` bytes.
+ Bytes(u64),
+}
+
+/// Classify the raw CLI value into a [`MemoryBudget`]. `fraction` is only used
+/// for the auto path.
+fn memory_budget_from_cli(opt: Option<u64>, fraction: f64) -> MemoryBudget {
+ match opt {
+ None => MemoryBudget::Auto { fraction },
+ Some(0) => MemoryBudget::Unbounded,
+ Some(n) => MemoryBudget::Bytes(n),
+ }
+}
+
+/// Where an auto/explicit pool size came from (for logging).
+#[derive(Debug, PartialEq, Eq)]
+enum PoolSource {
+ /// Explicit `--memory-pool-size N`.
+ Configured,
+ /// Auto: fraction of the cgroup memory limit.
+ AutoCgroup,
+ /// Auto: fraction of host memory.
+ AutoHost,
+}
+
+/// Why the executor is running without a bounded pool (for logging).
+#[derive(Debug, PartialEq, Eq)]
+enum UnboundedReason {
+ /// User asked for it via `--memory-pool-size 0`.
+ Explicit,
+ /// Neither host nor cgroup memory could be detected.
+ Undetected,
+}
+
+/// The resolved pool decision: a concrete byte budget or unbounded.
+#[derive(Debug, PartialEq, Eq)]
+enum ResolvedPool {
+ Bounded { bytes: u64, source: PoolSource },
+ Unbounded(UnboundedReason),
+}
+
+/// Pure resolver: turn the budget plus detected limits into a decision.
+/// `host_total` and `cgroup_limit` are only consulted for
[`MemoryBudget::Auto`].
+fn resolve_pool(
+ budget: MemoryBudget,
+ host_total: Option<u64>,
+ cgroup_limit: Option<u64>,
+) -> ResolvedPool {
+ match budget {
+ MemoryBudget::Unbounded =>
ResolvedPool::Unbounded(UnboundedReason::Explicit),
+ MemoryBudget::Bytes(n) => ResolvedPool::Bounded {
+ bytes: n,
+ source: PoolSource::Configured,
+ },
+ MemoryBudget::Auto { fraction } => {
+ let (base, source) = match (host_total, cgroup_limit) {
+ (Some(h), Some(c)) if c <= h => (c, PoolSource::AutoCgroup),
+ (Some(h), Some(_)) => (h, PoolSource::AutoHost),
+ (Some(h), None) => (h, PoolSource::AutoHost),
+ (None, Some(c)) => (c, PoolSource::AutoCgroup),
+ (None, None) => {
+ return
ResolvedPool::Unbounded(UnboundedReason::Undetected);
+ }
+ };
+ let bytes = (base as f64 * fraction) as u64;
+ ResolvedPool::Bounded { bytes, source }
+ }
+ }
+}
+
+/// Read the cgroup memory limit in bytes, or `None` if unlimited/absent.
+/// Tries cgroup v2 (`<root>/memory.max`) then v1
+/// (`<root>/memory/memory.limit_in_bytes`). A v1 "unlimited" sentinel (a value
+/// near `u64::MAX`) is returned as-is; callers clamp it with `min(host, _)`.
+fn read_cgroup_memory_limit(cgroup_root: &Path) -> Option<u64> {
+ // cgroup v2
+ if let Ok(s) = std::fs::read_to_string(cgroup_root.join("memory.max")) {
+ let s = s.trim();
+ if s == "max" {
+ return None;
+ }
+ if let Ok(v) = s.parse::<u64>() {
+ return Some(v);
+ }
+ }
+ // cgroup v1
+ if let Ok(s) =
+
std::fs::read_to_string(cgroup_root.join("memory").join("memory.limit_in_bytes"))
+ && let Ok(v) = s.trim().parse::<u64>()
+ {
+ return Some(v);
+ }
+ None
+}
+
+/// Total host memory in bytes via `sysinfo`, or `None` if unreadable.
+fn read_host_total_memory() -> Option<u64> {
+ let mut system = System::new_all();
+ system.refresh_memory_specifics(MemoryRefreshKind::nothing().with_ram());
+ let total = system.total_memory();
+ (total > 0).then_some(total)
+}
+
+/// Resolve the pool decision, reading host/cgroup limits only when
auto-sizing.
+fn detect_pool(budget: MemoryBudget) -> ResolvedPool {
+ match budget {
+ MemoryBudget::Auto { fraction } => resolve_pool(
+ MemoryBudget::Auto { fraction },
+ read_host_total_memory(),
+ // Fixed cgroup root: assumes a cgroup-namespaced container (the
+ // default for modern k8s/Docker), where this path *is* the
+ // container's own limit. In a non-namespaced setup the real limit
+ // may live under a sub-slice instead, in which case this reads
+ // the host's root limit and we fall back to host memory below.
+ read_cgroup_memory_limit(Path::new("/sys/fs/cgroup")),
+ ),
+ other => resolve_pool(other, None, None),
+ }
+}
+
/// Builds a per-task memory-pool policy: each task's runtime is rebuilt from
the
/// shared base env with a fresh [`FairSpillPool`] of size
/// `total_bytes / vcores`. The base env's disk manager, cache manager,
@@ -173,6 +306,9 @@ pub struct ExecutorProcessConfig {
/// `memory_pool_size / vcores`. When `None`, no pool is
/// installed and DataFusion falls back to its unbounded default.
pub memory_pool_size: Option<u64>,
+ /// Fraction (0.0, 1.0] of the detected host/cgroup memory limit used for
the
+ /// auto memory pool when `memory_pool_size` is `None`.
+ pub memory_pool_fraction: f64,
/// Maximum number of sessions whose shared base runtime env is retained on
/// the executor (LRU). Sharing reuses object-store clients and the Parquet
/// footer cache across a session's tasks and queries. `0` disables caching
@@ -245,6 +381,7 @@ impl Default for ExecutorProcessConfig {
executor_heartbeat_interval_seconds: 60,
metric_collection_policy:
ExecutorMetricCollectionPolicy::default(),
memory_pool_size: None,
+ memory_pool_fraction: DEFAULT_MEMORY_POOL_FRACTION,
session_runtime_cache_capacity: 16,
override_execution_engine: None,
override_function_registry: None,
@@ -330,16 +467,43 @@ pub async fn start_executor_process(
})
});
- let pool_policy: MemoryPoolPolicy = if let Some(total) =
opt.memory_pool_size {
- let policy = memory_pool_policy(total, vcores)?;
- let per_vcore = total / vcores as u64;
- info!(
- "Memory pool: total {total} bytes split into {vcores} tasks
({per_vcore} bytes each)"
- );
- policy
- } else {
- identity_pool_policy()
- };
+ let fraction = opt.memory_pool_fraction;
+ let pool_policy: MemoryPoolPolicy =
+ match detect_pool(memory_budget_from_cli(opt.memory_pool_size,
fraction)) {
+ ResolvedPool::Bounded { bytes, source } => {
+ let per_vcore = bytes / vcores as u64;
+ let source_desc = match source {
+ PoolSource::Configured => {
+ "configured via --memory-pool-size".to_string()
+ }
+ PoolSource::AutoCgroup => {
+ format!("auto: {:.0}% of cgroup memory limit",
fraction * 100.0)
+ }
+ PoolSource::AutoHost => {
+ format!("auto: {:.0}% of host memory", fraction *
100.0)
+ }
+ };
+ info!(
+ "Executor memory pool: {} ({source_desc}), \
+ split across {vcores} vcores ({}/vcore)",
+ bytesize::ByteSize::b(bytes),
+ bytesize::ByteSize::b(per_vcore),
+ );
+ memory_pool_policy(bytes, vcores)?
+ }
+ ResolvedPool::Unbounded(reason) => {
+ match reason {
+ UnboundedReason::Explicit => {
+ info!("Executor memory pool: unbounded
(--memory-pool-size 0)")
+ }
+ UnboundedReason::Undetected => warn!(
+ "Executor memory pool: unbounded (could not detect
host or \
+ cgroup memory limit; set --memory-pool-size to enable
spilling)"
+ ),
+ }
+ identity_pool_policy()
+ }
+ };
// Combined producer preserving the current per-task behavior: build a
fresh
// base env, then apply the pool policy. Used by
`Executor::produce_runtime`
@@ -1184,4 +1348,160 @@ mod memory_pool_tests {
let env = identity_pool_policy()(base.clone(), &SessionConfig::new(),
1).unwrap();
assert!(Arc::ptr_eq(&env, &base));
}
+
+ #[test]
+ fn budget_from_cli_classifies_none_zero_and_positive() {
+ assert_eq!(
+ memory_budget_from_cli(None, 0.7),
+ MemoryBudget::Auto { fraction: 0.7 }
+ );
+ assert_eq!(
+ memory_budget_from_cli(Some(0), 0.7),
+ MemoryBudget::Unbounded
+ );
+ assert_eq!(
+ memory_budget_from_cli(Some(1024), 0.7),
+ MemoryBudget::Bytes(1024)
+ );
+ }
+
+ const GB: u64 = 1024 * 1024 * 1024;
+
+ #[test]
+ fn resolve_explicit_bytes_ignores_detection() {
+ assert_eq!(
+ resolve_pool(MemoryBudget::Bytes(4 * GB), Some(32 * GB), Some(8 *
GB)),
+ ResolvedPool::Bounded {
+ bytes: 4 * GB,
+ source: PoolSource::Configured
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_unbounded_is_explicit() {
+ assert_eq!(
+ resolve_pool(MemoryBudget::Unbounded, Some(32 * GB), None),
+ ResolvedPool::Unbounded(UnboundedReason::Explicit)
+ );
+ }
+
+ #[test]
+ fn resolve_auto_prefers_cgroup_when_smaller() {
+ let expected = ((8 * GB) as f64 * 0.70) as u64;
+ assert_eq!(
+ resolve_pool(
+ MemoryBudget::Auto { fraction: 0.70 },
+ Some(32 * GB),
+ Some(8 * GB)
+ ),
+ ResolvedPool::Bounded {
+ bytes: expected,
+ source: PoolSource::AutoCgroup
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_auto_honors_custom_fraction() {
+ let expected = ((8 * GB) as f64 * 0.50) as u64;
+ assert_eq!(
+ resolve_pool(
+ MemoryBudget::Auto { fraction: 0.50 },
+ Some(32 * GB),
+ Some(8 * GB)
+ ),
+ ResolvedPool::Bounded {
+ bytes: expected,
+ source: PoolSource::AutoCgroup
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_auto_falls_back_to_host_when_cgroup_unlimited_sentinel() {
+ // cgroup v1 "unlimited" reports a value larger than host; min() picks
host.
+ let expected = ((32 * GB) as f64 * 0.70) as u64;
+ assert_eq!(
+ resolve_pool(
+ MemoryBudget::Auto { fraction: 0.70 },
+ Some(32 * GB),
+ Some(u64::MAX)
+ ),
+ ResolvedPool::Bounded {
+ bytes: expected,
+ source: PoolSource::AutoHost
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_auto_host_only() {
+ let expected = ((32 * GB) as f64 * 0.70) as u64;
+ assert_eq!(
+ resolve_pool(MemoryBudget::Auto { fraction: 0.70 }, Some(32 * GB),
None),
+ ResolvedPool::Bounded {
+ bytes: expected,
+ source: PoolSource::AutoHost
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_auto_cgroup_only() {
+ let expected = ((8 * GB) as f64 * 0.70) as u64;
+ assert_eq!(
+ resolve_pool(MemoryBudget::Auto { fraction: 0.70 }, None, Some(8 *
GB)),
+ ResolvedPool::Bounded {
+ bytes: expected,
+ source: PoolSource::AutoCgroup
+ }
+ );
+ }
+
+ #[test]
+ fn resolve_auto_undetected_is_unbounded() {
+ assert_eq!(
+ resolve_pool(MemoryBudget::Auto { fraction: 0.70 }, None, None),
+ ResolvedPool::Unbounded(UnboundedReason::Undetected)
+ );
+ }
+
+ use std::fs;
+ use tempfile::tempdir;
+
+ #[test]
+ fn cgroup_v2_reads_numeric_limit() {
+ let dir = tempdir().unwrap();
+ fs::write(dir.path().join("memory.max"), "8589934592").unwrap();
+ assert_eq!(
+ read_cgroup_memory_limit(dir.path()),
+ Some(8 * 1024 * 1024 * 1024)
+ );
+ }
+
+ #[test]
+ fn cgroup_v2_max_means_unlimited() {
+ let dir = tempdir().unwrap();
+ fs::write(dir.path().join("memory.max"), "max\n").unwrap();
+ assert_eq!(read_cgroup_memory_limit(dir.path()), None);
+ }
+
+ #[test]
+ fn cgroup_v1_reads_numeric_limit() {
+ let dir = tempdir().unwrap();
+ let v1 = dir.path().join("memory");
+ fs::create_dir_all(&v1).unwrap();
+ fs::write(v1.join("memory.limit_in_bytes"), "8589934592\n").unwrap();
+ assert_eq!(
+ read_cgroup_memory_limit(dir.path()),
+ Some(8 * 1024 * 1024 * 1024)
+ );
+ }
+
+ #[test]
+ fn cgroup_absent_returns_none() {
+ let dir = tempdir().unwrap();
+ assert_eq!(read_cgroup_memory_limit(dir.path()), None);
+ }
}
diff --git a/docs/source/user-guide/tuning-guide.md
b/docs/source/user-guide/tuning-guide.md
index 4100e0b5a..6e0ba0c51 100644
--- a/docs/source/user-guide/tuning-guide.md
+++ b/docs/source/user-guide/tuning-guide.md
@@ -102,10 +102,18 @@ decreasing the vcore count may help.
## Configuring Executor Memory Pool
-By default the executor uses DataFusion's unbounded memory pool, so spillable
-operators (sort, hash join, hash aggregate) grow until the host runs out of
-memory. To bound executor memory and let those operators spill to disk under
-pressure, pass `--memory-pool-size` when starting the executor:
+By default the executor auto-sizes a bounded memory pool from the detected
+host or cgroup memory limit, so spillable operators (sort, hash join, hash
+aggregate) spill to disk under pressure instead of growing until the host
+runs out of memory. The pool is sized to `--memory-pool-fraction` (default
+`0.70`, i.e. 70%) of the detected limit, so headroom is left for untracked
+overhead such as in-flight Arrow batches, shuffle writer buffers, and
+fragmentation. When running under a cgroup-limited container (the common
+Kubernetes/Docker case), the cgroup limit is used in preference to host
+memory when it is the tighter of the two.
+
+To use a specific budget instead of auto-detection, pass `--memory-pool-size`
+when starting the executor:
```sh
ballista-executor --memory-pool-size 8GB --vcores 8
@@ -113,20 +121,27 @@ ballista-executor --memory-pool-size 8GB --vcores 8
The argument accepts human-readable sizes (`8GB`, `512MiB`) or a plain byte
count. SI suffixes (`KB`/`MB`/`GB`) are powers of 10; IEC suffixes
-(`KiB`/`MiB`/`GiB`) are powers of 2.
-
-The total budget is divided equally across the executor's vcores: each task
-receives its own `FairSpillPool` of size `memory_pool_size / vcores`.
-With `--memory-pool-size 8GB --vcores 8`, every task sees a 1 GB
-pool, fully isolated from other tasks. Idle slots do not lend their share to
-busy ones, which keeps task memory predictable at the cost of some unused
-capacity when the executor is under-utilized.
+(`KiB`/`MiB`/`GiB`) are powers of 2. Pass `--memory-pool-size 0` to opt out of
+the pool entirely and run DataFusion's unbounded pool, matching the executor's
+pre-auto-sizing behavior.
+
+Whether the budget comes from auto-detection or `--memory-pool-size`, it is
+divided equally across the executor's vcores: each task receives its own
+`FairSpillPool` of size `memory_pool_size / vcores`. With
+`--memory-pool-size 8GB --vcores 8`, every task sees a 1 GB pool, fully
+isolated from other tasks. Idle slots do not lend their share to busy ones,
+which keeps task memory predictable at the cost of some unused capacity when
+the executor is under-utilized.
The executor refuses to start if the per-task share would round to zero (i.e.
`memory_pool_size < vcores`).
-When `--memory-pool-size` is not set, the executor behaves as before with no
-memory pool installed.
+Use `--memory-pool-fraction` to tune how much of the detected memory limit
+the auto-sized pool uses; it accepts a value in `(0.0, 1.0]` and defaults to
+`0.70`. It is ignored when `--memory-pool-size` is set.
+
+Because operators spill to disk rather than fail, make sure the executor's
+`--work-dir` has enough scratch space to absorb spills under memory pressure.
## Join Strategy
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
