Various cuDF CI fixes

cpp/src/utilities/host_memory.cpp

@@ -17,7 +17,9 @@
 
 #include <algorithm>
 #include <atomic>
+#include <cassert>
 #include <cstdlib>
+#include <memory>
 #include <mutex>
 #include <optional>
 #include <shared_mutex>
@@ -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<char*>(::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<char*>(aligned) - original;
+
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
+  *(static_cast<std::ptrdiff_t*>(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<std::ptrdiff_t*>(ptr) - 1);
+
+    // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
+    void* const original = static_cast<char*>(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,

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",

python/cudf_polars/cudf_polars/engine/core.py

@@ -581,12 +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,
-        statistics=Statistics(enable=False),
-    )
+    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))

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)
 
 

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:

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``,