Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Additional unit tests for GeneratedInternalRowToCudfRowIterator #10087

Merged
merged 5 commits into from
Jan 2, 2024
Merged

Additional unit tests for GeneratedInternalRowToCudfRowIterator #10087

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
61093fd
Additional unit tests for GeneratedInternalRowToCudfRowIterator
jbrennan333 Dec 20, 2023
5b8172c
Merge branch 'branch-24.02' into jtb-rtoc-test
jbrennan333 Dec 28, 2023
041ea1e
Fix bug found by premerge - needed to initialize HostAlloc
jbrennan333 Dec 29, 2023
5be42e8
Merge branch 'branch-24.02' into jtb-rtoc-test
jbrennan333 Jan 2, 2024
eeea304
update copyrights
jbrennan333 Jan 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 105 additions & 8 deletions .../test/scala/com/nvidia/spark/rapids/GeneratedInternalRowToCudfRowIteratorRetrySuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.{CpuSplitAndRetryOOM, RmmSpark}
import com.nvidia.spark.rapids.jni.{GpuSplitAndRetryOOM, RmmSpark}
import org.mockito.ArgumentMatchers.any
import org.mockito.Mockito.{doAnswer, spy, times, verify}
import org.mockito.invocation.InvocationOnMock
Expand Down Expand Up @@ -56,7 +56,8 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
// this forces a retry on the copy of the host column to a device column
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId)
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
TestUtils.compareBatches(expected, devBatch)
Expand All @@ -77,7 +78,8 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
val res = invocation.callRealMethod()
// we mock things this way due to code generation issues with mockito.
// when we add a table we have
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 3)
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 3,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
rapidsBufferSpy = spy(res.asInstanceOf[RapidsBuffer])
rapidsBufferSpy
}
Expand All @@ -90,7 +92,8 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
val myIter = spy(GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric))
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 2)
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 2,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
assertResult(0)(getAndResetNumRetryThrowCurrentTask)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
Expand Down Expand Up @@ -118,7 +121,8 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
val res = invocation.callRealMethod()
// we mock things this way due to code generation issues with mockito.
// when we add a table we have
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 3)
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 3,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
rapidsBufferSpy = spy(res.asInstanceOf[RapidsBuffer])
// at this point we have created a buffer in the Spill Framework
// lets spill it
Expand All @@ -134,7 +138,8 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
val myIter = spy(GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric))
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 2)
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 2,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
assertResult(0)(getAndResetNumRetryThrowCurrentTask)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
Expand Down Expand Up @@ -163,11 +168,103 @@ class GeneratedInternalRowToCudfRowIteratorRetrySuite
val myIter = GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(1),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId)
assertThrows[CpuSplitAndRetryOOM] {
RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.GPU.ordinal, 0)
assertThrows[GpuSplitAndRetryOOM] {
myIter.next()
}
assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
}
}

test("a retry when allocating dataBuffer is handled") {
val batch = buildBatch()
val batchIter = Seq(batch).iterator
withResource(new ColumnarToRowIterator(batchIter, NoopMetric, NoopMetric, NoopMetric,
NoopMetric)) { ctriter =>
val schema = Array(AttributeReference("longcol", LongType)().toAttribute)
val myIter = GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
// this forces a retry on the allocation of the dataBuffer
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.CPU.ordinal, 0)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
TestUtils.compareBatches(expected, devBatch)
}
}
assert(!GpuColumnVector.extractBases(batch).exists(_.getRefCount > 0))
assert(!myIter.hasNext)
assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
}
}

test("a retry when allocating offsetsBuffer is handled") {
val batch = buildBatch()
val batchIter = Seq(batch).iterator
withResource(new ColumnarToRowIterator(batchIter, NoopMetric, NoopMetric, NoopMetric,
NoopMetric)) { ctriter =>
val schema = Array(AttributeReference("longcol", LongType)().toAttribute)
val myIter = GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
// this forces a retry on the allocation of the dataBuffer
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.CPU.ordinal, 1)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
TestUtils.compareBatches(expected, devBatch)
}
}
assert(!GpuColumnVector.extractBases(batch).exists(_.getRefCount > 0))
assert(!myIter.hasNext)
assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
}
}

test("a split and retry when allocating dataBuffer is handled") {
val batch = buildBatch()
val batchIter = Seq(batch).iterator
withResource(new ColumnarToRowIterator(batchIter, NoopMetric, NoopMetric, NoopMetric,
NoopMetric)) { ctriter =>
val schema = Array(AttributeReference("longcol", LongType)().toAttribute)
val myIter = GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
// this forces a retry on the allocation of the dataBuffer
RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.CPU.ordinal, 0)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
TestUtils.compareBatches(expected, devBatch)
}
}
assert(!GpuColumnVector.extractBases(batch).exists(_.getRefCount > 0))
assert(!myIter.hasNext)
assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
}
}
test("a split and retry when allocating offsetsBuffer is handled") {
val batch = buildBatch()
val batchIter = Seq(batch).iterator
withResource(new ColumnarToRowIterator(batchIter, NoopMetric, NoopMetric, NoopMetric,
NoopMetric)) { ctriter =>
val schema = Array(AttributeReference("longcol", LongType)().toAttribute)
val myIter = GeneratedInternalRowToCudfRowIterator(
ctriter, schema, TargetSize(Int.MaxValue),
NoopMetric, NoopMetric, NoopMetric, NoopMetric, NoopMetric)
// this forces a retry on the allocation of the offsetsBuffer
RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId, 1,
RmmSpark.OomInjectionType.CPU.ordinal, 1)
withResource(myIter.next()) { devBatch =>
withResource(buildBatch()) { expected =>
TestUtils.compareBatches(expected, devBatch)
}
}
assert(!GpuColumnVector.extractBases(batch).exists(_.getRefCount > 0))
assert(!myIter.hasNext)
assertResult(0)(RapidsBufferCatalog.getDeviceStorage.currentSize)
}
}
}