batch small buffers when spilling via GDS

docs/configs.md

@@ -33,6 +33,7 @@ Name | Description | Default Value
 <a name="cloudSchemes"></a>spark.rapids.cloudSchemes|Comma separated list of additional URI schemes that are to be considered cloud based filesystems. Schemes already included: dbfs, s3, s3a, s3n, wasbs, gs. Cloud based stores generally would be total separate from the executors and likely have a higher I/O read cost. Many times the cloud filesystems also get better throughput when you have multiple readers in parallel. This is used with spark.rapids.sql.format.parquet.reader.type|None
 <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.direct.storage.spill.batchWriteBuffer.size"></a>spark.rapids.memory.gpu.direct.storage.spill.batchWriteBuffer.size|The size of the GPU memory buffer used to batch small buffers when spilling to GDS. Note that this buffer is mapped to the PCI Base Address Register (BAR) space, which may be very limited on some GPUs (e.g. the NVIDIA T4 only has 256 MiB), and it is also used by UCX bounce buffers.|8388608
 <a name="memory.gpu.direct.storage.spill.enabled"></a>spark.rapids.memory.gpu.direct.storage.spill.enabled|Should GPUDirect Storage (GDS) be used to spill GPU memory buffers directly to disk. GDS must be enabled and the directory `spark.local.dir` must support GDS. This is an experimental feature. For more information on GDS, see https://docs.nvidia.com/gpudirect-storage/.|false
 <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. Note that this limit will be reduced by the reserve memory configured in spark.rapids.memory.gpu.reserve.|1.0
 <a name="memory.gpu.oomDumpDir"></a>spark.rapids.memory.gpu.oomDumpDir|The path to a local directory where a heap dump will be created if the GPU encounters an unrecoverable out-of-memory (OOM) error. The filename will be of the form: "gpu-oom-<pid>.hprof" where <pid> is the process ID.|None

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

@@ -181,7 +181,7 @@ object RapidsBufferCatalog extends Logging with Arm {
     deviceStorage = new RapidsDeviceMemoryStore()
     val diskBlockManager = new RapidsDiskBlockManager(conf)
     if (rapidsConf.isGdsSpillEnabled) {
-      gdsStorage = new RapidsGdsStore(diskBlockManager)
+      gdsStorage = new RapidsGdsStore(diskBlockManager, rapidsConf.gdsSpillBatchWriteBufferSize)
       deviceStorage.setSpillStore(gdsStorage)
     } else {
       hostStorage = new RapidsHostMemoryStore(rapidsConf.hostSpillStorageSize)

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

@@ -369,6 +369,15 @@ object RapidsConf {
     .booleanConf
     .createWithDefault(false)
 
+  val GDS_SPILL_BATCH_WRITE_BUFFER_SIZE =
+    conf("spark.rapids.memory.gpu.direct.storage.spill.batchWriteBuffer.size")
+    .doc("The size of the GPU memory buffer used to batch small buffers when spilling to GDS. " +
+        "Note that this buffer is mapped to the PCI Base Address Register (BAR) space, which may " +
+        "be very limited on some GPUs (e.g. the NVIDIA T4 only has 256 MiB), and it is also used " +
+        "by UCX bounce buffers.")
+    .bytesConf(ByteUnit.BYTE)
+    .createWithDefault(ByteUnit.MiB.toBytes(8))
+
   val POOLED_MEM = conf("spark.rapids.memory.gpu.pooling.enabled")
     .doc("Should RMM act as a pooling allocator for GPU memory, or should it just pass " +
       "through to CUDA memory allocation directly. DEPRECATED: please use " +
@@ -1288,6 +1297,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val isGdsSpillEnabled: Boolean = get(GDS_SPILL)
 
+  lazy val gdsSpillBatchWriteBufferSize: Long = get(GDS_SPILL_BATCH_WRITE_BUFFER_SIZE)
+
   lazy val hasNans: Boolean = get(HAS_NANS)
 
   lazy val gpuTargetBatchSizeBytes: Long = get(GPU_BATCH_SIZE_BYTES)

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

@@ -18,19 +18,23 @@ package com.nvidia.spark.rapids
 
 import java.io.File
 import java.util.concurrent.ConcurrentHashMap
+import java.util.function.BiFunction
+
+import scala.collection.mutable.ArrayBuffer
 
 import ai.rapids.cudf._
 import com.nvidia.spark.rapids.StorageTier.StorageTier
 import com.nvidia.spark.rapids.format.TableMeta
 
-import org.apache.spark.sql.rapids.RapidsDiskBlockManager
+import org.apache.spark.sql.rapids.{RapidsDiskBlockManager, TempSpillBufferId}
 
 /** A buffer store using GPUDirect Storage (GDS). */
 class RapidsGdsStore(
     diskBlockManager: RapidsDiskBlockManager,
+    batchWriteBufferSize: Long,
     catalog: RapidsBufferCatalog = RapidsBufferCatalog.singleton)
     extends RapidsBufferStore(StorageTier.GDS, catalog) with Arm {
-  private[this] val sharedBufferFiles = new ConcurrentHashMap[RapidsBufferId, File]
+  private[this] val batchSpiller = new BatchSpiller()
 
   override protected def createBuffer(other: RapidsBuffer, otherBuffer: MemoryBuffer,
       stream: Cuda.Stream): RapidsBufferBase = {
@@ -39,46 +43,32 @@ class RapidsGdsStore(
         case d: DeviceMemoryBuffer => d
         case _ => throw new IllegalStateException("copying from buffer without device memory")
       }
-      val id = other.id
-      val path = if (id.canShareDiskPaths) {
-        sharedBufferFiles.computeIfAbsent(id, _ => id.getDiskPath(diskBlockManager))
-      } else {
-        id.getDiskPath(diskBlockManager)
-      }
-      // When sharing files, append to the file; otherwise, write from the beginning.
-      val fileOffset = if (id.canShareDiskPaths) {
-        // only one writer at a time for now when using shared files
-        path.synchronized {
-          CuFile.appendDeviceBufferToFile(path, deviceBuffer)
-        }
+      if (deviceBuffer.getLength < batchWriteBufferSize) {
+        batchSpiller.spill(other, deviceBuffer)
       } else {
-        CuFile.writeDeviceBufferToFile(path, 0, deviceBuffer)
-        0
+        singleShotSpill(other, deviceBuffer)
       }
-      logDebug(s"Spilled to $path $fileOffset:${other.size} via GDS")
-      new RapidsGdsBuffer(id, fileOffset, other.size, other.meta, other.getSpillPriority,
-        other.spillCallback)
     }
   }
 
-  class RapidsGdsBuffer(
-      id: RapidsBufferId,
-      val fileOffset: Long,
-      size: Long,
-      meta: TableMeta,
+  abstract class RapidsGdsBuffer(
+      override val id: RapidsBufferId,
+      override val size: Long,
+      override val meta: TableMeta,
       spillPriority: Long,
-      spillCallback: RapidsBuffer.SpillCallback)
+      override val spillCallback: RapidsBuffer.SpillCallback)
       extends RapidsBufferBase(id, size, meta, spillPriority, spillCallback) {
     override val storageTier: StorageTier = StorageTier.GDS
 
     override def getMemoryBuffer: MemoryBuffer = getDeviceMemoryBuffer
+  }
+
+  class RapidsGdsSingleShotBuffer(
+      id: RapidsBufferId, path: File, fileOffset: Long, size: Long, meta: TableMeta,
+      spillPriority: Long, spillCallback: RapidsBuffer.SpillCallback)
+      extends RapidsGdsBuffer(id, size, meta, spillPriority, spillCallback) {
 
     override def materializeMemoryBuffer: MemoryBuffer = {
-      val path = if (id.canShareDiskPaths) {
-        sharedBufferFiles.get(id)
-      } else {
-        id.getDiskPath(diskBlockManager)
-      }
       closeOnExcept(DeviceMemoryBuffer.allocate(size)) { buffer =>
         CuFile.readFileToDeviceBuffer(buffer, path, fileOffset)
         logDebug(s"Created device buffer for $path $fileOffset:$size via GDS")
@@ -88,15 +78,11 @@ class RapidsGdsStore(
 
     override def copyToMemoryBuffer(srcOffset: Long, dst: MemoryBuffer, dstOffset: Long,
         length: Long, stream: Cuda.Stream): Unit = {
-      val path = if (id.canShareDiskPaths) {
-        sharedBufferFiles.get(id)
-      } else {
-        id.getDiskPath(diskBlockManager)
-      }
       dst match {
         case dmOriginal: DeviceMemoryBuffer =>
           val dm = dmOriginal.slice(dstOffset, length)
           // TODO: switch to async API when it's released, using the passed in CUDA stream.
+          stream.sync()
           CuFile.readFileToDeviceBuffer(dm, path, fileOffset + srcOffset)
           logDebug(s"Created device buffer for $path $fileOffset:$size via GDS")
         case _ => throw new IllegalStateException(
@@ -105,15 +91,165 @@ class RapidsGdsStore(
     }
 
     override protected def releaseResources(): Unit = {
-      // Buffers that share paths must be cleaned up elsewhere
       if (id.canShareDiskPaths) {
-        sharedBufferFiles.remove(id)
+        // Buffers that share paths must be cleaned up elsewhere
       } else {
-        val path = id.getDiskPath(diskBlockManager)
         if (!path.delete() && path.exists()) {
           logWarning(s"Unable to delete GDS spill path $path")
         }
       }
     }
   }
-}
+
+  private def singleShotSpill(other: RapidsBuffer, deviceBuffer: DeviceMemoryBuffer)
+      : RapidsBufferBase = {
+    val id = other.id
+    val path = id.getDiskPath(diskBlockManager)
+    // When sharing files, append to the file; otherwise, write from the beginning.
+    val fileOffset = if (id.canShareDiskPaths) {
+      // only one writer at a time for now when using shared files
+      path.synchronized {
+        CuFile.appendDeviceBufferToFile(path, deviceBuffer)
+      }
+    } else {
+      CuFile.writeDeviceBufferToFile(path, 0, deviceBuffer)
+      0
+    }
+    logDebug(s"Spilled to $path $fileOffset:${other.size} via GDS")
+    new RapidsGdsSingleShotBuffer(id, path, fileOffset, other.size, other.meta,
+      other.getSpillPriority, other.spillCallback)
+  }
+
+  class BatchSpiller() {
+    private val blockSize = 4096
+    private[this] val spilledBuffers = new ConcurrentHashMap[File, Set[RapidsBufferId]]
+    private[this] val pendingBuffers = ArrayBuffer.empty[RapidsGdsBatchedBuffer]
+    private[this] val batchWriteBuffer = CuFileBuffer.allocate(batchWriteBufferSize, true)
+    private[this] var currentFile = TempSpillBufferId().getDiskPath(diskBlockManager)
+    private[this] var currentOffset = 0L
+
+    def spill(other: RapidsBuffer, deviceBuffer: DeviceMemoryBuffer): RapidsBufferBase =
+      this.synchronized {
+        if (deviceBuffer.getLength > batchWriteBufferSize - currentOffset) {
+          val path = currentFile.getAbsolutePath
+          withResource(new CuFileWriteHandle(path)) { handle =>
+            handle.write(batchWriteBuffer, currentOffset, 0)
+            logDebug(s"Spilled to $path 0:$currentOffset via GDS")
+          }
+          pendingBuffers.foreach(_.unsetPending())
+          pendingBuffers.clear
+          currentFile = TempSpillBufferId().getDiskPath(diskBlockManager)
+          currentOffset = 0
+        }
+
+        batchWriteBuffer.copyFromMemoryBuffer(
+          currentOffset, deviceBuffer, 0, deviceBuffer.getLength, Cuda.DEFAULT_STREAM)
+
+        val id = other.id
+        addBuffer(currentFile, id)
+        val gdsBuffer = new RapidsGdsBatchedBuffer(id, currentFile, currentOffset,
+          other.size, other.meta, other.getSpillPriority, other.spillCallback)
+        currentOffset += alignUp(deviceBuffer.getLength)
+        pendingBuffers += gdsBuffer
+        gdsBuffer
+      }
+
+    private def alignUp(length: Long): Long = {
+      (length + blockSize - 1) & ~(blockSize - 1)
+    }
+
+    private def copyToBuffer(
+        buffer: MemoryBuffer, offset: Long, size: Long, stream: Cuda.Stream): Unit = {
+      buffer.copyFromMemoryBuffer(0, batchWriteBuffer, offset, size, stream)
+    }
+
+    private def addBuffer(path: File, id: RapidsBufferId): Set[RapidsBufferId] = {
+      val updater = new BiFunction[File, Set[RapidsBufferId], Set[RapidsBufferId]] {
+        override def apply(key: File, value: Set[RapidsBufferId]): Set[RapidsBufferId] = {
+          if (value == null) {
+            Set(id)
+          } else {
+            value + id
+          }
+        }
+      }
+      spilledBuffers.compute(path, updater)
+    }
+
+    private def removeBuffer(path: File, id: RapidsBufferId): Set[RapidsBufferId] = {
+      val updater = new BiFunction[File, Set[RapidsBufferId], Set[RapidsBufferId]] {
+        override def apply(key: File, value: Set[RapidsBufferId]): Set[RapidsBufferId] = {
+          val newValue = value - id
+          if (newValue.isEmpty) {
+            null
+          } else {
+            newValue
+          }
+        }
+      }
+      spilledBuffers.computeIfPresent(path, updater)
+    }
+
+    class RapidsGdsBatchedBuffer(
+        id: RapidsBufferId,
+        path: File,
+        fileOffset: Long,
+        size: Long,
+        meta: TableMeta,
+        spillPriority: Long,
+        spillCallback: RapidsBuffer.SpillCallback,
+        var isPending: Boolean = true)
+        extends RapidsGdsBuffer(id, size, meta, spillPriority, spillCallback) {
+
+      override def materializeMemoryBuffer: MemoryBuffer = this.synchronized {
+        closeOnExcept(DeviceMemoryBuffer.allocate(size)) { buffer =>
+          if (isPending) {
+            copyToBuffer(buffer, fileOffset, size, Cuda.DEFAULT_STREAM)
+            Cuda.DEFAULT_STREAM.sync()
+            logDebug(s"Created device buffer $size from batch write buffer")
+          } else {
+            CuFile.readFileToDeviceBuffer(buffer, path, fileOffset)
+            logDebug(s"Created device buffer for $path $fileOffset:$size via GDS")
+          }
+          buffer
+        }
+      }
+
+      override def copyToMemoryBuffer(srcOffset: Long, dst: MemoryBuffer, dstOffset: Long,
+          length: Long, stream: Cuda.Stream): Unit = this.synchronized {
+        dst match {
+          case dmOriginal: DeviceMemoryBuffer =>
+            val dm = dmOriginal.slice(dstOffset, length)
+            if (isPending) {
+              copyToBuffer(dm, fileOffset + srcOffset, size, stream)
+              stream.sync()
+              logDebug(s"Created device buffer $size from batch write buffer")
+            } else {
+              // TODO: switch to async API when it's released, using the passed in CUDA stream.
+              stream.sync()
+              CuFile.readFileToDeviceBuffer(dm, path, fileOffset + srcOffset)
+              logDebug(s"Created device buffer for $path $fileOffset:$size via GDS")
+            }
+          case _ => throw new IllegalStateException(
+            s"GDS can only copy to device buffer, not ${dst.getClass}")
+        }
+      }
+
+      /**
+       * Mark this buffer as disk based, no longer in device memory.
+       */
+      def unsetPending(): Unit = this.synchronized {
+        isPending = false
+      }
+
+      override protected def releaseResources(): Unit = {
+        val ids = removeBuffer(path, id)
+        if (ids == null) {
+          if (!path.delete() && path.exists()) {
+            logWarning(s"Unable to delete GDS spill path $path")
+          }
+        }
+      }
+    }
+  }
+}