From bf47902b80e51d2639279a27be20c56af6d42a2f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 20 Sep 2026 09:49:21 -0600 Subject: [PATCH 1/4] feat: remove memory accounting from Comet's on-heap mode On-heap mode exists so that Spark's own SQL suite and the Iceberg suites can run against Comet without changing Spark's memory configuration. It is off by default, sits in CATEGORY_TESTING, and the driver plugin disables Comet entirely unless a test opts in. Comet's own suites run off-heap. The accounting that mode performed did not protect anything. Native memory is not on the JVM heap, so there is no Spark pool it can honestly be charged to: NativeMemoryConsumer is hardcoded to MemoryMode.OFF_HEAP and Spark sizes that pool from spark.memory.offHeap.size alone, which is 0 when off-heap is disabled. The fixed-size DataFusion pool that stood in for one was sized from spark.comet.memoryOverhead, a config that predates the unified pool and that the driver plugin could not fold into the container on three of five supported Spark versions. On-heap mode now gets UnboundedMemoryPool on the native side and an unbounded per-task Unsafe page allocator on the JVM side, which removes: - six of the nine MemoryPoolType variants and memory_limit_per_task - spark.comet.memoryOverhead, spark.comet.exec.onHeap.memoryPool, spark.comet.shuffle.jvm.memoryFactor and spark.comet.shuffle.jvm.memoryWaitTimeout - the shared-pool blocking-allocation protocol in what is now CometUnboundedShuffleMemoryAllocator, along with allocateBlocking - the driver plugin's spark.executor.memoryOverhead mutation and the ShimCometDriverPlugin files that only it needed spark.comet.exec.onHeap.enabled stays: it is still the switch that keeps Comet off in on-heap mode unless a test opts in. Closes #6063 --- .ai/skills/review-comet-memory-pr/SKILL.md | 12 +- .ai/skills/review-comet-shuffle-pr/SKILL.md | 20 +- dev/diffs/3.4.3.diff | 5 +- dev/diffs/3.5.9.diff | 5 +- dev/diffs/4.0.4.diff | 7 +- dev/diffs/4.1.3.diff | 7 +- .../contributor-guide/memory_management.md | 15 +- native/core/src/execution/jni_api.rs | 2 - .../core/src/execution/memory_pools/config.rs | 69 +-- native/core/src/execution/memory_pools/mod.rs | 25 +- .../CometBoundedShuffleMemoryAllocator.java | 352 ------------- .../comet/CometShuffleMemoryAllocator.java | 22 +- .../CometShuffleMemoryAllocatorTrait.java | 11 - .../CometUnboundedShuffleMemoryAllocator.java | 174 +++++++ .../CometUnifiedShuffleMemoryAllocator.java | 4 +- .../CometBypassMergeSortShuffleWriter.java | 1 - .../shuffle/CometUnsafeShuffleWriter.java | 1 - .../comet/execution/shuffle/SpillWriter.java | 4 +- .../scala/org/apache/comet/CometConf.scala | 48 +- .../org/apache/comet/CometExecIterator.scala | 32 +- .../comet/CometSparkSessionExtensions.scala | 48 -- .../scala/org/apache/comet/GenerateDocs.scala | 11 +- .../main/scala/org/apache/comet/Native.scala | 1 - .../main/scala/org/apache/spark/Plugins.scala | 2 +- .../CometSparkSessionExtensionsSuite.scala | 35 -- .../apache/comet/exec/CometExecSuite.scala | 50 +- .../CometExecIteratorLifecycleSuite.scala | 1 - ...nboundedShuffleMemoryAllocatorSuite.scala} | 49 +- .../spark/shuffle/sort/SpillSorterSuite.scala | 11 +- .../spark/sql/CometTPCDSQuerySuite.scala | 1 - .../org/apache/spark/sql/CometTestBase.scala | 1 - .../sql/benchmark/CometExecBenchmark.scala | 1 - .../sql/comet/CometPlanStabilitySuite.scala | 1 - .../sql/comet/CometTaskMetricsSuite.scala | 2 +- .../shuffle/CometDiskBlockWriterSuite.scala | 490 ++---------------- 35 files changed, 342 insertions(+), 1178 deletions(-) delete mode 100644 spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java create mode 100644 spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java rename spark/src/test/scala/org/apache/spark/shuffle/comet/{CometBoundedShuffleMemoryAllocatorSuite.scala => CometUnboundedShuffleMemoryAllocatorSuite.scala} (70%) diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md index 111118dcd37..71070d9b93c 100644 --- a/.ai/skills/review-comet-memory-pr/SKILL.md +++ b/.ai/skills/review-comet-memory-pr/SKILL.md @@ -110,9 +110,11 @@ configuration, from the inside out: so it only removes an entry that is still its own, which handles the race where an `acquire` observes an expired `Weak` and inserts a replacement first. Do not let that check be simplified away. -- [ ] **Only `fair_unified` and `greedy_unified` are valid in off-heap mode.** Other pool types are - on-heap only and belong to `CATEGORY_TESTING`, because on-heap mode exists so the Spark SQL - test suite can run against Comet and must not be used in production. +- [ ] **`fair_unified` and `greedy_unified` are the only pool types.** On-heap mode ignores the + pool-type string and always gets `UnboundedMemoryPool`: it exists so the Spark SQL test suite + can run against Comet, it accounts for nothing, and it must not be used in production. A PR + re-adding a sized on-heap pool is reintroducing a budget that bounds nothing real (see + issue #6063). ## 4. The Spark Bridge @@ -140,8 +142,6 @@ memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction native code. - [ ] Changing `memoryPool.fraction` semantics affects every deployment that tuned it as a haircut for the accounting gap. -- [ ] `memory_limit_per_task` is read only by the on-heap pool types. A PR wiring it into an - off-heap path is probably confused. - [ ] On Kubernetes, `spark.memory.offHeap.size` is **part of** the pod limit, not headroom on top of it. A PR whose fix is "raise the off-heap size" is asking for fewer executors per node. `spark.executor.memoryOverhead` is the only real slack in the container, and JVM non-heap @@ -177,7 +177,7 @@ what test was added. Ask for at least one of: is recoverable at task level. `spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala` and -`CometBoundedShuffleMemoryAllocatorSuite.scala` are the existing JVM-side tests. A change to the +`CometUnboundedShuffleMemoryAllocatorSuite.scala` are the existing JVM-side tests. A change to the bridge or an allocator should extend one of them. ## 8. Does the PR Make `memory_management.md` Stale? diff --git a/.ai/skills/review-comet-shuffle-pr/SKILL.md b/.ai/skills/review-comet-shuffle-pr/SKILL.md index b768923ae7f..88fbce3102f 100644 --- a/.ai/skills/review-comet-shuffle-pr/SKILL.md +++ b/.ai/skills/review-comet-shuffle-pr/SKILL.md @@ -184,16 +184,16 @@ crates before, and the "Key Classes" tables in the docs are exactly what goes st ## 7. Tests -| Suite | Covers | -| ------------------------------------------------------------------- | ------------------------------------------------ | -| `org.apache.comet.exec.CometNativeShuffleSuite` | Native shuffle end to end | -| `org.apache.comet.exec.CometColumnarShuffleSuite` | JVM columnar shuffle end to end | -| `CometShuffle4_0Suite` | Spark 4.x specific behavior | -| `CometDiskBlockWriterSuite` | JVM spill and page handling | -| `NativeBatchDecoderIteratorLifecycleChecks`, `...ConcurrencyChecks` | Reader lifetime and concurrency | -| `CometNativeShuffleInputRDDSuite` | The scheduling-anchor RDD | -| `CometCeleborn*Suite` | The Celeborn path, which is easy to forget | -| `CometShuffleBenchmark` | Throughput, needs `-Dspark.comet.memoryOverhead` | +| Suite | Covers | +| ------------------------------------------------------------------- | ------------------------------------------ | +| `org.apache.comet.exec.CometNativeShuffleSuite` | Native shuffle end to end | +| `org.apache.comet.exec.CometColumnarShuffleSuite` | JVM columnar shuffle end to end | +| `CometShuffle4_0Suite` | Spark 4.x specific behavior | +| `CometDiskBlockWriterSuite` | JVM spill and page handling | +| `NativeBatchDecoderIteratorLifecycleChecks`, `...ConcurrencyChecks` | Reader lifetime and concurrency | +| `CometNativeShuffleInputRDDSuite` | The scheduling-anchor RDD | +| `CometCeleborn*Suite` | The Celeborn path, which is easy to forget | +| `CometShuffleBenchmark` | Throughput | Ask specifically: diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 3c1eff6b4b1..52f952cd51d 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -3053,10 +3053,10 @@ index dd55fcfe42c..d9a3f2df535 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index ed2e309fa07..040013dc8ab 100644 +index ed2e309fa07..54d417624ff 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -74,6 +74,20 @@ trait SharedSparkSessionBase +@@ -74,6 +74,19 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3072,7 +3072,6 @@ index ed2e309fa07..040013dc8ab 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") -+ .set("spark.comet.memoryOverhead", "4g") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index f47e7f543f0..c870cb4597c 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -3064,10 +3064,10 @@ index e937173a590..263934fbe7b 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index c23bf4204f7..f91ba524e6e 100644 +index c23bf4204f7..07d215aad2b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -97,6 +97,20 @@ trait SharedSparkSessionBase +@@ -97,6 +97,19 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3083,7 +3083,6 @@ index c23bf4204f7..f91ba524e6e 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") -+ .set("spark.comet.memoryOverhead", "2g") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 1d9978f4186..8ee3f1786fd 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -3030,7 +3030,7 @@ index 30503af0fab..1491f4bc2d5 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala -index 08fd8a9ecb5..06967aec8e1 100644 +index 08fd8a9ecb5..e0b8cada307 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala @@ -27,6 +27,7 @@ import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE @@ -3847,10 +3847,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..93221ef4cf5 100644 +index 720b13b812e..e3ac2cebc6e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,21 @@ trait SharedSparkSessionBase +@@ -98,6 +98,20 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3866,7 +3866,6 @@ index 720b13b812e..93221ef4cf5 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") -+ .set("spark.comet.memoryOverhead", "2g") + + } conf.set( diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index af819e4a1b9..1e90ff80e0b 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -3196,7 +3196,7 @@ index 30503af0fab..1491f4bc2d5 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala -index 08fd8a9ecb5..06967aec8e1 100644 +index 08fd8a9ecb5..e0b8cada307 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala @@ -27,6 +27,7 @@ import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE @@ -4141,10 +4141,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..93221ef4cf5 100644 +index 720b13b812e..e3ac2cebc6e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,21 @@ trait SharedSparkSessionBase +@@ -98,6 +98,20 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4160,7 +4160,6 @@ index 720b13b812e..93221ef4cf5 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") -+ .set("spark.comet.memoryOverhead", "2g") + + } conf.set( diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index e7916f9c785..d8889777320 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -26,9 +26,12 @@ anyone debugging an out-of-memory report. For user-facing tuning advice, see the This page covers off-heap mode (`spark.memory.offHeap.enabled=true`) only. Comet also has an on-heap mode, but it exists so that the Spark SQL test suite can run against Comet without changing -Spark's memory configuration. It must not be used in production, and it is not described here. The -pool types that only on-heap mode exposes belong to the `CATEGORY_TESTING` config group for the -same reason. +Spark's memory configuration. Comet performs no memory accounting in it: the native side gets +DataFusion's `UnboundedMemoryPool` and the JVM shuffle allocator +(`CometUnboundedShuffleMemoryAllocator`) hands out `Unsafe` pages against no budget. Native memory +is not on the JVM heap, so there is no Spark pool it could honestly be charged to, and the +fixed-size pool that used to stand in for one bounded nothing the container cares about. On-heap +mode must not be used in production, and it is not described further here. ## Overview @@ -203,9 +206,6 @@ Comet's under-accounting (see [The accounting gap](#the-accounting-gap)). It hol the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage overshoots its reservations by roughly that slice. -A second value, `memory_limit_per_task`, is computed and passed alongside it, but only the on-heap -pool types read it. - ### Resolving the pool type `parse_memory_pool_config` (`native/core/src/execution/memory_pools/config.rs`) turns the pool-type @@ -216,7 +216,8 @@ string and the limit into a `MemoryPoolConfig`. Two pool types are valid in off- | `fair_unified` (default) | `memory_limit` | Delegates to Spark's `TaskMemoryManager`; task-shared | | `greedy_unified` | n/a (pool size `0`) | Spark owns the limit entirely; task-shared | -Any other pool type is rejected with a configuration error. +Any other pool type is rejected with a configuration error. In on-heap mode the pool-type string is +ignored and the pool is always `UnboundedMemoryPool`. ## The pool stack diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 7652428410e..f859e2f8688 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -520,7 +520,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( off_heap_mode: jboolean, memory_pool_type: JString, memory_limit: jlong, - memory_limit_per_task: jlong, task_attempt_id: jlong, task_cpus: jlong, key_unwrapper_obj: JObject, @@ -576,7 +575,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( off_heap_mode != JNI_FALSE, memory_pool_type, memory_limit, - memory_limit_per_task, )?; let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index 312a3604383..75a5033a0b4 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -21,12 +21,6 @@ use crate::errors::{CometError, CometResult}; pub(crate) enum MemoryPoolType { GreedyUnified, FairUnified, - Greedy, - FairSpill, - GreedyTaskShared, - FairSpillTaskShared, - GreedyGlobal, - FairSpillGlobal, Unbounded, } @@ -48,47 +42,30 @@ pub(crate) fn parse_memory_pool_config( off_heap_mode: bool, memory_pool_type: String, memory_limit: i64, - memory_limit_per_task: i64, ) -> CometResult { + if !off_heap_mode { + // On-heap mode exists so that the Spark SQL tests can run against Comet without changing + // Spark's memory configuration. Comet's native allocations are not on the JVM heap, so + // there is no Spark pool they can honestly be charged to, and the fixed-size pool that + // used to stand in for one bounded nothing the container cares about. It is not a + // production configuration, so it accounts for nothing. + return Ok(MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0)); + } + let pool_size = memory_limit as usize; - let memory_pool_config = if off_heap_mode { - match memory_pool_type.as_str() { - "fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size), - "greedy_unified" => { - // the `unified` memory pool interacts with Spark's memory pool to allocate - // memory therefore does not need a size to be explicitly set. The pool size - // shared with Spark is set by `spark.memory.offHeap.size`. - MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) - } - _ => { - return Err(CometError::Config(format!( - "Unsupported memory pool type for off-heap mode: {memory_pool_type}" - ))) - } - } - } else { - // Use the memory pool from DF - let pool_size_per_task = memory_limit_per_task as usize; - match memory_pool_type.as_str() { - "fair_spill_task_shared" => { - MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task) - } - "greedy_task_shared" => { - MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared, pool_size_per_task) - } - "fair_spill_global" => { - MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal, pool_size) - } - "greedy_global" => MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size), - "fair_spill" => MemoryPoolConfig::new(MemoryPoolType::FairSpill, pool_size_per_task), - "greedy" => MemoryPoolConfig::new(MemoryPoolType::Greedy, pool_size_per_task), - "unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0), - _ => { - return Err(CometError::Config(format!( - "Unsupported memory pool type for on-heap mode: {memory_pool_type}" - ))) - } + match memory_pool_type.as_str() { + "fair_unified" => Ok(MemoryPoolConfig::new( + MemoryPoolType::FairUnified, + pool_size, + )), + "greedy_unified" => { + // the `unified` memory pool interacts with Spark's memory pool to allocate + // memory therefore does not need a size to be explicitly set. The pool size + // shared with Spark is set by `spark.memory.offHeap.size`. + Ok(MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0)) } - }; - Ok(memory_pool_config) + _ => Err(CometError::Config(format!( + "Unsupported memory pool type: {memory_pool_type}" + ))), + } } diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index d7c2911f913..8cb4b3de4f1 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -21,12 +21,9 @@ pub mod logging_pool; mod task_shared; mod unified_pool; -use datafusion::execution::memory_pool::{ - FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool, -}; +use datafusion::execution::memory_pool::{MemoryPool, TrackConsumersPool, UnboundedMemoryPool}; use fair_pool::CometFairMemoryPool; use jni::objects::{Global, JObject}; -use once_cell::sync::OnceCell; use std::num::NonZeroUsize; use std::sync::Arc; use unified_pool::CometUnifiedMemoryPool; @@ -68,26 +65,6 @@ pub(crate) fn create_memory_pool( pool_size, )) }), - MemoryPoolType::GreedyTaskShared => acquire_task_shared_pool(task_attempt_id, || { - tracked(GreedyMemoryPool::new(pool_size)) - }), - MemoryPoolType::FairSpillTaskShared => { - acquire_task_shared_pool(task_attempt_id, || tracked(FairSpillPool::new(pool_size))) - } - MemoryPoolType::Greedy => tracked(GreedyMemoryPool::new(pool_size)), - MemoryPoolType::FairSpill => tracked(FairSpillPool::new(pool_size)), - MemoryPoolType::GreedyGlobal => { - static GLOBAL_MEMORY_POOL_GREEDY: OnceCell> = OnceCell::new(); - let memory_pool = - GLOBAL_MEMORY_POOL_GREEDY.get_or_init(|| tracked(GreedyMemoryPool::new(pool_size))); - Arc::clone(memory_pool) - } - MemoryPoolType::FairSpillGlobal => { - static GLOBAL_MEMORY_POOL_FAIR: OnceCell> = OnceCell::new(); - let memory_pool = - GLOBAL_MEMORY_POOL_FAIR.get_or_init(|| tracked(FairSpillPool::new(pool_size))); - Arc::clone(memory_pool) - } MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()), } } diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java deleted file mode 100644 index 9692c1efa24..00000000000 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java +++ /dev/null @@ -1,352 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.spark.shuffle.comet; - -import java.io.IOException; -import java.util.BitSet; -import java.util.HashMap; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import org.apache.spark.SparkConf; -import org.apache.spark.TaskContext; -import org.apache.spark.memory.MemoryConsumer; -import org.apache.spark.memory.MemoryMode; -import org.apache.spark.memory.SparkOutOfMemoryError; -import org.apache.spark.memory.TaskMemoryManager; -import org.apache.spark.sql.internal.SQLConf; -import org.apache.spark.unsafe.array.LongArray; -import org.apache.spark.unsafe.memory.MemoryBlock; -import org.apache.spark.unsafe.memory.UnsafeMemoryAllocator; - -import org.apache.comet.CometConf$; -import org.apache.comet.CometSparkSessionExtensions$; - -/** - * A simple memory allocator used by `CometShuffleExternalSorter` to allocate memory blocks which - * store serialized rows. We don't rely on Spark memory allocator because we need to allocate - * off-heap memory no matter memory mode is on-heap or off-heap. This allocator is configured with - * fixed size of memory, and it will throw `SparkOutOfMemoryError` if the memory is not enough. - * - *

Some methods are copied from `org.apache.spark.unsafe.memory.TaskMemoryManager` with - * modifications. Most modifications are to remove the dependency on the configured memory mode. - * - *

This allocator is only used by Comet Columnar Shuffle when running in on-heap mode. It is used - * when users run in on-heap mode as well as in the Spark tests which require on-heap memory - * configuration. - * - *

Thus, this allocator is used to allocate separate off-heap memory allocation for Comet - * Columnar Shuffle and execution apart from Spark's on-heap memory configuration. - */ -public final class CometBoundedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { - private static final Logger logger = - LoggerFactory.getLogger(CometBoundedShuffleMemoryAllocator.class); - - private final UnsafeMemoryAllocator allocator = new UnsafeMemoryAllocator(); - - private final long pageSize; - private final long totalMemory; - private long allocatedMemory = 0L; - - /** How often a thread blocked in {@link #allocateBlocking(long)} logs that it is waiting. */ - private static final long WAIT_LOG_INTERVAL_MS = 30_000L; - - /** How often a blocked thread checks for cooperative task cancellation. */ - private static final long TASK_KILL_POLL_INTERVAL_MS = 1_000L; - - /** The number of bits used to address the page table. */ - private static final int PAGE_NUMBER_BITS = 13; - - /** The number of entries in the page table. */ - private static final int PAGE_TABLE_SIZE = 1 << PAGE_NUMBER_BITS; - - private final MemoryBlock[] pageTable = new MemoryBlock[PAGE_TABLE_SIZE]; - private final BitSet allocatedPages = new BitSet(PAGE_TABLE_SIZE); - - /** The thread that allocated each page, used to decide whether a blocked wait can succeed. */ - private final Thread[] pageOwners = new Thread[PAGE_TABLE_SIZE]; - - /** Pool memory currently retained by each thread. */ - private final HashMap retainedMemory = new HashMap<>(); - - /** Threads currently blocked in {@link #allocateBlocking(long)} and their request sizes. */ - private final HashMap waitingThreads = new HashMap<>(); - - private static final int OFFSET_BITS = 51; - private static final long MASK_LONG_LOWER_51_BITS = 0x7FFFFFFFFFFFFL; - - CometBoundedShuffleMemoryAllocator( - SparkConf conf, TaskMemoryManager taskMemoryManager, long pageSize) { - super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); - this.pageSize = pageSize; - this.totalMemory = - CometSparkSessionExtensions$.MODULE$.getCometShuffleMemorySize(conf, SQLConf.get()); - } - - /** - * Returns the current allocation total in bytes. - * - *

Allocations bypass Spark's memory manager and use this allocator's own counter. Since the - * allocator is shared across tasks, this reports the shared total rather than per-task usage. - */ - @Override - public synchronized long getUsed() { - return allocatedMemory; - } - - private synchronized long _acquireMemory(long size) { - if (allocatedMemory >= totalMemory) { - throw new SparkOutOfMemoryError( - "UNABLE_TO_ACQUIRE_MEMORY", - java.util.Map.of( - "requestedBytes", String.valueOf(size), - "receivedBytes", String.valueOf(totalMemory - allocatedMemory))); - } - long allocationSize = Math.min(size, totalMemory - allocatedMemory); - allocatedMemory += allocationSize; - return allocationSize; - } - - public long spill(long l, MemoryConsumer memoryConsumer) throws IOException { - return 0; - } - - public synchronized LongArray allocateArray(long size) { - long required = size * 8L; - MemoryBlock page = allocateMemoryBlock(required); - return new LongArray(page); - } - - public synchronized void freeArray(LongArray array) { - if (array == null) { - return; - } - free(array.memoryBlock()); - } - - public synchronized MemoryBlock allocate(long required) { - long size = Math.max(pageSize, required); - return allocateMemoryBlock(size); - } - - /** - * Like {@link #allocate(long)}, but waits for other tasks of this shared pool to free memory, - * mirroring how Spark's unified memory manager blocks a task until memory becomes available. - * Callers must first spill buffered data they can cheaply release; memory this thread still - * retains (e.g. the sorter's pointer array or sibling writers' pages) is included in the liveness - * checks below. The wait fails fast when it can never succeed: when the request does not fit next - * to the requester's retained memory, or when all allocated memory is retained by blocked threads - * and none of their requests fits in the free pool. Because the holders it depends on may in turn - * be blocked on resources outside this pool that only a task waiting here can release, the wait - * is also bounded by `spark.comet.shuffle.jvm.memoryWaitTimeout`, after which the managed - * allocation error is thrown and Spark's task retry can recover. Task cancellation or Java - * interruption aborts the wait. - */ - @Override - public synchronized MemoryBlock allocateBlocking(long required) { - long memoryWaitTimeoutMs = - (long) CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().get(); - long size = Math.max(pageSize, required); - Thread self = Thread.currentThread(); - TaskContext taskContext = TaskContext.get(); - long waitStart = 0; - long lastLog = 0; - try { - while (true) { - if (taskContext != null) { - taskContext.killTaskIfInterrupted(); - } - try { - return allocateMemoryBlock(size); - } catch (SparkOutOfMemoryError e) { - if (waitingThreads.put(self, size) == null) { - // Wake existing waiters so they re-evaluate the deadlock check against the enlarged - // waiting set. - notifyAll(); - } - // This thread cannot free what it retains while it waits, so a request that does not - // fit next to its own retained memory can never be satisfied. - if (size > totalMemory - retainedMemory.getOrDefault(self, 0L)) { - throw e; - } - // The allocation just failed, so the request does not fit in the unallocated pool. - // Waiting can only succeed while some thread can still free memory: either a thread - // outside the waiting set retains pool memory, or another waiter's request fits in the - // free pool, in which case that waiter can proceed and eventually free what it retains. - if (allocatedMemory <= retainedByWaitingThreads() && !anyWaiterCanProceed()) { - throw e; - } - // The holders this wait depends on may themselves be blocked on resources outside - // this pool (Spark execution memory, locks, I/O) that only a task waiting here can - // release - a cycle this allocator cannot observe. Bound the wait so such cycles - // unwind with the managed allocation error instead of hanging the executor; Spark's - // task retry can then recover. - long now = System.currentTimeMillis(); - if (waitStart == 0) { - waitStart = now; - lastLog = now; - logger.warn( - "Waiting for other tasks to free up {} bytes of Comet shuffle pool memory", size); - } else if (now - waitStart >= memoryWaitTimeoutMs) { - logger.warn( - "Giving up after waiting {} ms for {} bytes of Comet shuffle pool memory " - + "(see {})", - now - waitStart, - size, - CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().key()); - throw e; - } else if (now - lastLog >= WAIT_LOG_INTERVAL_MS) { - lastLog = now; - logger.warn( - "Still waiting ({} ms so far) for {} bytes of Comet shuffle pool memory; " - + "{} bytes free, {} thread(s) waiting", - now - waitStart, - size, - totalMemory - allocatedMemory, - waitingThreads.size()); - } - try { - wait( - Math.max( - 1L, - Math.min(TASK_KILL_POLL_INTERVAL_MS, memoryWaitTimeoutMs - (now - waitStart)))); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - // Not an allocation failure: stay non-fatal so that an intentional task kill is - // classified as TaskKilled rather than ExceptionFailure (Spark's killed-task handler - // only matches `InterruptedException | NonFatal(_)`). - throw new RuntimeException( - "Interrupted while waiting for Comet shuffle pool memory", ie); - } - } - } - } finally { - if (waitingThreads.remove(self) != null) { - notifyAll(); - } - } - } - - private long retainedByWaitingThreads() { - long retained = 0; - for (Thread thread : waitingThreads.keySet()) { - retained += retainedMemory.getOrDefault(thread, 0L); - } - return retained; - } - - private boolean anyWaiterCanProceed() { - long free = totalMemory - allocatedMemory; - for (long requested : waitingThreads.values()) { - if (requested <= free) { - return true; - } - } - return false; - } - - private synchronized MemoryBlock allocateMemoryBlock(long required) { - if (required > TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES) { - throw new TooLargePageException(required); - } - - long got = _acquireMemory(required); - - if (got < required) { - allocatedMemory -= got; - - throw new SparkOutOfMemoryError( - "UNABLE_TO_ACQUIRE_MEMORY", - java.util.Map.of( - "requestedBytes", String.valueOf(required), - "receivedBytes", String.valueOf(totalMemory - allocatedMemory))); - } - - int pageNumber = allocatedPages.nextClearBit(0); - if (pageNumber >= PAGE_TABLE_SIZE) { - allocatedMemory -= got; - - throw new IllegalStateException( - "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages"); - } - - MemoryBlock block = allocator.allocate(got); - - block.pageNumber = pageNumber; - pageTable[pageNumber] = block; - allocatedPages.set(pageNumber); - pageOwners[pageNumber] = Thread.currentThread(); - retainedMemory.merge(Thread.currentThread(), got, Long::sum); - - return block; - } - - public synchronized long free(MemoryBlock block) { - if (block.pageNumber == MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER - || block.pageNumber == MemoryBlock.FREED_IN_TMM_PAGE_NUMBER) { - // Already freed block - return 0; - } - long blockSize = block.size(); - allocatedMemory -= blockSize; - - Thread owner = pageOwners[block.pageNumber]; - pageOwners[block.pageNumber] = null; - if (owner != null) { - retainedMemory.computeIfPresent(owner, (t, v) -> v - blockSize <= 0 ? null : v - blockSize); - } - - pageTable[block.pageNumber] = null; - allocatedPages.clear(block.pageNumber); - block.pageNumber = MemoryBlock.FREED_IN_TMM_PAGE_NUMBER; - - allocator.free(block); - // Wake up tasks waiting in `allocateBlocking`. - notifyAll(); - return blockSize; - } - - /** - * Returns the offset in the page for the given page plus base offset address. Note that this - * method assumes that the page number is valid. - */ - public long getOffsetInPage(long pagePlusOffsetAddress) { - long offsetInPage = decodeOffset(pagePlusOffsetAddress); - int pageNumber = TaskMemoryManager.decodePageNumber(pagePlusOffsetAddress); - assert (pageNumber >= 0 && pageNumber < PAGE_TABLE_SIZE); - MemoryBlock page = pageTable[pageNumber]; - assert (page != null); - return page.getBaseOffset() + offsetInPage; - } - - public long decodeOffset(long pagePlusOffsetAddress) { - return pagePlusOffsetAddress & MASK_LONG_LOWER_51_BITS; - } - - public long encodePageNumberAndOffset(int pageNumber, long offsetInPage) { - assert (pageNumber >= 0); - return ((long) pageNumber) << OFFSET_BITS | offsetInPage & MASK_LONG_LOWER_51_BITS; - } - - public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) { - return encodePageNumberAndOffset(page.pageNumber, offsetInPage - page.getBaseOffset()); - } -} diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java index e8f9525b667..acd1f1589ed 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java @@ -19,37 +19,27 @@ package org.apache.spark.shuffle.comet; -import org.apache.spark.SparkConf; import org.apache.spark.memory.MemoryMode; import org.apache.spark.memory.TaskMemoryManager; /** - * An interface to instantiate either CometBoundedShuffleMemoryAllocator (on-heap mode) or + * An interface to instantiate either CometUnboundedShuffleMemoryAllocator (on-heap mode) or * CometUnifiedShuffleMemoryAllocator (off-heap mode). */ public final class CometShuffleMemoryAllocator { - private static CometShuffleMemoryAllocatorTrait INSTANCE; /** - * Returns the singleton instance of `CometShuffleMemoryAllocator`. This method should be used - * instead of the constructor to ensure that only one instance of `CometShuffleMemoryAllocator` is - * created. For on-heap mode (Spark tests), this returns `CometBoundedShuffleMemoryAllocator`. + * Returns the shuffle memory allocator for the current task. Allocators store pages in the + * `TaskMemoryManager`, or in their own page table, so a new instance is created per task. For + * on-heap mode (Spark tests), this returns `CometUnboundedShuffleMemoryAllocator`. */ public static CometShuffleMemoryAllocatorTrait getInstance( - SparkConf conf, TaskMemoryManager taskMemoryManager, long pageSize) { + TaskMemoryManager taskMemoryManager, long pageSize) { if (taskMemoryManager.getTungstenMemoryMode() == MemoryMode.OFF_HEAP) { - // CometShuffleMemoryAllocator stores pages in TaskMemoryManager which is not singleton, - // but one instance per task. So we need to create a new instance for each task. return new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, pageSize); } - synchronized (CometShuffleMemoryAllocator.class) { - if (INSTANCE == null) { - // CometBoundedShuffleMemoryAllocator handles pages by itself so it can be a singleton. - INSTANCE = new CometBoundedShuffleMemoryAllocator(conf, taskMemoryManager, pageSize); - } - } - return INSTANCE; + return new CometUnboundedShuffleMemoryAllocator(taskMemoryManager, pageSize); } } diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java index 2c9af24a7d9..36fa9d2ff48 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java @@ -33,17 +33,6 @@ protected CometShuffleMemoryAllocatorTrait( public abstract MemoryBlock allocate(long required); - /** - * Like {@link #allocate(long)}, but may wait for memory freed by other tasks instead of failing - * immediately. Callers must first spill buffered data they can cheaply release; an implementation - * that waits must account for any memory the caller retains while blocked. The default - * implementation does not wait: `CometUnifiedShuffleMemoryAllocator` delegates to Spark's memory - * manager, which already arbitrates memory between tasks. - */ - public MemoryBlock allocateBlocking(long required) { - return allocate(required); - } - public abstract long free(MemoryBlock block); public abstract long getOffsetInPage(long pagePlusOffsetAddress); diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java new file mode 100644 index 00000000000..ace2c264358 --- /dev/null +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.shuffle.comet; + +import java.io.IOException; +import java.util.BitSet; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.spark.memory.MemoryConsumer; +import org.apache.spark.memory.MemoryMode; +import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.unsafe.array.LongArray; +import org.apache.spark.unsafe.memory.MemoryBlock; +import org.apache.spark.unsafe.memory.UnsafeMemoryAllocator; + +/** + * The memory allocator used by `CometShuffleExternalSorter` to allocate the memory blocks that hold + * serialized rows when Spark runs in on-heap mode. + * + *

Spark's own allocator cannot be used here. `TaskMemoryManager.allocatePage` hands out pages + * from `tungstenMemoryAllocator`, which is the on-heap allocator in this mode, and the row + * addresses derived from these pages are passed to `writeSortedFileNative` for Rust to dereference. + * The pages therefore have to be `Unsafe`-allocated regardless of Spark's memory mode. + * + *

Nothing bounds these allocations. On-heap mode exists so that the Spark SQL tests can run + * against Comet without changing Spark's memory configuration; it is not a production + * configuration, and Comet performs no memory accounting in it. See the memory management page in + * the contributor guide. The off-heap path (`CometUnifiedShuffleMemoryAllocator`) is the one that + * accounts, through Spark's unified memory manager. + * + *

The page table below is adapted from `org.apache.spark.unsafe.memory.TaskMemoryManager`, with + * the dependency on the configured memory mode removed. + */ +public final class CometUnboundedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { + private final UnsafeMemoryAllocator allocator = new UnsafeMemoryAllocator(); + + private final long pageSize; + + /** The number of bits used to address the page table. */ + private static final int PAGE_NUMBER_BITS = 13; + + /** The number of entries in the page table. */ + private static final int PAGE_TABLE_SIZE = 1 << PAGE_NUMBER_BITS; + + private final MemoryBlock[] pageTable = new MemoryBlock[PAGE_TABLE_SIZE]; + private final BitSet allocatedPages = new BitSet(PAGE_TABLE_SIZE); + + /** Bytes currently held in the page table, reported by {@link #getUsed()}. */ + private final AtomicLong allocatedMemory = new AtomicLong(); + + private static final int OFFSET_BITS = 51; + private static final long MASK_LONG_LOWER_51_BITS = 0x7FFFFFFFFFFFFL; + + CometUnboundedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { + super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); + this.pageSize = pageSize; + } + + /** + * Returns the current allocation total in bytes. Allocations bypass Spark's memory manager, and + * this allocator keeps no budget, so it reports the bytes currently held in its page table. + * + *

Read from an {@link AtomicLong} rather than under this allocator's monitor: + * `TaskMemoryManager` calls this while holding its own monitor, so taking a second lock here + * would invert the lock order against any future caller that allocates while holding it. + */ + @Override + public long getUsed() { + return allocatedMemory.get(); + } + + public long spill(long l, MemoryConsumer memoryConsumer) throws IOException { + return 0; + } + + public synchronized LongArray allocateArray(long size) { + long required = size * 8L; + MemoryBlock page = allocateMemoryBlock(required); + return new LongArray(page); + } + + public synchronized void freeArray(LongArray array) { + if (array == null) { + return; + } + free(array.memoryBlock()); + } + + public synchronized MemoryBlock allocate(long required) { + long size = Math.max(pageSize, required); + return allocateMemoryBlock(size); + } + + private synchronized MemoryBlock allocateMemoryBlock(long required) { + if (required > TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES) { + throw new TooLargePageException(required); + } + + int pageNumber = allocatedPages.nextClearBit(0); + if (pageNumber >= PAGE_TABLE_SIZE) { + throw new IllegalStateException( + "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages"); + } + + MemoryBlock block = allocator.allocate(required); + + block.pageNumber = pageNumber; + pageTable[pageNumber] = block; + allocatedPages.set(pageNumber); + allocatedMemory.addAndGet(block.size()); + + return block; + } + + public synchronized long free(MemoryBlock block) { + if (block.pageNumber == MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER + || block.pageNumber == MemoryBlock.FREED_IN_TMM_PAGE_NUMBER) { + // Already freed block + return 0; + } + long blockSize = block.size(); + + pageTable[block.pageNumber] = null; + allocatedPages.clear(block.pageNumber); + block.pageNumber = MemoryBlock.FREED_IN_TMM_PAGE_NUMBER; + allocatedMemory.addAndGet(-blockSize); + + allocator.free(block); + return blockSize; + } + + /** + * Returns the offset in the page for the given page plus base offset address. Note that this + * method assumes that the page number is valid. + */ + public long getOffsetInPage(long pagePlusOffsetAddress) { + long offsetInPage = decodeOffset(pagePlusOffsetAddress); + int pageNumber = TaskMemoryManager.decodePageNumber(pagePlusOffsetAddress); + assert (pageNumber >= 0 && pageNumber < PAGE_TABLE_SIZE); + MemoryBlock page = pageTable[pageNumber]; + assert (page != null); + return page.getBaseOffset() + offsetInPage; + } + + public long decodeOffset(long pagePlusOffsetAddress) { + return pagePlusOffsetAddress & MASK_LONG_LOWER_51_BITS; + } + + public long encodePageNumberAndOffset(int pageNumber, long offsetInPage) { + assert (pageNumber >= 0); + return ((long) pageNumber) << OFFSET_BITS | offsetInPage & MASK_LONG_LOWER_51_BITS; + } + + public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) { + return encodePageNumberAndOffset(page.pageNumber, offsetInPage - page.getBaseOffset()); + } +} diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index aa8de6f17fe..b7c7c58848e 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -33,8 +33,8 @@ * configured with `MemoryMode.OFF_HEAP`, i.e. it is using off-heap memory. * *

If the user does not enable off-heap memory then we want to use - * CometBoundedShuffleMemoryAllocator. The tests also need to default to using this because off-heap - * is not enabled when running the Spark SQL tests. + * CometUnboundedShuffleMemoryAllocator. The tests also need to default to using this because + * off-heap is not enabled when running the Spark SQL tests. */ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java index fdb4b289d5b..24929648b1d 100644 --- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java +++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java @@ -171,7 +171,6 @@ public void write(Iterator> records) throws IOException { allocator = CometShuffleMemoryAllocator.getInstance( - conf, memoryManager, Math.min( CometShuffleExternalSorter.MAXIMUM_PAGE_SIZE_BYTES, diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java index 078a4c34219..42ffb603b98 100644 --- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java +++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java @@ -283,7 +283,6 @@ private void open() { assert (sorter == null); allocator = CometShuffleMemoryAllocator.getInstance( - sparkConf, memoryManager, Math.min( CometShuffleExternalSorter.MAXIMUM_PAGE_SIZE_BYTES, memoryManager.pageSizeBytes())); diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java index 8d9331b586c..1683af3e35a 100644 --- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java +++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java @@ -169,9 +169,7 @@ public boolean acquireNewPageIfNecessary(int required) { public void initialCurrentPage(int required) { assert (currentPage == null); try { - // This writer has already spilled its own data, so on a shared pool it may wait for other - // tasks to free memory instead of failing right away. - currentPage = allocator.allocateBlocking(required); + currentPage = allocator.allocate(required); } catch (SparkOutOfMemoryError e) { logger.error("Unable to acquire {} bytes of memory", required); throw e; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d371f1ba44c..3bdae17ed32 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -317,14 +317,6 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) - val COMET_ONHEAP_MEMORY_OVERHEAD: ConfigEntry[Long] = conf("spark.comet.memoryOverhead") - .category(CATEGORY_TESTING) - .doc( - "The amount of additional memory to be allocated per executor process for Comet, in MiB, " + - "when running Spark in on-heap mode.") - .bytesConf(ByteUnit.MiB) - .createWithDefault(1024) - val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.enabled") @@ -578,30 +570,6 @@ object CometConf extends ShimCometConf { .intConf .createWithDefault(Int.MaxValue) - val COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT: ConfigEntry[Long] = - conf("spark.comet.shuffle.jvm.memoryWaitTimeout") - .category(CATEGORY_SHUFFLE) - .doc( - "How long a Comet JVM (columnar) shuffle task running in on-heap mode waits for other " + - "tasks to free shared shuffle pool memory before failing with an out-of-memory error " + - "(Spark may then retry the task). The wait ends earlier when it provably cannot " + - "succeed. This is an internal config for testing purpose or advanced tuning.") - .internal() - .timeConf(TimeUnit.MILLISECONDS) - .createWithDefault(TimeUnit.MINUTES.toMillis(5)) - - val COMET_SHUFFLE_JVM_MEMORY_FACTOR: ConfigEntry[Double] = - conf("spark.comet.shuffle.jvm.memoryFactor") - .withAlternative("spark.comet.columnar.shuffle.memory.factor") - .category(CATEGORY_TESTING) - .doc("Fraction of Comet memory to be allocated per executor process for JVM (columnar) " + - s"shuffle when running in on-heap mode. $TUNING_GUIDE.") - .doubleConf - .checkValue( - factor => factor > 0, - "Ensure that Comet shuffle memory overhead factor is a double greater than 0") - .createWithDefault(1.0) - val COMET_BATCH_SIZE: ConfigEntry[Int] = conf("spark.comet.batchSize") .category(CATEGORY_TUNING) .doc("The columnar batch size, i.e., the maximum number of rows that a batch can contain.") @@ -817,7 +785,10 @@ object CometConf extends ShimCometConf { val COMET_ONHEAP_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.onHeap.enabled") .category(CATEGORY_TESTING) - .doc("Whether to allow Comet to run in on-heap mode. Required for running Spark SQL tests.") + .doc( + "Whether to allow Comet to run in on-heap mode. Required for running Spark SQL tests. " + + "Comet performs no memory accounting in on-heap mode, so its allocations are bounded " + + "by nothing; this must not be used in production.") .booleanConf .createWithEnvVarOrDefault("ENABLE_COMET_ONHEAP", false) @@ -831,17 +802,6 @@ object CometConf extends ShimCometConf { .stringConf .createWithDefault("fair_unified") - val COMET_ONHEAP_MEMORY_POOL_TYPE: ConfigEntry[String] = conf( - "spark.comet.exec.onHeap.memoryPool") - .category(CATEGORY_TESTING) - .doc( - "The type of memory pool to be used for Comet native execution " + - "when running Spark in on-heap mode. Available pool types are `greedy`, `fair_spill`, " + - "`greedy_task_shared`, `fair_spill_task_shared`, `greedy_global`, `fair_spill_global`, " + - "and `unbounded`.") - .stringConf - .createWithDefault("greedy_task_shared") - val COMET_OFFHEAP_MEMORY_POOL_FRACTION: ConfigEntry[Double] = conf("spark.comet.exec.memoryPool.fraction") .category(CATEGORY_TUNING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 4de949b9de5..1eb861bd5c6 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -130,7 +130,6 @@ class CometExecIterator( memoryConfig.offHeapMode, memoryConfig.memoryPoolType, memoryConfig.memoryLimit, - memoryConfig.memoryLimitPerTask, taskAttemptId, taskCPUs, keyUnwrapper, @@ -393,8 +392,6 @@ object CometExecIterator extends Logging { } def getMemoryConfig(conf: SparkConf): MemoryConfig = { - val numCores = numDriverOrExecutorCores(conf) - val coresPerTask = conf.get("spark.task.cpus", "1").toInt // there are different paths for on-heap vs off-heap mode val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf) if (offHeapMode) { @@ -402,28 +399,19 @@ object CometExecIterator extends Logging { val offHeapSize = ByteUnit.MiB.toBytes(conf.getSizeAsMb("spark.memory.offHeap.size")) val memoryFraction = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.get() val memoryLimit = (offHeapSize * memoryFraction).toLong - val memoryLimitPerTask = (memoryLimit.toDouble * coresPerTask / numCores).toLong val memoryPoolType = COMET_OFFHEAP_MEMORY_POOL_TYPE.get() logDebug( s"memoryPoolType=$memoryPoolType, " + s"offHeapSize=${toMB(offHeapSize)}, " + s"memoryFraction=$memoryFraction, " + - s"memoryLimit=${toMB(memoryLimit)}, " + - s"memoryLimitPerTask=${toMB(memoryLimitPerTask)}") - MemoryConfig(offHeapMode, memoryPoolType = memoryPoolType, memoryLimit, memoryLimitPerTask) + s"memoryLimit=${toMB(memoryLimit)}") + MemoryConfig(offHeapMode, memoryPoolType, memoryLimit) } else { - // we'll use the built-in memory pool from DF, and initializes with `memory_limit` - // and `memory_fraction` below. - val memoryLimit = CometSparkSessionExtensions.getCometMemoryOverhead(conf) - // example 16GB maxMemory * 16 cores with 4 cores per task results - // in memory_limit_per_task = 16 GB * 4 / 16 = 16 GB / 4 = 4GB - val memoryLimitPerTask = (memoryLimit.toDouble * coresPerTask / numCores).toLong - val memoryPoolType = COMET_ONHEAP_MEMORY_POOL_TYPE.get() - logDebug( - s"memoryPoolType=$memoryPoolType, " + - s"memoryLimit=${toMB(memoryLimit)}, " + - s"memoryLimitPerTask=${toMB(memoryLimitPerTask)}") - MemoryConfig(offHeapMode, memoryPoolType = memoryPoolType, memoryLimit, memoryLimitPerTask) + // On-heap mode exists only so that the Spark SQL tests can run against Comet without + // changing Spark's memory configuration, and native memory cannot be charged to Spark's + // on-heap pool, so nothing is accounted. See the memory management contributor guide. + logDebug("on-heap mode: native memory is unbounded and unaccounted") + MemoryConfig(offHeapMode, memoryPoolType = "unbounded", memoryLimit = 0) } } @@ -455,8 +443,4 @@ object CometExecIterator extends Logging { } } -case class MemoryConfig( - offHeapMode: Boolean, - memoryPoolType: String, - memoryLimit: Long, - memoryLimitPerTask: Long) +case class MemoryConfig(offHeapMode: Boolean, memoryPoolType: String, memoryLimit: Long) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 4b37d7b61a8..1a2a8b95ad0 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -23,7 +23,6 @@ import java.nio.ByteOrder import org.apache.spark.{SparkConf, SparkEnv} import org.apache.spark.internal.Logging -import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.{SparkSession, SparkSessionExtensions} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag} @@ -247,53 +246,6 @@ object CometSparkSessionExtensions extends Logging { org.apache.spark.SPARK_VERSION >= "4.2" } - /** - * Determines required memory overhead in MB per executor process for Comet when running in - * on-heap mode. - */ - def getCometMemoryOverheadInMiB(sparkConf: SparkConf): Long = { - if (isOffHeapEnabled(sparkConf)) { - // off-heap mode sizes the native memory pool from spark.memory.offHeap.size instead - // (see CometExecIterator.getMemoryConfig), so this value does not apply - return 0 - } - ConfigHelpers.byteFromString( - sparkConf.get( - COMET_ONHEAP_MEMORY_OVERHEAD.key, - COMET_ONHEAP_MEMORY_OVERHEAD.defaultValueString), - ByteUnit.MiB) - } - - /** - * Calculates required memory overhead in bytes per executor process for Comet when running in - * on-heap mode. - */ - def getCometMemoryOverhead(sparkConf: SparkConf): Long = { - ByteUnit.MiB.toBytes(getCometMemoryOverheadInMiB(sparkConf)) - } - - /** - * Calculates required shuffle memory size in bytes per executor process for Comet when running - * in on-heap mode. - */ - def getCometShuffleMemorySize(sparkConf: SparkConf, conf: SQLConf = SQLConf.get): Long = { - assert(!isOffHeapEnabled(sparkConf)) - - val cometMemoryOverhead = getCometMemoryOverheadInMiB(sparkConf) - - val overheadFactor = COMET_SHUFFLE_JVM_MEMORY_FACTOR.get(conf) - - val shuffleMemorySize = (overheadFactor * cometMemoryOverhead).toLong - if (shuffleMemorySize > cometMemoryOverhead) { - logWarning( - s"Configured shuffle memory size $shuffleMemorySize is larger than Comet memory overhead " + - s"$cometMemoryOverhead, using Comet memory overhead instead.") - ByteUnit.MiB.toBytes(cometMemoryOverhead) - } else { - ByteUnit.MiB.toBytes(shuffleMemorySize) - } - } - def isOffHeapEnabled(sparkConf: SparkConf): Boolean = { sparkConf.getBoolean("spark.memory.offHeap.enabled", false) } diff --git a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala index 4eb29a6d58d..f36487080d3 100644 --- a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala +++ b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala @@ -27,7 +27,6 @@ import scala.collection.mutable.ListBuffer import org.apache.spark.sql.catalyst.analysis.FunctionRegistry import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, StringLPad, StringRPad} -import org.apache.comet.CometConf.COMET_ONHEAP_MEMORY_OVERHEAD import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.{CodegenDispatchFallback, CometAggregateExpressionSerde, CometCodegenDispatch, CometExpressionSerde, Compatible, Incompatible, NativeOptInAvailable, QueryPlanSerde, Unsupported} @@ -347,14 +346,8 @@ object GenerateDocs { if (conf.defaultValue.isEmpty) { w.write(s"| `${conf.key}` | $docWithEnvVar | |\n".getBytes) } else { - val isBytesConf = conf.key == COMET_ONHEAP_MEMORY_OVERHEAD.key - if (isBytesConf) { - val bytes = conf.defaultValue.get.asInstanceOf[Long] - w.write(s"| `${conf.key}` | $docWithEnvVar | $bytes MiB |\n".getBytes) - } else { - val defaultVal = conf.defaultValueString - w.write(s"| `${conf.key}` | $docWithEnvVar | $defaultVal |\n".getBytes) - } + val defaultVal = conf.defaultValueString + w.write(s"| `${conf.key}` | $docWithEnvVar | $defaultVal |\n".getBytes) } } w.write("\n".getBytes) diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 93b396ce0f2..2c41c602fdf 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -72,7 +72,6 @@ class Native extends NativeBase { offHeapMode: Boolean, memoryPoolType: String, memoryLimit: Long, - memoryLimitPerTask: Long, taskAttemptId: Long, taskCPUs: Long, keyUnwrapper: CometFileKeyUnwrapper, diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 24b6ab3eeaa..ca153d7aca3 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -36,7 +36,7 @@ import org.apache.comet.annotation.Public /** * Comet driver plugin. This class is loaded by Spark's plugin framework. It will be instantiated * on driver side only. It will update the SparkConf with the extra configuration provided by - * Comet, e.g., Comet memory configurations. + * Comet, e.g., the cache serializer and the session extension. * * Note that `SparkContext.conf` is spark package only. So this plugin must be in spark package. * Although `SparkContext.getConf` is public, it returns a copy of the SparkConf, so it cannot diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 331d4b4ece9..575d2796d8b 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -19,7 +19,6 @@ package org.apache.comet -import org.apache.spark.SparkConf import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.plans.logical.LocalRelation @@ -154,38 +153,4 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { NativeBase.setLoaded(true) SQLConf.get.setConfString(CometConf.COMET_DEBUG_ENABLED.key, "false") } - - def getBytesFromMib(mib: Long): Long = mib * 1024 * 1024 - - test("Default Comet memory overhead") { - val conf = new SparkConf() - assert(getCometMemoryOverhead(conf) == getBytesFromMib(1024)) - } - - test("Comet memory overhead") { - val sparkConf = new SparkConf() - sparkConf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "10g") - assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(1024 * 10)) - } - - test("Comet memory overhead (off heap)") { - val sparkConf = new SparkConf() - sparkConf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "64g") - sparkConf.set("spark.memory.offHeap.enabled", "true") - sparkConf.set("spark.memory.offHeap.size", "10g") - // off-heap mode sizes the native pool from spark.memory.offHeap.size instead - assert(getCometMemoryOverhead(sparkConf) == 0) - } - - test("Comet shuffle memory factor") { - val conf = new SparkConf() - - val sqlConf = new SQLConf - sqlConf.setConfString(CometConf.COMET_SHUFFLE_JVM_MEMORY_FACTOR.key, "0.2") - - // Minimum Comet memory overhead is 384MB - assert( - getCometShuffleMemorySize(conf, sqlConf) == - getBytesFromMib((1024 * 0.2).toLong)) - } } 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 387785dfc5b..af84468780c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3017,39 +3017,35 @@ class CometExecSuite extends CometTestBase { } test("spill sort with (multiple) dictionaries") { - withSQLConf(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key -> "15MB") { - withTempDir { dir => - val path = new Path(dir.toURI.toString, "part-r-0.parquet") - makeRawTimeParquetFileColumns(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) - readParquetFile(path.toString) { df => - Seq( - $"_0".desc_nulls_first, - $"_0".desc_nulls_last, - $"_0".asc_nulls_first, - $"_0".asc_nulls_last).foreach { colOrder => - val query = df.sortWithinPartitions(colOrder) - checkSparkAnswerAndOperator(query) - } + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + makeRawTimeParquetFileColumns(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) + readParquetFile(path.toString) { df => + Seq( + $"_0".desc_nulls_first, + $"_0".desc_nulls_last, + $"_0".asc_nulls_first, + $"_0".asc_nulls_last).foreach { colOrder => + val query = df.sortWithinPartitions(colOrder) + checkSparkAnswerAndOperator(query) } } } } test("spill sort with (multiple) dictionaries on mixed columns") { - withSQLConf(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key -> "15MB") { - withTempDir { dir => - val path = new Path(dir.toURI.toString, "part-r-0.parquet") - makeRawTimeParquetFile(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) - readParquetFile(path.toString) { df => - Seq( - $"_6".desc_nulls_first, - $"_6".desc_nulls_last, - $"_6".asc_nulls_first, - $"_6".asc_nulls_last).foreach { colOrder => - // TODO: We should be able to sort on dictionary timestamp column - val query = df.sortWithinPartitions(colOrder) - checkSparkAnswerAndOperator(query) - } + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + makeRawTimeParquetFile(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) + readParquetFile(path.toString) { df => + Seq( + $"_6".desc_nulls_first, + $"_6".desc_nulls_last, + $"_6".asc_nulls_first, + $"_6".asc_nulls_last).foreach { colOrder => + // TODO: We should be able to sort on dictionary timestamp column + val query = df.sortWithinPartitions(colOrder) + checkSparkAnswerAndOperator(query) } } } diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index a5b98dd5988..3d3b63f0fb2 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -109,7 +109,6 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { true, "fair_unified", 64L << 20, - 64L << 20, taskAttemptId, 1L, null, diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala similarity index 70% rename from spark/src/test/scala/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocatorSuite.scala rename to spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala index ab12d81a53f..81667684694 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocatorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala @@ -24,27 +24,16 @@ import scala.collection.mutable.ArrayBuffer import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.SparkConf -import org.apache.spark.memory.{SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager} -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.unsafe.memory.MemoryBlock -import org.apache.comet.CometConf - -class CometBoundedShuffleMemoryAllocatorSuite extends AnyFunSuite { +class CometUnboundedShuffleMemoryAllocatorSuite extends AnyFunSuite { private val pageSize = 4096L - private val memoryLimit = 1024L * 1024 - private def newAllocator(): CometBoundedShuffleMemoryAllocator = { - val conf = new SparkConf(false) - .set("spark.memory.offHeap.enabled", "false") - .set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "1m") + private def newAllocator(): CometUnboundedShuffleMemoryAllocator = { + val conf = new SparkConf(false).set("spark.memory.offHeap.enabled", "false") val taskMemoryManager = new TaskMemoryManager(new TestMemoryManager(conf), 0) - val sqlConf = new SQLConf - sqlConf.setConfString(CometConf.COMET_SHUFFLE_JVM_MEMORY_FACTOR.key, "1.0") - SQLConf.withExistingConf(sqlConf) { - // Avoid the executor singleton so each test owns its budget and allocated pages. - new CometBoundedShuffleMemoryAllocator(conf, taskMemoryManager, pageSize) - } + new CometUnboundedShuffleMemoryAllocator(taskMemoryManager, pageSize) } test("getUsed reports actual page sizes and ignores repeated frees") { @@ -97,29 +86,15 @@ class CometBoundedShuffleMemoryAllocatorSuite extends AnyFunSuite { } } - test("getUsed is unchanged when a partial grant is rolled back") { - val allocator = newAllocator() - val page = allocator.allocate(1) - try { - intercept[SparkOutOfMemoryError] { - allocator.allocate(memoryLimit) - } - assert(allocator.getUsed === page.size()) - } finally { - allocator.free(page) - } - assert(allocator.getUsed === 0L) - } - - test("getUsed is unchanged when the budget is exhausted") { + test("allocations are not bounded by any budget") { + // On-heap mode performs no memory accounting, so a request far larger than anything Comet + // would have been granted under the old fixed-size pool succeeds. val allocator = newAllocator() - val page = allocator.allocate(memoryLimit) + val huge = 64L * 1024 * 1024 + val page = allocator.allocate(huge) try { - assert(allocator.getUsed === memoryLimit) - intercept[SparkOutOfMemoryError] { - allocator.allocate(1) - } - assert(allocator.getUsed === memoryLimit) + assert(page.size() === huge) + assert(allocator.getUsed === huge) } finally { allocator.free(page) } diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala index c6330fb28e5..42ecc7611e8 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala @@ -27,7 +27,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.{SparkConf, TaskContext} import org.apache.spark.executor.ShuffleWriteMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} -import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocator +import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.Platform import org.apache.spark.unsafe.UnsafeAlignedOffset @@ -47,6 +47,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { private var conf: SparkConf = _ private var memoryManager: TestMemoryManager = _ private var taskMemoryManager: TaskMemoryManager = _ + private var allocator: CometShuffleMemoryAllocatorTrait = _ override def beforeEach(): Unit = { super.beforeEach() @@ -55,9 +56,13 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { memoryManager = new TestMemoryManager(conf) memoryManager.limit(100 * 1024 * 1024) // 100MB taskMemoryManager = new TaskMemoryManager(memoryManager, 0) + // One allocator per test: pages are addressed by a page number in the allocator's own table, + // so everything a sorter touches has to come from the same instance. + allocator = CometShuffleMemoryAllocator.getInstance(taskMemoryManager, PAGE_SIZE) } override def afterEach(): Unit = { + allocator = null if (taskMemoryManager != null) { taskMemoryManager.cleanUpAllAllocatedMemory() taskMemoryManager = null @@ -75,7 +80,6 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { spills: java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo] = new java.util.LinkedList[org.apache.spark.sql.comet.execution.shuffle.SpillInfo](), partitionChecksums: Array[Long] = new Array[Long](10)): SpillSorter = { - val allocator = CometShuffleMemoryAllocator.getInstance(conf, taskMemoryManager, PAGE_SIZE) val schema = createTestSchema() val writeMetrics = new ShuffleWriteMetrics() val taskContext = TaskContext.empty() @@ -225,7 +229,6 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { val sorter = createSpillSorter() try { val initialMemory = sorter.getMemoryUsage() - val allocator = CometShuffleMemoryAllocator.getInstance(conf, taskMemoryManager, PAGE_SIZE) val newArray = allocator.allocateArray(INITIAL_SIZE * 2) sorter.expandPointerArray(newArray) @@ -269,7 +272,7 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach { offHeapMemoryManager.limit(64L * 1024 * 1024) val offHeapTaskMemoryManager = new TaskMemoryManager(offHeapMemoryManager, 0) val allocator = - CometShuffleMemoryAllocator.getInstance(offHeapConf, offHeapTaskMemoryManager, PAGE_SIZE) + CometShuffleMemoryAllocator.getInstance(offHeapTaskMemoryManager, PAGE_SIZE) // The block manager is only touched when spilling, which this test never does. val sorter = new CometShuffleExternalSorter( allocator, diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala b/spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala index a82d6f99083..bf80f12e9bb 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala @@ -187,7 +187,6 @@ class CometTPCDSQuerySuite conf.set(CometConf.COMET_EXEC_ENABLED.key, "true") conf.set(CometConf.COMET_NATIVE_SCAN_ENABLED.key, "true") conf.set(CometConf.COMET_SHUFFLE_ENABLED.key, "true") - conf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "15g") conf.set(CometConf.COMET_EXPLAIN_TRANSFORMATIONS.key, "true") conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true") conf.set(MEMORY_OFFHEAP_ENABLED.key, "true") diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 9740a4a5468..9700e0fd076 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -94,7 +94,6 @@ abstract class CometTestBase conf.set(CometConf.COMET_NATIVE_SCAN_ENABLED.key, "true") conf.set(CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key, "false") conf.set(CometConf.COMET_SCAN_ALLOW_DISABLED_PARQUET_VECTORIZED_READER.key, "true") - conf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "2g") conf.set(CometConf.COMET_EXEC_SORT_MERGE_JOIN_WITH_JOIN_FILTER_ENABLED.key, "true") // Fail loudly if a serde declines an operator without stating why, rather than letting the // generic " is not supported" message mask the missing reason. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala index 2e7e143a472..615e456ca8d 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala @@ -59,7 +59,6 @@ object CometExecBenchmark extends CometBenchmarkBase { sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") - sparkSession.conf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "10g") // TODO: support dictionary encoding in vectorized execution sparkSession.conf.set("parquet.enable.dictionary", "false") sparkSession.conf.set("spark.sql.shuffle.partitions", "2") diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala index 8f309f5cd27..4d648574494 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala @@ -318,7 +318,6 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa conf.set(MEMORY_OFFHEAP_SIZE.key, "2g") conf.set(CometConf.COMET_ENABLED.key, "true") conf.set(CometConf.COMET_EXEC_ENABLED.key, "true") - conf.set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "1g") conf.set(CometConf.COMET_SHUFFLE_ENABLED.key, "true") conf.set(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key, "true") diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index 4679125541a..64a12c0b7e6 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -292,7 +292,7 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { memoryManager.limit(10 * 1024 * 1024) val taskMemoryManager = new TaskMemoryManager(memoryManager, 0) val allocator = - CometShuffleMemoryAllocator.getInstance(conf, taskMemoryManager, pageSize) + CometShuffleMemoryAllocator.getInstance(taskMemoryManager, pageSize) val taskContext = TaskContext.empty() val writeMetrics = taskContext.taskMetrics.shuffleWriteMetrics val sorter = new CometShuffleExternalSorter( diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala index 3b92acdfd69..782f9fdff7a 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala @@ -21,18 +21,16 @@ package org.apache.spark.sql.comet.execution.shuffle import java.io.File import java.util.{LinkedList => JLinkedList, Properties} -import java.util.concurrent.CountDownLatch import org.scalatest.concurrent.{Signaler, ThreadSignaler, TimeLimits} import org.scalatest.funsuite.AnyFunSuite -import org.scalatest.time.{Seconds, Span} -import org.apache.spark.{Partitioner, SparkConf, SparkContext, SparkEnv, TaskContext, TaskContextImpl, TaskKilledException} +import org.apache.spark.{Partitioner, SparkConf, SparkContext, SparkEnv, TaskContextImpl} import org.apache.spark.executor.{ShuffleWriteMetrics, TaskMetrics} -import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} +import org.apache.spark.memory.{SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager} import org.apache.spark.shuffle.api.{ShuffleExecutorComponents, ShuffleMapOutputWriter, ShufflePartitionWriter} import org.apache.spark.shuffle.api.metadata.MapOutputCommitMessage -import org.apache.spark.shuffle.comet.{CometBoundedShuffleMemoryAllocator, CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} +import org.apache.spark.shuffle.comet.{CometShuffleMemoryAllocator, CometShuffleMemoryAllocatorTrait} import org.apache.spark.shuffle.sort.SpillSorter import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow} @@ -81,8 +79,8 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { val tmmB = new TaskMemoryManager(memoryManager, 1L) val taskContextA = newTaskContext(tmmA, 0L) val taskContextB = newTaskContext(tmmB, 1L) - val allocatorA = CometShuffleMemoryAllocator.getInstance(conf, tmmA, pageSize) - val allocatorB = CometShuffleMemoryAllocator.getInstance(conf, tmmB, pageSize) + val allocatorA = CometShuffleMemoryAllocator.getInstance(tmmA, pageSize) + val allocatorB = CometShuffleMemoryAllocator.getInstance(tmmB, pageSize) val tempDir = Utils.createTempDir() try { @@ -200,316 +198,6 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { } } - test("on-heap shared pool: a task with nothing to spill waits for other tasks to free memory") { - // On-heap mode uses one executor-wide CometBoundedShuffleMemoryAllocator. A task whose own - // writers hold nothing spillable must wait for other tasks to free pool memory rather than - // fail with SparkOutOfMemoryError. - val conf = new SparkConf().set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool - resetOnHeapAllocatorSingleton() - val memoryManager = new TestMemoryManager(conf) - val tmmA = new TaskMemoryManager(memoryManager, 0L) - val tmmB = new TaskMemoryManager(memoryManager, 1L) - val waitTimeoutKey = "spark.comet.shuffle.jvm.memoryWaitTimeout" - val propertiesA = new Properties - propertiesA.setProperty(waitTimeoutKey, "5s") - val propertiesB = new Properties - propertiesB.setProperty(waitTimeoutKey, "100ms") - val taskContextA = newTaskContext(tmmA, 0L, propertiesA) - val taskContextB = newTaskContext(tmmB, 1L, propertiesB) - - // Initialize the shared allocator under B's short timeout, then reuse it under A's longer - // timeout. The wait below must use the requesting task's setting, not the first task's. - TaskContext.setTaskContext(taskContextB) - val allocatorB = - try { - CometShuffleMemoryAllocator.getInstance(conf, tmmB, pageSize) - } finally { - TaskContext.unset() - } - TaskContext.setTaskContext(taskContextA) - val allocatorA = - try { - CometShuffleMemoryAllocator.getInstance(conf, tmmA, pageSize) - } finally { - TaskContext.unset() - } - assert(allocatorA eq allocatorB) - - val tempDir = Utils.createTempDir() - try { - val serializer = new UnsafeRowSerializer(1).newInstance() - val fileA = new File(tempDir, "onheap-taskA") - val fileB = new File(tempDir, "onheap-taskB") - val writerA = new CometDiskBlockWriter( - fileA, - allocatorA, - taskContextA, - serializer, - schema, - new ShuffleWriteMetrics, - conf, - false, - new JLinkedList[CometDiskBlockWriter]()) - val writerB = new CometDiskBlockWriter( - fileB, - allocatorB, - taskContextB, - serializer, - schema, - new ShuffleWriteMetrics, - conf, - false, - new JLinkedList[CometDiskBlockWriter]()) - - val toUnsafe = UnsafeProjection.create(schema) - def insertOne(writer: CometDiskBlockWriter): Unit = { - writer.insertRow(toUnsafe(InternalRow(new Array[Byte](1024))), 0) - } - - // Task B fills the whole 1 MiB pool (4 pages) on its own thread, idles for a while, and - // then finishes, freeing the pool. - var rowsB = 0L - var segmentBLength = 0L - val poolFilled = new CountDownLatch(1) - val threadB = new Thread(() => { - TaskContext.setTaskContext(taskContextB) - try { - while (writerB.getActiveMemoryUsage < 4 * pageSize) { - insertOne(writerB) - rowsB += 1 - } - poolFilled.countDown() - Thread.sleep(500) - segmentBLength = writerB.close().length - } finally { - TaskContext.unset() - } - }) - threadB.start() - poolFilled.await() - - // Task A's first insert finds the pool exhausted and has nothing of its own to spill; it - // must wait for task B to finish instead of throwing SparkOutOfMemoryError. - var rowsA = 0L - TaskContext.setTaskContext(taskContextA) - try { - (0 until 10).foreach { _ => - insertOne(writerA) - rowsA += 1 - } - } finally { - TaskContext.unset() - } - threadB.join() - - val segmentA = writerA.close() - assert(writerA.getOutputRecords == rowsA) - assert(writerB.getOutputRecords == rowsB) - assert(segmentA.length > 0) - assert(segmentBLength > 0) - // Neither task spilled: A waited instead of stealing B's memory, and B was never touched. - assert(taskContextA.taskMetrics.diskBytesSpilled == 0) - assert(taskContextB.taskMetrics.diskBytesSpilled == 0) - } finally { - Utils.deleteRecursively(tempDir) - tmmA.cleanUpAllAllocatedMemory() - tmmB.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() - } - } - - test("on-heap shared pool: unsatisfiable allocations fail fast instead of waiting") { - val conf = new SparkConf().set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool - resetOnHeapAllocatorSingleton() - val memoryManager = new TestMemoryManager(conf) - val tmm = new TaskMemoryManager(memoryManager, 0L) - val taskContext = newTaskContext(tmm, 0L) - val allocator = CometShuffleMemoryAllocator.getInstance(conf, tmm, pageSize) - val tempDir = Utils.createTempDir() - try { - val writer = newWriter(new File(tempDir, "unsatisfiable"), allocator, taskContext, conf) - val toUnsafe = UnsafeProjection.create(schema) - failAfter(Span(60, Seconds)) { - // A row larger than the whole pool can never be satisfied. - intercept[SparkOutOfMemoryError] { - writer.insertRow(toUnsafe(InternalRow(new Array[Byte](2 * 1024 * 1024))), 0) - } - // A request that does not fit next to memory this thread itself retains (like the - // sorter's pointer array, which survives an empty spill) can never be satisfied either. - val bounded = allocator.asInstanceOf[CometBoundedShuffleMemoryAllocator] - val retained = bounded.allocateArray(pageSize / 8) // retain one page worth of longs - try { - intercept[SparkOutOfMemoryError] { - writer.insertRow(toUnsafe(InternalRow(new Array[Byte](900 * 1024))), 0) - } - } finally { - bounded.freeArray(retained) - } - } - writer.freeMemory() - } finally { - Utils.deleteRecursively(tempDir) - tmm.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() - } - } - - test("on-heap shared pool: interruption or cancellation aborts a waiting task") { - val conf = new SparkConf().set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool - resetOnHeapAllocatorSingleton() - val memoryManager = new TestMemoryManager(conf) - val tmmHolder = new TaskMemoryManager(memoryManager, 0L) - val tmmWaiter = new TaskMemoryManager(memoryManager, 1L) - val tmmCancelled = new TaskMemoryManager(memoryManager, 2L) - val taskContextHolder = newTaskContext(tmmHolder, 0L) - val taskContextWaiter = newTaskContext(tmmWaiter, 1L) - val taskContextCancelled = newTaskContext(tmmCancelled, 2L) - val allocatorHolder = CometShuffleMemoryAllocator.getInstance(conf, tmmHolder, pageSize) - val allocatorWaiter = CometShuffleMemoryAllocator.getInstance(conf, tmmWaiter, pageSize) - val allocatorCancelled = - CometShuffleMemoryAllocator.getInstance(conf, tmmCancelled, pageSize) - val tempDir = Utils.createTempDir() - try { - val writerHolder = - newWriter(new File(tempDir, "holder"), allocatorHolder, taskContextHolder, conf) - val writerWaiter = - newWriter(new File(tempDir, "waiter"), allocatorWaiter, taskContextWaiter, conf) - val writerCancelled = - newWriter(new File(tempDir, "cancelled"), allocatorCancelled, taskContextCancelled, conf) - val toUnsafe = UnsafeProjection.create(schema) - def insertOne(writer: CometDiskBlockWriter): Unit = { - writer.insertRow(toUnsafe(InternalRow(new Array[Byte](1024))), 0) - } - - failAfter(Span(60, Seconds)) { - // The holder fills the whole pool on its own thread and keeps it until released. - val poolFilled = new CountDownLatch(1) - val release = new CountDownLatch(1) - val holderThread = new Thread(() => { - while (writerHolder.getActiveMemoryUsage < 4 * pageSize) { - insertOne(writerHolder) - } - poolFilled.countDown() - release.await() - writerHolder.freeMemory() - }) - holderThread.start() - poolFilled.await() - - // A Java interrupt aborts the wait without surfacing a fatal allocation error. - @volatile var interruptedFailure: Throwable = null - val waiterThread = new Thread(() => { - try insertOne(writerWaiter) - catch { case t: Throwable => interruptedFailure = t } - }) - waiterThread.start() - while (!isBlockedInWait(waiterThread)) { - Thread.sleep(10) - } - waiterThread.interrupt() - waiterThread.join() - - assert(interruptedFailure != null) - assert(!interruptedFailure.isInstanceOf[OutOfMemoryError]) - assert(interruptedFailure.isInstanceOf[RuntimeException]) - assert(interruptedFailure.getCause.isInstanceOf[InterruptedException]) - - // With interruptOnCancel=false Spark only marks TaskContext; it does not interrupt the - // Java thread. The allocator must poll that flag and throw TaskKilledException promptly. - @volatile var cancelledFailure: Throwable = null - val cancelledThread = new Thread(() => { - TaskContext.setTaskContext(taskContextCancelled) - try insertOne(writerCancelled) - catch { case t: Throwable => cancelledFailure = t } - finally TaskContext.unset() - }) - cancelledThread.start() - while (!isBlockedInWait(cancelledThread)) { - Thread.sleep(10) - } - taskContextCancelled.markInterrupted("test cancellation") - cancelledThread.join(5000) - val exitedOnCancellation = !cancelledThread.isAlive - - release.countDown() - holderThread.join() - cancelledThread.join() - writerWaiter.freeMemory() - writerCancelled.freeMemory() - - assert(exitedOnCancellation) - assert(!cancelledThread.isInterrupted) - assert(cancelledFailure.isInstanceOf[TaskKilledException]) - } - } finally { - Utils.deleteRecursively(tempDir) - tmmHolder.cleanUpAllAllocatedMemory() - tmmWaiter.cleanUpAllAllocatedMemory() - tmmCancelled.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() - } - } - - test("on-heap shared pool: no false deadlock while another waiter can proceed") { - // Two waiters each retain a small pointer array while a holder owns most of the pool. When - // the holder frees, the large waiter's request still does not fit, but the small waiter's - // does: the small waiter must be allowed to proceed and finish, unblocking the large one, - // instead of either waiter being declared deadlocked. - val conf = new SparkConf().set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool - resetOnHeapAllocatorSingleton() - val memoryManager = new TestMemoryManager(conf) - val tmm = new TaskMemoryManager(memoryManager, 0L) - val bounded = CometShuffleMemoryAllocator - .getInstance(conf, tmm, pageSize) - .asInstanceOf[CometBoundedShuffleMemoryAllocator] - try { - failAfter(Span(60, Seconds)) { - val holderBlock = bounded.allocate(900 * 1024) - - @volatile var largeError: Throwable = null - @volatile var smallError: Throwable = null - val arraysAllocated = new CountDownLatch(2) - val largeWaiter = new Thread(() => { - val array = bounded.allocateArray(4096) // retains 32768 bytes - arraysAllocated.countDown() - try { - bounded.free(bounded.allocateBlocking(999448)) - } catch { - case t: Throwable => largeError = t - } finally { - bounded.freeArray(array) - } - }) - val smallWaiter = new Thread(() => { - val array = bounded.allocateArray(4096) // retains 32768 bytes - arraysAllocated.countDown() - try { - bounded.free(bounded.allocateBlocking(262144)) - } catch { - case t: Throwable => smallError = t - } finally { - bounded.freeArray(array) - } - }) - largeWaiter.start() - smallWaiter.start() - arraysAllocated.await() - while (!isBlockedInWait(largeWaiter) || !isBlockedInWait(smallWaiter)) { - Thread.sleep(10) - } - - bounded.free(holderBlock) - largeWaiter.join() - smallWaiter.join() - assert(largeError == null) - assert(smallError == null) - } - } finally { - tmm.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() - } - } - private def newWriter( file: File, allocator: CometShuffleMemoryAllocatorTrait, @@ -527,19 +215,21 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { new JLinkedList[CometDiskBlockWriter]()) } - test("on-heap shared pool: a fatal error during write() frees the task's buffered pages") { - // Spark's ShuffleWriteProcessor only calls stop(false) when write() throws an Exception, so - // a fatal error such as SparkOutOfMemoryError skips it. The buffered pages live in the - // executor-shared bounded pool where Spark's task-memory cleanup cannot see them, so - // write() itself must free them on the way out or they starve other tasks forever. + test("a fatal error during write() frees the task's buffered pages") { + // Spark's ShuffleWriteProcessor only calls stop(false) when write() throws an Exception, so a + // fatal error such as SparkOutOfMemoryError skips it. write() itself must therefore free the + // pages it buffered on the way out; until the task ends they would otherwise stay charged and + // starve the other tasks sharing the pool. val conf = new SparkConf() .setMaster("local[1]") .setAppName("CometDiskBlockWriterSuite") - .set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "1g") .set("spark.buffer.pageSize", "256k") - resetOnHeapAllocatorSingleton() val sc = new SparkContext(conf) val memoryManager = new TestMemoryManager(conf) + // Small enough that the rows below force real spilling. + memoryManager.limit(1024 * 1024) val tmm = new TaskMemoryManager(memoryManager, 0L) try { val taskContext = newTaskContext(tmm, 0L) @@ -587,39 +277,33 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { } assert(spillFiles.forall(!_.exists())) - // All pages buffered by the failed task were reclaimed, so a full-pool allocation - // succeeds; without the reclaim the orphaned pages would make it fail forever. - val bounded = CometShuffleMemoryAllocator - .getInstance(conf, tmm, pageSize) - .asInstanceOf[CometBoundedShuffleMemoryAllocator] - bounded.free(bounded.allocate(1024 * 1024)) + // Every page the failed task buffered was reclaimed by write() itself, before any + // task-level cleanup ran. + assert(tmm.getMemoryConsumptionForThisTask == 0) } finally { sc.stop() tmm.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() } } - test("on-heap shared pool: a failed SpillSorter constructor does not leak pool memory") { + test("a failed SpillSorter constructor does not leak pool memory") { // The unsafe sorter's constructor first allocates a one-entry array (8 bytes) inside // ShuffleInMemorySorter and then its real pointer array. If the second allocation fails, the - // first must be reclaimed: the writer is never handed to Spark, so no cleanup path would - // ever free it, and the orphaned bytes would make later full-pool allocations impossible. - val conf = new SparkConf().set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool - resetOnHeapAllocatorSingleton() + // first must be reclaimed: the writer is never handed to Spark, so no cleanup path short of + // the task ending would ever free it. + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "1g") val memoryManager = new TestMemoryManager(conf) val tmm = new TaskMemoryManager(memoryManager, 0L) val taskContext = newTaskContext(tmm, 0L) - val bounded = CometShuffleMemoryAllocator - .getInstance(conf, tmm, pageSize) - .asInstanceOf[CometBoundedShuffleMemoryAllocator] + val allocator = CometShuffleMemoryAllocator.getInstance(tmm, pageSize) try { - // A healthy task holds most of the pool, leaving room for the one-entry array but not for - // the 4096-entry pointer array. - val holderBlock = bounded.allocate(1008 * 1024) + // Room for the one-entry array but not for the 4096-entry (32 KiB) pointer array. + memoryManager.limit(1024) intercept[SparkOutOfMemoryError] { new SpillSorter( - bounded, + allocator, 4096, schema, UnsafeAlignedOffset.getUaoSize(), @@ -633,15 +317,15 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { new JLinkedList[SpillInfo](), () => ()) } - bounded.free(holderBlock) - // With the constructor cleanup the pool is empty again, so a full-pool allocation - // succeeds; a leaked constructor allocation would make it fail forever. - bounded.free(bounded.allocate(1024 * 1024)) + // With the constructor cleanup nothing is left charged; a leaked constructor allocation + // would show up here. + assert(allocator.getUsed == 0) // A failure after the pointer array is adopted (here: schema serialization rejecting an // out-of-range parquet.field.id) must free the adopted array as well. The enclosing // sorter field is never assigned in this case, so not even the unsafe writer's // task-completion listener could see the allocation. + memoryManager.limit(1024 * 1024) val badField = StructField( "a", IntegerType, @@ -650,7 +334,7 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { val badSchema = StructType(Seq(StructField("s", StructType(Seq(badField))))) intercept[IllegalArgumentException] { new SpillSorter( - bounded, + allocator, 4096, badSchema, UnsafeAlignedOffset.getUaoSize(), @@ -664,98 +348,27 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { new JLinkedList[SpillInfo](), () => ()) } - bounded.free(bounded.allocate(1024 * 1024)) + assert(allocator.getUsed == 0) } finally { tmm.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() } } - test("on-heap shared pool: a waiter unwinds when the holder is blocked on Spark memory") { - // Cross-pool cycle: a task blocked in Spark's execution-memory pool retains Comet pool - // memory, while another task waits in the Comet pool for memory that only the blocked task - // could free. Spark's pool may in turn only be freed by the Comet waiter, and dependencies - // outside this pool are invisible to the allocator, so the wait is bounded: the Comet - // waiter must unwind at the timeout so its task can release memory. - val conf = new SparkConf() - .set("spark.comet.memoryOverhead", "1") // 1 MiB Comet pool - .set("spark.testing.memory", (8 * 1024 * 1024).toString) - .set("spark.testing.reservedMemory", "0") - .set("spark.memory.fraction", "1.0") - resetOnHeapAllocatorSingleton() - SQLConf.get.setConfString("spark.comet.shuffle.jvm.memoryWaitTimeout", "1s") - val unified = UnifiedMemoryManager(conf, numCores = 2) - val tmmSparkHog = new TaskMemoryManager(unified, 0L) - val tmmHolder = new TaskMemoryManager(unified, 1L) - val bounded = CometShuffleMemoryAllocator - .getInstance(conf, new TaskMemoryManager(new TestMemoryManager(conf), 2L), pageSize) - .asInstanceOf[CometBoundedShuffleMemoryAllocator] - def newConsumer(tmm: TaskMemoryManager): MemoryConsumer = - new MemoryConsumer(tmm, 1024 * 1024, MemoryMode.ON_HEAP) { - override def spill(size: Long, trigger: MemoryConsumer): Long = 0 - } - try { - failAfter(Span(60, Seconds)) { - // Another task owns the whole Spark execution pool. - val sparkHog = newConsumer(tmmSparkHog) - assert(sparkHog.acquireMemory(8 * 1024 * 1024) == 8 * 1024 * 1024) - - // The holder retains Comet pool memory and then blocks acquiring Spark memory. - val holderReleased = new CountDownLatch(1) - val holderThread = new Thread(() => { - val cometBlock = bounded.allocate(500 * 1024) - val holderConsumer = newConsumer(tmmHolder) - try { - val got = holderConsumer.acquireMemory(1024 * 1024) // blocks until sparkHog frees - holderConsumer.freeMemory(got) - } finally { - bounded.free(cometBlock) - holderReleased.countDown() - } - }) - holderThread.start() - def parkedOnSparkMemory: Boolean = - holderThread.getState == Thread.State.WAITING && - holderThread.getStackTrace - .exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") - while (!parkedOnSparkMemory) { - Thread.sleep(10) - } - - // The Comet waiter cannot fit next to the blocked holder's memory and must unwind - // instead of waiting for a release that can never come. - intercept[SparkOutOfMemoryError] { - bounded.allocateBlocking(800 * 1024) - } - - // Once Spark memory frees, the holder resumes and releases its Comet memory. - sparkHog.freeMemory(8 * 1024 * 1024) - holderReleased.await() - holderThread.join() - bounded.free(bounded.allocate(1024 * 1024)) - } - } finally { - SQLConf.get.unsetConf("spark.comet.shuffle.jvm.memoryWaitTimeout") - tmmSparkHog.cleanUpAllAllocatedMemory() - tmmHolder.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() - } - } - - test("on-heap shared pool: the unsafe writer allocates nothing before write()") { + test("the unsafe writer allocates nothing before write()") { // Spark evaluates the shuffle input iterator between constructing the writer and calling // write(), and that evaluation can block on Spark's execution-memory pool (e.g. an eager - // input sort). The writer must not retain Comet pool memory across that window, or two - // tasks can deadlock across the two pools; and when write() fails, even with a fatal - // error, everything it allocated must be reclaimed. + // input sort). The writer must not retain pool memory across that window, or two tasks can + // deadlock across the two pools; and when write() fails, even with a fatal error, everything + // it allocated must be reclaimed. val conf = new SparkConf() .setMaster("local[1]") .setAppName("CometDiskBlockWriterSuite") - .set("spark.comet.memoryOverhead", "1") // 1 MiB shared pool + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "1g") .set("spark.buffer.pageSize", "256k") - resetOnHeapAllocatorSingleton() val sc = new SparkContext(conf) val memoryManager = new TestMemoryManager(conf) + memoryManager.limit(1024 * 1024) val tmm = new TaskMemoryManager(memoryManager, 0L) try { val taskContext = newTaskContext(tmm, 0L) @@ -801,11 +414,8 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { // Construct the writer exactly as Spark does before evaluating the input iterator. val writer = newUnsafeWriter() - val bounded = CometShuffleMemoryAllocator - .getInstance(conf, tmm, pageSize) - .asInstanceOf[CometBoundedShuffleMemoryAllocator] - // Construction must not have taken anything from the shared pool. - bounded.free(bounded.allocate(1024 * 1024)) + // Construction must not have taken anything from the pool. + assert(tmm.getMemoryConsumptionForThisTask == 0) // A fatal error from the record iterator mid-write must not leak the sorter's pages or // pointer array either, and the task-completion listener stays a no-op afterwards. @@ -825,11 +435,10 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { writer.write(rows) } taskContext.markTaskCompleted(None) - bounded.free(bounded.allocate(1024 * 1024)) + assert(tmm.getMemoryConsumptionForThisTask == 0) } finally { sc.stop() tmm.cleanUpAllAllocatedMemory() - resetOnHeapAllocatorSingleton() } } @@ -851,17 +460,4 @@ class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { } } } - - /** Whether the thread is parked in `Object.wait` (the allocator uses a timed wait). */ - private def isBlockedInWait(thread: Thread): Boolean = { - val state = thread.getState - state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING - } - - /** Clears the CometShuffleMemoryAllocator singleton so this suite controls its pool size. */ - private def resetOnHeapAllocatorSingleton(): Unit = { - val field = classOf[CometShuffleMemoryAllocator].getDeclaredField("INSTANCE") - field.setAccessible(true) - field.set(null, null) - } } From fda45acd68f034617c6840374c49c94c6184a230 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 20 Sep 2026 10:22:00 -0600 Subject: [PATCH 2/4] fix: register renamed allocator suite and make page exhaustion spillable Four follow-ups from a self review of this PR. The renamed `CometUnboundedShuffleMemoryAllocatorSuite` was never renamed in `pr_build_linux.yml` or `pr_build_macos.yml`, so `check-suites` failed Preflight and the suite ran in neither job. `CometUnboundedShuffleMemoryAllocator` never refuses a page, which leaves its page table as the only limit it has. Exhausting it threw `IllegalStateException`, which nothing catches, so the task failed. Report `SparkOutOfMemoryError` instead, which `SpillWriter.acquireNewPageIfNecessary` and `CometShuffleExternalSorter.growPointerArrayIfNecessary` already answer by spilling and retrying. A page number indexes one allocator's own table, so an address encoded by one instance cannot be resolved by another. That was previously guaranteed by the executor-wide singleton and is now a requirement on callers, so say so on both the allocator and the factory. The JVM shuffle guide described a single allocator that spills when an allocation fails, which is true of the off-heap path only. --- .github/workflows/pr_build_linux.yml | 2 +- .github/workflows/pr_build_macos.yml | 2 +- docs/source/contributor-guide/jvm_shuffle.md | 16 ++++++++++++++-- .../comet/CometShuffleMemoryAllocator.java | 4 ++++ .../CometUnboundedShuffleMemoryAllocator.java | 19 +++++++++++++++++-- ...UnboundedShuffleMemoryAllocatorSuite.scala | 15 ++++++++++++--- 6 files changed, 49 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index a9c184a20b0..3efc65bb26f 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -513,7 +513,7 @@ jobs: org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.comet.exec.DisableAQECometAsyncShuffleSuite - org.apache.spark.shuffle.comet.CometBoundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b47ed5a46f6..1150b132dab 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -161,7 +161,7 @@ jobs: org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.comet.exec.DisableAQECometAsyncShuffleSuite - org.apache.spark.shuffle.comet.CometBoundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index a4d0b37361c..1226b4c82b4 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -190,8 +190,20 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies ## Memory Management -- `CometShuffleMemoryAllocator`: Custom allocator for off-heap memory pages -- Memory is allocated in pages; when allocation fails, writers spill to disk +- `CometShuffleMemoryAllocator.getInstance` returns the allocator for the task. Pages are always + `Unsafe`-allocated, because their addresses are handed to native code, but what bounds them + depends on Spark's memory mode. +- Off-heap mode gets `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer` + drawing from `spark.memory.offHeap.size`. When it cannot acquire a page the writer spills to + disk. +- On-heap mode gets `CometUnboundedShuffleMemoryAllocator`, which keeps no budget and so never + refuses a page. Nothing bounds these allocations, and memory pressure never triggers a spill. + That mode exists only so the Spark SQL tests can run against Comet. See + [Memory Management](memory_management.md). +- Row count still triggers spilling in either mode. `CometDiskBlockWriter` spills at + `min(spark.comet.shuffle.jvm.spillThreshold, spark.comet.shuffle.jvm.batchSize)`. + `CometShuffleExternalSorter` spills at `spark.comet.shuffle.jvm.spillThreshold` alone, which + defaults to `Int.MaxValue`, so on the sort path that trigger is effectively off by default. - `CometDiskBlockWriter` coordinates spilling across all partition writers (largest first) ## Configuration diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java index acd1f1589ed..61c0c48275b 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java @@ -32,6 +32,10 @@ public final class CometShuffleMemoryAllocator { * Returns the shuffle memory allocator for the current task. Allocators store pages in the * `TaskMemoryManager`, or in their own page table, so a new instance is created per task. For * on-heap mode (Spark tests), this returns `CometUnboundedShuffleMemoryAllocator`. + * + *

Call this once per task and share the result. `CometUnboundedShuffleMemoryAllocator` numbers + * pages within its own table, so a record address produced by one instance cannot be resolved by + * another, and two instances in one task would decode each other's addresses to the wrong memory. */ public static CometShuffleMemoryAllocatorTrait getInstance( TaskMemoryManager taskMemoryManager, long pageSize) { diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java index ace2c264358..4ce62af6e8c 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java @@ -25,6 +25,7 @@ import org.apache.spark.memory.MemoryConsumer; import org.apache.spark.memory.MemoryMode; +import org.apache.spark.memory.SparkOutOfMemoryError; import org.apache.spark.memory.TaskMemoryManager; import org.apache.spark.unsafe.array.LongArray; import org.apache.spark.unsafe.memory.MemoryBlock; @@ -47,6 +48,12 @@ * *

The page table below is adapted from `org.apache.spark.unsafe.memory.TaskMemoryManager`, with * the dependency on the configured memory mode removed. + * + *

A page number indexes this instance's own table, so an address encoded by one allocator means + * nothing to another. Everything that addresses a page therefore has to come from the same + * instance, which is why each caller creates one allocator and shares it for the life of the task + * rather than creating one per writer. The off-heap allocator has no such constraint, because it + * stores its pages in the `TaskMemoryManager` instead. */ public final class CometUnboundedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { private final UnsafeMemoryAllocator allocator = new UnsafeMemoryAllocator(); @@ -115,8 +122,16 @@ private synchronized MemoryBlock allocateMemoryBlock(long required) { int pageNumber = allocatedPages.nextClearBit(0); if (pageNumber >= PAGE_TABLE_SIZE) { - throw new IllegalStateException( - "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages"); + // The page table is the only limit this allocator has. Report it the way a memory manager + // reports a refused acquisition, so that the callers which already handle that by spilling + // and retrying (`SpillWriter.acquireNewPageIfNecessary`, + // `CometShuffleExternalSorter.growPointerArrayIfNecessary`) can free pages and make + // progress, rather than failing the task with an exception nothing catches. + throw new SparkOutOfMemoryError( + "UNABLE_TO_ACQUIRE_MEMORY", + java.util.Map.of( + "requestedBytes", String.valueOf(required), + "receivedBytes", String.valueOf(0))); } MemoryBlock block = allocator.allocate(required); diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala index 81667684694..7b2be4fbdff 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocatorSuite.scala @@ -24,7 +24,7 @@ import scala.collection.mutable.ArrayBuffer import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.SparkConf -import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} +import org.apache.spark.memory.{SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager} import org.apache.spark.unsafe.memory.MemoryBlock class CometUnboundedShuffleMemoryAllocatorSuite extends AnyFunSuite { @@ -101,7 +101,10 @@ class CometUnboundedShuffleMemoryAllocatorSuite extends AnyFunSuite { assert(allocator.getUsed === 0L) } - test("getUsed is unchanged when the page table is exhausted") { + test("an exhausted page table is reported as a refused acquisition") { + // The page table is the only limit this allocator has. It has to surface as + // SparkOutOfMemoryError so that the writers, which respond to that by spilling and retrying, + // can free pages instead of failing the task. val allocator = newAllocator() val pages = ArrayBuffer.empty[MemoryBlock] val maxPages = 1 << 13 @@ -111,10 +114,16 @@ class CometUnboundedShuffleMemoryAllocatorSuite extends AnyFunSuite { } val allocatedBytes = pages.map(_.size()).sum assert(allocator.getUsed === allocatedBytes) - intercept[IllegalStateException] { + intercept[SparkOutOfMemoryError] { allocator.allocateArray(1) } assert(allocator.getUsed === allocatedBytes) + + // Freeing a page makes room again, which is what lets a spilling writer make progress. + allocator.free(pages.remove(pages.length - 1)) + val page = allocator.allocateArray(1).memoryBlock() + pages += page + assert(allocator.getUsed === allocatedBytes) } finally { pages.foreach(allocator.free) } From 12f1f1f1edb8c89e7d1ad76013d134667b487a5b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 20 Sep 2026 10:36:54 -0600 Subject: [PATCH 3/4] fix: drop the dead TimeLimits mixin from CometDiskBlockWriterSuite Scalafix `RemoveUnused` fails the Lint Java jobs on the 3.4, 3.5 and 4.0 profiles: - private implicit val signaler: Signaler = ThreadSignaler + ThreadSignaler This PR removed the five tests that covered the blocking-allocation protocol, and they were the only callers of `failAfter`. With those gone the `TimeLimits` mixin and the `Signaler` that served it are unreferenced, so remove both rather than the val alone. Lint did not catch this earlier in the PR because Preflight was failing on the renamed allocator suite and gated every downstream job. Spark 4.1 is excluded from the lint matrix, so formatting locally on the default profile does not run scalafix at all. --- .../comet/execution/shuffle/CometDiskBlockWriterSuite.scala | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala index 782f9fdff7a..7a9a4cb0637 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala @@ -22,7 +22,6 @@ package org.apache.spark.sql.comet.execution.shuffle import java.io.File import java.util.{LinkedList => JLinkedList, Properties} -import org.scalatest.concurrent.{Signaler, ThreadSignaler, TimeLimits} import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.{Partitioner, SparkConf, SparkContext, SparkEnv, TaskContextImpl} @@ -41,9 +40,7 @@ import org.apache.spark.sql.types.{BinaryType, IntegerType, MetadataBuilder, Str import org.apache.spark.unsafe.UnsafeAlignedOffset import org.apache.spark.util.Utils -class CometDiskBlockWriterSuite extends AnyFunSuite with TimeLimits { - - private implicit val signaler: Signaler = ThreadSignaler +class CometDiskBlockWriterSuite extends AnyFunSuite { private val schema = StructType(Seq(StructField("a", BinaryType))) private val pageSize: Long = 256 * 1024 From 97217c9cbbf943ba1b161a7586cc698f49caefd1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 21 Sep 2026 13:15:43 -0600 Subject: [PATCH 4/4] fix: reconcile the driver plugin comment with on-heap mode having no pool #6054 landed on main and its comment above warnIfExecutorMemoryOverheadUnset named spark.comet.memoryOverhead, which this branch removes, and said the share operators reserve is charged against a memory pool, which is now true in off-heap mode only. --- spark/src/main/scala/org/apache/spark/Plugins.scala | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index ca153d7aca3..ed60deaaabb 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -150,13 +150,14 @@ object CometDriverPlugin extends Logging { } // Comet's native allocations are made by the Rust global allocator and live in the native heap. - // The share that operators reserve is charged against a memory pool, but everything else -- - // expression kernels and Arrow array builders, decompression buffers, Parquet reader structures, - // object store buffers, the tokio runtime, allocator overhead -- is covered by no budget at all, - // and neither is Comet's JVM-side Arrow allocator. The only slack the executor container has for - // that is spark.executor.memoryOverhead, which the JVM's own non-heap usage already draws on. + // In off-heap mode the share that operators reserve is charged against a memory pool, but + // everything else -- expression kernels and Arrow array builders, decompression buffers, Parquet + // reader structures, object store buffers, the tokio runtime, allocator overhead -- is covered by + // no budget at all, and neither is Comet's JVM-side Arrow allocator. In on-heap mode the pool is + // unbounded and nothing is bounded at all. The only slack the executor container has for that is + // spark.executor.memoryOverhead, which the JVM's own non-heap usage already draws on. // - // Comet used to add spark.comet.memoryOverhead to it here, but a driver plugin cannot: on Spark + // Comet used to add an overhead of its own to it here, but a driver plugin cannot: on Spark // 3.4, 3.5 and 4.0, SparkContext builds the default ResourceProfile before it creates the plugin // container, and the cluster managers size executors from that profile rather than re-reading // the conf, so the new value never reached the container. Say so while the application is still