From b24484397f28b97523aba4bd8054f2f23005a8b1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 08:46:28 -0700 Subject: [PATCH 1/9] Remove passing xfails from cudf-pnadas list These tests marked as xfails recently started passing, causing CI failures in jobs like https://github.com/rapidsai/cudf/actions/runs/26497271908/job/78028482500#step:13:1669 --- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index 482ddf104dcf..d976eedf4dec 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -284,10 +284,6 @@ def pytest_unconfigure(config): "tests/arrays/floating/test_to_numpy.py::test_to_numpy[False-series]": "AssertionError: numpy array are different", "tests/arrays/floating/test_to_numpy.py::test_to_numpy_copy": "TODO: Add a reason for failure", "tests/arrays/integer/test_arithmetic.py::test_values_multiplying_large_series_by_NA": "TODO: Add a reason for failure", - "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[ge-UInt64Dtype--1]": "TODO: Add a reason for failure", - "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[gt-UInt64Dtype--1]": "TODO: Add a reason for failure", - "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[le-UInt64Dtype--1]": "TODO: Add a reason for failure", - "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[lt-UInt64Dtype--1]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int16Dtype]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int32Dtype]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int64Dtype]": "TODO: Add a reason for failure", From d80b26629d4a81205449d3be8f2601aba5d3f523 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 10:16:15 -0700 Subject: [PATCH 2/9] Remove another --- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index d976eedf4dec..57fd64f7d7be 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -2792,7 +2792,6 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby.py::test_as_index_select_column": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_group_name_available_in_inference_pass": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_group_on_two_row_multiindex_returns_one_tuple_key": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_all_nan_groups_drop": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_as_index_select_column_sum_empty_df": "AssertionError: Expected type , found instead", "tests/groupby/test_groupby.py::test_groupby_column_index_name_lost[any]": "RuntimeError: Fast-to-slow transfer is blocked", "tests/groupby/test_groupby.py::test_groupby_column_index_name_lost[shift]": "TODO: Add a reason for failure", From 0c5fbc96f24740cee638e44249a2daabc190d9f1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 11:56:56 -0700 Subject: [PATCH 3/9] xfail_strict=false --- python/cudf/cudf/pandas/scripts/run-pandas-tests.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh b/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh index c15cae3a3aaf..57561a1fb9a5 100755 --- a/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh +++ b/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh @@ -54,7 +54,8 @@ if [ ! -d "pandas-tests" ]; then # Vendored from pandas/pyproject.toml cat > pandas-tests/pyproject.toml << \EOF [tool.pytest.ini_options] -xfail_strict = true +# flaky xpasses tracked at https://github.com/rapidsai/cudf/issues/22681 +xfail_strict = false markers = [ "single_cpu: tests that should run on a single cpu only", "slow: mark a test as slow", From 1f1fd8f75f52f1429a1116158a92ffdcbc37b80a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 11:57:10 -0700 Subject: [PATCH 4/9] Revert "Remove another" This reverts commit d80b26629d4a81205449d3be8f2601aba5d3f523. --- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index 57fd64f7d7be..d976eedf4dec 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -2792,6 +2792,7 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby.py::test_as_index_select_column": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_group_name_available_in_inference_pass": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_group_on_two_row_multiindex_returns_one_tuple_key": "TODO: Add a reason for failure", + "tests/groupby/test_groupby.py::test_groupby_all_nan_groups_drop": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_as_index_select_column_sum_empty_df": "AssertionError: Expected type , found instead", "tests/groupby/test_groupby.py::test_groupby_column_index_name_lost[any]": "RuntimeError: Fast-to-slow transfer is blocked", "tests/groupby/test_groupby.py::test_groupby_column_index_name_lost[shift]": "TODO: Add a reason for failure", From 1aec11e9c6f7df690d4a4dfb5d70de651b0298de Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 11:57:39 -0700 Subject: [PATCH 5/9] Revert "Remove passing xfails from cudf-pnadas list" This reverts commit b24484397f28b97523aba4bd8054f2f23005a8b1. --- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index d976eedf4dec..482ddf104dcf 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -284,6 +284,10 @@ def pytest_unconfigure(config): "tests/arrays/floating/test_to_numpy.py::test_to_numpy[False-series]": "AssertionError: numpy array are different", "tests/arrays/floating/test_to_numpy.py::test_to_numpy_copy": "TODO: Add a reason for failure", "tests/arrays/integer/test_arithmetic.py::test_values_multiplying_large_series_by_NA": "TODO: Add a reason for failure", + "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[ge-UInt64Dtype--1]": "TODO: Add a reason for failure", + "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[gt-UInt64Dtype--1]": "TODO: Add a reason for failure", + "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[le-UInt64Dtype--1]": "TODO: Add a reason for failure", + "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_scalar[lt-UInt64Dtype--1]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int16Dtype]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int32Dtype]": "TODO: Add a reason for failure", "tests/arrays/integer/test_comparison.py::TestComparisonOps::test_ufunc_with_out[Int64Dtype]": "TODO: Add a reason for failure", From 2660976f4640a193435f6a98e98947d953c1dd25 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 13:23:52 -0700 Subject: [PATCH 6/9] Vendor rmm aligned_host_allocate and aligned_host_deallocate This vendors some private rmm methods that were recently moved in https://github.com/rapidsai/rmm/pull/2416/. The implementations are *mostly* the same, but they're not templated (generic?) over `Alloc`. --- cpp/src/utilities/host_memory.cpp | 53 ++++++++++++++++++++++++++++--- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index 9bae40e8395f..76bfca2ebcad 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -17,7 +17,9 @@ #include #include +#include #include +#include #include #include #include @@ -27,6 +29,51 @@ namespace cudf { namespace { +// Inlined from RMM internals after public MR definitions moved to source files: +// https://github.com/rapidsai/rmm/pull/2416 +void* aligned_host_allocate(std::size_t bytes, std::size_t alignment) +{ + assert(rmm::is_supported_alignment(alignment)); + + // allocate memory for bytes, plus potential alignment correction, + // plus store of the correction offset + std::size_t padded_allocation_size{bytes + alignment + sizeof(std::ptrdiff_t)}; + char* const original = static_cast(::operator new(padded_allocation_size)); + + // account for storage of offset immediately prior to the aligned pointer + // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic) + void* aligned{original + sizeof(std::ptrdiff_t)}; + + // std::align modifies `aligned` to point to the first aligned location + std::align(alignment, bytes, aligned, padded_allocation_size); + + // Compute the offset between the original and aligned pointers + std::ptrdiff_t const offset = static_cast(aligned) - original; + + // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic) + *(static_cast(aligned) - 1) = offset; + + return aligned; +} + +void aligned_host_deallocate(void* ptr, + [[maybe_unused]] std::size_t bytes, + [[maybe_unused]] std::size_t alignment) noexcept +{ + assert(rmm::is_supported_alignment(alignment)); + + if (ptr != nullptr) { + // Get offset from the location immediately prior to the aligned pointer + // NOLINTNEXTLINE + std::ptrdiff_t const offset = *(reinterpret_cast(ptr) - 1); + + // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic) + void* const original = static_cast(ptr) - offset; + + ::operator delete(original); + } +} + class pinned_pool_with_fallback_memory_resource { using upstream_mr = rmm::mr::pinned_host_memory_resource; using host_pooled_mr = rmm::mr::pool_memory_resource; @@ -219,8 +266,7 @@ class new_delete_memory_resource { void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { try { - return rmm::detail::aligned_host_allocate( - bytes, alignment, [](std::size_t size) { return ::operator new(size); }); + return aligned_host_allocate(bytes, alignment); } catch (std::bad_alloc const& e) { CUDF_FAIL("Failed to allocate memory: " + std::string{e.what()}, rmm::out_of_memory); } @@ -237,8 +283,7 @@ class new_delete_memory_resource { std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { - rmm::detail::aligned_host_deallocate( - ptr, bytes, alignment, [](void* ptr) { ::operator delete(ptr); }); + aligned_host_deallocate(ptr, bytes, alignment); } void deallocate([[maybe_unused]] cuda::stream_ref stream, From b278e508a381e728d015ae1d2c441dd4791f21b9 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 27 May 2026 16:56:44 -0700 Subject: [PATCH 7/9] new API changes Signed-off-by: niranda perera --- python/cudf_polars/cudf_polars/engine/core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 0a22ab9cdeac..2fff0ade80fd 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -584,8 +584,7 @@ def all_gather_host_data( allgather = AllGather( comm=comm, op_id=op_id, - br=br, - statistics=Statistics(enable=False), + br=br ) # TODO: Make AllGather (bulk) a context manager so this becomes # with AllGather(...) as ag: From ac4a7490d5c47e996a20c5929ad244cd27f46d8d Mon Sep 17 00:00:00 2001 From: Niranda Perera Date: Wed, 27 May 2026 17:04:25 -0700 Subject: [PATCH 8/9] style check --- python/cudf_polars/cudf_polars/engine/core.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 2fff0ade80fd..cbfa023c3c06 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -581,11 +581,7 @@ def all_gather_host_data( ------- List of bytes, one element per rank, ordered by rank index. """ - allgather = AllGather( - comm=comm, - op_id=op_id, - br=br - ) + allgather = AllGather(comm=comm, op_id=op_id, br=br) # TODO: Make AllGather (bulk) a context manager so this becomes # with AllGather(...) as ag: # ag.insert(0, PackedData.from_host_bytes(data, br)) From bcf0939bd18ffb392bfa683d447ccf5d52fcec9b Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 27 May 2026 18:38:21 -0700 Subject: [PATCH 9/9] statistics --- python/cudf_polars/cudf_polars/engine/dask.py | 10 +++++++--- python/cudf_polars/cudf_polars/engine/ray.py | 14 +++++++++++--- python/cudf_polars/cudf_polars/engine/spmd.py | 19 +++++++++++++++++-- 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index fcb0b9ce280a..014627b82c48 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -21,6 +21,7 @@ from rapidsmpf.config import Options from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context import polars as pl @@ -44,7 +45,6 @@ from collections.abc import Callable from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from cudf_polars.dsl.ir import IR @@ -240,7 +240,8 @@ def _setup_worker( comm = mp_ctx.comm barrier(comm) - ctx = Context.from_options(comm.logger, mr, options) + statistics = Statistics.from_options(options) + ctx = Context.from_options(comm.logger, mr, options, statistics) # Set the current RMM device resource so all temporary allocations # in libcudf also use the same memory resource. rmm.mr.set_current_device_resource(ctx.br().device_mr) @@ -329,7 +330,10 @@ def _reset_worker( mp_ctx.ctx.shutdown() mp_ctx.ctx = None options = Options.deserialize(rapidsmpf_options_as_bytes) - mp_ctx.ctx = Context.from_options(mp_ctx.comm.logger, mp_ctx.mr, options) + statistics = Statistics.from_options(options) + mp_ctx.ctx = Context.from_options( + mp_ctx.comm.logger, mp_ctx.mr, options, statistics + ) rmm.mr.set_current_device_resource(mp_ctx.ctx.br().device_mr) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 25b5206b514c..edf1905e916d 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -17,6 +17,7 @@ from rapidsmpf.config import Options from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context import polars as pl @@ -41,7 +42,6 @@ from collections.abc import Callable from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from ray.actor import ActorHandle @@ -186,6 +186,7 @@ def __init__( self._rapidsmpf_options: Options = Options.deserialize( rapidsmpf_options_as_bytes ) + self._rapidsmpf_statistics = Statistics.from_options(self._rapidsmpf_options) self._nranks: int = nranks self._py_executor = ThreadPoolExecutor( max_workers=num_py_executors, @@ -242,7 +243,10 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: ) barrier(self._comm) self._ctx = Context.from_options( - self._comm.logger, self._mr, self._rapidsmpf_options + self._comm.logger, + self._mr, + self._rapidsmpf_options, + self._rapidsmpf_statistics, ) # Set the current RMM device resource so all temporary allocations # in libcudf also use the same memory resource. @@ -269,8 +273,12 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes) -> None: self._ctx.shutdown() self._ctx = None self._rapidsmpf_options = Options.deserialize(rapidsmpf_options_as_bytes) + self._rapidsmpf_statistics = Statistics.from_options(self._rapidsmpf_options) self._ctx = Context.from_options( - self._comm.logger, self._mr, self._rapidsmpf_options + self._comm.logger, + self._mr, + self._rapidsmpf_options, + self._rapidsmpf_statistics, ) def shutdown(self) -> None: diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 18804f410ef8..d334e9f1c95c 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -376,7 +376,14 @@ def __init__( # Register `_cleanup_ctx`, which shuts down whatever `self._ctx` points # to at engine shutdown time, i.e. the `Context` from the latest reset. - self._ctx = Context.from_options(comm.logger, mr, self.rapidsmpf_options) + if self.rapidsmpf_options is not None: + statistics = Statistics.from_options(self.rapidsmpf_options) + else: + statistics = None + + self._ctx = Context.from_options( + comm.logger, mr, self.rapidsmpf_options, statistics + ) exit_stack.callback(self._cleanup_ctx) # Register after `_cleanup_ctx` so on teardown (LIFO) the @@ -486,7 +493,15 @@ def _reset( # Context (the test driver's main thread). The per-engine RMM # resource is kept alive across resets, see :meth:`_cleanup_ctx`. self._ctx.shutdown() - self._ctx = Context.from_options(self._comm.logger, self._mr, rapidsmpf_options) + + if rapidsmpf_options is not None: + statistics = Statistics.from_options(rapidsmpf_options) + else: + statistics = None + + self._ctx = Context.from_options( + self._comm.logger, self._mr, rapidsmpf_options, statistics + ) # Re-run ``StreamingEngine.__init__`` on the existing instance to # reconfigure the polars ``GPUEngine`` layer (``self.config``,