Spill framework refactor for better performance and extensibility [databricks]

sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVector.java

@@ -1052,6 +1052,10 @@ public final int numNulls() {
 
   public static long getTotalDeviceMemoryUsed(ColumnarBatch batch) {
     long sum = 0;
+    if (batch.numCols() == 1 && batch.column(0) instanceof GpuPackedTableColumn) {
+      // this is a special case for a packed batch
+      return ((GpuPackedTableColumn) batch.column(0)).getTableBuffer().getLength();
+    }
     if (batch.numCols() > 0) {
       if (batch.column(0) instanceof WithTableBuffer) {
         WithTableBuffer wtb = (WithTableBuffer) batch.column(0);

sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVectorFromBuffer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -113,6 +113,29 @@ public GpuColumnVectorFromBuffer(DataType type, ColumnVector cudfColumn,
     this.tableMeta = meta;
   }
 
+  public static boolean isFromBuffer(ColumnarBatch cb) {
+    if (cb.numCols() > 0) {
+      long bufferAddr = 0L;
+      boolean isSet = false;
+      for (int i = 0; i < cb.numCols(); ++i) {
+        GpuColumnVectorFromBuffer gcvfb = null;
+        if (!(cb.column(i) instanceof GpuColumnVectorFromBuffer)) {
+          return false;
+        } else {
+          gcvfb = (GpuColumnVectorFromBuffer) cb.column(i);
+          if (!isSet) {
+            bufferAddr = gcvfb.buffer.getAddress();
+            isSet = true;
+          } else if (bufferAddr != gcvfb.buffer.getAddress()) {
+            return false;
+          }
+        }
+      }
+      return true;
+    }
+    return false;
+  }
+
   /**
    * Get the underlying contiguous buffer, shared between columns of the original
    * `ContiguousTable`

sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuCompressedColumnVector.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,6 +47,15 @@ public static boolean isBatchCompressed(ColumnarBatch batch) {
     return batch.numCols() == 1 && batch.column(0) instanceof GpuCompressedColumnVector;
   }
 
+  public static ColumnarBatch incRefCounts(ColumnarBatch batch) {
+    if (!isBatchCompressed(batch)) {
+      throw new IllegalStateException(
+              "Attempted to incRefCount for a compressed batch, but the batch was not compressed.");
+    }
+    ((GpuCompressedColumnVector)batch.column(0)).buffer.incRefCount();
+    return batch;
+  }
+
   /**
    * Build a columnar batch from a compressed data buffer and specified table metadata
    * NOTE: The data remains compressed and cannot be accessed directly from the columnar batch.

sql-plugin/src/main/java/com/nvidia/spark/rapids/InternalRowToColumnarBatchIterator.java

@@ -21,19 +21,11 @@
 import java.util.NoSuchElementException;
 import java.util.Optional;
 
-import com.nvidia.spark.Retryable;
 import scala.Option;
 import scala.Tuple2;
 import scala.collection.Iterator;
 
-import ai.rapids.cudf.ColumnVector;
-import ai.rapids.cudf.DType;
-import ai.rapids.cudf.HostColumnVector;
-import ai.rapids.cudf.HostColumnVectorCore;
-import ai.rapids.cudf.HostMemoryBuffer;
-import ai.rapids.cudf.NvtxColor;
-import ai.rapids.cudf.NvtxRange;
-import ai.rapids.cudf.Table;
+import ai.rapids.cudf.*;
 import com.nvidia.spark.rapids.jni.RowConversion;
 import com.nvidia.spark.rapids.shims.CudfUnsafeRow;
 
@@ -236,8 +228,7 @@ private HostMemoryBuffer[] getHostBuffersWithRetry(
     try {
       hBuf = HostAlloc$.MODULE$.alloc((dataBytes + offsetBytes),true);
       SpillableHostBuffer sBuf = SpillableHostBuffer$.MODULE$.apply(hBuf, hBuf.getLength(),
-          SpillPriorities$.MODULE$.ACTIVE_ON_DECK_PRIORITY(),
-          RapidsBufferCatalog$.MODULE$.singleton());
+          SpillPriorities$.MODULE$.ACTIVE_ON_DECK_PRIORITY());
       hBuf = null;  // taken over by spillable host buffer
       return Tuple2.apply(sBuf, numRowsWrapper);
     } finally {

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

@@ -63,26 +63,6 @@ object RapidsPluginImplicits {
     }
   }
 
-  implicit class RapidsBufferColumn(rapidsBuffer: RapidsBuffer) {
-
-    /**
-     * safeFree: Is an implicit on RapidsBuffer class that tries to free the resource, if an
-     * Exception was thrown prior to this free, it adds the new exception to the suppressed
-     * exceptions, otherwise just throws
-     *
-     * @param e Exception which we don't want to suppress
-     */
-    def safeFree(e: Throwable = null): Unit = {
-      if (rapidsBuffer != null) {
-        try {
-          rapidsBuffer.free()
-        } catch {
-          case suppressed: Throwable if e != null => e.addSuppressed(suppressed)
-        }
-      }
-    }
-  }
-
   implicit class AutoCloseableSeq[A <: AutoCloseable](val in: collection.SeqLike[A, _]) {
     /**
      * safeClose: Is an implicit on a sequence of AutoCloseable classes that tries to close each
@@ -111,46 +91,12 @@ object RapidsPluginImplicits {
     }
   }
 
-  implicit class RapidsBufferSeq[A <: RapidsBuffer](val in: collection.SeqLike[A, _]) {
-    /**
-     * safeFree: Is an implicit on a sequence of RapidsBuffer classes that tries to free each
-     * element of the sequence, even if prior free calls fail. In case of failure in any of the
-     * free calls, an Exception is thrown containing the suppressed exceptions (getSuppressed),
-     * if any.
-     */
-    def safeFree(error: Throwable = null): Unit = if (in != null) {
-      var freeException: Throwable = null
-      in.foreach { element =>
-        if (element != null) {
-          try {
-            element.free()
-          } catch {
-            case e: Throwable if error != null => error.addSuppressed(e)
-            case e: Throwable if freeException == null => freeException = e
-            case e: Throwable => freeException.addSuppressed(e)
-          }
-        }
-      }
-      if (freeException != null) {
-        // an exception happened while we were trying to safely free
-        // resources, throw the exception to alert the caller
-        throw freeException
-      }
-    }
-  }
-
   implicit class AutoCloseableArray[A <: AutoCloseable](val in: Array[A]) {
     def safeClose(e: Throwable = null): Unit = if (in != null) {
       in.toSeq.safeClose(e)
     }
   }
 
-  implicit class RapidsBufferArray[A <: RapidsBuffer](val in: Array[A]) {
-    def safeFree(e: Throwable = null): Unit = if (in != null) {
-      in.toSeq.safeFree(e)
-    }
-  }
-
   class MapsSafely[A, Repr] {
     /**
      * safeMap: safeMap implementation that is leveraged by other type-specific implicits.

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

@@ -163,20 +163,6 @@ object Arm extends ArmScalaSpecificImpl {
     }
   }
 
-  /** Executes the provided code block, freeing the RapidsBuffer only if an exception occurs */
-  def freeOnExcept[T <: RapidsBuffer, V](r: T)(block: T => V): V = {
-    try {
-      block(r)
-    } catch {
-      case t: ControlThrowable =>
-        // Don't close for these cases..
-        throw t
-      case t: Throwable =>
-        r.safeFree(t)
-        throw t
-    }
-  }
-
   /** Executes the provided code block and then closes the resource */
   def withResource[T <: AutoCloseable, V](h: CloseableHolder[T])
       (block: CloseableHolder[T] => V): V = {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ import java.lang.management.ManagementFactory
 import java.util.concurrent.atomic.AtomicLong
 
 import ai.rapids.cudf.{Cuda, Rmm, RmmEventHandler}
+import com.nvidia.spark.rapids.spill.SpillableDeviceStore
 import com.sun.management.HotSpotDiagnosticMXBean
 
 import org.apache.spark.internal.Logging
@@ -34,8 +35,7 @@ import org.apache.spark.sql.rapids.execution.TrampolineUtil
  *                            depleting the device store
  */
 class DeviceMemoryEventHandler(
-    catalog: RapidsBufferCatalog,
-    store: RapidsDeviceMemoryStore,
+    store: SpillableDeviceStore,
     oomDumpDir: Option[String],
     maxFailedOOMRetries: Int) extends RmmEventHandler with Logging {
 
@@ -92,8 +92,8 @@ class DeviceMemoryEventHandler(
      * from cuDF. If we succeed, cuDF resets `retryCount`, and so the new count sent to us
      * must be <= than what we saw last, so we can reset our tracking.
      */
-    def resetIfNeeded(retryCount: Int, storeSpillableSize: Long): Unit = {
-      if (storeSpillableSize != 0 || retryCount <= retryCountLastSynced) {
+    def resetIfNeeded(retryCount: Int, couldSpill: Boolean): Unit = {
+      if (couldSpill || retryCount <= retryCountLastSynced) {
         reset()
       }
     }
@@ -114,22 +114,20 @@ class DeviceMemoryEventHandler(
       s"onAllocFailure invoked with invalid retryCount $retryCount")
 
     try {
-      val storeSize = store.currentSize
-      val storeSpillableSize = store.currentSpillableSize
-
       val attemptMsg = if (retryCount > 0) {
         s"Attempt ${retryCount}. "
       } else {
         "First attempt. "
       }
 
       val retryState = oomRetryState.get()
-      retryState.resetIfNeeded(retryCount, storeSpillableSize)
 
-      logInfo(s"Device allocation of $allocSize bytes failed, device store has " +
-        s"$storeSize total and $storeSpillableSize spillable bytes. $attemptMsg" +
-        s"Total RMM allocated is ${Rmm.getTotalBytesAllocated} bytes. ")
-      if (storeSpillableSize == 0) {
+      val amountSpilled = store.spill(allocSize)
+      retryState.resetIfNeeded(retryCount, amountSpilled > 0)
+      logInfo(s"Device allocation of $allocSize bytes failed. " +
+        s"Device store spilled $amountSpilled bytes. $attemptMsg" +
+        s"Total RMM allocated is ${Rmm.getTotalBytesAllocated} bytes.")
+      val shouldRetry = if (amountSpilled == 0) {
         if (retryState.shouldTrySynchronizing(retryCount)) {
           Cuda.deviceSynchronize()
           logWarning(s"[RETRY ${retryState.getRetriesSoFar}] " +
@@ -149,15 +147,11 @@ class DeviceMemoryEventHandler(
           false
         }
       } else {
-        val targetSize = Math.max(storeSpillableSize - allocSize, 0)
-        logDebug(s"Targeting device store size of $targetSize bytes")
-        val maybeAmountSpilled = catalog.synchronousSpill(store, targetSize, Cuda.DEFAULT_STREAM)
-        maybeAmountSpilled.foreach { amountSpilled =>
-          logInfo(s"Spilled $amountSpilled bytes from the device store")
-          TrampolineUtil.incTaskMetricsMemoryBytesSpilled(amountSpilled)
-        }
+        TrampolineUtil.incTaskMetricsMemoryBytesSpilled(amountSpilled)
         true
       }
+
+      shouldRetry
     } catch {
       case t: Throwable =>
         logError(s"Error handling allocation failure", t)

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

@@ -22,6 +22,8 @@ import scala.collection.mutable.ArrayBuffer
 import scala.util.control.NonFatal
 
 import ai.rapids.cudf._
+import com.nvidia.spark.rapids.jni.RmmSpark
+import com.nvidia.spark.rapids.spill.SpillFramework
 
 import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
@@ -169,7 +171,9 @@ object GpuDeviceManager extends Logging {
     chunkedPackMemoryResource = None
     poolSizeLimit = 0L
 
-    RapidsBufferCatalog.close()
+    SpillFramework.shutdown()
+    RmmSpark.clearEventHandler()
+    Rmm.clearEventHandler()
     GpuShuffleEnv.shutdown()
     // try to avoid segfault on RMM shutdown
     val timeout = System.nanoTime() + TimeUnit.SECONDS.toNanos(10)
@@ -278,6 +282,8 @@ object GpuDeviceManager extends Logging {
     }
   }
 
+  private var memoryEventHandler: DeviceMemoryEventHandler = _
+
   private def initializeRmm(gpuId: Int, rapidsConf: Option[RapidsConf]): Unit = {
     if (!Rmm.isInitialized) {
       val conf = rapidsConf.getOrElse(new RapidsConf(SparkEnv.get.conf))
@@ -385,8 +391,25 @@ object GpuDeviceManager extends Logging {
         }
       }
 
-      RapidsBufferCatalog.init(conf)
-      GpuShuffleEnv.init(conf, RapidsBufferCatalog.getDiskBlockManager())
+      SpillFramework.initialize(conf)
+
+      memoryEventHandler = new DeviceMemoryEventHandler(
+        SpillFramework.stores.deviceStore,
+        conf.gpuOomDumpDir,
+        conf.gpuOomMaxRetries)
+
+      if (conf.sparkRmmStateEnable) {
+        val debugLoc = if (conf.sparkRmmDebugLocation.isEmpty) {
+          null
+        } else {
+          conf.sparkRmmDebugLocation
+        }
+        RmmSpark.setEventHandler(memoryEventHandler, debugLoc)
+      } else {
+        logWarning("SparkRMM retry has been disabled")
+        Rmm.setEventHandler(memoryEventHandler)
+      }
+      GpuShuffleEnv.init(conf)
     }
   }
 

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

@@ -1110,7 +1110,7 @@ class CudfSpillableHostConcatResult(
     val hmb: HostMemoryBuffer) extends SpillableHostConcatResult {
 
   override def toBatch: ColumnarBatch = {
-    closeOnExcept(buffer.getHostBuffer()) { hostBuf =>
+    closeOnExcept(buffer.getHostBuffer) { hostBuf =>
       SerializedTableColumn.from(header, hostBuf)
     }
   }

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

@@ -123,7 +123,7 @@ trait GpuRowBasedUserDefinedFunction extends GpuExpression
       val retConverter = GpuRowToColumnConverter.getConverterForType(dataType, nullable)
       val retType = GpuColumnVector.convertFrom(dataType, nullable)
       val retRow = new GenericInternalRow(size = 1)
-      closeOnExcept(new RapidsHostColumnBuilder(retType, batch.numRows)) { builder =>
+      withResource(new RapidsHostColumnBuilder(retType, batch.numRows)) { builder =>
         /**
          * This `nullSafe` is for https://github.com/NVIDIA/spark-rapids/issues/3942.
          * And more details can be found from