zhztheplayer commented on code in PR #9115:
URL: https://github.com/apache/incubator-gluten/pull/9115#discussion_r2012125644
##########
backends-velox/src/main/scala/org/apache/spark/memory/GlobalOffHeapMemory.scala:
##########
@@ -48,28 +53,72 @@ object GlobalOffHeapMemory {
f
}
- def acquire(numBytes: Long): Boolean = {
- memoryManager().acquireStorageMemory(
- BlockId(s"test_${UUID.randomUUID()}"),
- numBytes,
- MemoryMode.OFF_HEAP)
+ def acquire(numBytes: Long): Unit = memoryManagerOption().foreach {
+ mm =>
+ val succeeded =
+ mm.acquireStorageMemory(
+ BlockId(s"test_${UUID.randomUUID()}"),
+ numBytes,
+ MemoryMode.OFF_HEAP)
+
+ if (succeeded) {
+ recorder.inc(numBytes)
+ return
+ }
+
+ // Throw OOM.
+ val offHeapMemoryTotal =
+ mm.maxOffHeapStorageMemory + mm.offHeapExecutionMemoryUsed
+ throw new GlutenException(
+ s"Spark off-heap memory is exhausted." +
+ s" Storage: ${mm.offHeapStorageMemoryUsed} / $offHeapMemoryTotal," +
+ s" execution: ${mm.offHeapExecutionMemoryUsed} /
$offHeapMemoryTotal")
+ }
+
+ def release(numBytes: Long): Unit = memoryManagerOption().foreach {
+ mm =>
+ mm.releaseStorageMemory(numBytes, MemoryMode.OFF_HEAP)
+ recorder.inc(-numBytes)
+ }
+
+ def currentBytes(): Long = {
+ recorder.current()
}
- def release(numBytes: Long): Unit = {
- memoryManager().releaseStorageMemory(numBytes, MemoryMode.OFF_HEAP)
+ def newReservationListener(): ReservationListener = {
+ new ReservationListener {
+ private val recorder: MemoryUsageRecorder = new
SimpleMemoryUsageRecorder()
+
+ override def reserve(size: Long): Long = {
+ acquire(size)
+ recorder.inc(size)
+ size
+ }
+
+ override def unreserve(size: Long): Long = {
+ release(size)
+ recorder.inc(-size)
+ size
+ }
+
+ override def getUsedBytes: Long = {
+ recorder.current()
+ }
+ }
}
- private def memoryManager(): MemoryManager = {
+ private def memoryManagerOption(): Option[MemoryManager] = {
val env = SparkEnv.get
if (env != null) {
- return env.memoryManager
+ return Some(env.memoryManager)
}
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]
+ return
Some(FIELD_MEMORY_MANAGER.get(tc.taskMemoryManager()).asInstanceOf[MemoryManager])
}
- throw new GlutenException(
+ logWarning(
Review Comment:
I prefer throwing here but it will lead to crash issues so far when the
executor process is shutting down. So temporarily we can rely on warning
messages.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
