diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index cf32f9dad66..bb8f07192da 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -518,6 +518,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.comet.exec.CometMemoryPoolNativeUsageSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index af4adab8634..813f64d8732 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -199,6 +199,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.comet.exec.CometMemoryPoolNativeUsageSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index c252afedc45..d7961bdd317 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -97,10 +97,12 @@ rather than asking for a separate allocation: memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction ``` -`spark.comet.exec.memoryPool.fraction` defaults to `1.0`. Lowering it is the current workaround for -Comet's under-accounting (see [The accounting gap](#the-accounting-gap)). It holds back a slice of -the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage -overshoots its reservations by roughly that slice. +`spark.comet.exec.memoryPool.fraction` defaults to `1.0`. Lowering it is one workaround for Comet's +under-accounting (see [The accounting gap](#the-accounting-gap)): it holds back a slice of the +off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage +overshoots its reservations by roughly that slice. It bounds only what Comet may reserve; the +off-heap pools compare real allocator usage against the whole `spark.memory.offHeap.size`, so +lowering the fraction to provoke spilling does not also lower the ceiling on real usage. A second value, `memory_limit_per_task`, is computed and passed alongside it, but only the on-heap pool types read it. @@ -260,8 +262,11 @@ diverge for several structural reasons: JVM closes them (see [Crossing the FFI boundary](#crossing-the-ffi-boundary)). The practical consequence is that `reserved()` is a lower bound on Comet's real footprint, and the -gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` exists purely so operators can -hand-tune a haircut that covers the gap for their workload. +gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` lets operators hold back a slice +of the pool that covers the gap for their workload, and the off-heap pools additionally compare +the allocator's real usage against `spark.memory.offHeap.size`, logging a crossing by default and +refusing the reservation when `spark.comet.exec.memoryPool.enforceNativeUsage` is set, so the gap +is at least measured rather than only estimated. To measure the gap on a real query, enable tracing with the `jemalloc` feature and compare `jemalloc_allocated` against the summed `thread_NNN_comet_memory_reserved` values; see @@ -316,11 +321,17 @@ has bounds _declared reservations_, and the sections above describe several stru declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how much they matter: -- **No signal for real native usage.** The only way to observe the gap today is to enable tracing - with the `jemalloc` feature and compare `jemalloc_allocated` against summed reservations after - the fact. There is no runtime value that an operator, a metric, or a policy could read. -- **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to - guess a per-workload haircut rather than measuring anything. +- **Real native usage is process-wide, with no per-task attribution.** `alloc_accounting` reports + one balance for the whole executor, so the off-heap pools' check cannot tell which task caused + an overrun: once any task pushes real usage past the budget, every task's next reservation sees + it, and under enforcement every one of them is denied. +- **The check gates reservations, not allocations.** An allocation that never goes through the + pool is counted after the fact and is never refused, so real usage can still exceed the budget + between reservations. Enforcement therefore falls on the operators that do reserve, which are + not necessarily the ones responsible for the overshoot, and spilling releases only reserved + bytes so it may not relieve an overshoot that lives in untracked allocations. +- **`spark.comet.exec.memoryPool.fraction` is still set by hand.** It asks operators to guess how + much of the pool to hold back, even though the overrun it guards against is now measured. - **`CometArrowAllocator` is unbounded** and participates in no budget. - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be resident on either side with no reservation covering it, because reservations are made and diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 9757bd8a57a..ddbf9395ed2 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -37,12 +37,17 @@ The `alloc-accounting` feature adds a second, allocator-independent measure of n wraps whichever global allocator the build selected and reports the bytes it has handed out as `native_allocated`. Unlike `jemalloc_allocated` it does not require jemalloc, and it counts only what Rust code allocated, so it can be compared against the memory pool's reservations without the -allocator's own caching in the way. The two features are independent and can be combined: +allocator's own caching in the way. It also backs the off-heap memory pools' check of real usage +against Comet's budget, described in the [tuning guide]. It is on by default, so the command above +already includes it; the two features are independent, and a build that drops the defaults can +re-add either one: ```shell -make release COMET_FEATURES="jemalloc,alloc-accounting" +cd native && cargo build --release --no-default-features --features hdfs-opendal,jemalloc ``` +[tuning guide]: ../user-guide/latest/tuning.md + Example output: ```json @@ -121,10 +126,10 @@ not being tracked by the pool. ## Definition of Labels -| Label | Meaning | -| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| jvm_heap_used | JVM heap memory usage of live objects for the executor process | -| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | -| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate to within 64 KiB of un-flushed delta per live thread. | -| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | -| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | +| Label | Meaning | +| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| jvm_heap_used | JVM heap memory usage of live objects for the executor process | +| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | +| native_allocated | Bytes handed out by the Rust global allocator, process-wide (`alloc-accounting` feature, on by default). Approximate to within 64 KiB of un-flushed delta per live thread. | +| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | +| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 25b805d1ccb..cca53645470 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -114,6 +114,28 @@ when there is sufficient memory in order to leave enough memory for other operat The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not need to spill or have a single spillable operator. +Both pools reserve memory against Spark's ledger, which only counts what operators explicitly asked for. Native memory +that never went through the pool, such as scratch buffers inside kernels or intermediate Arrow arrays, stays invisible +until the executor exceeds its container limit and is killed. To make it visible, both pools also compare the memory +the native allocator has actually handed out against `spark.memory.offHeap.size`. + +By default a crossing is only logged, once per task. Setting `spark.comet.exec.memoryPool.enforceNativeUsage` to +`true` makes the pools refuse the reservation instead, so operators that can spill do so and those that cannot fail +the task rather than the executor. Enforcement is off by default while the rate of false positives on real workloads +is established, and because spilling releases only reserved bytes: a denial provoked by untracked allocations may not +relieve the pressure it reports. + +Two things about the comparison are worth knowing. The budget is the whole off-heap size, not the reservable portion: +`spark.comet.exec.memoryPool.fraction` bounds what Comet may reserve, and lowering it to provoke spilling +deliberately does not lower this ceiling too. The two sides also do not measure the same population, because +`spark.memory.offHeap.size` is shared with Spark's own off-heap allocations and with Comet's JVM-side shuffle pages +while the measured usage counts only Comet's native allocations. It is a loose backstop against losing the executor, +not a bound on total off-heap usage. + +The comparison is process-wide on both sides, so there is no per-task attribution: once any task pushes real usage +past the budget, every task's next reservation sees it. It gates reservations only; allocations themselves are never +refused. + [shuffle]: #shuffle [Advanced Memory Tuning]: #advanced-memory-tuning diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 501592945d2..91c300af157 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -105,7 +105,7 @@ datafusion-functions-nested = { version = "55.1.0" } [features] backtrace = ["datafusion/backtrace"] -default = ["hdfs-opendal"] +default = ["hdfs-opendal", "alloc-accounting"] contrib-lance = ["dep:comet-contrib-lance"] hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] @@ -114,10 +114,11 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] # Default builds carry zero Delta surface. contrib-delta = ["dep:comet-contrib-delta"] -# Observability for real native memory usage. Wraps the global allocator to track the bytes it -# hands out, and reports the total as the `native_allocated` tracing metric so it can be compared -# against the memory pool's reservations. Never rejects an allocation. Off by default; a build -# without it has no wrapper and no per-allocation work. +# Real native memory usage. Wraps the global allocator to track the bytes it hands out, reports +# the total as the `native_allocated` tracing metric, and backs the off-heap memory pools' check +# of real usage against Comet's budget. Never rejects an allocation. On by default; build with +# `--no-default-features` (re-adding the other defaults) to drop the wrapper and its +# per-allocation work, which also leaves the pools' check reporting zero bytes in use. alloc-accounting = [] # exclude optional packages from cargo machete verifications diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index 394bbe774d3..12b9f58c3a4 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -181,10 +181,7 @@ unsafe impl GlobalAlloc for AccountingAllocator { } #[cfg(test)] -mod tests { - use super::*; - use std::alloc::System; - use std::sync::atomic::AtomicUsize; +pub(crate) mod test_support { use std::sync::{Mutex, MutexGuard}; /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it @@ -192,13 +189,24 @@ mod tests { /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by /// making each window microseconds wide and each expected move far larger than anything else /// allocates in that time. + /// + /// Lives outside the tests module because the memory pool's own gate test moves the balance + /// the same way and has to share the lock. static SERIAL: Mutex<()> = Mutex::new(()); - fn serial() -> MutexGuard<'static, ()> { + pub(crate) fn serial() -> MutexGuard<'static, ()> { SERIAL .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } +} + +#[cfg(test)] +mod tests { + use super::test_support::serial; + use super::*; + use std::alloc::System; + use std::sync::atomic::AtomicUsize; const MIB: usize = 1024 * 1024; diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 0107847c9e0..da06bf5aee5 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -112,8 +112,9 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED, - COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_MEMORY_POOL_ENFORCE_NATIVE_USAGE, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + SPARK_MEMORY_OFFHEAP_SIZE, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use datafusion_comet_proto::spark_operator::operator::OpStruct; @@ -495,6 +496,15 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( let max_temp_directory_size = spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024); let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY); + // Defaults to true, so it has to be read with an explicit default: the config map only + // carries values Spark actually holds, and `CometExecIterator` injects this one for that + // reason. + // Defaults to false: the pools observe and log a crossing but do not act on it yet. + let enforce_native_usage = + spark_config.get_bool_with_default(COMET_MEMORY_POOL_ENFORCE_NATIVE_USAGE, false); + // Injected by `CometExecIterator` because it is a Spark config rather than a Comet one. + // Absent (0) in on-heap mode, which leaves the check off. + let off_heap_size = spark_config.get_usize(SPARK_MEMORY_OFFHEAP_SIZE, 0); with_trace("createPlan", tracing_enabled, || { // Init JVM classes @@ -527,6 +537,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( memory_pool_type, memory_limit, memory_limit_per_task, + enforce_native_usage, + off_heap_size, )?; let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); diff --git a/native/core/src/execution/memory_pools/checked_pool.rs b/native/core/src/execution/memory_pools/checked_pool.rs new file mode 100644 index 00000000000..9ae814d5100 --- /dev/null +++ b/native/core/src/execution/memory_pools/checked_pool.rs @@ -0,0 +1,310 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{Debug, Display, Formatter, Result as FmtResult}; +use std::sync::atomic::{AtomicBool, Ordering}; + +use datafusion::{ + common::{resources_datafusion_err, DataFusionError}, + execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}, +}; +use log::warn; + +use crate::alloc_accounting; + +/// Wraps an off-heap memory pool so that the bytes the native allocator has actually handed out +/// are compared against a budget, and not just the bytes operators declared. +/// +/// A pool on its own only counts what operators voluntarily reserve, so native memory that +/// bypasses it is invisible until the executor exceeds its container limit and is killed. This +/// wrapper consults the `alloc-accounting` balance before the inner pool's own check, which makes +/// that overrun visible at the next reservation instead. +/// +/// `enforce` decides what a crossing does. By default it is false and the crossing is only logged, +/// once per pool, because the rate of false positives on real workloads is not yet established and +/// a spurious denial is worse than a late one: spilling releases *reserved* bytes, so if the +/// overshoot is in untracked allocations a denial may not relieve it and the task fails anyway. +/// With `spark.comet.exec.memoryPool.enforceNativeUsage` the reservation is refused, operators +/// that can spill do so, and those that cannot fail the task rather than the executor. +/// +/// Both `fair_unified` and `greedy_unified` wear this. The budget is `spark.memory.offHeap.size`, +/// deliberately not the pool's own limit: that limit is the off-heap size times +/// `spark.comet.exec.memoryPool.fraction`, and the fraction is how operators hold back reservable +/// memory to force spilling. Deriving this budget from it too would turn a small fraction into +/// denied reservations rather than the spills it was set to cause. +/// +/// Note that the budget and the balance are not measuring the same population: +/// `spark.memory.offHeap.size` is shared with Spark's own Tungsten off-heap allocations and with +/// Comet's JVM-side shuffle pages, while the balance counts only Comet's Rust allocations. The +/// comparison is therefore a loose backstop against the executor being killed, not a bound on +/// total off-heap usage. +/// +/// Both the balance and the budget are process-wide, so there is no per-task attribution: once any +/// task pushes real usage past the budget, every task's next non-zero reservation sees it. +/// Allocations themselves are never refused; this is a reservation gate, not a hard limit. +/// +/// A build without the `alloc-accounting` feature reports a balance of zero, which leaves the +/// wrapper a passthrough. +pub struct CheckedMemoryPool { + inner: P, + budget: usize, + enforce: bool, + /// Set once this pool has logged, whether that was an observed crossing or a refusal, so that + /// neither mode floods the log: `try_grow` is called constantly and the condition is sticky + /// once real usage is high. One line per pool means one line per task that hit the budget. + reported: AtomicBool, +} + +impl CheckedMemoryPool

{ + pub fn new(inner: P, budget: usize, enforce: bool) -> Self { + Self { + inner, + budget, + enforce, + reported: AtomicBool::new(false), + } + } +} + +impl Debug for CheckedMemoryPool

{ + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.debug_struct("CheckedMemoryPool") + .field("budget", &self.budget) + .field("enforce", &self.enforce) + .field("inner", &self.inner) + .finish() + } +} + +impl Display for CheckedMemoryPool

{ + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!( + f, + "CheckedMemoryPool(budget={}, enforce={}, inner={})", + self.budget, self.enforce, self.inner + ) + } +} + +impl MemoryPool for CheckedMemoryPool

{ + fn name(&self) -> &str { + "CheckedMemoryPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer) + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer) + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.try_grow(reservation, additional).unwrap() + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<(), DataFusionError> { + if additional == 0 { + return Ok(()); + } + // Checked first because it is a single atomic load, whereas the inner unified pools cross + // JNI to ask Spark. A request denied here never reaches Spark's ledger. + let in_use = alloc_accounting::current_balance(); + if in_use.saturating_add(additional) > self.budget { + if self.enforce { + // Log the first refusal too, not just the first observed crossing. A refusal is + // returned as an error, and DataFusion usually answers it by spilling and + // retrying, which swallows the error: without this line the check can be + // refusing reservations on every task and leave no trace anywhere that it ran. + if !self.reported.swap(true, Ordering::Relaxed) { + warn!( + "Refusing to reserve {additional} bytes for {}: it would take Comet's \ + real native memory usage ({in_use} bytes) past \ + spark.memory.offHeap.size ({} bytes). Operators that can spill will \ + spill; those that cannot will fail the task. Set \ + spark.comet.exec.memoryPool.enforceNativeUsage=false to only log this.", + reservation.consumer().name(), + self.budget + ); + } + return Err(resources_datafusion_err!( + "Failed to reserve {additional} bytes for {}: native memory in use is \ + {in_use} bytes of a {} byte budget (spark.memory.offHeap.size). Reserved: \ + {}. Raise spark.memory.offHeap.size, or disable this check with \ + spark.comet.exec.memoryPool.enforceNativeUsage=false", + reservation.consumer().name(), + self.budget, + self.reserved() + )); + } + if !self.reported.swap(true, Ordering::Relaxed) { + // Both quantities are named because either can be what crosses the budget: a + // large single request against modest usage, or a small request against usage + // that is already near the ceiling. + warn!( + "Reserving {additional} bytes for {} would take Comet's real native memory \ + usage ({in_use} bytes) past spark.memory.offHeap.size ({} bytes). This \ + memory is not covered by the pool's reservations and the executor may be \ + killed for exceeding its container limit. Set \ + spark.comet.exec.memoryPool.enforceNativeUsage=true to refuse such \ + reservations instead, or raise spark.memory.offHeap.size.", + reservation.consumer().name(), + self.budget + ); + } + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + // Always the inner pool's limit, never the budget, in both modes. The budget is a + // process-wide backstop compared against process-wide allocator usage, not this pool's + // reservable limit, so reporting it here would be a lie about the pool. DataFusion also + // acts on this value: `AggregateExec::should_use_partial_reduce_hash_stream` bails out + // whenever the pool reports `Finite`, so returning the budget would silently change the + // aggregation strategy as a side effect of enabling the check, which is not something a + // memory guard should do. + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::UnboundedMemoryPool; + use std::sync::Arc; + + fn enforcing(budget: usize) -> Arc { + Arc::new(CheckedMemoryPool::new( + UnboundedMemoryPool::default(), + budget, + true, + )) + } + + fn observing(budget: usize) -> Arc { + Arc::new(CheckedMemoryPool::new( + UnboundedMemoryPool::default(), + budget, + false, + )) + } + + #[test] + fn a_zero_byte_grow_never_fails() { + let pool = enforcing(0); + let reservation = MemoryConsumer::new("zero").register(&pool); + reservation.try_grow(0).unwrap(); + } + + /// Never `Finite(budget)`, in either mode. DataFusion changes its aggregation strategy when a + /// pool reports a finite limit, so reporting the backstop here would make enabling the check + /// alter execution even when it never refuses anything. + #[test] + fn always_reports_the_inner_limit_whether_enforcing_or_not() { + assert!(matches!( + enforcing(4096).memory_limit(), + MemoryLimit::Infinite + )); + assert!(matches!( + observing(4096).memory_limit(), + MemoryLimit::Infinite + )); + } + + #[test] + fn successful_grows_and_shrinks_reach_the_inner_pool() { + let pool = enforcing(usize::MAX); + let reservation = MemoryConsumer::new("delegate").register(&pool); + reservation.try_grow(1024).unwrap(); + assert_eq!(pool.reserved(), 1024); + reservation.shrink(1024); + assert_eq!(pool.reserved(), 0); + } + + /// The gate compares real bytes, not reservations: a block this pool never heard about is + /// enough to deny a one-byte request, and freeing it is enough to allow the same request. + /// + /// The block is touched so it is really allocated, and the margins are far wider than anything + /// the rest of the crate allocates in the microseconds between the checks. The serial lock + /// keeps the accounting tests that move the balance by tens of megabytes out of that window. + #[test] + #[cfg(feature = "alloc-accounting")] + fn denies_when_real_bytes_plus_request_exceed_the_budget() { + use std::hint::black_box; + + const HEADROOM: usize = 64 * 1024 * 1024; + const BLOCK: usize = 256 * 1024 * 1024; + + let _guard = alloc_accounting::test_support::serial(); + let budget = alloc_accounting::current_balance() + HEADROOM; + let pool = enforcing(budget); + let reservation = MemoryConsumer::new("checked").register(&pool); + + let held: Vec = black_box(vec![1u8; BLOCK]); + let denied = reservation.try_grow(1).unwrap_err(); + black_box(&held); + assert!( + matches!(denied, DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted, got {denied:?}" + ); + let message = denied.to_string(); + assert!( + message.contains("native memory in use is") && message.contains("checked"), + "message should name the real bytes in use and the consumer: {message}" + ); + assert_eq!(pool.reserved(), 0, "a denied request must not be reserved"); + + drop(held); + reservation.try_grow(1).unwrap(); + assert_eq!(pool.reserved(), 1); + } + + /// The same crossing, with enforcement off: the reservation still succeeds and reaches the + /// inner pool, which is what makes the default safe to ship. + #[test] + #[cfg(feature = "alloc-accounting")] + fn allows_the_same_crossing_when_only_observing() { + use std::hint::black_box; + + const HEADROOM: usize = 64 * 1024 * 1024; + const BLOCK: usize = 256 * 1024 * 1024; + + let _guard = alloc_accounting::test_support::serial(); + let budget = alloc_accounting::current_balance() + HEADROOM; + let pool = observing(budget); + let reservation = MemoryConsumer::new("observed").register(&pool); + + let held: Vec = black_box(vec![1u8; BLOCK]); + reservation.try_grow(1).unwrap(); + black_box(&held); + assert_eq!(pool.reserved(), 1, "observing must not withhold the bytes"); + drop(held); + } +} diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index 312a3604383..35790282852 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -33,6 +33,14 @@ pub(crate) enum MemoryPoolType { pub(crate) struct MemoryPoolConfig { pub(crate) pool_type: MemoryPoolType, pub(crate) pool_size: usize, + /// Budget the off-heap pools compare real native usage against, or `None` when there is + /// nothing to compare to. Always `None` in on-heap mode, where Comet has no off-heap + /// allotment. Set whether or not the crossing is enforced, because observing it is the + /// default behaviour. + pub(crate) native_usage_budget: Option, + /// Whether crossing `native_usage_budget` refuses the reservation. When false the pools only + /// log the crossing. + pub(crate) enforce_native_usage: bool, } impl MemoryPoolConfig { @@ -40,8 +48,16 @@ impl MemoryPoolConfig { Self { pool_type, pool_size, + native_usage_budget: None, + enforce_native_usage: false, } } + + fn with_native_usage_budget(mut self, budget: Option, enforce: bool) -> Self { + self.native_usage_budget = budget; + self.enforce_native_usage = enforce; + self + } } pub(crate) fn parse_memory_pool_config( @@ -49,16 +65,26 @@ pub(crate) fn parse_memory_pool_config( memory_pool_type: String, memory_limit: i64, memory_limit_per_task: i64, + enforce_native_usage: bool, + off_heap_size: usize, ) -> CometResult { let pool_size = memory_limit as usize; let memory_pool_config = if off_heap_mode { + // Deliberately the whole off-heap size rather than `pool_size`, which is that size times + // `spark.comet.exec.memoryPool.fraction`. The fraction bounds what Comet may *reserve*, + // and lowering it is how operators force spilling; reusing it here would also lower the + // ceiling on *real* usage, so a small fraction would deny reservations outright instead + // of provoking the spills it was set to cause. + let native_usage_budget = (off_heap_size > 0).then_some(off_heap_size); match memory_pool_type.as_str() { - "fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size), + "fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size) + .with_native_usage_budget(native_usage_budget, enforce_native_usage), "greedy_unified" => { // the `unified` memory pool interacts with Spark's memory pool to allocate // memory therefore does not need a size to be explicitly set. The pool size // shared with Spark is set by `spark.memory.offHeap.size`. MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) + .with_native_usage_budget(native_usage_budget, enforce_native_usage) } _ => { return Err(CometError::Config(format!( diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index d7c2911f913..eaf8e85aa89 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -15,12 +15,14 @@ // specific language governing permissions and limitations // under the License. +mod checked_pool; mod config; mod fair_pool; pub mod logging_pool; mod task_shared; mod unified_pool; +use checked_pool::CheckedMemoryPool; use datafusion::execution::memory_pool::{ FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool, }; @@ -52,21 +54,40 @@ pub(crate) fn create_memory_pool( )) } + /// Wraps an off-heap pool in the real-native-usage check when there is a budget to compare + /// against, so that Comet's actual allocations are measured rather than only its declared + /// reservations. Whether a crossing is refused or merely logged is `enforce`. `tracked` stays + /// outermost so a denial is still annotated with the largest consumers. + fn checked( + pool: impl MemoryPool + 'static, + budget: Option, + enforce: bool, + ) -> Arc { + match budget { + Some(budget) => tracked(CheckedMemoryPool::new(pool, budget, enforce)), + None => tracked(pool), + } + } + let pool_type = memory_pool_config.pool_type; let pool_size = memory_pool_config.pool_size; + let native_usage_budget = memory_pool_config.native_usage_budget; + let enforce_native_usage = memory_pool_config.enforce_native_usage; match pool_type { MemoryPoolType::GreedyUnified => acquire_task_shared_pool(task_attempt_id, || { - tracked(CometUnifiedMemoryPool::new( - comet_task_memory_manager, - task_attempt_id, - )) + checked( + CometUnifiedMemoryPool::new(comet_task_memory_manager, task_attempt_id), + native_usage_budget, + enforce_native_usage, + ) }), MemoryPoolType::FairUnified => acquire_task_shared_pool(task_attempt_id, || { - tracked(CometFairMemoryPool::new( - comet_task_memory_manager, - pool_size, - )) + checked( + CometFairMemoryPool::new(comet_task_memory_manager, pool_size), + native_usage_budget, + enforce_native_usage, + ) }), MemoryPoolType::GreedyTaskShared => acquire_task_shared_pool(task_attempt_id, || { tracked(GreedyMemoryPool::new(pool_size)) diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5de..c0e81c4a6d4 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -24,10 +24,14 @@ pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirec pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; +pub(crate) const COMET_MEMORY_POOL_ENFORCE_NATIVE_USAGE: &str = + "spark.comet.exec.memoryPool.enforceNativeUsage"; +pub(crate) const SPARK_MEMORY_OFFHEAP_SIZE: &str = "spark.memory.offHeap.size"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; + fn get_bool_with_default(&self, name: &str, default_value: bool) -> bool; fn get_u64(&self, name: &str, default_value: u64) -> u64; fn get_usize(&self, name: &str, default_value: usize) -> usize; } @@ -39,6 +43,12 @@ impl SparkConfig for HashMap { .unwrap_or(false) } + fn get_bool_with_default(&self, name: &str, default_value: bool) -> bool { + self.get(name) + .and_then(|str_val| str_val.parse::().ok()) + .unwrap_or(default_value) + } + fn get_u64(&self, name: &str, default_value: u64) -> u64 { self.get(name) .and_then(|str_val| str_val.parse::().ok()) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index fc858a75a18..2a059a3e688 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -850,6 +850,22 @@ object CometConf extends ShimCometConf { .doubleConf .createWithDefault(1.0) + val COMET_OFFHEAP_MEMORY_POOL_ENFORCE_NATIVE_USAGE: ConfigEntry[Boolean] = + conf("spark.comet.exec.memoryPool.enforceNativeUsage") + .category(CATEGORY_TUNING) + .doc( + "Comet's off-heap memory pools compare the memory the native allocator has actually " + + "handed out against `spark.memory.offHeap.size`, which bounds native memory that " + + "operators never reserved and that the pools cannot otherwise see. By default a " + + "crossing is only logged, once per task. When this is enabled the pools refuse the " + + "reservation instead, so operators that can spill do so and those that cannot fail " + + "the task rather than the executor. Enforcement is off by default while the rate of " + + "false positives on real workloads is still being established. " + + "Only applies to off-heap mode. " + + s"$TUNING_GUIDE.") + .booleanConf + .createWithDefault(false) + val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired") .category(CATEGORY_EXEC) .doc( diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e2c132904d5..329610976a2 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -357,6 +357,19 @@ object CometExecIterator extends Logging { builder.putEntries( CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_OFFHEAP_MEMORY_POOL_ENFORCE_NATIVE_USAGE.key, + CometConf.COMET_OFFHEAP_MEMORY_POOL_ENFORCE_NATIVE_USAGE.get(SQLConf.get).toString) + // The off-heap pools check real native usage against the whole off-heap size, which is a + // Spark config rather than a Comet one and so is not carried by `cometSqlConfs`. Deliberately + // not the memory limit the pools already receive: that has the pool fraction applied, and the + // fraction bounds reservations rather than real usage. + val sparkConf = SparkEnv.get.conf + if (CometSparkSessionExtensions.isOffHeapEnabled(sparkConf)) { + builder.putEntries( + "spark.memory.offHeap.size", + ByteUnit.MiB.toBytes(sparkConf.getSizeAsMb("spark.memory.offHeap.size")).toString) + } builder.build().toByteArray } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometMemoryPoolNativeUsageSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometMemoryPoolNativeUsageSuite.scala new file mode 100644 index 00000000000..06b2faeb35d --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometMemoryPoolNativeUsageSuite.scala @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import org.apache.spark.SparkConf +import org.apache.spark.sql.CometTestBase + +import org.apache.comet.CometConf + +/** + * Exercises the off-heap memory pools' check of real native memory usage. + * + * The budget is `spark.memory.offHeap.size`, which Spark fixes when the session starts, so these + * tests need a session of their own rather than a `withSQLConf` override. The size below is + * smaller than the memory Comet's native code already holds before a query runs, so every + * reservation crosses the budget and the first one is refused once enforcement is turned on. + * Deliberately not achieved by lowering `spark.comet.exec.memoryPool.fraction`: that bounds what + * Comet may reserve, not what the check measures. + */ +class CometMemoryPoolNativeUsageSuite extends CometTestBase { + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.memory.offHeap.size", "2m") + conf + } + + /** A sort, so that an operator actually reserves. */ + private def sortSmallInput(): Unit = + spark.range(0, 1000).selectExpr("id", "id % 7 AS m").sort("m", "id").collect() + + private def failureMessages(run: => Unit): Seq[String] = + causeChain(intercept[Throwable](run)).map(t => s"${t.getClass.getName}: ${t.getMessage}") + + test("off-heap pools deny reservations once real native usage exceeds the off-heap size") { + Seq("fair_unified", "greedy_unified").foreach { poolType => + withSQLConf( + CometConf.COMET_OFFHEAP_MEMORY_POOL_ENFORCE_NATIVE_USAGE.key -> "true", + CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key -> poolType) { + val messages = failureMessages(sortSmallInput()) + assert( + messages.exists(_.contains("native memory in use is")), + s"expected $poolType to deny the reservation on real native usage, but got:\n " + + messages.mkString("\n ")) + } + } + } + + test("the check only observes by default") { + // The same query under the same off-heap size, so enforcement is the only difference. It is + // too small for the sort either way, which is what makes the budget bite in the test above; + // what changes here is who refuses the reservation. Left at its default the check only logs, + // so the reservation is not held back in Comet and reaches Spark's ledger, which fails it. + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key -> "greedy_unified") { + val messages = failureMessages(sortSmallInput()) + assert( + !messages.exists(_.contains("native memory in use is")), + "the check is not enforcing by default but still refused the reservation:\n " + + messages.mkString("\n ")) + assert( + messages.exists(_.contains("failed to acquire")), + "expected the reservation to reach Spark's ledger, but got:\n " + + messages.mkString("\n ")) + } + } +}