Refactor memory pool setup

Author: andygrove

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -519,6 +519,14 @@ object CometConf extends ShimCometConf {
     .stringConf
     .createWithDefault("default")
 
+  val COMET_EXEC_MEMORY_POOL_FRACTION: ConfigEntry[Double] =
+    conf("spark.comet.exec.memoryPool.limit")
+      .doc("Fraction of pool that is available to Comet. Only applies to off-heap mode. " +
+        s"$TUNING_GUIDE.")
+      .internal()
+      .doubleConf
+      .createWithDefault(0.9)
+
   val COMET_SCAN_PREFETCH_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.scan.preFetch.enabled")
       .doc("Whether to enable pre-fetching feature of CometScan.")

native/core/src/execution/memory_pools/config.rs

@@ -59,11 +59,12 @@ pub(crate) fn parse_memory_pool_config(
     memory_limit: i64,
     memory_limit_per_task: i64,
 ) -> CometResult<MemoryPoolConfig> {
-    let pool_size = memory_limit as usize;
+    let pool_size_global = memory_limit as usize;
+    let pool_size_per_task = memory_limit_per_task as usize;
     let memory_pool_config = if off_heap_mode {
         match memory_pool_type.as_str() {
             "default" | "fair_unified" => {
-                MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size)
+                MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size_per_task)
             }
             "greedy_unified" => {
                 // the `unified` memory pool interacts with Spark's memory pool to allocate
@@ -79,7 +80,6 @@ pub(crate) fn parse_memory_pool_config(
         }
     } else {
         // Use the memory pool from DF
-        let pool_size_per_task = memory_limit_per_task as usize;
         match memory_pool_type.as_str() {
             "fair_spill_task_shared" => {
                 MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task)
@@ -88,9 +88,11 @@ pub(crate) fn parse_memory_pool_config(
                 MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared, pool_size_per_task)
             }
             "fair_spill_global" => {
-                MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal, pool_size)
+                MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal, pool_size_global)
+            }
+            "greedy_global" => {
+                MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size_global)
             }
-            "greedy_global" => MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size),
             "fair_spill" => MemoryPoolConfig::new(MemoryPoolType::FairSpill, pool_size_per_task),
             "greedy" => MemoryPoolConfig::new(MemoryPoolType::Greedy, pool_size_per_task),
             "unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0),

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -86,31 +86,13 @@ class CometExecIterator(
     val conf = SparkEnv.get.conf
     val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs
 
-    val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf)
-    val memoryLimit = if (offHeapMode) {
-      // in unified mode we share off-heap memory with Spark
-      ByteUnit.MiB.toBytes(conf.getSizeAsMb("spark.memory.offHeap.size"))
-    } else {
-      // we'll use the built-in memory pool from DF, and initializes with `memory_limit`
-      // and `memory_fraction` below.
-      CometSparkSessionExtensions.getCometMemoryOverhead(conf)
-    }
-
     // serialize Spark conf in protobuf format
     val builder = ConfigMap.newBuilder()
     conf.getAll.foreach { case (k, v) =>
       builder.putEntries(k, v)
     }
     val protobufSparkConfigs = builder.build().toByteArray
 
-    val memoryLimitPerTask = if (offHeapMode) {
-      // this per-task limit is not used in native code when using unified memory
-      // so we can skip calculating it and avoid logging irrelevant information
-      0
-    } else {
-      getMemoryLimitPerTask(conf)
-    }
-
     // Create keyUnwrapper if encryption is enabled
     val keyUnwrapper = if (encryptedFilePaths.nonEmpty) {
       val unwrapper = new CometFileKeyUnwrapper()
@@ -124,6 +106,8 @@ class CometExecIterator(
       null
     }
 
+    val memoryConfig = CometExecIterator.getMemoryConfig(conf)
+
     nativeLib.createPlan(
       id,
       cometBatchIterators,
@@ -135,10 +119,10 @@ class CometExecIterator(
       cometTaskMemoryManager,
       localDiskDirs,
       batchSize = COMET_BATCH_SIZE.get(),
-      offHeapMode,
-      memoryPoolType = COMET_EXEC_MEMORY_POOL_TYPE.get(),
-      memoryLimit,
-      memoryLimitPerTask,
+      memoryConfig.offHeapMode,
+      memoryConfig.memoryPoolType,
+      memoryConfig.memoryLimit,
+      memoryConfig.memoryLimitPerTask,
       taskAttemptId,
       debug = COMET_DEBUG_ENABLED.get(),
       explain = COMET_EXPLAIN_NATIVE_ENABLED.get(),
@@ -152,34 +136,6 @@ class CometExecIterator(
   private var currentBatch: ColumnarBatch = null
   private var closed: Boolean = false
 
-  private def getMemoryLimitPerTask(conf: SparkConf): Long = {
-    val numCores = numDriverOrExecutorCores(conf).toFloat
-    val maxMemory = CometSparkSessionExtensions.getCometMemoryOverhead(conf)
-    val coresPerTask = conf.get("spark.task.cpus", "1").toFloat
-    // example 16GB maxMemory * 16 cores with 4 cores per task results
-    // in memory_limit_per_task = 16 GB * 4 / 16 = 16 GB / 4 = 4GB
-    val limit = (maxMemory.toFloat * coresPerTask / numCores).toLong
-    logInfo(
-      s"Calculated per-task memory limit of $limit ($maxMemory * $coresPerTask / $numCores)")
-    limit
-  }
-
-  private def numDriverOrExecutorCores(conf: SparkConf): Int = {
-    def convertToInt(threads: String): Int = {
-      if (threads == "*") Runtime.getRuntime.availableProcessors() else threads.toInt
-    }
-
-    val LOCAL_N_REGEX = """local\[([0-9]+|\*)\]""".r
-    val LOCAL_N_FAILURES_REGEX = """local\[([0-9]+|\*)\s*,\s*([0-9]+)\]""".r
-    val master = conf.get("spark.master")
-    master match {
-      case "local" => 1
-      case LOCAL_N_REGEX(threads) => convertToInt(threads)
-      case LOCAL_N_FAILURES_REGEX(threads, _) => convertToInt(threads)
-      case _ => conf.get("spark.executor.cores", "1").toInt
-    }
-  }
-
   private def getNextBatch: Option[ColumnarBatch] = {
     assert(partitionIndex >= 0 && partitionIndex < numParts)
 
@@ -289,25 +245,11 @@ class CometExecIterator(
         traceMemoryUsage()
       }
 
-      // The allocator thoughts the exported ArrowArray and ArrowSchema structs are not released,
-      // so it will report:
-      // Caused by: java.lang.IllegalStateException: Memory was leaked by query.
-      // Memory leaked: (516) Allocator(ROOT) 0/516/808/9223372036854775807 (res/actual/peak/limit)
-      // Suspect this seems a false positive leak, because there is no reported memory leak at JVM
-      // when profiling. `allocator` reports a leak because it calculates the accumulated number
-      // of memory allocated for ArrowArray and ArrowSchema. But these exported ones will be
-      // released in native side later.
-      // More to clarify it. For ArrowArray and ArrowSchema, Arrow will put a release field into the
-      // memory region which is a callback function pointer (C function) that could be called to
-      // release these structs in native code too. Once we wrap their memory addresses at native
-      // side using FFI ArrowArray and ArrowSchema, and drop them later, the callback function will
-      // be called to release the memory.
-      // But at JVM, the allocator doesn't know about this fact so it still keeps the accumulated
-      // number.
-      // Tried to manually do `release` and `close` that can make the allocator happy, but it will
-      // cause JVM runtime failure.
-
-      // allocator.close()
+      val memInUse = cometTaskMemoryManager.getUsed
+      if (memInUse != 0) {
+        logWarning(s"CometExecIterator closed with non-zero memory usage : $memInUse")
+      }
+
       closed = true
     }
   }
@@ -320,5 +262,67 @@ class CometExecIterator(
     val threadId = Thread.currentThread().getId
     nativeLib.logMemoryUsage(s"task_memory_comet_$threadId", cometTaskMemory)
     nativeLib.logMemoryUsage(s"task_memory_spark_$threadId", sparkTaskMemory)
+
+    logInfo(
+      s"Task $taskAttemptId is using $totalTaskMemory bytes of memory " +
+        s"(Comet: $cometTaskMemory, Spark: $sparkTaskMemory)")
   }
 }
+
+object CometExecIterator extends Logging {
+
+  def getMemoryConfig(conf: SparkConf): MemoryConfig = {
+    val numCores = numDriverOrExecutorCores(conf).toFloat
+    val coresPerTask = conf.get("spark.task.cpus", "1").toFloat
+    // there are different paths for on-heap vs off-heap mode
+    val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf)
+    if (offHeapMode) {
+      // in off-heap mode, Comet uses unified memory management to share off-heap memory with Spark
+      val memoryLimit = ByteUnit.MiB.toBytes(conf.getSizeAsMb("spark.memory.offHeap.size"))
+      val maxMemory = memoryLimit * CometConf.COMET_EXEC_MEMORY_POOL_FRACTION.get()
+      val memoryLimitPerTask = (maxMemory.toFloat * coresPerTask / numCores).toLong
+      MemoryConfig(
+        offHeapMode,
+        memoryPoolType = COMET_EXEC_MEMORY_POOL_TYPE.get(),
+        memoryLimit,
+        memoryLimitPerTask)
+    } else {
+      // we'll use the built-in memory pool from DF, and initializes with `memory_limit`
+      // and `memory_fraction` below.
+      val memoryLimit = CometSparkSessionExtensions.getCometMemoryOverhead(conf)
+      val maxMemory = CometSparkSessionExtensions.getCometMemoryOverhead(conf)
+      // example 16GB maxMemory * 16 cores with 4 cores per task results
+      // in memory_limit_per_task = 16 GB * 4 / 16 = 16 GB / 4 = 4GB
+      val memoryLimitPerTask = (maxMemory.toFloat * coresPerTask / numCores).toLong
+      logInfo(
+        s"Calculated per-task memory limit of $memoryLimitPerTask " +
+          s"($maxMemory * $coresPerTask / $numCores)")
+      MemoryConfig(
+        offHeapMode,
+        memoryPoolType = COMET_EXEC_MEMORY_POOL_TYPE.get(),
+        memoryLimit,
+        memoryLimitPerTask)
+    }
+  }
+
+  private def numDriverOrExecutorCores(conf: SparkConf): Int = {
+    def convertToInt(threads: String): Int = {
+      if (threads == "*") Runtime.getRuntime.availableProcessors() else threads.toInt
+    }
+    val LOCAL_N_REGEX = """local\[([0-9]+|\*)\]""".r
+    val LOCAL_N_FAILURES_REGEX = """local\[([0-9]+|\*)\s*,\s*([0-9]+)\]""".r
+    val master = conf.get("spark.master")
+    master match {
+      case "local" => 1
+      case LOCAL_N_REGEX(threads) => convertToInt(threads)
+      case LOCAL_N_FAILURES_REGEX(threads, _) => convertToInt(threads)
+      case _ => conf.get("spark.executor.cores", "1").toInt
+    }
+  }
+}
+
+case class MemoryConfig(
+    offHeapMode: Boolean,
+    memoryPoolType: String,
+    memoryLimit: Long,
+    memoryLimitPerTask: Long)