From 92fc1c415e79dcab660e4d205adfff67a3697a7f Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 28 Oct 2026 10:15:50 -0700 Subject: [PATCH] feat: trim the native heap when an executor's last native plan finishes glibc keeps much of the memory that Comet's native operators free in its per-thread arenas instead of returning it to the operating system. Spills free a lot of memory, and one stage's native plans run on different threads than the next stage's, so an executor's resident memory can stay near one stage's peak while the next stage allocates on top of it. The container can then run out of memory while Comet's memory pools stay within their limits. Add spark.comet.memory.trimNativeHeapWhenIdle.enabled, off by default. When it is set, releasing the last native plan running on an executor calls glibc's malloc_trim(0) and logs how long that took and the anonymous resident memory before and after. It only acts on Linux builds that use glibc and the default system allocator, and it is read from the executor's configuration. --- .../source/user-guide/latest/tuning/memory.md | 14 +++ native/core/src/execution/jni_api.rs | 28 ++++- native/core/src/execution/mod.rs | 1 + native/core/src/execution/native_heap_trim.rs | 111 ++++++++++++++++++ native/core/src/execution/spark_config.rs | 2 + .../scala/org/apache/comet/CometConf.scala | 20 ++++ .../org/apache/comet/CometExecIterator.scala | 33 ++++++ .../apache/comet/exec/CometExecSuite.scala | 8 ++ .../CometExecIteratorLifecycleSuite.scala | 11 ++ 9 files changed, 226 insertions(+), 2 deletions(-) create mode 100644 native/core/src/execution/native_heap_trim.rs diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 2dd730e9365..732f52e88e7 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -273,6 +273,20 @@ these events, and one without Comet on its classpath logs once that it dropped t `spark.eventLog.excludedPatterns=org.apache.comet.CometExecutorMemoryUsage` leaves them out of the event log. +## Returning Freed Native Memory Between Stages + +On Linux, glibc's allocator keeps much of the memory that Comet's native operators free instead of +returning it to the operating system, and it keeps it per thread. Spills free a lot of memory, and +one stage's native plans run on different threads than the next stage's, so an executor's resident +memory can stay near the peak of one stage while the next stage allocates on top of it. The executor +can then outgrow its container while the memory usage log shows little allocated. + +Setting `spark.comet.memory.trimNativeHeapWhenIdle.enabled=true` makes an executor call glibc's +`malloc_trim` when its last running native plan finishes, which is typically between stages, and log +at INFO level how long that took and its resident memory before and after. Set it when the +application is submitted. It does nothing on other platforms, or in builds with the `jemalloc` or +`mimalloc` feature. + ## Batch Size Comet processes data in columnar batches. The batch size is controlled by `spark.comet.batchSize` (default diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 494e7acdff9..13c1499c6f5 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -125,6 +125,7 @@ use tokio::task::JoinHandle; use crate::execution::memory_pools::{ create_memory_pool, overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, }; +use crate::execution::native_heap_trim; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, ShuffleWriterExec, @@ -141,8 +142,9 @@ use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE, COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_exec::CometParquetScanOptions; @@ -284,6 +286,15 @@ fn sum_reserved(pools: &[Arc]) -> usize { pools.iter().map(|pool| pool.reserved()).sum() } +/// The number of native plans registered on this executor, that is created and not yet released. +fn live_plan_count() -> usize { + get_thread_memory_pools() + .lock() + .values() + .map(HashMap::len) + .sum() +} + /// Bytes reserved across `pools`, less the part that Spark has not granted them; see /// [`MemoryUsage::pools_reserved`]. fn sum_reserved_less_overcommit(pools: &[Arc]) -> usize { @@ -597,6 +608,9 @@ struct ExecutionContext { pub output_stream: Option>, /// Removes this context's tracing memory-pool entry on every exit path. memory_pool_registration: Option, + /// Whether releasing this plan trims the native heap when no other plan is left running on + /// the executor; see [`crate::execution::native_heap_trim`]. + trim_native_heap_when_idle: bool, } /// Accept serialized query plan and return the address of the native query plan. @@ -728,6 +742,7 @@ struct PlanSettings { tracing_enabled: bool, max_temp_directory_size: u64, logging_memory_pool: bool, + trim_native_heap_when_idle: bool, } /// Deserializes the Spark configs of a plan and initializes the Tokio runtime from them. The @@ -750,6 +765,7 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult bool { + supported && live_plans == 0 +} + +#[cfg(all(target_os = "linux", target_env = "gnu"))] +mod platform { + use log::info; + use std::time::Instant; + + pub(super) const SUPPORTED: bool = true; + + pub(super) fn trim_and_log() { + extern "C" { + fn malloc_trim(pad: usize) -> std::os::raw::c_int; + } + let before = anonymous_rss_mib(); + let start = Instant::now(); + // SAFETY: glibc's malloc_trim takes the lock of each arena it trims and can be called from + // any thread. A pad of 0 keeps nothing beyond what the allocator needs at the top of the + // main heap. + unsafe { + malloc_trim(0); + } + let elapsed_ms = start.elapsed().as_millis(); + match (before, anonymous_rss_mib()) { + (Some(before), Some(after)) => info!( + "Trimmed the native heap after the last Comet native plan on this executor \ + finished, in {elapsed_ms} ms: anonymous resident memory {before:.1} MiB before, \ + {after:.1} MiB after" + ), + _ => info!( + "Trimmed the native heap after the last Comet native plan on this executor \ + finished, in {elapsed_ms} ms" + ), + } + } + + /// The process's anonymous resident memory, which the container's memory limit counts, or + /// None when `/proc` cannot be read. + fn anonymous_rss_mib() -> Option { + let status = procfs::process::Process::myself().ok()?.status().ok()?; + status.rssanon.map(|kib| kib as f64 / 1024.0) + } +} + +#[cfg(not(all(target_os = "linux", target_env = "gnu")))] +mod platform { + pub(super) const SUPPORTED: bool = false; + + pub(super) fn trim_and_log() {} +} + +#[cfg(test)] +mod tests { + use super::should_trim; + + #[test] + fn trims_only_when_supported_and_no_plan_is_left() { + assert!(should_trim(true, 0)); + assert!(!should_trim(true, 1), "another plan is still running"); + assert!(!should_trim(false, 0), "the build cannot trim"); + } +} diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 511b5733388..745187e0103 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -22,6 +22,8 @@ pub(crate) const COMET_DEBUG_ENABLED: &str = "spark.comet.debug.enabled"; pub(crate) const COMET_EXPLAIN_NATIVE_ENABLED: &str = "spark.comet.explain.native.enabled"; pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirectorySize"; pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; +pub(crate) const COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE: &str = + "spark.comet.memory.trimNativeHeapWhenIdle.enabled"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED: &str = diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..3ad0c9bd679 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -583,6 +583,26 @@ object CometConf extends ShimCometConf { .checkValue(_ >= 0, "The memory usage log interval must not be negative") .createWithDefault(TimeUnit.SECONDS.toMillis(10)) + val COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE: ConfigEntry[Boolean] = + conf("spark.comet.memory.trimNativeHeapWhenIdle.enabled") + .category(CATEGORY_TUNING) + .doc( + "Whether an executor hands the memory that the C allocator keeps cached back to the " + + "operating system when its last running Comet native plan finishes, which is " + + "typically between stages. glibc keeps much of the memory that native operators free " + + "instead of returning it, so the executor's resident memory can stay near the peak " + + "of one stage while the next stage allocates on top of it, and the container can run " + + "out of memory while Comet's memory pools stay within their limits. With this " + + "enabled, Comet calls glibc's malloc_trim when no native plan is left running on the " + + "executor, and logs at INFO level how long it took and the resident memory before " + + "and after. Only Linux builds that use glibc and the default system allocator trim. " + + "Elsewhere, and in builds with the jemalloc or mimalloc feature, this does nothing. " + + "This is an executor setting, read from the application's configuration, so it must " + + "be set when the application is submitted. An invalid value disables trimming with " + + s"a warning. $TUNING_GUIDE.") + .booleanConf + .createWithDefault(false) + val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.enabled") diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index d0db9d9dfcf..728685ed62f 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -553,6 +553,32 @@ object CometExecIterator extends Logging { } } + private val invalidTrimSettingWarned = new AtomicBoolean(false) + + /** + * Whether the executor trims its native heap when its last Comet native plan finishes, for the + * executor's configured value, if any. It is read for every native plan, so a value that does + * not parse disables trimming with a warning, logged once, rather than failing every Comet + * task. + */ + def trimNativeHeapWhenIdle(configured: Option[String]): Boolean = + configured match { + case None => COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.defaultValue.get + case Some(value) => + try { + COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.valueConverter(value) + } catch { + case NonFatal(e) => + if (invalidTrimSettingWarned.compareAndSet(false, true)) { + logWarning( + s"Not trimming the native heap: invalid value '$value' for " + + s"${COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key}. Expected true or false. " + + e.getMessage) + } + false + } + } + /** * Whether the cluster manager for `master` runs each executor in a container whose memory it * sizes from the memory overhead settings: YARN and Kubernetes. Local mode, local-cluster @@ -891,6 +917,13 @@ object CometExecIterator extends Logging { val executorCores = numDriverOrExecutorCores(SparkEnv.get.conf) builder.putEntries("spark.executor.cores", executorCores.toString) + // Trimming acts on the whole executor, so it is read from the executor's configuration rather + // than the session's, and overrides a value set in the session. + val trimKey = COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key + builder.putEntries( + trimKey, + trimNativeHeapWhenIdle(SparkEnv.get.conf.getOption(trimKey)).toString) + // Any Comet config that the native side reads must be added here manually, resolved. // `cometSqlConfs` only carries values that were explicitly set, exactly as they were // written, so defaults from `createWithDefault(...)` would otherwise not cross JNI, and diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index c1d03d6a40b..9242c1faa24 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -116,6 +116,14 @@ class CometExecSuite extends CometTestBase { withSQLConf(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key -> "TRUE") { assert(entries.get(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key) == "true") } + + // Trimming the native heap is an executor setting: it crosses resolved from the executor's + // configuration, and a value set in the session does not change it. + val trimKey = CometConf.COMET_MEMORY_TRIM_NATIVE_HEAP_WHEN_IDLE.key + assert(defaults.get(trimKey) == "false") + withSQLConf(trimKey -> "true") { + assert(entries.get(trimKey) == "false") + } } test("native sort spill files are compressed with the default spill codec") { diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index 1caaec887fa..a2003da74f2 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -845,6 +845,17 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { assert(memoryUsageLogInterval(Some("10 seconds please")) == 0L) } + test("trimming the native heap is off by default and on a value it cannot use") { + import CometExecIterator.trimNativeHeapWhenIdle + assert(!trimNativeHeapWhenIdle(None)) + assert(trimNativeHeapWhenIdle(Some("true"))) + assert(trimNativeHeapWhenIdle(Some(" TRUE "))) + assert(!trimNativeHeapWhenIdle(Some("false"))) + // Values that would otherwise throw from the creation of every native plan, failing its task. + assert(!trimNativeHeapWhenIdle(Some("yes"))) + assert(!trimNativeHeapWhenIdle(Some(""))) + } + test("the native runtime gets a worker thread for every core the executor runs tasks on") { import CometExecIterator.numDriverOrExecutorCores def cores(master: String, executorCores: Option[String] = None): Option[Int] =