diff --git a/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp b/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp index 36e47d111d17..6d0f06208e16 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/integrations/partition.hpp @@ -138,52 +138,4 @@ partition_and_pack( rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); -/** - * @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, - rapidsmpf::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, - rapidsmpf::BufferResource* br, - rapidsmpf::AllowOverbooking allow_overbooking); - } // namespace cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/src/integrations/partition.cpp b/cpp/libcudf_streaming/src/integrations/partition.cpp index f9af79f6fd76..f258273672ad 100644 --- a/cpp/libcudf_streaming/src/integrations/partition.cpp +++ b/cpp/libcudf_streaming/src/integrations/partition.cpp @@ -183,44 +183,4 @@ std::unique_ptr unpack_and_concat(std::vectordevice_mr()); } -std::vector spill_partitions(std::vector&& partitions, - rapidsmpf::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() == rapidsmpf::MemoryType::DEVICE) { device_size += data->size; } - } - // Spill each partition to host memory. - auto reservation = br->reserve_or_fail(device_size, rapidsmpf::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, - rapidsmpf::BufferResource* br, - rapidsmpf::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() != rapidsmpf::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 cudf_streaming::integrations diff --git a/cpp/libcudf_streaming/src/streaming/partition.cpp b/cpp/libcudf_streaming/src/streaming/partition.cpp index b39e39d68e66..7ef41cbc1804 100644 --- a/cpp/libcudf_streaming/src/streaming/partition.cpp +++ b/cpp/libcudf_streaming/src/streaming/partition.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -76,7 +77,7 @@ rapidsmpf::streaming::Actor unpack_and_concat(std::shared_ptrbr()->stream_pool().get_stream(); std::unique_ptr ret = cudf_streaming::integrations::unpack_and_concat( - cudf_streaming::integrations::unspill_partitions( + rapidsmpf::unspill_partitions( std::move(data), ctx->br().get(), rapidsmpf::AllowOverbooking::NO), stream, ctx->br().get()); diff --git a/cpp/libcudf_streaming/tests/test_partition.cpp b/cpp/libcudf_streaming/tests/test_partition.cpp index ede097a15c4f..d261968a27bd 100644 --- a/cpp/libcudf_streaming/tests/test_partition.cpp +++ b/cpp/libcudf_streaming/tests/test_partition.cpp @@ -91,42 +91,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 = rapidsmpf::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 = unspill_partitions(std::move(input), br.get(), rapidsmpf::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 = 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, rapidsmpf::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/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index 8c942db2f02e..bee723fb413a 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -152,21 +152,6 @@ template return sort_table(table->view(), column_indices); } -/// @brief Create a PackedData object from a host buffer -[[nodiscard]] inline rapidsmpf::PackedData create_packed_data( - std::span metadata, - std::span data, - rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br) -{ - auto metadata_ptr = std::make_unique>(metadata.begin(), metadata.end()); - - auto reservation = - br->reserve(rapidsmpf::MemoryType::DEVICE, data.size(), rapidsmpf::AllowOverbooking::YES); - auto data_ptr = std::make_unique(data.data(), data.size(), stream); - return rapidsmpf::PackedData{std::move(metadata_ptr), br->move(std::move(data_ptr), stream)}; -} - /** * @brief Generate a packed data object with the given number of elements and offset. * diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pxd b/python/cudf_streaming/cudf_streaming/integrations/partition.pxd index 9d674b1bd4c0..91d88bf694f4 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/partition.pxd +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pxd @@ -22,12 +22,6 @@ cpdef object split_and_pack( BufferResource br, ) cpdef object unpack_and_concat(object partitions, Stream stream, BufferResource br) -cpdef object spill_partitions(object partitions, BufferResource br) -cpdef object unspill_partitions( - object partitions, - BufferResource br, - object allow_overbooking, -) cpdef object packed_data_from_cudf_packed_columns( PackedColumns packed_columns, Stream stream, diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pyi b/python/cudf_streaming/cudf_streaming/integrations/partition.pyi index efb30be09176..39ec33c51bda 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/partition.pyi +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyi @@ -29,17 +29,6 @@ def unpack_and_concat( stream: Stream, br: BufferResource, ) -> Table: ... -def spill_partitions( - partitions: Iterable[PackedData], - *, - br: BufferResource, -) -> list[PackedData]: ... -def unspill_partitions( - partitions: Iterable[PackedData], - *, - br: BufferResource, - allow_overbooking: bool, -) -> list[PackedData]: ... def packed_data_from_cudf_packed_columns( packed_columns: PackedColumns, stream: Stream, diff --git a/python/cudf_streaming/cudf_streaming/integrations/partition.pyx b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx index f9fdc5b62203..29f281316f8b 100644 --- a/python/cudf_streaming/cudf_streaming/integrations/partition.pyx +++ b/python/cudf_streaming/cudf_streaming/integrations/partition.pyx @@ -19,11 +19,8 @@ from rmm.librmm.device_buffer cimport device_buffer from rmm.pylibrmm.stream cimport Stream 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) +from rapidsmpf.memory.buffer_resource cimport BufferResource, cpp_BufferResource +from rapidsmpf.memory.packed_data cimport PackedData, cpp_PackedData cdef extern from "" nogil: @@ -256,68 +253,6 @@ cpdef object unpack_and_concat( ) return Table.from_libcudf(move(_ret), stream, br._device_mr) - -cdef extern from "" nogil: - cdef vector[cpp_PackedData] cpp_spill_partitions \ - "cudf_streaming::integrations::spill_partitions"( - vector[cpp_PackedData] partitions, - cpp_BufferResource* br, - ) 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) - - -cdef extern from "" nogil: - cdef vector[cpp_PackedData] cpp_unspill_partitions \ - "cudf_streaming::integrations::unspill_partitions"( - vector[cpp_PackedData] partitions, - cpp_BufferResource* br, - AllowOverbooking allow_overbooking, - ) except +ex_handler - - cdef extern from *: """ #include @@ -344,61 +279,6 @@ cdef extern from *: ) except +ex_handler nogil -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) - - cpdef object packed_data_from_cudf_packed_columns( PackedColumns packed_columns, Stream stream,