This is an automated email from the ASF dual-hosted git repository.
hongze pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new ee99bf0082 [VL] Acquire off-heap global memory via the new API for
off-heap broadcast exchange (#9075)
ee99bf0082 is described below
commit ee99bf008204296958611966042dfaa8bdda072d
Author: Hongze Zhang <[email protected]>
AuthorDate: Fri Mar 21 19:19:05 2025 +0000
[VL] Acquire off-heap global memory via the new API for off-heap broadcast
exchange (#9075)
---
.../apache/spark/memory/GlobalOffHeapMemory.scala | 22 +++++++++++++--
.../execution/unsafe/UnsafeBytesBufferArray.scala | 32 +++++++++++++---------
.../unsafe/UnsafeColumnarBuildSideRelation.scala | 17 +++---------
.../spark/memory/GlobalOffHeapMemorySuite.scala | 2 +-
.../UnsafeColumnarBuildSideRelationTest.scala | 7 +----
5 files changed, 45 insertions(+), 35 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/spark/memory/GlobalOffHeapMemory.scala
b/backends-velox/src/main/scala/org/apache/spark/memory/GlobalOffHeapMemory.scala
index cfaf797869..3e25fdfc74 100644
---
a/backends-velox/src/main/scala/org/apache/spark/memory/GlobalOffHeapMemory.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/memory/GlobalOffHeapMemory.scala
@@ -24,9 +24,26 @@ import org.apache.spark.storage.BlockId
import java.lang.reflect.Field
import java.util.UUID
+/**
+ * API #acuqire is for reserving some global off-heap memory from Spark memory
manager. Once
+ * reserved, Spark tasks will have less off-heap memory to use because of the
reservation.
+ *
+ * Note the API #acuqire doesn't trigger spills on Spark tasks although OOM
may be encountered.
+ *
+ * The utility internally relies on the Spark storage memory pool. As Spark
doesn't expect trait
+ * BlockId to be extended by user, TestBlockId is chosen for the storage
memory reservations.
+ */
object GlobalOffHeapMemory {
private val FIELD_MEMORY_MANAGER: Field = {
- val f = classOf[TaskMemoryManager].getDeclaredField("memoryManager")
+ val f =
+ try {
+ classOf[TaskMemoryManager].getDeclaredField("memoryManager")
+ } catch {
+ case e: Exception =>
+ throw new GlutenException(
+ "Unable to find field TaskMemoryManager#memoryManager via
reflection",
+ e)
+ }
f.setAccessible(true)
f
}
@@ -38,7 +55,7 @@ object GlobalOffHeapMemory {
MemoryMode.OFF_HEAP)
}
- def free(numBytes: Long): Unit = {
+ def release(numBytes: Long): Unit = {
memoryManager().releaseStorageMemory(numBytes, MemoryMode.OFF_HEAP)
}
@@ -49,6 +66,7 @@ object GlobalOffHeapMemory {
}
val tc = TaskContext.get()
if (tc != null) {
+ // This may happen in test code that mocks the task context without
booting up SparkEnv.
return
FIELD_MEMORY_MANAGER.get(tc.taskMemoryManager()).asInstanceOf[MemoryManager]
}
throw new GlutenException(
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
index aafc76ac08..ecfb777178 100644
---
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
@@ -16,11 +16,15 @@
*/
package org.apache.spark.sql.execution.unsafe
+import org.apache.gluten.exception.GlutenException
+
+import org.apache.spark.SparkEnv
import org.apache.spark.annotation.Experimental
import org.apache.spark.internal.Logging
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
+import org.apache.spark.memory.GlobalOffHeapMemory
import org.apache.spark.unsafe.Platform
import org.apache.spark.unsafe.array.LongArray
+import org.apache.spark.unsafe.memory.MemoryAllocator
/**
* Used to store broadcast variable off-heap memory for broadcast variable.
The underlying data
@@ -35,14 +39,8 @@ import org.apache.spark.unsafe.array.LongArray
*/
// scalastyle:off no.finalize
@Experimental
-case class UnsafeBytesBufferArray(
- arraySize: Int,
- bytesBufferLengths: Array[Int],
- totalBytes: Long,
- tmm: TaskMemoryManager)
- extends MemoryConsumer(tmm, MemoryMode.OFF_HEAP)
- with Logging {
-
+case class UnsafeBytesBufferArray(arraySize: Int, bytesBufferLengths:
Array[Int], totalBytes: Long)
+ extends Logging {
{
assert(
arraySize == bytesBufferLengths.length,
@@ -50,6 +48,7 @@ case class UnsafeBytesBufferArray(
"not equal to buffer lengths!")
assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be
negative!")
}
+ private val allocatedBytes = (totalBytes + 7) / 8 * 8
/**
* A single array to store all bytesBufferArray's value, it's inited once
when first time get
@@ -64,8 +63,6 @@ case class UnsafeBytesBufferArray(
bytesBufferLengths.init.scanLeft(0)(_ + _)
}
- override def spill(l: Long, memoryConsumer: MemoryConsumer): Long = 0L
-
/**
* Put bytesBuffer at specified array index.
*
@@ -79,7 +76,16 @@ case class UnsafeBytesBufferArray(
assert(bytesBuffer.length == bytesBufferLengths(index))
// first to allocate underlying long array
if (null == longArray && index == 0) {
- longArray = allocateArray((totalBytes + 7) / 8)
+ if (!GlobalOffHeapMemory.acquire(allocatedBytes)) {
+ val memoryManager = SparkEnv.get.memoryManager
+ val offHeapMemoryTotal =
+ memoryManager.maxOffHeapStorageMemory +
memoryManager.offHeapExecutionMemoryUsed
+ throw new GlutenException(
+ s"Spark off-heap memory is exhausted." +
+ s" Storage: ${memoryManager.offHeapStorageMemoryUsed} /
$offHeapMemoryTotal," +
+ s" execution: ${memoryManager.offHeapExecutionMemoryUsed} /
$offHeapMemoryTotal")
+ }
+ longArray = new
LongArray(MemoryAllocator.UNSAFE.allocate(allocatedBytes))
}
Platform.copyMemory(
@@ -132,8 +138,8 @@ case class UnsafeBytesBufferArray(
override def finalize(): Unit = {
try {
if (longArray != null) {
- freeArray(longArray)
longArray = null
+ GlobalOffHeapMemory.release(allocatedBytes)
}
} finally {
super.finalize()
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala
index 47f659d189..09ff5dabd9 100644
---
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala
@@ -25,10 +25,8 @@ import org.apache.gluten.sql.shims.SparkShimLoader
import org.apache.gluten.utils.ArrowAbiUtil
import org.apache.gluten.vectorized.{ColumnarBatchSerializerJniWrapper,
NativeColumnarToRowJniWrapper}
-import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.annotation.Experimental
import org.apache.spark.internal.Logging
-import org.apache.spark.memory.TaskMemoryManager
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
UnsafeProjection, UnsafeRow}
import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode,
IdentityBroadcastMode}
@@ -78,8 +76,7 @@ case class UnsafeColumnarBuildSideRelation(
UnsafeBytesBufferArray(
bytesBufferArray.length,
bytesBufferArray.map(_.length),
- bytesBufferArray.map(_.length.toLong).sum,
- TaskContext.get().taskMemoryManager
+ bytesBufferArray.map(_.length.toLong).sum
),
mode
)
@@ -123,15 +120,12 @@ case class UnsafeColumnarBuildSideRelation(
// scalastyle:off
/**
- * This is used in Broadcast, shared by multiple tasks, we use off-heap
memory to reduce on-heap
- * pressure Similar to
+ * We use off-heap memory to reduce on-heap pressure Similar to
*
https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala#L389-L410
*/
// scalastyle:on
- val taskMemoryManager = new TaskMemoryManager(SparkEnv.get.memoryManager,
0)
- batches =
- UnsafeBytesBufferArray(totalArraySize, bytesBufferLengths, totalBytes,
taskMemoryManager)
+ batches = UnsafeBytesBufferArray(totalArraySize, bytesBufferLengths,
totalBytes)
for (i <- 0 until totalArraySize) {
val length = bytesBufferLengths(i)
@@ -148,10 +142,7 @@ case class UnsafeColumnarBuildSideRelation(
val bytesBufferLengths = kryo.readObject(in, classOf[Array[Int]])
val totalBytes = in.readLong()
- val taskMemoryManager = new TaskMemoryManager(SparkEnv.get.memoryManager,
0)
-
- batches =
- UnsafeBytesBufferArray(totalArraySize, bytesBufferLengths, totalBytes,
taskMemoryManager)
+ batches = UnsafeBytesBufferArray(totalArraySize, bytesBufferLengths,
totalBytes)
for (i <- 0 until totalArraySize) {
val length = bytesBufferLengths(i)
diff --git
a/backends-velox/src/test/scala/org/apache/spark/memory/GlobalOffHeapMemorySuite.scala
b/backends-velox/src/test/scala/org/apache/spark/memory/GlobalOffHeapMemorySuite.scala
index 2a25d09d3b..9ac890e87a 100644
---
a/backends-velox/src/test/scala/org/apache/spark/memory/GlobalOffHeapMemorySuite.scala
+++
b/backends-velox/src/test/scala/org/apache/spark/memory/GlobalOffHeapMemorySuite.scala
@@ -95,7 +95,7 @@ class GlobalOffHeapMemorySuite extends AnyFunSuite with
BeforeAndAfterAll {
Collections.emptyMap())
Assert.assertTrue(GlobalOffHeapMemory.acquire(300))
Assert.assertEquals(100, consumer.borrow(200))
- GlobalOffHeapMemory.free(10)
+ GlobalOffHeapMemory.release(10)
Assert.assertEquals(10, consumer.borrow(50))
}
}
diff --git
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
index f47c8bd562..9aed31476e 100644
---
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
+++
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
@@ -17,7 +17,6 @@
package org.apache.spark.sql.execution.unsafe;
import org.apache.spark.{SparkConf, SparkEnv}
-import org.apache.spark.memory.{TaskMemoryManager, UnifiedMemoryManager}
import org.apache.spark.serializer.{JavaSerializer, KryoSerializer}
import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.catalyst.plans.physical.IdentityBroadcastMode
@@ -37,9 +36,6 @@ class UnsafeColumnarBuildSideRelationTest extends
SharedSparkSession {
override def beforeAll(): Unit = {
super.beforeAll()
- val taskMemoryManager = new TaskMemoryManager(
- new UnifiedMemoryManager(SparkEnv.get.conf, Long.MaxValue, Long.MaxValue
/ 2, 1),
- 0)
val a = AttributeReference("a", StringType, nullable = false, null)()
val output = Seq(a)
val totalArraySize = 1
@@ -48,8 +44,7 @@ class UnsafeColumnarBuildSideRelationTest extends
SharedSparkSession {
val bytesArray = UnsafeBytesBufferArray(
1,
perArraySize,
- 10,
- taskMemoryManager
+ 10
)
bytesArray.putBytesBuffer(0, "1234567890".getBytes())
unsafeRelWithIdentityMode = UnsafeColumnarBuildSideRelation(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
