[FEA] Zero Conf - test/prototype chunked compress/encrypt

[revans2]

Is your feature request related to a problem? Please describe.
We want to try and improve our CPU utilization for shuffle with compression/encryption on shuffle partitions larger than 200. This was the main reason for the multi-threaded shuffle. To be able to do the CPU computation heavy processing in a thread pool instead of doing it in a single thread. I think we can do even better still if we can chunk the input data at arbitrary points and then encrypt/compress the data in a thread pool before it is written out to a file/etc.

The goal of this is to see if we can build a thread pool implementation that can take one large buffer of data to be output, chunk it into decently large chunks, optionally encrypt/compress that data and then slice it back together again as it is written out.

Along with this have the reader do the opposite. As it pulls in the data, detect that there are separate chunks of information, so that they get farmed out to a thread pool to do the decryption/decompression.

We want to make sure that we can prototype this for an internal shuffle where we don't care about the file format where the data is stored in, but also if we need to write the data out in a format that is compatible with the external shuffle file format. I think we should be able to make it work in both cases, but we might need a bit of extra metadata inline to know how the chunks are split up.

[ttnghia]

Since cudf has already implemented host-side compression/decompression (codec) as internal APIs, we can leverage them to implement this functionality in spark-rapids-jni. In order to do so we need a few steps:

• Expose the codec APIs such that they can be called outside of cudf: Add cudf/io/codec.hpp to expose compression/decompression APIs rapidsai/cudf#19113
• Implement additional host-side codec APIs to take in host buffers and also output into host buffers, then implement the serialization/deserialization function in spark-rapids-jni to use the codec function from cudf for this FEA: Implement Cookie serialization format spark-rapids-jni#3463

[binmahone]

Caveat: The below discussion only covers CPU shuffle, not taking GPU shuffle into account

Hi @revans2 , I think #13479 has almost implemented this idea, except that #13479 only targets for single batch case (as of today), and we need to extend it to multiple batch cases (we might need some merge sort capabilities as well), so that it will be a general option to replace vanilla spark's Sort based Shuffle r/w routines.

For shuffle read part, I think it's less important because the previous work #12674

So in the new shuffle writer based on #13479, we'll:

1. Regardless of partition number LT or GT 200, we will ALWAYS:

   a. exploit the fact that a GPU batch is sorted on partitioning key. This means that we won't open N (N = partition number) files, so we don't have to worry about the scaling issue past 200 shuffle partitions
   b. enable threaded pool to do encoding/compression
   c. having another separate thread to streaming serialized&compressed data to disk, and interleaving the encoding/compression process as much as possible

2. If the task is single-batch case, then we'll seek opportunities to avoid write amplification (e.g. avoid writing the data to disk for more than once). Remember, when the serialized&compressed steps are parallelized, they won't be the biggest bottleneck, disk IO will be. See the first snapshot in reduce NDS benchmark e2e time by optimizing RapidsShuffleManager #13479 (comment)

3. If the task is multi-batch case, we'll have a final sort merge step to copy all the M (M = data size / batch size) files to the final shuffle write stream.

What you think?