Register default allocator for host memory

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

@@ -22,7 +22,7 @@ import scala.collection.mutable.Queue
 import ai.rapids.cudf.{Cuda, HostColumnVector, NvtxColor, Table}
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
-import com.nvidia.spark.rapids.RmmRapidsRetryIterator.splitSpillableInHalfByRows
+import com.nvidia.spark.rapids.RmmRapidsRetryIterator.{splitSpillableInHalfByRows, withRetryNoSplit}
 import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion
 import com.nvidia.spark.rapids.jni.RowConversion
 import com.nvidia.spark.rapids.shims.ShimUnaryExecNode
@@ -258,18 +258,21 @@ class ColumnarToRowIterator(batches: Iterator[ColumnarBatch],
     // perform conversion
     try {
       devCb.foreach { devCb =>
-        withResource(devCb) { _ =>
-          withResource(new NvtxWithMetrics("ColumnarToRow: batch", NvtxColor.RED, opTime)) { _ =>
-            cb = new ColumnarBatch(GpuColumnVector.extractColumns(devCb).safeMap(toHost),
-              devCb.numRows())
-            it = cb.rowIterator()
-            // In order to match the numOutputRows metric in the generated code we update
-            // numOutputRows for each batch. This is less accurate than doing it at output
-            // because it will over count the number of rows output in the case of a limit,
-            // but it is more efficient.
-            numOutputRows += cb.numRows()
+        val sDevCb = SpillableColumnarBatch(devCb, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)
+        cb = withRetryNoSplit(sDevCb) { _ =>
+          withResource(sDevCb.getColumnarBatch()) { devCb =>
+            withResource(new NvtxWithMetrics("ColumnarToRow: batch", NvtxColor.RED, opTime)) { _ =>
+              new ColumnarBatch(GpuColumnVector.extractColumns(devCb).safeMap(toHost),
+                devCb.numRows())
+            }
           }
         }
+        it = cb.rowIterator()
+        // In order to match the numOutputRows metric in the generated code we update
+        // numOutputRows for each batch. This is less accurate than doing it at output
+        // because it will over count the number of rows output in the case of a limit,
+        // but it is more efficient.
+        numOutputRows += cb.numRows()
       }
     } finally {
       // Leaving the GPU for a while: if this iterator is configured to release

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

@@ -450,10 +450,8 @@ object GpuDeviceManager extends Logging {
       logInfo(s"Initializing pinned memory pool (${pinnedSize / 1024 / 1024.0} MiB)")
       PinnedMemoryPool.initialize(pinnedSize, gpuId)
     }
-    if (nonPinnedLimit >= 0) {
-      // Host memory limits must be set after the pinned memory pool is initialized
-      HostAlloc.initialize(nonPinnedLimit)
-    }
+    // Host memory limits must be set after the pinned memory pool is initialized
+    HostAlloc.initialize(nonPinnedLimit)
   }
 
   /**

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

@@ -21,6 +21,7 @@ import scala.collection.mutable.ArrayBuffer
 import ai.rapids.cudf.{ContiguousTable, Cuda, NvtxColor, NvtxRange, Table}
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
+import com.nvidia.spark.rapids.RmmRapidsRetryIterator.withRetryNoSplit
 
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
@@ -126,7 +127,9 @@ trait GpuPartitioning extends Partitioning {
     val mightNeedToSplit = totalInputSize > GpuPartitioning.MaxCpuBatchSize
 
     val hostPartColumns = withResource(partitionColumns) { _ =>
-      partitionColumns.map(_.copyToHost())
+      withRetryNoSplit {
+        partitionColumns.safeMap(_.copyToHost())
+      }
     }
     try {
       // Leaving the GPU for a while

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

@@ -16,12 +16,12 @@
 
 package com.nvidia.spark.rapids
 
-import ai.rapids.cudf.{HostMemoryBuffer, MemoryBuffer, PinnedMemoryPool}
+import ai.rapids.cudf.{DefaultHostMemoryAllocator, HostMemoryAllocator, HostMemoryBuffer, MemoryBuffer, PinnedMemoryPool}
 import com.nvidia.spark.rapids.jni.RmmSpark
 
 import org.apache.spark.internal.Logging
 
-private class HostAlloc(nonPinnedLimit: Long) extends Logging {
+private class HostAlloc(nonPinnedLimit: Long) extends HostMemoryAllocator with Logging {
   private var currentNonPinnedAllocated: Long = 0L
   private val pinnedLimit: Long = PinnedMemoryPool.getTotalPoolSizeBytes
   // For now we are going to assume that we are the only ones calling into the pinned pool
@@ -219,6 +219,12 @@ private class HostAlloc(nonPinnedLimit: Long) extends Logging {
     }
     ret.get
   }
+
+  override def allocate(amount: Long, preferPinned: Boolean): HostMemoryBuffer =
+    alloc(amount, preferPinned)
+
+  override def allocate(amount: Long): HostMemoryBuffer =
+    alloc(amount)
 }
 
 /**
@@ -233,6 +239,7 @@ object HostAlloc {
 
   def initialize(nonPinnedLimit: Long): Unit = synchronized {
     singleton = new HostAlloc(nonPinnedLimit)
+    DefaultHostMemoryAllocator.set(singleton)
   }
 
   def tryAlloc(amount: Long, preferPinned: Boolean = true): Option[HostMemoryBuffer] = {

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

@@ -18,7 +18,7 @@ package com.nvidia.spark.rapids
 
 import ai.rapids.cudf.Table
 import com.nvidia.spark.rapids.Arm.withResource
-import com.nvidia.spark.rapids.jni.{GpuSplitAndRetryOOM, RmmSpark}
+import com.nvidia.spark.rapids.jni.{CpuSplitAndRetryOOM, RmmSpark}
 import org.mockito.ArgumentMatchers.any
 import org.mockito.Mockito.{doAnswer, spy, times, verify}
 import org.mockito.invocation.InvocationOnMock
@@ -97,7 +97,7 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
           TestUtils.compareBatches(expected, devBatch)
         }
       }
-      assertResult(5)(getAndResetNumRetryThrowCurrentTask)
+      assertResult(6)(getAndResetNumRetryThrowCurrentTask)
       assert(!myIter.hasNext)
       assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
       // This is my wrap around of checking that we did retry the last part
@@ -141,7 +141,7 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
           TestUtils.compareBatches(expected, devBatch)
         }
       }
-      assertResult(5)(getAndResetNumRetryThrowCurrentTask)
+      assertResult(6)(getAndResetNumRetryThrowCurrentTask)
       assert(!myIter.hasNext)
       assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
       // This is my wrap around of checking that we did retry the last part
@@ -164,7 +164,7 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
         ctriter, schema, TargetSize(1),
         NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
       RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId)
-      assertThrows[GpuSplitAndRetryOOM] {
+      assertThrows[CpuSplitAndRetryOOM] {
         myIter.next()
       }
       assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)

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

@@ -23,7 +23,7 @@ import com.nvidia.spark.rapids.GpuColumnVector.GpuColumnarBatchBuilder
 import org.apache.spark.sql.types.{BinaryType, StringType, StructField, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
-class GpuColumnarToRowSuite extends SparkQueryCompareTestSuite {
+class GpuColumnarToRowSuite extends RmmSparkRetrySuiteBase {
   test("iterate past empty input batches") {
     val batchIter: Iterator[ColumnarBatch] = new Iterator[ColumnarBatch] {
       private[this] var batchCount = 0