Skip to content

Scope binning memory resource fixture to class in Python tests #2284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 2 commits into from
Mar 11, 2026
Merged

Scope binning memory resource fixture to class in Python tests #2284

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
db4f3f6
Scope binning memory resource fixture to class in Python tests
bdice Mar 10, 2026
475a745
Apply suggestions from code review
bdice Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 43 additions & 26 deletions python/rmm/rmm/tests/test_binning_memory_resource.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0

"""Tests for BinningMemoryResource."""
Expand All @@ -15,28 +15,8 @@
import rmm


@pytest.mark.parametrize("dtype", _dtypes)
@pytest.mark.parametrize("nelem", _nelems)
@pytest.mark.parametrize("alloc", _allocs)
@pytest.mark.parametrize(
"upstream_mr",
[
lambda: rmm.mr.CudaMemoryResource(),
lambda: rmm.mr.ManagedMemoryResource(),
lambda: rmm.mr.PoolMemoryResource(
rmm.mr.CudaMemoryResource(), 1 << 20
),
]
+ (
[
lambda: rmm.mr.SystemMemoryResource(),
lambda: rmm.mr.SamHeadroomMemoryResource(headroom=1 << 20),
]
if _SYSTEM_MEMORY_SUPPORTED
else []
),
)
def test_binning_memory_resource(dtype, nelem, alloc, upstream_mr):
def _make_binning_mr(upstream_mr):
"""Build a BinningMemoryResource from the given upstream factory."""
upstream = upstream_mr()

# Add fixed-size bins 256KiB, 512KiB, 1MiB, 2MiB, 4MiB
Expand All @@ -47,7 +27,44 @@ def test_binning_memory_resource(dtype, nelem, alloc, upstream_mr):
cuda_mr = rmm.mr.CudaMemoryResource()
mr.add_bin(1 << 10, fixed_mr) # 1KiB bin
mr.add_bin(1 << 23, cuda_mr) # 8MiB bin
return mr


_upstream_mrs = [
Comment thread
bdice marked this conversation as resolved.
Outdated
lambda: rmm.mr.CudaMemoryResource(),
lambda: rmm.mr.ManagedMemoryResource(),
lambda: rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource(), 1 << 20),
] + (
[
lambda: rmm.mr.SystemMemoryResource(),
lambda: rmm.mr.SamHeadroomMemoryResource(headroom=1 << 20),
]
if _SYSTEM_MEMORY_SUPPORTED
else []
)


# Create the BinningMemoryResource once per upstream_mr (class-scoped),
# avoiding the expensive ManagedMemoryResource/FixedSizeMemoryResource slab
# allocation on every test invocation.
@pytest.fixture(scope="class", params=_upstream_mrs)
Comment thread
bdice marked this conversation as resolved.
Outdated
def binning_mr(request):
"""Create the BinningMemoryResource once per upstream_mr."""
return _make_binning_mr(request.param)


class TestBinningMemoryResource:
@pytest.fixture(autouse=True)
def _apply_mr(self, binning_mr):
"""Set the device resource before each test."""
rmm.mr.set_current_device_resource(binning_mr)

rmm.mr.set_current_device_resource(mr)
assert rmm.mr.get_current_device_resource_type() is type(mr)
array_tester(dtype, nelem, alloc)
@pytest.mark.parametrize("dtype", _dtypes)
@pytest.mark.parametrize("nelem", _nelems)
@pytest.mark.parametrize("alloc", _allocs)
def test_binning_memory_resource(self, binning_mr, dtype, nelem, alloc):
assert (
rmm.mr.get_current_device_resource_type()
is rmm.mr.BinningMemoryResource
)
array_tester(dtype, nelem, alloc)
Loading