Prevent potential GPU OOM in R2C with split retry - #14073
Conversation
|
@greptile full review |
There was a problem hiding this comment.
Pull request overview
This PR adds split retry functionality to Row-to-Columnar (R2C) conversion to prevent GPU out-of-memory (OOM) errors by implementing a mechanism to split batches in half when GPU OOM occurs during host-to-GPU data transfer.
Key Changes
- Introduced a new
HostColumnarBatchWithRowRangeclass that wraps host columns with row range tracking and supports splitting - Modified
RowToColumnarIteratorto use split retry logic for GPU OOM scenarios, allowing single input batches to produce multiple output batches - Added
buildHostColumnsWithoutOwnership()method to transfer ownership of host columns to the retry framework
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala |
New class that wraps host columns with row range support, implements splitting logic for GPU OOM retry, and handles slicing of various column types (LIST, STRUCT, STRING, fixed-width) |
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuRowToColumnarExec.scala |
Updated RowToColumnarIterator to build host columns separately, use split retry for TargetSize goals, and maintain a pending batch iterator for split outputs |
sql-plugin/src/main/java/com/nvidia/spark/rapids/GpuColumnVector.java |
Added buildHostColumnsWithoutOwnership() method to transfer host column ownership to caller |
tests/src/test/scala/com/nvidia/spark/rapids/RowToColumnarIteratorRetrySuite.scala |
Added comprehensive test coverage for GPU OOM split retry scenarios including single batch requirement enforcement, multiple batch production, multiple consecutive splits, and single row edge case |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 176-184 (link)logic: Workaround condition may not cover all cases. If
dataLen == 0butnullCount > 0(empty strings with some null values), the code goes to the else branch and creates a 0-length slice, which could trigger the samecopyToDeviceissue mentioned in the comment.
4 files reviewed, 1 comment
|
build |
1 similar comment
|
build |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 178-188 (link)style: Check that this empty string workaround handles all edge cases correctly. The 1-byte allocation when
dataLen == 0 && nullCount == 0works aroundcopyToDevicebehavior, but verify it handles: (1) all strings in slice are empty (not null), (2) mixed empty strings at different positions. This edge case could cause issues with certain data distributions.
4 files reviewed, 1 comment
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 180 (link)syntax: typo: "Existsing" should be "Existing"
4 files reviewed, 1 comment
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 178-186 (link)style: Consider if allocating a 1-byte dummy buffer for empty strings when
dataLen == 0butnumRows > nullCountis necessary, or if cudf would acceptnullhere to save allocations.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
4 files reviewed, 1 comment
|
@greptileai can you review it again ? |
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 180-188 (link)style: Edge case handling for empty strings is correct: when
dataLen == 0but there are non-null rows (numRows > nullCount), allocates 1-byte buffer to satisfy cuDF requirements. However, consider validating that this 1-byte allocation doesn't cause issues if we're in a tight memory situation (though this is unlikely given the small size).Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
4 files reviewed, 1 comment
|
build |
There was a problem hiding this comment.
Additional Comments (1)
-
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostColumnarBatchWithRowRange.scala, line 271-282 (link)logic: Documentation claims split ownership transfer but implementation doesn't match - both splits increment reference count, so memory isn't transferred to last split as stated. Should the memory ownership transfer logic be implemented, or should the documentation be corrected to reflect the current reference counting approach?
3 files reviewed, 1 comment
|
build |
|
@thirtiseven I should have time to look at this tomorrow, if it can wait until then. |
| totalOutputBytes += GpuColumnVector.getTotalDeviceMemoryUsed(batch) | ||
| totalOutputRows += batch.numRows() | ||
| if (totalOutputRows > 0 && totalOutputBytes > 0) { | ||
| targetRows = |
There was a problem hiding this comment.
should we call this "nextBatchTargetRows"?
| // Return the first batch now and keep the iterator for subsequent output batches. | ||
| // This ensures we only transfer one split at a time (avoid multiple device allocations). | ||
| closeOnExcept(it.next()) { first => | ||
| pendingBatchIter = it |
There was a problem hiding this comment.
should we check if it.hasNext before we set pendingBatchIter? e.g. set to None otherwise?
There was a problem hiding this comment.
yes good catch
| * | ||
| * Memory management uses reference counting: each instance increments the reference count | ||
| * of the host columns on construction and decrements it on close. The host columns are | ||
| * freed when the last reference is closed. |
There was a problem hiding this comment.
we should add that validity and offset buffers are copied, not logically sliced, in this description.
|
build |
|
Hi @abellina could you take another look? |
|
NOTE: release/26.02 has been created from main. Please retarget your PR to release/26.02 if it should be included in the release. |
|
build |
| if (totalOutputRows > 0 && totalOutputBytes > 0) { | ||
| targetRows = | ||
| GpuBatchUtils.estimateRowCount(targetSizeBytes, totalOutputBytes, totalOutputRows) | ||
| val dataTypes = localSchema.fields.map(_.dataType) |
There was a problem hiding this comment.
it would be better if this was a class val.
| HostColumnarBatchWithRowRange(hostColumns, rowCount, dataTypes) | ||
| } | ||
|
|
||
| if (localGoal.isInstanceOf[RequireSingleBatchLike]) { |
There was a problem hiding this comment.
localGoal match {
case RequireSingleBatchLike =>
...
case => // other goals
...
}
| // Return the first batch now and keep the iterator for subsequent output batches. | ||
| // This ensures we only transfer one split at a time (avoid multiple device allocations). | ||
| closeOnExcept(it.next()) { first => | ||
| pendingBatchIter = if (it.hasNext) it else Iterator.empty |
There was a problem hiding this comment.
nit, could just set pendingBatchIter = it, without checking, since we should be gating all the other logic in that it.hasNext returns false
There was a problem hiding this comment.
thanks for the review, done.
|
Let's waiting #14428 merging in first |
|
NOTE: release/26.04 has been created from main. Please retarget your PR to release/26.04 if it should be included in the release. |
Rebased on top of NVIDIA#14428's per-batch retry mechanism. Adds GPU OOM split retry for the host-to-GPU transfer step using HostColumnarBatchWithRowRange, which supports logical slicing of host columns without copying underlying memory. Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
a6ea3b9 to
c821e95
Compare
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
NOTE: release/26.06 has been created from main. Please retarget your PR to release/26.06 if it should be included in the release. |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release. |
Fixes #14018
Description
#13842 added retry for R2C
convertto prevent Host OOM. Following on from that, this pr aimed to add split and retry for GPU OOM when copying the converted results to GPU.This PR:
HostColumnarBatchWithRowRange, a wrapper for host columns that supports logical slicing without copying underlying host memory. This allows splitting a large host batch into smaller chunks for transfer.GpuRowToColumnarExecto use split and retry when a GPU OOM occurs during transfer.Note that when we split a batch into two halves, we can't free them until both halves are processed. So, the first half just borrows the data, but we pass the ownership of the host columns to the second half. This ensures the host memory stays alive exactly as long as needed and is freed when the last split is closed.
Performance tests:
Benchmark measures end-to-end time for R2C conversion + GPU aggregation on 100M rows. Data is created via sc.parallelize + createDataFrame to force CPU rows through GpuRowToColumnarExec. Three test groups:
Each group: 5 warmup + 10 timed runs, reporting median. spark-shell --master local[4] --driver-memory 8g.
benchmark code:
r2c_split_retry_benchmark.scala.zip
Checklists
Documentation
Testing
(Please provide the names of the existing tests in the PR description.)
Performance