This is an automated email from the ASF dual-hosted git repository.
jackylee-ch pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 4876f4e551 [GLUTEN-13093][CORE] Fix two dynamic off-heap sizing memory
bugs (#13094)
4876f4e551 is described below
commit 4876f4e551b21a2db4995c0d7f4023ca6481d88f
Author: YangJie <[email protected]>
AuthorDate: Wed Oct 21 03:09:20 2026 -0400
[GLUTEN-13093][CORE] Fix two dynamic off-heap sizing memory bugs (#13094)
---
.../DynamicOffHeapSizingMemoryTarget.java | 21 ++++++--
.../org/apache/spark/memory/SparkMemoryUtil.scala | 23 +++++++--
.../DynamicOffHeapSizingMemoryTargetTest.java | 41 +++++++++++++++
.../gluten/GlutenDynamicOffHeapSizingSuite.scala | 59 +++++++++++++++++++++-
4 files changed, 136 insertions(+), 8 deletions(-)
diff --git
a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java
b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java
index ed5be45e8b..1341f18984 100644
---
a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java
+++
b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java
@@ -385,6 +385,7 @@ public class DynamicOffHeapSizingMemoryTarget implements
MemoryTarget, KnownName
public static long shrinkOnHeapMemory(long totalMemory, long freeMemory,
boolean isAsyncGc) {
boolean updateMaxHeapFreeRatio = false;
+ boolean updateMinHeapFreeRatio = false;
Object hotSpotBean = null;
String maxHeapFreeRatioName = "MaxHeapFreeRatio";
String minHeapFreeRatioName = "MinHeapFreeRatio";
@@ -401,6 +402,7 @@ public class DynamicOffHeapSizingMemoryTarget implements
MemoryTarget, KnownName
if (newValue < ORIGINAL_MIN_HEAP_FREE_RATIO) {
// Adjust the MinHeapFreeRatio to avoid the violation of the
MaxHeapFreeRatio.
setOption.invoke(hotSpotBean, minHeapFreeRatioName,
Integer.toString(newValue));
+ updateMinHeapFreeRatio = true;
}
if (newValue < ORIGINAL_MAX_HEAP_FREE_RATIO) {
setOption.invoke(hotSpotBean, maxHeapFreeRatioName,
Integer.toString(newValue));
@@ -416,16 +418,29 @@ public class DynamicOffHeapSizingMemoryTarget implements
MemoryTarget, KnownName
"Failed to update JVM heap free ratio via HotSpotDiagnosticMXBean:
{}", e.toString());
return totalMemory;
} finally {
- // Reset the MaxHeapFreeRatio to the original values.
+ // Reset both heap free ratios to the original values. Each revert is
+ // best-effort on its own: a failing Max revert must not silently skip
+ // the Min revert, which would leave the same permanent mutation behind.
if (hotSpotBean != null && updateMaxHeapFreeRatio) {
try {
Class<?> beanClass =
Class.forName("com.sun.management.HotSpotDiagnosticMXBean");
Method setOption = beanClass.getMethod("setVMOption", String.class,
String.class);
setOption.invoke(
hotSpotBean, maxHeapFreeRatioName,
Integer.toString(ORIGINAL_MAX_HEAP_FREE_RATIO));
- LOG.info("Reverted VM flags back.");
+ LOG.info("Reverted MaxHeapFreeRatio back.");
} catch (Exception ignore) {
- // best‐effort revert
+ // best-effort revert
+ }
+ }
+ if (hotSpotBean != null && updateMinHeapFreeRatio) {
+ try {
+ Class<?> beanClass =
Class.forName("com.sun.management.HotSpotDiagnosticMXBean");
+ Method setOption = beanClass.getMethod("setVMOption", String.class,
String.class);
+ setOption.invoke(
+ hotSpotBean, minHeapFreeRatioName,
Integer.toString(ORIGINAL_MIN_HEAP_FREE_RATIO));
+ LOG.info("Reverted MinHeapFreeRatio back.");
+ } catch (Exception ignore) {
+ // best-effort revert
}
}
}
diff --git
a/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala
b/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala
index a902c448fe..372e27debf 100644
--- a/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala
+++ b/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala
@@ -16,6 +16,7 @@
*/
package org.apache.spark.memory
+import org.apache.gluten.config.GlutenCoreConfig
import org.apache.gluten.memory.memtarget._
import org.apache.gluten.memory.memtarget.spark.{RegularMemoryConsumer,
TreeMemoryConsumer}
import org.apache.gluten.proto.MemoryUsageStats
@@ -35,8 +36,9 @@ object SparkMemoryUtil {
private val mmClazz = classOf[MemoryManager]
private val smpField = mmClazz.getDeclaredField("offHeapStorageMemoryPool")
private val empField = mmClazz.getDeclaredField("offHeapExecutionMemoryPool")
- smpField.setAccessible(true)
- empField.setAccessible(true)
+ private val onHeapSmpField =
mmClazz.getDeclaredField("onHeapStorageMemoryPool")
+ private val onHeapEmpField =
mmClazz.getDeclaredField("onHeapExecutionMemoryPool")
+ Seq(smpField, empField, onHeapSmpField,
onHeapEmpField).foreach(_.setAccessible(true))
private val tmmClazz = classOf[TaskMemoryManager]
private val consumersField = tmmClazz.getDeclaredField("consumers")
@@ -51,8 +53,21 @@ object SparkMemoryUtil {
// We assume storage memory can be fully transferred to execution memory so
far
def getCurrentAvailableOffHeapMemory: Long = {
val mm = SparkEnv.get.memoryManager
- val smp = smpField.get(mm).asInstanceOf[StorageMemoryPool]
- val emp = empField.get(mm).asInstanceOf[ExecutionMemoryPool]
+ // With dynamic off-heap sizing enabled, Gluten's global reservations are
+ // charged to the ON-heap pools (see GlobalOffHeapMemoryTarget), so the
+ // available figure must be read from the same pools. Read the flag the
same
+ // way GlobalOffHeapMemoryTarget.mode does, so the two cannot pick
different
+ // pool families.
+ val (smp, emp) =
+ if (GlutenCoreConfig.get.dynamicOffHeapSizingEnabled) {
+ (
+ onHeapSmpField.get(mm).asInstanceOf[StorageMemoryPool],
+ onHeapEmpField.get(mm).asInstanceOf[ExecutionMemoryPool])
+ } else {
+ (
+ smpField.get(mm).asInstanceOf[StorageMemoryPool],
+ empField.get(mm).asInstanceOf[ExecutionMemoryPool])
+ }
smp.memoryFree + emp.memoryFree
}
diff --git
a/gluten-core/src/test/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTargetTest.java
b/gluten-core/src/test/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTargetTest.java
index c2f39b5fd0..bace71dbc6 100644
---
a/gluten-core/src/test/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTargetTest.java
+++
b/gluten-core/src/test/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTargetTest.java
@@ -16,11 +16,18 @@
*/
package org.apache.gluten.memory.memtarget;
+import com.sun.management.VMOption;
import org.junit.After;
import org.junit.Assert;
+import org.junit.Assume;
import org.junit.Before;
import org.junit.Test;
+import javax.management.MBeanServer;
+
+import java.lang.management.ManagementFactory;
+import java.lang.reflect.Method;
+
/**
* Ledger-correctness tests for {@link DynamicOffHeapSizingMemoryTarget}.
*
@@ -62,6 +69,40 @@ public class DynamicOffHeapSizingMemoryTargetTest {
DynamicOffHeapSizingMemoryTarget.usedOffHeapBytesForTesting());
}
+ @Test
+ public void shrinkOnHeapMemoryRestoresHeapFreeRatios() throws Exception {
+ final MBeanServer mbs = ManagementFactory.getPlatformMBeanServer();
+ final Class<?> beanClass =
Class.forName("com.sun.management.HotSpotDiagnosticMXBean");
+ final Object bean =
+ ManagementFactory.newPlatformMXBeanProxy(
+ mbs, "com.sun.management:type=HotSpotDiagnostic", beanClass);
+ final Method getOption = beanClass.getMethod("getVMOption", String.class);
+ final String minBefore = ((VMOption) getOption.invoke(bean,
"MinHeapFreeRatio")).getValue();
+ final String maxBefore = ((VMOption) getOption.invoke(bean,
"MaxHeapFreeRatio")).getValue();
+
+ // The shrink only lowers the ratios when the original MinHeapFreeRatio is
+ // above the target (5); skip on JVMs configured below it, where there is
+ // nothing to restore. On JVMs where the VM option is not writable (e.g.
+ // JDK 8) the production code leaves the flags untouched and the assertions
+ // below hold trivially.
+ Assume.assumeTrue(
+ "JVM launched with MinHeapFreeRatio <= 5; nothing to restore",
+ Integer.parseInt(minBefore) > 5);
+
+ // isAsyncGc=true runs a single System.gc() and may set the async-GC
suspend
+ // flag for the remainder of this fork; no other test reads it.
+ DynamicOffHeapSizingMemoryTarget.shrinkOnHeapMemory(
+ Runtime.getRuntime().totalMemory(), Runtime.getRuntime().freeMemory(),
true);
+
+ final String minAfter = ((VMOption) getOption.invoke(bean,
"MinHeapFreeRatio")).getValue();
+ final String maxAfter = ((VMOption) getOption.invoke(bean,
"MaxHeapFreeRatio")).getValue();
+ // The finally block lowered the heap free ratios for the shrink and must
+ // restore both; leaving MinHeapFreeRatio lowered makes the JVM shrink its
+ // heap aggressively forever after.
+ Assert.assertEquals(minBefore, minAfter);
+ Assert.assertEquals(maxBefore, maxAfter);
+ }
+
@Test
public void borrowRecordsOnlyTheGrantedAmountForFullGrants() {
final CountingTarget wrapped = new CountingTarget(Long.MAX_VALUE);
diff --git
a/gluten-core/src/test/scala/org/apache/gluten/GlutenDynamicOffHeapSizingSuite.scala
b/gluten-core/src/test/scala/org/apache/gluten/GlutenDynamicOffHeapSizingSuite.scala
index 6095eca27b..4337f6ed53 100644
---
a/gluten-core/src/test/scala/org/apache/gluten/GlutenDynamicOffHeapSizingSuite.scala
+++
b/gluten-core/src/test/scala/org/apache/gluten/GlutenDynamicOffHeapSizingSuite.scala
@@ -19,7 +19,9 @@ package org.apache.gluten
import org.apache.gluten.component.WithDummyBackend
import org.apache.gluten.config.GlutenCoreConfig
-import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.{SparkConf, SparkContext, SparkEnv}
+import org.apache.spark.memory.SparkMemoryUtil
+import org.apache.spark.sql.internal.SQLConf
import org.scalatest.funsuite.AnyFunSuite
@@ -52,4 +54,59 @@ class GlutenDynamicOffHeapSizingSuite extends AnyFunSuite
with WithDummyBackend
sc.stop()
}
}
+
+ // Pure reflection: the pool classes are private[spark] and not referable
+ // from this package at compile time.
+ private def poolFamilyFreeSum(onHeap: Boolean): Long = {
+ val mm = SparkEnv.get.memoryManager
+ def poolFree(fieldName: String): Long = {
+ // scalastyle:off classforname
+ val mmClazz = Class.forName("org.apache.spark.memory.MemoryManager")
+ // scalastyle:on classforname
+ val field = mmClazz.getDeclaredField(fieldName)
+ field.setAccessible(true)
+ val pool = field.get(mm)
+ val memoryFree = pool.getClass.getMethod("memoryFree")
+ memoryFree.setAccessible(true)
+ memoryFree.invoke(pool).asInstanceOf[Long]
+ }
+ poolFree(if (onHeap) "onHeapStorageMemoryPool" else
"offHeapStorageMemoryPool") +
+ poolFree(if (onHeap) "onHeapExecutionMemoryPool" else
"offHeapExecutionMemoryPool")
+ }
+
+ test("available off-heap metric follows the pool family the dynamic mode
reserves from") {
+ val conf = new SparkConf(false)
+ .setAppName("GlutenDynamicOffHeapSizingSuite")
+ .set("spark.master", "local[1]")
+ .set("spark.plugins", classOf[GlutenPlugin].getName)
+ .set("spark.testing", "true")
+ .set("spark.ui.enabled", "false")
+ .set("spark.memory.offHeap.enabled", "true")
+ .set("spark.memory.offHeap.size", "512m")
+ val sc = new SparkContext(conf)
+ try {
+ // Default mode: global reservations are charged to the off-heap pools,
so
+ // the metric must equal exactly those pools' free memory.
+ assert(SparkMemoryUtil.getCurrentAvailableOffHeapMemory ==
poolFamilyFreeSum(false))
+ } finally {
+ sc.stop()
+ }
+
+ val dynamicConf =
conf.clone().set(GlutenCoreConfig.DYNAMIC_OFFHEAP_SIZING_ENABLED.key, "true")
+ val dynamicSc = new SparkContext(dynamicConf)
+ try {
+ // Dynamic mode: reservations land in the ON-heap pools
+ // (GlobalOffHeapMemoryTarget), so the metric must equal the on-heap free
+ // sum - reading the off-heap pools (the old behavior) diverges as soon
+ // as any reservation is charged. Both read the flag through
GlutenCoreConfig,
+ // i.e. the active SQLConf; with no session here, supply it in scope.
+ val sqlConf = new SQLConf
+
sqlConf.setConfString(GlutenCoreConfig.DYNAMIC_OFFHEAP_SIZING_ENABLED.key,
"true")
+ SQLConf.withExistingConf(sqlConf) {
+ assert(SparkMemoryUtil.getCurrentAvailableOffHeapMemory ==
poolFamilyFreeSum(true))
+ }
+ } finally {
+ dynamicSc.stop()
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
