Add retry with SpillableHostBuffer to GPU scans

[firestarman]

contribute to #8890

This PR adds the support of retry with SpillableHostBuffer into GPU scans when reading the data from files to host, to try to resolve some CPU OOMs in GPU scans, which are reported from some customer queries. And these host memory allocations run without the protection of the retry framework.

The OOM error is like as the below.

Caused by: com.nvidia.spark.rapids.jni.CpuRetryOOM: Could not complete allocation after 1000 retries
        at com.nvidia.spark.rapids.HostAlloc.alloc(HostAlloc.scala:244)
        at com.nvidia.spark.rapids.HostAlloc.allocate(HostAlloc.scala:250)
        at ai.rapids.cudf.HostMemoryBuffer.allocate(HostMemoryBuffer.java:138)
        at ai.rapids.cudf.HostMemoryBuffer.allocate(HostMemoryBuffer.java:149)
        at com.nvidia.spark.rapids.ParquetPartitionReaderBase.$anonfun$readPartFile$1(GpuParquetScan.scala:1938)
        at com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:30)
        at com.nvidia.spark.rapids.ParquetPartitionReaderBase.readPartFile(GpuParquetScan.scala:1936)

This change can be covered by existing tests.

[binmahone]

related to #8890

[firestarman]

build

[binmahone]

Not directly related to this PR, but it reminds me that generally speaking, SpillableHostBuffer on a sliced HostMemoryBuffer could be, IMHO, buggy in some cases @abellina @revans2 . Let's say originally I have a HostMemoryBuffer x, then another y sliced from x. When y is spilled and read from disk again, y is no longer sharing underlying buffer with x anymore, right? If the code writer assumes any content written to x is visible to y, then SpillingFramework is causing a bug? Should we make sliced HostMemoryBuffer non-spillable at all?

[binmahone]

But marking sliced HostMemoryBuffer non-spillable at all has side effects, too. This will essentially make the original HostMemoryBuffer x also unspillable, because it always have another reference coming from the sliced HostMemoryBuffer

[binmahone]

Also, SpillableHostBuffer on a sliced HostMemoryBuffer would mistakenly increase totalSize of SpillableHostStore:

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

Line 1363 in b1e0fd6

def trackNoSpill(handle: HostSpillableHandle[_]): Unit = {

[abellina]

In this case, we are slicing only to "right size" the buffer. The alternative is to memcpy it out to another right-sized buffer or to create a spillable buffer that knows about the slice, or track the "right size" alongside it for later slicing.

In terms of tracking the memory, we'll be undercounting it because we have 0->finalBufferSize, not the whole length. We will let it go and spill it because we hold onto the only reference in the spill framework. When we bring it back after spill it will be of size finalBufferSize, not the original, but I think this is OK because we didn't mean to keep the original.

So I think the main issue is undercounting.

[firestarman]

Thx for the discussion, so should I make any update to fix this "undercounting"?

[binmahone]

@abellina For this case I think it's okay to keep the code as it is, because we can assume finalBufferSize is close to totalBufferSize. But I brought up this question to discuss the general cases ?

[firestarman]

OK, I fully get what is going on here. i will resolve this by "tracking the 'right size' alongside it for later slicing".

[firestarman]

Updated

[firestarman]

build