Repository: apache/incubator-gluten
Pull request #9075: [VL] Acquire off-heap global memory via the new API for off-heap broadcast exchange
Author: login: zhztheplayer
Labels: name: VELOX
Created: 2025-03-20T16:12:17Z | Merged: 2025-03-21T19:19:05Z

A follow-up for #8762 to avoid creating a dummy TaskMemoryManager (which was hack) for reserving off-heap memory for off-heap broadcast exchange. 

The code now calls new API introduced in https://github.com/apache/incubator-gluten/pull/9066 to acquire for off-heap storage memory.


Diff:
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 cfaf7978692f..3e25fdfc74a2 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 aafc76ac08bf..ecfb77717871 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 47f659d189a7..09ff5dabd990 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 2a25d09d3b19..9ac890e87a01 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 f47c8bd562b8..9aed31476e53 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(


Reviews:
user: login: jinchengchenghh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T10:25:48Z
user: login: zhztheplayer
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T10:44:15Z
user: login: zhztheplayer
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T10:50:12Z
user: login: jinchengchenghh
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T11:45:48Z
user: login: jinchengchenghh
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T12:28:27Z
user: login: zhztheplayer
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T13:32:34Z
user: login: jinchengchenghh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T14:38:47Z
user: login: jinchengchenghh
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T14:38:57Z
user: login: zhztheplayer
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-21T14:59:33Z

Review comments:
diff_hunk: @@ -79,7 +75,17 @@ case class UnsafeBytesBufferArray(
     assert(bytesBuffer.length == bytesBufferLengths(index))
     // first to allocate underlying long array
     if (null == longArray && index == 0) {
-      longArray = allocateArray((totalBytes + 7) / 8)
+      val numBytes = totalBytes + 7
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: jinchengchenghh
body: Why the numBytes is not (totalBytes + 7) / 8
created_at: 2025-03-21T10:24:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -132,8 +138,8 @@ case class UnsafeBytesBufferArray(
   override def finalize(): Unit = {
     try {
       if (longArray != null) {
-        freeArray(longArray)
         longArray = null
+        GlobalOffHeapMemory.release(longArray.size())
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: jinchengchenghh
body:  call longArray.size() after set it to null
created_at: 2025-03-21T10:24:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -132,8 +138,8 @@ case class UnsafeBytesBufferArray(
   override def finalize(): Unit = {
     try {
       if (longArray != null) {
-        freeArray(longArray)
         longArray = null
+        GlobalOffHeapMemory.release(longArray.size())
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: zhztheplayer
body: Good catch. Thanks.
created_at: 2025-03-21T10:44:15Z
author_association: MEMBER
diff_hunk: @@ -79,7 +75,17 @@ case class UnsafeBytesBufferArray(
     assert(bytesBuffer.length == bytesBufferLengths(index))
     // first to allocate underlying long array
     if (null == longArray && index == 0) {
-      longArray = allocateArray((totalBytes + 7) / 8)
+      val numBytes = totalBytes + 7
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: zhztheplayer
body: `allocateArray` accepts number of long elements so number of bytes will be multiplied by 8
created_at: 2025-03-21T10:50:12Z
author_association: MEMBER
diff_hunk: @@ -48,6 +48,7 @@ case class UnsafeBytesBufferArray(arraySize: Int, bytesBufferLengths: Array[Int]
         "not equal to buffer lengths!")
     assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be negative!")
   }
+  private val allocatedBytes = totalBytes + 7;
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: jinchengchenghh
body: I think it is not totalBytes + 7, it is the mutiples of 8
created_at: 2025-03-21T12:28:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -48,6 +48,7 @@ case class UnsafeBytesBufferArray(arraySize: Int, bytesBufferLengths: Array[Int]
         "not equal to buffer lengths!")
     assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be negative!")
   }
+  private val allocatedBytes = totalBytes + 7;
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: zhztheplayer
body: I am not sure I am following? `allocatedBytes` is the number of bytes, not the number of long integers. 
created_at: 2025-03-21T13:32:34Z
author_association: MEMBER
diff_hunk: @@ -48,6 +48,7 @@ case class UnsafeBytesBufferArray(arraySize: Int, bytesBufferLengths: Array[Int]
         "not equal to buffer lengths!")
     assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be negative!")
   }
+  private val allocatedBytes = totalBytes + 7;
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: jinchengchenghh
body: If better change it in a follow up PR to `totalBytes + 7）/8 * 8`
created_at: 2025-03-21T14:38:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -48,6 +48,7 @@ case class UnsafeBytesBufferArray(arraySize: Int, bytesBufferLengths: Array[Int]
         "not equal to buffer lengths!")
     assert(totalBytes >= 0, "Unsafe buffer array total bytes can't be negative!")
   }
+  private val allocatedBytes = totalBytes + 7;
path: backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeBytesBufferArray.scala
user: login: zhztheplayer
body: No problem. I am updating this PR.
created_at: 2025-03-21T14:59:32Z
author_association: MEMBER

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-20T16:12:35Z
body: <!--
  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/LICEGJ-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.
-->

Thanks for opening a pull request!

Could you open an issue for this pull request on Github Issues?

https://github.com/apache/incubator-gluten/issues

Then could you also rename ***commit message*** and ***pull request title*** in the following format?

    [GLUTEN-${ISSUES_ID}][COMPONENT]feat/fix: ${detailed message}

See also:

  * [Other pull requests](https://github.com/apache/incubator-gluten/pulls/)


author_association: NONE
user: login: zhztheplayer
created_at: 2025-03-20T16:29:43Z
body: cc @JkSelf ptal

The new API tends to help with broadcast's off-heap memory management.
author_association: MEMBER
user: login: zhztheplayer
created_at: 2025-03-21T09:48:00Z
body: cc @zjuwangg
author_association: MEMBER
user: login: zjuwangg
created_at: 2025-03-26T06:27:28Z
body: > cc @zjuwangg

Sorry for the missed message. Glad to see such improvement. I'll investigate this later.
author_association: CONTRIBUTOR