GpuCoalesceBatches should throw SplitAndRetyOOM on GPU OOM error. (#9374

) Signed-off-by: Jason Lowe <jlowe@nvidia.com>
NVIDIA · Oct 4, 2023 · 7e899f4 · 7e899f4
1 parent 54f5073
commit 7e899f4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala
@@ -23,6 +23,7 @@ import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.RmmRapidsRetryIterator.{withRetry, withRetryNoSplit}
 import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion
+import com.nvidia.spark.rapids.jni.SplitAndRetryOOM
 import com.nvidia.spark.rapids.shims.{ShimExpression, ShimUnaryExecNode}
 
 import org.apache.spark.TaskContext
@@ -670,7 +671,7 @@ abstract class AbstractGpuCoalesceIterator(
         val it = batchesToCoalesce.batches
         val numBatches = it.length
         if (numBatches <= 1) {
-          throw new OutOfMemoryError(s"Cannot split a sequence of $numBatches batches")
+          throw new SplitAndRetryOOM(s"Cannot split a sequence of $numBatches batches")
         }
         val res = it.splitAt(numBatches / 2)
         Seq(BatchesToCoalesce(res._1), BatchesToCoalesce(res._2))

diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuCoalesceBatchesRetrySuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuCoalesceBatchesRetrySuite.scala
@@ -175,7 +175,7 @@ class GpuCoalesceBatchesRetrySuite
   test("coalesce gpu batches fails with OOM if it cannot split enough") {
     val iters = getIters(mockInjectSplitAndRetry = true)
     iters.foreach { iter =>
-      assertThrows[OutOfMemoryError] {
+      assertThrows[SplitAndRetryOOM] {
         iter.next() // throws
       }
       val batches = iter.asInstanceOf[CoalesceIteratorMocks].getBatches()