Add config to limit maximum RMM pool size

docs/configs.md

@@ -29,8 +29,9 @@ scala> spark.conf.set("spark.rapids.sql.incompatibleOps.enabled", true)
 
 Name | Description | Default Value
 -----|-------------|--------------
-<a name="memory.gpu.allocFraction"></a>spark.rapids.memory.gpu.allocFraction|The fraction of total GPU memory that should be initially allocated for pooled memory. Extra memory will be allocated as needed, but it may result in more fragmentation.|0.9
+<a name="memory.gpu.allocFraction"></a>spark.rapids.memory.gpu.allocFraction|The fraction of total GPU memory that should be initially allocated for pooled memory. Extra memory will be allocated as needed, but it may result in more fragmentation. This must be less than or equal to the maximum limit configured via spark.rapids.memory.gpu.maxAllocFraction.|0.9
 <a name="memory.gpu.debug"></a>spark.rapids.memory.gpu.debug|Provides a log of GPU memory allocations and frees. If set to STDOUT or STDERR the logging will go there. Setting it to NONE disables logging. All other values are reserved for possible future expansion and in the mean time will disable logging.|NONE
+<a name="memory.gpu.maxAllocFraction"></a>spark.rapids.memory.gpu.maxAllocFraction|The fraction of total GPU memory that limits the maximum size of the RMM pool. The value must be greater than or equal to the setting for spark.rapids.memory.gpu.allocFraction.|1.0
 <a name="memory.gpu.pooling.enabled"></a>spark.rapids.memory.gpu.pooling.enabled|Should RMM act as a pooling allocator for GPU memory, or should it just pass through to CUDA memory allocation directly.|true
 <a name="memory.host.spillStorageSize"></a>spark.rapids.memory.host.spillStorageSize|Amount of off-heap host memory to use for buffering spilled GPU data before spilling to local disk|1073741824
 <a name="memory.pinnedPool.size"></a>spark.rapids.memory.pinnedPool.size|The size of the pinned memory pool in bytes unless otherwise specified. Use 0 to disable the pool.|0

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDeviceManager.scala

@@ -126,6 +126,11 @@ object GpuDeviceManager extends Logging {
     initializeMemory(addr)
   }
 
+  def shutdown(): Unit = synchronized {
+    Rmm.shutdown()
+    singletonMemoryInitialized = false
+  }
+
   def getResourcesFromTaskContext: Map[String, ResourceInformation] = {
     val tc = TaskContext.get()
     if (tc == null) Map.empty[String, ResourceInformation] else tc.resources()
@@ -158,6 +163,18 @@ object GpuDeviceManager extends Logging {
         logWarning(s"Initial RMM allocation(${initialAllocation / 1024 / 1024.0} MB) is " +
           s"larger than free memory(${info.free / 1024 / 1024.0} MB)")
       }
+      val maxAllocation = if (conf.rmmAllocMaxFraction < 1) {
+        (conf.rmmAllocMaxFraction * info.total).toLong
+      } else {
+        // Do not attempt to enforce any artificial pool limit based on queried GPU memory size
+        // if config indicates all GPU memory should be used.
+        0
+      }
+      if (maxAllocation > 0 && maxAllocation < initialAllocation) {
+        throw new IllegalArgumentException(s"${RapidsConf.RMM_ALLOC_MAX_FRACTION} " +
+            s"configured as ${conf.rmmAllocMaxFraction} which is less than the " +
+            s"${RapidsConf.RMM_ALLOC_FRACTION} setting of ${conf.rmmAllocFraction}")
+      }
       var init = RmmAllocationMode.CUDA_DEFAULT
       val features = ArrayBuffer[String]()
       if (conf.isPooledMemEnabled) {
@@ -189,7 +206,7 @@ object GpuDeviceManager extends Logging {
 
       try {
         Cuda.setDevice(gpuId)
-        Rmm.initialize(init, logConf, initialAllocation)
+        Rmm.initialize(init, logConf, initialAllocation, maxAllocation)
         GpuShuffleEnv.init(conf, info)
       } catch {
         case e: Exception => logError("Could not initialize RMM", e)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala

@@ -252,14 +252,24 @@ object RapidsConf {
     .stringConf
     .createWithDefault("NONE")
 
+  private val RMM_ALLOC_MAX_FRACTION_KEY = "spark.rapids.memory.gpu.maxAllocFraction"
+
   val RMM_ALLOC_FRACTION = conf("spark.rapids.memory.gpu.allocFraction")
     .doc("The fraction of total GPU memory that should be initially allocated " +
       "for pooled memory. Extra memory will be allocated as needed, but it may " +
-      "result in more fragmentation.")
+      "result in more fragmentation. This must be less than or equal to the maximum limit " +
+      s"configured via $RMM_ALLOC_MAX_FRACTION_KEY.")
     .doubleConf
     .checkValue(v => v >= 0 && v <= 1, "The fraction value must be in [0, 1].")
     .createWithDefault(0.9)
 
+  val RMM_ALLOC_MAX_FRACTION = conf(RMM_ALLOC_MAX_FRACTION_KEY)
+    .doc("The fraction of total GPU memory that limits the maximum size of the RMM pool. " +
+        s"The value must be greater than or equal to the setting for $RMM_ALLOC_FRACTION.")
+    .doubleConf
+    .checkValue(v => v >= 0 && v <= 1, "The fraction value must be in [0, 1].")
+    .createWithDefault(1)
+
   val HOST_SPILL_STORAGE_SIZE = conf("spark.rapids.memory.host.spillStorageSize")
     .doc("Amount of off-heap host memory to use for buffering spilled GPU data " +
         "before spilling to local disk")
@@ -767,6 +777,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val rmmAllocFraction: Double = get(RMM_ALLOC_FRACTION)
 
+  lazy val rmmAllocMaxFraction: Double = get(RMM_ALLOC_MAX_FRACTION)
+
   lazy val hostSpillStorageSize: Long = get(HOST_SPILL_STORAGE_SIZE)
 
   lazy val hasNans: Boolean = get(HAS_NANS)

tests/src/test/scala/com/nvidia/spark/rapids/GpuDeviceManagerSuite.scala

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids
+
+import ai.rapids.cudf.{Cuda, DeviceMemoryBuffer}
+import org.scalatest.FunSuite
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
+
+class GpuDeviceManagerSuite extends FunSuite with Arm {
+
+  test("RMM pool limit") {
+    SparkSession.getActiveSession.foreach(_.close())
+    GpuDeviceManager.shutdown()
+    val totalGpuSize = Cuda.memGetInfo().total
+    val initPoolFraction = 0.1
+    val maxPoolFraction = 0.2
+    val conf = new SparkConf()
+        .set(RapidsConf.POOLED_MEM.key, "true")
+        .set(RapidsConf.RMM_ALLOC_FRACTION.key, initPoolFraction.toString)
+        .set(RapidsConf.RMM_ALLOC_MAX_FRACTION.key, maxPoolFraction.toString)
+    try {
+      TestUtils.withGpuSparkSession(conf) { _ =>
+        val initPoolSize = (totalGpuSize * initPoolFraction).toLong
+        val allocSize = Math.max(initPoolSize - 1024 * 1024, 0)
+        // initial allocation should fit within initial pool size
+        withResource(DeviceMemoryBuffer.allocate(allocSize)) { _ =>
+          // this should grow the pool
+          withResource(DeviceMemoryBuffer.allocate(allocSize)) { _ =>
+            assertThrows[OutOfMemoryError] {
+              // this should exceed the specified pool limit
+              DeviceMemoryBuffer.allocate(allocSize).close()
+            }
+          }
+        }
+      }
+    } finally {
+      GpuDeviceManager.shutdown()
+    }
+  }
+}