Add HostMemoryAllocator interface

[gerashegalov]

Description

Creates an interface to intercept calls to HostMemoryBuffer.allocate

Fixes NVIDIA/spark-rapids#8884

Signed-off-by: Gera Shegalov [email protected]

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[revans2]

The current set of changes look good, but how do we actually inject our own HostMemoryAllocator for something like copying data from the device to the host?

We also will need to look at the C++ code at some point.

cudf/java/src/main/native/src/cudf_jni_apis.hpp

Line 103 in 55a4ecf

jobject allocate_host_buffer(JNIEnv *env, jlong amount, jboolean prefer_pinned);

That is really only used for I/O operations like arrow and the writer APIs.

[gerashegalov]

@revans2

how do we actually inject our own HostMemoryAllocator for something like copying data from the device to the host?

If we need full flexibility of initializing in the layer above cuDF/Java, and if it's a singleton I was thinking of an API like Cuda.setHostMemoryAllocator(HostMemoryAllocator hma).

We also will need to look at the C++ code at some point.

cudf/java/src/main/native/src/cudf_jni_apis.hpp

Line 103 in 55a4ecf

jobject allocate_host_buffer(JNIEnv *env, jlong amount, jboolean prefer_pinned);

That is really only used for I/O operations like arrow and the writer APIs.

I can add this once it's clear where the singleton allocator lives.

[revans2]

Lets take the example of ColumnView.copyToHost, which the Spark plugin could use for columnar to row conversion. One of the rules we need to follow to avoid deadlocks is that all host memory is spillable if we do a blocking allocation call. But because ColumnView.copyToHost can allocate multiple buffers before returning we cannot make them spillable in between each allocation without even more major changes to copyToHost. The plan to work around this was to ask the ColumnView how much memory we needed to allocate (#13919) and then reserve that amount so we can be guaranteed that the memory is available for us to allocate.

I see two ways to do this.

1. Change every single API that allocates large amounts of host memory to take an HostMemoryAllocator as a parameter.
2. Have a singleton, but the singleton we put in would have enough smarts that the plugin could call into it to say for this block of code allocate memory from this reservation instead.

Originally I thought we would do option 1, but the singleton idea is growing on me, and I think we could have the APIs be just as clean. If you do want to do a singleton that is fine. I will just need to make sure that we have an issue to implement a HostMemoryAllocator that would give us the flexibility we need.

[gerashegalov]

I added explicit HostMemoryAllocator parameter passing because it is indeed simple and gives us most flexibility.

still need to think about allocate_host_buffer, maybe in a separate PR after all.

[revans2]

Looks good. If we want to have a follow on for the native code that is fine.

[gerashegalov]

/merge