Add type annotations to RMM tests#2332
Merged
rapids-bot[bot] merged 5 commits intorapidsai:mainfrom Apr 1, 2026
Merged
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdded and tightened static type annotations across many test modules and one type stub; refined CUDA array-interface typing in a stub; strengthened a statistics test with additional non-None assertions. No runtime behavior or public API signatures changed. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
Author
Collaborator
|
@TomAugspurger I am happy to merge annotations if you want to add them. I don't view this as a high priority for RMM but am supportive if it helps cuDF typing efforts. |
TomAugspurger
force-pushed
the
tom/tests-typing
branch
from
04770b9 to
508326a
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/rmm/rmm/tests/test_device_buffer.py (1)
145-147: Use the stricter union type instead ofAnyin this test parameter.
hb: Anyweakens type checking in a typing-focused PR. The type aliasHOST_BUFFER_OR_INVALID_T(defined at line 43 asHOST_BUFFER_T | str | int | None) is the appropriate annotation here and aligns with the function's design intent to handle both valid and invalid buffer types.Proposed change
def test_rmm_device_buffer_bytes_roundtrip( - hb: Any, # Intentionally tests invalid types + hb: HOST_BUFFER_OR_INVALID_T, ) -> None:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/rmm/rmm/tests/test_device_buffer.py` around lines 145 - 147, Replace the weak Any annotation on the hb parameter in test_rmm_device_buffer_bytes_roundtrip with the stricter union type alias HOST_BUFFER_OR_INVALID_T; update the function signature (test_rmm_device_buffer_bytes_roundtrip) to use hb: HOST_BUFFER_OR_INVALID_T so the test enforces the intended valid/invalid buffer types and restores proper static typing checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@python/rmm/rmm/tests/test_device_buffer.py`:
- Around line 145-147: Replace the weak Any annotation on the hb parameter in
test_rmm_device_buffer_bytes_roundtrip with the stricter union type alias
HOST_BUFFER_OR_INVALID_T; update the function signature
(test_rmm_device_buffer_bytes_roundtrip) to use hb: HOST_BUFFER_OR_INVALID_T so
the test enforces the intended valid/invalid buffer types and restores proper
static typing checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 39c56b95-38a2-4e55-8a5d-161b344cad5e
📒 Files selected for processing (11)
python/rmm/rmm/pylibrmm/device_buffer.pyipython/rmm/rmm/tests/test_allocations.pypython/rmm/rmm/tests/test_arena_memory_resource.pypython/rmm/rmm/tests/test_binning_memory_resource.pypython/rmm/rmm/tests/test_device_buffer.pypython/rmm/rmm/tests/test_helpers.pypython/rmm/rmm/tests/test_initialization.pypython/rmm/rmm/tests/test_lifecycle.pypython/rmm/rmm/tests/test_pool_memory_resource.pypython/rmm/rmm/tests/test_statistics.pypython/rmm/rmm/tests/test_stream.py
✅ Files skipped from review due to trivial changes (9)
- python/rmm/rmm/tests/test_initialization.py
- python/rmm/rmm/pylibrmm/device_buffer.pyi
- python/rmm/rmm/tests/test_binning_memory_resource.py
- python/rmm/rmm/tests/test_stream.py
- python/rmm/rmm/tests/test_allocations.py
- python/rmm/rmm/tests/test_arena_memory_resource.py
- python/rmm/rmm/tests/test_helpers.py
- python/rmm/rmm/tests/test_lifecycle.py
- python/rmm/rmm/tests/test_pool_memory_resource.py
TomAugspurger
commented
Mar 30, 2026
TomAugspurger
added
non-breaking
Contributor
Author
|
/ok to test c379ed9 |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python/rmm/rmm/tests/test_device_buffer.py (2)
377-382: Consider cleanup for allocated memory in negative test.If
pytest.raisesfails (i.e., no TypeError is raised), the allocated memory at line 380 would leak. While this is unlikely given the test's purpose, using a try/finally or ensuring cleanup would be more defensive.def test_memory_resource_deallocate_stream_none() -> None: """Test that DeviceMemoryResource.deallocate raises TypeError for stream=None""" mr = rmm.mr.get_current_device_resource() ptr = mr.allocate(1024) - with pytest.raises(TypeError, match="stream argument cannot be None"): - mr.deallocate(ptr, 1024, stream=None) # type: ignore[arg-type] + try: + with pytest.raises(TypeError, match="stream argument cannot be None"): + mr.deallocate(ptr, 1024, stream=None) # type: ignore[arg-type] + finally: + mr.deallocate(ptr, 1024)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/rmm/rmm/tests/test_device_buffer.py` around lines 377 - 382, The test allocates memory with mr.allocate(1024) but if pytest.raises does not raise, that allocation will leak; update test_memory_resource_deallocate_stream_none to ensure allocated pointer is always freed by surrounding the call with a try/finally (or equivalent) so that mr.deallocate(ptr, 1024) is called in the finally block using a valid stream argument (i.e., omit stream or pass a proper stream) and guard against ptr being None, referencing the existing mr.allocate and mr.deallocate calls.
41-48: Unused type aliasHOST_BUFFER_OR_INVALID_T.
HOST_BUFFER_OR_INVALID_Tis defined on line 44 buttest_rmm_device_buffer_bytes_roundtripusesAnyinstead (line 140). Consider either using this alias or removing it.def test_rmm_device_buffer_bytes_roundtrip( - hb: Any, # Intentionally tests invalid types + hb: HOST_BUFFER_OR_INVALID_T, # Intentionally tests invalid types ) -> None:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/rmm/rmm/tests/test_device_buffer.py` around lines 41 - 48, The type alias HOST_BUFFER_OR_INVALID_T is unused; update the test to use it or remove the alias: either change the type annotation in test_rmm_device_buffer_bytes_roundtrip from Any to HOST_BUFFER_OR_INVALID_T (ensuring the alias is imported/visible where the test lives) or delete the HOST_BUFFER_OR_INVALID_T definition if you prefer not to use it; this involves editing the alias declaration near HOST_BUFFER_T and the test function signature/annotations in test_rmm_device_buffer_bytes_roundtrip to keep types consistent.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@python/rmm/rmm/tests/test_device_buffer.py`:
- Around line 377-382: The test allocates memory with mr.allocate(1024) but if
pytest.raises does not raise, that allocation will leak; update
test_memory_resource_deallocate_stream_none to ensure allocated pointer is
always freed by surrounding the call with a try/finally (or equivalent) so that
mr.deallocate(ptr, 1024) is called in the finally block using a valid stream
argument (i.e., omit stream or pass a proper stream) and guard against ptr being
None, referencing the existing mr.allocate and mr.deallocate calls.
- Around line 41-48: The type alias HOST_BUFFER_OR_INVALID_T is unused; update
the test to use it or remove the alias: either change the type annotation in
test_rmm_device_buffer_bytes_roundtrip from Any to HOST_BUFFER_OR_INVALID_T
(ensuring the alias is imported/visible where the test lives) or delete the
HOST_BUFFER_OR_INVALID_T definition if you prefer not to use it; this involves
editing the alias declaration near HOST_BUFFER_T and the test function
signature/annotations in test_rmm_device_buffer_bytes_roundtrip to keep types
consistent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b3a9e339-1fe5-406d-b8a4-6cb3e347a046
📒 Files selected for processing (4)
python/rmm/rmm/tests/test_allocations.pypython/rmm/rmm/tests/test_binning_memory_resource.pypython/rmm/rmm/tests/test_device_buffer.pypython/rmm/rmm/tests/test_stream.py
✅ Files skipped from review due to trivial changes (3)
- python/rmm/rmm/tests/test_stream.py
- python/rmm/rmm/tests/test_binning_memory_resource.py
- python/rmm/rmm/tests/test_allocations.py
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python/rmm/rmm/tests/test_device_buffer.py (2)
149-151: Unreachable condition:mv.strides is Noneis always False for Python memoryview.Python's
memoryview.stridesproperty always returns a tuple, neverNone. This condition branch can never be reached with the current test cases.If this was intended to guard against a hypothetical edge case, the test inputs (lines 120-135) don't include any type that would produce a memoryview with
strides is None. Consider removing the unreachable branch or documenting why it's preserved.Suggested simplification (if the None check is unnecessary)
- elif mv.strides is None or len(mv.strides) != 1: + elif len(mv.strides) != 1: with pytest.raises(ValueError): rmm.DeviceBuffer.to_device(hb)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/rmm/rmm/tests/test_device_buffer.py` around lines 149 - 151, The test contains an unreachable branch checking mv.strides is None (memoryview.strides is always a tuple); remove the redundant None check in the condition or replace the whole branch with a single check that asserts len(mv.strides) != 1 before calling rmm.DeviceBuffer.to_device(hb). Update the test around the mv.strides usage (reference mv.strides and rmm.DeviceBuffer.to_device) so it either deletes the `mv.strides is None or` clause or provides a memory object that legitimately exercises the intended edge case if you truly need a None-like behavior.
185-190: Consider adding CuPy array coverage (optional).The tests currently cover
rmm.DeviceBufferand NumbaDeviceNDArray. For comprehensive CUDA array interface testing, consider adding a CuPy array factory:pytest.param( lambda: cupy.array([97, 98, 99], dtype="u1"), marks=pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available"), ),This is outside the scope of this type-annotation PR but would improve array type coverage. Based on learnings: "Flag missing tests for different array types including CuPy and Numba array types."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/rmm/rmm/tests/test_device_buffer.py` around lines 185 - 190, Add an optional CuPy case to the existing pytest.parametrize for "cuda_ary" by inserting a pytest.param that returns a CuPy array via a lambda (e.g., lambda: cupy.array([97,98,99], dtype="u1")) and mark it with pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available") so tests skip when CuPy is absent; ensure you import or reference HAS_CUPY and cupy consistently with the test module and place the new pytest.param alongside the existing rmm.DeviceBuffer and cuda.to_device entries in the "cuda_ary" parameter list.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@python/rmm/rmm/tests/test_device_buffer.py`:
- Around line 149-151: The test contains an unreachable branch checking
mv.strides is None (memoryview.strides is always a tuple); remove the redundant
None check in the condition or replace the whole branch with a single check that
asserts len(mv.strides) != 1 before calling rmm.DeviceBuffer.to_device(hb).
Update the test around the mv.strides usage (reference mv.strides and
rmm.DeviceBuffer.to_device) so it either deletes the `mv.strides is None or`
clause or provides a memory object that legitimately exercises the intended edge
case if you truly need a None-like behavior.
- Around line 185-190: Add an optional CuPy case to the existing
pytest.parametrize for "cuda_ary" by inserting a pytest.param that returns a
CuPy array via a lambda (e.g., lambda: cupy.array([97,98,99], dtype="u1")) and
mark it with pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available") so
tests skip when CuPy is absent; ensure you import or reference HAS_CUPY and cupy
consistently with the test module and place the new pytest.param alongside the
existing rmm.DeviceBuffer and cuda.to_device entries in the "cuda_ary" parameter
list.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9636589f-1249-4f4b-9834-933d1c645b71
📒 Files selected for processing (1)
python/rmm/rmm/tests/test_device_buffer.py
TomAugspurger
commented
Apr 1, 2026
| with pytest.raises(ValueError): | ||
| rmm.DeviceBuffer.to_device(hb) | ||
| elif len(mv.strides) != 1: | ||
| elif mv.strides is None or len(mv.strides) != 1: |
Contributor
Author
There was a problem hiding this comment.
Older versions of python could return None for an empty memoryview. But at least in my version the typestub indicates that .strides can still be None. We can keep this or add a type-ignore. Since this is a test, I think this is fine.
bdice
approved these changes
Apr 1, 2026
Collaborator
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
Thanks @TomAugspurger. I defer to your typing knowledge here, everything seems fine from my side. Please merge when you're ready.
| """Test that DeviceBuffer.__init__ raises TypeError for stream=None""" | ||
| with pytest.raises(TypeError, match="stream argument cannot be None"): | ||
| rmm.DeviceBuffer(size=10, stream=None) | ||
| rmm.DeviceBuffer(size=10, stream=None) # type: ignore[arg-type] |
Collaborator
There was a problem hiding this comment.
What is this type: ignore doing? Is it because we're testing something that isn't allowed to be None?
Contributor
Author
There was a problem hiding this comment.
Exactly: DeviceBuffer(..., stream=None) is invalid since stream is declared with a type Stream rather than Stream | None. This test validates that it raises a TypeError at runtime.
Contributor
Author
|
/merge |
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.
Description
#2331 fixed an issue with RMM using a deprecated type in its type annotations for
CudaStreamPool. Our used the non-deprecatedrmm.pylibrmm.CudaStreamFlags, which was technically incorrect according to mypy but worked at runtime because the two implementations were identical.Our CI (specifically our type-checking pre-commit check) didn't catch these because our tests don't have type annotations. This PR is a POC showing how this should have caught the issue, if we had type annotations on our tests.