diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md index 21ed92015f9..398adf7472b 100644 --- a/.ai/skills/review-comet-memory-pr/SKILL.md +++ b/.ai/skills/review-comet-memory-pr/SKILL.md @@ -115,9 +115,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 @@ -145,8 +147,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 @@ -182,7 +182,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/.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/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/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/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 93ba4e8754c..eea389205dd 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 58011c4a652..fa7e2e375c8 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -565,7 +565,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, @@ -621,7 +620,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 82f985cef95..69311698b2e 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -22,12 +22,9 @@ mod spark_memory; 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; @@ -70,26 +67,6 @@ pub(crate) fn create_memory_pool( task_attempt_id, )) }), - 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..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 @@ -19,37 +19,31 @@ 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`. + * + *

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( - 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..4ce62af6e8c --- /dev/null +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnboundedShuffleMemoryAllocator.java @@ -0,0 +1,189 @@ +/* + * 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.SparkOutOfMemoryError; +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. + * + *

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(); + + 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) { + // 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); + + 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 3ccf4341420..f91c92d1255 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -333,14 +333,6 @@ object CometConf extends ShimCometConf { .checkValue(_ >= 0, "The memory usage log interval must not be negative") .createWithDefault(TimeUnit.SECONDS.toMillis(10)) - 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") @@ -594,30 +586,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.") @@ -833,7 +801,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) @@ -847,17 +818,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 6a51cd05eaa..c9078d5385d 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -135,7 +135,6 @@ class CometExecIterator( memoryConfig.offHeapMode, memoryConfig.memoryPoolType, memoryConfig.memoryLimit, - memoryConfig.memoryLimitPerTask, taskAttemptId, taskCPUs, keyUnwrapper, @@ -608,8 +607,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) { @@ -617,28 +614,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) } } @@ -670,8 +658,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 749a3d6ecb0..55cf7ba187f 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} @@ -241,53 +240,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 d38efb12e50..664cab7959c 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..ed60deaaabb 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 @@ -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 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 c1f29c7447e..e0d8705d3b1 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 71% 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..7b2be4fbdff 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 @@ -25,26 +25,15 @@ 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.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,36 +86,25 @@ class CometBoundedShuffleMemoryAllocatorSuite extends AnyFunSuite { } } - test("getUsed is unchanged when a partial grant is rolled back") { + 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(1) + val huge = 64L * 1024 * 1024 + val page = allocator.allocate(huge) try { - intercept[SparkOutOfMemoryError] { - allocator.allocate(memoryLimit) - } - assert(allocator.getUsed === page.size()) + assert(page.size() === huge) + assert(allocator.getUsed === huge) } finally { allocator.free(page) } assert(allocator.getUsed === 0L) } - test("getUsed is unchanged when the budget is exhausted") { - val allocator = newAllocator() - val page = allocator.allocate(memoryLimit) - try { - assert(allocator.getUsed === memoryLimit) - intercept[SparkOutOfMemoryError] { - allocator.allocate(1) - } - assert(allocator.getUsed === memoryLimit) - } finally { - allocator.free(page) - } - 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 @@ -136,10 +114,16 @@ class CometBoundedShuffleMemoryAllocatorSuite 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) } 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..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 @@ -21,18 +21,15 @@ 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} @@ -43,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 @@ -81,8 +76,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 +195,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 +212,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 +274,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 +314,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 +331,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 +345,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 +411,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 +432,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 +457,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) - } }