You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/09/29 23:13:41 UTC

(datafusion-comet) branch branch-1.0 updated: fix: log partial memory grants at DEBUG and drop the memory usage dump (#6269) (#6419)

This is an automated email from the ASF dual-hosted git repository.

andygrove pushed a commit to branch branch-1.0
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git


The following commit(s) were added to refs/heads/branch-1.0 by this push:
     new efe6b03269 fix: log partial memory grants at DEBUG and drop the memory usage dump (#6269) (#6419)
efe6b03269 is described below

commit efe6b032699f000b9a1211542c165ad6db733ed9
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Sep 29 17:13:35 2026 -0600

    fix: log partial memory grants at DEBUG and drop the memory usage dump (#6269) (#6419)
    
    CometTaskMemoryManager.acquireMemory logged a warning, then called
    TaskMemoryManager.showMemoryUsage, every time Spark granted less memory than a
    native pool asked for. A partial grant is how a native operator learns to spill,
    so a query that spills logged hundreds of them. Log the partial grant at DEBUG
    instead. A refusal that fails the task already says in its error what Spark
    granted and lists the pool's top consumers.
    
    Drop the dump rather than moving it to DEBUG. showMemoryUsage takes the
    TaskMemoryManager monitor. Another acquire of the same task can hold that
    monitor while it waits inside Spark for memory, and this thread still holds its
    partial grant. The task then hangs until some other task frees memory.
    greedy_unified can reach this on main, because it calls Spark without a lock.
    sunchao found the cycle while reviewing #5613.
    
    Closes #6257.
    
    Adapted for branch-1.0: CometTaskMemoryManagerSuite does not exist here,
    because main added it in #5408. It is added with only this fix's two tests
    and the helpers they use, and listed in both PR workflows. The suite's other
    tests on main check a NativeMemoryConsumer.getUsed override from #5408.
    
    (cherry picked from commit e1d2c11729c2fc60a5def4e87bb17e5b28df2a29)
---
 .github/workflows/pr_build_linux.yml               |   1 +
 .github/workflows/pr_build_macos.yml               |   1 +
 docs/source/contributor-guide/debugging.md         |  11 +++
 .../org/apache/spark/CometTaskMemoryManager.java   |  15 ++-
 .../apache/spark/CometTaskMemoryManagerSuite.scala | 108 +++++++++++++++++++++
 5 files changed, 132 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 87683b57df..d42f3226dd 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -350,6 +350,7 @@ jobs:
               org.apache.comet.CometSetOpWithGroupBySuite
               org.apache.comet.CometSparkSessionExtensionsSuite
               org.apache.spark.CometPluginsSuite
+              org.apache.spark.CometTaskMemoryManagerSuite
               org.apache.spark.CometPluginsDefaultSuite
               org.apache.spark.CometPluginsNonOverrideSuite
               org.apache.spark.CometPluginsUnifiedModeOverrideSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index d775b0e86e..dfa337c2a3 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -166,6 +166,7 @@ jobs:
               org.apache.comet.CometSetOpWithGroupBySuite
               org.apache.comet.CometSparkSessionExtensionsSuite
               org.apache.spark.CometPluginsSuite
+              org.apache.spark.CometTaskMemoryManagerSuite
               org.apache.spark.CometPluginsDefaultSuite
               org.apache.spark.CometPluginsNonOverrideSuite
               org.apache.spark.CometPluginsUnifiedModeOverrideSuite
diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md
index ce988cac21..94458e0381 100644
--- a/docs/source/contributor-guide/debugging.md
+++ b/docs/source/contributor-guide/debugging.md
@@ -221,6 +221,17 @@ Example log output:
 
 When backtraces are enabled (see earlier section) then backtraces will be included for failed allocations.
 
+Under memory pressure, Spark often grants a native memory pool less than it asked for, so Comet
+logs this only at `DEBUG`. The operator that needed the memory usually spills and carries on. If it
+cannot, the task fails with an error that says how much Spark granted and lists the largest memory
+consumers. To log every partial grant, set this logger to `DEBUG`, for example in
+`log4j2.properties`:
+
+```properties
+logger.cometTaskMemoryManager.name = org.apache.spark.CometTaskMemoryManager
+logger.cometTaskMemoryManager.level = debug
+```
+
 ### Dumping native stream output with the `debug` module
 
 The `native/core/src/debug/` module ships small wrappers that print every
diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
index 3f344da68f..1d8108f866 100644
--- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
+++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
@@ -61,8 +61,17 @@ public class CometTaskMemoryManager {
     }
     long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer);
     long newUsed = used.addAndGet(acquired);
-    if (acquired < size) {
-      logger.warn(
+    // A partial grant is routine, not an error: the native pool either refuses the reservation,
+    // which tells the operator to spill, or carries the shortfall as overcommit. A refusal that
+    // fails the task says in its error what Spark granted and which consumers hold the most.
+    //
+    // Nothing here may take the TaskMemoryManager monitor, which rules out showMemoryUsage.
+    // Another acquire of this task can hold that monitor while it waits inside Spark for memory,
+    // and this thread holds the partial grant until native code hands it back.
+    // getMemoryConsumptionForThisTask takes only the memory manager's monitor, which a waiting
+    // acquire gives up.
+    if (acquired < size && logger.isDebugEnabled()) {
+      logger.debug(
           "Task {} requested {} bytes but only received {} bytes. Current allocation is {} and "
               + "the total memory consumption is {} bytes.",
           taskAttemptId,
@@ -70,8 +79,6 @@ public class CometTaskMemoryManager {
           acquired,
           newUsed,
           internal.getMemoryConsumptionForThisTask());
-      // If memory manager is not able to acquire the requested size, log memory usage
-      internal.showMemoryUsage();
     }
     return acquired;
   }
diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
new file mode 100644
index 0000000000..88bc52af1a
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
@@ -0,0 +1,108 @@
+/*
+ * 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
+
+import java.util.Properties
+
+import org.apache.logging.log4j.{Level, LogManager}
+import org.apache.logging.log4j.core.{LogEvent, LoggerContext}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+
+class CometTaskMemoryManagerSuite extends SparkFunSuite {
+
+  test("partial and zero grants log nothing at INFO or above") {
+    // Spark refuses native reservations routinely under memory pressure, each time with one of
+    // these.
+    withTaskMemoryManager { _ =>
+      val manager = new CometTaskMemoryManager(1L, 0L)
+      val events = logEvents(Level.INFO) {
+        assert(manager.acquireMemory(768L) == 768L)
+        assert(manager.acquireMemory(512L) == 256L)
+        assert(manager.acquireMemory(128L) == 0L)
+      }
+      assert(events.isEmpty, events.map(_.getMessage.getFormattedMessage).mkString("\n"))
+      manager.releaseMemory(1024L)
+    }
+  }
+
+  test("a partial grant is logged at DEBUG, without the task's memory usage dump") {
+    withTaskMemoryManager { _ =>
+      val manager = new CometTaskMemoryManager(1L, 0L)
+      val messages = logEvents(Level.DEBUG) {
+        assert(manager.acquireMemory(768L) == 768L)
+        assert(manager.acquireMemory(512L) == 256L)
+      }.map(_.getMessage.getFormattedMessage)
+      val log = messages.mkString("\n")
+      assert(messages.exists(_.contains("requested 512 bytes but only received 256 bytes")), log)
+      // TaskMemoryManager.showMemoryUsage takes the task's monitor; see acquireMemory.
+      assert(!messages.exists(_.contains("Memory used in task")), log)
+      manager.releaseMemory(1024L)
+    }
+  }
+
+  /** The events that Comet's and Spark's task memory managers log at `level` or above. */
+  private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = {
+    val appender = new LogAppender("task memory manager")
+    appender.setThreshold(level)
+    val loggers = Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName)
+    val context = LogManager.getContext(false).asInstanceOf[LoggerContext]
+    val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey)
+    try withLogAppender(appender, loggers, Some(level))(f)
+    finally {
+      // For a logger with no config of its own, withLogAppender adds one and never removes it.
+      // It copies the root config's additivity, which is off, so left in place it would keep the
+      // logger's events out of the test log for the rest of the run.
+      unconfigured.foreach(context.getConfiguration.removeLogger)
+      context.updateLoggers()
+    }
+    appender.loggingEvents.toSeq
+  }
+
+  private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = {
+    val memoryManager = new TestMemoryManager(new SparkConf())
+    memoryManager.limit(1024)
+    val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L)
+    val taskContext = new TaskContextImpl(
+      stageId = 0,
+      stageAttemptNumber = 0,
+      partitionId = 0,
+      numPartitions = 1,
+      taskAttemptId = 0L,
+      attemptNumber = 0,
+      taskMemoryManager = taskMemoryManager,
+      localProperties = new Properties,
+      metricsSystem = null,
+      taskMetrics = TaskMetrics.empty,
+      cpus = 1,
+      resources = Map.empty)
+
+    TaskContext.setTaskContext(taskContext)
+    try {
+      f(taskMemoryManager)
+    } finally {
+      try {
+        taskMemoryManager.cleanUpAllAllocatedMemory()
+      } finally {
+        TaskContext.unset()
+      }
+    }
+  }
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org