comphead opened a new pull request, #7304: URL: https://github.com/apache/datafusion-comet/pull/7304
## Which issue does this PR close? None. Related to #7108, which stays open: it asks whether release builds should default to a different allocator, and this PR adds an opt-in mitigation for the default glibc build instead. ## Rationale for this change An executor can be OOMKilled by its container while Comet's memory pools and allocation counter stay well within their limits. Much of the difference is native memory that has already been freed: glibc keeps a large part of what Comet's operators free in its per-thread arenas instead of returning it to the operating system. #7108 measured this on Linux. Live native memory dropped to about 2 MiB between queries while RSS stayed near 12 GiB, and one `malloc_trim` released 5 to 6 GiB. It gets worse between stages: - Comet's operators free a lot of memory. Each sort spill and each shuffle writer spill drops the batches it buffered, which are Arrow buffers of tens of KiB to a few MiB. glibc serves those from its arenas rather than with `mmap`. - One stage's native plans run on different threads than the next stage's. A plan with no JVM input runs on Tokio workers (`BatchProducer::spawn`), and a plan that reads shuffle runs on the Spark task thread. The next stage does not reuse what the previous one left in the other threads' arenas, so its allocations add to the retained memory. Spark's own off-heap memory does not build up this way, because its pages are 64 MiB and glibc maps and unmaps each of them directly. I saw this pattern on a wide group-by (87 sort keys, sort-based aggregation) that Spark finishes and Comet does not: every result-stage wave was OOMKilled. A local reproduction at the same per-task size shows the mechanism. Comet's result stage alone needs about what Spark's does (9.1 vs 8.9 GiB), but right after the map stage it reaches 16.2 GiB, because about 7 GiB that the map stage freed is still resident when it starts. ## What changes are included in this PR? - A new setting, `spark.comet.memory.trimNativeHeapWhenIdle.enabled`, `false` by default. When it is set, releasing the last native plan running on an executor calls glibc's `malloc_trim(0)` and logs at INFO level how long that took and the process's anonymous resident memory before and after. It only runs when no plan is left, which is typically between stages, so it does not take the arena locks of running operators. - It only acts on Linux builds that use glibc and the default system allocator. On other platforms, and in builds with the `jemalloc` or `mimalloc` feature, it does nothing. - It is an executor setting. `CometExecIterator.serializeCometSQLConfs` reads it from the executor's `SparkConf` rather than the session's, so it applies however the job is launched (see #7301). An invalid value logs a warning once and leaves trimming off. - The code is a new `native_heap_trim.rs` module and a hook at the end of `release_plan` in `jni_api.rs`, with a small `live_plan_count()` helper. The memory tuning guide gets a section on the setting. ## How are these changes tested? - Unit tests: `native_heap_trim::tests` checks that it trims only when the build supports it and no plan is left. A `CometExecIteratorLifecycleSuite` test covers parsing the setting. `CometExecSuite` "SQLConf serde resolves the configs that native code parses" checks that the value always crosses JNI resolved, from the executor's configuration. These have not run locally yet. - `malloc_trim` does not exist on macOS, where I work, so I type-checked and linted the Linux path for `aarch64-unknown-linux-gnu` (clippy with `-D warnings`). It has not run on Linux yet. A Linux run with the setting on is the next step, and the INFO line will show its effect. - To exercise the hook itself, I ran the reproduction with a jemalloc build in which the hook purged jemalloc's arenas instead of calling `malloc_trim` (a test-only change, not in this PR), and with jemalloc keeping freed memory as glibc does. With the setting off, 4 of 4 runs exceeded the memory limit in the result stage. With it on, 4 of 4 passed. The hook ran twice per job, once after the map stage's last task and once at the end, and each purge took 0.3 to 0.7 s. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
