diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d71a2ac8f..e0608e114 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -168,10 +168,6 @@ if(BUILD_SLURM_SUPPORT) endif() endif() -if(BUILD_CUDF_TESTS) - include(../cmake/thirdparty/get_cudf_streaming.cmake) -endif() - # ################################################################################################## # * library targets -------------------------------------------------------------------------------- @@ -206,6 +202,7 @@ add_library( src/memory/memory_type.cpp src/memory/pinned_memory_resource.cpp src/memory/scoped_memory_record.cpp + src/memory/spill.cpp src/memory/spill_manager.cpp src/pausable_thread_loop.cpp src/progress_thread.cpp @@ -347,6 +344,10 @@ if(NOT TARGET rapidsmpf::rapidsmpf) add_library(rapidsmpf::rapidsmpf ALIAS rapidsmpf) endif() +if(BUILD_CUDF_TESTS) + include(../cmake/thirdparty/get_cudf_streaming.cmake) +endif() + # ################################################################################################## # * linter configuration --------------------------------------------------------------------------- if(RAPIDSMPF_CLANG_TIDY) diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index c6c918209..25164fcfd 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -320,7 +321,7 @@ rapidsmpf::Duration do_run( for (auto finished_partition : shuffler.local_partitions()) { auto packed_chunks = shuffler.extract(finished_partition); auto output_partition = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES ), stream, diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 30470066e..ab7c46519 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -121,7 +122,7 @@ coro::task broadcast( 0, std::make_unique( cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(result), ctx->br().get(), AllowOverbooking::YES ), stream, @@ -602,7 +603,7 @@ streaming::Actor shuffle( pid, std::make_unique( cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(packed_data), ctx->br().get(), AllowOverbooking::YES ), stream, diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp index 1e322b476..98bcc6868 100644 --- a/cpp/examples/example_shuffle.cpp +++ b/cpp/examples/example_shuffle.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -112,7 +113,7 @@ int main(int argc, char** argv) { // convenience function. local_outputs.push_back( cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(packed_chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, diff --git a/cpp/include/rapidsmpf/memory/spill.hpp b/cpp/include/rapidsmpf/memory/spill.hpp new file mode 100644 index 000000000..39fbf4b60 --- /dev/null +++ b/cpp/include/rapidsmpf/memory/spill.hpp @@ -0,0 +1,64 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include +#include + +namespace rapidsmpf { + +/** + * @brief Spill partitions from device memory to host memory. + * + * Moves the buffer of each `PackedData` from device memory to host memory using + * the provided buffer resource and the buffer's CUDA stream. Partitions that are + * already in host memory are passed through unchanged. + * + * For device-resident partitions, a host memory reservation is made before moving + * the buffer. If the reservation fails due to insufficient host memory, an exception + * is thrown. Overbooking is not allowed. + * + * @param partitions The partitions to spill. + * @param br Buffer resource used to reserve host memory and perform the move. + * + * @return A vector of `PackedData`, where each buffer resides in host memory. + * + * @throws rapidsmpf::reservation_error If host memory reservation fails. + */ +std::vector spill_partitions( + std::vector&& partitions, BufferResource* br +); + +/** + * @brief Move spilled partitions (i.e., packed tables in host memory) back to device + * memory. + * + * Each partition is inspected to determine whether its buffer resides in device memory. + * Buffers already in device memory are left untouched. Host-resident buffers are moved + * to device memory using the provided buffer resource and the buffer's CUDA stream. + * + * If insufficient device memory is available, the buffer resource's spill manager is + * invoked to free memory. If overbooking occurs and spilling fails to reclaim enough + * memory, behavior depends on the `allow_overbooking` flag. + * + * @param partitions The partitions to unspill, potentially containing host-resident data. + * @param br Buffer resource responsible for memory reservation and spills. + * @param allow_overbooking If false, ensures enough memory is freed to satisfy the + * reservation; otherwise, allows overbooking even if spilling was insufficient. + * + * @return A vector of `PackedData`, each with a buffer in device memory. + * + * @throws rapidsmpf::reservation_error If overbooking exceeds the amount spilled and + * `allow_overbooking` is false. + */ +std::vector unspill_partitions( + std::vector&& partitions, + BufferResource* br, + AllowOverbooking allow_overbooking +); + +} // namespace rapidsmpf diff --git a/cpp/src/memory/spill.cpp b/cpp/src/memory/spill.cpp new file mode 100644 index 000000000..f79b33162 --- /dev/null +++ b/cpp/src/memory/spill.cpp @@ -0,0 +1,62 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include + +namespace rapidsmpf { + +std::vector spill_partitions( + std::vector&& partitions, BufferResource* br +) { + // Sum the total size of all packed data in device memory. + std::size_t device_size{0}; + for (auto& [_, data] : partitions) { + if (data->mem_type() == MemoryType::DEVICE) { + device_size += data->size; + } + } + // Spill each partition to host memory. + auto reservation = br->reserve_or_fail(device_size, SPILL_TARGET_MEMORY_TYPES); + std::vector ret; + ret.reserve(partitions.size()); + for (auto& [metadata, data] : partitions) { + ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); + } + return ret; +} + +std::vector unspill_partitions( + std::vector&& partitions, + BufferResource* br, + AllowOverbooking allow_overbooking +) { + auto statistics = br->statistics(); + // Sum the total size of all packed data not in device memory already. + std::size_t non_device_size{0}; + for (auto& [_, data] : partitions) { + if (data->mem_type() != MemoryType::DEVICE) { + non_device_size += data->size; + } + } + + // Unspill each partition. + auto reservation = + br->reserve_device_memory_and_spill(non_device_size, allow_overbooking); + std::vector ret; + ret.reserve(partitions.size()); + for (auto& [metadata, data] : partitions) { + ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); + } + + return ret; +} + +} // namespace rapidsmpf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index dd3f4a8de..12421123a 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -88,6 +88,7 @@ target_sources( test_rmm_resource_adaptor.cpp test_sparse_alltoall.cpp test_spill_manager.cpp + test_spilling.cpp test_statistics.cpp test_stream_ordered_timing.cpp test_system_info.cpp diff --git a/cpp/tests/test_partition.cpp b/cpp/tests/test_partition.cpp index 2100890f7..8ba164e9a 100644 --- a/cpp/tests/test_partition.cpp +++ b/cpp/tests/test_partition.cpp @@ -101,43 +101,3 @@ TEST_P(NumOfPartitions, split_and_pack) { // Compare the input table with the result. CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result); } - -class SpillingTest : public ::testing::Test { - protected: - void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); - stream = cudf::get_default_stream(); - } - - std::shared_ptr br; - rmm::cuda_stream_view stream; -}; - -TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata) { - std::vector metadata{42, 99}; - std::vector payload{10, 20, 30}; - - // Create device input. - std::vector input; - input.push_back(create_packed_data(metadata, payload, stream, br.get())); - - // Device -> Device (moves data) - auto on_gpu = cudf_streaming::integrations::unspill_partitions( - std::move(input), br.get(), AllowOverbooking::YES - ); - ASSERT_EQ(on_gpu.size(), 1); - EXPECT_EQ(on_gpu[0].data->mem_type(), rapidsmpf::MemoryType::DEVICE); - EXPECT_EQ(*on_gpu[0].metadata, metadata); - - // Device -> Host - auto back_on_host = - cudf_streaming::integrations::spill_partitions(std::move(on_gpu), br.get()); - ASSERT_EQ(back_on_host.size(), 1); - EXPECT_EQ(back_on_host[0].data->mem_type(), rapidsmpf::MemoryType::HOST); - EXPECT_EQ(*back_on_host[0].metadata, metadata); - - // Check that contents match original - auto res = br->reserve_or_fail(back_on_host[0].data->size, MemoryType::HOST); - auto actual = br->move_to_host_buffer(std::move(back_on_host[0].data), res); - EXPECT_EQ(actual->copy_to_uint8_vector(), payload); -} diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index cd3dc9d21..9c8786929 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -207,7 +208,7 @@ void test_shuffler( for (auto finished_partition : shuffler.local_partitions()) { auto packed_chunks = shuffler.extract(finished_partition); auto result = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES ), stream, @@ -456,10 +457,9 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { { // Now extract triggers spilling of the partition not being extracted. - std::vector output_chunks = - cudf_streaming::integrations::unspill_partitions( - shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES - ); + std::vector output_chunks = rapidsmpf::unspill_partitions( + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES + ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); // And insert also triggers spilling. We end up with zero device allocations. @@ -470,15 +470,13 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { } // Extract and unspill both partitions. - std::vector out0 = - cudf_streaming::integrations::unspill_partitions( - shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES - ); + std::vector out0 = rapidsmpf::unspill_partitions( + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES + ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); - std::vector out1 = - cudf_streaming::integrations::unspill_partitions( - shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES - ); + std::vector out1 = rapidsmpf::unspill_partitions( + shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES + ); EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); // Disable spilling and insert the first partition. @@ -905,7 +903,7 @@ TEST(Shuffler, concurrent_wait) { EXPECT_NO_THROW(shuffler.wait(wait_timeout)); auto chunks = shuffler.extract(pid); auto result = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, @@ -1001,7 +999,7 @@ TEST(Shuffler, opid_reuse) { for (auto pid : shuffler.local_partitions()) { auto chunks = shuffler.extract(pid); auto result = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, @@ -1098,7 +1096,7 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { for (auto pid : shuffler.local_partitions()) { auto chunks = shuffler.extract(pid); auto result = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, diff --git a/cpp/tests/test_spilling.cpp b/cpp/tests/test_spilling.cpp new file mode 100644 index 000000000..2c2698b85 --- /dev/null +++ b/cpp/tests/test_spilling.cpp @@ -0,0 +1,59 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include + +#include "utils.hpp" + +using namespace rapidsmpf; + +class SpillingTest : public ::testing::Test { + protected: + void SetUp() override { + br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + stream = rmm::cuda_stream_default; + } + + std::shared_ptr br; + rmm::cuda_stream_view stream; +}; + +TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata) { + std::vector metadata{42, 99}; + std::vector payload{10, 20, 30}; + + // Create device input. + std::vector input; + input.push_back(create_packed_data(metadata, payload, stream, br.get())); + + // Device -> Device (moves data) + auto on_gpu = unspill_partitions(std::move(input), br.get(), AllowOverbooking::YES); + ASSERT_EQ(on_gpu.size(), 1); + EXPECT_EQ(on_gpu[0].data->mem_type(), MemoryType::DEVICE); + EXPECT_EQ(*on_gpu[0].metadata, metadata); + + // Device -> Host + auto back_on_host = spill_partitions(std::move(on_gpu), br.get()); + ASSERT_EQ(back_on_host.size(), 1); + EXPECT_EQ(back_on_host[0].data->mem_type(), MemoryType::HOST); + EXPECT_EQ(*back_on_host[0].metadata, metadata); + + // Check that contents match original + auto res = br->reserve_or_fail(back_on_host[0].data->size, MemoryType::HOST); + auto actual = br->move_to_host_buffer(std::move(back_on_host[0].data), res); + EXPECT_EQ(actual->copy_to_uint8_vector(), payload); +} diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py index d907b1023..6be1a0d29 100644 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py @@ -14,7 +14,6 @@ from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, - unspill_partitions, ) from mpi4py import MPI @@ -26,6 +25,7 @@ from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.spill import unspill_partitions from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.shuffler import Shuffler diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py index 4b63df070..8508e0a5a 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py @@ -16,7 +16,6 @@ from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, - unspill_partitions, ) import rmm.mr @@ -24,6 +23,7 @@ from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.spill import unspill_partitions from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.shuffler import Shuffler from rapidsmpf.statistics import Statistics diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py index 34c3dc0fc..743038b55 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py @@ -13,13 +13,13 @@ from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, - unspill_partitions, ) import rmm from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.spill import unspill_partitions from rapidsmpf.shuffler import Shuffler from rapidsmpf.testing import assert_eq diff --git a/python/rapidsmpf/rapidsmpf/memory/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/memory/CMakeLists.txt index 651425f11..e9b3a66ed 100644 --- a/python/rapidsmpf/rapidsmpf/memory/CMakeLists.txt +++ b/python/rapidsmpf/rapidsmpf/memory/CMakeLists.txt @@ -1,13 +1,13 @@ # ================================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ================================================================================= set(cython_modules buffer.pyx buffer_resource.pyx content_description.pyx memory_reservation.pyx packed_data.pyx - pinned_memory_resource.pyx scoped_memory_record.pyx spill_manager.pyx + pinned_memory_resource.pyx scoped_memory_record.pyx spill.pyx spill_manager.pyx ) rapids_cython_create_modules( diff --git a/python/rapidsmpf/rapidsmpf/memory/spill.pxd b/python/rapidsmpf/rapidsmpf/memory/spill.pxd new file mode 100644 index 000000000..b053483ee --- /dev/null +++ b/python/rapidsmpf/rapidsmpf/memory/spill.pxd @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from rapidsmpf.memory.buffer_resource cimport BufferResource + + +cpdef object spill_partitions(object partitions, BufferResource br) +cpdef object unspill_partitions( + object partitions, + BufferResource br, + object allow_overbooking, +) diff --git a/python/rapidsmpf/rapidsmpf/memory/spill.pyi b/python/rapidsmpf/rapidsmpf/memory/spill.pyi new file mode 100644 index 000000000..2810f9ed9 --- /dev/null +++ b/python/rapidsmpf/rapidsmpf/memory/spill.pyi @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Iterable + +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.packed_data import PackedData + +def spill_partitions( + partitions: Iterable[PackedData], + *, + br: BufferResource, +) -> list[PackedData]: ... +def unspill_partitions( + partitions: Iterable[PackedData], + *, + br: BufferResource, + allow_overbooking: bool, +) -> list[PackedData]: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/spill.pyx b/python/rapidsmpf/rapidsmpf/memory/spill.pyx new file mode 100644 index 000000000..36de0e4a9 --- /dev/null +++ b/python/rapidsmpf/rapidsmpf/memory/spill.pyx @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Spill and unspill packed partitions between device and host memory.""" + +from cython.operator cimport dereference as deref +from libcpp.utility cimport move +from libcpp.vector cimport vector + +from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.memory.buffer_resource cimport (AllowOverbooking, + BufferResource, + cpp_BufferResource) +from rapidsmpf.memory.packed_data cimport (PackedData, cpp_PackedData, + packed_data_vector_to_list) + + +cdef vector[cpp_PackedData] _partitions_py_to_cpp(partitions): + cdef vector[cpp_PackedData] ret + for part in partitions: + if not (part).c_obj: + raise ValueError("PackedData was empty") + ret.push_back(move(deref((part).c_obj))) + return move(ret) + + +cdef extern from "" nogil: + cdef vector[cpp_PackedData] cpp_spill_partitions \ + "rapidsmpf::spill_partitions"( + vector[cpp_PackedData] partitions, + cpp_BufferResource* br, + ) except +ex_handler + + cdef vector[cpp_PackedData] cpp_unspill_partitions \ + "rapidsmpf::unspill_partitions"( + vector[cpp_PackedData] partitions, + cpp_BufferResource* br, + AllowOverbooking allow_overbooking, + ) except +ex_handler + + +cpdef object spill_partitions( + object partitions, + BufferResource br, +): + """ + Spill partitions from device memory to host memory. + + Moves the buffer of each ``PackedData`` from device memory to host memory using + the provided buffer resource and the buffer's CUDA stream. Partitions already + in host memory are returned unchanged. + + For device-resident partitions, a host memory reservation is made before moving + the buffer. If the reservation fails due to insufficient host memory, an + exception is raised. Overbooking is not allowed. + + The input partitions are released and are left empty on return. + + Parameters + ---------- + partitions + The partitions to spill. + br + Buffer resource used to reserve host memory and perform the move. + + Returns + ------- + A list of partitions whose buffers reside in host memory. + + Raises + ------ + ReservationError + If host memory reservation fails. + """ + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) + cdef vector[cpp_PackedData] _ret + with nogil: + _ret = cpp_spill_partitions( + move(_partitions), + _br, + ) + return packed_data_vector_to_list(move(_ret), br) + + +cpdef object unspill_partitions( + object partitions, + BufferResource br, + object allow_overbooking, +): + """ + Move spilled partitions back to device memory. + + Each partition is inspected to determine whether its buffer resides in device + memory. Buffers already in device memory are left untouched. Host-resident buffers + are moved to device memory using the provided buffer resource and the buffer's CUDA + stream. + + If insufficient device memory is available, the buffer resource's spill manager is + invoked to free memory. If overbooking occurs and spilling fails to reclaim enough + memory, behavior depends on ``allow_overbooking``. + + The input partitions are released and are left empty on return. + + Parameters + ---------- + partitions + The partitions to unspill, potentially containing host-resident data. + br + Buffer resource responsible for memory reservation and spills. + allow_overbooking + If False, ensures enough memory is freed to satisfy the reservation; + otherwise, allows overbooking even if spilling was insufficient. + + Returns + ------- + A list of partitions whose buffers reside in device memory. + + Raises + ------ + ReservationError + If overbooking exceeds the amount spilled and ``allow_overbooking is False``. + """ + if not isinstance(allow_overbooking, bool): + raise TypeError("allow_overbooking must be a bool") + cdef cpp_BufferResource* _br = br.ptr() + cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) + cdef vector[cpp_PackedData] _ret + cdef AllowOverbooking ab = ( + AllowOverbooking.YES if allow_overbooking else AllowOverbooking.NO + ) + with nogil: + _ret = cpp_unspill_partitions( + move(_partitions), + _br, + ab, + ) + return packed_data_vector_to_list(move(_ret), br) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_partition.py b/python/rapidsmpf/rapidsmpf/tests/test_partition.py index 8c37b5bed..38ce185bc 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_partition.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_partition.py @@ -11,10 +11,8 @@ pytest.importorskip("cudf_streaming") from cudf_streaming.integrations.partition import ( partition_and_pack, - spill_partitions, split_and_pack, unpack_and_concat, - unspill_partitions, ) from rmm.pylibrmm.stream import DEFAULT_STREAM @@ -107,31 +105,3 @@ def test_split_and_pack_unpack_out_of_range( br=br, stream=DEFAULT_STREAM, ) - - -@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_spill_unspill_roundtrip( - device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int -) -> None: - br = BufferResource(device_mr) - expect = _make_table(cols) - partitions = partition_and_pack( - expect, - columns_to_hash=(1,), - num_partitions=num_partitions, - br=br, - stream=DEFAULT_STREAM, - ) - - # Spill roundtrip - spilled = spill_partitions(partitions.values(), br=br) - unspilled = unspill_partitions(spilled, br=br, allow_overbooking=False) - - got = unpack_and_concat( - unspilled, - br=br, - stream=DEFAULT_STREAM, - ) - # Since the row order isn't preserved, we sort the rows by the first column. - assert_eq(expect, got, sort_rows=0) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py index c6c8d08f7..e4cdb6cb8 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py @@ -13,10 +13,10 @@ from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, - unspill_partitions, ) from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.spill import unspill_partitions from rapidsmpf.shuffler import ( Shuffler, ) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_spill.py b/python/rapidsmpf/rapidsmpf/tests/test_spill.py new file mode 100644 index 000000000..7652ebf99 --- /dev/null +++ b/python/rapidsmpf/rapidsmpf/tests/test_spill.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, +) + +from rmm.pylibrmm.stream import DEFAULT_STREAM + +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.spill import spill_partitions, unspill_partitions +from rapidsmpf.testing import assert_eq + +cudf = pytest.importorskip("cudf") + +if TYPE_CHECKING: + import rmm.mr + + +def _make_table(cols: list[list[int]]) -> plc.Table: + # Assigns empty column inputs as int64 + return plc.Table( + [ + plc.Column.from_iterable_of_py(col, plc.DataType(plc.TypeId.INT64)) + for col in cols + ] + ) + + +@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) +@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) +def test_spill_unspill_roundtrip( + device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int +) -> None: + br = BufferResource(device_mr) + expect = _make_table(cols) + partitions = partition_and_pack( + expect, + columns_to_hash=(1,), + num_partitions=num_partitions, + br=br, + stream=DEFAULT_STREAM, + ) + + # Spill roundtrip + spilled = spill_partitions(partitions.values(), br=br) + unspilled = unspill_partitions(spilled, br=br, allow_overbooking=False) + + got = unpack_and_concat( + unspilled, + br=br, + stream=DEFAULT_STREAM, + ) + # Since the row order isn't preserved, we sort the rows by the first column. + assert_eq(expect, got, sort_rows=0)