From c0178781a66aa72101ede4a226beb55695d30b68 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 27 Jul 2026 23:39:50 +0000 Subject: [PATCH 1/5] Add cucascade::io::datasource engine wrappers and Python bindings --- .../cucascade/cudf/rest_datasource_engine.hpp | 92 +++++++ .../cudf/uring_datasource_engine.hpp | 82 ++++++ python/CMakeLists.txt | 60 +++++ python/README.md | 78 ++++++ python/cucascade/CMakeLists.txt | 31 +++ python/cucascade/__init__.py | 6 + python/cucascade/datasource.pxd | 47 ++++ python/cucascade/datasource.pyi | 39 +++ python/cucascade/datasource.pyx | 254 ++++++++++++++++++ python/pdsh_results.jsonl | 1 + python/pyproject.toml | 27 ++ src/cudf/CMakeLists.txt | 7 +- src/cudf/rest_datasource_engine.cpp | 53 ++++ src/cudf/uring_datasource_engine.cpp | 35 +++ 14 files changed, 810 insertions(+), 2 deletions(-) create mode 100644 include/cucascade/cudf/rest_datasource_engine.hpp create mode 100644 include/cucascade/cudf/uring_datasource_engine.hpp create mode 100644 python/CMakeLists.txt create mode 100644 python/README.md create mode 100644 python/cucascade/CMakeLists.txt create mode 100644 python/cucascade/__init__.py create mode 100644 python/cucascade/datasource.pxd create mode 100644 python/cucascade/datasource.pyi create mode 100644 python/cucascade/datasource.pyx create mode 100644 python/pdsh_results.jsonl create mode 100644 python/pyproject.toml create mode 100644 src/cudf/rest_datasource_engine.cpp create mode 100644 src/cudf/uring_datasource_engine.cpp diff --git a/include/cucascade/cudf/rest_datasource_engine.hpp b/include/cucascade/cudf/rest_datasource_engine.hpp new file mode 100644 index 0000000..679b53e --- /dev/null +++ b/include/cucascade/cudf/rest_datasource_engine.hpp @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace cucascade::io { + +class ioctx; + +/** + * @brief Self-contained REST datasource engine for S3/HTTP object-store reads. + * + * Owns a NUMA-local pinned host staging pool and a pool of libcurl reactor + * threads backed by SigV4 presigned-URL signing. Callers open individual + * @c datasource instances via @c open(); each datasource shares the engine's + * @c ioctx and memory pool but carries its own per-scan @c prefetching_handle. + * + * The engine must outlive every datasource it produces. + * + * @code{.cpp} + * auto engine = cucascade::io::rest_datasource_engine::make_s3(...); + * auto ds = engine->open("s3://my-bucket/data/lineitem.parquet"); + * ds->fadvise(byte_ranges, device_id); + * @endcode + */ +class rest_datasource_engine { + public: + static constexpr std::size_t default_block_size = 1UL << 20; + static constexpr std::size_t default_pool_capacity = 20UL * 128UL * (1UL << 20); + + /** + * @brief Construct an S3-backed REST engine with SigV4 presigned-URL signing. + * + * Credentials are static for the lifetime of the engine. For short-lived STS + * tokens, reconstruct the engine before the token expires. + * + * @param access_key_id AWS access key ID. + * @param secret_access_key AWS secret access key. + * @param session_token STS session token; empty for long-lived credentials. + * @param region AWS region (e.g. @c "us-east-1"). + * @param endpoint S3-compatible endpoint host (e.g. @c "s3.amazonaws.com" + * or a MinIO host:port). Leave empty to derive from region. + * @param n_reactors Number of libcurl reactor threads. + * @param tls_verify Whether to verify TLS peer certificates. + * @param pool_capacity Total capacity of the pinned host staging pool in bytes. + * @param block_size Fixed block size in bytes for the staging pool. + * @return A ready-to-use engine. The engine's @c ioctx is started before returning. + */ + explicit rest_datasource_engine(std::string access_key_id, + std::string secret_access_key, + std::string session_token, + std::string region, + std::string endpoint, + std::size_t n_reactors = 4, + bool tls_verify = true, + std::size_t pool_capacity = default_pool_capacity, + std::size_t block_size = default_block_size); + + ~rest_datasource_engine(); + + rest_datasource_engine(rest_datasource_engine const&) = delete; + rest_datasource_engine& operator=(rest_datasource_engine const&) = delete; + + /** + * @brief Open a datasource for the given S3 URI. + * + * Issues an HTTP HEAD request to resolve the object size. + * + * @param path S3 URI of the form @c "s3://bucket/key". + * @return A @c datasource bound to this engine's @c ioctx. The returned + * datasource must not outlive this engine. + * @throw std::runtime_error if the HEAD request fails or the URI is malformed. + */ + [[nodiscard]] std::unique_ptr open(std::string path) const; + + private: + cucascade::memory::numa_region_pinned_host_memory_resource _upstream; + cucascade::memory::fixed_size_host_memory_resource _host_mr; + std::shared_ptr _io_ctx; +}; + +} // namespace cucascade::io diff --git a/include/cucascade/cudf/uring_datasource_engine.hpp b/include/cucascade/cudf/uring_datasource_engine.hpp new file mode 100644 index 0000000..765b6d7 --- /dev/null +++ b/include/cucascade/cudf/uring_datasource_engine.hpp @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cucascade::io { + +/** + * @brief Self-contained io_uring datasource engine for local NVMe reads. + * + * Owns the full stack needed for O_DIRECT io_uring reads: a NUMA-local pinned + * host memory pool, a pool of io_uring reactor threads, and an @c ioctx. + * Callers open individual @c datasource instances via @c open(); each datasource + * shares the engine's @c ioctx and memory pool but carries its own per-scan + * @c prefetching_handle. + * + * The engine must outlive every datasource it produces. + * + * @code{.cpp} + * cucascade::io::uring_datasource_engine engine; + * auto ds = engine.open("/mnt/nvme/data/lineitem.parquet"); + * ds->fadvise(byte_ranges, device_id); + * // ... read through ds as a cudf::io::datasource ... + * @endcode + */ +class uring_datasource_engine { + public: + static constexpr std::size_t default_block_size = 1UL << 20; ///< 1 MiB + static constexpr std::size_t default_pool_capacity = 20UL * 128UL * (1UL << 20); ///< ~2.5 GiB + + /** + * @brief Construct a uring datasource engine. + * + * @param n_reactors Number of io_uring reactor threads. + * @param pool_capacity Total capacity of the pinned host staging pool in bytes. + * @param block_size Size of each fixed-size block in the staging pool in bytes. + * Must be a power of two and at least the alignment required + * by O_DIRECT on the target filesystem. + * @param use_odirect Whether to open files with @c O_DIRECT (bypasses page cache). + * @param numa_node NUMA node from which to allocate the pinned staging pool. + */ + explicit uring_datasource_engine(std::size_t n_reactors = 2, + std::size_t pool_capacity = default_pool_capacity, + std::size_t block_size = default_block_size, + bool use_odirect = true, + int numa_node = 0); + + ~uring_datasource_engine(); + + uring_datasource_engine(uring_datasource_engine const&) = delete; + uring_datasource_engine& operator=(uring_datasource_engine const&) = delete; + + /** + * @brief Open a datasource for the given local file path. + * + * @param path Absolute or relative path to a local file. + * @return A @c datasource bound to this engine's @c ioctx. The returned + * datasource must not outlive this engine. + * @throw std::runtime_error if the file cannot be opened. + */ + [[nodiscard]] std::unique_ptr open(std::string path) const; + + private: + cucascade::memory::numa_region_pinned_host_memory_resource _upstream; + cucascade::memory::fixed_size_host_memory_resource _host_mr; + std::shared_ptr _reactor_ctx; + std::shared_ptr _io_ctx; +}; + +} // namespace cucascade::io diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt new file mode 100644 index 0000000..b100808 --- /dev/null +++ b/python/CMakeLists.txt @@ -0,0 +1,60 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +cmake_minimum_required(VERSION 4.0 FATAL_ERROR) + +project( + cucascade_python + VERSION 0.1.0 + LANGUAGES CXX +) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +find_package(PkgConfig REQUIRED) +pkg_check_modules(LIBURING REQUIRED IMPORTED_TARGET liburing) +pkg_check_modules(CURL REQUIRED IMPORTED_TARGET libcurl) +find_package(OpenSSL REQUIRED) + +find_package(cudf REQUIRED) +find_package(cuCascade REQUIRED) + +set(rapids-cmake-dir + "" + CACHE PATH "Optional path to an already-fetched rapids-cmake source tree (skips download)" +) +if(rapids-cmake-dir) + list(APPEND CMAKE_MODULE_PATH "${rapids-cmake-dir}/rapids-cmake") +else() + set(rapids-cmake-version "24.12") + include(FetchContent) + FetchContent_Declare( + rapids-cmake + GIT_REPOSITORY https://github.com/rapidsai/rapids-cmake.git + GIT_TAG "branch-${rapids-cmake-version}" + GIT_SHALLOW TRUE + ) + FetchContent_MakeAvailable(rapids-cmake) + list(APPEND CMAKE_MODULE_PATH "${rapids-cmake_SOURCE_DIR}/rapids-cmake") +endif() + +execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import pylibcudf, pathlib; print(pathlib.Path(pylibcudf.__file__).parent.parent)" + OUTPUT_VARIABLE PYLIBCUDF_SOURCE_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + COMMAND_ERROR_IS_FATAL ANY +) + +include(rapids-cython-core) +rapids_cython_init() + +set(CYTHON_FLAGS "${CYTHON_FLAGS} -I${PYLIBCUDF_SOURCE_DIR}" CACHE STRING "" FORCE) +message(STATUS "cucascade: pylibcudf source dir: ${PYLIBCUDF_SOURCE_DIR}") + +add_subdirectory(cucascade) diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..42525af --- /dev/null +++ b/python/README.md @@ -0,0 +1,78 @@ +# cucascade Python bindings + +Python bindings for cuCascade's cudf datasource layer, exposing `UringEngine` +(local NVMe via io_uring) and `RestEngine` (S3/HTTP) with advisory prefetch +support via `CuCascadeDatasource.fadvise()`. + +## Prerequisites + +- CUDA toolkit +- cudf built or installed (cmake config must be findable) +- cuCascade C++ libraries built and installed +- pylibcudf installed +- Python >= 3.11 + +## 1. Build and install cuCascade C++ + +```bash +CUCASCADE_SRC=/path/to/cuCascade +CUCASCADE_BUILD=${CUCASCADE_SRC}/build +CUDF_CMAKE_DIR=/path/to/cudf/build # directory containing cudf-config.cmake + +cmake -S "${CUCASCADE_SRC}" -B "${CUCASCADE_BUILD}" \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=native \ + -DCMAKE_INSTALL_PREFIX="${CONDA_PREFIX}" \ + -DCUCASCADE_BUILD_CUDF=ON \ + -DCUCASCADE_BUILD_IO=ON \ + -DCUCASCADE_BUILD_TESTS=OFF \ + -DCUCASCADE_BUILD_BENCHMARKS=OFF \ + -DCUCASCADE_BUILD_SHARED_LIBS=ON \ + -DCUCASCADE_BUILD_STATIC_LIBS=OFF \ + -Dcudf_DIR="${CUDF_CMAKE_DIR}" \ + -DCMAKE_PREFIX_PATH="${CONDA_PREFIX}" + +cmake --build "${CUCASCADE_BUILD}" -j$(nproc) +cmake --install "${CUCASCADE_BUILD}" +``` + +Install to `${CONDA_PREFIX}` so the shared libraries are on the default +dynamic linker search path at runtime. + +## 2. Build and install the Python package + +```bash +CUCASCADE_SRC=/path/to/cuCascade +CUCASCADE_BUILD=${CUCASCADE_SRC}/build +CUDF_CMAKE_DIR=/path/to/cudf/build + +pip install --no-build-isolation --no-deps \ + --config-settings "cmake.args=-DCMAKE_PREFIX_PATH=${CONDA_PREFIX};${CUDF_CMAKE_DIR}" \ + --config-settings "cmake.args=-Dconcurrentqueue_dir=${CUCASCADE_BUILD}/_deps/concurrentqueue-src" \ + "${CUCASCADE_SRC}/python" +``` + +**`CMAKE_PREFIX_PATH`** must include: +- `${CONDA_PREFIX}` — finds cuCascade, liburing, libcurl, OpenSSL +- the cudf cmake build directory — lets `find_package(cudf)` resolve + +**`concurrentqueue_dir`** — the moodycamel concurrentqueue headers are +fetched by the cuCascade C++ build and live at +`${CUCASCADE_BUILD}/_deps/concurrentqueue-src`. + +**`rapids-cmake-dir`** (optional) — if rapids-cmake is already available +(e.g. inside a pylibcudf build directory), pass it to avoid a network fetch: + +```bash + --config-settings "cmake.args=-Drapids-cmake-dir=/path/to/rapids-cmake-src" \ +``` + +Otherwise the build fetches rapids-cmake automatically from GitHub. + +## Verify + +```python +import cucascade +engine = cucascade.UringEngine() +ds = engine.open("/path/to/file.parquet") +``` diff --git a/python/cucascade/CMakeLists.txt b/python/cucascade/CMakeLists.txt new file mode 100644 index 0000000..dbb903b --- /dev/null +++ b/python/cucascade/CMakeLists.txt @@ -0,0 +1,31 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +set(cython_sources datasource.pyx) +set(linked_libraries cuCascade::cucascade_cudf cudf::cudf) + +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" + MODULE_PREFIX cucascade_ + ASSOCIATED_TARGETS cucascade_cudf_shared +) + +get_target_property(_cudf_loc cudf::cudf LOCATION) +get_filename_component(_cudf_lib_dir "${_cudf_loc}" DIRECTORY) +set_target_properties(cucascade_datasource + PROPERTIES INSTALL_RPATH "${_cudf_lib_dir}" +) + +set(concurrentqueue_dir + "" + CACHE PATH "Path to the moodycamel concurrentqueue source directory" +) +if(concurrentqueue_dir) + target_include_directories(cucascade_datasource PRIVATE "${concurrentqueue_dir}") +endif() diff --git a/python/cucascade/__init__.py b/python/cucascade/__init__.py new file mode 100644 index 0000000..0097545 --- /dev/null +++ b/python/cucascade/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from cucascade.datasource import CuCascadeDatasource, RestEngine, UringEngine + +__all__ = ["CuCascadeDatasource", "RestEngine", "UringEngine"] diff --git a/python/cucascade/datasource.pxd b/python/cucascade/datasource.pxd new file mode 100644 index 0000000..bec7ffe --- /dev/null +++ b/python/cucascade/datasource.pxd @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport int64_t +from libcpp cimport bool as cpp_bool +from libcpp.memory cimport unique_ptr +from libcpp.string cimport string +from libcpp.vector cimport vector + +from pylibcudf.libcudf.io.datasource cimport datasource as cudf_datasource + + +cdef extern from "cudf/io/text/byte_range_info.hpp" namespace "cudf::io::text" nogil: + cdef cppclass byte_range_info: + byte_range_info(int64_t offset, int64_t size) except + + int64_t offset() const + int64_t size() const + + +cdef extern from "cucascade/cudf/datasource.hpp" namespace "cucascade::io" nogil: + cdef cppclass cc_datasource "cucascade::io::datasource"(cudf_datasource): + unique_ptr[cc_datasource] duplicate() except + + void fadvise(const vector[byte_range_info]& ranges, int dev_id) except + + + +cdef extern from "cucascade/cudf/uring_datasource_engine.hpp" namespace "cucascade::io" nogil: + cdef cppclass uring_datasource_engine: + uring_datasource_engine(size_t n_reactors, + size_t pool_capacity, + size_t block_size, + cpp_bool use_odirect, + int numa_node) except + + unique_ptr[cc_datasource] open(string path) except + + + +cdef extern from "cucascade/cudf/rest_datasource_engine.hpp" namespace "cucascade::io" nogil: + cdef cppclass rest_datasource_engine: + rest_datasource_engine(string access_key_id, + string secret_access_key, + string session_token, + string region, + string endpoint, + size_t n_reactors, + cpp_bool tls_verify, + size_t pool_capacity, + size_t block_size) except + + unique_ptr[cc_datasource] open(string path) except + diff --git a/python/cucascade/datasource.pyi b/python/cucascade/datasource.pyi new file mode 100644 index 0000000..895d407 --- /dev/null +++ b/python/cucascade/datasource.pyi @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pylibcudf.io.datasource import Datasource + + +class CuCascadeDatasource(Datasource): + def fadvise(self, ranges: list[tuple[int, int]], dev_id: int = -1) -> None: ... + def duplicate(self) -> CuCascadeDatasource: ... + + +class UringEngine: + def __init__( + self, + n_reactors: int = 2, + pool_capacity: int = 2684354560, + block_size: int = 1048576, + use_odirect: bool = True, + numa_node: int = 0, + ) -> None: ... + def open(self, path: str) -> CuCascadeDatasource: ... + + +class RestEngine: + def __init__( + self, + access_key_id: str = "", + secret_access_key: str = "", + session_token: str = "", + region: str = "us-east-1", + endpoint: str = "", + n_reactors: int = 4, + tls_verify: bool = True, + pool_capacity: int = 2684354560, + block_size: int = 1048576, + ) -> None: ... + def open(self, path: str) -> CuCascadeDatasource: ... diff --git a/python/cucascade/datasource.pyx b/python/cucascade/datasource.pyx new file mode 100644 index 0000000..964f9d2 --- /dev/null +++ b/python/cucascade/datasource.pyx @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""cuCascade datasource bindings for use with pylibcudf readers. + +Provides :class:`UringEngine` (local NVMe via io_uring) and +:class:`RestEngine` (S3/HTTP via libcurl), both producing +:class:`CuCascadeDatasource` instances that implement +``cudf::io::datasource`` and support advisory prefetch via +:meth:`CuCascadeDatasource.fadvise`. +""" + +from cython.operator cimport dereference as deref +from libc.stdint cimport int64_t +from libcpp.memory cimport unique_ptr +from libcpp.string cimport string +from libcpp.utility cimport move +from libcpp.vector cimport vector + +from pylibcudf.io.datasource cimport Datasource +from pylibcudf.libcudf.io.datasource cimport datasource as cudf_datasource + +from cucascade.datasource cimport ( + byte_range_info, + cc_datasource, + rest_datasource_engine, + uring_datasource_engine, +) + +__all__ = ["CuCascadeDatasource", "RestEngine", "UringEngine"] + + +cdef class CuCascadeDatasource(Datasource): + """A pylibcudf :class:`~pylibcudf.io.datasource.Datasource` backed by + a ``cucascade::io::datasource``. + + Instances are produced by :meth:`UringEngine.open` or + :meth:`RestEngine.open` and are not constructed directly. + + The owning engine must outlive every datasource it produces. + When a single file is read by multiple concurrent splits, call + :meth:`duplicate` to obtain an independent datasource per split so + that their :meth:`fadvise` calls do not interfere with each other. + """ + + cdef unique_ptr[cc_datasource] _ds + + cdef cudf_datasource* get_datasource(self) except * nogil: + return self._ds.get() + + def fadvise(self, list ranges, int dev_id=-1): + """Hint the IO layer about byte ranges this scan will read soon. + + Queues the ranges into the engine's prefetch cache so that the + data is staged into pinned host memory before the caller blocks + on :meth:`~pylibcudf.io.types.SourceInfo`-based reads. + + Parameters + ---------- + ranges : list[tuple[int, int]] + Byte ranges as ``(offset, size)`` pairs, for example from + :meth:`~pylibcudf.io.experimental.HybridScanReader.filter_column_chunks_byte_ranges` + or + :meth:`~pylibcudf.io.experimental.HybridScanReader.payload_column_chunks_byte_ranges`. + dev_id : int, optional + Preferred CUDA device id for pinned-host staging placement. + Pass ``-1`` (the default) to express no preference. + """ + cdef vector[byte_range_info] c_ranges + cdef int64_t off, sz + for off, sz in ranges: + c_ranges.emplace_back(off, sz) + with nogil: + deref(self._ds).fadvise(c_ranges, dev_id) + + def duplicate(self): + """Return a datasource sharing the same file handle with an independent + prefetch handle. + + Use one duplicate per split when a single file is read by several + concurrent :class:`~pylibcudf.io.experimental.HybridScanReader` + instances so that per-split :meth:`fadvise` calls do not overwrite + each other's prefetch state. + + Returns + ------- + CuCascadeDatasource + A new datasource over the same underlying file. + """ + cdef CuCascadeDatasource out = CuCascadeDatasource.__new__(CuCascadeDatasource) + with nogil: + out._ds = move(deref(self._ds).duplicate()) + return out + + +cdef class UringEngine: + """io_uring-backed datasource engine for local NVMe reads. + + Owns a NUMA-local pinned host staging pool and a pool of io_uring + reactor threads. All datasources produced by :meth:`open` share the + engine's resources; the engine must therefore outlive every datasource + it produces. + + Parameters + ---------- + n_reactors : int, optional + Number of io_uring reactor threads. Default is 2. + pool_capacity : int, optional + Total capacity of the pinned host staging pool in bytes. + Default is approximately 2.5 GiB (20 × 128 MiB). + block_size : int, optional + Fixed block size in bytes for the staging pool. Must be a power + of two and at least the O_DIRECT alignment requirement of the + target filesystem. Default is 1 MiB. + use_odirect : bool, optional + Whether to open files with ``O_DIRECT`` to bypass the page cache. + Default is ``True``. + numa_node : int, optional + NUMA node from which to allocate the pinned staging pool. + Default is 0. + """ + + cdef unique_ptr[uring_datasource_engine] _engine + + def __cinit__( + self, + size_t n_reactors=2, + size_t pool_capacity=2684354560, + size_t block_size=1048576, + bint use_odirect=True, + int numa_node=0, + ): + with nogil: + self._engine.reset( + new uring_datasource_engine( + n_reactors, pool_capacity, block_size, use_odirect, numa_node + ) + ) + + def open(self, str path): + """Open a datasource for a local file. + + Parameters + ---------- + path : str + Path to the local file. + + Returns + ------- + CuCascadeDatasource + A datasource bound to this engine. Must not outlive the engine. + + Raises + ------ + RuntimeError + If the file cannot be opened. + """ + cdef CuCascadeDatasource ds = CuCascadeDatasource.__new__(CuCascadeDatasource) + cdef string c_path = path.encode() + with nogil: + ds._ds = move(deref(self._engine).open(c_path)) + return ds + + +cdef class RestEngine: + """libcurl-backed datasource engine for S3/HTTP object-store reads. + + Owns a NUMA-local pinned host staging pool and a pool of libcurl + reactor threads with SigV4 presigned-URL signing. All datasources + produced by :meth:`open` share the engine's resources; the engine + must therefore outlive every datasource it produces. + + Parameters + ---------- + access_key_id : str, optional + AWS access key ID. Default is ``""`` (reads from environment). + secret_access_key : str, optional + AWS secret access key. Default is ``""`` (reads from environment). + session_token : str, optional + STS session token; leave empty for long-lived credentials. + region : str, optional + AWS region. Default is ``"us-east-1"``. + endpoint : str, optional + S3-compatible endpoint host (e.g. ``"s3.amazonaws.com"`` or a + MinIO ``host:port``). Leave empty to derive from region. + n_reactors : int, optional + Number of libcurl reactor threads. Default is 4. + tls_verify : bool, optional + Whether to verify TLS peer certificates. Default is ``True``. + pool_capacity : int, optional + Total capacity of the pinned host staging pool in bytes. + Default is approximately 2.5 GiB. + block_size : int, optional + Fixed block size in bytes for the staging pool. Default is 1 MiB. + """ + + cdef unique_ptr[rest_datasource_engine] _engine + + def __cinit__( + self, + str access_key_id="", + str secret_access_key="", + str session_token="", + str region="us-east-1", + str endpoint="", + size_t n_reactors=4, + bint tls_verify=True, + size_t pool_capacity=2684354560, + size_t block_size=1048576, + ): + cdef string c_access_key_id = access_key_id.encode() + cdef string c_secret_access_key = secret_access_key.encode() + cdef string c_session_token = session_token.encode() + cdef string c_region = region.encode() + cdef string c_endpoint = endpoint.encode() + with nogil: + self._engine.reset( + new rest_datasource_engine( + c_access_key_id, + c_secret_access_key, + c_session_token, + c_region, + c_endpoint, + n_reactors, + tls_verify, + pool_capacity, + block_size, + ) + ) + + def open(self, str path): + """Open a datasource for an S3 URI. + + Issues an HTTP HEAD request to resolve the object size. + + Parameters + ---------- + path : str + S3 URI of the form ``s3://bucket/key``. + + Returns + ------- + CuCascadeDatasource + A datasource bound to this engine. Must not outlive the engine. + + Raises + ------ + RuntimeError + If the HEAD request fails or the URI is malformed. + """ + cdef CuCascadeDatasource ds = CuCascadeDatasource.__new__(CuCascadeDatasource) + cdef string c_path = path.encode() + with nogil: + ds._ds = move(deref(self._engine).open(c_path)) + return ds diff --git a/python/pdsh_results.jsonl b/python/pdsh_results.jsonl new file mode 100644 index 0000000..50744f2 --- /dev/null +++ b/python/pdsh_results.jsonl @@ -0,0 +1 @@ +{"engine_name": "cudf-polars", "queries": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], "query_set": "pdsh", "dataset_path": "/home/coder/data/tpch-rs-float/scale-1000-float-v2", "scale_factor": 1000, "suffix": "", "qualification": false, "frontend": "spmd", "iterations": 2, "io_mode": "lukewarm", "collect_traces": false, "native_parquet": false, "max_io_threads": 4, "n_workers": 1, "extra_info": {"environment": {"CUDF_POLARS_LOG_TRACES_MEMORY": null, "CUDF_POLARS_LOG_TRACES": null, "DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT": null, "DASK_DISTRIBUTED__COMM__UCX__CONNECT_TIMEOUT": null, "KVIKIO_NTHREADS": "8", "LIBCUDF_NUM_HOST_WORKERS": null, "OMP_NUM_THREADS": "1", "POLARS_MAX_THREADS": "1", "RAPIDSMPF_NUM_STREAMING_THREADS": "8", "UCX_MAX_RNDV_RAILS": "1", "UCX_PROTO_ENABLE": null, "UCX_RNDV_FRAG_MEM_TYPES": null, "UCX_RNDV_MTYPE_WORKER_FC_ENABLE": null, "UCX_RNDV_MTYPE_WORKER_MAX_MEM": null, "UCX_RNDV_PIPELINE_ERROR_HANDLING": null}}, "run_id": "b6bb83ff-4e5d-48d0-bdd8-506e6e7a33a7", "timestamp": "2026-07-27T23:01:23.951881+00:00", "command_line": "/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py --path /home/coder/data/tpch-rs-float/scale-1000-float-v2 --suffix '' --no-explain --no-print-results --iterations 2 --frontend spmd --pinned-memory --pinned-initial-pool-size 4000000000 all", "streaming_options": {"rapidsmpf": {"pinned_memory": "True", "pinned_initial_pool_size": "4000000000"}, "executor": {}, "engine": {}}, "records": {"1": [{"query": 1, "iteration": 0, "duration": 9.603747435845435, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 1, "iteration": 1, "duration": 8.959347929805517, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "2": [{"query": 2, "iteration": 0, "duration": 1.0147004406899214, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 2, "iteration": 1, "duration": 0.9327058754861355, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "3": [{"query": 3, "iteration": 0, "duration": 13.419135103933513, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 3, "iteration": 1, "duration": 20.19711658731103, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "4": [{"query": 4, "iteration": 0, "duration": 5.699714717455208, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 4, "iteration": 1, "duration": 9.655277168378234, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "5": [{"query": 5, "iteration": 0, "duration": 13.382477785460651, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 5, "iteration": 1, "duration": 15.713967258110642, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "6": [{"query": 6, "iteration": 0, "duration": 5.351213915273547, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 6, "iteration": 1, "duration": 7.692393693141639, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "7": [{"query": 7, "iteration": 0, "duration": 37.38136589154601, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 7, "iteration": 1, "duration": 38.60627434030175, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "8": [{"query": 8, "iteration": 0, "duration": 19.7513942765072, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 8, "iteration": 1, "duration": 20.178146742284298, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "9": [{"query": 9, "iteration": 0, "duration": 61.46330539043993, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 9, "iteration": 1, "duration": 64.16345161385834, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "10": [{"query": 10, "iteration": 0, "duration": 38.552414486184716, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 10, "iteration": 1, "duration": 45.9994532270357, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "11": [{"query": 11, "iteration": 0, "duration": 1.566706725396216, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 11, "iteration": 1, "duration": 1.5286435969173908, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "12": [{"query": 12, "iteration": 0, "duration": 10.606999724172056, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 12, "iteration": 1, "duration": 9.88353905826807, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "13": [{"query": 13, "iteration": 0, "duration": 7.336470068432391, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 13, "iteration": 1, "duration": 6.694976732134819, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "14": [{"query": 14, "iteration": 0, "duration": 15.797284464351833, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 14, "iteration": 1, "duration": 16.315921252593398, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "15": [{"query": 15, "iteration": 0, "duration": 7.4547947738319635, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 15, "iteration": 1, "duration": 10.861559779383242, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "16": [{"query": 16, "iteration": 0, "duration": 8.040832746773958, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 16, "iteration": 1, "duration": 4.459938563406467, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "17": [{"query": 17, "iteration": 0, "duration": 12.303235572762787, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 17, "iteration": 1, "duration": 10.637650010176003, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "18": [{"query": 18, "iteration": 0, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 745, in groupby_actor\n | await _shuffle_reduce(\n | ...<14 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 412, in _shuffle_reduce\n | aggregated, input_drained, _ = await _local_aggregation(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<5 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 218, in _local_aggregation\n | chunk = await evaluate_chunk(\n | ^^^^^^^^^^^^^^^^^^^^^\n | ...<4 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 792, in evaluate_chunk\n | chunk = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | _evaluate_chunk_sync, chunk, single_ir, ir_context, context.br()\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 748, in _evaluate_chunk_sync\n | df = ir.do_evaluate(\n | *ir._non_child_args,\n | DataFrame.from_table(chunk.table_view(), names, dtypes, chunk.stream),\n | context=ir_context,\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 2225, in do_evaluate\n | group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"pylibcudf/groupby.pyx\", line 165, in pylibcudf.groupby.GroupBy.aggregate\n | cpdef tuple aggregate(\n | ^^^^^^^\n | File \"pylibcudf/groupby.pyx\", line 201, in pylibcudf.groupby.GroupBy.aggregate\n | c_res = dereference(self.c_obj).aggregate(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 400000424 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 400000424 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}, {"query": 18, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py\", line 486, in fanout_node_unbounded\n | mid = sm.insert(msg.copy(memory_reservation))\n | ~~~~~~~~^^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/message.pyx\", line 146, in rapidsmpf.streaming.core.message.Message.copy\n | ret = self._handle.copy(deref(res))\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 799948384 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 799948384 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}], "19": [{"query": 19, "iteration": 0, "duration": 33.98043346218765, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 19, "iteration": 1, "duration": 29.311996046453714, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "20": [{"query": 20, "iteration": 0, "duration": 16.99104288686067, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 20, "iteration": 1, "duration": 16.40389629174024, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "21": [{"query": 21, "iteration": 0, "duration": 29.779883013106883, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 21, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 1462, in join_actor\n | await gather_in_task_group(*actor_tasks)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 984, in _shuffle_join\n | await gather_in_task_group(*actor_tasks)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 507, in _join_chunks\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<5 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 2908, in do_evaluate\n | plc.copying.gather(left.table, lg, left_policy, stream=stream),\n | ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"pylibcudf/copying.pyx\", line 64, in pylibcudf.copying.gather\n | cpdef Table gather(\n | ^^^^^^^^^^^\n | File \"pylibcudf/copying.pyx\", line 105, in pylibcudf.copying.gather\n | c_result = cpp_copying.gather(\n | ^^^\n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 731460136 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 731460136 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}], "22": [{"query": 22, "iteration": 0, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 724, in scan_node\n | await gather_in_task_group(\n | ...<5 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 709, in _producer\n | await read_chunk(\n | ...<8 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 584, in read_chunk\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<3 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/io.py\", line 682, in do_evaluate\n | return Scan.do_evaluate(\n | ~~~~~~~~~~~~~~~~^\n | schema,\n | ^^^^^^^\n | ...<11 lines>...\n | context=context,\n | ^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 1164, in do_evaluate\n | tbl_w_meta = plc.io.parquet.read_parquet(\n | parquet_reader_options,\n | parquet_metadatas=parquet_metadatas,\n | stream=stream,\n | )\n | File \"pylibcudf/io/parquet.pyx\", line 684, in pylibcudf.io.parquet.read_parquet\n | cpdef read_parquet(\n | ^^^^^^^^^^^\n | File \"pylibcudf/io/parquet.pyx\", line 725, in pylibcudf.io.parquet.read_parquet\n | cpp_read_parquet(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 569804625 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 569804625 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}, {"query": 22, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 724, in scan_node\n | await gather_in_task_group(\n | ...<5 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 709, in _producer\n | await read_chunk(\n | ...<8 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 584, in read_chunk\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<3 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/io.py\", line 682, in do_evaluate\n | return Scan.do_evaluate(\n | ~~~~~~~~~~~~~~~~^\n | schema,\n | ^^^^^^^\n | ...<11 lines>...\n | context=context,\n | ^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 1164, in do_evaluate\n | tbl_w_meta = plc.io.parquet.read_parquet(\n | parquet_reader_options,\n | parquet_metadatas=parquet_metadatas,\n | stream=stream,\n | )\n | File \"pylibcudf/io/parquet.pyx\", line 684, in pylibcudf.io.parquet.read_parquet\n | cpdef read_parquet(\n | ^^^^^^^^^^^\n | File \"pylibcudf/io/parquet.pyx\", line 725, in pylibcudf.io.parquet.read_parquet\n | cpp_read_parquet(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 875963952 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 875963952 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}]}, "plans": {}, "versions": {"cudf_polars": {"version": "26.10.00", "commit": ""}, "polars": "1.42.1", "python": "3.13.14", "rapidsmpf": {"version": "26.10.00", "commit": ""}, "duckdb": "1.5.4"}, "hardware": {"gpus": [{"name": "NVIDIA B200", "index": 0, "free_memory": 190840766464, "used_memory": 1425080320, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 1, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 2, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 3, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 4, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 5, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 6, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 7, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}], "cpu": {"model": "INTEL(R) XEON(R) PLATINUM 8570", "physical_cores": 112, "logical_cores": 224}}, "validation_method": null, "roles": [], "config_options": {"config_options": {"raise_on_fail": false, "parquet_options": {"chunked": true, "n_output_chunks": 1, "chunk_read_limit": 0, "pass_read_limit": 0, "max_footer_samples": 3, "max_row_group_samples": 1, "use_rapidsmpf_native": false, "use_hybrid_scan": true, "prefetch_file_metadata": true, "hybrid_scan_stats_pruning": true, "prefetch_backend": "cucascade", "use_jit_filter": false}, "executor": {"name": "streaming", "cluster": "default_singleton", "fallback_mode": "warn", "max_rows_per_partition": 1000000, "target_partition_size": 1500000000, "broadcast_limit": 16000000000, "client_device_threshold": 0.5, "sink_to_directory": false, "dynamic_planning": {"sample_chunk_count": 2, "join_prefilter_threshold": 0.5, "join_prefilter_max_key_columns": 1, "join_prefilter_trace": false}, "join_filter_pushdown": null, "max_io_threads": 4, "num_prefetch_workers": 2, "spill_to_pinned_memory": false, "num_py_executors": 8, "min_device_size": 192265846784, "spmd_context": null, "ray_context": null, "dask_context": null}, "device": null, "memory_resource_config": null}, "rapidsmpf_options": {"num_streaming_threads": "4", "pinned_max_pool_size": "", "statistics": "", "memory_reserve_timeout": "", "pinned_initial_pool_size": "4000000000", "spill_device_limit": "", "pinned_memory": "True", "periodic_spill_check": "", "num_streams": "", "log": "", "allow_overbooking_by_default": ""}}} diff --git a/python/pyproject.toml b/python/pyproject.toml new file mode 100644 index 0000000..72d8f45 --- /dev/null +++ b/python/pyproject.toml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +build-backend = "scikit_build_core.build" +requires = [ + "scikit-build-core[pyproject]>=0.11.0", + "cmake>=4.0", + "cython>=3.0.0", + "ninja", +] + +[project] +name = "cucascade" +version = "0.1.0" +description = "Python bindings for cuCascade" +requires-python = ">=3.11" +dependencies = [ + "pylibcudf", +] + +[tool.scikit-build] +build-dir = "build/{wheel_tag}" +cmake.build-type = "Release" +ninja.make-fallback = false +wheel.packages = ["cucascade"] + diff --git a/src/cudf/CMakeLists.txt b/src/cudf/CMakeLists.txt index 3302dd9..9918dac 100644 --- a/src/cudf/CMakeLists.txt +++ b/src/cudf/CMakeLists.txt @@ -25,6 +25,9 @@ target_sources( # The cudf::io::datasource bridge over the cudf-free io core — only built when # the io library is enabled (the io core itself carries no cudf dependency). if(CUCASCADE_BUILD_IO) - target_sources(cucascade_cudf_objects - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/datasource.cpp) + target_sources( + cucascade_cudf_objects + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/datasource.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/uring_datasource_engine.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/rest_datasource_engine.cpp) endif() diff --git a/src/cudf/rest_datasource_engine.cpp b/src/cudf/rest_datasource_engine.cpp new file mode 100644 index 0000000..45e313d --- /dev/null +++ b/src/cudf/rest_datasource_engine.cpp @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +namespace cucascade::io { + +rest_datasource_engine::rest_datasource_engine(std::string access_key_id, + std::string secret_access_key, + std::string session_token, + std::string region, + std::string endpoint, + std::size_t n_reactors, + bool tls_verify, + std::size_t pool_capacity, + std::size_t block_size) + : _upstream(0, true), + _host_mr(0, _upstream, pool_capacity, pool_capacity, block_size, 128, 1) +{ + rest::s3::static_credentials creds{.access_key_id = std::move(access_key_id), + .secret_access_key = std::move(secret_access_key), + .session_token = std::move(session_token), + .expires_at = std::nullopt}; + + auto authorizer = std::make_shared( + std::move(creds), std::move(region), std::move(endpoint)); + + rest::config rest_cfg{}; + rest_cfg.bounce_block_size = _host_mr.get_block_size(); + rest_cfg.tls_verify = tls_verify; + + auto rest_ctx = std::make_shared( + std::move(rest_cfg), std::move(authorizer), &_host_mr); + + auto io_ctx = std::make_shared(n_reactors, std::move(rest_ctx)); + io_ctx->start(); + _io_ctx = std::move(io_ctx); +} + +rest_datasource_engine::~rest_datasource_engine() { _io_ctx->shutdown(); } + +std::unique_ptr rest_datasource_engine::open(std::string path) const +{ + return open_datasource(_io_ctx, std::move(path)); +} + +} // namespace cucascade::io diff --git a/src/cudf/uring_datasource_engine.cpp b/src/cudf/uring_datasource_engine.cpp new file mode 100644 index 0000000..ff63324 --- /dev/null +++ b/src/cudf/uring_datasource_engine.cpp @@ -0,0 +1,35 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +namespace cucascade::io { + +uring_datasource_engine::uring_datasource_engine(std::size_t n_reactors, + std::size_t pool_capacity, + std::size_t block_size, + bool use_odirect, + int numa_node) + : _upstream(numa_node, true), + _host_mr(0, _upstream, pool_capacity, pool_capacity, block_size, 128, 1), + _reactor_ctx(std::make_shared( + uring::uring_reactor::reactor_config_type{.bounce_size = _host_mr.get_block_size(), + .use_odirect = use_odirect}, + &_host_mr)), + _io_ctx(std::make_shared(n_reactors, _reactor_ctx)) +{ + _io_ctx->start(); +} + +uring_datasource_engine::~uring_datasource_engine() { _io_ctx->shutdown(); } + +std::unique_ptr uring_datasource_engine::open(std::string path) const +{ + return open_datasource(_io_ctx, std::move(path)); +} + +} // namespace cucascade::io From fd8f5af5843c163d9525c5e478de1012dcb48c42 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 29 Jul 2026 17:13:36 +0000 Subject: [PATCH 2/5] Add read_ranges_async to Python bindings --- python/cucascade/__init__.py | 4 +-- python/cucascade/datasource.pxd | 4 ++- python/cucascade/datasource.pyi | 7 ++++ python/cucascade/datasource.pyx | 61 +++++++++++++++++++++++++++++++-- python/pdsh_results.jsonl | 1 - 5 files changed, 71 insertions(+), 6 deletions(-) delete mode 100644 python/pdsh_results.jsonl diff --git a/python/cucascade/__init__.py b/python/cucascade/__init__.py index 0097545..dba196c 100644 --- a/python/cucascade/__init__.py +++ b/python/cucascade/__init__.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from cucascade.datasource import CuCascadeDatasource, RestEngine, UringEngine +from cucascade.datasource import CuCascadeDatasource, ReadFuture, RestEngine, UringEngine -__all__ = ["CuCascadeDatasource", "RestEngine", "UringEngine"] +__all__ = ["CuCascadeDatasource", "ReadFuture", "RestEngine", "UringEngine"] diff --git a/python/cucascade/datasource.pxd b/python/cucascade/datasource.pxd index bec7ffe..b29f3f7 100644 --- a/python/cucascade/datasource.pxd +++ b/python/cucascade/datasource.pxd @@ -1,8 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from libc.stdint cimport int64_t +from libc.stdint cimport int64_t, uint8_t from libcpp cimport bool as cpp_bool +from libcpp.future cimport future from libcpp.memory cimport unique_ptr from libcpp.string cimport string from libcpp.vector cimport vector @@ -21,6 +22,7 @@ cdef extern from "cucascade/cudf/datasource.hpp" namespace "cucascade::io" nogil cdef cppclass cc_datasource "cucascade::io::datasource"(cudf_datasource): unique_ptr[cc_datasource] duplicate() except + void fadvise(const vector[byte_range_info]& ranges, int dev_id) except + + future[size_t] host_read_async(size_t offset, size_t size, uint8_t* dst) except + cdef extern from "cucascade/cudf/uring_datasource_engine.hpp" namespace "cucascade::io" nogil: diff --git a/python/cucascade/datasource.pyi b/python/cucascade/datasource.pyi index 895d407..e22e69a 100644 --- a/python/cucascade/datasource.pyi +++ b/python/cucascade/datasource.pyi @@ -6,9 +6,16 @@ from __future__ import annotations from pylibcudf.io.datasource import Datasource +class ReadFuture: + def get(self) -> int: ... + + class CuCascadeDatasource(Datasource): def fadvise(self, ranges: list[tuple[int, int]], dev_id: int = -1) -> None: ... def duplicate(self) -> CuCascadeDatasource: ... + def read_ranges_async( + self, ranges: list[tuple[int, int]], buffer: memoryview + ) -> list[ReadFuture]: ... class UringEngine: diff --git a/python/cucascade/datasource.pyx b/python/cucascade/datasource.pyx index 964f9d2..04d4244 100644 --- a/python/cucascade/datasource.pyx +++ b/python/cucascade/datasource.pyx @@ -10,7 +10,8 @@ Provides :class:`UringEngine` (local NVMe via io_uring) and """ from cython.operator cimport dereference as deref -from libc.stdint cimport int64_t +from libc.stdint cimport int64_t, uint8_t +from libcpp.future cimport future from libcpp.memory cimport unique_ptr from libcpp.string cimport string from libcpp.utility cimport move @@ -26,7 +27,29 @@ from cucascade.datasource cimport ( uring_datasource_engine, ) -__all__ = ["CuCascadeDatasource", "RestEngine", "UringEngine"] +__all__ = ["CuCascadeDatasource", "ReadFuture", "RestEngine", "UringEngine"] + + +cdef class ReadFuture: + """A pending async host read issued by :meth:`CuCascadeDatasource.read_ranges_async`. + + Call :meth:`get` to block until the read completes. + """ + + cdef future[size_t] _fut + + def get(self): + """Block until the read completes. + + Returns + ------- + int + Number of bytes transferred. + """ + cdef size_t result + with nogil: + result = self._fut.get() + return result cdef class CuCascadeDatasource(Datasource): @@ -91,6 +114,40 @@ cdef class CuCascadeDatasource(Datasource): out._ds = move(deref(self._ds).duplicate()) return out + def read_ranges_async(self, list ranges, object buffer): + """Submit async host reads for each range into a contiguous buffer. + + The reads are submitted immediately and run on the engine's reactor + threads. Call :meth:`ReadFuture.get` on each returned future to wait + for the corresponding range to complete. + + Parameters + ---------- + ranges : list[tuple[int, int]] + Byte ranges as ``(offset, size)`` pairs in the same order that data + should appear in ``buffer``. + buffer : memoryview + Contiguous writable host buffer sized to hold the sum of all range + sizes. Typically a slice of a :class:`PinnedBuffer` array. + + Returns + ------- + list[ReadFuture] + One future per range, in the same order as ``ranges``. + """ + cdef uint8_t[::1] c_buf = buffer + cdef uint8_t* base = &c_buf[0] + cdef list futures = [] + cdef ReadFuture rf + cdef int64_t off, sz + cdef size_t dst_offset = 0 + for off, sz in ranges: + rf = ReadFuture.__new__(ReadFuture) + rf._fut = deref(self._ds).host_read_async(off, sz, base + dst_offset) + futures.append(rf) + dst_offset += sz + return futures + cdef class UringEngine: """io_uring-backed datasource engine for local NVMe reads. diff --git a/python/pdsh_results.jsonl b/python/pdsh_results.jsonl deleted file mode 100644 index 50744f2..0000000 --- a/python/pdsh_results.jsonl +++ /dev/null @@ -1 +0,0 @@ -{"engine_name": "cudf-polars", "queries": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], "query_set": "pdsh", "dataset_path": "/home/coder/data/tpch-rs-float/scale-1000-float-v2", "scale_factor": 1000, "suffix": "", "qualification": false, "frontend": "spmd", "iterations": 2, "io_mode": "lukewarm", "collect_traces": false, "native_parquet": false, "max_io_threads": 4, "n_workers": 1, "extra_info": {"environment": {"CUDF_POLARS_LOG_TRACES_MEMORY": null, "CUDF_POLARS_LOG_TRACES": null, "DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT": null, "DASK_DISTRIBUTED__COMM__UCX__CONNECT_TIMEOUT": null, "KVIKIO_NTHREADS": "8", "LIBCUDF_NUM_HOST_WORKERS": null, "OMP_NUM_THREADS": "1", "POLARS_MAX_THREADS": "1", "RAPIDSMPF_NUM_STREAMING_THREADS": "8", "UCX_MAX_RNDV_RAILS": "1", "UCX_PROTO_ENABLE": null, "UCX_RNDV_FRAG_MEM_TYPES": null, "UCX_RNDV_MTYPE_WORKER_FC_ENABLE": null, "UCX_RNDV_MTYPE_WORKER_MAX_MEM": null, "UCX_RNDV_PIPELINE_ERROR_HANDLING": null}}, "run_id": "b6bb83ff-4e5d-48d0-bdd8-506e6e7a33a7", "timestamp": "2026-07-27T23:01:23.951881+00:00", "command_line": "/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py --path /home/coder/data/tpch-rs-float/scale-1000-float-v2 --suffix '' --no-explain --no-print-results --iterations 2 --frontend spmd --pinned-memory --pinned-initial-pool-size 4000000000 all", "streaming_options": {"rapidsmpf": {"pinned_memory": "True", "pinned_initial_pool_size": "4000000000"}, "executor": {}, "engine": {}}, "records": {"1": [{"query": 1, "iteration": 0, "duration": 9.603747435845435, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 1, "iteration": 1, "duration": 8.959347929805517, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "2": [{"query": 2, "iteration": 0, "duration": 1.0147004406899214, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 2, "iteration": 1, "duration": 0.9327058754861355, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "3": [{"query": 3, "iteration": 0, "duration": 13.419135103933513, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 3, "iteration": 1, "duration": 20.19711658731103, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "4": [{"query": 4, "iteration": 0, "duration": 5.699714717455208, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 4, "iteration": 1, "duration": 9.655277168378234, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "5": [{"query": 5, "iteration": 0, "duration": 13.382477785460651, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 5, "iteration": 1, "duration": 15.713967258110642, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "6": [{"query": 6, "iteration": 0, "duration": 5.351213915273547, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 6, "iteration": 1, "duration": 7.692393693141639, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "7": [{"query": 7, "iteration": 0, "duration": 37.38136589154601, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 7, "iteration": 1, "duration": 38.60627434030175, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "8": [{"query": 8, "iteration": 0, "duration": 19.7513942765072, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 8, "iteration": 1, "duration": 20.178146742284298, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "9": [{"query": 9, "iteration": 0, "duration": 61.46330539043993, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 9, "iteration": 1, "duration": 64.16345161385834, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "10": [{"query": 10, "iteration": 0, "duration": 38.552414486184716, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 10, "iteration": 1, "duration": 45.9994532270357, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "11": [{"query": 11, "iteration": 0, "duration": 1.566706725396216, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 11, "iteration": 1, "duration": 1.5286435969173908, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "12": [{"query": 12, "iteration": 0, "duration": 10.606999724172056, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 12, "iteration": 1, "duration": 9.88353905826807, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "13": [{"query": 13, "iteration": 0, "duration": 7.336470068432391, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 13, "iteration": 1, "duration": 6.694976732134819, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "14": [{"query": 14, "iteration": 0, "duration": 15.797284464351833, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 14, "iteration": 1, "duration": 16.315921252593398, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "15": [{"query": 15, "iteration": 0, "duration": 7.4547947738319635, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 15, "iteration": 1, "duration": 10.861559779383242, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "16": [{"query": 16, "iteration": 0, "duration": 8.040832746773958, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 16, "iteration": 1, "duration": 4.459938563406467, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "17": [{"query": 17, "iteration": 0, "duration": 12.303235572762787, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 17, "iteration": 1, "duration": 10.637650010176003, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "18": [{"query": 18, "iteration": 0, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 745, in groupby_actor\n | await _shuffle_reduce(\n | ...<14 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 412, in _shuffle_reduce\n | aggregated, input_drained, _ = await _local_aggregation(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<5 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py\", line 218, in _local_aggregation\n | chunk = await evaluate_chunk(\n | ^^^^^^^^^^^^^^^^^^^^^\n | ...<4 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 792, in evaluate_chunk\n | chunk = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | _evaluate_chunk_sync, chunk, single_ir, ir_context, context.br()\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 748, in _evaluate_chunk_sync\n | df = ir.do_evaluate(\n | *ir._non_child_args,\n | DataFrame.from_table(chunk.table_view(), names, dtypes, chunk.stream),\n | context=ir_context,\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 2225, in do_evaluate\n | group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"pylibcudf/groupby.pyx\", line 165, in pylibcudf.groupby.GroupBy.aggregate\n | cpdef tuple aggregate(\n | ^^^^^^^\n | File \"pylibcudf/groupby.pyx\", line 201, in pylibcudf.groupby.GroupBy.aggregate\n | c_res = dereference(self.c_obj).aggregate(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 400000424 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 400000424 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}, {"query": 18, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py\", line 486, in fanout_node_unbounded\n | mid = sm.insert(msg.copy(memory_reservation))\n | ~~~~~~~~^^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/message.pyx\", line 146, in rapidsmpf.streaming.core.message.Message.copy\n | ret = self._handle.copy(deref(res))\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 799948384 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 799948384 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}], "19": [{"query": 19, "iteration": 0, "duration": 33.98043346218765, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 19, "iteration": 1, "duration": 29.311996046453714, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "20": [{"query": 20, "iteration": 0, "duration": 16.99104288686067, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 20, "iteration": 1, "duration": 16.40389629174024, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}], "21": [{"query": 21, "iteration": 0, "duration": 29.779883013106883, "statistics": {}, "traces": null, "validation_result": null, "status": "success"}, {"query": 21, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 1462, in join_actor\n | await gather_in_task_group(*actor_tasks)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 984, in _shuffle_join\n | await gather_in_task_group(*actor_tasks)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py\", line 507, in _join_chunks\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<5 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 2908, in do_evaluate\n | plc.copying.gather(left.table, lg, left_policy, stream=stream),\n | ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"pylibcudf/copying.pyx\", line 64, in pylibcudf.copying.gather\n | cpdef Table gather(\n | ^^^^^^^^^^^\n | File \"pylibcudf/copying.pyx\", line 105, in pylibcudf.copying.gather\n | c_result = cpp_copying.gather(\n | ^^^\n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 731460136 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 731460136 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}], "22": [{"query": 22, "iteration": 0, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 724, in scan_node\n | await gather_in_task_group(\n | ...<5 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 709, in _producer\n | await read_chunk(\n | ...<8 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 584, in read_chunk\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<3 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/io.py\", line 682, in do_evaluate\n | return Scan.do_evaluate(\n | ~~~~~~~~~~~~~~~~^\n | schema,\n | ^^^^^^^\n | ...<11 lines>...\n | context=context,\n | ^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 1164, in do_evaluate\n | tbl_w_meta = plc.io.parquet.read_parquet(\n | parquet_reader_options,\n | parquet_metadatas=parquet_metadatas,\n | stream=stream,\n | )\n | File \"pylibcudf/io/parquet.pyx\", line 684, in pylibcudf.io.parquet.read_parquet\n | cpdef read_parquet(\n | ^^^^^^^^^^^\n | File \"pylibcudf/io/parquet.pyx\", line 725, in pylibcudf.io.parquet.read_parquet\n | cpp_read_parquet(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 569804625 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 569804625 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}, {"query": 22, "iteration": 1, "status": "error", "traceback": " + Exception Group Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 513, in execute_ir_on_rank\n | run_actor_network(ctx, actors=nodes)\n | ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 361, in rapidsmpf.streaming.core.actor.run_actor_network\n | with ThreadPoolExecutor(max_workers=1, thread_name_prefix=\"rapidsmpf-actor\") as executor:\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 369, in rapidsmpf.streaming.core.actor.run_actor_network\n | raise\n | \n | File \"rapidsmpf/streaming/core/actor.pyx\", line 366, in rapidsmpf.streaming.core.actor.run_actor_network\n | worker.result()\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 460, in result\n | return self.__get_result()\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/_base.py\", line 402, in __get_result\n | raise self._exception\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 268, in rapidsmpf.streaming.core.actor.sync_wait\n | raise\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 263, in rapidsmpf.streaming.core.actor.sync_wait\n | with asyncio.Runner() as runner:\n | ^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 264, in rapidsmpf.streaming.core.actor.sync_wait\n | return runner.run(run_and_publish_task(coro, task_ready))\n | \n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/runners.py\", line 119, in run\n | return self._loop.run_until_complete(task)\n | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/base_events.py\", line 725, in run_until_complete\n | return future.result()\n | ~~~~~~~~~~~~~^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 246, in run_and_publish_task\n | return await coro\n | ^^^^^^^\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 172, in run_py_actors\n | async with asyncio.TaskGroup() as tg:\n | ^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Exception Group Traceback (most recent call last):\n | File \"rapidsmpf/streaming/core/actor.pyx\", line 158, in py_actor\n | return await func(*args, **kwargs)\n | ^^^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 724, in scan_node\n | await gather_in_task_group(\n | ...<5 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py\", line 220, in gather_in_task_group\n | async with asyncio.TaskGroup() as tg:\n | ~~~~~~~~~~~~~~~~~^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 71, in __aexit__\n | return await self._aexit(et, exc)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/asyncio/taskgroups.py\", line 173, in _aexit\n | raise BaseExceptionGroup(\n | ...<2 lines>...\n | ) from None\n | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)\n +-+---------------- 1 ----------------\n | Traceback (most recent call last):\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 709, in _producer\n | await read_chunk(\n | ...<8 lines>...\n | )\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py\", line 584, in read_chunk\n | df = await ir_context.to_thread(\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | ...<3 lines>...\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 180, in to_thread\n | return await loop.run_in_executor(self.py_executor, func_call)\n | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n | File \"/home/coder/.conda/envs/rapids/lib/python3.13/concurrent/futures/thread.py\", line 59, in run\n | result = self.fn(*self.args, **self.kwargs)\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/io.py\", line 682, in do_evaluate\n | return Scan.do_evaluate(\n | ~~~~~~~~~~~~~~~~^\n | schema,\n | ^^^^^^^\n | ...<11 lines>...\n | context=context,\n | ^^^^^^^^^^^^^^^^\n | )\n | ^\n | File \"/home/coder/cudf/python/cudf_polars/cudf_polars/dsl/ir.py\", line 1164, in do_evaluate\n | tbl_w_meta = plc.io.parquet.read_parquet(\n | parquet_reader_options,\n | parquet_metadatas=parquet_metadatas,\n | stream=stream,\n | )\n | File \"pylibcudf/io/parquet.pyx\", line 684, in pylibcudf.io.parquet.read_parquet\n | cpdef read_parquet(\n | ^^^^^^^^^^^\n | File \"pylibcudf/io/parquet.pyx\", line 725, in pylibcudf.io.parquet.read_parquet\n | cpp_read_parquet(\n | \n | MemoryError: std::bad_alloc: out_of_memory: CUDA error (failed to allocate 875963952 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n +------------------------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 1088, in run_polars_query\n record = run_polars_query_iteration(\n q_id=q_id,\n ...<8 lines>...\n result_casts=casts if casts else None,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 952, in run_polars_query_iteration\n result, duration = execute_query(q_id, iteration, q, run_config, args, engine)\n ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py\", line 858, in execute_query\n result = q.collect(engine=engine)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/_utils/deprecation.py\", line 97, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py\", line 344, in wrapper\n return function(*args, **kwargs)\n File \"/home/coder/.conda/envs/rapids/lib/python3.13/site-packages/polars/lazyframe/frame.py\", line 2630, in collect\n return wrap_df(ldf.collect(engine, callback))\n ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/callback.py\", line 322, in _callback\n return evaluate_streaming(ir, config_options)\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/parallel.py\", line 262, in evaluate_streaming\n result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)\n ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py\", line 111, in evaluate_logical_plan\n _gpu_result, metadata_collector = evaluate_pipeline_spmd_mode(\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~^\n ir,\n ^^^\n ...<2 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/spmd.py\", line 143, in evaluate_pipeline_spmd_mode\n df, metadata = evaluate_on_rank(\n ~~~~~~~~~~~~~~~~^\n context,\n ^^^^^^^^\n ...<5 lines>...\n query_id=query_id,\n ^^^^^^^^^^^^^^^^^^\n )\n ^\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 790, in evaluate_on_rank\n return execute_ir_on_rank(\n ctx,\n ...<6 lines>...\n collective_id_map,\n )\n File \"/home/coder/cudf/python/cudf_polars/cudf_polars/engine/core.py\", line 524, in execute_ir_on_rank\n raise MemoryError(hint) from e\nMemoryError: Try lowering `target_partition_size` (current 1500000000) and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory.\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ for troubleshooting guidance.\nOriginal error:\nstd::bad_alloc: out_of_memory: CUDA error (failed to allocate 875963952 bytes) at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/mr/cuda_async_view_memory_resource.cpp:43: cudaErrorMemoryAllocation out of memory\n"}]}, "plans": {}, "versions": {"cudf_polars": {"version": "26.10.00", "commit": ""}, "polars": "1.42.1", "python": "3.13.14", "rapidsmpf": {"version": "26.10.00", "commit": ""}, "duckdb": "1.5.4"}, "hardware": {"gpus": [{"name": "NVIDIA B200", "index": 0, "free_memory": 190840766464, "used_memory": 1425080320, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 1, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 2, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 3, "free_memory": 191499337728, "used_memory": 766509056, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 4, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 5, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 6, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}, {"name": "NVIDIA B200", "index": 7, "free_memory": 190846992384, "used_memory": 1418854400, "total_memory": 192265846784}], "cpu": {"model": "INTEL(R) XEON(R) PLATINUM 8570", "physical_cores": 112, "logical_cores": 224}}, "validation_method": null, "roles": [], "config_options": {"config_options": {"raise_on_fail": false, "parquet_options": {"chunked": true, "n_output_chunks": 1, "chunk_read_limit": 0, "pass_read_limit": 0, "max_footer_samples": 3, "max_row_group_samples": 1, "use_rapidsmpf_native": false, "use_hybrid_scan": true, "prefetch_file_metadata": true, "hybrid_scan_stats_pruning": true, "prefetch_backend": "cucascade", "use_jit_filter": false}, "executor": {"name": "streaming", "cluster": "default_singleton", "fallback_mode": "warn", "max_rows_per_partition": 1000000, "target_partition_size": 1500000000, "broadcast_limit": 16000000000, "client_device_threshold": 0.5, "sink_to_directory": false, "dynamic_planning": {"sample_chunk_count": 2, "join_prefilter_threshold": 0.5, "join_prefilter_max_key_columns": 1, "join_prefilter_trace": false}, "join_filter_pushdown": null, "max_io_threads": 4, "num_prefetch_workers": 2, "spill_to_pinned_memory": false, "num_py_executors": 8, "min_device_size": 192265846784, "spmd_context": null, "ray_context": null, "dask_context": null}, "device": null, "memory_resource_config": null}, "rapidsmpf_options": {"num_streaming_threads": "4", "pinned_max_pool_size": "", "statistics": "", "memory_reserve_timeout": "", "pinned_initial_pool_size": "4000000000", "spill_device_limit": "", "pinned_memory": "True", "periodic_spill_check": "", "num_streams": "", "log": "", "allow_overbooking_by_default": ""}}} From 6822a6eb699866fcf0388ac4591eca2374b9ee87 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 30 Jul 2026 12:11:40 +0000 Subject: [PATCH 3/5] more knobs to turn --- include/cucascade/cudf/rest_datasource_engine.hpp | 11 +++++++---- python/cucascade/datasource.pxd | 5 ++++- python/cucascade/datasource.pyi | 3 +++ python/cucascade/datasource.pyx | 15 +++++++++++++++ src/cudf/rest_datasource_engine.cpp | 8 +++++++- 5 files changed, 36 insertions(+), 6 deletions(-) diff --git a/include/cucascade/cudf/rest_datasource_engine.hpp b/include/cucascade/cudf/rest_datasource_engine.hpp index 679b53e..a16331c 100644 --- a/include/cucascade/cudf/rest_datasource_engine.hpp +++ b/include/cucascade/cudf/rest_datasource_engine.hpp @@ -61,10 +61,13 @@ class rest_datasource_engine { std::string session_token, std::string region, std::string endpoint, - std::size_t n_reactors = 4, - bool tls_verify = true, - std::size_t pool_capacity = default_pool_capacity, - std::size_t block_size = default_block_size); + std::size_t n_reactors = 4, + bool tls_verify = true, + std::size_t pool_capacity = default_pool_capacity, + std::size_t block_size = default_block_size, + std::size_t max_connections = 16, + std::size_t chunk_size = 8UL << 20, + std::size_t max_n_chunks = 16); ~rest_datasource_engine(); diff --git a/python/cucascade/datasource.pxd b/python/cucascade/datasource.pxd index b29f3f7..c66dc78 100644 --- a/python/cucascade/datasource.pxd +++ b/python/cucascade/datasource.pxd @@ -45,5 +45,8 @@ cdef extern from "cucascade/cudf/rest_datasource_engine.hpp" namespace "cucascad size_t n_reactors, cpp_bool tls_verify, size_t pool_capacity, - size_t block_size) except + + size_t block_size, + size_t max_connections, + size_t chunk_size, + size_t max_n_chunks) except + unique_ptr[cc_datasource] open(string path) except + diff --git a/python/cucascade/datasource.pyi b/python/cucascade/datasource.pyi index e22e69a..3ca22a7 100644 --- a/python/cucascade/datasource.pyi +++ b/python/cucascade/datasource.pyi @@ -42,5 +42,8 @@ class RestEngine: tls_verify: bool = True, pool_capacity: int = 2684354560, block_size: int = 1048576, + max_connections: int = 16, + chunk_size: int = 8388608, + max_n_chunks: int = 16, ) -> None: ... def open(self, path: str) -> CuCascadeDatasource: ... diff --git a/python/cucascade/datasource.pyx b/python/cucascade/datasource.pyx index 04d4244..be1531b 100644 --- a/python/cucascade/datasource.pyx +++ b/python/cucascade/datasource.pyx @@ -248,6 +248,15 @@ cdef class RestEngine: Default is approximately 2.5 GiB. block_size : int, optional Fixed block size in bytes for the staging pool. Default is 1 MiB. + max_connections : int, optional + Maximum concurrent in-flight HTTP connections per reactor. + Default is 16. + chunk_size : int, optional + Maximum bytes per ranged GET request. Adjacent segments are fused + up to this size; oversized segments are split. Default is 8 MiB. + max_n_chunks : int, optional + Maximum number of destination buffers fused into a single scatter + GET. Default is 16. """ cdef unique_ptr[rest_datasource_engine] _engine @@ -263,6 +272,9 @@ cdef class RestEngine: bint tls_verify=True, size_t pool_capacity=2684354560, size_t block_size=1048576, + size_t max_connections=16, + size_t chunk_size=8388608, + size_t max_n_chunks=16, ): cdef string c_access_key_id = access_key_id.encode() cdef string c_secret_access_key = secret_access_key.encode() @@ -281,6 +293,9 @@ cdef class RestEngine: tls_verify, pool_capacity, block_size, + max_connections, + chunk_size, + max_n_chunks, ) ) diff --git a/src/cudf/rest_datasource_engine.cpp b/src/cudf/rest_datasource_engine.cpp index 45e313d..36328a8 100644 --- a/src/cudf/rest_datasource_engine.cpp +++ b/src/cudf/rest_datasource_engine.cpp @@ -19,7 +19,10 @@ rest_datasource_engine::rest_datasource_engine(std::string access_key_id, std::size_t n_reactors, bool tls_verify, std::size_t pool_capacity, - std::size_t block_size) + std::size_t block_size, + std::size_t max_connections, + std::size_t chunk_size, + std::size_t max_n_chunks) : _upstream(0, true), _host_mr(0, _upstream, pool_capacity, pool_capacity, block_size, 128, 1) { @@ -34,6 +37,9 @@ rest_datasource_engine::rest_datasource_engine(std::string access_key_id, rest::config rest_cfg{}; rest_cfg.bounce_block_size = _host_mr.get_block_size(); rest_cfg.tls_verify = tls_verify; + rest_cfg.max_connections = max_connections; + rest_cfg.chunk_size = chunk_size; + rest_cfg.max_n_chunks = max_n_chunks; auto rest_ctx = std::make_shared( std::move(rest_cfg), std::move(authorizer), &_host_mr); From 7b5e5142a0b61f0fd87aadd90fa2bc17e9d62d33 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 30 Jul 2026 19:20:22 +0000 Subject: [PATCH 4/5] add bindings for host_read_ranges_async_io --- include/cucascade/cudf/datasource.hpp | 18 ++++++++++++++ python/cucascade/datasource.pxd | 7 ++++++ python/cucascade/datasource.pyi | 3 +++ python/cucascade/datasource.pyx | 34 +++++++++++++++++++++++++++ src/cudf/datasource.cpp | 10 ++++++++ 5 files changed, 72 insertions(+) diff --git a/include/cucascade/cudf/datasource.hpp b/include/cucascade/cudf/datasource.hpp index 9ec79a8..2a8fb90 100644 --- a/include/cucascade/cudf/datasource.hpp +++ b/include/cucascade/cudf/datasource.hpp @@ -20,11 +20,13 @@ #include #include +#include #include #include #include +#include namespace cucascade::io { @@ -138,6 +140,22 @@ class datasource : public cudf::io::datasource { /// @c disposable call at consume time. void fadvise(std::span ranges, std::optional dev_id); + /** + * @brief Submit all byte ranges as a single vectorized host read. + * + * Uses the ioctx's scatter-read backend to fetch all segments in as few + * HTTP requests as possible, writing each range directly into the + * caller-supplied destination pointer in its @p segment. + * + * @param segments File offsets, sizes, and destination buffers. + * @return A future that resolves when every segment has been written. + */ + [[nodiscard]] std::future host_read_ranges_async( + std::span segments); + + [[nodiscard]] std::future host_read_ranges_async( + std::vector& segments); + void prefetch(cache::prefetching_stage site); private: diff --git a/python/cucascade/datasource.pxd b/python/cucascade/datasource.pxd index c66dc78..8e1f31a 100644 --- a/python/cucascade/datasource.pxd +++ b/python/cucascade/datasource.pxd @@ -8,6 +8,7 @@ from libcpp.memory cimport unique_ptr from libcpp.string cimport string from libcpp.vector cimport vector + from pylibcudf.libcudf.io.datasource cimport datasource as cudf_datasource @@ -18,11 +19,17 @@ cdef extern from "cudf/io/text/byte_range_info.hpp" namespace "cudf::io::text" n int64_t size() const +cdef extern from "cucascade/io/types.hpp" namespace "cucascade::io" nogil: + cdef cppclass io_object_segment: + io_object_segment(size_t offset, size_t size, uint8_t* buffer) except + + + cdef extern from "cucascade/cudf/datasource.hpp" namespace "cucascade::io" nogil: cdef cppclass cc_datasource "cucascade::io::datasource"(cudf_datasource): unique_ptr[cc_datasource] duplicate() except + void fadvise(const vector[byte_range_info]& ranges, int dev_id) except + future[size_t] host_read_async(size_t offset, size_t size, uint8_t* dst) except + + future[size_t] host_read_ranges_async(vector[io_object_segment]& segments) except + cdef extern from "cucascade/cudf/uring_datasource_engine.hpp" namespace "cucascade::io" nogil: diff --git a/python/cucascade/datasource.pyi b/python/cucascade/datasource.pyi index 3ca22a7..bbe31cf 100644 --- a/python/cucascade/datasource.pyi +++ b/python/cucascade/datasource.pyi @@ -16,6 +16,9 @@ class CuCascadeDatasource(Datasource): def read_ranges_async( self, ranges: list[tuple[int, int]], buffer: memoryview ) -> list[ReadFuture]: ... + def read_all_ranges_async( + self, ranges: list[tuple[int, int]], buffer: memoryview + ) -> ReadFuture: ... class UringEngine: diff --git a/python/cucascade/datasource.pyx b/python/cucascade/datasource.pyx index be1531b..e5db248 100644 --- a/python/cucascade/datasource.pyx +++ b/python/cucascade/datasource.pyx @@ -23,6 +23,7 @@ from pylibcudf.libcudf.io.datasource cimport datasource as cudf_datasource from cucascade.datasource cimport ( byte_range_info, cc_datasource, + io_object_segment, rest_datasource_engine, uring_datasource_engine, ) @@ -148,6 +149,39 @@ cdef class CuCascadeDatasource(Datasource): dst_offset += sz return futures + def read_all_ranges_async(self, list ranges, object buffer): + """Submit all byte ranges as a single vectorized host read. + + Uses the engine's scatter-read backend to fetch all ranges in as few + HTTP requests as possible, writing each range into ``buffer`` at the + corresponding offset. Returns a single future that resolves when every + range has been written. + + Parameters + ---------- + ranges : list[tuple[int, int]] + Byte ranges as ``(offset, size)`` pairs in file order. + buffer : memoryview + Contiguous writable host buffer sized to hold the sum of all range + sizes. + + Returns + ------- + ReadFuture + A single future that resolves when all ranges have been written. + """ + cdef uint8_t[::1] c_buf = buffer + cdef uint8_t* base = &c_buf[0] + cdef vector[io_object_segment] segments + cdef int64_t off, sz + cdef size_t dst_offset = 0 + for off, sz in ranges: + segments.emplace_back(off, sz, base + dst_offset) + dst_offset += sz + cdef ReadFuture rf = ReadFuture.__new__(ReadFuture) + rf._fut = deref(self._ds).host_read_ranges_async(segments) + return rf + cdef class UringEngine: """io_uring-backed datasource engine for local NVMe reads. diff --git a/src/cudf/datasource.cpp b/src/cudf/datasource.cpp index 76143c4..b6bd597 100644 --- a/src/cudf/datasource.cpp +++ b/src/cudf/datasource.cpp @@ -106,6 +106,16 @@ std::future datasource::host_read_async(size_t offset, size_t size, uint _io_ctx->host_read_async(*_io_object, offset, size, dst, &_prefetch_handle)); } +std::future datasource::host_read_ranges_async(std::span segments) +{ + return bridge_semi_to_std(_io_ctx->host_read_ranges_async_io(*_io_object, segments)); +} + +std::future datasource::host_read_ranges_async(std::vector& segments) +{ + return host_read_ranges_async(std::span{segments}); +} + std::future> datasource::host_read_async( size_t offset, size_t size) { From e62350f5f66e8f9afbe17e1ad3c798ae59fa4e48 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 31 Jul 2026 04:14:28 +0000 Subject: [PATCH 5/5] enable prefetch cache --- .../cucascade/cudf/rest_datasource_engine.hpp | 11 ++++-- python/cucascade/datasource.pxd | 3 +- python/cucascade/datasource.pyi | 1 + python/cucascade/datasource.pyx | 7 ++++ src/cudf/rest_datasource_engine.cpp | 38 +++++++++++++++++-- 5 files changed, 52 insertions(+), 8 deletions(-) diff --git a/include/cucascade/cudf/rest_datasource_engine.hpp b/include/cucascade/cudf/rest_datasource_engine.hpp index a16331c..d4fd7a6 100644 --- a/include/cucascade/cudf/rest_datasource_engine.hpp +++ b/include/cucascade/cudf/rest_datasource_engine.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -67,7 +68,8 @@ class rest_datasource_engine { std::size_t block_size = default_block_size, std::size_t max_connections = 16, std::size_t chunk_size = 8UL << 20, - std::size_t max_n_chunks = 16); + std::size_t max_n_chunks = 16, + bool enable_cache = false); ~rest_datasource_engine(); @@ -87,9 +89,10 @@ class rest_datasource_engine { [[nodiscard]] std::unique_ptr open(std::string path) const; private: - cucascade::memory::numa_region_pinned_host_memory_resource _upstream; - cucascade::memory::fixed_size_host_memory_resource _host_mr; - std::shared_ptr _io_ctx; + cucascade::memory::numa_region_pinned_host_memory_resource _upstream; + cucascade::memory::fixed_size_host_memory_resource _host_mr; + std::shared_ptr _io_ctx; + std::unique_ptr _reservation_manager; }; } // namespace cucascade::io diff --git a/python/cucascade/datasource.pxd b/python/cucascade/datasource.pxd index 8e1f31a..e2623b9 100644 --- a/python/cucascade/datasource.pxd +++ b/python/cucascade/datasource.pxd @@ -55,5 +55,6 @@ cdef extern from "cucascade/cudf/rest_datasource_engine.hpp" namespace "cucascad size_t block_size, size_t max_connections, size_t chunk_size, - size_t max_n_chunks) except + + size_t max_n_chunks, + cpp_bool enable_cache) except + unique_ptr[cc_datasource] open(string path) except + diff --git a/python/cucascade/datasource.pyi b/python/cucascade/datasource.pyi index bbe31cf..d5a293b 100644 --- a/python/cucascade/datasource.pyi +++ b/python/cucascade/datasource.pyi @@ -48,5 +48,6 @@ class RestEngine: max_connections: int = 16, chunk_size: int = 8388608, max_n_chunks: int = 16, + enable_cache: bool = False, ) -> None: ... def open(self, path: str) -> CuCascadeDatasource: ... diff --git a/python/cucascade/datasource.pyx b/python/cucascade/datasource.pyx index e5db248..fa6323b 100644 --- a/python/cucascade/datasource.pyx +++ b/python/cucascade/datasource.pyx @@ -291,6 +291,11 @@ cdef class RestEngine: max_n_chunks : int, optional Maximum number of destination buffers fused into a single scatter GET. Default is 16. + enable_cache : bool, optional + Whether to enable cuCascade's internal prefetch cache. When ``True``, + :meth:`CuCascadeDatasource.fadvise` queues S3 downloads into cuCascade's + bounce buffer pool so that subsequent reads may be served from cache + rather than S3. Default is ``False``. """ cdef unique_ptr[rest_datasource_engine] _engine @@ -309,6 +314,7 @@ cdef class RestEngine: size_t max_connections=16, size_t chunk_size=8388608, size_t max_n_chunks=16, + bint enable_cache=False, ): cdef string c_access_key_id = access_key_id.encode() cdef string c_secret_access_key = secret_access_key.encode() @@ -330,6 +336,7 @@ cdef class RestEngine: max_connections, chunk_size, max_n_chunks, + enable_cache, ) ) diff --git a/src/cudf/rest_datasource_engine.cpp b/src/cudf/rest_datasource_engine.cpp index 36328a8..d30d2bb 100644 --- a/src/cudf/rest_datasource_engine.cpp +++ b/src/cudf/rest_datasource_engine.cpp @@ -4,10 +4,14 @@ */ #include +#include #include #include #include #include +#include +#include +#include namespace cucascade::io { @@ -22,7 +26,8 @@ rest_datasource_engine::rest_datasource_engine(std::string access_key_id, std::size_t block_size, std::size_t max_connections, std::size_t chunk_size, - std::size_t max_n_chunks) + std::size_t max_n_chunks, + bool enable_cache) : _upstream(0, true), _host_mr(0, _upstream, pool_capacity, pool_capacity, block_size, 128, 1) { @@ -46,10 +51,37 @@ rest_datasource_engine::rest_datasource_engine(std::string access_key_id, auto io_ctx = std::make_shared(n_reactors, std::move(rest_ctx)); io_ctx->start(); - _io_ctx = std::move(io_ctx); + _io_ctx = io_ctx; + + if (enable_cache) { + memory::topology_discovery discovery; + static_cast(discovery.discover()); + auto const& topology = discovery.get_topology(); + + auto configs = memory::reservation_manager_configurator{} + .set_number_of_gpus(1) + .use_host_per_numa() + .set_total_host_capacity(pool_capacity) + .build(topology); + + _reservation_manager = + std::make_unique(std::move(configs)); + + io::cache::config cache_cfg{}; + cache_cfg.dispose_after_use = true; + + auto topo_index = + std::make_shared(topology, std::vector{0}); + + _io_ctx->initialize_cache(*_reservation_manager, cache_cfg, std::move(topo_index)); + } } -rest_datasource_engine::~rest_datasource_engine() { _io_ctx->shutdown(); } +rest_datasource_engine::~rest_datasource_engine() +{ + _io_ctx->pre_destroy(); + _io_ctx->shutdown(); +} std::unique_ptr rest_datasource_engine::open(std::string path) const {