Moving spill/unspill_partitions from cudf_streaming - #1087
Merged
rapids-bot[bot] merged 3 commits intoJun 9, 2026
Conversation
Signed-off-by: niranda perera <niranda.perera@gmail.com>
pentschev
added
bug
This was referenced Jun 8, 2026
Signed-off-by: niranda perera <niranda.perera@gmail.com>
nirandaperera
force-pushed
the
moving-back-spill-unspill
branch
from
June 9, 2026 05:13
fb1c467 to
6ea2bf9
Compare
Contributor
Author
|
/merge |
8 tasks
rapids-bot Bot
pushed a commit
to NVIDIA/cudf
that referenced
this pull request
Jun 10, 2026
Removed `spill/unspill_partitions` from partition.hpp and use rapidsmpf utils. Depends on rapidsai/rapidsmpf#1087 Authors: - Niranda Perera (https://github.com/nirandaperera) Approvers: - Peter Andreas Entschev (https://github.com/pentschev) URL: #22810
rapids-bot Bot
pushed a commit
that referenced
this pull request
Jun 11, 2026
## `[C++]` Make shuffler distributed tests cudf-free ## Summary The shuffler round-trip tests depended on cudf for input generation and result checking (`random_table_with_index`, `partition_and_pack`/`unpack_and_concat`, `CUDF_TEST_EXPECT_TABLES_EQUIVALENT`, `cudf::test::BaseFixture`), which made them heavy and coupled to the cudf integration. This replaces that with a self-contained, conservation-preserving data model built directly from `PackedData`, so the shuffler tests no longer depend on cudf and run as part of the default (non-`BUILD_CUDF_TESTS`) build. ## Changes ### C++ utilities - Add `ceil_div` and a `chunk_indices` view to `utils/misc.hpp`. `chunk_indices` tiles `[0, count)` into exactly `num_chunks` contiguous (possibly empty) ranges whose sizes sum to `count`. - `ceil_div` is computed as `x / y + (x % y != 0)` rather than the naive `(x + y - 1) / y`, which overflows for large unsigned values and is UB on signed overflow near the type maximum. - Add compile-time edge-case coverage for `ceil_div` in `test_misc.cpp` (`static_assert`s for exact/remainder division, zero numerator, denominator of one, and values at `uint64`/`int64` maxima). ### C++ shuffler tests (now cudf-free) - Rework the shuffler tests (`test_shuffler.cpp`, `streaming/test_shuffler.cpp`) around `make_partition_data` (produces routed chunks) and `validate_partition_data` (verifies received chunks), dropping all `cudf`/`cudf_test` usage and switching to `rmm::` stream/MR helpers and `::testing::Test`. - Generalize the test helpers `generate_packed_data`/`validate_packed_data` in `tests/utils.hpp` over the element type and build device buffers via the buffer resource + `cuda_memcpy_async` instead of cudf. - De-couple `streaming/test_allgather.cpp` from `cudf_streaming/integrations/partition.hpp`. - `cpp/tests/CMakeLists.txt`: move `test_shuffler.cpp`, `streaming/test_shuffler.cpp`, and `streaming/test_allgather.cpp` out of the `BUILD_CUDF_TESTS` block into the default build; the cudf block now only carries `test_partition.cpp`, `test_shuffler_many_streams.cpp`, and `streaming/test_leaf_actor.cpp`. ### Python (now cudf-free) - Add a `PackedData.from_device_buffer(gpu_data, metadata, stream, br)` factory (plus `.pyi` stub) so packed data can be built from an rmm `DeviceBuffer` without cudf. - Add shared, cudf-free shuffle test helpers to `rapidsmpf/testing.py`: `chunk_indices`, `generate_packed_data`, `validate_packed_data`, `make_partition_data`, and `validate_partition_data`. - Rewrite the Python shuffler tests (`tests/test_shuffler.py`, `tests/streaming/test_shuffler.py`) around those helpers and the `Shuffler`/`ShufflerAsync` object interfaces, removing the `cudf` / `cudf_streaming` / `pylibcudf` dependencies. Depends on #1087 Authors: - Niranda Perera (https://github.com/nirandaperera) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) - Peter Andreas Entschev (https://github.com/pentschev) URL: #1084
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Moves the cuDF-independent spill/unspill helpers from
cudf_streaming::integrations::partition(in thecudfrepo) back intolibrapidsmpf. They only needPackedDataandBufferResource, so they don't belong behind a libcudf dependency.Changes
C++
cpp/include/rapidsmpf/memory/spill.hppandcpp/src/memory/spill.cppdefiningrapidsmpf::spill_partitions/rapidsmpf::unspill_partitions(wired intocpp/CMakeLists.txt).cpp/tests/test_spilling.cppfor a roundtrip test (extracted fromtest_partition.cpp).cpp/tests/test_shuffler.cpp,cpp/examples/example_shuffle.cpp,cpp/benchmarks/bench_shuffle.cpp, andcpp/benchmarks/streaming/ndsh/join.cppswitched fromcudf_streaming::integrations::unspill_partitionstorapidsmpf::unspill_partitions, with the new header included.Python
python/rapidsmpf/rapidsmpf/memory/spill.{pyx,pxd,pyi}exposingspill_partitions/unspill_partitionsviarapidsmpf.memory.spill(added tomemory/CMakeLists.txt).python/rapidsmpf/rapidsmpf/tests/test_spill.py(extracted fromtest_partition.py).bulk_mpi_shuffle.py,ray/bulk_ray_shuffle.py,ray/ray_shuffle_example.py, andtests/test_shuffler.pyswitched to importing fromrapidsmpf.memory.spill.