feat: Add fair unified memory pool #1369

kazuyukitanimura · 2025-02-05T18:00:08Z

Which issue does this PR close?

Rationale for this change

Current Comet unified memory pool is a greedy pool. One thread (consumer) can take a large amount of memory that causes OOM for other threads, especially for aggregation.

What changes are included in this PR?

Added a fair version of unified memory pool similar to DataFusion FairSpilPool that caps the memory usage at pool_size/num
The fair unified memory pool is the default for off-heap mode with this PR

How are these changes tested?

Exisiting tests

codecov-commenter · 2025-02-05T19:10:48Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 39.06%. Comparing base (f09f8af) to head (05fd0f1).
Report is 22 commits behind head on main.

Additional details and impacted files

@@              Coverage Diff              @@
##               main    #1369       +/-   ##
=============================================
- Coverage     56.12%   39.06%   -17.06%     
- Complexity      976     2071     +1095     
=============================================
  Files           119      263      +144     
  Lines         11743    60746    +49003     
  Branches       2251    12909    +10658     
=============================================
+ Hits           6591    23733    +17142     
- Misses         4012    32530    +28518     
- Partials       1140     4483     +3343

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

kazuyukitanimura · 2025-02-05T21:43:45Z

@andygrove @comphead @himadripal

docs/source/user-guide/configs.md

common/src/main/scala/org/apache/comet/CometConf.scala

native/core/src/execution/fair_memory_pool.rs

native/core/src/execution/jni_api.rs

himadripal · 2025-02-06T19:10:39Z

docs/source/user-guide/configs.md

@@ -48,7 +48,7 @@ Comet provides the following configuration settings.
 | spark.comet.exec.hashJoin.enabled | Whether to enable hashJoin by default. | true |
 | spark.comet.exec.initCap.enabled | Whether to enable initCap by default. | false |
 | spark.comet.exec.localLimit.enabled | Whether to enable localLimit by default. | true |
-| spark.comet.exec.memoryPool | The type of memory pool to be used for Comet native execution. Available memory pool types are 'greedy', 'fair_spill', 'greedy_task_shared', 'fair_spill_task_shared', 'greedy_global' and 'fair_spill_global', By default, this config is 'greedy_task_shared'. | greedy_task_shared |
+| spark.comet.exec.memoryPool | The type of memory pool to be used for Comet native execution. Available memory pool types are 'greedy', 'fair_spill', 'greedy_task_shared', 'fair_spill_task_shared', 'greedy_global' and 'fair_spill_global'. For off-eap types are 'unified' and `fair_unified`. | greedy_task_shared |


add default is fair_unified

We tend to omit what the default value is. When default value is changed, the option description is often forgotten to be updated because there is no checking for the description. So we deleted default value explanations in the past

Co-authored-by: Zhen Wang <[email protected]>

native/core/src/execution/fair_memory_pool.rs

comphead · 2025-02-06T19:20:56Z

native/core/src/execution/fair_memory_pool.rs

+impl MemoryPool for CometFairMemoryPool {
+    fn register(&self, _: &MemoryConsumer) {
+        let mut state = self.state.lock();
+        state.num = state.num.checked_add(1).unwrap();


this unwraps can be super confusing if they happen

This won't happen in reality, num is usize. I wouldn't worry about it.

changed to expect() with a message

native/core/src/execution/fair_memory_pool.rs

comphead

Thanks @kazuyukitanimura nice job
unwraps probably could be wrapped into user friendly messages or its gonna super confusing to unwrap panic when it is supposed to be a memory or memory manager issue.

native/core/src/execution/fair_memory_pool.rs

Co-authored-by: Oleks V <[email protected]>

wForget · 2025-02-07T07:42:32Z

native/core/src/execution/fair_memory_pool.rs

+        state.num = state.num.checked_sub(1).unwrap();
+    }
+
+    fn grow(&self, reservation: &MemoryReservation, additional: usize) {


This method implementation seems to be different from memory_pool, but I'm not sure which one is reasonable.

https://github.com/apache/datafusion-comet/blob/main/native/core/src/execution/memory_pool.rs#L83-L87

kazuyukitanimura · 2025-02-07T22:15:44Z

@andygrove I would like to make sure to include this for the 0.6.0 release

andygrove · 2025-02-07T22:32:10Z

@andygrove I would like to make sure to include this for the 0.6.0 release

I added this PR to #1361

andygrove · 2025-02-08T19:09:22Z

I tried running TPC-H locally with this PR but there is excessive spilling and the benchmark is extremely slow. q3 took 3.4 minutes compared to 16 seconds on main branch.

kazuyukitanimura · 2025-02-10T18:49:54Z

Thanks @andygrove for trying. You may have to increase the memory because the fair memory pool limits memory usage earlier in order to leave the the rest for other threads.
Or I can keep the default to the original memory pool and make this fair pool optional

kazuyukitanimura · 2025-02-10T19:00:08Z

Went ahead and kept the default pool choice @andygrove

andygrove

Thanks @kazuyukitanimura

kazuyukitanimura · 2025-02-10T22:01:44Z

Merged, thanks @andygrove

## Which issue does this PR close? ## Rationale for this change Current Comet unified memory pool is a greedy pool. One thread (consumer) can take a large amount of memory that causes OOM for other threads, especially for aggregation. ## What changes are included in this PR? Added a fair version of unified memory pool similar to DataFusion `FairSpilPool` that caps the memory usage at `pool_size/num` The fair unified memory pool is the default for off-heap mode with this PR ## How are these changes tested? Exisiting tests

## Which issue does this PR close? Closes #1388 ## Rationale for this change Following up on #1369 and #1386 ## What changes are included in this PR? Updated the doc ## How are these changes tested?

kazuyukitanimura added 3 commits February 5, 2025 09:50

feat: Add fair memory pool

6ffcb8d

feat: Add fair unified memory pool

5cd0dcf

feat: Add fair unified memory pool

11477be

mem

16bbc8b

wForget reviewed Feb 6, 2025

View reviewed changes

docs/source/user-guide/configs.md Outdated Show resolved Hide resolved

common/src/main/scala/org/apache/comet/CometConf.scala Outdated Show resolved Hide resolved

wForget reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/fair_memory_pool.rs Show resolved Hide resolved

himadripal reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/jni_api.rs Outdated Show resolved Hide resolved

himadripal reviewed Feb 6, 2025

View reviewed changes

comphead and others added 2 commits February 6, 2025 11:13

Update common/src/main/scala/org/apache/comet/CometConf.scala

ca6f74c

Co-authored-by: Zhen Wang <[email protected]>

Update docs/source/user-guide/configs.md

4dae6ca

Co-authored-by: Zhen Wang <[email protected]>

comphead reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/fair_memory_pool.rs Show resolved Hide resolved

comphead reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/fair_memory_pool.rs Outdated Show resolved Hide resolved

comphead reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/fair_memory_pool.rs Outdated Show resolved Hide resolved

comphead reviewed Feb 6, 2025

View reviewed changes

native/core/src/execution/fair_memory_pool.rs Show resolved Hide resolved

kazuyukitanimura and others added 4 commits February 6, 2025 15:04

Update native/core/src/execution/fair_memory_pool.rs

59810c5

Co-authored-by: Oleks V <[email protected]>

Update native/core/src/execution/fair_memory_pool.rs

9f4d110

Co-authored-by: Oleks V <[email protected]>

feat: Add fair unified memory pool

7c83de3

feat: Add fair unified memory pool

319ac0d

wForget reviewed Feb 7, 2025

View reviewed changes

andygrove mentioned this pull request Jan 31, 2025

Comet 0.6.0 Release (Feb 2025) #1361

Open

kazuyukitanimura added 2 commits February 10, 2025 10:56

address review comments

6cfafd2

address review comments

05fd0f1

andygrove approved these changes Feb 10, 2025

View reviewed changes

kazuyukitanimura merged commit f099e6e into apache:main Feb 10, 2025
74 checks passed

kazuyukitanimura mentioned this pull request Feb 12, 2025

doc: update memory tuning guide #1394

Merged

kazuyukitanimura mentioned this pull request Feb 15, 2025

bug: Fix memory reservation and allocation problems for SortExec apache/datafusion#14644

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add fair unified memory pool #1369

feat: Add fair unified memory pool #1369

kazuyukitanimura commented Feb 5, 2025

codecov-commenter commented Feb 5, 2025 •

edited

Loading

kazuyukitanimura commented Feb 5, 2025

himadripal Feb 6, 2025

kazuyukitanimura Feb 6, 2025

comphead Feb 6, 2025

kazuyukitanimura Feb 6, 2025

kazuyukitanimura Feb 10, 2025

comphead left a comment

wForget Feb 7, 2025

kazuyukitanimura commented Feb 7, 2025

andygrove commented Feb 7, 2025

andygrove commented Feb 8, 2025

kazuyukitanimura commented Feb 10, 2025

kazuyukitanimura commented Feb 10, 2025

andygrove left a comment

kazuyukitanimura commented Feb 10, 2025

feat: Add fair unified memory pool #1369

feat: Add fair unified memory pool #1369

Conversation

kazuyukitanimura commented Feb 5, 2025

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

codecov-commenter commented Feb 5, 2025 • edited Loading

Codecov Report

kazuyukitanimura commented Feb 5, 2025

himadripal Feb 6, 2025

Choose a reason for hiding this comment

kazuyukitanimura Feb 6, 2025

Choose a reason for hiding this comment

comphead Feb 6, 2025

Choose a reason for hiding this comment

kazuyukitanimura Feb 6, 2025

Choose a reason for hiding this comment

kazuyukitanimura Feb 10, 2025

Choose a reason for hiding this comment

comphead left a comment

Choose a reason for hiding this comment

wForget Feb 7, 2025

Choose a reason for hiding this comment

kazuyukitanimura commented Feb 7, 2025

andygrove commented Feb 7, 2025

andygrove commented Feb 8, 2025

kazuyukitanimura commented Feb 10, 2025

kazuyukitanimura commented Feb 10, 2025

andygrove left a comment

Choose a reason for hiding this comment

kazuyukitanimura commented Feb 10, 2025

codecov-commenter commented Feb 5, 2025 •

edited

Loading