feat: Add a spark.comet.exec.memoryPool configuration for experimenting with various datafusion memory pool setups.

[Kontinuation]

Which issue does this PR close?

This issue relates to #996 and #1004

Rationale for this change

This is for investigating various approaches to simplify memory-related configuration and reduce the memory required to run large queries. @andygrove

What changes are included in this PR?

This PR adds a spark.comet.exec.memoryPool configuration for easily running queries using various memory pool setups.

• greedy: Each operator has its own GreedyMemoryPool, which is the same as the current situation.
• fair_spill: Each operator has its own FairSpillPool
• greedy_task_shared (default): All operators for the same task attempt share the same GreedyMemoryPool.
• fair_spill_task_shared: All operators for the same task attempt share the same FairSpillPool
• greedy_global: All operators in the same executor instance share the same GreedyMemoryPool
• fair_spill_global: All operators in the same executor instance share the same FairSpillPool

How are these changes tested?

TODO: add tests running in native memory management mode.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 77.27273% with 5 lines in your changes missing coverage. Please review.

Project coverage is 34.27%. Comparing base (591f45a) to head (bd5a0c7).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...ain/scala/org/apache/comet/CometExecIterator.scala	73.68%	2 Missing and 3 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1021      +/-   ##
============================================
- Coverage     34.30%   34.27%   -0.03%     
- Complexity      887      889       +2     
============================================
  Files           112      112              
  Lines         43429    43502      +73     
  Branches       9623     9615       -8     
============================================
+ Hits          14897    14912      +15     
- Misses        25473    25542      +69     
+ Partials       3059     3048      -11     

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

[andygrove]

Thanks for filing this @Kontinuation. I will close my PRs.

Also, there is a suggestion in #1017 for always using unified memory.

[Kontinuation]

Thanks for filing this @Kontinuation. I will close my PRs.

Also, there is a suggestion in #1017 for always using unified memory.

I agree that using the unified memory manager is a better approach. Vanilla Spark operators and comet operators are governed by the same memory manager and they are all using offheap memory. Vanilla Spark operators can free some memory when comet operators are under pressure. I'll also put more work into improving the unified memory management.

I think the native memory management approach may still be relevant when users don't want vanilla Spark to use off-heap memory. We can set the default value of spark.comet.exec.memoryPool as greedy_task_shared to make comet memory overhead configuration more intuitive. If users want memory-intensive operators to spill properly, they can try out fair_spill_task_shared.

[kazuyukitanimura]

Thank you @Kontinuation
Do you mind rebasing onto the latest main?

[Kontinuation]

Thank you @Kontinuation Do you mind rebasing onto the latest main?

I have noticed that for now Comet biases towards always using the unified memory manager, so this patch for native memory management may not be relevant, so I prefer closing this PR. If you still want to merge it into the upstream then I'll take some time to rebase it.

[andygrove]

I have noticed that for now Comet biases towards always using the unified memory manager, so this patch for native memory management may not be relevant, so I prefer closing this PR. If you still want to merge it into the upstream then I'll take some time to rebase it.

Just to add some context here, we have discovered a need to support existing deployments that are not using offheap memory, so would like to add a non-offheap option back in, and this PR seems like the best starting point for deciding how we want to configure the memory pools in that case.

[Kontinuation]

Just to add some context here, we have discovered a need to support existing deployments that are not using offheap memory, so would like to add a non-offheap option back in, and this PR seems like the best starting point for deciding how we want to configure the memory pools in that case.

OK. I have rebased this branch onto main.

[andygrove]

Thanks @Kontinuation. I will review this today.

[andygrove]

We should add that this config only takes effect when off-heap memory is disabled.

[andygrove]

We can do that in a follow up PR though, and add some more detailed documentation in the tuning guide.

[viirya]

Yea, I am not sure if all users can understand the difference between these pool configs. It is better we can add more description.

[andygrove]

I had to remove the following code from CometSparkSessionExtensions to test with off-heap disabled. We can remove this in a follow up PR as well.

      if (!isOffHeapEnabled(conf) && !isTesting) {
        logWarning("Comet native exec disabled because spark.memory.offHeap.enabled=false")
        withInfo(plan, "Comet native exec disabled because spark.memory.offHeap.enabled=false")
        return plan
      }

[andygrove]

I also had to remove the off-heap check here:

  private[comet] def isCometShuffleEnabled(conf: SQLConf): Boolean =
    COMET_EXEC_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) &&
      (isOffHeapEnabled(conf) || isTesting)

[andygrove]

I also had to make change in CometShuffleMemoryAllocatorTrait and I am now able to test. I will post results later today.

-    if (isSparkTesting && !useUnifiedMemAllocator) {
+    if (!useUnifiedMemAllocator) {

[andygrove]

We should be using the total executor memory here rather than just the overhead memory. Without this change I was unable to run any benchmarks.

Suggested change

	val maxMemory = CometSparkSessionExtensions.getCometMemoryOverhead(conf)
	val maxMemory = ConfigHelpers
	.byteFromString(conf.get("spark.executor.memory", "1024MB"), ByteUnit.BYTE)

[Kontinuation]

I thought that overhead memory is the amount of native memory reserved for comet to use. The actual memory usage will be doubled if we reserve the same amount of executor memory for comet.

[andygrove]

That's a good point. The default overhead amount was way too low for the benchmarks to run. Using the executor amount is probably overkill. I'll give this some more thought.

[andygrove]

I tested all memory pool options using the follow-on PR #1210. I tested with TPC-H @ 100 GB single node.

unified: completed in 367 seconds
greedy: completed in 359 seconds
greedy_task_shared: completed in 369 seconds
greedy_global: failed due to OOM
fair_spill_task_shared: too slow due to spilling in shuffle
fair_spill_global: too slow due to spilling in shuffle
fair_spill: too slow due to spilling in shuffle

I would like to go ahead and merge this PR because it provides a nice foundation for experimenting with different memory pool options.

These memory pool options are not available to users until the follow on PR #1210 is merged and I will add documentation as part of that PR.

Thanks again @Kontinuation for your work on this.

[andygrove]

@viirya @kazuyukitanimura If there are no objections I would like to go ahead and merge this soon.

[viirya]

I will take another look today.

[viirya]

Why we don't need to remove the pool for MemoryPoolType::GreedyTaskShared from TASK_SHARED_MEMORY_POOLS?

[Kontinuation]

It should be removed. I'll fix it in a follow-up commit.