diff --git a/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp index 0fbab9c2777b..ddbc7c11a38a 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp @@ -502,10 +502,14 @@ int main(int argc, char** argv) // We're only going to measure the last run, so disable initially. stats->disable(); + RAPIDSMPF_EXPECTS(args.pinned_mem_disable || rapidsmpf::is_pinned_memory_resources_supported(), + "pinned host memory is not supported on this system; pass `-L` to disable it.", + std::runtime_error); + auto pinned_pool_properties = + args.pinned_mem_disable ? rapidsmpf::PinnedMemoryDisabled : rapidsmpf::PinnedPoolProperties{}; auto br = rapidsmpf::BufferResource::create( rmm_mr, - args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled - : rapidsmpf::PinnedMemoryResource::make_if_available(), + std::move(pinned_pool_properties), std::move(memory_limits), std::chrono::milliseconds{1}, std::make_shared(16, rmm::cuda_stream::flags::non_blocking), diff --git a/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp index 511b9a022135..1d7f1604498e 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -331,11 +331,14 @@ int main(int argc, char** argv) auto stats = rapidsmpf::Statistics::create(); - auto pinned_mr = args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled - : rapidsmpf::PinnedMemoryResource::make_if_available(); - auto br = rapidsmpf::BufferResource::create( + RAPIDSMPF_EXPECTS(args.pinned_mem_disable || rapidsmpf::is_pinned_memory_resources_supported(), + "pinned host memory is not supported on this system; pass `-L` to disable it.", + std::runtime_error); + auto pinned_pool_properties = + args.pinned_mem_disable ? rapidsmpf::PinnedMemoryDisabled : rapidsmpf::PinnedPoolProperties{}; + auto br = rapidsmpf::BufferResource::create( rmm_mr, - pinned_mr, + std::move(pinned_pool_properties), std::move(memory_limits), std::nullopt, std::make_shared(16, rmm::cuda_stream::flags::non_blocking), @@ -425,7 +428,7 @@ int main(int argc, char** argv) if (args.enable_memory_profiler) { log->print(statistics->report({ .mr = stat_enabled_mr, - .pinned_mr = pinned_mr, + .pinned_mr = br->try_pinned_mr(), .header = "Statistics (of the last run):", })); } else { diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp index 2261b9f4607f..49556081e849 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp @@ -195,10 +195,10 @@ std::pair, std::shared_ptr> cr "noting that this may significantly degrade spilling performance.", std::invalid_argument); + auto pinned_pool_properties = + arguments.no_pinned_host_memory ? PinnedMemoryDisabled : PinnedPoolProperties{}; auto br = BufferResource::create(std::move(mr), - arguments.no_pinned_host_memory - ? PinnedMemoryResource::Disabled - : PinnedMemoryResource::make_if_available(), + std::move(pinned_pool_properties), std::move(memory_limits), arguments.periodic_spill, std::make_shared( diff --git a/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp index 03576f816ef4..8a849bd22d98 100644 --- a/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp @@ -46,7 +46,7 @@ class BaseStreamingFixture : public ::testing::Test { stream = cudf::get_default_stream(); br = rapidsmpf::BufferResource::create( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits)); + mr_cuda, rapidsmpf::PinnedMemoryDisabled, std::move(memory_limits)); ctx = std::make_shared( std::move(options), GlobalEnvironment->comm_->logger(), br); } diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp index 7d8347296583..35719072c67b 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -38,13 +38,18 @@ class StreamingTableChunk : public BaseStreamingFixture, auto stream_pool = std::make_shared(16, rmm::cuda_stream::flags::non_blocking); stream = cudf::get_default_stream(); - br = rapidsmpf::BufferResource::create( - mr_cuda, // device_mr - rapidsmpf::PinnedMemoryResource::make_if_available(), // pinned_mr - memory_limits, // memory_limits - std::chrono::milliseconds{1}, // periodic_spill_check - stream_pool, // stream_pool - rapidsmpf::Statistics::disabled() // statistics + // Enable pinned host memory only when supported; otherwise the non-pinned + // params still run and the PINNED_HOST cases skip in the test bodies. + auto pinned_pool_properties = rapidsmpf::is_pinned_memory_resources_supported() + ? rapidsmpf::PinnedPoolProperties{} + : rapidsmpf::PinnedMemoryDisabled; + br = rapidsmpf::BufferResource::create( + mr_cuda, // device_mr + std::move(pinned_pool_properties), // pinned_pool_properties + memory_limits, // memory_limits + std::chrono::milliseconds{1}, // periodic_spill_check + stream_pool, // stream_pool + rapidsmpf::Statistics::disabled() // statistics ); ctx = std::make_shared( options, GlobalEnvironment->comm_->logger(), br); diff --git a/cpp/libcudf_streaming/tests/test_shuffler.cpp b/cpp/libcudf_streaming/tests/test_shuffler.cpp index 8b4165f2219c..8a9274ed39c9 100644 --- a/cpp/libcudf_streaming/tests/test_shuffler.cpp +++ b/cpp/libcudf_streaming/tests/test_shuffler.cpp @@ -139,8 +139,7 @@ class MemoryLimits_NumPartition memory_limits = std::get<0>(GetParam()); total_num_partitions = std::get<1>(GetParam()); total_num_rows = std::get<2>(GetParam()); - br = rapidsmpf::BufferResource::create( - mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits); + br = rapidsmpf::BufferResource::create(mr(), rapidsmpf::PinnedMemoryDisabled, memory_limits); shuffler = std::make_unique(GlobalEnvironment->comm_, 0, // op_id @@ -267,7 +266,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) // exposed via `device_mr_adaptor()`, so the test can observe per-rank // allocation counts via `get_main_record().num_current_allocs()`. auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref(), - rapidsmpf::PinnedMemoryResource::Disabled, + rapidsmpf::PinnedMemoryDisabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, std::nullopt // disable periodic spill check );