Skip to content

Commit e1d2c11

Browse files
andygroveandygrove
authored
fix: log partial memory grants at DEBUG and drop the memory usage dump (#6269)
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.
1 parent e897f8a commit e1d2c11

3 files changed

Lines changed: 73 additions & 7 deletions

File tree

‎docs/source/contributor-guide/debugging.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,17 @@ Example log output:
221221

222222
When backtraces are enabled (see earlier section) then backtraces will be included for failed allocations.
223223

224+
Under memory pressure, Spark often grants a native memory pool less than it asked for, so Comet
225+
logs this only at `DEBUG`. The operator that needed the memory usually spills and carries on. If it
226+
cannot, the task fails with an error that says how much Spark granted and lists the largest memory
227+
consumers. To log every partial grant, set this logger to `DEBUG`, for example in
228+
`log4j2.properties`:
229+
230+
```properties
231+
logger.cometTaskMemoryManager.name = org.apache.spark.CometTaskMemoryManager
232+
logger.cometTaskMemoryManager.level = debug
233+
```
234+
224235
### Dumping native stream output with the `debug` module
225236

226237
The `native/core/src/debug/` module ships small wrappers that print every

‎spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,17 +71,24 @@ public long acquireMemory(long size) {
7171
}
7272
long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer);
7373
long newUsed = used.addAndGet(acquired);
74-
if (acquired < size) {
75-
logger.warn(
74+
// A partial grant is routine, not an error: the native pool either refuses the reservation,
75+
// which tells the operator to spill, or carries the shortfall as overcommit. A refusal that
76+
// fails the task says in its error what Spark granted and which consumers hold the most.
77+
//
78+
// Nothing here may take the TaskMemoryManager monitor, which rules out showMemoryUsage.
79+
// Another acquire of this task can hold that monitor while it waits inside Spark for memory,
80+
// and this thread holds the partial grant until native code hands it back.
81+
// getMemoryConsumptionForThisTask takes only the memory manager's monitor, which a waiting
82+
// acquire gives up.
83+
if (acquired < size && logger.isDebugEnabled()) {
84+
logger.debug(
7685
"Task {} requested {} bytes but only received {} bytes. Current allocation is {} and "
7786
+ "the total memory consumption is {} bytes.",
7887
taskAttemptId,
7988
size,
8089
acquired,
8190
newUsed,
8291
internal.getMemoryConsumptionForThisTask());
83-
// If memory manager is not able to acquire the requested size, log memory usage
84-
internal.showMemoryUsage();
8592
}
8693
return acquired;
8794
}

‎spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala‎

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,12 @@ package org.apache.spark
2121

2222
import java.util.Properties
2323

24-
import org.scalatest.funsuite.AnyFunSuite
25-
24+
import org.apache.logging.log4j.{Level, LogManager}
25+
import org.apache.logging.log4j.core.{LogEvent, LoggerContext}
2626
import org.apache.spark.executor.TaskMetrics
2727
import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager}
2828

29-
class CometTaskMemoryManagerSuite extends AnyFunSuite {
29+
class CometTaskMemoryManagerSuite extends SparkFunSuite {
3030

3131
test("native memory usage is visible to Spark's memory consumer") {
3232
withTaskMemoryManager { taskMemoryManager =>
@@ -104,6 +104,54 @@ class CometTaskMemoryManagerSuite extends AnyFunSuite {
104104
}
105105
}
106106

107+
test("partial and zero grants log nothing at INFO or above") {
108+
// Spark refuses native reservations routinely under memory pressure, each time with one of
109+
// these.
110+
withTaskMemoryManager { _ =>
111+
val manager = new CometTaskMemoryManager(1L, 0L)
112+
val events = logEvents(Level.INFO) {
113+
assert(manager.acquireMemory(768L) == 768L)
114+
assert(manager.acquireMemory(512L) == 256L)
115+
assert(manager.acquireMemory(128L) == 0L)
116+
}
117+
assert(events.isEmpty, events.map(_.getMessage.getFormattedMessage).mkString("\n"))
118+
manager.releaseMemory(1024L)
119+
}
120+
}
121+
122+
test("a partial grant is logged at DEBUG, without the task's memory usage dump") {
123+
withTaskMemoryManager { _ =>
124+
val manager = new CometTaskMemoryManager(1L, 0L)
125+
val messages = logEvents(Level.DEBUG) {
126+
assert(manager.acquireMemory(768L) == 768L)
127+
assert(manager.acquireMemory(512L) == 256L)
128+
}.map(_.getMessage.getFormattedMessage)
129+
val log = messages.mkString("\n")
130+
assert(messages.exists(_.contains("requested 512 bytes but only received 256 bytes")), log)
131+
// TaskMemoryManager.showMemoryUsage takes the task's monitor; see acquireMemory.
132+
assert(!messages.exists(_.contains("Memory used in task")), log)
133+
manager.releaseMemory(1024L)
134+
}
135+
}
136+
137+
/** The events that Comet's and Spark's task memory managers log at `level` or above. */
138+
private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = {
139+
val appender = new LogAppender("task memory manager")
140+
appender.setThreshold(level)
141+
val loggers = Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName)
142+
val context = LogManager.getContext(false).asInstanceOf[LoggerContext]
143+
val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey)
144+
try withLogAppender(appender, loggers, Some(level))(f)
145+
finally {
146+
// For a logger with no config of its own, withLogAppender adds one and never removes it.
147+
// It copies the root config's additivity, which is off, so left in place it would keep the
148+
// logger's events out of the test log for the rest of the run.
149+
unconfigured.foreach(context.getConfiguration.removeLogger)
150+
context.updateLoggers()
151+
}
152+
appender.loggingEvents.toSeq
153+
}
154+
107155
private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = {
108156
val memoryManager = new TestMemoryManager(new SparkConf())
109157
memoryManager.limit(1024)

0 commit comments

Comments
 (0)