Pre-touch pages concurrently during pinned memory pool initialization - #23457
Pre-touch pages concurrently during pinned memory pool initialization#23457rishic3 wants to merge 11 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesPinned-memory pool initialization now accepts a positive thread count. Java selects single-thread or parallel RMM resources. JNI adds parallel page touching, CUDA registration, cleanup, and validation. Tests cover allocation, alignment, and invalid arguments. Parallel pinned-memory pool
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java (1)
56-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a case where thread count exceeds page count.
touch_pagesclampsthread_countto the page count; a tiny pool with a largeinitializationThreads(e.g. 4 KiB pool with 16 threads) would exercise that clamp and the single-page partition path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java` around lines 56 - 68, Extend initParallelPoolWithNonPageAlignedSize to use a small page-aligned pool and an initializationThreads value greater than its page count, such as a 4 KiB pool with 16 threads. Keep the allocation and byte read/write assertions, while ensuring the test exercises thread-count clamping and the single-page partition path.java/src/main/native/src/RmmJni.cpp (1)
531-546: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueConsider a thread-safe strerror.
std::strerroris not thread-safe and these helpers can be invoked from concurrent deallocation paths.strerror_r/a small local buffer avoids interleaved or torn messages.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/native/src/RmmJni.cpp` around lines 531 - 546, Update log_system_error_noexcept to avoid calling the non-thread-safe std::strerror from concurrent cleanup paths. Resolve the error text with strerror_r (or an equivalent thread-safe conversion) into a local buffer, then pass that buffer to the existing CUDF_LOG_WARN/CUDF_LOG_ERROR calls while preserving the noexcept protection and warning/error branching.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@java/src/main/native/src/RmmJni.cpp`:
- Around line 531-546: Update log_system_error_noexcept to avoid calling the
non-thread-safe std::strerror from concurrent cleanup paths. Resolve the error
text with strerror_r (or an equivalent thread-safe conversion) into a local
buffer, then pass that buffer to the existing CUDF_LOG_WARN/CUDF_LOG_ERROR calls
while preserving the noexcept protection and warning/error branching.
In `@java/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java`:
- Around line 56-68: Extend initParallelPoolWithNonPageAlignedSize to use a
small page-aligned pool and an initializationThreads value greater than its page
count, such as a 4 KiB pool with 16 threads. Keep the allocation and byte
read/write assertions, while ensuring the test exercises thread-count clamping
and the single-page partition path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e26c1850-9a4f-4fd8-880e-7946f528a8f1
📒 Files selected for processing (4)
java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.javajava/src/main/java/ai/rapids/cudf/Rmm.javajava/src/main/native/src/RmmJni.cppjava/src/test/java/ai/rapids/cudf/PinnedMemoryPoolTest.java
|
|
||
| // Note that these threads will inherit the caller's CPU affinity but are not individually | ||
| // affinity bound. If the calling thread is not bound to a NUMA node, the resultant pages | ||
| // could span multiple NUMA nodes. |
There was a problem hiding this comment.
Codex provided some suggestions to consider
- Choose one target NUMA node before starting the first-touch workers. Derive
it from either the selected GPU's local NUMA node or the initiating thread's
existing CPU/memory placement. - Make every worker allocate its pages on that node. Bind each worker's CPU
affinity to CPUs on the target node, or apply an equivalent NUMA memory
policy to the mapping while it is first-touched. Restore any temporary
affinity or policy afterward. - If the process spans multiple NUMA nodes and no reliable target node can be
selected, do not silently scatter the pool. Fall back to the existing
initialization path, require an explicit opt-in, or clearly document and
warn about the possible long-term DMA penalty. - Add a multi-NUMA test that verifies the touched pages reside on the intended
node and measures sustained H2D and D2H bandwidth against the existing
cudaHostAllocpool. The current tests establish only that allocation
succeeds and that the CPU can read and write the bytes; they do not validate
page placement or GPU-transfer performance.
There was a problem hiding this comment.
This is naive and probably wrong in places, but the way I see it there are two cases to consider.
- the calling thread is intentionally constrained (given a non-default memory policy
MPOL_BIND,MPOL_INTERLEAVE,MPOL_PREFERRED, and/or bound to CPUs on specific nodes). - the calling thread is unconstrained (default memory policy
MPOL_DEFAULTand unbound CPU affinity).
Case 1 is straightforward; we should respect the calling thread's constraints, and that requires no explicit action on our part since it is inherited by threads. And it would work seamlessly with a future feature that binds JVM threads on startup.
Case 2, where the caller is unbound, is where I suppose we might want to do something. Taking Codex's two suggested options:
- Bind threads to the same node as the caller. It is reasonable to say that this is probably closer to what a single-threaded
cudaHostAllocwould do, but as far as I knowcudaHostAllocmakes no guarantee about sticking to the caller's node (e.g., the thread could migrate mid-call). Furthermore, there is no guarantee that the node of the caller is actually the NUMA node that is good for the GPU later on. - Bind threads to node(s) with affinity to the given GPU ID. That would be doing a lot more than what
cudaHostAllocdoes. And (as Codex pointed out) if there are multiple nodes it would not be clear, what to do, e.g., whether the user wants all memory on one node, interleaved, etc. It seems reasonable to leave that responsibility to a higher caller.
Finally this all assumes we can distinguish case 1 from case 2, but if the only lever is CPU affinity, it is not clear. E.g., if we see a mask that spans nodes {0,1}, I don't know if we should respect that assuming the user knows the GPU has affinity to both those CPUs, or if we should try to do something better because our machine only has two nodes.
So for this change I don't think there is any init time decision that is correct in general (so I think do-nothing is reasonable), and it is already opt-in on the Spark side.
Also note that if this eventually goes into RMM, I think the NUMA ID would be made an explicit part of the API (as suggested by @bdice).
There was a problem hiding this comment.
I didn't response to suggestions 3. and 4.
clearly document and warn about the possible long-term DMA penalty
I'll surface the risk beyond this comment, to a docstring.
Add a multi-NUMA test that verifies the touched pages reside on the intended node and measures sustained H2D and D2H bandwidth
Assuming we don't do any explicit binding per above, I'll ignore this one.
abellina
left a comment
abellina
left a comment
There was a problem hiding this comment.
I have some comments, but overall looks good to me.
| } | ||
| // munmap always unmaps whole pages, so we can pass the unaligned bytes here. | ||
| if (::munmap(ptr, bytes) != 0) { log_system_error_noexcept("munmap", errno); } | ||
| } |
There was a problem hiding this comment.
it would remove ambiguity if we aligned the bytes argument, and we can remove the comment as well. + if we decide to align differently later, or request huge pages at mmap time we can factor that into the alignment code?
There was a problem hiding this comment.
Makes sense, though I kept an unaligned fallback to not skip the unmap.. let me know how it looks.
| [[nodiscard]] bool operator==( | ||
| [[maybe_unused]] parallel_init_pinned_host_memory_resource const& other) const noexcept | ||
| { | ||
| return true; |
There was a problem hiding this comment.
how about:
return this == std::addressof(other);
There was a problem hiding this comment.
Hm. I copied this from RMM. I presume equality means we can do
void* ptr = a.allocate_sync(bytes);
b.deallocate_sync(ptr, bytes);on different pools a and b.. is that violated here? I mean I think the suggestion is harmless but I wonder if it makes it appear as though we are not cross-deallocation-compatible.
| auto const status = cudaHostUnregister(ptr); | ||
| if (status != cudaSuccess) { | ||
| cudaGetLastError(); | ||
| log_cuda_error_noexcept("cudaHostUnregister", status); |
There was a problem hiding this comment.
this skips the munmap when cudaHostUnregister fails. I think we should try it anyway.
There was a problem hiding this comment.
Done. Should we be concerned about giving back a virtual memory range that CUDA still thinks is pinned? I'm guessing not if the context is screwed already...
| start_workers.store(true); | ||
| start_workers.notify_all(); | ||
| for (auto& worker : workers) { | ||
| if (worker.joinable()) { worker.join(); } |
There was a problem hiding this comment.
we should catch any exception the worker.join() could throw so we call join() on all workers. Unlikely to trigger in practice, but still would be nice to do.
igorpeshansky
left a comment
igorpeshansky
left a comment
There was a problem hiding this comment.
Good implementation. My main question is structural: should this land in RMM instead? You filed rapidsai/rmm#2496 for exactly this, and NVIDIA/cudf-spark-jni#4593 needs the same touch loop but has already diverged from this one (e.g., it has neither a start barrier nor MADV_HUGEPAGE).
Separately, one thing that needs to be addressed before merge is the init failure path (#23457 (comment)) — the new backing allocation can fail in ways cudaHostAlloc couldn't, and the current getSingleton() makes any such failure permanent.
| public static native long newPinnedPoolMemoryResource(long initSize, long maxSize); | ||
|
|
||
| static native long newParallelPinnedPoolMemoryResource( | ||
| long poolSize, int parallelInitializationThreads); |
There was a problem hiding this comment.
newPinnedPoolMemoryResource takes both initSize and maxSize, while this API hard-codes both to poolSize. Is dropping growable controls deliberate, or is it deferred until it's actually needed (related to the visibility question #23457 (comment))?
Related: @revans2 raised concerns of fixed-size pools being stuck (unable to grow, shrink, or be returned) in NVIDIA/cudf-spark-jni#4593 (review) — making initialization cheaper is an invitation to have larger pinned pools, none of which are reclaimable. Worth calling out in the javadoc alongside the NUMA note?
There was a problem hiding this comment.
The Java pinned pool was already always fixed size (passing poolSize for both init and max sizes), afaik. Whether we want it to be growable I'm not sure, but the concerns seem different than a fixed-size pageable pool, since pinned allocations can always fall back to pageable.
| explicit parallel_init_pinned_host_memory_resource(std::size_t initialization_threads) | ||
| : initialization_threads_{initialization_threads} | ||
| { | ||
| CUDF_EXPECTS(initialization_threads_ > 0, |
There was a problem hiding this comment.
This does not set an upper limit, and pretouch_parallel only compares to page_count, so nothing prevents the code from trying to spawn millions of threads. Since the numbers in the PR description flatten after ~4 threads, does it make sense to constrain initialization_threads_ to, e.g., std::thread::hardware_concurrency()? Also in Java_ai_rapids_cudf_Rmm_newParallelPinnedPoolMemoryResource…
There was a problem hiding this comment.
Good point. hardware_concurrency is a good idea, I've gone with that.
| * page faults concurrently. After the pages are physically backed the allocation is then registered | ||
| * with CUDA. The RMM pool using this upstream resource is otherwise unchanged. | ||
| */ | ||
| class parallel_init_pinned_host_memory_resource final { |
There was a problem hiding this comment.
This class is added within a JNI file, which seems like an anti-pattern to me (though already present here, sadly)… This would really benefit from C++-level tests (like in NVIDIA/cudf-spark-jni#4593), which require it to live in its own C++ source in a directory with a testing target (related to the RMM re-homing discussed in #23457 (comment)).
There was a problem hiding this comment.
Makes sense. Perhaps we can defer to a follow-up where we clean this up along with the other classes. Unless you think there is a blocking need for C++ tests in this one.
There was a problem hiding this comment.
This seems to be heavily related to NVIDIA/cudf-spark-jni#4593. Most of the concerns raised there no longer apply here, but a few still do — I've called them out separately.
Since you've already said that RMM is the eventual destination for this (#23457 (comment)), why not start there, given that rapidsai/rmm#2496 is already open, and NVIDIA/cudf-spark-jni#4593 needs the same code?
4 participants
Description
Contributes to NVIDIA/cudf-spark#15145. Closes #23477.
This adds a parallel pre-touch initialization mode to the Java pinned memory pool. By default, the RMM pinned pool makes a single
cudaHostAllocfor the requested region, wherein the CUDA driver is responsible for physically backing and pinning all the pages. In the new mode, we parallelize the physical backing of the pages by having many threads pre-touch the pages, serving page faults concurrently; the allocation can then be pinned and mapped withcudaHostRegister. Furthermore, the new initialization mode requests huge pages for the pinned allocation so when possible there are even fewer page faults to service.Here are the initialization times on a single EC2 g6.4xlarge (16 vcpus):
...and the same numbers in graph form:
Here is the marginal speedup we get from enabling huge pages on top of parallelism (green lines are huge-page-enabled, pink are not) -- huge pages help at basically all thread counts and it evens out at 16 threads:
Checklist