Update to nvcomp-2.x JNI APIs #3757

jbrennan333 · 2021-10-06T14:03:14Z

Closes #3754.

The nvcomp JNI API in CUDF is being updated to 2.x via rapidsai/cudf#9384
Once that change goes into CUDF, we will need to merge these changes to update the plugin to use the new nvcomp-2.x APIs. This PR also includes changes (from Alessandro) to add a config option to specify the lz4 chunk size.

Note that this PR will not build without the CUDF changes, so we can't build it yet.

Signed-off-by: Jim Brennan <[email protected]>

revans2 · 2021-10-06T14:05:16Z

Converted to Draft because rapidsai/cudf#9384 is not in, and is even still in draft

jbrennan333 · 2021-10-06T14:09:43Z

Thanks Bobby.

abellina

I think main thing is if there's a way to do the lz4 config slightly differently. I know that's how I had prototyped it... (sorry ahead of time).

abellina · 2021-10-11T22:59:19Z

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

@@ -31,7 +32,8 @@ class NvcompLZ4CompressionCodec extends TableCompressionCodec with Arm {
      contigTable: ContiguousTable,
      stream: Cuda.Stream): CompressedTable = {
    val tableBuffer = contigTable.getBuffer
-    val (compressedSize, oversizedBuffer) = NvcompLZ4CompressionCodec.compress(tableBuffer, stream)
+    val (compressedSize, oversizedBuffer) =
+      NvcompLZ4CompressionCodec.compress(tableBuffer, codecConfigs.lz4ChunkSize, stream)


nit, since we have codecConfigs everywhere, may as well pass it to compress?

I made the change to replace condecConfigs with rapidsConf as a constructor argument, and then I set a val in the class for the lz4ChunkSize, which is what I pass here now.

Using RapidsConf didn't work in general because it is not serializable, so I did change this to pass codecConfigs.

abellina · 2021-10-11T23:05:24Z

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

+      input: DeviceMemoryBuffer,
+      lz4ChunkSize: Int,
+      stream: Cuda.Stream): (Long, DeviceMemoryBuffer) = {
+    val lz4Config = LZ4Compressor.configure(lz4ChunkSize, input.getLength())


does this need to be wrapped in withResource?

Not for the LZ4Compressor. The Configuration returned by this method is not closeable, and doesn't need to be. The configuration for LZ4Decompressor does need to be closed, because there is a metadata object that needs to be destroyed.

abellina · 2021-10-11T23:06:48Z

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

@@ -1659,6 +1665,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {

  lazy val shuffleCompressionCodec: String = get(SHUFFLE_COMPRESSION_CODEC)

+  lazy val  shuffleCompressionLz4ChunkSize: Int = get(SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE)


Suggested change

lazy val shuffleCompressionLz4ChunkSize: Int = get(SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE)

lazy val shuffleCompressionLz4ChunkSize: Int = get(SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE)

Good catch! Fixed.

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

abellina · 2021-10-15T22:17:39Z

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

+      maxBatchMemorySize)
+    val inputBuffers: Array[BaseDeviceMemoryBuffer] = tables.map { table =>
+      val buffer = table.getBuffer
+      // cudf compressor will try to close this batch but this interface does not close inputs


Something like:

Suggested change

// cudf compressor will try to close this batch but this interface does not close inputs

// cudf compressor guarantees that close will be called for `inputBuffers` and will not throw before

// but this interface does not close inputs

i.e. could be a leak if compress throws before it wraps it, but that's not the case as written.

wonder if compress should not close?

I will follow up with Jason on this. This was part of the patch he gave me. My guess is that the compress code is trying to free up the input buffers before allocating the final output buffers, to reduce the memory pressure. But doing this incRef here defeats that.

I did update the comment for now.

abellina · 2021-10-15T22:21:14Z

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

+    }
+    closeOnExcept(batchCompressor.compress(inputBuffers, stream)) { compressedBuffers =>
+      withResource(new NvtxRange("lz4 post process", NvtxColor.YELLOW)) { _ =>
+        require(compressedBuffers.length == tables.length)


require is nice, I would add a message.

abellina · 2021-10-15T22:21:27Z

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

    extends BatchedBufferDecompressor(maxBatchMemory, stream) {
  override val codecId: Byte = CodecType.NVCOMP_LZ4

  override def decompressAsync(
      inputBuffers: Array[BaseDeviceMemoryBuffer],
      bufferMetas: Array[BufferMeta],
      stream: Cuda.Stream): Array[DeviceMemoryBuffer] = {
-    BatchedLZ4Decompressor.decompressAsync(inputBuffers, stream)
+    require(inputBuffers.length == bufferMetas.length)


same here, so we have a human readable message.

abellina · 2021-10-15T22:22:02Z

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

+    closeOnExcept(batchCompressor.compress(inputBuffers, stream)) { compressedBuffers =>
+      withResource(new NvtxRange("lz4 post process", NvtxColor.YELLOW)) { _ =>
+        require(compressedBuffers.length == tables.length)
+        compressedBuffers.zipWithIndex.map { case (buffer, i) =>


compressedBuffers.zip(tables) { case (buffer, table) =>

Since we don't need the index.

Thanks - Improved as suggested.

abellina · 2021-10-15T22:24:01Z

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

+      bufferMetas.zip(inputBuffers).safeMap { case (meta, input) =>
+        // cudf decompressor will try to close inputs but this interface does not close inputs
+        input.incRefCount()


should decompressor not close its inputs?

Going to follow up with Jason on this.

This is to be consistent with the plugin compressor batch API which also does not close its inputs. We could consider changing this, but we'd need to update all callers accordingly.

abellina · 2021-10-15T22:25:19Z

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

+  val SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE = conf("spark.rapids.shuffle.compression.lz4.chunkSize")
+    .doc("A configurable chunk size to use when compressing with LZ4.")
+    .internal()
+    .integerConf


nit, this should be bytesConf(ByteUnit.BYTE) (my fault)

When I use bytesConf(), it forces the type to be Long, but the APIs for chunksize all take an integer, so I wanted to define this as an integer.

Actually, it is just in the lz4compressConfigure JNI that we are using an int. I can change this to use longs for chunksize everywhere - but will need to make changes in cudf as well. Do you think it is worth it?

jbrennan333 · 2021-10-19T16:22:10Z

Thanks for the review @abellina! I put up another commit that addresses most of them. This does not include changing the chunkSize config to a Long and it does not change the compress/decompress closing inputs behavior.

jbrennan333 · 2021-10-19T18:31:34Z

I pushed another commit to change the chunkSize to long after pushing a corresponding commit in the cudf PR.

abellina

This LGTM

abellina · 2021-10-20T18:49:46Z

I +1'ed the cuDF PR for this. Thanks @jbrennan333

jlowe · 2021-10-22T20:47:25Z

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

+      bufferMetas.zip(inputBuffers).safeMap { case (meta, input) =>
+        // cudf decompressor will try to close inputs but this interface does not close inputs
+        input.incRefCount()


This is to be consistent with the plugin compressor batch API which also does not close its inputs. We could consider changing this, but we'd need to update all callers accordingly.

jbrennan333 · 2021-10-26T15:45:42Z

build

Update to nvcomp 2.x APIs

cd40825

Signed-off-by: Jim Brennan <[email protected]>

jbrennan333 added feature request New feature or request shuffle things that impact the shuffle plugin cudf_dependency An issue or PR with this label depends on a new feature in cudf labels Oct 6, 2021

revans2 marked this pull request as draft October 6, 2021 14:04

abellina requested changes Oct 11, 2021

View reviewed changes

jbrennan333 added 3 commits October 14, 2021 09:50

Merge branch 'branch-21.12' into jtb-nvcomp-2.x

31873ee

Address review comments

4ef0b03

Pass codecConfigs to compress()

959f2eb

abellina requested changes Oct 15, 2021

View reviewed changes

Update based on review comments

8585124

jbrennan333 mentioned this pull request Oct 19, 2021

Update Java nvcomp JNI bindings to nvcomp 2.x API rapidsai/cudf#9384

Merged

Change rapidsConf.shuffleCompressionLz4ChunkSize to a long

8fa8ba2

abellina reviewed Oct 19, 2021

View reviewed changes

abellina approved these changes Oct 20, 2021

View reviewed changes

jlowe approved these changes Oct 22, 2021

View reviewed changes

jbrennan333 marked this pull request as ready for review October 26, 2021 15:45

jbrennan333 merged commit ee44d6e into NVIDIA:branch-21.12 Oct 26, 2021

abellina mentioned this pull request Nov 12, 2021

[FEA] Make LZ4_CHUNK_SIZE configurable #2867

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update to nvcomp-2.x JNI APIs #3757

Update to nvcomp-2.x JNI APIs #3757

jbrennan333 commented Oct 6, 2021

revans2 commented Oct 6, 2021

jbrennan333 commented Oct 6, 2021

abellina left a comment

abellina Oct 11, 2021

jbrennan333 Oct 13, 2021

jbrennan333 Oct 14, 2021

abellina Oct 11, 2021

jbrennan333 Oct 13, 2021

abellina Oct 11, 2021

jbrennan333 Oct 13, 2021

abellina Oct 15, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

jbrennan333 Oct 19, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

jlowe Oct 22, 2021

abellina Oct 15, 2021

jbrennan333 Oct 19, 2021

jbrennan333 Oct 19, 2021

jbrennan333 commented Oct 19, 2021

jbrennan333 commented Oct 19, 2021

abellina left a comment

abellina commented Oct 20, 2021

jlowe Oct 22, 2021

jbrennan333 commented Oct 26, 2021

		@@ -1659,6 +1665,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {

		lazy val shuffleCompressionCodec: String = get(SHUFFLE_COMPRESSION_CODEC)

		lazy val shuffleCompressionLz4ChunkSize: Int = get(SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE)

	// cudf compressor will try to close this batch but this interface does not close inputs
	// cudf compressor guarantees that close will be called for `inputBuffers` and will not throw before
	// but this interface does not close inputs

Update to nvcomp-2.x JNI APIs #3757

Update to nvcomp-2.x JNI APIs #3757

Conversation

jbrennan333 commented Oct 6, 2021

revans2 commented Oct 6, 2021

jbrennan333 commented Oct 6, 2021

abellina left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jbrennan333 commented Oct 19, 2021

jbrennan333 commented Oct 19, 2021

abellina left a comment

Choose a reason for hiding this comment

abellina commented Oct 20, 2021

Choose a reason for hiding this comment

jbrennan333 commented Oct 26, 2021