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] =