diff --git a/python/REVIEW_GUIDELINES.md b/python/REVIEW_GUIDELINES.md index 274af488b..d4847c517 100644 --- a/python/REVIEW_GUIDELINES.md +++ b/python/REVIEW_GUIDELINES.md @@ -58,7 +58,6 @@ For general development guidance including build commands, test commands, code s - Missing edge case coverage (zero-size, alignment) - **Using external datasets** (tests must not depend on external resources) - Missing tests for different array types (CuPy, Numba) -- **Using test classes instead of standalone functions** (RMM prefers `test_foo_bar()` functions over `class TestFoo`) ### Documentation - Missing or incorrect docstrings for public methods @@ -189,7 +188,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): - Test allocation/deallocation pairs - Test edge cases: zero-size, large allocations - Test different array types -- Use standalone `test_foo_bar()` functions, not test classes +- Prefer standalone `test_foo_bar()` functions; test classes are acceptable when needed for class-scoped fixture sharing - Use synthetic data, never external resources --- @@ -253,7 +252,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): - [ ] Are allocation/deallocation pairs tested? - [ ] Are edge cases tested (zero-size, large allocations)? - [ ] Are different array types tested? -- [ ] Are tests written as standalone functions? --- diff --git a/python/rmm/rmm/tests/test_binning_memory_resource.py b/python/rmm/rmm/tests/test_binning_memory_resource.py index cffbe91df..19f34824b 100644 --- a/python/rmm/rmm/tests/test_binning_memory_resource.py +++ b/python/rmm/rmm/tests/test_binning_memory_resource.py @@ -3,37 +3,43 @@ """Tests for BinningMemoryResource.""" +import numpy as np import pytest from test_helpers import ( _SYSTEM_MEMORY_SUPPORTED, _allocs, _dtypes, - _nelems, array_tester, ) import rmm +# Sized so float64 allocations (1 KiB - 128 KiB) hit distinct auto bins. +_BINNING_NELEMS = [128, 256, 512, 1024, 4096, 16384] + +# 16777216 float64 elements = 128 MiB; exercises the explicit CudaMR bin. +_LARGE_NELEM = 16777216 + 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 - mr = rmm.mr.BinningMemoryResource(upstream, 18, 22) + # Auto-create fixed-size bins: 1 KiB, 2 KiB, 4 KiB, … 128 KiB + mr = rmm.mr.BinningMemoryResource(upstream, 10, 17) - # Test adding some explicit bin MRs + # Test adding explicit bin MRs fixed_mr = rmm.mr.FixedSizeMemoryResource(upstream, 1 << 10) cuda_mr = rmm.mr.CudaMemoryResource() - mr.add_bin(1 << 10, fixed_mr) # 1KiB bin - mr.add_bin(1 << 23, cuda_mr) # 8MiB bin + mr.add_bin(1 << 10, fixed_mr) # 1 KiB bin (replaces auto bin) + mr.add_bin(1 << 27, cuda_mr) # 128 MiB bin return mr _UPSTREAM_MRS = [ lambda: rmm.mr.CudaMemoryResource(), lambda: rmm.mr.ManagedMemoryResource(), - lambda: rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource(), 1 << 20), + lambda: rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource(), 1 << 28), ] + ( [ lambda: rmm.mr.SystemMemoryResource(), @@ -60,7 +66,7 @@ def _apply_mr(self, binning_mr): rmm.mr.set_current_device_resource(binning_mr) @pytest.mark.parametrize("dtype", _dtypes) - @pytest.mark.parametrize("nelem", _nelems) + @pytest.mark.parametrize("nelem", _BINNING_NELEMS) @pytest.mark.parametrize("alloc", _allocs) def test_binning_memory_resource(self, binning_mr, dtype, nelem, alloc): assert ( @@ -68,3 +74,12 @@ def test_binning_memory_resource(self, binning_mr, dtype, nelem, alloc): is rmm.mr.BinningMemoryResource ) array_tester(dtype, nelem, alloc) + + @pytest.mark.parametrize("alloc", _allocs) + def test_binning_large_allocation(self, binning_mr, alloc): + """Allocate 128 MiB to exercise the explicit CudaMemoryResource bin.""" + assert ( + rmm.mr.get_current_device_resource_type() + is rmm.mr.BinningMemoryResource + ) + array_tester(np.float64, _LARGE_NELEM, alloc) diff --git a/python/rmm/rmm/tests/test_fixed_size_memory_resource.py b/python/rmm/rmm/tests/test_fixed_size_memory_resource.py index 86c2e9371..1d8659459 100644 --- a/python/rmm/rmm/tests/test_fixed_size_memory_resource.py +++ b/python/rmm/rmm/tests/test_fixed_size_memory_resource.py @@ -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 FixedSizeMemoryResource.""" @@ -15,28 +15,49 @@ import rmm -@pytest.mark.parametrize("dtype", _dtypes) -@pytest.mark.parametrize("nelem", _nelems) -@pytest.mark.parametrize("alloc", _allocs) -@pytest.mark.parametrize( - "upstream", +def _make_fixed_size_mr(upstream): + """Build a FixedSizeMemoryResource from the given upstream factory.""" + return rmm.mr.FixedSizeMemoryResource( + upstream(), block_size=1 << 20, blocks_to_preallocate=128 + ) + + +_UPSTREAMS = [ + lambda: rmm.mr.CudaMemoryResource(), + lambda: rmm.mr.ManagedMemoryResource(), +] + ( [ - lambda: rmm.mr.CudaMemoryResource(), - lambda: rmm.mr.ManagedMemoryResource(), + lambda: rmm.mr.SystemMemoryResource(), + lambda: rmm.mr.SamHeadroomMemoryResource(headroom=1 << 20), ] - + ( - [ - lambda: rmm.mr.SystemMemoryResource(), - lambda: rmm.mr.SamHeadroomMemoryResource(headroom=1 << 20), - ] - if _SYSTEM_MEMORY_SUPPORTED - else [] - ), + if _SYSTEM_MEMORY_SUPPORTED + else [] ) -def test_fixed_size_memory_resource(dtype, nelem, alloc, upstream): - mr = rmm.mr.FixedSizeMemoryResource( - upstream(), block_size=1 << 20, blocks_to_preallocate=128 - ) - rmm.mr.set_current_device_resource(mr) - assert rmm.mr.get_current_device_resource_type() is type(mr) - array_tester(dtype, nelem, alloc) + + +# Create the FixedSizeMemoryResource once per upstream (class-scoped), +# avoiding the expensive ManagedMemoryResource slab allocation on every +# test invocation. +@pytest.fixture(scope="class", params=_UPSTREAMS) +def fixed_size_mr(request): + """Create the FixedSizeMemoryResource once per upstream.""" + return _make_fixed_size_mr(request.param) + + +class TestFixedSizeMemoryResource: + @pytest.fixture(autouse=True) + def _apply_mr(self, fixed_size_mr): + """Set the device resource before each test.""" + rmm.mr.set_current_device_resource(fixed_size_mr) + + @pytest.mark.parametrize("dtype", _dtypes) + @pytest.mark.parametrize("nelem", _nelems) + @pytest.mark.parametrize("alloc", _allocs) + def test_fixed_size_memory_resource( + self, fixed_size_mr, dtype, nelem, alloc + ): + assert ( + rmm.mr.get_current_device_resource_type() + is rmm.mr.FixedSizeMemoryResource + ) + array_tester(dtype, nelem, alloc)