Fix callback resource alignment forwarding - #2485
Open
nethum529 wants to merge 7 commits into
Open
Conversation
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe callback memory resource API now passes stream and alignment through C++ and Python callbacks, validates allocation alignment before invocation, updates documentation and type declarations, and adjusts C++ and Python tests and lint configuration. ChangesCallback alignment API
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx (1)
662-675: 🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy liftPreserve the prior Python callback contract through a deprecation cycle.
Existing
(size, stream)/(ptr, size, stream)callbacks now fail at their first invocation because this public constructor unconditionally calls the new arities. Provide a temporary legacy adapter with a deprecation warning before making the stream-first, alignment-aware signatures mandatory.As per coding guidelines, “Treat Python API changes that break backward compatibility … without a deprecation cycle, as CRITICAL.”
🤖 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 `@python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx` around lines 662 - 675, The callback_memory_resource constructor now invokes only the new stream-first, alignment-aware callback signatures, breaking existing callbacks using (size, stream) and (ptr, size, stream). Update __init__ to detect and adapt legacy callback arities through a temporary compatibility wrapper, emit a deprecation warning when that adapter is used, and preserve the existing behavior for new signatures while retaining the callback references safely.Source: Coding guidelines
🤖 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.
Inline comments:
In `@cpp/include/rmm/mr/callback_memory_resource.hpp`:
- Around line 27-61: Preserve backward compatibility for the public
allocate_callback_t and deallocate_callback_t callback declarations by retaining
deprecated legacy aliases or constructor overloads with the previous signatures.
Adapt legacy callbacks to the new alignment-aware callbacks internally, emit
deprecation warnings, and add migration guidance documenting the new signatures
and removal plan.
---
Outside diff comments:
In `@python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx`:
- Around line 662-675: The callback_memory_resource constructor now invokes only
the new stream-first, alignment-aware callback signatures, breaking existing
callbacks using (size, stream) and (ptr, size, stream). Update __init__ to
detect and adapt legacy callback arities through a temporary compatibility
wrapper, emit a deprecation warning when that adapter is used, and preserve the
existing behavior for new signatures while retaining the callback references
safely.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d804c145-8cb1-469a-a1df-1ad6db2d9733
📒 Files selected for processing (12)
.pre-commit-config.yamlcpp/include/rmm/mr/callback_memory_resource.hppcpp/include/rmm/mr/detail/callback_memory_resource_impl.hppcpp/src/mr/detail/callback_memory_resource_impl.cppcpp/tests/mr/callback_mr_tests.cppcpp/tests/mr/cccl_adaptor_tests.cppcpp/tests/mr/mr_ref_callback_tests.cpppython/rmm/rmm/librmm/memory_resource.pxdpython/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyipython/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyxpython/rmm/rmm/tests/test_callback_memory_resource.pypython/rmm/rmm/tests/test_failure_callback_resource_adaptor.py
Comment on lines
+27
to
+61
| * `void* allocate_callback_t(cuda_stream_view stream, std::size_t bytes, std::size_t alignment, | ||
| * void* arg);` | ||
| * | ||
| * * The callback receives only valid power-of-two alignment values. An invalid alignment raises | ||
| * `rmm::logic_error` before callback invocation. | ||
| * | ||
| * * Returns a pointer to an allocation of at least `bytes` usable immediately on | ||
| * `stream`. The stream-ordered behavior requirements are identical to | ||
| * `allocate`. | ||
| * `stream` with the requested `alignment`. The stream-ordered behavior requirements are | ||
| * identical to `allocate`. | ||
| * | ||
| * * The `arg` is provided to the constructor of the `callback_memory_resource` | ||
| * and will be forwarded along to every invocation of the callback function. | ||
| */ | ||
| using allocate_callback_t = std::function<void*(std::size_t, cuda_stream_view, void*)>; | ||
| using allocate_callback_t = std::function<void*(cuda_stream_view, std::size_t, std::size_t, void*)>; | ||
|
|
||
| /** | ||
| * @brief Callback function type used by callback_memory_resource for deallocation. | ||
| * | ||
| * The signature of the callback function is: | ||
| * `void deallocate_callback_t(void* ptr, std::size_t bytes, cuda_stream_view stream, void* arg);` | ||
| * `void deallocate_callback_t(cuda_stream_view stream, void* ptr, std::size_t bytes, | ||
| * std::size_t alignment, void* arg);` | ||
| * | ||
| * * Deallocates memory pointed to by `ptr`. `bytes` specifies the size of the allocation | ||
| * in bytes, and must equal the value of `bytes` that was passed to the allocate callback | ||
| * function. The stream-ordered behavior requirements are identical to | ||
| * `deallocate`. | ||
| * function. `alignment` must equal the value of `alignment` that was passed to the allocate | ||
| * callback function. The stream-ordered behavior requirements are identical to `deallocate`. | ||
| * | ||
| * * The callback must not throw. An exception escaping the callback causes termination because | ||
| * deallocation is `noexcept`. | ||
| * | ||
| * * The `arg` is provided to the constructor of the `callback_memory_resource` | ||
| * and will be forwarded along to every invocation of the callback function. | ||
| */ | ||
| using deallocate_callback_t = std::function<void(void*, std::size_t, cuda_stream_view, void*)>; | ||
| using deallocate_callback_t = | ||
| std::function<void(cuda_stream_view, void*, std::size_t, std::size_t, void*)>; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
Preserve the previous callback API through a deprecation cycle.
Replacing these public aliases immediately breaks existing callback declarations and constructor call sites. Retain a deprecated legacy overload/alias and provide migration guidance before removal.
As per path instructions, “Breaking changes require deprecation warnings and migration guide updates.”
🤖 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 `@cpp/include/rmm/mr/callback_memory_resource.hpp` around lines 27 - 61,
Preserve backward compatibility for the public allocate_callback_t and
deallocate_callback_t callback declarations by retaining deprecated legacy
aliases or constructor overloads with the previous signatures. Adapt legacy
callbacks to the new alignment-aware callbacks internally, emit deprecation
warnings, and add migration guidance documenting the new signatures and removal
plan.
Sources: Coding guidelines, Path instructions
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Testing
pre-commit run --all-filesBreaking change:
allocate_callback_tanddeallocate_callback_treorder parameters (stream first) and add analignmentparameter, so existing C++ and Python callbacks must update their signatures; the break is loud (old callbacks fail to compile / raise TypeError) rather than silent. Needs thebreakinglabel.Closes #2397